Testing/update ax dep (#1291)

* Update get_MTGP for Ax >= 0.3.4

* Re-enable test_persistent_gp_multitask_ax.py on all Python versions

* Fix is_ok() error

* add line

* Re-add restriction on Ax build for python 3.11+

* Try Ax again in later python versions

* Do not install gpcam on python 3.12

* fix

* fix again

* Fix gpcam install

* Dont do tasmanian on 3.12

* and tasmanian_async

.github/workflows/extra.yml

@@ -114,6 +114,10 @@ jobs:
             conda env update --file install/gen_deps_environment.yml
 
         - name: Install ax-platform
+          run: |
+            pip install ax-platform
+
+        - name: Install gpcam
           if: matrix.python-version != '3.12'
           run: |
             pip install gpcam
@@ -212,7 +216,12 @@ jobs:
             rm ./libensemble/tests/regression_tests/test_ytopt_heffte.py
             rm ./libensemble/tests/regression_tests/test_gpCAM.py
             rm ./libensemble/tests/regression_tests/test_persistent_gp.py
-            rm ./libensemble/tests/regression_tests/test_persistent_gp_multitask_ax.py
+
+        - name: Remove test for persistent Tasmanian on Python 3.12
+          if: matrix.python-version >= '3.12'
+          run: |
+            rm ./libensemble/tests/regression_tests/test_persistent_tasmanian.py
+            rm ./libensemble/tests/regression_tests/test_persistent_tasmanian_async.py
 
         - name: Remove Balsam/Globus-compute tests on Pydantic 2
           if: matrix.pydantic-version == '2.6.4'

libensemble/gen_funcs/persistent_ax_multitask.py

@@ -23,11 +23,18 @@
 from ax.core.optimization_config import OptimizationConfig
 from ax.core.parameter import ParameterType, RangeParameter
 from ax.core.search_space import SearchSpace
-from ax.modelbridge.factory import get_MTGP, get_sobol
+from ax.modelbridge.factory import get_sobol
 from ax.runners import SyntheticRunner
 from ax.storage.json_store.save import save_experiment
 from ax.storage.metric_registry import register_metric
 from ax.storage.runner_registry import register_runner
+from ax.utils.common.result import Ok
+
+try:
+    from ax.modelbridge.factory import get_MTGP
+except ImportError:
+    # For Ax >= 0.3.4
+    from ax.modelbridge.factory import get_MTGP_LEGACY as get_MTGP
 
 from libensemble.message_numbers import EVAL_GEN_TAG, FINISHED_PERSISTENT_GEN_TAG, PERSIS_STOP, STOP_TAG
 from libensemble.tools.persistent_support import PersistentSupport
@@ -245,7 +252,8 @@ def fetch_trial_data(self, trial):
                     "sem": 0.0,
                 }
             )
-        return Data(df=pd.DataFrame.from_records(records))
+        data = Data(df=pd.DataFrame.from_records(records))
+        return Ok(data)
 
 
 def max_utility_from_GP(n, m, gr, hifi_task):