Rename choose_generation_strategy to choose_generation_strategy_legacy

Summary: We have a new choose_generation_strategy in ax.preview.modelbridge that will be migrated out of preview in the next diff. Rename existing choose_generation_strategy to avoid name conflict

Differential Revision: D70647194

Author: mpolson64

ax/analysis/plotly/tests/test_predicted_effects.py

@@ -15,7 +15,7 @@
 from ax.core.observation import ObservationFeatures
 from ax.core.trial import Trial
 from ax.exceptions.core import UserInputError
-from ax.generation_strategy.dispatch_utils import choose_generation_strategy
+from ax.generation_strategy.dispatch_utils import choose_generation_strategy_legacy
 from ax.modelbridge.prediction_utils import predict_at_point
 from ax.modelbridge.registry import Generators
 from ax.utils.common.testutils import TestCase
@@ -50,7 +50,7 @@ def test_compute_for_requires_a_gs(self) -> None:
     def test_compute_for_requires_trials(self) -> None:
         analysis = PredictedEffectsPlot(metric_name="branin")
         experiment = get_branin_experiment()
-        generation_strategy = choose_generation_strategy(
+        generation_strategy = choose_generation_strategy_legacy(
             search_space=experiment.search_space,
             experiment=experiment,
         )
@@ -62,7 +62,7 @@ def test_compute_for_requires_trials(self) -> None:
     def test_compute_for_requires_a_model_that_predicts(self) -> None:
         analysis = PredictedEffectsPlot(metric_name="branin")
         experiment = get_branin_experiment(with_batch=True, with_completed_batch=True)
-        generation_strategy = choose_generation_strategy(
+        generation_strategy = choose_generation_strategy_legacy(
             search_space=experiment.search_space,
             experiment=experiment,
         )
@@ -311,7 +311,7 @@ def test_it_does_not_plot_abandoned_trials(self) -> None:
     def test_it_works_for_non_batch_experiments(self) -> None:
         # GIVEN an experiment with the default generation strategy
         experiment = get_branin_experiment(with_batch=False)
-        generation_strategy = choose_generation_strategy(
+        generation_strategy = choose_generation_strategy_legacy(
             search_space=experiment.search_space,
             experiment=experiment,
         )

ax/generation_strategy/dispatch_utils.py

@@ -294,7 +294,7 @@ def calculate_num_initialization_trials(
     return max(ret, 5)
 
 
-def choose_generation_strategy(
+def choose_generation_strategy_legacy(
     search_space: SearchSpace,
     *,
     use_batch_trials: bool = False,

ax/generation_strategy/tests/test_dispatch_utils.py

@@ -21,7 +21,7 @@
     SearchSpace,
 )
 
-from ax.generation_strategy.dispatch_utils import choose_generation_strategy
+from ax.generation_strategy.dispatch_utils import choose_generation_strategy_legacy
 from ax.metrics.torchx import TorchXMetric
 from ax.runners.torchx import TorchXRunner
 from ax.service.scheduler import FailureRateExceededError, Scheduler, SchedulerOptions
@@ -87,7 +87,7 @@ def test_run_experiment_locally(self) -> None:
         scheduler = Scheduler(
             experiment=experiment,
             generation_strategy=(
-                choose_generation_strategy(
+                choose_generation_strategy_legacy(
                     search_space=experiment.search_space,
                 )
             ),
@@ -117,7 +117,7 @@ def test_stop_trials(self) -> None:
         scheduler = Scheduler(
             experiment=experiment,
             generation_strategy=(
-                choose_generation_strategy(
+                choose_generation_strategy_legacy(
                     search_space=experiment.search_space,
                 )
             ),
@@ -155,7 +155,7 @@ def test_run_experiment_locally_in_batches(self) -> None:
         scheduler = Scheduler(
             experiment=experiment,
             generation_strategy=(
-                choose_generation_strategy(
+                choose_generation_strategy_legacy(
                     search_space=experiment.search_space,
                     max_parallelism_cap=parallelism,
                 )

ax/runners/tests/test_torchx.py

@@ -55,7 +55,7 @@
     UserInputError,
 )
 from ax.exceptions.generation_strategy import MaxParallelismReachedException
-from ax.generation_strategy.dispatch_utils import choose_generation_strategy
+from ax.generation_strategy.dispatch_utils import choose_generation_strategy_legacy
 from ax.generation_strategy.generation_strategy import GenerationStrategy
 from ax.global_stopping.strategies.base import BaseGlobalStoppingStrategy
 from ax.global_stopping.strategies.improvement import constraint_satisfaction
@@ -1768,7 +1768,7 @@ def _set_generation_strategy(
             "enforce_sequential_optimization", self._enforce_sequential_optimization
         )
         if self._generation_strategy is None:
-            self._generation_strategy = choose_generation_strategy(
+            self._generation_strategy = choose_generation_strategy_legacy(
                 search_space=self.experiment.search_space,
                 optimization_config=self.experiment.optimization_config,
                 enforce_sequential_optimization=enforce_sequential_optimization,

ax/service/ax_client.py

@@ -27,7 +27,7 @@
 from ax.core.utils import get_pending_observation_features
 from ax.exceptions.constants import CHOLESKY_ERROR_ANNOTATION
 from ax.exceptions.core import SearchSpaceExhausted, UserInputError
-from ax.generation_strategy.dispatch_utils import choose_generation_strategy
+from ax.generation_strategy.dispatch_utils import choose_generation_strategy_legacy
 from ax.generation_strategy.generation_strategy import GenerationStrategy
 from ax.modelbridge.base import Adapter
 from ax.service.utils.best_point import (
@@ -75,7 +75,7 @@ def __init__(
         self.experiment = experiment
         if generation_strategy is None:
             # pyre-fixme[4]: Attribute must be annotated.
-            self.generation_strategy = choose_generation_strategy(
+            self.generation_strategy = choose_generation_strategy_legacy(
                 search_space=experiment.search_space,
                 use_batch_trials=self.arms_per_trial > 1,
                 random_seed=self.random_seed,

ax/service/managed_loop.py

@@ -1,4 +1,5 @@
 #!/usr/bin/env python3
+
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 #
 # This source code is licensed under the MIT license found in the
@@ -43,7 +44,7 @@
     UserInputError,
 )
 from ax.exceptions.generation_strategy import AxGenerationException
-from ax.generation_strategy.dispatch_utils import choose_generation_strategy
+from ax.generation_strategy.dispatch_utils import choose_generation_strategy_legacy
 from ax.generation_strategy.generation_strategy import (
     GenerationStep,
     GenerationStrategy,
@@ -352,7 +353,7 @@ def setUp(self) -> None:
             ),
             name="branin_experiment_no_impl_runner_or_metrics",
         )
-        self.sobol_MBM_GS = choose_generation_strategy(
+        self.sobol_MBM_GS = choose_generation_strategy_legacy(
             search_space=get_branin_search_space()
         )
         self.two_sobol_steps_GS = GenerationStrategy(  # Contrived GS to ensure

ax/service/tests/scheduler_test_utils.py

@@ -1,4 +1,5 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
+
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
@@ -9,7 +10,7 @@
 from ax.core.multi_type_experiment import MultiTypeExperiment
 from ax.core.objective import Objective
 from ax.core.optimization_config import OptimizationConfig
-from ax.generation_strategy.dispatch_utils import choose_generation_strategy
+from ax.generation_strategy.dispatch_utils import choose_generation_strategy_legacy
 from ax.generation_strategy.generation_strategy import (
     GenerationStep,
     GenerationStrategy,
@@ -103,7 +104,7 @@ def setUp(self) -> None:
             default_runner=None,
             name="branin_experiment_no_impl_runner_or_metrics",
         )
-        self.sobol_MBM_GS = choose_generation_strategy(
+        self.sobol_MBM_GS = choose_generation_strategy_legacy(
             search_space=get_branin_search_space()
         )
         self.two_sobol_steps_GS = GenerationStrategy(  # Contrived GS to ensure

ax/service/tests/test_scheduler.py

@@ -34,7 +34,7 @@
 from ax.core.types import ComparisonOp
 from ax.exceptions.core import ObjectNotFoundError
 from ax.exceptions.storage import JSONDecodeError, SQADecodeError, SQAEncodeError
-from ax.generation_strategy.dispatch_utils import choose_generation_strategy
+from ax.generation_strategy.dispatch_utils import choose_generation_strategy_legacy
 from ax.metrics.branin import BraninMetric
 from ax.modelbridge.registry import Generators
 from ax.models.torch.botorch_modular.surrogate import Surrogate, SurrogateSpec
@@ -1867,7 +1867,7 @@ def test_GeneratorRunGenMetadata(self) -> None:
 
     def test_UpdateGenerationStrategyIncrementally(self) -> None:
         experiment = get_branin_experiment()
-        generation_strategy = choose_generation_strategy(experiment.search_space)
+        generation_strategy = choose_generation_strategy_legacy(experiment.search_space)
         save_experiment(experiment=experiment)
         save_generation_strategy(generation_strategy=generation_strategy)
 
@@ -2258,15 +2258,17 @@ def test_AnalysisCard(self) -> None:
     def test_delete_generation_strategy(self) -> None:
         # GIVEN an experiment with a generation strategy
         experiment = get_branin_experiment()
-        generation_strategy = choose_generation_strategy(experiment.search_space)
+        generation_strategy = choose_generation_strategy_legacy(experiment.search_space)
         generation_strategy.experiment = experiment
         save_experiment(experiment)
         save_generation_strategy(generation_strategy=generation_strategy)
 
         # AND GIVEN another experiment with a generation strategy
         experiment2 = get_branin_experiment()
         experiment2.name = "experiment2"
-        generation_strategy2 = choose_generation_strategy(experiment2.search_space)
+        generation_strategy2 = choose_generation_strategy_legacy(
+            experiment2.search_space
+        )
         generation_strategy2.experiment = experiment2
         save_experiment(experiment2)
         save_generation_strategy(generation_strategy=generation_strategy2)
@@ -2288,7 +2290,7 @@ def test_delete_generation_strategy(self) -> None:
     def test_delete_generation_strategy_max_gs_to_delete(self) -> None:
         # GIVEN an experiment with a generation strategy
         experiment = get_branin_experiment()
-        generation_strategy = choose_generation_strategy(experiment.search_space)
+        generation_strategy = choose_generation_strategy_legacy(experiment.search_space)
         generation_strategy.experiment = experiment
         save_experiment(experiment)
         save_generation_strategy(generation_strategy=generation_strategy)

ax/storage/sqa_store/tests/test_sqa_store.py

@@ -21,7 +21,7 @@
     ReductionCriterion,
     SingleDiagnosticBestModelSelector,
 )
-from ax.generation_strategy.dispatch_utils import choose_generation_strategy
+from ax.generation_strategy.dispatch_utils import choose_generation_strategy_legacy
 from ax.generation_strategy.generation_node import GenerationNode
 
 from ax.generation_strategy.generation_node_input_constructors import (
@@ -204,7 +204,7 @@ def get_generation_strategy(
                 get_sobol
             )
     else:
-        gs = choose_generation_strategy(
+        gs = choose_generation_strategy_legacy(
             search_space=get_search_space(), should_deduplicate=True
         )
         if with_callable_model_kwarg: