Method to consolidate Experiment.status from generator runs (#4900)

Summary:

Add a new static method `experiment_status_from_generator_runs()` to `GenerationStrategy` that extracts and validates a suggested ExperimentStatus from a list of GeneratorRun objects. It collects all unique suggested_experiment_status values from the runs and:

- Returns None with a warning if there are conflicting statuses across runs
- Returns None with an info log if no statuses are found
- Returns the single agreed-upon status otherwise

Differential Revision: D92985915

Author: CristianLara

ax/generation_strategy/generation_strategy.py

@@ -16,6 +16,7 @@
 from ax.adapter.base import Adapter
 from ax.core.data import Data
 from ax.core.experiment import Experiment
+from ax.core.experiment_status import ExperimentStatus
 from ax.core.generator_run import GeneratorRun
 from ax.core.observation import ObservationFeatures
 from ax.core.utils import extend_pending_observations, extract_pending_observations
@@ -312,6 +313,43 @@ def gen(
             )
         return grs_for_multiple_trials
 
+    @staticmethod
+    def experiment_status_from_generator_runs(
+        generator_runs: list[GeneratorRun],
+    ) -> ExperimentStatus | None:
+        """Extract and validate suggested experiment status from generator runs.
+
+        Collects the suggested_experiment_status directly from the GeneratorRun
+        objects, validates that all runs suggest the same status, and returns
+        that status.
+
+        Args:
+            generator_runs: List of generator runs to extract statuses from.
+
+        Returns:
+            The suggested experiment status that all generator runs agree on,
+            or None if no statuses were found or if there are conflicting statuses.
+        """
+        suggested_statuses: set[ExperimentStatus] = set()
+        for gr in generator_runs:
+            if gr.suggested_experiment_status is not None:
+                suggested_statuses.add(gr.suggested_experiment_status)
+
+        if len(suggested_statuses) > 1:
+            logger.warning(
+                "Multiple different suggested experiment statuses found: "
+                f"{suggested_statuses}. "
+                "All generator runs used in a single gen() call should suggest the "
+                "same experiment status. Skipping updating experiment status."
+            )
+            return None
+
+        if len(suggested_statuses) == 0:
+            logger.info("No suggested_experiment_status found on any generator runs.")
+            return None
+
+        return suggested_statuses.pop()
+
     def current_generator_run_limit(
         self,
     ) -> tuple[int, bool]:

ax/generation_strategy/tests/test_generation_strategy.py

@@ -24,6 +24,7 @@
 from ax.adapter.torch import TorchAdapter
 from ax.core.arm import Arm
 from ax.core.experiment import Experiment
+from ax.core.experiment_status import ExperimentStatus
 from ax.core.generator_run import GeneratorRun
 from ax.core.observation import ObservationFeatures
 from ax.core.parameter import ChoiceParameter, FixedParameter, Parameter, ParameterType
@@ -2022,6 +2023,74 @@ def test_optimization_complete_single_node_no_criteria(self) -> None:
 
         self.assertFalse(gs.optimization_complete)
 
+    def test_experiment_status_from_generation_strategy(self) -> None:
+        """Test that experiment status is correctly propagated through
+        generator runs and extracted via experiment_status_from_generator_runs."""
+
+        with self.subTest("gen returns GRs with correct suggested_experiment_status"):
+            for status in [
+                ExperimentStatus.INITIALIZATION,
+                ExperimentStatus.OPTIMIZATION,
+            ]:
+                with self.subTest(status=status):
+                    exp = get_branin_experiment()
+                    node_with_status = GenerationNode(
+                        name="test_node",
+                        generator_specs=[self.sobol_generator_spec],
+                        suggested_experiment_status=status,
+                    )
+                    gs = GenerationStrategy(nodes=[node_with_status])
+                    gs.experiment = exp
+
+                    grs = gs.gen(experiment=exp, num_trials=1)
+                    flat_grs = [gr for trial_grs in grs for gr in trial_grs]
+
+                    extracted_status = (
+                        GenerationStrategy.experiment_status_from_generator_runs(
+                            flat_grs
+                        )
+                    )
+                    self.assertEqual(extracted_status, status)
+
+        with self.subTest("conflicting statuses return None"):
+            gr1 = GeneratorRun(
+                arms=[Arm(name="0_0", parameters={"x1": 0.0, "x2": 0.0})],
+                suggested_experiment_status=ExperimentStatus.INITIALIZATION,
+            )
+            gr2 = GeneratorRun(
+                arms=[Arm(name="0_1", parameters={"x1": 1.0, "x2": 1.0})],
+                suggested_experiment_status=ExperimentStatus.OPTIMIZATION,
+            )
+            mixed_grs = [gr1, gr2]
+
+            result = GenerationStrategy.experiment_status_from_generator_runs(mixed_grs)
+            self.assertIsNone(result)
+
+        with self.subTest("multiple trials all carry experiment status"):
+            exp = get_branin_experiment()
+            node_with_status = GenerationNode(
+                name="multi_trial_node",
+                generator_specs=[self.sobol_generator_spec],
+                suggested_experiment_status=ExperimentStatus.INITIALIZATION,
+            )
+            gs = GenerationStrategy(nodes=[node_with_status])
+            gs.experiment = exp
+
+            grs = gs.gen(experiment=exp, num_trials=3)
+
+            self.assertEqual(len(grs), 3)
+            for gr_list in grs:
+                self.assertEqual(len(gr_list), 1)
+                self.assertEqual(gr_list[0]._generation_node_name, "multi_trial_node")
+                self.assertEqual(
+                    gr_list[0].suggested_experiment_status,
+                    ExperimentStatus.INITIALIZATION,
+                )
+            extracted_status = GenerationStrategy.experiment_status_from_generator_runs(
+                [gr for trial_grs in grs for gr in trial_grs]
+            )
+            self.assertEqual(extracted_status, ExperimentStatus.INITIALIZATION)
+
     # ------------- Testing helpers (put tests above this line) -------------
 
     def _run_GS_for_N_rounds(