Move all step-based GS functionality to nodes (facebook#4746)

Summary:
X-link: facebookincubator/MCGrad#94

Pull Request resolved: facebook#4746

Pull Request resolved: facebook#4731

A lot going on here:

1. Make `GenerationStep` a factory for `GenerationNode` by replacing its `__init__(self, ...)` constructor (which would have to return a `GStep`) with a `__new__(cls, ...)` construtor, which can return a `GNode` (magic, thanks Devmate!)
2. Adapt storage: stop storing steps and just treat them as nodes. No step-only fields will be saved going forward. Backward compatibility is handled though.
3. Change a bazillion tests and checks in downstream applications.

NOTE: need to remember next steps here: https://www.internalfb.com/diff/D86066476?dst_version_fbid=732904879835890&transaction_fbid=1400151778204976, cc mgarrard

sorrybigdiff

Differential Revision: D80128678

Author: Lena Kashtelyan

ax/benchmark/tests/methods/test_methods.py

@@ -37,9 +37,9 @@ def _test_mbm_acquisition(self, batch_size: int) -> None:
         )
         self.assertEqual(method.name, expected_name)
         gs = method.generation_strategy
-        sobol, kg = gs._steps
-        self.assertEqual(kg.generator, Generators.BOTORCH_MODULAR)
-        generator_kwargs = none_throws(kg.generator_kwargs)
+        sobol, kg = gs._nodes
+        self.assertEqual(kg.generator_spec.generator_enum, Generators.BOTORCH_MODULAR)
+        generator_kwargs = none_throws(kg.generator_spec.generator_kwargs)
         self.assertEqual(generator_kwargs["botorch_acqf_class"], qKnowledgeGradient)
         surrogate_spec = generator_kwargs["surrogate_spec"]
         self.assertEqual(
@@ -64,11 +64,11 @@ def _test_benchmark_replication_runs(
             num_sobol_trials=2,
             name="test",
         )
-        n_sobol_trials = method.generation_strategy._steps[0].num_trials
-        self.assertEqual(n_sobol_trials, 2)
+        n_sobol_trials_tc = method.generation_strategy._nodes[0].num_trials
+        self.assertEqual(n_sobol_trials_tc, 2)
         self.assertEqual(method.name, "test")
         # Only run one non-Sobol trial
-        n_total_trials = n_sobol_trials + 1
+        n_total_trials = n_sobol_trials_tc + 1
         problem = get_benchmark_problem(
             problem_key="ackley4", num_trials=n_total_trials
         )
@@ -104,5 +104,5 @@ def test_sobol(self) -> None:
         method = get_sobol_benchmark_method()
         self.assertEqual(method.name, "Sobol")
         gs = method.generation_strategy
-        self.assertEqual(len(gs._steps), 1)
-        self.assertEqual(gs._steps[0].generator, Generators.SOBOL)
+        self.assertEqual(len(gs._nodes), 1)
+        self.assertEqual(gs._nodes[0].generator_spec.generator_enum, Generators.SOBOL)

ax/benchmark/tests/test_benchmark_method.py

@@ -18,9 +18,11 @@ def test_benchmark_method(self) -> None:
         self.assertEqual(method.name, "Sobol10")
 
         # test that `fit_tracking_metrics` has been correctly set to False
-        for step in method.generation_strategy._steps:
+        for step in method.generation_strategy._nodes:
             self.assertFalse(
-                none_throws(step.generator_kwargs).get("fit_tracking_metrics")
+                none_throws(step.generator_specs[0].generator_kwargs).get(
+                    "fit_tracking_metrics"
+                )
             )
 
         method = BenchmarkMethod(generation_strategy=gs)

ax/core/base_trial.py

@@ -128,10 +128,6 @@ def __init__(
         # Counter to maintain how many arms have been named by this BatchTrial
         self._num_arms_created = 0
 
-        # If generator run(s) in this trial were generated from a generation
-        # strategy, this property will be set to the generation step that produced
-        # the generator run(s).
-        self._generation_step_index: int | None = None
         # NOTE: Please do not store any data related to trial deployment or data-
         # fetching in properties. It is intended to only store properties related
         # to core Ax functionality and not to any third-system that the trials
@@ -475,22 +471,6 @@ def _get_default_name(self, arm_index: int | None = None) -> str:
             arm_index = self._num_arms_created
         return f"{self.index}_{arm_index}"
 
-    def _set_generation_step_index(self, generation_step_index: int | None) -> None:
-        """Sets the `generation_step_index` property of the trial, to reflect which
-        generation step of a given generation strategy (if any) produced the generator
-        run(s) attached to this trial.
-        """
-        if (
-            self._generation_step_index is not None
-            and generation_step_index is not None
-            and self._generation_step_index != generation_step_index
-        ):
-            raise ValueError(
-                "Cannot add generator runs from different generation steps to a "
-                "single trial."
-            )
-        self._generation_step_index = generation_step_index
-
     @property
     def active_arms(self) -> list[Arm]:
         """All non abandoned arms associated with this trial."""

ax/core/batch_trial.py

@@ -298,10 +298,6 @@ def add_generator_run(self, generator_run: GeneratorRun) -> BatchTrial:
         self._add_generator_run(generator_run=generator_run)
 
         self._generator_runs.append(generator_run)
-        if generator_run._generation_step_index is not None:
-            self._set_generation_step_index(
-                generation_step_index=generator_run._generation_step_index
-            )
         self._refresh_arms_by_name()
         return self
 

ax/core/generator_run.py

@@ -97,7 +97,6 @@ def __init__(
         adapter_kwargs: dict[str, Any] | None = None,
         gen_metadata: TGenMetadata | None = None,
         generator_state_after_gen: dict[str, Any] | None = None,
-        generation_step_index: int | None = None,
         candidate_metadata_by_arm_signature: None
         | (dict[str, TCandidateMetadata]) = None,
         generation_node_name: str | None = None,
@@ -135,12 +134,6 @@ def __init__(
                 model when reinstantiating it to continue generation from it,
                 rather than to reproduce the conditions, in which this generator
                 run was created.
-            generation_step_index: Deprecated in favor of generation_node_name.
-                Optional index of the generation step that produced this generator run.
-                Applicable only if the generator run was created via a generation
-                strategy (in which case this index should reflect the index of
-                generation step in a generation strategy) or a standalone generation
-                step (in which case this index should be ``-1``).
             candidate_metadata_by_arm_signature: Optional dictionary of arm signatures
                 to model-produced candidate metadata that corresponds to that arm in
                 this generator run.
@@ -197,16 +190,6 @@ def __init__(
                     "candidate metadata, but not among the arms on this GeneratorRun."
                 )
         self._candidate_metadata_by_arm_signature = candidate_metadata_by_arm_signature
-
-        # Validate that generation step index is either not set (not from generation
-        # strategy or ste), is non-negative (from generation step) or is -1 (from a
-        # standalone generation step that was not a part of a generation strategy).
-        assert (
-            generation_step_index is None  # Not generation strategy/step
-            or generation_step_index == -1  # Standalone generation step
-            or generation_step_index >= 0  # Generation strategy
-        )
-        self._generation_step_index = generation_step_index
         self._generation_node_name = generation_node_name
 
     @property
@@ -342,7 +325,6 @@ def clone(self) -> GeneratorRun:
             adapter_kwargs=self._adapter_kwargs,
             gen_metadata=self._gen_metadata,
             generator_state_after_gen=self._generator_state_after_gen,
-            generation_step_index=self._generation_step_index,
             candidate_metadata_by_arm_signature=cand_metadata,
             generation_node_name=self._generation_node_name,
         )

ax/core/tests/test_batch_trial.py

@@ -115,9 +115,6 @@ def test_add_generator_run(self) -> None:
         self.assertEqual(len(self.batch.generator_runs), 1)
         self.assertEqual(sum(self.batch.weights), sum(self.weights))
 
-        # Overwrite the GS index to not-None.
-        self.batch._generation_step_index = 0
-
         # one of these arms already exists on the BatchTrial,
         # so we should just update its weight
         new_arms = [
@@ -135,8 +132,6 @@ def test_add_generator_run(self) -> None:
             float(sum(self.batch.weights)),
             float(sum(self.weights)) + new_gr_total_weight,
         )
-        # Check the GS index was not overwritten to None.
-        self.assertEqual(self.batch._generation_step_index, 0)
 
     def test_InitWithGeneratorRun(self) -> None:
         generator_run = GeneratorRun(arms=self.arms, weights=self.weights)

ax/generation_strategy/dispatch_utils.py

@@ -466,7 +466,7 @@ def choose_generation_strategy_legacy(
             logger.debug(
                 f"calculated num_initialization_trials={num_initialization_trials}"
             )
-        steps = []
+        nodes = []
         # `disable_progbar` and jit_compile defaults and overrides
         model_is_saasbo = use_saasbo and (suggested_model is Generators.BOTORCH_MODULAR)
         if disable_progbar is not None and not model_is_saasbo:
@@ -493,7 +493,7 @@ def choose_generation_strategy_legacy(
             # Note: With `use_all_trials_in_exp=True`, the Sobol step will automatically
             # account for any existing trials in the experiment, so we can use the total
             # number of initialization trials directly.
-            steps.append(
+            nodes.append(
                 _make_sobol_step(
                     num_trials=num_initialization_trials,
                     min_trials_observed=min_sobol_trials_observed,
@@ -503,7 +503,7 @@ def choose_generation_strategy_legacy(
                     should_deduplicate=should_deduplicate,
                 )
             )
-        steps.append(
+        nodes.append(
             _make_botorch_step(
                 generator=suggested_model,
                 winsorization_config=winsorization_config,
@@ -518,17 +518,17 @@ def choose_generation_strategy_legacy(
             ),
         )
         # set name for GS
-        bo_step = steps[-1]
-        surrogate_spec = bo_step.generator_kwargs.get("surrogate_spec")
+        bo_step = nodes[-1]
+        surrogate_spec = bo_step.generator_spec.generator_kwargs.get("surrogate_spec")
         name = None
         if (
-            bo_step.generator is Generators.BOTORCH_MODULAR
+            bo_step.generator_spec.generator_enum is Generators.BOTORCH_MODULAR
             and surrogate_spec is not None
             and (model_config := surrogate_spec.model_configs[0]).botorch_model_class
             == SaasFullyBayesianSingleTaskGP
         ):
             name = f"Sobol+{model_config.name}"
-        gs = GenerationStrategy(steps=steps, name=name)
+        gs = GenerationStrategy(nodes=nodes, name=name)
         logger.info(
             f"Using Bayesian Optimization generation strategy: {gs}. Iterations after"
             f" {num_initialization_trials} will take longer to generate due"