Speed up test_ax_sweep_impl_minimal by mocking Sobol sensitivity MC samples

ax/utils/testing/mock.py

@@ -11,6 +11,7 @@
 from typing import Any
 from unittest import mock
 
+from ax.utils.sensitivity import sobol_measures
 from botorch.fit import fit_fully_bayesian_model_nuts
 from botorch.optim.optimize_mixed import optimize_acqf_mixed_alternating
 from botorch.test_utils.mock import mock_optimize_context_manager
@@ -135,6 +136,58 @@ def inner(*args: Any, **kwargs: Any) -> Any:
     return inner
 
 
+@contextmanager
+def mock_sensitivity_analysis_context_manager(
+    num_mc_samples: int = 32,
+) -> Generator[None, None, None]:
+    """A context manager that shrinks the Monte-Carlo sample count used by
+    variance-based sensitivity analysis (Sobol indices).
+
+    ``ax_parameter_sens`` defaults to ``num_mc_samples=10**4``, which means
+    ``O(num_mc_samples * num_params)`` posterior-mean evaluations per call. On a
+    fully Bayesian model each of those is multiplied by the number of MCMC
+    samples, so a single call can take tens of seconds even on a handful of
+    training points. Tests that exercise reporting/analysis code paths
+    (``get_standard_plots``, ``SensitivityAnalysisPlot``, ``OverviewAnalysis``)
+    pay this cost once per call and are dominated by it.
+
+    Note that ``mock_botorch_optimize`` does not help here: it makes *fitting*
+    cheap, not posterior evaluation.
+
+    Do not use this in tests that assert on sensitivity values, since a small
+    sample count makes the estimates noisy.
+
+    Args:
+        num_mc_samples: The Monte-Carlo sample count to force.
+    """
+    with ExitStack() as es:
+        for fn_name in (
+            "compute_sobol_indices_from_model_list",
+            "compute_derivatives_from_model_list",
+        ):
+            original = getattr(sobol_measures, fn_name)
+
+            def few_samples(
+                *args: Any, __original: Any = original, **kwargs: Any
+            ) -> Tensor:
+                kwargs.setdefault("num_mc_samples", num_mc_samples)
+                return __original(*args, **kwargs)
+
+            es.enter_context(mock.patch.object(sobol_measures, fn_name, few_samples))
+        yield
+
+
+def mock_sensitivity_analysis(f: Callable) -> Callable:
+    """Wraps `f` in `mock_sensitivity_analysis_context_manager` as a decorator."""
+
+    @wraps(f)
+    def inner(*args: Any, **kwargs: Any) -> Any:
+        with mock_sensitivity_analysis_context_manager():
+            return f(*args, **kwargs)
+
+    return inner
+
+
 @contextmanager
 def skip_fit_gpytorch_mll_context_manager() -> Generator[None, None, None]:
     """A context manager that makes `fit_gpytorch_mll` a no-op.