Propagate reduced_state setting when loading the target experiment, to how the auxiliary ones are loaded (facebook#3597)

Summary:
Pull Request resolved: facebook#3597

This not being a propagated was a bug.

Reviewed By: Cesar-Cardoso

Differential Revision: D72266152

fbshipit-source-id: a6f2937555db9cd34f627e450b0c469db054099b

Author: Lena Kashtelyan

ax/storage/sqa_store/decoder.py

@@ -103,7 +103,7 @@ def get_enum_name(
             raise SQADecodeError(f"Value {value} is invalid for enum {enum}.")
 
     def _auxiliary_experiments_by_purpose_from_experiment_sqa(
-        self, experiment_sqa: SQAExperiment
+        self, experiment_sqa: SQAExperiment, reduced_state: bool = False
     ) -> dict[AuxiliaryExperimentPurpose, list[AuxiliaryExperiment]] | None:
         auxiliary_experiments_by_purpose = None
         if experiment_sqa.auxiliary_experiments_by_purpose:
@@ -122,7 +122,9 @@ def _auxiliary_experiments_by_purpose_from_experiment_sqa(
                     if isinstance(aux_exp_json, str):
                         aux_exp_json = {"experiment_name": aux_exp_json}
                     aux_experiment = auxiliary_experiment_from_json(
-                        json=aux_exp_json, config=self.config
+                        json=aux_exp_json,
+                        config=self.config,
+                        reduced_state=reduced_state,
                     )
                     auxiliary_experiments_by_purpose[aux_exp_purpose].append(
                         aux_experiment
@@ -133,6 +135,7 @@ def _init_experiment_from_sqa(
         self,
         experiment_sqa: SQAExperiment,
         load_auxiliary_experiments: bool = True,
+        reduced_state: bool = False,
     ) -> Experiment:
         """First step of conversion within experiment_from_sqa."""
         opt_config, tracking_metrics = self.opt_config_and_tracking_metrics_from_sqa(
@@ -170,7 +173,8 @@ def _init_experiment_from_sqa(
         auxiliary_experiments_by_purpose = (
             (
                 self._auxiliary_experiments_by_purpose_from_experiment_sqa(
-                    experiment_sqa=experiment_sqa
+                    experiment_sqa=experiment_sqa,
+                    reduced_state=reduced_state,
                 )
             )
             if load_auxiliary_experiments
@@ -281,6 +285,7 @@ def experiment_from_sqa(
             experiment = self._init_experiment_from_sqa(
                 experiment_sqa,
                 load_auxiliary_experiments=load_auxiliary_experiments,
+                reduced_state=reduced_state,
             )
         trials = [
             self.trial_from_sqa(
@@ -1333,7 +1338,9 @@ def _get_scalarized_outcome_constraint_children_metrics(
 
 
 def auxiliary_experiment_from_json(
-    json: dict[str, Any], config: SQAConfig
+    json: dict[str, Any],
+    config: SQAConfig,
+    reduced_state: bool = False,
 ) -> AuxiliaryExperiment:
     """
     Load an ``AuxiliaryExperiment`` from JSON.
@@ -1353,5 +1360,6 @@ def auxiliary_experiment_from_json(
         config=config,
         skip_runners_and_metrics=True,
         load_auxiliary_experiments=False,
+        reduced_state=reduced_state,
     )
     return AuxiliaryExperiment(experiment)

ax/storage/sqa_store/tests/test_sqa_store.py

@@ -30,6 +30,7 @@
 from ax.core.outcome_constraint import OutcomeConstraint
 from ax.core.parameter import ParameterType, RangeParameter
 from ax.core.runner import Runner
+from ax.core.trial import Trial
 from ax.core.trial_status import TrialStatus
 from ax.core.types import ComparisonOp
 from ax.exceptions.core import ObjectNotFoundError
@@ -275,6 +276,8 @@ def test_saving_and_loading_experiment_with_aux_exp(self) -> None:
             is_test=True,
         )
         save_experiment(aux_experiment, config=self.config)
+        # pyre-ignore[16]: `AuxiliaryExperimentPurpose` has no attribute
+        purpose = self.config.auxiliary_experiment_purpose_enum.MyAuxExpPurpose
 
         experiment_w_aux_exp = Experiment(
             name="test_experiment_w_aux_exp_in_SQAStoreTest",
@@ -284,10 +287,7 @@ def test_saving_and_loading_experiment_with_aux_exp(self) -> None:
             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)
-                ]
+                purpose: [AuxiliaryExperiment(experiment=aux_experiment)]
             },
         )
         self.assertIsNone(experiment_w_aux_exp.db_id)
@@ -299,6 +299,66 @@ def test_saving_and_loading_experiment_with_aux_exp(self) -> None:
         self.assertEqual(experiment_w_aux_exp, loaded_experiment)
         self.assertEqual(len(loaded_experiment.auxiliary_experiments_by_purpose), 1)
 
+    def test_saving_and_loading_experiment_with_aux_exp_reduced_state(self) -> None:
+        aux_exp = Experiment(
+            name="test_aux_exp_in_SQAStoreTest_reduced_state",
+            search_space=get_search_space(),
+            optimization_config=get_optimization_config(),
+            description="test description",
+            tracking_metrics=[Metric(name="tracking")],
+            is_test=True,
+        )
+        aux_exp_gs = get_generation_strategy(with_callable_model_kwarg=False)
+        aux_exp.new_trial(aux_exp_gs.gen(experiment=aux_exp))
+        save_experiment(aux_exp, config=self.config)
+        # pyre-ignore[16]: `AuxiliaryExperimentPurpose` has no attribute
+        purpose = self.config.auxiliary_experiment_purpose_enum.MyAuxExpPurpose
+
+        target_exp = Experiment(
+            name="test_experiment_w_aux_exp_in_SQAStoreTest_reduced_state",
+            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={
+                purpose: [AuxiliaryExperiment(experiment=aux_exp)]
+            },
+        )
+        target_exp_gs = get_generation_strategy(with_callable_model_kwarg=False)
+        target_exp.new_trial(target_exp_gs.gen(experiment=target_exp))
+        self.assertIsNone(target_exp.db_id)
+        save_experiment(target_exp, config=self.config)
+        self.assertIsNotNone(target_exp.db_id)
+        loaded_target_exp = load_experiment(
+            target_exp.name, config=self.config, reduced_state=True
+        )
+        self.assertNotEqual(target_exp, loaded_target_exp)
+        self.assertIsNotNone(  # State of the original aux experiment is not reduced.
+            none_throws(
+                assert_is_instance(aux_exp.trials[0], Trial).generator_run
+            ).gen_metadata
+        )
+        self.assertIsNotNone(  # State of the original target experiment is not reduced.
+            none_throws(
+                assert_is_instance(target_exp.trials[0], Trial).generator_run
+            ).gen_metadata
+        )
+        self.assertIsNone(  # State of the loaded target experiment *is reduced*.
+            none_throws(
+                assert_is_instance(loaded_target_exp.trials[0], Trial).generator_run
+            ).gen_metadata
+        )
+        loaded_aux_exp = loaded_target_exp.auxiliary_experiments_by_purpose[purpose][0]
+        self.assertIsNone(  # State of the loaded target experiment *is reduced*.
+            none_throws(
+                assert_is_instance(
+                    loaded_aux_exp.experiment.trials[0], Trial
+                ).generator_run
+            ).gen_metadata
+        )
+        self.assertEqual(len(loaded_target_exp.auxiliary_experiments_by_purpose), 1)
+
     def test_saving_with_aux_exp_not_in_db(self) -> None:
         aux_experiment = Experiment(
             name="aux_experiment_not_in_db", search_space=get_search_space()