From de06e691a6db36306b478cb2e4f32d0df598848b Mon Sep 17 00:00:00 2001
From: Madeline Scyphers <madelinescyphers@gmail.com>
Date: Mon, 22 Jul 2024 14:35:53 -0400
Subject: [PATCH] Make modular botorch config require Ax 0.3.5 and above

---
 boa/config/converters.py                      |  8 ++++++++
 tests/1unit_tests/test_generation_strategy.py |  4 ++++
 tests/integration_tests/test_cli.py           | 16 ++++++++++++++--
 3 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/boa/config/converters.py b/boa/config/converters.py
index 543139e..0b42b9d 100644
--- a/boa/config/converters.py
+++ b/boa/config/converters.py
@@ -15,6 +15,8 @@
 from ax.service.utils.instantiation import TParameterRepresentation
 from ax.service.utils.scheduler_options import SchedulerOptions
 
+from boa.utils import check_min_package_version
+
 if TYPE_CHECKING:
     from .config import BOAMetric
 
@@ -54,6 +56,12 @@ def _gen_strat_converter(gs: Optional[dict] = None) -> dict:
                 gs["steps"][i] = step
                 steps.append(step)
                 continue
+            if step["model"] == "BOTORCH_MODULAR" and not check_min_package_version("ax-platform", "0.3.5"):
+                raise ValueError(
+                    "BOTORCH_MODULAR model is not available in BOA with Ax version < 0.3.5. "
+                    "Please upgrade to a newer version of Ax."
+                )
+
             if "model_kwargs" in step:
                 if "botorch_acqf_class" in step["model_kwargs"] and not isinstance(
                     step["model_kwargs"]["botorch_acqf_class"], botorch.acquisition.AcquisitionFunction
diff --git a/tests/1unit_tests/test_generation_strategy.py b/tests/1unit_tests/test_generation_strategy.py
index faf5a06..84a8610 100644
--- a/tests/1unit_tests/test_generation_strategy.py
+++ b/tests/1unit_tests/test_generation_strategy.py
@@ -2,6 +2,7 @@
 import botorch.models
 import gpytorch.kernels
 import gpytorch.mlls
+import pytest
 from ax.modelbridge.generation_strategy import GenerationStep, GenerationStrategy
 from ax.modelbridge.registry import Models
 
@@ -47,6 +48,9 @@ def test_auto_gen_use_saasbo(saasbo_config, tmp_path):
         assert "FullyBayesian" in gs.name
 
 
+@pytest.importorskip(
+    "ax-platform", minversion="0.3.5", reason="BOTORCH_MODULAR model is not available in BOA with Ax version < 0.3.5."
+)
 def test_modular_botorch(gen_strat_modular_botorch_config, tmp_path):
     controller = Controller(
         config=gen_strat_modular_botorch_config,
diff --git a/tests/integration_tests/test_cli.py b/tests/integration_tests/test_cli.py
index 230d963..824005d 100644
--- a/tests/integration_tests/test_cli.py
+++ b/tests/integration_tests/test_cli.py
@@ -12,7 +12,6 @@
     get_trial_dir,
     load_jsonlike,
     scheduler_from_json_file,
-    scheduler_to_json_file,
     split_shell_command,
 )
 from boa.cli import main as cli_main
@@ -80,7 +79,20 @@ def test_calling_command_line_test_script_doesnt_error_out_and_produces_correct_
 # (which can customize the GP process even more)
 @pytest.mark.parametrize(
     "r_scripts_run",
-    ["r_full", "r_light", "r_streamlined", "r_streamlined_botorch_modular"],
+    [
+        "r_full",
+        "r_light",
+        "r_streamlined",
+        "r_streamlined_botorch_modular",
+        pytest.param(
+            "r_streamlined_botorch_modular",
+            marks=pytest.importorskip(
+                "ax-platform",
+                minversion="0.3.5",
+                reason="BOTORCH_MODULAR model is not available in BOA with Ax version < 0.3.5.",
+            ),
+        ),
+    ],
 )
 @pytest.mark.skipif(not R_INSTALLED, reason="requires R to be installed")
 def test_calling_command_line_r_test_scripts(r_scripts_run, request):