Add caching to common methods#4830
Closed
mgarrard wants to merge 3 commits into
Closed
Conversation
meta-cla
Bot
added
the
CLA Signed
mgarrard
force-pushed
the
export-D91552553
branch
from
ad21e82 to
ef74f20
Compare
mgarrard
added a commit
to mgarrard/Ax
that referenced
this pull request
Jan 27, 2026
Summary: This method is called many, many times during generation and it's computational cost adds up over time. By cacheing it we can significant improvements in computation time, especially in high trial count regimes. Differential Revision: D91552553
mgarrard
force-pushed
the
export-D91552553
branch
from
ef74f20 to
12c53d3
Compare
mgarrard
added a commit
to mgarrard/Ax
that referenced
this pull request
Jan 28, 2026
Summary: This method is called many, many times during generation and it's computational cost adds up over time. By cacheing it we can significant improvements in computation time, especially in high trial count regimes. Differential Revision: D91552553
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #4830 +/- ##
=======================================
Coverage 96.75% 96.76%
=======================================
Files 591 591
Lines 61874 61882 +8
=======================================
+ Hits 59869 59882 +13
+ Misses 2005 2000 -5 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
mgarrard
added a commit
to mgarrard/Ax
that referenced
this pull request
Feb 4, 2026
Summary: This method is called many, many times during generation and it's computational cost adds up over time. By cacheing it we can significant improvements in computation time, especially in high trial count regimes. Reviewed By: mpolson64 Differential Revision: D91552553
mgarrard
added a commit
to mgarrard/Ax
that referenced
this pull request
Feb 5, 2026
Summary: This method is called many, many times during generation and it's computational cost adds up over time. By cacheing it we can significant improvements in computation time, especially in high trial count regimes. Reviewed By: mpolson64 Differential Revision: D91552553
mgarrard
added a commit
to mgarrard/Ax
that referenced
this pull request
Feb 5, 2026
Summary: This method is called many, many times during generation and it's computational cost adds up over time. By cacheing it we can significant improvements in computation time, especially in high trial count regimes. Reviewed By: mpolson64 Differential Revision: D91552553
mgarrard
added a commit
to mgarrard/Ax
that referenced
this pull request
Feb 5, 2026
Summary: This method is called many, many times during generation and it's computational cost adds up over time. By cacheing it we can significant improvements in computation time, especially in high trial count regimes. Reviewed By: mpolson64 Differential Revision: D91552553
added 2 commits
February 9, 2026 12:18
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
Summary: Since transition_to is now required on transition criterion, we can remove checks/asserts related to none checks. this is a basic no-op simplification. Futher restructuring seperated into a different diff for ease of review Differential Revision: D91398877
mgarrard
force-pushed
the
export-D91552553
branch
from
12c53d3 to
8a8bfca
Compare
mgarrard
added a commit
to mgarrard/Ax
that referenced
this pull request
Feb 9, 2026
Summary: This method is called many, many times during generation and it's computational cost adds up over time. By cacheing it we can significant improvements in computation time, especially in high trial count regimes. Reviewed By: mpolson64 Differential Revision: D91552553
Summary: Pull Request resolved: facebook#4830 This method is called many, many times during generation and it's computational cost adds up over time. By cacheing it we can significant improvements in computation time, especially in high trial count regimes. Reviewed By: mpolson64 Differential Revision: D91552553
mgarrard
force-pushed
the
export-D91552553
branch
from
8a8bfca to
a57e9c0
Compare
|
This pull request has been merged in 404850a. |
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
3 participants
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 method is called many, many times during generation and it's computational cost adds up over time. By cacheing it we can significant improvements in computation time, especially in high trial count regimes.
Differential Revision: D91552553