Add suggested_experiment_status to legacy GS dispatch --> have AxSweep experiment status appear in the UI (#5056)

Summary:
Pull Request resolved: #5056

Add `suggested_experiment_status` to the `GenerationStep`s created by
`choose_generation_strategy_legacy`. The Sobol step gets
`ExperimentStatus.INITIALIZATION` and the BayesOpt step gets
`ExperimentStatus.OPTIMIZATION`, matching the convention used by
all other GS factory functions.

Reviewed By: CristianLara

Differential Revision: D96784737

fbshipit-source-id: 50a4205fc41b02556614d976db4e35594f069d7c

Author: Lena Kashtelyan

ax/generation_strategy/dispatch_utils.py

@@ -13,6 +13,7 @@
 import torch
 from ax.adapter.registry import GeneratorRegistryBase, Generators
 from ax.core.experiment import Experiment
+from ax.core.experiment_status import ExperimentStatus
 from ax.core.optimization_config import OptimizationConfig
 from ax.core.parameter import ChoiceParameter, ParameterType, RangeParameter
 from ax.core.search_space import SearchSpace
@@ -66,6 +67,7 @@ def _make_sobol_step(
         generator_kwargs={"deduplicate": True, "seed": seed},
         use_all_trials_in_exp=True,
         should_deduplicate=should_deduplicate,
+        suggested_experiment_status=ExperimentStatus.INITIALIZATION,
     )
 
 
@@ -133,6 +135,7 @@ def _make_botorch_step(
         max_parallelism=max_concurrency,
         generator_kwargs=generator_kwargs,
         should_deduplicate=should_deduplicate,
+        suggested_experiment_status=ExperimentStatus.OPTIMIZATION,
     )
 
 

ax/generation_strategy/tests/test_dispatch_utils.py

@@ -13,6 +13,7 @@
 import torch
 from ax.adapter.registry import Generators
 from ax.adapter.transforms.log_y import LogY
+from ax.core.experiment_status import ExperimentStatus
 from ax.core.objective import Objective
 from ax.core.optimization_config import MultiObjectiveOptimizationConfig
 from ax.generation_strategy.dispatch_utils import (
@@ -1088,3 +1089,32 @@ def test_non_saasbo_discards_irrelevant_generator_kwargs(self) -> None:
                 run_branin_experiment_with_generation_strategy(
                     generation_strategy=gp_saasbo,
                 )
+
+    def test_suggested_experiment_status(self) -> None:
+        with self.subTest(
+            "choose_generation_strategy_legacy sets statuses on both nodes"
+        ):
+            gs = choose_generation_strategy_legacy(
+                search_space=get_branin_search_space(),
+            )
+            # Sobol step should be INITIALIZATION
+            self.assertEqual(
+                gs._nodes[0].suggested_experiment_status,
+                ExperimentStatus.INITIALIZATION,
+            )
+            # BoTorch step should be OPTIMIZATION
+            self.assertEqual(
+                gs._nodes[1].suggested_experiment_status,
+                ExperimentStatus.OPTIMIZATION,
+            )
+
+        with self.subTest("force_random_search sets INITIALIZATION on Sobol-only GS"):
+            gs = choose_generation_strategy_legacy(
+                search_space=get_branin_search_space(),
+                force_random_search=True,
+            )
+            self.assertEqual(len(gs._nodes), 1)
+            self.assertEqual(
+                gs._nodes[0].suggested_experiment_status,
+                ExperimentStatus.INITIALIZATION,
+            )

ax/orchestration/tests/test_orchestrator.py

@@ -164,12 +164,14 @@ class TestAxOrchestrator(TestCase):
         "transition_criteria="
         "[MinTrials(transition_to='GenerationStep_1_BoTorch'), "
         "MinTrials(transition_to='GenerationStep_1_BoTorch')], "
+        "suggested_experiment_status=ExperimentStatus.INITIALIZATION, "
         "pausing_criteria="
         "[MaxTrialsAwaitingData(threshold=5)]), "
         "GenerationNode(name='GenerationStep_1_BoTorch', "
         "generator_specs=[GeneratorSpec(generator_enum=BoTorch, "
         "generator_key_override=None)], "
         "transition_criteria=None, "
+        "suggested_experiment_status=ExperimentStatus.OPTIMIZATION, "
         "pausing_criteria="
         "[MaxGenerationParallelism(threshold=3)])]), "
         "options=OrchestratorOptions(max_pending_trials=10, "
@@ -2913,12 +2915,14 @@ class TestAxOrchestratorMultiTypeExperiment(TestAxOrchestrator):
         "transition_criteria="
         "[MinTrials(transition_to='GenerationStep_1_BoTorch'), "
         "MinTrials(transition_to='GenerationStep_1_BoTorch')], "
+        "suggested_experiment_status=ExperimentStatus.INITIALIZATION, "
         "pausing_criteria="
         "[MaxTrialsAwaitingData(threshold=5)]), "
         "GenerationNode(name='GenerationStep_1_BoTorch', "
         "generator_specs=[GeneratorSpec(generator_enum=BoTorch, "
         "generator_key_override=None)], "
         "transition_criteria=None, "
+        "suggested_experiment_status=ExperimentStatus.OPTIMIZATION, "
         "pausing_criteria="
         "[MaxGenerationParallelism(threshold=3)])]), "
         "options=OrchestratorOptions(max_pending_trials=10, "