Get rid of SurrogateRunner and SurrogateBenchmarkProblem (#2954)

Summary:

**Context**:  This diff cuts over over uses of `SurrogateRunner` to use `ParamBasedTestProblemRunner` with a `test_problem` that is the newly introduced `SurrogateTestFunction`, and the following diff after that will bring us down to only one runner class for benchmarking by merging `ParamBasedTestProblemRunner` into `BenchmarkRunner`. Having only one runner will make it easier to enable asynchronous benchmarks. Currently, SurrogateRunner had its own logic for tracking when trials are completed, which would make it difficult to work in with asynchronicity.

**Note on naming**: Some names have become non-intuitive in the process of benchmarking. To accord with some future changes I hope to make, I called a new class SurrogateTestFunction, whereas SurrogateParamBasedTestProblem would be more in line with the old naming.

The name changes I hope to make:
* ParamBasedTestProblemRunner -> nothing, absorbed into BenchmarkRunner
* ParamBasedTestProblem -> TestFunction, to emphasize that all it does is generate data (rather than more generally specify the problem we are solving) and that it is deterministic, and to differentiate it from BenchmarkProblem. BenchmarkTestFunction would also be a candidate.
* BoTorchTestProblem -> BoTorchTestFunction

**Changes in this diff**:
* Introduces SurrogateTestFunction, a ParamBasedTestProblem for * Removes SurrogateRunner. There is some loss of (unused) functionality here; dict-valued 'noise_std' is no longer supported. I expect we will eventually want that, but it doesn't need to be done now.
* Changed `predict_for_tensor` utility function to use a `SurrogateTestFunction` rather than a `SurrogateRunner`
* Removed `SurrogateRunner` from the json_store registry. There is no point in keeping it around for backward compatibility since it was just deserializing to a dummy `SyntheticRunner` anyway.

Differential Revision: D64855010

ax/benchmark/tests/runners/test_botorch_test_problem.py

@@ -150,7 +150,7 @@ def test_synthetic_runner(self) -> None:
                     outcome_names = ["objective_0"]
             elif isinstance(test_problem, TestParamBasedTestProblem):
                 outcome_names = [f"objective_{i}" for i in range(num_outcomes)]
-            elif isinstance(test_problem, SurrogateTestFunction):
+            else:
                 outcome_names = ["branin"]
 
             # Set up runner

ax/benchmark/tests/runners/test_surrogate_runner.py

@@ -8,15 +8,10 @@
 from unittest.mock import MagicMock, patch
 
 import torch
-from ax.benchmark.runners.surrogate import SurrogateRunner, SurrogateTestFunction
-from ax.core.parameter import ParameterType, RangeParameter
-from ax.core.search_space import SearchSpace
+from ax.benchmark.runners.surrogate import SurrogateTestFunction
 from ax.modelbridge.torch import TorchModelBridge
 from ax.utils.common.testutils import TestCase
-from ax.utils.testing.benchmark_stubs import (
-    get_soo_surrogate_legacy,
-    get_soo_surrogate_test_function,
-)
+from ax.utils.testing.benchmark_stubs import get_soo_surrogate_test_function
 
 
 class TestSurrogateTestFunction(TestCase):
@@ -84,84 +79,3 @@ def _construct_test_function(name: str) -> SurrogateTestFunction:
         self.assertEqual(runner_1, runner_1a)
         self.assertNotEqual(runner_1, runner_2)
         self.assertNotEqual(runner_1, 1)
-
-
-class TestSurrogateRunner(TestCase):
-    def setUp(self) -> None:
-        super().setUp()
-        self.search_space = SearchSpace(
-            parameters=[
-                RangeParameter("x", ParameterType.FLOAT, 0.0, 5.0),
-                RangeParameter("y", ParameterType.FLOAT, 1.0, 10.0, log_scale=True),
-                RangeParameter("z", ParameterType.INT, 1.0, 5.0, log_scale=True),
-            ]
-        )
-
-    def test_surrogate_runner(self) -> None:
-        # Construct a search space with log-scale parameters.
-        for noise_std in (0.0, 0.1, {"dummy_metric": 0.2}):
-            with self.subTest(noise_std=noise_std):
-                surrogate = MagicMock()
-                mock_mean = torch.tensor([[0.1234]], dtype=torch.double)
-                surrogate.predict = MagicMock(return_value=(mock_mean, 0))
-                surrogate.device = torch.device("cpu")
-                surrogate.dtype = torch.double
-                runner = SurrogateRunner(
-                    name="test runner",
-                    _surrogate=surrogate,
-                    _datasets=[],
-                    outcome_names=["dummy_metric"],
-                    noise_stds=noise_std,
-                )
-                self.assertEqual(runner.name, "test runner")
-                self.assertIs(runner.surrogate, surrogate)
-                self.assertEqual(runner.outcome_names, ["dummy_metric"])
-                self.assertEqual(runner.noise_stds, noise_std)
-
-    def test_lazy_instantiation(self) -> None:
-        runner = get_soo_surrogate_legacy().runner
-
-        self.assertIsNone(runner._surrogate)
-        self.assertIsNone(runner._datasets)
-
-        # Accessing `surrogate` sets datasets and surrogate
-        self.assertIsInstance(runner.surrogate, TorchModelBridge)
-        self.assertIsInstance(runner._surrogate, TorchModelBridge)
-        self.assertIsInstance(runner._datasets, list)
-
-        # Accessing `datasets` also sets datasets and surrogate
-        runner = get_soo_surrogate_legacy().runner
-        self.assertIsInstance(runner.datasets, list)
-        self.assertIsInstance(runner._surrogate, TorchModelBridge)
-        self.assertIsInstance(runner._datasets, list)
-
-        with patch.object(
-            runner,
-            "get_surrogate_and_datasets",
-            wraps=runner.get_surrogate_and_datasets,
-        ) as mock_get_surrogate_and_datasets:
-            runner.surrogate
-        mock_get_surrogate_and_datasets.assert_not_called()
-
-    def test_instantiation_raises_with_missing_args(self) -> None:
-        with self.assertRaisesRegex(
-            ValueError, "If `get_surrogate_and_datasets` is None, `_surrogate` and "
-        ):
-            SurrogateRunner(name="test runner", outcome_names=[], noise_stds=0.0)
-
-    def test_equality(self) -> None:
-        def _construct_runner(name: str) -> SurrogateRunner:
-            return SurrogateRunner(
-                name=name,
-                _surrogate=MagicMock(),
-                _datasets=[],
-                outcome_names=["dummy_metric"],
-                noise_stds=0.0,
-            )
-
-        runner_1 = _construct_runner("test 1")
-        runner_2 = _construct_runner("test 2")
-        runner_1a = _construct_runner("test 1")
-        self.assertEqual(runner_1, runner_1a)
-        self.assertNotEqual(runner_1, runner_2)
-        self.assertNotEqual(runner_1, 1)

ax/storage/json_store/registry.py

@@ -370,8 +370,6 @@
     "SumConstraint": SumConstraint,
     "Surrogate": Surrogate,
     "SurrogateMetric": BenchmarkMetric,  # backward-compatiblity
-    # NOTE: SurrogateRunners -> SyntheticRunner on load due to complications
-    "SurrogateRunner": SyntheticRunner,
     "SobolQMCNormalSampler": SobolQMCNormalSampler,
     "SyntheticRunner": SyntheticRunner,
     "SurrogateSpec": SurrogateSpec,

ax/utils/testing/benchmark_stubs.py

@@ -15,12 +15,11 @@
 from ax.benchmark.benchmark_metric import BenchmarkMetric
 from ax.benchmark.benchmark_problem import BenchmarkProblem, create_problem_from_botorch
 from ax.benchmark.benchmark_result import AggregatedBenchmarkResult, BenchmarkResult
-from ax.benchmark.problems.surrogate import SurrogateBenchmarkProblem
 from ax.benchmark.runners.botorch_test import (
     ParamBasedTestProblem,
     ParamBasedTestProblemRunner,
 )
-from ax.benchmark.runners.surrogate import SurrogateRunner, SurrogateTestFunction
+from ax.benchmark.runners.surrogate import SurrogateTestFunction
 from ax.core.experiment import Experiment
 from ax.core.objective import MultiObjective, Objective
 from ax.core.optimization_config import (
@@ -129,41 +128,7 @@ def get_soo_surrogate() -> BenchmarkProblem:
     )
 
 
-def get_soo_surrogate_legacy() -> SurrogateBenchmarkProblem:
-    experiment = get_branin_experiment(with_completed_trial=True)
-    surrogate = TorchModelBridge(
-        experiment=experiment,
-        search_space=experiment.search_space,
-        model=BoTorchModel(surrogate=Surrogate(botorch_model_class=SingleTaskGP)),
-        data=experiment.lookup_data(),
-        transforms=[],
-    )
-    runner = SurrogateRunner(
-        name="test",
-        outcome_names=["branin"],
-        get_surrogate_and_datasets=lambda: (surrogate, []),
-    )
-
-    observe_noise_sd = True
-    objective = Objective(
-        metric=BenchmarkMetric(
-            name="branin", lower_is_better=True, observe_noise_sd=observe_noise_sd
-        ),
-    )
-    optimization_config = OptimizationConfig(objective=objective)
-
-    return SurrogateBenchmarkProblem(
-        name="test",
-        search_space=experiment.search_space,
-        optimization_config=optimization_config,
-        num_trials=6,
-        observe_noise_stds=observe_noise_sd,
-        optimal_value=0.0,
-        runner=runner,
-    )
-
-
-def get_moo_surrogate() -> SurrogateBenchmarkProblem:
+def get_moo_surrogate() -> BenchmarkProblem:
     experiment = get_branin_experiment_with_multi_objective(with_completed_trial=True)
     surrogate = TorchModelBridge(
         experiment=experiment,
@@ -173,11 +138,15 @@ def get_moo_surrogate() -> SurrogateBenchmarkProblem:
         transforms=[],
     )
 
-    runner = SurrogateRunner(
+    outcome_names = ["branin_a", "branin_b"]
+    test_function = SurrogateTestFunction(
         name="test",
-        outcome_names=["branin_a", "branin_b"],
+        outcome_names=outcome_names,
         get_surrogate_and_datasets=lambda: (surrogate, []),
     )
+    runner = ParamBasedTestProblemRunner(
+        test_problem=test_function, outcome_names=outcome_names
+    )
     observe_noise_sd = True
     optimization_config = MultiObjectiveOptimizationConfig(
         objective=MultiObjective(
@@ -199,7 +168,7 @@ def get_moo_surrogate() -> SurrogateBenchmarkProblem:
             ],
         )
     )
-    return SurrogateBenchmarkProblem(
+    return BenchmarkProblem(
         name="test",
         search_space=experiment.search_space,
         optimization_config=optimization_config,