Update trials_as_df method to support GenNode based GS (#2095)

Summary:
Pull Request resolved: #2095

This diff modifies the trials as df method to support representation of GenerationNodes as well as GenerationSteps in GSs

In coming diffs:
(5) Do a final pass of the generationStrategy/GenerationNode files to see what else can be migrated/condensed
(6) rename transiton criterion to action criterion
(7) remove conditionals for legacy usecase
( clean up any lingering todos

Reviewed By: bernardbeckerman

Differential Revision: D52273582

fbshipit-source-id: 740ef0e17456e1a2a41d4472ae6432552b4b7977

ax/modelbridge/generation_strategy.py

@@ -276,10 +276,11 @@ def uses_non_registered_models(self) -> bool:
         return not self._uses_registered_models
 
     @property
-    @step_based_gs_only
     def trials_as_df(self) -> Optional[pd.DataFrame]:
         """Puts information on individual trials into a data frame for easy
-        viewing. For example:
+        viewing.
+
+        For example for a GenerationStrategy composed of GenerationSteps:
         Gen. Step | Model | Trial Index | Trial Status | Arm Parameterizations
         0         | Sobol | 0           | RUNNING      | {"0_0":{"x":9.17...}}
         """
@@ -293,11 +294,15 @@ def trials_as_df(self) -> Optional[pd.DataFrame]:
             len(step.trials_from_node) == 0 for step in self._nodes
         ):
             return None
+
+        step_or_node_col = (
+            "Generation Node" if self.is_node_based else "Generation Step"
+        )
         records = [
             {
-                "Generation Step": step.node_name,
+                step_or_node_col: node.node_name,
                 "Generation Model": self._nodes[
-                    step_idx
+                    node_idx
                 ].model_spec_to_gen_from.model_key,
                 "Trial Index": trial_idx,
                 "Trial Status": self.experiment.trials[trial_idx].status.name,
@@ -306,12 +311,13 @@ def trials_as_df(self) -> Optional[pd.DataFrame]:
                     for arm in self.experiment.trials[trial_idx].arms
                 },
             }
-            for step_idx, step in enumerate(self._nodes)
-            for trial_idx in step.trials_from_node
+            for node_idx, node in enumerate(self._nodes)
+            for trial_idx in node.trials_from_node
         ]
+
         return pd.DataFrame.from_records(records).reindex(
             columns=[
-                "Generation Step",
+                step_or_node_col,
                 "Generation Model",
                 "Trial Index",
                 "Trial Status",

ax/modelbridge/tests/test_generation_strategy.py

@@ -566,6 +566,69 @@ def test_trials_as_df(self) -> None:
             "GenerationStep_1",
         )
 
+        # construct the same GS as above but directly with nodes
+        sobol_model_spec = ModelSpec(
+            model_enum=Models.SOBOL,
+            model_kwargs={},
+            model_gen_kwargs={},
+        )
+        node_gs = GenerationStrategy(
+            nodes=[
+                GenerationNode(
+                    node_name="sobol_2_trial",
+                    model_specs=[sobol_model_spec],
+                    transition_criteria=[
+                        MaxTrials(
+                            threshold=2,
+                            not_in_statuses=[TrialStatus.FAILED, TrialStatus.ABANDONED],
+                            block_gen_if_met=True,
+                            block_transition_if_unmet=True,
+                            transition_to="sobol_3_trial",
+                        )
+                    ],
+                    gen_unlimited_trials=False,
+                ),
+                GenerationNode(
+                    node_name="sobol_3_trial",
+                    model_specs=[sobol_model_spec],
+                    transition_criteria=[
+                        MaxTrials(
+                            threshold=2,
+                            not_in_statuses=[TrialStatus.FAILED, TrialStatus.ABANDONED],
+                            block_gen_if_met=True,
+                            block_transition_if_unmet=True,
+                            transition_to=None,
+                        )
+                    ],
+                    gen_unlimited_trials=False,
+                ),
+            ]
+        )
+        self.assertIsNone(node_gs.trials_as_df)
+        # Now the trial should appear in the DF.
+        trial = exp.new_trial(node_gs.gen(experiment=exp))
+
+        self.assertFalse(node_gs.trials_as_df.empty)
+        self.assertEqual(
+            node_gs.trials_as_df.head()["Trial Status"][0],
+            "CANDIDATE",
+        )
+        # Changes in trial status should be reflected in the DF.
+        trial._status = TrialStatus.RUNNING
+        self.assertEqual(node_gs.trials_as_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(node_gs.gen(experiment=exp))
+        self.assertEqual(
+            node_gs.trials_as_df.head()["Generation Node"][0],
+            "sobol_2_trial",
+        )
+        self.assertEqual(
+            node_gs.trials_as_df.head()["Generation Node"][2],
+            "sobol_3_trial",
+        )
+
     def test_max_parallelism_reached(self) -> None:
         exp = get_branin_experiment()
         sobol_generation_strategy = GenerationStrategy(