Add caching to common methods (#4830)

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

Author: mgarrard

ax/generation_strategy/generation_node.py

@@ -190,6 +190,10 @@ def __init__(
         self.fallback_specs = (
             fallback_specs if fallback_specs is not None else DEFAULT_FALLBACK
         )
+        # Cache for trials_from_node property to avoid recomputation
+        # on every access. Invalidated when trial count changes.
+        self._trials_from_node_cache: set[int] | None = None
+        self._cached_trial_count: int = -1
 
     @property
     def name(self) -> str:
@@ -724,17 +728,28 @@ def _pick_fitted_adapter_to_gen_from(self) -> GeneratorSpec:
     def trials_from_node(self) -> set[int]:
         """Returns a set containing the indices of trials generated by this node.
 
+        Results are cached and invalidated when the experiment's trial count changes.
+
         Returns:
             Set[int]: A set containing all the indices of trials generated by this node.
         """
+        current_trial_count = len(self.experiment.trials)
+        if (
+            self._trials_from_node_cache is not None
+            and self._cached_trial_count == current_trial_count
+        ):
+            return self._trials_from_node_cache
+
+        # (re)-build cache
         trials_from_node = set()
-        for _idx, trial in self.experiment.trials.items():
+        for trial in self.experiment.trials.values():
             for gr in trial.generator_runs:
-                if (
-                    gr._generation_node_name is not None
-                    and gr._generation_node_name == self.name
-                ):
+                if gr._generation_node_name == self.name:
                     trials_from_node.add(trial.index)
+                    break
+
+        self._trials_from_node_cache = trials_from_node
+        self._cached_trial_count = current_trial_count
         return trials_from_node
 
     @property

ax/generation_strategy/generation_strategy.py

@@ -361,6 +361,10 @@ def _unset_non_persistent_state_fields(self) -> None:
             n._step_index = None
             if len(n.generator_specs) > 1:
                 n._generator_spec_to_gen_from = None
+            # Reset cache fields that are used for performance optimization only
+            # and should not affect equality comparisons.
+            n._trials_from_node_cache = None
+            n._cached_trial_count = -1
 
     # TODO: Deprecate `steps` argument fully in Q1'26.
     def _validate_and_set_step_sequence(self, steps: list[GenerationNode]) -> None: