Make GenNode._fit a helper called from GenNode.gen

ax/analysis/plotly/arm_effects/predicted_effects.py

@@ -96,7 +96,9 @@ def compute(
             )
 
         if generation_strategy.model is None:
-            generation_strategy._fit_current_model(data=experiment.lookup_data())
+            generation_strategy._curr._fit(
+                experiment=experiment, data=experiment.lookup_data()
+            )
 
         model = none_throws(generation_strategy.model)
         if not is_predictive(model=model):

ax/benchmark/tests/test_benchmark.py

@@ -316,10 +316,10 @@ def _test_replication_async(self, map_data: bool) -> None:
             "Complete out of order": [0, 0, 1, 2],
         }
         expected_pending_in_each_gen = {
-            "All complete at different times": [[], [0], [1], [2]],
-            "Trials complete immediately": [[], [0], [], [2]],
-            "Trials complete at same time": [[], [0], [], [2]],
-            "Complete out of order": [[], [0], [0], [2]],
+            "All complete at different times": [[None], [0], [1], [2]],
+            "Trials complete immediately": [[None], [0], [None], [2]],
+            "Trials complete at same time": [[None], [0], [None], [2]],
+            "Complete out of order": [[None], [0], [0], [2]],
         }
         # When two trials complete at the same time, the inference trace uses
         # data from both to get the best point, and repeats it.

ax/core/utils.py

@@ -397,32 +397,31 @@ def get_pending_observation_features_based_on_trial_status(
 def extend_pending_observations(
     experiment: Experiment,
     pending_observations: dict[str, list[ObservationFeatures]],
-    generator_runs: list[GeneratorRun],
-) -> dict[str, list[ObservationFeatures]]:
+    generator_run: GeneratorRun,
+) -> None:
     """Extend given pending observations dict (from metric name to observations
     that are pending for that metric), with arms in a given generator run.
 
+    Note: This function performs this operation in-place for performance reasons.
+    It is only used within the ``GenerationStrategy`` class, and is not intended
+    for wide re-use. Please use caution when re-using this function.
+
     Args:
         experiment: Experiment, for which the generation strategy is producing
             ``GeneratorRun``s.
         pending_observations: Dict from metric name to pending observations for
             that metric, used to avoid resuggesting arms that will be explored soon.
-        generator_runs: List of ``GeneratorRun``s currently produced by the
-            ``GenerationStrategy``.
+        generator_run: ``GeneratorRun`` currently produced by the
+            ``GenerationStrategy`` to add to the pending points.
 
-    Returns:
-        A new dictionary of pending observations to avoid in-place modification
     """
-    pending_observations = deepcopy(pending_observations)
-    extended_observations: dict[str, list[ObservationFeatures]] = {}
     for m in experiment.metrics:
-        extended_obs_set = set(pending_observations.get(m, []))
-        for generator_run in generator_runs:
-            for a in generator_run.arms:
-                ob_ft = ObservationFeatures.from_arm(a)
-                extended_obs_set.add(ob_ft)
-        extended_observations[m] = list(extended_obs_set)
-    return extended_observations
+        if m not in pending_observations:
+            pending_observations[m] = []
+        pending_observations[m].extend(
+            ObservationFeatures.from_arm(a) for a in generator_run.arms
+        )
+    return
 
 
 # -------------------- Get target trial utils. ---------------------

ax/exceptions/generation_strategy.py

@@ -25,19 +25,14 @@ class MaxParallelismReachedException(AxGenerationException):
 
     def __init__(
         self,
-        model_name: str,
         num_running: int,
         step_index: int | None = None,
         node_name: str | None = None,
     ) -> None:
         if node_name is not None:
-            msg_start = (
-                f"Maximum parallelism for generation node #{node_name} ({model_name})"
-            )
+            msg_start = f"Maximum parallelism for generation node #{node_name}"
         else:
-            msg_start = (
-                f"Maximum parallelism for generation step #{step_index} ({model_name})"
-            )
+            msg_start = f"Maximum parallelism for generation step #{step_index}"
         super().__init__(
             msg_start
             + f" has been reached: {num_running} trials are currently 'running'. Some "

ax/generation_strategy/external_generation_node.py

@@ -125,12 +125,11 @@ def _fitted_model(self) -> None:
     def model_spec_to_gen_from(self) -> None:
         return None
 
-    def fit(
+    def _fit(
         self,
         experiment: Experiment,
-        data: Data,
-        search_space: SearchSpace | None = None,
-        optimization_config: OptimizationConfig | None = None,
+        data: Data | None = None,
+        status_quo_features: ObservationFeatures | None = None,
         **kwargs: Any,
     ) -> None:
         """A method used to initialize or update the experiment state / data
@@ -144,26 +143,29 @@ def fit(
             data: The experiment data used to fit the model.
             search_space: UNSUPPORTED. An optional override for the experiment
                 search space.
-            optimization_config: UNSUPPORTED. An optional override for the experiment
-                optimization config.
+            status_quo_features: UNSUPPORTED. An optional specification of which
+                status quo (aka "control") features should be used. Currently not
+                yet used in ``ExternalGenerationNode``.
             kwargs: UNSUPPORTED. Additional keyword arguments for model fitting.
         """
-        if search_space is not None or optimization_config is not None or kwargs:
+        if status_quo_features is not None or kwargs:
             raise UnsupportedError(
-                "Unexpected arguments encountered. `ExternalGenerationNode.fit` only "
+                "Unexpected arguments encountered. `ExternalGenerationNode._fit` only "
                 "supports `experiment` and `data` arguments. "
                 "Each of the following arguments should be None / empty. "
-                f"{search_space=}, {optimization_config=}, {kwargs=}."
+                f"{status_quo_features=}, {kwargs=}."
             )
         t_fit_start = time.monotonic()
         self.update_generator_state(
             experiment=experiment,
-            data=data,
+            data=data if data is not None else experiment.lookup_data(),
         )
         self.fit_time_since_gen += time.monotonic() - t_fit_start
 
     def _gen(
         self,
+        experiment: Experiment,
+        data: Data | None = None,
         n: int | None = None,
         pending_observations: dict[str, list[ObservationFeatures]] | None = None,
         **model_gen_kwargs: Any,
@@ -198,6 +200,8 @@ def _gen(
                         pending_parameters.append(o.parameters)
         generated_params: list[TParameterization] = []
         for _ in range(n):
+            # NOTE: We could pass `experiment` and `data` to `get_next_candidate`
+            # to make it possible for it to be stateless in more cases.
             params = self.get_next_candidate(pending_parameters=pending_parameters)
             generated_params.append(params)
             pending_parameters.append(params)