From df4afbfd961c3b721304e8f3e62fddd829d72ad2 Mon Sep 17 00:00:00 2001 From: Elizabeth Santorella Date: Mon, 16 Mar 2026 09:24:16 -0700 Subject: [PATCH] Update skip_fit_gpytorch_mll mock to patch at call site Summary: `FitGPyTorchMLL` is going away, so a mock needs to be updated. Note that `skip_fit_gpytorch_mll_context_manager` will now only work with MBM. It used to mock a function used by `fit_gpytorch_mll`, so it would work whereever `fit_gpytorch_mll` was called; now does so by mocking `fit_gpytorch_mll` itself, through its import at `ax.generators.torch.botorch_modular.utils`. Is this OK for current use cases? There are only two places in Ax that `fit_gpytorch_mll` is used: One in `botorch_modular/utils`, where this will work, and one in `tutorials/modular_botorch/modular_botorch.ipynb`, where this won't work since `fit_gpytorch_mll` is imported from `botorch.fit`. So I updated the tutorial to import from `ax.generators.torch.botorch_modular.utils` Reviewed By: saitcakmak, mpolson64 Differential Revision: D96645791 --- ax/utils/testing/mock.py | 3 ++- tutorials/modular_botorch/modular_botorch.ipynb | 5 +++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/ax/utils/testing/mock.py b/ax/utils/testing/mock.py index 905ec88cb5a..0593143518a 100644 --- a/ax/utils/testing/mock.py +++ b/ax/utils/testing/mock.py @@ -142,7 +142,8 @@ def skip_fit_gpytorch_mll_context_manager() -> Generator[None, None, None]: This should only be used to speed up slow tests. """ with mock.patch( - "botorch.fit.FitGPyTorchMLL", side_effect=lambda *args, **kwargs: args[0] + "ax.generators.torch.botorch_modular.utils.fit_gpytorch_mll", + side_effect=lambda mll, **kwargs: mll.eval(), ) as mock_fit: yield if mock_fit.call_count < 1: diff --git a/tutorials/modular_botorch/modular_botorch.ipynb b/tutorials/modular_botorch/modular_botorch.ipynb index aeb05f8d335..50a024955c4 100644 --- a/tutorials/modular_botorch/modular_botorch.ipynb +++ b/tutorials/modular_botorch/modular_botorch.ipynb @@ -465,9 +465,10 @@ "metadata": {}, "outputs": [], "source": [ - "from botorch.fit import fit_gpytorch_mll\n", "from gpytorch.mlls.marginal_log_likelihood import MarginalLogLikelihood\n", - "from ax.generators.torch.botorch_modular.utils import fit_botorch_model\n", + "from ax.generators.torch.botorch_modular.utils import (\n", + " fit_botorch_model, fit_gpytorch_mll\n", + ")\n", "\n", "@fit_botorch_model.register(SimpleCustomGP)\n", "def _fit_botorch_model_test(\n",