Migrate GS usages to new module location (#3311)

Summary:
Pull Request resolved: #3311

Removing the following old files from `ax/modelbridge/` in favor of the new `ax/generation_strategy` directory.

```
best_model_selector.py
dispatch_utils.py
external_generation_node.py
generation_node_input_constructors.py
generation_node.py
generation_strategy.py
model_spec.py
transition_criterion.py
```

Differential Revision: D68645075

ax/analysis/healthcheck/constraints_feasibility.py

@@ -23,8 +23,8 @@
 from ax.core.generation_strategy_interface import GenerationStrategyInterface
 from ax.core.optimization_config import OptimizationConfig
 from ax.exceptions.core import UserInputError
+from ax.generation_strategy.generation_strategy import GenerationStrategy
 from ax.modelbridge.base import Adapter
-from ax.modelbridge.generation_strategy import GenerationStrategy
 from ax.modelbridge.transforms.derelativize import Derelativize
 from pyre_extensions import assert_is_instance, none_throws
 

ax/analysis/healthcheck/tests/test_constraints_feasibility.py

@@ -23,10 +23,10 @@
 from ax.core.objective import Objective
 from ax.core.optimization_config import OptimizationConfig
 from ax.exceptions.core import UserInputError
+from ax.generation_strategy.generation_node import GenerationNode
+from ax.generation_strategy.generation_strategy import GenerationStrategy
+from ax.generation_strategy.model_spec import GeneratorSpec
 from ax.modelbridge.factory import get_sobol
-from ax.modelbridge.generation_node import GenerationNode
-from ax.modelbridge.generation_strategy import GenerationStrategy
-from ax.modelbridge.model_spec import GeneratorSpec
 from ax.modelbridge.registry import Generators
 from ax.utils.common.testutils import TestCase
 from ax.utils.testing.core_stubs import (

ax/analysis/plotly/arm_effects/insample_effects.py

@@ -22,8 +22,8 @@
 from ax.core.generator_run import GeneratorRun
 from ax.core.outcome_constraint import OutcomeConstraint
 from ax.exceptions.core import DataRequiredError, UserInputError
+from ax.generation_strategy.generation_strategy import GenerationStrategy
 from ax.modelbridge.base import Adapter
-from ax.modelbridge.generation_strategy import GenerationStrategy
 from ax.modelbridge.registry import Generators
 from ax.modelbridge.transforms.derelativize import Derelativize
 from ax.utils.common.logger import get_logger

ax/analysis/plotly/arm_effects/predicted_effects.py

@@ -23,8 +23,8 @@
 from ax.core.experiment import Experiment
 from ax.core.generation_strategy_interface import GenerationStrategyInterface
 from ax.exceptions.core import UserInputError
+from ax.generation_strategy.generation_strategy import GenerationStrategy
 from ax.modelbridge.base import Adapter
-from ax.modelbridge.generation_strategy import GenerationStrategy
 from ax.modelbridge.transforms.derelativize import Derelativize
 from pyre_extensions import assert_is_instance, none_throws
 

ax/analysis/plotly/cross_validation.py

@@ -14,8 +14,8 @@
 from ax.core.experiment import Experiment
 from ax.core.generation_strategy_interface import GenerationStrategyInterface
 from ax.exceptions.core import UserInputError
+from ax.generation_strategy.generation_strategy import GenerationStrategy
 from ax.modelbridge.cross_validation import cross_validate
-from ax.modelbridge.generation_strategy import GenerationStrategy
 from plotly import express as px, graph_objects as go
 from pyre_extensions import assert_is_instance, none_throws
 

ax/analysis/plotly/surface/contour.py

@@ -22,8 +22,8 @@
 from ax.core.generation_strategy_interface import GenerationStrategyInterface
 from ax.core.observation import ObservationFeatures
 from ax.exceptions.core import UserInputError
+from ax.generation_strategy.generation_strategy import GenerationStrategy
 from ax.modelbridge.base import Adapter
-from ax.modelbridge.generation_strategy import GenerationStrategy
 from plotly import graph_objects as go
 from pyre_extensions import none_throws
 

ax/analysis/plotly/surface/slice.py

@@ -22,8 +22,8 @@
 from ax.core.generation_strategy_interface import GenerationStrategyInterface
 from ax.core.observation import ObservationFeatures
 from ax.exceptions.core import UserInputError
+from ax.generation_strategy.generation_strategy import GenerationStrategy
 from ax.modelbridge.base import Adapter
-from ax.modelbridge.generation_strategy import GenerationStrategy
 from plotly import express as px, graph_objects as go
 from pyre_extensions import none_throws
 

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.modelbridge.dispatch_utils import choose_generation_strategy
+from ax.generation_strategy.dispatch_utils import choose_generation_strategy
 from ax.modelbridge.prediction_utils import predict_at_point
 from ax.modelbridge.registry import Generators
 from ax.utils.common.testutils import TestCase

ax/benchmark/benchmark_method.py

@@ -15,7 +15,7 @@
 from ax.core.types import TParameterization
 from ax.early_stopping.strategies.base import BaseEarlyStoppingStrategy
 
-from ax.modelbridge.generation_strategy import GenerationStrategy
+from ax.generation_strategy.generation_strategy import GenerationStrategy
 from ax.service.utils.best_point_mixin import BestPointMixin
 from ax.utils.common.base import Base
 from pyre_extensions import none_throws

ax/benchmark/methods/modular_botorch.py

@@ -8,8 +8,8 @@
 from typing import Any
 
 from ax.benchmark.benchmark_method import BenchmarkMethod
-from ax.modelbridge.generation_node import GenerationStep
-from ax.modelbridge.generation_strategy import GenerationStrategy
+from ax.generation_strategy.generation_node import GenerationStep
+from ax.generation_strategy.generation_strategy import GenerationStrategy
 from ax.modelbridge.registry import Generators
 from ax.models.torch.botorch_modular.surrogate import SurrogateSpec
 from botorch.acquisition.acquisition import AcquisitionFunction

ax/benchmark/methods/sobol.py

@@ -7,7 +7,10 @@
 
 
 from ax.benchmark.benchmark_method import BenchmarkMethod
-from ax.modelbridge.generation_strategy import GenerationStep, GenerationStrategy
+from ax.generation_strategy.generation_strategy import (
+    GenerationStep,
+    GenerationStrategy,
+)
 from ax.modelbridge.registry import Generators
 
 

ax/benchmark/tests/test_benchmark.py

@@ -47,9 +47,12 @@
 from ax.core.parameter import ChoiceParameter, ParameterType, RangeParameter
 from ax.core.search_space import SearchSpace
 from ax.early_stopping.strategies.threshold import ThresholdEarlyStoppingStrategy
-from ax.modelbridge.external_generation_node import ExternalGenerationNode
-from ax.modelbridge.generation_strategy import GenerationNode, GenerationStrategy
-from ax.modelbridge.model_spec import GeneratorSpec
+from ax.generation_strategy.external_generation_node import ExternalGenerationNode
+from ax.generation_strategy.generation_strategy import (
+    GenerationNode,
+    GenerationStrategy,
+)
+from ax.generation_strategy.model_spec import GeneratorSpec
 from ax.modelbridge.registry import Generators
 from ax.service.utils.scheduler_options import TrialType
 from ax.storage.json_store.load import load_experiment

ax/modelbridge/tests/test_aepsych_criterion.py → ax/generation_strategy/tests/test_aepsych_criterion.py

@@ -10,9 +10,15 @@
 import pandas as pd
 from ax.core.base_trial import TrialStatus
 from ax.core.data import Data
-from ax.modelbridge.generation_strategy import GenerationStep, GenerationStrategy
+from ax.generation_strategy.generation_strategy import (
+    GenerationStep,
+    GenerationStrategy,
+)
+from ax.generation_strategy.transition_criterion import (
+    MinimumPreferenceOccurances,
+    MinTrials,
+)
 from ax.modelbridge.registry import Generators
-from ax.modelbridge.transition_criterion import MinimumPreferenceOccurances, MinTrials
 from ax.utils.common.testutils import TestCase
 from ax.utils.testing.core_stubs import get_experiment
 

ax/modelbridge/tests/test_best_model_selector.py → ax/generation_strategy/tests/test_best_model_selector.py

@@ -9,11 +9,11 @@
 from unittest.mock import Mock, patch
 
 from ax.exceptions.core import UserInputError
-from ax.modelbridge.best_model_selector import (
+from ax.generation_strategy.best_model_selector import (
     ReductionCriterion,
     SingleDiagnosticBestModelSelector,
 )
-from ax.modelbridge.model_spec import GeneratorSpec
+from ax.generation_strategy.model_spec import GeneratorSpec
 from ax.modelbridge.registry import Generators
 from ax.utils.common.testutils import TestCase
 
@@ -85,10 +85,10 @@ def test_SingleDiagnosticBestModelSelector_model_cv_kwargs(self) -> None:
         for ms in self.model_specs:
             ms._fitted_model = Mock()
         with patch(
-            "ax.modelbridge.model_spec.cross_validate",
+            "ax.generation_strategy.model_spec.cross_validate",
             return_value=Mock(),
         ) as mock_cv, patch(
-            "ax.modelbridge.model_spec.compute_diagnostics",
+            "ax.generation_strategy.model_spec.compute_diagnostics",
             side_effect=self.diagnostics,
         ):
             # Max/mean picks index 2 since it has the largest mean (0.55 vs 0.1 & 0.2).

ax/modelbridge/tests/test_dispatch_utils.py → ax/generation_strategy/tests/test_dispatch_utils.py

@@ -13,7 +13,7 @@
 import torch
 from ax.core.objective import MultiObjective
 from ax.core.optimization_config import MultiObjectiveOptimizationConfig
-from ax.modelbridge.dispatch_utils import (
+from ax.generation_strategy.dispatch_utils import (
     _make_botorch_step,
     calculate_num_initialization_trials,
     choose_generation_strategy,

ax/modelbridge/tests/test_external_generation_node.py → ax/generation_strategy/tests/test_external_generation_node.py

@@ -14,8 +14,8 @@
 from ax.core.observation import ObservationFeatures
 from ax.core.types import TParameterization
 from ax.exceptions.core import UnsupportedError
-from ax.modelbridge.external_generation_node import ExternalGenerationNode
-from ax.modelbridge.generation_strategy import GenerationStrategy
+from ax.generation_strategy.external_generation_node import ExternalGenerationNode
+from ax.generation_strategy.generation_strategy import GenerationStrategy
 from ax.modelbridge.random import RandomAdapter
 from ax.utils.common.testutils import TestCase
 from ax.utils.testing.core_stubs import (

ax/modelbridge/tests/test_generation_node.py → ax/generation_strategy/tests/test_generation_node.py

@@ -14,23 +14,26 @@
 from ax.core.base_trial import TrialStatus
 from ax.core.observation import ObservationFeatures
 from ax.exceptions.core import UserInputError
-from ax.modelbridge.best_model_selector import (
+from ax.generation_strategy.best_model_selector import (
     ReductionCriterion,
     SingleDiagnosticBestModelSelector,
 )
-from ax.modelbridge.factory import get_sobol
-from ax.modelbridge.generation_node import (
+from ax.generation_strategy.generation_node import (
     GenerationNode,
     GenerationStep,
     MISSING_MODEL_SELECTOR_MESSAGE,
 )
-from ax.modelbridge.generation_node_input_constructors import (
+from ax.generation_strategy.generation_node_input_constructors import (
     InputConstructorPurpose,
     NodeInputConstructors,
 )
-from ax.modelbridge.model_spec import FactoryFunctionGeneratorSpec, GeneratorSpec
+from ax.generation_strategy.model_spec import (
+    FactoryFunctionGeneratorSpec,
+    GeneratorSpec,
+)
+from ax.generation_strategy.transition_criterion import MinTrials
+from ax.modelbridge.factory import get_sobol
 from ax.modelbridge.registry import Generators
-from ax.modelbridge.transition_criterion import MinTrials
 from ax.utils.common.constants import Keys
 from ax.utils.common.logger import get_logger
 from ax.utils.common.testutils import TestCase

ax/modelbridge/tests/test_generation_node_input_constructors.py → ax/generation_strategy/tests/test_generation_node_input_constructors.py

@@ -16,13 +16,13 @@
 from ax.core.generator_run import GeneratorRun
 from ax.core.observation import ObservationFeatures
 from ax.exceptions.generation_strategy import AxGenerationException
-from ax.modelbridge.generation_node import GenerationNode
-from ax.modelbridge.generation_node_input_constructors import (
+from ax.generation_strategy.generation_node import GenerationNode
+from ax.generation_strategy.generation_node_input_constructors import (
     InputConstructorPurpose,
     NodeInputConstructors,
 )
-from ax.modelbridge.generation_strategy import GenerationStrategy
-from ax.modelbridge.model_spec import GeneratorSpec
+from ax.generation_strategy.generation_strategy import GenerationStrategy
+from ax.generation_strategy.model_spec import GeneratorSpec
 from ax.modelbridge.registry import Generators
 from ax.utils.common.constants import Keys
 from ax.utils.common.testutils import TestCase

ax/modelbridge/tests/test_generation_strategy.py → ax/generation_strategy/tests/test_generation_strategy.py

@@ -29,16 +29,24 @@
     GenerationStrategyRepeatedPoints,
     MaxParallelismReachedException,
 )
-from ax.modelbridge.best_model_selector import SingleDiagnosticBestModelSelector
-from ax.modelbridge.discrete import DiscreteAdapter
-from ax.modelbridge.factory import get_sobol
-from ax.modelbridge.generation_node import GenerationNode
-from ax.modelbridge.generation_node_input_constructors import (
+from ax.generation_strategy.best_model_selector import SingleDiagnosticBestModelSelector
+from ax.generation_strategy.generation_node import GenerationNode
+from ax.generation_strategy.generation_node_input_constructors import (
     InputConstructorPurpose,
     NodeInputConstructors,
 )
-from ax.modelbridge.generation_strategy import GenerationStep, GenerationStrategy
-from ax.modelbridge.model_spec import GeneratorSpec
+from ax.generation_strategy.generation_strategy import (
+    GenerationStep,
+    GenerationStrategy,
+)
+from ax.generation_strategy.model_spec import GeneratorSpec
+from ax.generation_strategy.transition_criterion import (
+    AutoTransitionAfterGen,
+    MaxGenerationParallelism,
+    MinTrials,
+)
+from ax.modelbridge.discrete import DiscreteAdapter
+from ax.modelbridge.factory import get_sobol
 from ax.modelbridge.random import RandomAdapter
 from ax.modelbridge.registry import (
     _extract_model_state_after_gen,
@@ -48,11 +56,6 @@
     MODEL_KEY_TO_MODEL_SETUP,
 )
 from ax.modelbridge.torch import TorchAdapter
-from ax.modelbridge.transition_criterion import (
-    AutoTransitionAfterGen,
-    MaxGenerationParallelism,
-    MinTrials,
-)
 from ax.models.random.sobol import SobolGenerator
 from ax.utils.common.constants import Keys
 from ax.utils.common.equality import same_elements
@@ -78,7 +81,7 @@ class TestGenerationStrategyWithoutAdapterMocks(TestCase):
 
     @mock_botorch_optimize
     @patch(
-        "ax.modelbridge.generation_node._extract_model_state_after_gen",
+        "ax.generation_strategy.generation_node._extract_model_state_after_gen",
         wraps=_extract_model_state_after_gen,
     )
     def test_with_model_selection(self, mock_model_state: Mock) -> None:
@@ -806,7 +809,9 @@ def test_deduplication(self) -> None:
 
         with self.assertRaisesRegex(
             GenerationStrategyRepeatedPoints, "exceeded `MAX_GEN_DRAWS`"
-        ), mock.patch("ax.modelbridge.generation_node.logger.info") as mock_logger:
+        ), mock.patch(
+            "ax.generation_strategy.generation_node.logger.info"
+        ) as mock_logger:
             g = sobol.gen(exp)
         self.assertEqual(mock_logger.call_count, 5)
         self.assertIn(

ax/modelbridge/tests/test_transition_criterion.py → ax/generation_strategy/tests/test_transition_criterion.py

@@ -15,14 +15,13 @@
 from ax.core.base_trial import TrialStatus
 from ax.core.data import Data
 from ax.exceptions.core import UserInputError
-from ax.modelbridge.generation_strategy import (
+from ax.generation_strategy.generation_strategy import (
     GenerationNode,
     GenerationStep,
     GenerationStrategy,
 )
-from ax.modelbridge.model_spec import GeneratorSpec
-from ax.modelbridge.registry import Generators
-from ax.modelbridge.transition_criterion import (
+from ax.generation_strategy.model_spec import GeneratorSpec
+from ax.generation_strategy.transition_criterion import (
     AutoTransitionAfterGen,
     AuxiliaryExperimentCheck,
     IsSingleObjective,
@@ -32,6 +31,7 @@
     MinimumTrialsInStatus,
     MinTrials,
 )
+from ax.modelbridge.registry import Generators
 from ax.utils.common.logger import get_logger
 from ax.utils.common.testutils import TestCase
 from ax.utils.testing.core_stubs import (