improve optimization complete logic

[mgarrard]

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

[meta-codesync]

@mgarrard has exported this pull request. If you are a Meta employee, you can view the originating Diff in D91549954.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 96.29630% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 96.71%. Comparing base (1eb22a8) to head (fd87f1a).

Files with missing lines	Patch %	Lines
ax/generation_strategy/transition_criterion.py	83.33%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #4828      +/-   ##
==========================================
- Coverage   96.71%   96.71%   -0.01%     
==========================================
  Files         586      585       -1     
  Lines       61323    61220     -103     
==========================================
- Hits        59309    59208     -101     
+ Misses       2014     2012       -2     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.