Add custom method with model_config to choose_gs

Summary:
This diff adds the ability to use custom `ModelConfig` objects with `choose_generation_strategy` by introducing a new `"custom"` method option.

## Motivation

After D83062555, the internal `choose_gs_internal` function now uses the OSS `choose_generation_strategy` for default GS dispatch. However, there was no way to pass custom model configurations (like MAP SAAS) through this function - they had to be handled as special cases outside the dispatch flow.

## Solution

Rather than adding specific method types for each model variant (which would turn the Literal type into a "zoo" of options), we add a single `"custom"` method that allows passing a `ModelConfig` as a separate function argument. This keeps the API clean while providing full flexibility.

Key changes:
- Add `"custom"` to the `method` Literal type in `GenerationStrategyDispatchStruct`
- Add optional `model_config` parameter to `choose_generation_strategy()`
- Add validation: `model_config` must be provided iff `method="custom"`
- Update `_get_mbm_node()` to handle custom model configs

Differential Revision: D89906836

Author: saitcakmak

ax/api/utils/generation_strategy_dispatch.py

@@ -11,7 +11,7 @@
 from ax.adapter.registry import Generators
 from ax.api.utils.structs import GenerationStrategyDispatchStruct
 from ax.core.trial_status import TrialStatus
-from ax.exceptions.core import UnsupportedError
+from ax.exceptions.core import UnsupportedError, UserInputError
 from ax.generation_strategy.center_generation_node import CenterGenerationNode
 from ax.generation_strategy.dispatch_utils import get_derelativize_config
 from ax.generation_strategy.generation_strategy import (
@@ -22,6 +22,7 @@
 from ax.generation_strategy.transition_criterion import MinTrials
 from ax.generators.torch.botorch_modular.surrogate import ModelConfig, SurrogateSpec
 from botorch.models.fully_bayesian import SaasFullyBayesianSingleTaskGP
+from pyre_extensions import none_throws
 
 
 def _get_sobol_node(
@@ -96,17 +97,25 @@ def _get_mbm_node(
     method: str,
     torch_device: str | None,
     simplify_parameter_changes: bool,
-) -> GenerationNode:
+    model_config: ModelConfig | None = None,
+) -> tuple[GenerationNode, str]:
     """Constructs an MBM node based on the method specified in
     ``struct``.
 
     The ``SurrogateSpec`` takes the following form for the given method:
     - BALANCED: Two model configs: one with MBM defaults, the other with
         linear kernel with input warping.
     - FAST: An empty model config that utilizes MBM defaults.
+    - CUSTOM: Uses the provided ``model_config``.
     """
     # Construct the surrogate spec.
-    if method == "quality":
+    if method == "custom":
+        model_config = none_throws(model_config)
+        model_configs = [model_config]
+        mbm_name = (
+            model_config.name if model_config.name is not None else "custom_config"
+        )
+    elif method == "quality":
         model_configs = [
             ModelConfig(
                 botorch_model_class=SaasFullyBayesianSingleTaskGP,
@@ -117,8 +126,10 @@ def _get_mbm_node(
                 name="WarpedSAAS",
             )
         ]
+        mbm_name = method
     elif method == "fast":
         model_configs = [ModelConfig(name="MBM defaults")]
+        mbm_name = method
     else:
         raise UnsupportedError(f"Unsupported generation method: {method}.")
 
@@ -142,11 +153,12 @@ def _get_mbm_node(
             )
         ],
         should_deduplicate=True,
-    )
+    ), mbm_name
 
 
 def choose_generation_strategy(
     struct: GenerationStrategyDispatchStruct,
+    model_config: ModelConfig | None = None,
 ) -> GenerationStrategy:
     """
     Choose a generation strategy based on the properties of the experiment and the
@@ -159,10 +171,23 @@ def choose_generation_strategy(
         struct: A ``GenerationStrategyDispatchStruct``
             object that informs
             the choice of generation strategy.
+        model_config: An optional ``ModelConfig`` to use for the Bayesian optimization
+            phase. This must be provided when ``struct.method`` is ``"custom"``, and
+            must not be provided otherwise.
 
     Returns:
         A generation strategy.
     """
+    # Validate model_config usage.
+    if struct.method == "custom":
+        if model_config is None:
+            raise UserInputError("model_config must be provided when method='custom'.")
+    elif model_config is not None:
+        raise UserInputError(
+            "model_config should only be provided when method='custom'. "
+            f"Got method='{struct.method}'."
+        )
+
     # Handle the random search case.
     if struct.method == "random_search":
         nodes = [
@@ -178,10 +203,11 @@ def choose_generation_strategy(
         ]
         gs_name = "QuasiRandomSearch"
     else:
-        mbm_node = _get_mbm_node(
+        mbm_node, mbm_name = _get_mbm_node(
             method=struct.method,
             torch_device=struct.torch_device,
             simplify_parameter_changes=struct.simplify_parameter_changes,
+            model_config=model_config,
         )
         if (
             struct.initialization_budget is None
@@ -198,10 +224,10 @@ def choose_generation_strategy(
                 ),
                 mbm_node,
             ]
-            gs_name = f"Sobol+MBM:{struct.method}"
+            gs_name = f"Sobol+MBM:{mbm_name}"
         else:
             nodes = [mbm_node]
-            gs_name = f"MBM:{struct.method}"
+            gs_name = f"MBM:{mbm_name}"
     if struct.initialize_with_center and (
         struct.initialization_budget is None or struct.initialization_budget > 0
     ):

ax/api/utils/structs.py

@@ -53,6 +53,11 @@ class GenerationStrategyDispatchStruct:
                 - ``"random_search"``, primarily intended for pure exploration
                     experiments, this method utilizes quasi-random Sobol sequences
                     for candidate generation.
+                - ``"custom"``, allows using a custom ``ModelConfig`` for the
+                    Bayesian optimization phase. When using this method, the
+                    ``model_config`` argument must be provided to
+                    ``choose_generation_strategy``. This is an advanced option
+                    and should not be considered a part of the public API.
         initialization_budget: The number of trials to use for initialization.
             If ``None``, a default budget of 5 trials is used.
         initialization_random_seed: The random seed to use with the Sobol generator
@@ -87,7 +92,7 @@ class GenerationStrategyDispatchStruct:
             irrelevant parameters.
     """
 
-    method: Literal["quality", "fast", "random_search"] = "fast"
+    method: Literal["quality", "fast", "random_search", "custom"] = "fast"
     # Initialization options
     initialization_budget: int | None = None
     initialization_random_seed: int | None = None

ax/api/utils/tests/test_generation_strategy_dispatch.py

@@ -16,6 +16,7 @@
 from ax.api.utils.structs import GenerationStrategyDispatchStruct
 from ax.core.trial import Trial
 from ax.core.trial_status import TrialStatus
+from ax.exceptions.core import UserInputError
 from ax.generation_strategy.center_generation_node import CenterGenerationNode
 from ax.generation_strategy.dispatch_utils import get_derelativize_config
 from ax.generation_strategy.transition_criterion import MinTrials
@@ -28,6 +29,7 @@
 from ax.utils.testing.mock import mock_botorch_optimize
 from ax.utils.testing.utils import run_trials_with_gs
 from botorch.models.fully_bayesian import SaasFullyBayesianSingleTaskGP
+from botorch.models.map_saas import EnsembleMapSaasSingleTaskGP
 from pyre_extensions import assert_is_instance, none_throws
 
 
@@ -252,3 +254,69 @@ def test_gs_simplify_parameter_changes(self) -> None:
                 mbm_spec.generator_kwargs["acquisition_options"],
                 {"prune_irrelevant_parameters": simplify},
             )
+
+    def test_choose_gs_custom_with_model_config(self) -> None:
+        """Test that custom method works with a provided ModelConfig."""
+        custom_model_config = ModelConfig(
+            botorch_model_class=EnsembleMapSaasSingleTaskGP,
+            name="MAPSAAS",
+        )
+        struct = GenerationStrategyDispatchStruct(
+            method="custom",
+            initialization_budget=3,
+            initialize_with_center=False,
+            torch_device="cpu",
+        )
+        gs = choose_generation_strategy(struct=struct, model_config=custom_model_config)
+        self.assertEqual(len(gs._nodes), 2)
+        self.assertEqual(gs.name, "Sobol+MBM:MAPSAAS")
+
+        # Check the MBM node uses the custom model config.
+        mbm_node = gs._nodes[1]
+        self.assertEqual(len(mbm_node.generator_specs), 1)
+        mbm_spec = mbm_node.generator_specs[0]
+        self.assertEqual(mbm_spec.generator_enum, Generators.BOTORCH_MODULAR)
+        expected_ss = SurrogateSpec(model_configs=[custom_model_config])
+        self.assertEqual(
+            mbm_spec.generator_kwargs["surrogate_spec"],
+            expected_ss,
+        )
+        self.assertEqual(
+            mbm_spec.generator_kwargs["torch_device"],
+            torch.device("cpu"),
+        )
+
+    def test_choose_gs_custom_without_name(self) -> None:
+        """Test that custom method works with unnamed ModelConfig."""
+        custom_model_config = ModelConfig(
+            botorch_model_class=SaasFullyBayesianSingleTaskGP,
+            # No name provided.
+        )
+        struct = GenerationStrategyDispatchStruct(
+            method="custom",
+            initialization_budget=3,
+            initialize_with_center=False,
+        )
+        gs = choose_generation_strategy(struct=struct, model_config=custom_model_config)
+        # Should use "custom_config" as the default name.
+        self.assertEqual(gs.name, "Sobol+MBM:custom_config")
+
+    def test_choose_gs_custom_model_config_validation(self) -> None:
+        """Test validation of model_config and custom method pairing."""
+        # Test that custom method raises an error when model_config is not provided.
+        struct = GenerationStrategyDispatchStruct(method="custom")
+        with self.assertRaisesRegex(
+            UserInputError,
+            "model_config must be provided when method='custom'.",
+        ):
+            choose_generation_strategy(struct=struct)
+
+        # Test that providing model_config without custom method raises an error.
+        custom_model_config = ModelConfig(name="SomeConfig")
+        struct = GenerationStrategyDispatchStruct(method="fast")
+        with self.assertRaisesRegex(
+            UserInputError,
+            "model_config should only be provided when method='custom'. "
+            "Got method='fast'.",
+        ):
+            choose_generation_strategy(struct=struct, model_config=custom_model_config)