improve optimization complete logic#4828
Closed
mgarrard wants to merge 1 commit intofacebook:mainfrom
Closed
Conversation
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #4828 +/- ##
=======================================
Coverage 96.75% 96.75%
=======================================
Files 591 591
Lines 61874 61882 +8
=======================================
+ Hits 59869 59877 +8
Misses 2005 2005 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
fd87f1a to
c77d15a
Compare
mgarrard
added a commit
to mgarrard/Ax
that referenced
this pull request
Jan 27, 2026
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
mgarrard
added a commit
to mgarrard/Ax
that referenced
this pull request
Jan 27, 2026
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
c77d15a to
964c8ab
Compare
mgarrard
added a commit
to mgarrard/Ax
that referenced
this pull request
Jan 27, 2026
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
mgarrard
added a commit
to mgarrard/Ax
that referenced
this pull request
Jan 27, 2026
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
mgarrard
added a commit
to mgarrard/Ax
that referenced
this pull request
Jan 27, 2026
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
mgarrard
added a commit
to mgarrard/Ax
that referenced
this pull request
Jan 28, 2026
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
mgarrard
added a commit
to mgarrard/Ax
that referenced
this pull request
Feb 4, 2026
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
mgarrard
added a commit
to mgarrard/Ax
that referenced
this pull request
Feb 5, 2026
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
mgarrard
added a commit
to mgarrard/Ax
that referenced
this pull request
Feb 5, 2026
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
mgarrard
added a commit
to mgarrard/Ax
that referenced
this pull request
Feb 5, 2026
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
964c8ab to
5c24950
Compare
mgarrard
added a commit
to mgarrard/Ax
that referenced
this pull request
Feb 9, 2026
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)
Reviewed By: lena-kashtelyan
Differential Revision: D91549954
mgarrard
added a commit
to mgarrard/Ax
that referenced
this pull request
Feb 9, 2026
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)
Reviewed By: lena-kashtelyan
Differential Revision: D91549954
mgarrard
added a commit
to mgarrard/Ax
that referenced
this pull request
Feb 9, 2026
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)
Reviewed By: lena-kashtelyan
Differential Revision: D91549954
mgarrard
added a commit
to mgarrard/Ax
that referenced
this pull request
Feb 9, 2026
Summary: Pull Request resolved: facebook#4828 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) Reviewed By: lena-kashtelyan Differential Revision: D91549954
5c24950 to
e1596d1
Compare
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)
Reviewed By: lena-kashtelyan
Differential Revision: D91549954
e1596d1 to
9e28603
Compare
|
This pull request has been merged in ed975ea. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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
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