Ax Sweeper Plugin support for Python 3.11

plugins/hydra_ax_sweeper/example/conf/config.yaml

@@ -18,7 +18,9 @@ hydra:
 
       experiment:
         # Default to minimize, set to false to maximize
-        minimize: true
+        objectives:
+          result:
+            minimize: true
 
       early_stop:
         # Number of epochs without a significant improvement from

plugins/hydra_ax_sweeper/hydra_plugins/hydra_ax_sweeper/_core.py

@@ -184,7 +184,12 @@ def sweep(self, arguments: List[str]) -> None:
                 num_trials_so_far += len(list_of_trials_to_launch)
                 num_trials_left -= len(list_of_trials_to_launch)
 
-                best_parameters, predictions = ax_client.get_best_parameters()
+                result = ax_client.get_best_parameters()
+                if result is None:
+                    log.info("Ax has not yet found any parameters")
+                    continue
+
+                best_parameters, predictions = result
                 metric = predictions[0][ax_client.objective_name]
 
                 if self.early_stopper.should_stop(metric, best_parameters):
@@ -336,8 +341,8 @@ def create_choice_param_from_range_override(override: Override) -> Dict[str, Any
         "name": key,
         "type": "choice",
         "values": [val for val in override.sweep_iterator()],
+        "is_ordered": False,
     }
-
     return param
 
 

plugins/hydra_ax_sweeper/hydra_plugins/hydra_ax_sweeper/config.py

@@ -2,6 +2,7 @@
 from dataclasses import dataclass, field
 from typing import Any, Dict, List, Optional
 
+from ax.service.utils.instantiation import ObjectiveProperties
 from hydra.core.config_store import ConfigStore
 
 
@@ -20,11 +21,7 @@ class ExperimentConfig:
     # Experiment name
     name: Optional[str] = None
 
-    # Name of metric to optimize or null if you only return one metric
-    objective_name: str = "objective"
-
-    # Defaults to minimize, set to false to maximize
-    minimize: bool = True
+    objectives: Optional[Dict[str, ObjectiveProperties]] = None
 
     # For the remaining parameters, refer the Ax documentation: https://ax.dev/api/core.html#experiment
     parameter_constraints: Optional[List[str]] = None

plugins/hydra_ax_sweeper/setup.py

@@ -21,15 +21,14 @@
         "Programming Language :: Python :: 3.8",
         "Programming Language :: Python :: 3.9",
         "Programming Language :: Python :: 3.10",
+        "Programming Language :: Python :: 3.11",
         "Operating System :: POSIX :: Linux",
         "Operating System :: MacOS",
         "Development Status :: 4 - Beta",
     ],
     install_requires=[
-        "hydra-core>=1.1.0.dev7",
-        "ax-platform>=0.1.20,<0.2.1",  # https://github.com/facebookresearch/hydra/issues/1767
-        "torch",
-        "gpytorch<=1.8.1",  # avoid deprecation warnings. This can probably be removed when ax-platform is unpinned.
+        "hydra-core>=1.2",
+        "ax-platform>=0.3.4,<0.4",
     ],
     include_package_data=True,
 )

plugins/hydra_ax_sweeper/tests/apps/polynomial.yaml

@@ -12,7 +12,9 @@ hydra:
       max_trials: 10
 
       experiment:
-        minimize: true
+        objectives:
+          result: 
+            minimize: true
 
       early_stop:
         max_epochs_without_improvement: 2

plugins/hydra_ax_sweeper/tests/apps/polynomial_with_coefficients.yaml

@@ -1,16 +1,19 @@
+---
 defaults:
-  - override hydra/sweeper: ax
+    - override hydra/sweeper: ax
 
 polynomial:
-  coefficients: ???
+    coefficients: ???
 
 hydra:
-  sweeper:
-    ax_config:
-      max_trials: 10
+    sweeper:
+        ax_config:
+            max_trials: 10
 
-      experiment:
-        minimize: true
+            experiment:
+                objectives:
+                    result:
+                        minimize: true
 
-      early_stop:
-        max_epochs_without_improvement: 2
+            early_stop:
+                max_epochs_without_improvement: 2

plugins/hydra_ax_sweeper/tests/apps/polynomial_with_constraint.yaml

@@ -1,25 +1,27 @@
+---
 defaults:
-  - override hydra/sweeper: ax
+    - override hydra/sweeper: ax
 
 polynomial:
-  x: ???
-  y: ???
-  z: ???
+    x: ???
+    y: ???
+    z: ???
 
 hydra:
-  sweeper:
-    ax_config:
-      max_trials: 10
+    sweeper:
+        ax_config:
+            max_trials: 10
 
-      experiment:
-        minimize: true
-        objective_name: result
-        outcome_constraints: ['constraint_metric >= 2.1']  # Optional.
+            experiment:
+                objectives:
+                    result:
+                        minimize: true
+                outcome_constraints: [constraint_metric >= 2.1]  # Optional.
 
-      early_stop:
-        max_epochs_without_improvement: 2
+            early_stop:
+                max_epochs_without_improvement: 2
 
-      params:
-        polynomial.x:
-          type: range
-          bounds: [-1, 1]
+            params:
+                polynomial.x:
+                    type: range
+                    bounds: [-1, 1]

plugins/hydra_ax_sweeper/tests/apps/polynomial_with_dict_coefficients.yaml

@@ -1,16 +1,19 @@
+---
 defaults:
-  - override hydra/sweeper: ax
+    - override hydra/sweeper: ax
 
 polynomial:
-  coefficients: ???
+    coefficients: ???
 
 hydra:
-  sweeper:
-    ax_config:
-      max_trials: 10
+    sweeper:
+        ax_config:
+            max_trials: 10
 
-      experiment:
-        minimize: true
+            experiment:
+                objectives:
+                    result:
+                        minimize: true
 
-      early_stop:
-        max_epochs_without_improvement: 2
+            early_stop:
+                max_epochs_without_improvement: 2

plugins/hydra_ax_sweeper/tests/config/config.yaml

@@ -17,7 +17,9 @@ hydra:
       max_trials: 5
 
       experiment:
-        minimize: true
+        objectives:
+          result:
+            minimize: true
 
       early_stop:
         max_epochs_without_improvement: 2

plugins/hydra_ax_sweeper/tests/test_ax_sweeper_plugin.py

@@ -4,6 +4,7 @@
 from pathlib import Path
 from typing import Any, List
 
+import pytest
 from hydra.core.plugins import Plugins
 from hydra.plugins.sweeper import Sweeper
 from hydra.test_utils.test_utils import (
@@ -18,6 +19,15 @@
 
 chdir_plugin_root()
 
+os.environ["HYDRA_FULL_ERROR"] = "1"
+
+filter_ax_internal_warnings = pytest.mark.filterwarnings(
+    "ignore::DeprecationWarning",
+    "ignore::FutureWarning",
+    "ignore::RuntimeWarning",
+    "ignore::linear_operator.utils.warnings.NumericalWarning",
+)
+
 
 def test_discovery() -> None:
     """
@@ -62,6 +72,7 @@ def test_chunk_method_for_invalid_inputs(n: int) -> None:
         list(chunk_func(batch, n))
 
 
+@filter_ax_internal_warnings
 def test_jobs_dirs(hydra_sweep_runner: TSweepRunner) -> None:
     # Verify that the spawned jobs are not overstepping the directories of one another.
     sweep = hydra_sweep_runner(
@@ -144,6 +155,7 @@ def test_jobs_configured_via_cmd(
         assert math.isclose(best_parameters["quadratic.y"], -2.0, abs_tol=1e-4)
 
 
+@filter_ax_internal_warnings
 def test_jobs_configured_via_cmd_and_config(hydra_sweep_runner: TSweepRunner) -> None:
     sweep = hydra_sweep_runner(
         calling_file="tests/test_ax_sweeper_plugin.py",
@@ -169,6 +181,7 @@ def test_jobs_configured_via_cmd_and_config(hydra_sweep_runner: TSweepRunner) ->
         assert math.isclose(best_parameters["quadratic.y"], 1.0, abs_tol=1e-4)
 
 
+@filter_ax_internal_warnings
 def test_configuration_set_via_cmd_and_default_config(
     hydra_sweep_runner: TSweepRunner,
 ) -> None:
@@ -222,7 +235,7 @@ def test_ax_logging(tmpdir: Path, cmd_arg: str, expected_str: str) -> None:
         "polynomial.z=10",
         "hydra.sweeper.ax_config.max_trials=2",
     ] + [cmd_arg]
-    result, _ = run_python_script(cmd)
+    result, _ = run_python_script(cmd, allow_warnings=True)
     assert "polynomial.x: range=[-5.0, -2.0]" in result
     assert expected_str in result
     assert "polynomial.z: fixed=10" in result
@@ -243,7 +256,7 @@ def test_search_space_exhausted_exception(tmpdir: Path, cmd_args: List[str]) ->
         "hydra.job.chdir=True",
         "hydra.sweeper.ax_config.max_trials=2",
     ] + cmd_args
-    run_python_script(cmd)
+    run_python_script(cmd, allow_warnings=True)
 
 
 @mark.parametrize(
@@ -264,7 +277,7 @@ def test_search_space_with_constraint_metric(tmpdir: Path, cmd_args: List[str])
         "hydra.job.chdir=True",
         "hydra.sweeper.ax_config.max_trials=2",
     ] + cmd_args
-    results, _ = run_python_script(cmd)
+    results, _ = run_python_script(cmd, allow_warnings=True)
 
 
 @mark.parametrize(
@@ -298,7 +311,7 @@ def test_jobs_using_choice_between_lists(
         "hydra.job.chdir=True",
         "hydra.sweeper.ax_config.max_trials=3",
     ] + [cmd_arg]
-    result, _ = run_python_script(cmd)
+    result, _ = run_python_script(cmd, allow_warnings=True)
     assert f"polynomial.coefficients: {serialized_encoding}" in result
     assert f"'polynomial.coefficients': {best_coefficients}" in result
     assert f"New best value: {best_value}" in result
@@ -335,7 +348,7 @@ def test_jobs_using_choice_between_dicts(
         "hydra.job.chdir=True",
         "hydra.sweeper.ax_config.max_trials=3",
     ] + [cmd_arg]
-    result, _ = run_python_script(cmd)
+    result, _ = run_python_script(cmd, allow_warnings=True)
     assert f"polynomial.coefficients: {serialized_encoding}" in result
     assert f"'+polynomial.coefficients': {best_coefficients}" in result
     assert f"New best value: {best_value}" in result
@@ -351,6 +364,6 @@ def test_example_app(tmpdir: Path) -> None:
         "banana.y=interval(-5, 10.1)",
         "hydra.sweeper.ax_config.max_trials=2",
     ]
-    result, _ = run_python_script(cmd)
+    result, _ = run_python_script(cmd, allow_warnings=True)
     assert "banana.x: range=[-5, 5]" in result
     assert "banana.y: range=[-5.0, 10.1]" in result