improve optimization complete logic (facebook#4828)

Summary:

This criteria updates the completion state logic to assume if a node can transition, and that transition is to itself, then the optimization is complete.

This works because should_transition_to_next_node only considers transtion blocking criteria (ie not max parallelism) when thinking about should transition or not. And if a node points to itself, we can assume that signifies the end of the optimiztion (steps are initialized this way earlier in this stack). this allows allows for the gs to be re-called into, and the tc criterion to change thus putting it back into a non-complete state.

An alternative I considered is to check if all transition edges are completed, and at least one points to self. This would look something like the below snippet. It would be much more expensive to evaluate, and is guarding against a malformed strategy. Edges are already known to be created in order of importance, and self transition edges should be considered ending edges when their importance is considered

```
property
def optimization_complete(self) -> bool:
    if len(self._curr.transition_criteria) == 0:
        return False

    # Check ALL transition edges, not just the first matching one
    for next_node, all_tc in self._curr.transition_edges.items():
        transition_blocking = [tc for tc in all_tc if tc.block_transition_if_unmet]
        if not transition_blocking:
            continue
        
        all_met = all(
            tc.is_met(experiment=self.experiment, curr_node=self._curr)
            for tc in transition_blocking
        )
        
        if all_met:
            # An edge's criteria are met - check where it points
            if next_node != self._curr.name:
                return False  # Can transition to different node, not complete
    
    # All met edges (if any) point to self
    # Check if we actually have any met criteria pointing to self
    can_transition, next_node = self._curr.should_transition_to_next_node(
        raise_data_required_error=False
    )
    return can_transition and next_node == self._curr.name
```

The thrid alternative is to instate "compeletion node", which i think could be viable in the future if we have more complex generation strategies than we currently support, and the self generation logic is too cumbersome.

For now though, I think this is a pretty nice simplification that also should have some compute wins. Going from O (number of nodes * number of TC per node), to O(number of tc on current node)

Differential Revision: D91549954

Author: mgarrard

ax/generation_strategy/generation_node.py

@@ -41,7 +41,6 @@
 )
 from ax.generation_strategy.generator_spec import GeneratorSpec
 from ax.generation_strategy.transition_criterion import (
-    AutoTransitionAfterGen,
     MaxGenerationParallelism,
     MinTrials,
     TransitionCriterion,
@@ -249,21 +248,6 @@ def experiment(self) -> Experiment:
         """Returns the experiment associated with this GenerationStrategy"""
         return self.generation_strategy.experiment
 
-    @property
-    def is_completed(self) -> bool:
-        """Returns True if this GenerationNode is complete and should transition to
-        the next node.
-        """
-        # TODO: @mgarrard make this logic more robust and general
-        # We won't mark a node completed if it has an AutoTransitionAfterGen criterion
-        # as this is typically used in cyclic generation strategies
-        should_transition, _ = self.should_transition_to_next_node(
-            raise_data_required_error=False
-        )
-        return should_transition and not any(
-            isinstance(tc, AutoTransitionAfterGen) for tc in self.transition_criteria
-        )
-
     @property
     def previous_node(self) -> GenerationNode | None:
         """Returns the previous ``GenerationNode``, if any."""

ax/generation_strategy/generation_strategy.py

@@ -180,8 +180,20 @@ def last_generator_run(self) -> GeneratorRun | None:
 
     @property
     def optimization_complete(self) -> bool:
-        """Checks whether all nodes are completed in the generation strategy."""
-        return all(node.is_completed for node in self._nodes)
+        """Checks whether optimization is complete.
+
+        A strategy is complete when the current node's transition criteria
+        are met and point back to itself (self-transition).
+
+        Nodes with no transition_criteria are infinite by design and never complete.
+        """
+        if len(self._curr.transition_criteria) == 0:
+            return False
+
+        can_transition, next_node = self._curr.should_transition_to_next_node(
+            raise_data_required_error=False
+        )
+        return can_transition and next_node == self._curr.name
 
     def gen_single_trial(
         self,
@@ -612,13 +624,13 @@ 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 is completed,
-        and it is not the last node in this generation strategy. 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.
+        """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 current
-        generation node is complete, but it is also the last in generation strategy.
+        NOTE: this method raises ``GenerationStrategyCompleted`` error if the
+        optimization is complete
 
         Args:
             raise_data_required_error: Whether to raise ``DataRequiredError`` in the

ax/generation_strategy/tests/test_generation_strategy.py

@@ -2000,6 +2000,28 @@ def test_gs_with_input_constructor(self) -> None:
             self.assertEqual(trial.generator_runs[1]._generation_node_name, "sobol_3")
             self.assertEqual(len(trial.generator_runs[1].arms), 8)
 
+    def test_optimization_complete_single_node_no_criteria(self) -> None:
+        """Test that a single node with no transition_criteria never completes."""
+        exp = get_branin_experiment()
+        gs = GenerationStrategy(
+            nodes=[
+                GenerationNode(
+                    name="infinite sobol",
+                    generator_specs=[self.sobol_generator_spec],
+                    transition_criteria=[],  # No criteria = infinite by design
+                ),
+            ]
+        )
+        gs.experiment = exp
+
+        # Generate many trials - never completes
+        for _ in range(3):
+            self.assertFalse(gs.optimization_complete)
+            gr = gs.gen_single_trial(experiment=exp)
+            exp.new_trial(generator_run=gr).mark_running(no_runner_required=True)
+
+        self.assertFalse(gs.optimization_complete)
+
     # ------------- Testing helpers (put tests above this line) -------------
 
     def _run_GS_for_N_rounds(