add auxiliary experiments to SQAExperiment (#2658)

Summary:
Pull Request resolved: #2658

SQA storage change for auxiliary_expriments in experiments

Reviewed By: mgarrard, Cesar-Cardoso

Differential Revision: D60927108

ax/storage/sqa_store/decoder.py

@@ -17,6 +17,7 @@
 from ax.analysis.analysis import AnalysisCard, AnalysisCardLevel
 
 from ax.core.arm import Arm
+from ax.core.auxiliary import AuxiliaryExperiment, AuxiliaryExperimentPurpose
 from ax.core.base_trial import BaseTrial, TrialStatus
 from ax.core.batch_trial import AbandonedArm, BatchTrial, GeneratorRunStruct
 from ax.core.data import Data
@@ -103,6 +104,32 @@ def get_enum_name(
         except ValueError:
             raise SQADecodeError(f"Value {value} is invalid for enum {enum}.")
 
+    def _auxiliary_experiments_by_purpose_from_experiment_sqa(
+        self, experiment_sqa: SQAExperiment
+    ) -> Optional[dict[AuxiliaryExperimentPurpose, list[AuxiliaryExperiment]]]:
+        auxiliary_experiments_by_purpose = None
+        if experiment_sqa.auxiliary_experiments_by_purpose:
+            from ax.storage.sqa_store.load import load_experiment
+
+            auxiliary_experiments_by_purpose = {}
+            aux_exp_name_dict = not_none(
+                experiment_sqa.auxiliary_experiments_by_purpose
+            )
+            for aux_exp_purpose_str, aux_exp_names in aux_exp_name_dict.items():
+                aux_exp_purpose = next(
+                    member
+                    for member in self.config.auxiliary_experiment_purpose_enum
+                    if member.value == aux_exp_purpose_str
+                )
+                auxiliary_experiments_by_purpose[aux_exp_purpose] = []
+                for aux_exp_name in aux_exp_names:
+                    auxiliary_experiments_by_purpose[aux_exp_purpose].append(
+                        AuxiliaryExperiment(
+                            experiment=load_experiment(aux_exp_name, config=self.config)
+                        )
+                    )
+        return auxiliary_experiments_by_purpose
+
     def _init_experiment_from_sqa(
         self,
         experiment_sqa: SQAExperiment,
@@ -147,6 +174,13 @@ def _init_experiment_from_sqa(
         # so need to convert it to regular dict.
         properties = dict(experiment_sqa.properties or {})
         default_data_type = experiment_sqa.default_data_type
+
+        auxiliary_experiments_by_purpose = (
+            self._auxiliary_experiments_by_purpose_from_experiment_sqa(
+                experiment_sqa=experiment_sqa
+            )
+        )
+
         return Experiment(
             name=experiment_sqa.name,
             description=experiment_sqa.description,
@@ -158,6 +192,7 @@ def _init_experiment_from_sqa(
             is_test=experiment_sqa.is_test,
             properties=properties,
             default_data_type=default_data_type,
+            auxiliary_experiments_by_purpose=auxiliary_experiments_by_purpose,
         )
 
     def _init_mt_experiment_from_sqa(

ax/storage/sqa_store/encoder.py

@@ -174,6 +174,15 @@ def experiment_to_sqa(self, experiment: Experiment) -> SQAExperiment:
             value=experiment.experiment_type, enum=self.config.experiment_type_enum
         )
 
+        auxiliary_experiments_by_purpose = {}
+        for (
+            aux_exp_type_enum,
+            aux_exps,
+        ) in experiment.auxiliary_experiments_by_purpose.items():
+            aux_exp_type = aux_exp_type_enum.value
+            aux_exp_jsons = [aux_exp.experiment.name for aux_exp in aux_exps]
+            auxiliary_experiments_by_purpose[aux_exp_type] = aux_exp_jsons
+
         properties = experiment._properties
         runners = []
         if isinstance(experiment, MultiTypeExperiment):
@@ -213,6 +222,7 @@ def experiment_to_sqa(self, experiment: Experiment) -> SQAExperiment:
             properties=properties,
             default_trial_type=experiment.default_trial_type,
             default_data_type=experiment.default_data_type,
+            auxiliary_experiments_by_purpose=auxiliary_experiments_by_purpose,
         )
         return exp_sqa
 

ax/storage/sqa_store/sqa_classes.py

@@ -49,7 +49,6 @@
 )
 from sqlalchemy.orm import backref, relationship
 
-
 ONLY_ONE_FIELDS = ["experiment_id", "generator_run_id"]
 
 
@@ -510,6 +509,12 @@ class SQAExperiment(Base):
     default_trial_type: Optional[str] = Column(String(NAME_OR_TYPE_FIELD_LENGTH))
     # pyre-fixme[8]: Attribute has type `DataType`; used as `Column[typing.Any]`.
     default_data_type: DataType = Column(IntEnum(DataType), nullable=True)
+    # pyre-fixme[8]: Incompatible attribute type [8]: Attribute
+    # `auxiliary_experiments_by_purpose` declared in class `SQAExperiment` has
+    # type `Optional[Dict[str, List[str]]]` but is used as type `Column[typing.Any]`
+    auxiliary_experiments_by_purpose: Optional[dict[str, list[str]]] = Column(
+        JSONEncodedTextDict, nullable=True, default={}
+    )
 
     # relationships
     # Trials and experiments are mutable, so the children relationships need

ax/storage/sqa_store/sqa_config.py

@@ -13,6 +13,7 @@
 from ax.analysis.analysis import AnalysisCard
 
 from ax.core.arm import Arm
+from ax.core.auxiliary import AuxiliaryExperimentPurpose
 from ax.core.batch_trial import AbandonedArm
 from ax.core.data import Data
 from ax.core.experiment import Experiment
@@ -86,6 +87,7 @@ def _default_class_to_sqa_class(self=None) -> dict[type[Base], type[SQABase]]:
     )
     experiment_type_enum: Optional[Union[Enum, type[Enum]]] = None
     generator_run_type_enum: Optional[Union[Enum, type[Enum]]] = GeneratorRunType
+    auxiliary_experiment_purpose_enum: type[Enum] = AuxiliaryExperimentPurpose
 
     # pyre-fixme[4]: Attribute annotation cannot contain `Any`.
     # pyre-fixme[24]: Generic type `type` expects 1 type parameter, use

ax/storage/sqa_store/tests/test_sqa_store.py

@@ -8,7 +8,7 @@
 
 import logging
 from datetime import datetime
-from enum import Enum
+from enum import Enum, unique
 from logging import Logger
 from typing import Any
 from unittest import mock
@@ -20,7 +20,9 @@
 from ax.analysis.markdown.markdown_analysis import MarkdownAnalysisCard
 from ax.analysis.plotly.plotly_analysis import PlotlyAnalysisCard
 from ax.core.arm import Arm
+from ax.core.auxiliary import AuxiliaryExperiment, AuxiliaryExperimentPurpose
 from ax.core.batch_trial import BatchTrial, LifecycleStage
+from ax.core.experiment import Experiment
 from ax.core.generator_run import GeneratorRun
 from ax.core.metric import Metric
 from ax.core.objective import MultiObjective, Objective
@@ -208,6 +210,46 @@ def test_ExperimentSaveAndLoad(self) -> None:
             loaded_experiment = load_experiment(exp.name)
             self.assertEqual(loaded_experiment, exp)
 
+    def test_saving_and_loading_experiment_with_aux_exp(self) -> None:
+        @unique
+        class TestAuxiliaryExperimentPurpose(AuxiliaryExperimentPurpose):
+            MyAuxExpPurpose = "my_auxiliary_experiment_purpose"
+
+        self.config.auxiliary_experiment_purpose_enum = TestAuxiliaryExperimentPurpose
+
+        aux_experiment = Experiment(
+            name="test_aux_exp_in_SQAStoreTest",
+            search_space=get_search_space(),
+            optimization_config=get_optimization_config(),
+            description="test description",
+            tracking_metrics=[Metric(name="tracking")],
+            is_test=True,
+        )
+        save_experiment(aux_experiment, config=self.config)
+
+        experiment_w_aux_exp = Experiment(
+            name="test_experiment_w_aux_exp_in_SQAStoreTest",
+            search_space=get_search_space(),
+            optimization_config=get_optimization_config(),
+            description="test description",
+            tracking_metrics=[Metric(name="tracking")],
+            is_test=True,
+            auxiliary_experiments_by_purpose={
+                # pyre-ignore[16]: `AuxiliaryExperimentPurpose` has no attribute
+                self.config.auxiliary_experiment_purpose_enum.MyAuxExpPurpose: [
+                    AuxiliaryExperiment(experiment=aux_experiment)
+                ]
+            },
+        )
+        self.assertIsNone(experiment_w_aux_exp.db_id)
+        save_experiment(experiment_w_aux_exp, config=self.config)
+        self.assertIsNotNone(experiment_w_aux_exp.db_id)
+        loaded_experiment = load_experiment(
+            experiment_w_aux_exp.name, config=self.config
+        )
+        self.assertEqual(experiment_w_aux_exp, loaded_experiment)
+        self.assertEqual(len(loaded_experiment.auxiliary_experiments_by_purpose), 1)
+
     def test_saving_an_experiment_with_type_requires_an_enum(self) -> None:
         self.experiment.experiment_type = "TEST"
         with self.assertRaises(SQAEncodeError):