Fix warm start from old experiment with mapdata (#839)

Summary: Pull Request resolved: #839 Reviewed By: bernardbeckerman Differential Revision: D34725796 fbshipit-source-id: f9bad4c1163530277176a53aab9be503fe83f1c2
facebook · Mar 9, 2022 · b4a1adc · b4a1adc
1 parent b0510ee
commit b4a1adc
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 11 deletions.
diff --git a/ax/core/experiment.py b/ax/core/experiment.py
@@ -1000,18 +1000,17 @@ def warm_start_from_old_experiment(
                     )
                     # Attach updated data to new trial on experiment and mark trial
                     # as completed.
-                    if self.default_data_type == DataType.MAP_DATA:
-                        map_key_infos = checked_cast(
-                            MapData, old_experiment.lookup_data()
-                        ).map_key_infos
-                        self.attach_data(
-                            data=self.default_data_constructor(
-                                df=new_df,
-                                map_key_infos=map_key_infos,
-                            )
+                    old_data = (
+                        old_experiment.default_data_constructor(
+                            df=new_df,
+                            map_key_infos=checked_cast(
+                                MapData, old_experiment.lookup_data()
+                            ).map_key_infos,
                         )
-                    else:
-                        self.attach_data(data=self.default_data_constructor(df=new_df))
+                        if old_experiment.default_data_type == DataType.MAP_DATA
+                        else old_experiment.default_data_constructor(df=new_df)
+                    )
+                    self.attach_data(data=old_data)
                     new_trial.mark_completed()
                 else:
                     new_trial.mark_abandoned(reason=trial.abandoned_reason)

diff --git a/ax/core/tests/test_experiment.py b/ax/core/tests/test_experiment.py
@@ -805,6 +805,15 @@ def testWarmStartFromOldExperiment(self):
         )
         self.assertEqual(len(new_experiment.trials), len(old_experiment.trials) - 2)
 
+        # Warm start from an experiment with only a subset of metrics
+        map_data_experiment = get_branin_experiment_with_timestamp_map_metric()
+        map_data_experiment.warm_start_from_old_experiment(
+            old_experiment=old_experiment
+        )
+        self.assertEqual(
+            len(map_data_experiment.trials), len(old_experiment.trials) - 1
+        )
+
     def test_is_test_warning(self):
         experiments_module = "ax.core.experiment"
         with self.subTest("it warns on construction for a test"):