call maybe transition to next node in Analyses and enable skipping center GN when trials are provided (#4922)

Summary:

see title. This important for using a new GS on an experiment with data in model-based analyses.

This adds support for
* skipping the center GN when trials with data are provided
* Transitioning the GS when possible when extracting the adapter in analyses (important for fast-forwarding the GS for an experiment with data).

Differential Revision: D93804600

Author: sdaulton

ax/analysis/utils.py

@@ -76,7 +76,7 @@ def extract_relevant_adapter(
             "Provided GenerationStrategy has no adapter, but no Experiment was "
             "provided to source data to fit the adapter."
         )
-
+    generation_strategy.maybe_transition_to_next_node(raise_data_required_error=False)
     generation_strategy.current_node._fit(experiment=experiment)
     adapter = generation_strategy.adapter
 

ax/api/utils/generation_strategy_dispatch.py

@@ -270,7 +270,12 @@ def choose_generation_strategy(
     if struct.initialize_with_center and (
         struct.initialization_budget is None or struct.initialization_budget > 0
     ):
-        center_node = CenterGenerationNode(next_node_name=nodes[0].name)
+        center_node = CenterGenerationNode(
+            next_node_name=nodes[0].name,
+            use_existing_trials_for_initialization=(
+                struct.use_existing_trials_for_initialization
+            ),
+        )
         nodes.insert(0, center_node)
         gs_name = f"Center+{gs_name}"
     return GenerationStrategy(name=gs_name, nodes=nodes)

ax/generation_strategy/center_generation_node.py

@@ -29,12 +29,14 @@
 @dataclass(init=False)
 class CenterGenerationNode(ExternalGenerationNode):
     next_node_name: str
+    use_existing_trials_for_initialization: bool
 
     def __init__(
         self,
         next_node_name: str,
         suggested_experiment_status: ExperimentStatus
         | None = ExperimentStatus.INITIALIZATION,
+        use_existing_trials_for_initialization: bool = False,
     ) -> None:
         """A generation node that samples the center of the search space.
         This generation node is only used to generate the first point of the experiment.
@@ -49,6 +51,9 @@ def __init__(
                 the center point.
             suggested_experiment_status: Optional suggested experiment status for this
                 node.
+            use_existing_trials_for_initialization: If True and the experiment already
+                has trials, this node will be skipped during transition checks
+                outside the gen flow (e.g., when fitting a model for analysis).
         """
         super().__init__(
             name="CenterOfSearchSpace",
@@ -63,6 +68,9 @@ def __init__(
         )
         self.search_space: SearchSpace | None = None
         self.next_node_name = next_node_name
+        self.use_existing_trials_for_initialization = (
+            use_existing_trials_for_initialization
+        )
         self.fallback_specs: dict[type[Exception], GeneratorSpec] = {
             AxGenerationException: GeneratorSpec(
                 generator_enum=Generators.SOBOL, generator_key_override="Fallback_Sobol"
@@ -72,6 +80,23 @@ def __init__(
         # custom property to enable single center point computation
         self._center_params: TParameterization | None = None
 
+    def should_transition_to_next_node(
+        self, raise_data_required_error: bool = True
+    ) -> tuple[bool, str]:
+        # Lazily evaluate skip condition for cases where gen() hasn't been
+        # called (e.g., maybe_transition_to_next_node from analysis code).
+        # When use_existing_trials_for_initialization is True, existing trials
+        # count toward initialization, so the center node can be skipped.
+        if (
+            not self._should_skip
+            and self.use_existing_trials_for_initialization
+            and len(self.experiment.trials) > 0
+        ):
+            self._should_skip = True
+        return super().should_transition_to_next_node(
+            raise_data_required_error=raise_data_required_error
+        )
+
     def update_generator_state(self, experiment: Experiment, data: Data) -> None:
         # State is already set in gen() and will persist during generation
         pass

ax/generation_strategy/generation_strategy.py

@@ -328,7 +328,7 @@ def current_generator_run_limit(
                 generate any more generator runs at all.
         """
         try:
-            self._maybe_transition_to_next_node(raise_data_required_error=False)
+            self.maybe_transition_to_next_node(raise_data_required_error=False)
         except GenerationStrategyCompleted:
             return 0, True
 
@@ -543,7 +543,7 @@ def _gen_with_multiple_nodes(
                 # reset should skip as conditions may have changed, do not reset
                 # until now so node properties can be as up to date as possible
                 node_to_gen_from._should_skip = False
-            transitioned = self._maybe_transition_to_next_node()
+            transitioned = self._transition_to_next_node()
             try:
                 gr = self._curr.gen(
                     experiment=experiment,
@@ -604,24 +604,11 @@ def _should_continue_gen_for_trial(self) -> bool:
 
     # ------------------------- Node selection logic helpers. -------------------------
 
-    def _maybe_transition_to_next_node(
-        self,
-        raise_data_required_error: bool = True,
-    ) -> bool:
-        """Moves this generation strategy to next node if the current node's
-        transition criteria are met. This method is safe to use both when generating
-        candidates or simply checking how many generator runs (to be made into trials)
-        can currently be produced.
-
-        NOTE: this method raises ``GenerationStrategyCompleted`` error if the
-        optimization is complete
-
-        Args:
-            raise_data_required_error: Whether to raise ``DataRequiredError`` in the
-                maybe_step_completed method in GenerationNode class.
+    def _transition_to_next_node(self, raise_data_required_error: bool = True) -> bool:
+        """Attempts a single transition to the next node if criteria are met.
 
         Returns:
-            Whether generation strategy moved to the next node.
+            Whether the generation strategy moved to the next node.
         """
         move_to_next_node, next_node = self._curr.should_transition_to_next_node(
             raise_data_required_error=raise_data_required_error
@@ -634,3 +621,29 @@ def _maybe_transition_to_next_node(
                 )
             self._curr = self.nodes_by_name[next_node]
         return move_to_next_node
+
+    def maybe_transition_to_next_node(
+        self, raise_data_required_error: bool = True
+    ) -> bool:
+        """Moves this generation strategy to next node if the current node's
+        transition criteria are met, advancing through multiple nodes if
+        possible. This method is safe to use both when generating candidates or
+        simply checking how many generator runs (to be made into trials) can
+        currently be produced.
+
+        NOTE: this method raises ``GenerationStrategyCompleted`` error if the
+        optimization is complete
+
+        Args:
+            raise_data_required_error: Whether to raise ``DataRequiredError`` in the
+                maybe_step_completed method in GenerationNode class.
+
+        Returns:
+            Whether generation strategy moved to the next node.
+        """
+        moved = False
+        while self._transition_to_next_node(
+            raise_data_required_error=raise_data_required_error
+        ):
+            moved = True
+        return moved

ax/generation_strategy/tests/test_center_generation_node.py

@@ -219,7 +219,8 @@ def test_deduplication(self) -> None:
     def test_repr(self) -> None:
         self.assertEqual(
             repr(self.node),
-            "CenterGenerationNode(next_node_name='test')",
+            "CenterGenerationNode(next_node_name='test',"
+            " use_existing_trials_for_initialization=False)",
         )
 
     def test_equality(self) -> None:

ax/service/ax_client.py

@@ -1226,7 +1226,7 @@ def fit_model(self) -> None:
                 "At least one trial must be completed with data to fit a model."
             )
         # Check if we should transition before generating the next candidate.
-        self.generation_strategy._maybe_transition_to_next_node()
+        self.generation_strategy.maybe_transition_to_next_node()
         self.generation_strategy._curr._fit(experiment=self.experiment)
 
     def verify_trial_parameterization(

ax/service/tests/test_ax_client.py

@@ -2597,7 +2597,7 @@ def test_get_pareto_optimal_points_objective_threshold_inference(
         ax_client, _ = get_branin_currin_optimization_with_N_sobol_trials(
             num_trials=20, include_objective_thresholds=False
         )
-        ax_client.generation_strategy._maybe_transition_to_next_node()
+        ax_client.generation_strategy.maybe_transition_to_next_node()
         ax_client.generation_strategy._curr._fit(experiment=ax_client.experiment)
         with with_rng_seed(seed=RANDOM_SEED):
             predicted_pareto = ax_client.get_pareto_optimal_parameters()