Deprecate GenerationStrategy.trials_as_df (#3275)

Summary:
Pull Request resolved: #3275

Deprecates `GS.trials_as_df` in favor of the more informative `Experiment.to_df`.

Reviewed By: lena-kashtelyan

Differential Revision: D68629386

fbshipit-source-id: d07e8dcc06f7469548858ac1f1ed76c22747ce4c

ax/modelbridge/generation_strategy.py

@@ -8,8 +8,8 @@
 
 from __future__ import annotations
 
+import warnings
 from collections.abc import Callable
-
 from copy import deepcopy
 from functools import wraps
 from logging import Logger
@@ -33,7 +33,7 @@
 from ax.modelbridge.generation_node_input_constructors import InputConstructorPurpose
 from ax.modelbridge.model_spec import FactoryFunctionModelSpec
 from ax.modelbridge.transition_criterion import TrialBasedCriterion
-from ax.utils.common.logger import _round_floats_for_logging, get_logger
+from ax.utils.common.logger import get_logger
 from ax.utils.common.typeutils import assert_is_instance_list
 from pyre_extensions import none_throws
 
@@ -277,54 +277,18 @@ def trials_as_df(self) -> pd.DataFrame | None:
         """Puts information on individual trials into a data frame for easy
         viewing.
 
-        For example for a GenerationStrategy composed of GenerationSteps:
-        Gen. Step | Models   | Trial Index | Trial Status | Arm Parameterizations
-        [0]       | [Sobol]  | 0           | RUNNING      | {"0_0":{"x":9.17...}}
+        THIS METHOD IS DEPRECATED AND WILL BE REMOVED IN A FUTURE RELEASE.
+        Please use `Experiment.to_df()` instead.
         """
-        logger.info(
-            "Note that parameter values in dataframe are rounded to 2 decimal "
-            "points; the values in the dataframe are thus not the exact ones "
-            "suggested by Ax in trials."
+        warnings.warn(
+            "`GenerationStrategy.trials_as_df` is deprecated and will be removed in "
+            "a future release. Please use `Experiment.to_df()` instead.",
+            DeprecationWarning,
+            stacklevel=2,
         )
-        if self._experiment is None or len(self.experiment.trials) == 0:
+        if self._experiment is None:
             return None
-
-        step_or_node_col = (
-            "Generation Nodes" if self.is_node_based else "Generation Step"
-        )
-        records = [
-            {
-                step_or_node_col: [
-                    (
-                        gr._generation_node_name
-                        if gr.generator_run_type != "MANUAL"
-                        else "MANUAL"
-                    )
-                    for gr in trial.generator_runs
-                ],
-                "Generation Model(s)": [
-                    (gr._model_key if gr.generator_run_type != "MANUAL" else "MANUAL")
-                    for gr in trial.generator_runs
-                ],
-                "Trial Index": trial_idx,
-                "Trial Status": trial.status.name,
-                "Arm Parameterizations": {
-                    arm.name: _round_floats_for_logging(arm.parameters)
-                    for arm in trial.arms
-                },
-            }
-            for trial_idx, trial in self.experiment.trials.items()
-        ]
-
-        return pd.DataFrame.from_records(records).reindex(
-            columns=[
-                step_or_node_col,
-                "Generation Model(s)",
-                "Trial Index",
-                "Trial Status",
-                "Arm Parameterizations",
-            ]
-        )
+        return self.experiment.to_df()
 
     @property
     def optimization_complete(self) -> bool:

ax/modelbridge/tests/test_generation_strategy.py

@@ -743,24 +743,14 @@ def test_trials_as_df(self) -> None:
                 GenerationStep(model=Models.SOBOL, num_trials=3),
             ]
         )
-        # No trials yet, so the DF will be None.
-        self.assertIsNone(sobol_generation_strategy.trials_as_df)
-        # Now the trial should appear in the DF.
-        trial = exp.new_trial(sobol_generation_strategy.gen(experiment=exp))
-        trials_df = none_throws(sobol_generation_strategy.trials_as_df)
-        self.assertFalse(trials_df.empty)
-        self.assertEqual(trials_df.head()["Trial Status"][0], "CANDIDATE")
-        # Changes in trial status should be reflected in the DF.
-        trial._status = TrialStatus.RUNNING
-        trials_df = none_throws(sobol_generation_strategy.trials_as_df)
-        self.assertEqual(trials_df.head()["Trial Status"][0], "RUNNING")
-        # Check that rows are present for step 0 and 1 after moving to step 1
-        for _i in range(3):
-            # attach necessary trials to fill up the Generation Strategy
-            trial = exp.new_trial(sobol_generation_strategy.gen(experiment=exp))
-        trials_df = none_throws(sobol_generation_strategy.trials_as_df)
-        self.assertEqual(trials_df.head()["Generation Step"][0], ["GenerationStep_0"])
-        self.assertEqual(trials_df.head()["Generation Step"][2], ["GenerationStep_1"])
+        # No experiment attached to the GS, should be None.
+        with self.assertWarnsRegex(DeprecationWarning, "trials_as_df"):
+            self.assertIsNone(sobol_generation_strategy.trials_as_df)
+        # Experiment attached with a trial, should match Experiment.to_df().
+        exp.new_trial(sobol_generation_strategy.gen(experiment=exp))
+        with self.assertWarnsRegex(DeprecationWarning, "trials_as_df"):
+            trials_df = none_throws(sobol_generation_strategy.trials_as_df)
+        self.assertTrue(trials_df.equals(exp.to_df()))
 
     def test_max_parallelism_reached(self) -> None:
         exp = get_branin_experiment()
@@ -1847,41 +1837,6 @@ def test_node_gs_with_auto_transitions_three_phase(self) -> None:
         trial = exp.new_batch_trial(generator_runs=gs_2._gen_with_multiple_nodes(exp))
         self.assertEqual(trial.generator_runs[0]._generation_node_name, "sobol_4")
 
-    def test_trials_as_df_node_gs(self) -> None:
-        exp = get_branin_experiment()
-        gs = self.complex_multinode_per_trial_gs
-        arms_per_node = {
-            "sobol": 3,
-            "mbm": 1,
-            "sobol_2": 2,
-            "sobol_3": 3,
-            "sobol_4": 4,
-        }
-        gs.experiment = exp
-        self.assertIsNone(gs.trials_as_df)
-        # Now the trial should appear in the DF.
-        trial = exp.new_batch_trial(
-            generator_runs=gs._gen_with_multiple_nodes(exp, arms_per_node=arms_per_node)
-        )
-        trials_df = none_throws(gs.trials_as_df)
-        self.assertFalse(trials_df.empty)
-        self.assertEqual(trials_df.head()["Trial Status"][0], "CANDIDATE")
-        self.assertEqual(trials_df.head()["Generation Model(s)"][0], ["Sobol"])
-        # Changes in trial status should be reflected in the DF.
-        trial.run()
-        self.assertEqual(
-            none_throws(gs.trials_as_df).head()["Trial Status"][0], "RUNNING"
-        )
-        # Add a new trial which will be generated from multiple nodes, and check that
-        # is properly reflected in the DF.
-        trial = exp.new_batch_trial(
-            generator_runs=gs._gen_with_multiple_nodes(exp, arms_per_node=arms_per_node)
-        )
-        self.assertEqual(
-            none_throws(gs.trials_as_df).head()["Generation Nodes"][1],
-            ["mbm", "sobol_2", "sobol_3"],
-        )
-
     def test_gs_with_fixed_features_constructor(self) -> None:
         exp = get_branin_experiment()
         sobol_criterion = [