Clean up center generation node implementation (#4734)

Summary:

This diff cleans up center generation node logic to hopefully make it a bit more eloquent. It does the following:
* Removes AutoTransitionAfterGenOrExhaustion as it is duplicative of AutoTransitionAfterGen, it just has a different defalt
* Updates custom gen override to skip if either (1) the center point already exists or (2) we were unable to find a suitable center point
* makes update_generator_state to be a pass through since we'll set the property during gen override
*only compute center params once by adding a property on the center generation node

Notably, it would be only in a misconfigured state that we would use sobol fallback now, this is consistent with David's initial changes, see D87797277 for explaination of that chage -- we could even remove sobol fallback from the node if we want

Reviewed By: saitcakmak

Differential Revision:
D90040731

Privacy Context Container: L1307644

Author: mgarrard

ax/generation_strategy/center_generation_node.py

@@ -21,9 +21,7 @@
 from ax.exceptions.generation_strategy import AxGenerationException
 from ax.generation_strategy.external_generation_node import ExternalGenerationNode
 from ax.generation_strategy.generator_spec import GeneratorSpec
-from ax.generation_strategy.transition_criterion import (
-    AutoTransitionAfterGenOrExhaustion,
-)
+from ax.generation_strategy.transition_criterion import AutoTransitionAfterGen
 from pyre_extensions import none_throws
 
 
@@ -43,7 +41,7 @@ def __init__(self, next_node_name: str) -> None:
         super().__init__(
             name="CenterOfSearchSpace",
             transition_criteria=[
-                AutoTransitionAfterGenOrExhaustion(
+                AutoTransitionAfterGen(
                     transition_to=next_node_name,
                     continue_trial_generation=False,
                 )
@@ -58,9 +56,12 @@ def __init__(self, next_node_name: str) -> None:
             ),
             **self.fallback_specs,  # This includes the default fallbacks.
         }
+        # custom property to enable single center point computation
+        self._center_params: TParameterization | None = None
 
     def update_generator_state(self, experiment: Experiment, data: Data) -> None:
-        self.search_space = experiment.search_space
+        # State is already set in gen() and will persist during generation
+        pass
 
     def gen(
         self,
@@ -79,27 +80,26 @@ def gen(
         before attempting generation. If so, it sets _should_skip to True and
         returns None, allowing the generation strategy to transition to the next node.
         """
-        # Check if center already exists or is infeasible
         self.search_space = experiment.search_space
-        center_params = self.compute_center_params()
+        self._center_params = self.compute_center_params()
 
         # Check if unable to find a suitable center
-        if center_params is None:
+        if self._center_params is None:
             self._should_skip = True
             return None
 
         # Check if center already exists in experiment
-        center_arm = Arm(parameters=center_params)
+        center_arm = Arm(parameters=self._center_params)
         if center_arm.signature in experiment.arms_by_signature:
             self._should_skip = True
             return None
 
-        # Otherwise, proceed with normal generation
         return super().gen(
             experiment=experiment,
             pending_observations=pending_observations,
             skip_fit=skip_fit,
             data=data,
+            n=n,
             arms_per_node=arms_per_node,
             **gs_gen_kwargs,
         )
@@ -164,12 +164,4 @@ def get_next_candidate(
         parameter bounds and parameter constraints w.r.t non-log range parameters.
         This finds the center of the largest inscribed ball in the feasible region.
         """
-        center_params = self.compute_center_params()
-        if center_params is None:
-            # raising an exception here will cause fallback to sobol, currently
-            # it should be very unlikely to hit this case
-            raise AxGenerationException(
-                "Center of the search space does not satisfy parameter "
-                "constraints. The generation strategy will fallback to Sobol. "
-            )
-        return center_params
+        return none_throws(self._center_params)

ax/generation_strategy/external_generation_node.py

@@ -23,8 +23,6 @@
 from ax.generation_strategy.transition_criterion import TransitionCriterion
 
 
-# TODO[drfreund]: Introduce a `GenerationNodeInterface` to
-# make inheritance/overriding of `GenNode` methods cleaner.
 class ExternalGenerationNode(GenerationNode, ABC):
     """A generation node intended to be used with non-Ax methods for
     candidate generation.

ax/generation_strategy/tests/test_center_generation_node.py

@@ -10,7 +10,6 @@
 from unittest.mock import patch
 
 from ax.adapter.registry import Generators
-
 from ax.core.arm import Arm
 from ax.core.experiment import Experiment
 from ax.core.parameter import (
@@ -30,9 +29,7 @@
 from ax.generation_strategy.generation_node import GenerationNode
 from ax.generation_strategy.generation_strategy import GenerationStrategy
 from ax.generation_strategy.generator_spec import GeneratorSpec
-from ax.generation_strategy.transition_criterion import (
-    AutoTransitionAfterGenOrExhaustion,
-)
+from ax.generation_strategy.transition_criterion import AutoTransitionAfterGen
 from ax.utils.common.testutils import TestCase
 from ax.utils.testing.core_stubs import get_branin_experiment
 from pyre_extensions import none_throws
@@ -123,7 +120,7 @@ def test_center_generation(self) -> None:
         self.assertEqual(
             self.node.transition_criteria,
             [
-                AutoTransitionAfterGenOrExhaustion(
+                AutoTransitionAfterGen(
                     transition_to="test", continue_trial_generation=False
                 )
             ],

ax/generation_strategy/tests/test_transition_criterion.py

@@ -11,12 +11,10 @@
 
 import pandas as pd
 from ax.adapter.registry import Generators
-from ax.core.arm import Arm
 from ax.core.auxiliary import AuxiliaryExperiment, AuxiliaryExperimentPurpose
 from ax.core.data import Data
 from ax.core.trial_status import TrialStatus
 from ax.exceptions.core import UserInputError
-from ax.generation_strategy.center_generation_node import CenterGenerationNode
 from ax.generation_strategy.generation_strategy import (
     GenerationNode,
     GenerationStep,
@@ -25,7 +23,6 @@
 from ax.generation_strategy.generator_spec import GeneratorSpec
 from ax.generation_strategy.transition_criterion import (
     AutoTransitionAfterGen,
-    AutoTransitionAfterGenOrExhaustion,
     AuxiliaryExperimentCheck,
     IsSingleObjective,
     MaxGenerationParallelism,
@@ -376,69 +373,6 @@ def test_auto_with_should_skip_node(self) -> None:
             .is_met(experiment=experiment, curr_node=gs._nodes[0])
         )
 
-    def test_auto_transition_after_gen_or_exhaustion(self) -> None:
-        """Test AutoTransitionAfterGenOrExhaustion transitions after generation
-        or when search space is exhausted.
-        """
-        experiment = self.branin_experiment
-
-        # Test 1: Transition after successful generation (like AutoTransitionAfterGen)
-        gs = GenerationStrategy(
-            name="test",
-            nodes=[
-                GenerationNode(
-                    name="sobol_1",
-                    generator_specs=[self.sobol_generator_spec],
-                    transition_criteria=[
-                        AutoTransitionAfterGenOrExhaustion(transition_to="sobol_2")
-                    ],
-                ),
-                GenerationNode(
-                    name="sobol_2", generator_specs=[self.sobol_generator_spec]
-                ),
-            ],
-        )
-        gs.experiment = experiment
-
-        # Generate from first node
-        gs.gen(experiment=experiment)
-        self.assertEqual(gs.current_node_name, "sobol_1")
-
-        # Should transition to next node on next gen after generating
-        gs.gen(experiment=experiment)
-        self.assertEqual(gs.current_node_name, "sobol_2")
-
-        # Test 2: Transition immediately when search space is exhausted
-        # Use CenterGenerationNode which can only generate one unique candidate
-        experiment2 = get_branin_experiment()
-        # Add the center point so it's already in the experiment
-        center_arm = Arm(parameters={"x1": 2.5, "x2": 7.5})
-        experiment2.new_trial().add_arm(arm=center_arm)
-
-        gs2 = GenerationStrategy(
-            name="test_exhaustion",
-            nodes=[
-                CenterGenerationNode(next_node_name="sobol"),
-                GenerationNode(
-                    name="sobol", generator_specs=[self.sobol_generator_spec]
-                ),
-            ],
-        )
-        gs2.experiment = experiment2
-
-        # Since center already exists, should skip CenterGenerationNode
-        # and transition directly to sobol
-        gr = gs2.gen(experiment=experiment2, n=1)[0]
-        self.assertEqual(gr[0]._generation_node_name, "sobol")
-        self.assertEqual(gs2.current_node_name, "sobol")
-
-        # Test 3: Call block_continued_generation_error
-        criterion = AutoTransitionAfterGenOrExhaustion(transition_to="sobol_2")
-        # This method has a pass statement, so calling it should not raise an error
-        criterion.block_continued_generation_error(
-            node_name="sobol_1", experiment=experiment, trials_from_node=set()
-        )
-
     def test_is_single_objective_does_not_transition(self) -> None:
         exp = self.branin_experiment
         exp.optimization_config = get_branin_multi_objective_optimization_config()

ax/generation_strategy/transition_criterion.py

@@ -603,73 +603,6 @@ def block_continued_generation_error(
         pass
 
 
-class AutoTransitionAfterGenOrExhaustion(TransitionCriterion):
-    """A class to designate automatic transition from one GenerationNode to another
-    after generating a candidate or when the search space is exhausted.
-
-    This criterion is met when either:
-    1. A GeneratorRun is generated by this GenerationNode, OR
-    2. The search space is exhausted (no more unique candidates can be generated)
-
-    This is particularly useful for nodes that have a limited search space, such as
-    CenterGenerationNode, which can only generate one unique candidate (the center).
-    If that candidate already exists or is infeasible, the search space is exhausted and
-    the node should transition to the next node.
-
-    Args:
-        transition_to: The name of the GenerationNode the GenerationStrategy should
-            transition to next.
-        block_transition_if_unmet: A flag to prevent the node from completing and
-            being able to transition to another node. This criterion defaults to
-            setting this to True to ensure we validate that either a GeneratorRun was
-            generated or the search space is exhausted.
-        continue_trial_generation: A flag to indicate that all generation for a given
-            trial is not completed, and thus even after transition, the next node will
-            continue to generate arms for the same trial. Example usage: in
-            ``BatchTrial``s we may  enable generation of arms within a batch from
-            different ``GenerationNodes`` by setting this flag to True.
-    """
-
-    def __init__(
-        self,
-        transition_to: str,
-        block_transition_if_unmet: bool | None = True,
-        continue_trial_generation: bool | None = False,
-    ) -> None:
-        super().__init__(
-            transition_to=transition_to,
-            block_transition_if_unmet=block_transition_if_unmet,
-            continue_trial_generation=continue_trial_generation,
-        )
-
-    def is_met(
-        self,
-        experiment: Experiment,
-        curr_node: GenerationNode,
-    ) -> bool:
-        """Return True if any GeneratorRun is generated by this GenerationNode
-        or if the node should be skipped due to search space exhaustion.
-        """
-        if curr_node._should_skip:
-            return True
-        last_gr_from_gs = curr_node.generation_strategy.last_generator_run
-        if (
-            last_gr_from_gs is not None
-            and last_gr_from_gs._generation_node_name == curr_node.name
-        ):
-            return True
-        return False
-
-    def block_continued_generation_error(
-        self,
-        node_name: str,
-        experiment: Experiment,
-        trials_from_node: set[int],
-    ) -> None:
-        """Error to be raised if the `block_gen_if_met` flag is set to True."""
-        pass
-
-
 class AuxiliaryExperimentCheck(TransitionCriterion):
     """A class to transition from one GenerationNode to another by checking if certain
     types of Auxiliary Experiment purposes exists.