hydra-ax-sweeper 1.3.0.dev0__tar.gz → 1.4.0.dev4__tar.gz

hydra_ax_sweeper-1.4.0.dev4/MANIFEST.in

@@ -0,0 +1,3 @@
+global-exclude *.pyc
+global-exclude __pycache__
+recursive-include hydra_plugins/* *.yaml py.typed

hydra_ax_sweeper-1.4.0.dev4/PKG-INFO

@@ -0,0 +1,35 @@
+Metadata-Version: 2.4
+Name: hydra-ax-sweeper
+Version: 1.4.0.dev4
+Summary: Hydra Ax Sweeper plugin
+Home-page: https://github.com/facebookresearch/hydra/
+Author: Omry Yadan, Shagun Sodhani
+Author-email: omry@fb.com, sshagunsodhani@gmail.com
+License: MIT
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Classifier: Operating System :: POSIX :: Linux
+Classifier: Operating System :: MacOS
+Classifier: Development Status :: 4 - Beta
+Requires-Python: >=3.11,<3.15
+Description-Content-Type: text/markdown
+Requires-Dist: hydra-core>=1.1.0.dev7
+Requires-Dist: ax-platform<1.3.0,>=1.2.4
+Requires-Dist: torch>=2.2
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: license
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
+
+# Hydra Ax Sweeper
+Provides a [`Ax Sweeper`](https://ax.dev/) based Hydra Sweeper supporting parallel execution.
+
+See [website](https://hydra.cc/docs/plugins/ax_sweeper/) for more information

hydra_ax_sweeper-1.4.0.dev4/hydra_ax_sweeper.egg-info/PKG-INFO

@@ -0,0 +1,35 @@
+Metadata-Version: 2.4
+Name: hydra-ax-sweeper
+Version: 1.4.0.dev4
+Summary: Hydra Ax Sweeper plugin
+Home-page: https://github.com/facebookresearch/hydra/
+Author: Omry Yadan, Shagun Sodhani
+Author-email: omry@fb.com, sshagunsodhani@gmail.com
+License: MIT
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Classifier: Operating System :: POSIX :: Linux
+Classifier: Operating System :: MacOS
+Classifier: Development Status :: 4 - Beta
+Requires-Python: >=3.11,<3.15
+Description-Content-Type: text/markdown
+Requires-Dist: hydra-core>=1.1.0.dev7
+Requires-Dist: ax-platform<1.3.0,>=1.2.4
+Requires-Dist: torch>=2.2
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: license
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
+
+# Hydra Ax Sweeper
+Provides a [`Ax Sweeper`](https://ax.dev/) based Hydra Sweeper supporting parallel execution.
+
+See [website](https://hydra.cc/docs/plugins/ax_sweeper/) for more information

{hydra-ax-sweeper-1.3.0.dev0 → hydra_ax_sweeper-1.4.0.dev4}/hydra_ax_sweeper.egg-info/SOURCES.txt

@@ -11,4 +11,6 @@ hydra_plugins/hydra_ax_sweeper/__init__.py
 hydra_plugins/hydra_ax_sweeper/_core.py
 hydra_plugins/hydra_ax_sweeper/_earlystopper.py
 hydra_plugins/hydra_ax_sweeper/ax_sweeper.py
-hydra_plugins/hydra_ax_sweeper/config.py
+hydra_plugins/hydra_ax_sweeper/config.py
+hydra_plugins/hydra_ax_sweeper/py.typed
+tests/test_ax_sweeper_plugin.py

hydra_ax_sweeper-1.4.0.dev4/hydra_ax_sweeper.egg-info/requires.txt

@@ -0,0 +1,3 @@
+hydra-core>=1.1.0.dev7
+ax-platform<1.3.0,>=1.2.4
+torch>=2.2

{hydra-ax-sweeper-1.3.0.dev0 → hydra_ax_sweeper-1.4.0.dev4}/hydra_plugins/hydra_ax_sweeper/__init__.py

@@ -1,3 +1,3 @@
 # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
 
-__version__ = "1.3.0.dev0"
+__version__ = "1.4.0.dev4"

{hydra-ax-sweeper-1.3.0.dev0 → hydra_ax_sweeper-1.4.0.dev4}/hydra_plugins/hydra_ax_sweeper/_core.py

@@ -1,11 +1,23 @@
 # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
 import logging
 from dataclasses import dataclass
-from typing import Any, Dict, Iterable, List, Mapping, Optional, Tuple, Union
-
+from typing import (
+    Any,
+    Dict,
+    Iterable,
+    List,
+    Literal,
+    Mapping,
+    Optional,
+    Tuple,
+    Union,
+    cast,
+)
+
+from ax.api.client import Client  # type: ignore
+from ax.api.configs import ChoiceParameterConfig, RangeParameterConfig  # type: ignore
 from ax.core import types as ax_types  # type: ignore
-from ax.exceptions.core import SearchSpaceExhausted  # type: ignore
-from ax.service.ax_client import AxClient  # type: ignore
+from ax.exceptions.core import SearchSpaceExhausted, UnsupportedError  # type: ignore
 from hydra.core.override_parser.overrides_parser import OverridesParser
 from hydra.core.override_parser.types import IntervalSweep, Override, Transformer
 from hydra.core.plugins import Plugins
@@ -19,6 +31,10 @@ from .config import AxConfig, ClientConfig, ExperimentConfig
 
 log = logging.getLogger(__name__)
 
+AxRangeParameterType = Literal["float", "int"]
+AxChoiceParameterType = Literal["float", "int", "str", "bool"]
+AxParameterConfig = Union[RangeParameterConfig, ChoiceParameterConfig]
+
 
 @dataclass
 class Trial:
@@ -34,7 +50,7 @@ class TrialBatch:
 
 def encoder_parameters_into_string(parameters: List[Dict[str, Any]]) -> str:
     """Convert a list of params into a string"""
-    mandatory_keys = set(["name", "type", "bounds", "values", "value"])
+    mandatory_keys = {"name", "type", "bounds", "values", "value"}
     parameter_log_string = ""
     for parameter in parameters:
         parameter_log_string += "\n"
@@ -52,7 +68,9 @@ def encoder_parameters_into_string(parameters: List[Dict[str, Any]]) -> str:
     return parameter_log_string
 
 
-def map_params_to_arg_list(params: Mapping[str, Union[str, float, int]]) -> List[str]:
+def map_params_to_arg_list(
+    params: Mapping[str, Union[str, float, int, bool]],
+) -> List[str]:
     """Method to map a dictionary of params to a list of string arguments"""
     arg_list = []
     for key in params:
@@ -60,45 +78,76 @@ def map_params_to_arg_list(params: Mapping[str, Union[str, float, int]]) -> List
     return arg_list
 
 
+def get_ax_choice_parameter_type(values: Iterable[Any]) -> AxChoiceParameterType:
+    value_types = {type(value) for value in values}
+    if value_types == {bool}:
+        return "bool"
+    if value_types <= {int}:
+        return "int"
+    if value_types <= {int, float}:
+        return "float"
+    if value_types == {str}:
+        return "str"
+    raise ValueError(f"Unsupported mixed Ax parameter value types: {value_types}")
+
+
+def create_ax_parameter_config(param: Dict[Any, Any]) -> AxParameterConfig:
+    name = param["name"]
+    if param["type"] == "range":
+        bounds = param["bounds"]
+        range_parameter_type: AxRangeParameterType = (
+            "int" if all(type(bound) is int for bound in bounds) else "float"
+        )
+        return RangeParameterConfig(
+            name=name,
+            bounds=cast(Tuple[float, float], tuple(bounds)),
+            parameter_type=range_parameter_type,
+            scaling="log" if param.get("log_scale") else None,
+        )
+
+    if param["type"] == "choice":
+        values = param["values"]
+        choice_parameter_type = get_ax_choice_parameter_type(values)
+        is_ordered = param.get("is_ordered", choice_parameter_type != "str")
+    elif param["type"] == "fixed":
+        values = [param["value"]]
+        choice_parameter_type = get_ax_choice_parameter_type(values)
+        is_ordered = None
+    else:
+        raise ValueError(f"Unexpected Ax parameter type: {param['type']}")
+
+    if choice_parameter_type == "float":
+        values = [float(value) for value in values]
+    return ChoiceParameterConfig(
+        name=name,
+        values=cast(Any, values),
+        parameter_type=choice_parameter_type,
+        is_ordered=is_ordered,
+    )
+
+
 def get_one_batch_of_trials(
-    ax_client: AxClient,
-    parallelism: Tuple[int, int],
-    num_trials_so_far: int,
+    ax_client: Client,
     num_max_trials_to_do: int,
 ) -> TrialBatch:
     """Returns a TrialBatch that contains a list of trials that can be
     run in parallel. TrialBatch also flags if the search space is exhausted."""
 
-    is_search_space_exhausted = False
-    # Ax throws an exception if the search space is exhausted. We catch
-    # the exception and set the flag to True
-    (num_trials, max_parallelism_setting) = parallelism
-    if max_parallelism_setting == -1:
-        # Special case, we can group all the trials into one batch
-        max_parallelism_setting = num_trials - num_trials_so_far
-
-        if num_trials == -1:
-            # This is a special case where we can run as many trials in parallel as we want.
-            # Given that num_trials is also -1, we can run all the trials in parallel.
-            max_parallelism_setting = num_max_trials_to_do
-
-    list_of_trials = []
-    for _ in range(max_parallelism_setting):
-        try:
-            parameters, trial_index = ax_client.get_next_trial()
-            list_of_trials.append(
-                Trial(
-                    overrides=map_params_to_arg_list(params=parameters),
-                    trial_index=trial_index,
-                )
-            )
-        except SearchSpaceExhausted:
-            is_search_space_exhausted = True
-            break
+    try:
+        trials = ax_client.get_next_trials(max_trials=num_max_trials_to_do)
+    except SearchSpaceExhausted:
+        return TrialBatch(list_of_trials=[], is_search_space_exhausted=True)
 
+    list_of_trials = [
+        Trial(
+            overrides=map_params_to_arg_list(params=parameters),
+            trial_index=trial_index,
+        )
+        for trial_index, parameters in trials.items()
+    ]
     return TrialBatch(
         list_of_trials=list_of_trials,
-        is_search_space_exhausted=is_search_space_exhausted,
+        is_search_space_exhausted=len(list_of_trials) == 0,
     )
 
 
@@ -146,56 +195,40 @@ class CoreAxSweeper(Sweeper):
         ax_client = self.setup_ax_client(arguments)
 
         num_trials_left = self.max_trials
-        max_parallelism = ax_client.get_max_parallelism()
-        current_parallelism_index = 0
-        # Index to track the parallelism value we are using right now.
         is_search_space_exhausted = False
         # Ax throws an exception if the search space is exhausted. We catch
         # the exception and set the flag to True
 
         best_parameters = {}
         while num_trials_left > 0 and not is_search_space_exhausted:
-            current_parallelism = max_parallelism[current_parallelism_index]
-            num_trials, max_parallelism_setting = current_parallelism
-            num_trials_so_far = 0
-            while (
-                num_trials > num_trials_so_far or num_trials == -1
-            ) and num_trials_left > 0:
-                trial_batch = get_one_batch_of_trials(
-                    ax_client=ax_client,
-                    parallelism=current_parallelism,
-                    num_trials_so_far=num_trials_so_far,
-                    num_max_trials_to_do=num_trials_left,
-                )
-
-                list_of_trials_to_launch = trial_batch.list_of_trials[:num_trials_left]
-                is_search_space_exhausted = trial_batch.is_search_space_exhausted
+            num_trials_to_request = min(num_trials_left, self.max_batch_size or 5)
+            trial_batch = get_one_batch_of_trials(
+                ax_client=ax_client,
+                num_max_trials_to_do=num_trials_to_request,
+            )
 
-                log.info(
-                    "AxSweeper is launching {} jobs".format(
-                        len(list_of_trials_to_launch)
-                    )
-                )
+            list_of_trials_to_launch = trial_batch.list_of_trials[:num_trials_left]
+            is_search_space_exhausted = trial_batch.is_search_space_exhausted
 
-                self.sweep_over_batches(
-                    ax_client=ax_client, list_of_trials=list_of_trials_to_launch
-                )
+            log.info(
+                "AxSweeper is launching {} jobs".format(len(list_of_trials_to_launch))
+            )
 
-                num_trials_so_far += len(list_of_trials_to_launch)
-                num_trials_left -= len(list_of_trials_to_launch)
+            self.sweep_over_batches(
+                ax_client=ax_client, list_of_trials=list_of_trials_to_launch
+            )
 
-                best_parameters, predictions = ax_client.get_best_parameters()
-                metric = predictions[0][ax_client.objective_name]
+            num_trials_left -= len(list_of_trials_to_launch)
 
+            best_point = self.get_best_point(ax_client)
+            if best_point is not None:
+                best_parameters, metric = best_point
                 if self.early_stopper.should_stop(metric, best_parameters):
-                    num_trials_left = -1
                     break
 
-                if is_search_space_exhausted:
-                    log.info("Ax has exhausted the search space")
-                    break
-
-            current_parallelism_index += 1
+            if is_search_space_exhausted:
+                log.info("Ax has exhausted the search space")
+                break
 
         results_to_serialize = {"optimizer": "ax", "ax": best_parameters}
         OmegaConf.save(
@@ -205,7 +238,7 @@ class CoreAxSweeper(Sweeper):
         log.info("Best parameters: " + str(best_parameters))
 
     def sweep_over_batches(
-        self, ax_client: AxClient, list_of_trials: List[Trial]
+        self, ax_client: Client, list_of_trials: List[Trial]
     ) -> None:
         assert self.launcher is not None
         assert self.job_idx is not None
@@ -234,11 +267,30 @@ class CoreAxSweeper(Sweeper):
                         val = (val, None)  # specify unknown noise
                     else:
                         val = (val, 0)  # specify no noise
+                if isinstance(val, tuple):
+                    val = {self.experiment.objective_name: val}
                 ax_client.complete_trial(
                     trial_index=batch[idx].trial_index, raw_data=val
                 )
 
-    def setup_ax_client(self, arguments: List[str]) -> AxClient:
+    def get_best_point(
+        self, ax_client: Client
+    ) -> Optional[Tuple[Mapping[str, Any], float]]:
+        try:
+            best_parameters, metrics, _, _ = ax_client.get_best_parameterization(
+                use_model_predictions=False
+            )
+        except (AssertionError, UnsupportedError):
+            return None
+
+        metric = metrics.get(self.experiment.objective_name)
+        if metric is None:
+            return None
+        if isinstance(metric, tuple):
+            metric = metric[0]
+        return best_parameters, float(metric)
+
+    def setup_ax_client(self, arguments: List[str]) -> Client:
         """Method to setup the Ax Client"""
         parameters: List[Dict[Any, Any]] = []
         for key, value in self.ax_params.items():
@@ -249,7 +301,6 @@ class CoreAxSweeper(Sweeper):
                 if not (all(isinstance(x, int) for x in bounds)):
                     # Type mismatch. Promote all to float
                     param["bounds"] = [float(x) for x in bounds]
-
             parameters.append(param)
             parameters[-1]["name"] = key
         commandline_params = self.parse_commandline_args(arguments)
@@ -265,11 +316,27 @@ class CoreAxSweeper(Sweeper):
         log.info(
             f"AxSweeper is optimizing the following parameters: {encoder_parameters_into_string(parameters)}"
         )
-        ax_client = AxClient(
-            verbose_logging=self.ax_client_config.verbose_logging,
-            random_seed=self.ax_client_config.random_seed,
+
+        if not self.ax_client_config.verbose_logging:
+            logging.getLogger("ax.api.client").setLevel(logging.WARNING)
+        ax_client = Client(random_seed=self.ax_client_config.random_seed)
+        ax_client.configure_experiment(
+            parameters=[create_ax_parameter_config(param) for param in parameters],
+            parameter_constraints=self.experiment.parameter_constraints,
+            name=self.experiment.name,
+        )
+        if self.experiment.status_quo is not None:
+            ax_client.attach_baseline(
+                parameters=cast(Mapping[str, Any], self.experiment.status_quo)
+            )
+        ax_client.configure_optimization(
+            objective=(
+                f"-{self.experiment.objective_name}"
+                if self.experiment.minimize
+                else self.experiment.objective_name
+            ),
+            outcome_constraints=self.experiment.outcome_constraints,
         )
-        ax_client.create_experiment(parameters=parameters, **self.experiment)
 
         return ax_client
 
@@ -288,8 +355,12 @@ class CoreAxSweeper(Sweeper):
                     param = create_choice_param_from_range_override(override)
                 elif override.is_interval_sweep():
                     param = create_range_param_using_interval_override(override)
+                else:
+                    raise ValueError(f"Unsupported sweep override: {override}")
             elif not override.is_hydra_override():
                 param = create_fixed_param_from_element_override(override)
+            else:
+                continue
             parameters.append(param)
 
         return parameters

{hydra-ax-sweeper-1.3.0.dev0 → hydra_ax_sweeper-1.4.0.dev4}/hydra_plugins/hydra_ax_sweeper/_earlystopper.py

@@ -1,8 +1,6 @@
 # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
 import logging
-from typing import Optional
-
-from ax import ParameterType  # type: ignore
+from typing import Any, Mapping, Optional
 
 log = logging.getLogger(__name__)
 
@@ -24,7 +22,7 @@ class EarlyStopper:
         self.current_epochs_without_improvement = 0
 
     def should_stop(
-        self, potential_best_value: float, best_parameters: ParameterType
+        self, potential_best_value: float, best_parameters: Mapping[str, Any]
     ) -> bool:
         """Check if the optimization process should be stopped."""
         is_improving = True

{hydra-ax-sweeper-1.3.0.dev0 → hydra_ax_sweeper-1.4.0.dev4}/hydra_plugins/hydra_ax_sweeper/config.py

@@ -34,7 +34,6 @@ class ExperimentConfig:
 
 @dataclass
 class ClientConfig:
-
     verbose_logging: bool = False
 
     # set random seed here to make Ax results reproducible
@@ -43,12 +42,11 @@ class ClientConfig:
 
 @dataclass
 class AxConfig:
-
     # max_trials is application-specific. Tune it for your use case
     max_trials: int = 10
-    early_stop: EarlyStopConfig = EarlyStopConfig()
-    experiment: ExperimentConfig = ExperimentConfig()
-    client: ClientConfig = ClientConfig()
+    early_stop: EarlyStopConfig = field(default_factory=EarlyStopConfig)
+    experiment: ExperimentConfig = field(default_factory=ExperimentConfig)
+    client: ClientConfig = field(default_factory=ClientConfig)
     params: Dict[str, Any] = field(default_factory=dict)
     # is_noisy = True indicates measurements have unknown uncertainty
     # is_noisy = False indicates measurements have an uncertainty of zero
@@ -60,7 +58,7 @@ class AxSweeperConf:
     _target_: str = "hydra_plugins.hydra_ax_sweeper.ax_sweeper.AxSweeper"
     # Maximum number of trials to run in parallel
     max_batch_size: Optional[int] = None
-    ax_config: AxConfig = AxConfig()
+    ax_config: AxConfig = field(default_factory=AxConfig)
 
 
 ConfigStore.instance().store(

{hydra-ax-sweeper-1.3.0.dev0 → hydra_ax_sweeper-1.4.0.dev4}/setup.py

@@ -11,25 +11,25 @@ setup(
     author="Omry Yadan, Shagun Sodhani",
     author_email="omry@fb.com, sshagunsodhani@gmail.com",
     description="Hydra Ax Sweeper plugin",
+    license="MIT",
     long_description=(Path(__file__).parent / "README.md").read_text(),
     long_description_content_type="text/markdown",
     url="https://github.com/facebookresearch/hydra/",
     packages=find_namespace_packages(include=["hydra_plugins.*"]),
     classifiers=[
-        "License :: OSI Approved :: MIT License",
-        "Programming Language :: Python :: 3.7",
-        "Programming Language :: Python :: 3.8",
-        "Programming Language :: Python :: 3.9",
-        "Programming Language :: Python :: 3.10",
+        "Programming Language :: Python :: 3.11",
+        "Programming Language :: Python :: 3.12",
+        "Programming Language :: Python :: 3.13",
+        "Programming Language :: Python :: 3.14",
         "Operating System :: POSIX :: Linux",
         "Operating System :: MacOS",
         "Development Status :: 4 - Beta",
     ],
+    python_requires=">=3.11,<3.15",
     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.
+        "ax-platform>=1.2.4,<1.3.0",
+        "torch>=2.2",
     ],
     include_package_data=True,
 )

hydra_ax_sweeper-1.4.0.dev4/tests/test_ax_sweeper_plugin.py

@@ -0,0 +1,461 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+import math
+import os
+from pathlib import Path
+from typing import Any, List, Tuple
+
+from hydra.core.plugins import Plugins
+from hydra.plugins.sweeper import Sweeper
+from hydra.test_utils.test_utils import (
+    TSweepRunner,
+    chdir_plugin_root,
+    run_python_script,
+)
+from omegaconf import DictConfig, OmegaConf
+from pytest import mark, raises
+
+from hydra_plugins.hydra_ax_sweeper.ax_sweeper import AxSweeper
+
+chdir_plugin_root()
+
+WARNING_FILTERS = [
+    # 2026-05-15: linear_operator 0.6.1 imports torch.jit.script via
+    # Ax -> Botorch -> GPyTorch on Python 3.14 with Torch 2.12.
+    # Remove when the current Ax stack no longer emits it under -Werror.
+    "ignore:`torch.jit.script` is deprecated:DeprecationWarning",
+    # 2026-05-15: same linear_operator import path as above, but Torch emits this
+    # alternate wording in some import paths. Remove with the filter above.
+    "ignore:`torch.jit.script` is not supported:DeprecationWarning",
+    # 2026-05-15: Ax 1.2.4 uses asyncio.iscoroutinefunction in retry helpers.
+    # Remove when Ax no longer emits it on Python 3.14 under -Werror.
+    "ignore:'asyncio.iscoroutinefunction' is deprecated:DeprecationWarning",
+    # 2026-05-15: Ax 1.2.4's JSON storage registry imports this moved shim.
+    # Remove when importing ax.api.Client no longer touches the shim.
+    "ignore:ax.service.utils.orchestrator_options has been moved:DeprecationWarning",
+    # 2026-05-15: Ax 1.2.4's overview analysis imports this moved shim.
+    # Remove when importing ax.api.Client no longer touches the shim.
+    "ignore:ax.service.orchestrator has been moved:DeprecationWarning",
+    # 2026-05-15: Botorch/linear_operator can add Cholesky jitter while fitting
+    # the Ax surrogate for log-scale tests. Remove when the current Ax stack no
+    # longer emits it under -Werror for the plugin's log-scale sweep.
+    "ignore:A not p.d., added jitter:Warning",
+]
+
+pytestmark = [mark.filterwarnings(warning_filter) for warning_filter in WARNING_FILTERS]
+PYTHON_WARNING_FILTERS = [f"-W{warning_filter}" for warning_filter in WARNING_FILTERS]
+
+
+def run_ax_python_script(cmd: List[str]) -> Tuple[str, str]:
+    return run_python_script(PYTHON_WARNING_FILTERS + cmd)
+
+
+def test_discovery() -> None:
+    """
+    Tests that this plugin can be discovered via the plugins subsystem when looking for Sweeper
+    :return:
+    """
+    assert AxSweeper.__name__ in [
+        x.__name__ for x in Plugins.instance().discover(Sweeper)
+    ]
+
+
+def quadratic(cfg: DictConfig) -> Any:
+    return 100 * (cfg.quadratic.x**2) + 1 * cfg.quadratic.y
+
+
+@mark.parametrize(
+    "n,expected",
+    [
+        (None, [[1, 2, 3, 4, 5]]),
+        (1, [[1], [2], [3], [4], [5]]),
+        (2, [[1, 2], [3, 4], [5]]),
+        (5, [[1, 2, 3, 4, 5]]),
+        (6, [[1, 2, 3, 4, 5]]),
+    ],
+)
+def test_chunk_method_for_valid_inputs(n: int, expected: List[List[int]]) -> None:
+    from hydra_plugins.hydra_ax_sweeper._core import CoreAxSweeper
+
+    chunk_func = CoreAxSweeper.chunks
+    batch = [1, 2, 3, 4, 5]
+    out = list(chunk_func(batch, n))
+    assert out == expected
+
+
+@mark.parametrize("n", [-1, -11, 0])
+def test_chunk_method_for_invalid_inputs(n: int) -> None:
+    from hydra_plugins.hydra_ax_sweeper._core import CoreAxSweeper
+
+    chunk_func = CoreAxSweeper.chunks
+    batch = [1, 2, 3, 4, 5]
+    with raises(ValueError):
+        list(chunk_func(batch, n))
+
+
+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(
+        calling_file="tests/test_ax_sweeper_plugin.py",
+        calling_module=None,
+        task_function=quadratic,
+        config_path="config",
+        config_name="config.yaml",
+        overrides=[
+            "hydra/launcher=basic",
+            "hydra.sweeper.ax_config.max_trials=6",
+            "hydra.sweeper.ax_config.early_stop.max_epochs_without_improvement=100",
+            "hydra.sweeper.max_batch_size=2",
+            "params=basic",
+        ],
+    )
+    with sweep:
+        assert isinstance(sweep.temp_dir, str)
+        dirs = [
+            x
+            for x in os.listdir(sweep.temp_dir)
+            if os.path.isdir(os.path.join(sweep.temp_dir, x))
+        ]
+        assert len(dirs) == 6  # and a total of 6 unique output directories
+
+
+@mark.parametrize("test_conf", ["basic", "logscale"])
+def test_jobs_configured_via_config(
+    hydra_sweep_runner: TSweepRunner, test_conf: str
+) -> None:
+    sweep = hydra_sweep_runner(
+        calling_file="tests/test_ax_sweeper_plugin.py",
+        calling_module=None,
+        task_function=quadratic,
+        config_path="config",
+        config_name="config.yaml",
+        overrides=["hydra/launcher=basic", f"params={test_conf}"],
+    )
+    with sweep:
+        assert sweep.returns is None
+        returns = OmegaConf.load(f"{sweep.temp_dir}/optimization_results.yaml")
+        assert isinstance(returns, DictConfig)
+        assert returns["optimizer"] == "ax"
+        assert len(returns) == 2
+        best_parameters = returns.ax
+        assert math.isclose(best_parameters["quadratic.x"], 0.0, abs_tol=1e-4)
+        expected_y = 0.0 if test_conf == "basic" else -1.0
+        assert math.isclose(best_parameters["quadratic.y"], expected_y, abs_tol=1e-4)
+
+
+@mark.parametrize(
+    "test_conf, override, expected_x",
+    [
+        ("basic", "int(interval(1, 5))", 1.0),
+    ],
+)
+def test_jobs_configured_via_cmd(
+    hydra_sweep_runner: TSweepRunner, test_conf: str, override: str, expected_x: float
+) -> None:
+    sweep = hydra_sweep_runner(
+        calling_file="tests/test_ax_sweeper_plugin.py",
+        calling_module=None,
+        task_function=quadratic,
+        config_path="config",
+        config_name="config.yaml",
+        overrides=[
+            "hydra/launcher=basic",
+            "hydra.sweeper.ax_config.max_trials=5",
+            f"quadratic.x={override}",
+            "quadratic.y=-2",
+            f"params={test_conf}",
+        ],
+    )
+    with sweep:
+        assert sweep.returns is None
+        returns = OmegaConf.load(f"{sweep.temp_dir}/optimization_results.yaml")
+        assert isinstance(returns, DictConfig)
+        assert returns["optimizer"] == "ax"
+        assert len(returns) == 2
+        best_parameters = returns.ax
+        assert math.isclose(best_parameters["quadratic.x"], expected_x, abs_tol=1e-4)
+        assert math.isclose(best_parameters["quadratic.y"], -2.0, abs_tol=1e-4)
+
+
+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",
+        calling_module=None,
+        task_function=quadratic,
+        config_path="config",
+        config_name="config.yaml",
+        overrides=[
+            "hydra/launcher=basic",
+            "hydra.sweeper.ax_config.max_trials=2",
+            "quadratic.x=int(interval(-5, -2))",
+            "params=basic",
+        ],
+    )
+    with sweep:
+        assert sweep.returns is None
+        returns = OmegaConf.load(f"{sweep.temp_dir}/optimization_results.yaml")
+        assert isinstance(returns, DictConfig)
+        assert returns["optimizer"] == "ax"
+        assert len(returns) == 2
+        best_parameters = returns.ax
+        assert math.isclose(best_parameters["quadratic.x"], -3.0, abs_tol=1e-4)
+        assert math.isclose(best_parameters["quadratic.y"], 0.0, abs_tol=1e-4)
+
+
+def test_command_line_int_interval_optimizes_integer_parabola(
+    hydra_sweep_runner: TSweepRunner,
+) -> None:
+    sweep = hydra_sweep_runner(
+        calling_file="tests/test_ax_sweeper_plugin.py",
+        calling_module=None,
+        task_function=quadratic,
+        config_path="config",
+        config_name="config.yaml",
+        overrides=[
+            "hydra/launcher=basic",
+            "hydra.sweeper.ax_config.max_trials=12",
+            "quadratic.x=int(interval(-5, 5))",
+            "quadratic.y=0",
+        ],
+    )
+    with sweep:
+        assert sweep.returns is None
+        returns = OmegaConf.load(f"{sweep.temp_dir}/optimization_results.yaml")
+        assert isinstance(returns, DictConfig)
+        best_parameters = returns.ax
+        assert isinstance(best_parameters["quadratic.x"], int)
+        assert best_parameters["quadratic.x"] == 0
+        assert best_parameters["quadratic.y"] == 0
+
+
+def test_configuration_set_via_cmd_and_default_config(
+    hydra_sweep_runner: TSweepRunner,
+) -> None:
+    sweep = hydra_sweep_runner(
+        calling_file="tests/test_ax_sweeper_plugin.py",
+        calling_module=None,
+        task_function=quadratic,
+        config_path="config",
+        config_name="config.yaml",
+        overrides=[
+            "hydra/launcher=basic",
+            "hydra.sweeper.ax_config.max_trials=2",
+            "hydra.sweeper.ax_config.early_stop.max_epochs_without_improvement=2",
+            "quadratic=basic",
+            "quadratic.x=interval(-5, -2)",
+            "quadratic.y=interval(-1, 1)",
+        ],
+    )
+    with sweep:
+        assert sweep.returns is None
+        returns = OmegaConf.load(f"{sweep.temp_dir}/optimization_results.yaml")
+        assert isinstance(returns, DictConfig)
+        best_parameters = returns.ax
+        assert "quadratic.x" in best_parameters
+        assert "quadratic.y" in best_parameters
+
+
+@mark.parametrize(
+    "cmd_arg, expected_str",
+    [
+        ("polynomial.y=choice(-1, 0, 1)", "polynomial.y: choice=[-1, 0, 1]"),
+        ("polynomial.y=range(-1, 2)", "polynomial.y: choice=[-1, 0, 1]"),
+        ("polynomial.y=range(-1, 3, 1)", "polynomial.y: choice=[-1, 0, 1, 2]"),
+        (
+            "polynomial.y=range(-1, 2, 0.5)",
+            "polynomial.y: choice=[-1.0, -0.5, 0.0, 0.5, 1.0, 1.5]",
+        ),
+        ("polynomial.y=int(interval(-1, 2))", "polynomial.y: range=[-1, 2]"),
+        ("polynomial.y=interval(-1, 2)", "polynomial.y: range=[-1.0, 2.0]"),
+        (
+            "polynomial.y=tag(log, interval(0.00001, 1))",
+            "polynomial.y: range=[1e-05, 1.0], log_scale = True",
+        ),
+        ("polynomial.y=2", "polynomial.y: fixed=2"),
+        ("polynomial.y=2.0", "polynomial.y: fixed=2.0"),
+    ],
+)
+def test_ax_logging(cmd_arg: str, expected_str: str) -> None:
+    from hydra_plugins.hydra_ax_sweeper._core import (
+        CoreAxSweeper,
+        encoder_parameters_into_string,
+    )
+    from hydra_plugins.hydra_ax_sweeper.config import AxConfig
+
+    sweeper = CoreAxSweeper(AxConfig(), max_batch_size=None)
+    parameters = sweeper.parse_commandline_args(
+        [
+            "polynomial.x=interval(-5, -2)",
+            "polynomial.z=10",
+            cmd_arg,
+        ]
+    )
+    result = encoder_parameters_into_string(parameters)
+    assert "polynomial.x: range=[-5.0, -2.0]" in result
+    assert expected_str in result
+    assert "polynomial.z: fixed=10" in result
+
+
+def test_command_line_log_interval_configures_ax_log_range() -> None:
+    from ax.api.configs import RangeParameterConfig  # type: ignore
+
+    from hydra_plugins.hydra_ax_sweeper._core import (
+        CoreAxSweeper,
+        create_ax_parameter_config,
+    )
+    from hydra_plugins.hydra_ax_sweeper.config import AxConfig
+
+    sweeper = CoreAxSweeper(AxConfig(), max_batch_size=None)
+    (parameter,) = sweeper.parse_commandline_args(
+        ["polynomial.y=tag(log, interval(0.00001, 1))"]
+    )
+    ax_parameter = create_ax_parameter_config(parameter)
+    assert isinstance(ax_parameter, RangeParameterConfig)
+    assert ax_parameter.name == "polynomial.y"
+    assert ax_parameter.parameter_type == "float"
+    assert ax_parameter.bounds == (1e-05, 1.0)
+    assert ax_parameter.scaling == "log"
+
+
+def test_ax_logging_from_hydra_app(tmpdir: Path) -> None:
+    cmd = [
+        "tests/apps/polynomial.py",
+        "-m",
+        f'hydra.run.dir="{str(tmpdir)}"',
+        "hydra.job.chdir=True",
+        "polynomial.x=interval(-5, -2)",
+        "polynomial.z=10",
+        "hydra.sweeper.ax_config.max_trials=2",
+        "polynomial.y=int(interval(-1, 2))",
+    ]
+    result, _ = run_ax_python_script(cmd)
+    assert "polynomial.x: range=[-5.0, -2.0]" in result
+    assert "polynomial.y: range=[-1, 2]" in result
+    assert "polynomial.z: fixed=10" in result
+
+
+@mark.parametrize(
+    "cmd_args",
+    [
+        ["polynomial.y=choice(-1, 0, 1)", "polynomial.x=range(2,4)"],
+        ["polynomial.y=1", "polynomial.x=range(2,4)"],
+    ],
+)
+def test_search_space_exhausted_exception(tmpdir: Path, cmd_args: List[str]) -> None:
+    cmd = [
+        "tests/apps/polynomial.py",
+        "-m",
+        f'hydra.run.dir="{str(tmpdir)}"',
+        "hydra.job.chdir=True",
+        "hydra.sweeper.ax_config.max_trials=2",
+    ] + cmd_args
+    run_ax_python_script(cmd)
+
+
+@mark.parametrize(
+    "cmd_args",
+    [
+        ["polynomial.y=choice(-1, 0, 1)", "polynomial.x=range(2,4)"],
+        ["polynomial.y=1", "polynomial.x=range(2,4)"],
+    ],
+)
+def test_search_space_with_constraint_metric(tmpdir: Path, cmd_args: List[str]) -> None:
+    # test that outcome_constraints experiment parameter `outcome_constraints`
+    # works correctly, and that the ax_sweeper supports outputting a dictionary
+    # from the evaluation function so that multiple metrics can be supported.
+    cmd = [
+        "tests/apps/polynomial_with_constraint.py",
+        "-m",
+        f'hydra.run.dir="{str(tmpdir)}"',
+        "hydra.job.chdir=True",
+        "hydra.sweeper.ax_config.max_trials=2",
+    ] + cmd_args
+    results, _ = run_ax_python_script(cmd)
+
+
+@mark.parametrize(
+    "cmd_arg, serialized_encoding, best_coefficients, best_value",
+    [
+        (
+            "polynomial.coefficients=[-1, 0, 1],[2, 3, 4],[5, 6, 7]",
+            "choice=['[-1,0,1]', '[2,3,4]', '[5,6,7]']",
+            "'[-1,0,1]'",
+            101.0,
+        ),
+        (
+            "polynomial.coefficients=choice([8, 12, 11],[-1, -1, 1000], [-2, 4, 7])",
+            "choice=['[8,12,11]', '[-1,-1,1000]', '[-2,4,7]']",
+            "'[-2,4,7]'",
+            447,
+        ),
+    ],
+)
+def test_jobs_using_choice_between_lists(
+    tmpdir: Path,
+    cmd_arg: str,
+    serialized_encoding: str,
+    best_coefficients: str,
+    best_value: float,
+) -> None:
+    cmd = [
+        "tests/apps/polynomial_with_list_coefficients.py",
+        "-m",
+        f'hydra.run.dir="{str(tmpdir)}"',
+        "hydra.job.chdir=True",
+        "hydra.sweeper.ax_config.max_trials=3",
+    ] + [cmd_arg]
+    result, _ = run_ax_python_script(cmd)
+    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
+
+
+@mark.parametrize(
+    "cmd_arg, serialized_encoding, best_coefficients, best_value",
+    [
+        (
+            "+polynomial.coefficients=choice({x:-1, y:0, z:1},{x:2, y:3, z:4},{x:5, y:6, z:7})",
+            "choice=['{x:-1,y:0,z:1}', '{x:2,y:3,z:4}', '{x:5,y:6,z:7}']",
+            "'{x:-1,y:0,z:1}'",
+            101.0,
+        ),
+        (
+            "+polynomial.coefficients=choice({x:8, y:12, z:11},{x:-1, y:-1, z:1000}, {x:-2, y:4, z:7})",
+            "choice=['{x:8,y:12,z:11}', '{x:-1,y:-1,z:1000}', '{x:-2,y:4,z:7}']",
+            "'{x:-2,y:4,z:7}'}",
+            447,
+        ),
+    ],
+)
+def test_jobs_using_choice_between_dicts(
+    tmpdir: Path,
+    cmd_arg: str,
+    serialized_encoding: str,
+    best_coefficients: str,
+    best_value: float,
+) -> None:
+    cmd = [
+        "tests/apps/polynomial_with_dict_coefficients.py",
+        "-m",
+        f'hydra.run.dir="{str(tmpdir)}"',
+        "hydra.job.chdir=True",
+        "hydra.sweeper.ax_config.max_trials=3",
+    ] + [cmd_arg]
+    result, _ = run_ax_python_script(cmd)
+    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
+
+
+def test_example_app(tmpdir: Path) -> None:
+    cmd = [
+        "example/banana.py",
+        "-m",
+        f'hydra.run.dir="{str(tmpdir)}"',
+        "hydra.job.chdir=True",
+        "banana.x=int(interval(-5, 5))",
+        "banana.y=interval(-5, 10.1)",
+        "hydra.sweeper.ax_config.max_trials=2",
+    ]
+    result, _ = run_ax_python_script(cmd)
+    assert "banana.x: range=[-5, 5]" in result
+    assert "banana.y: range=[-5.0, 10.1]" in result

hydra-ax-sweeper-1.3.0.dev0/MANIFEST.in

@@ -1,3 +0,0 @@
-global-exclude *.pyc
-global-exclude __pycache__
-recursive-include hydra_plugins/* *.yaml

hydra-ax-sweeper-1.3.0.dev0/PKG-INFO

@@ -1,21 +0,0 @@
-Metadata-Version: 2.1
-Name: hydra-ax-sweeper
-Version: 1.3.0.dev0
-Summary: Hydra Ax Sweeper plugin
-Home-page: https://github.com/facebookresearch/hydra/
-Author: Omry Yadan, Shagun Sodhani
-Author-email: omry@fb.com, sshagunsodhani@gmail.com
-Classifier: License :: OSI Approved :: MIT License
-Classifier: Programming Language :: Python :: 3.7
-Classifier: Programming Language :: Python :: 3.8
-Classifier: Programming Language :: Python :: 3.9
-Classifier: Programming Language :: Python :: 3.10
-Classifier: Operating System :: POSIX :: Linux
-Classifier: Operating System :: MacOS
-Classifier: Development Status :: 4 - Beta
-Description-Content-Type: text/markdown
-
-# Hydra Ax Sweeper
-Provides a [`Ax Sweeper`](https://ax.dev/) based Hydra Sweeper supporting parallel execution.
-
-See [website](https://hydra.cc/docs/plugins/ax_sweeper/) for more information

hydra-ax-sweeper-1.3.0.dev0/hydra_ax_sweeper.egg-info/PKG-INFO

@@ -1,21 +0,0 @@
-Metadata-Version: 2.1
-Name: hydra-ax-sweeper
-Version: 1.3.0.dev0
-Summary: Hydra Ax Sweeper plugin
-Home-page: https://github.com/facebookresearch/hydra/
-Author: Omry Yadan, Shagun Sodhani
-Author-email: omry@fb.com, sshagunsodhani@gmail.com
-Classifier: License :: OSI Approved :: MIT License
-Classifier: Programming Language :: Python :: 3.7
-Classifier: Programming Language :: Python :: 3.8
-Classifier: Programming Language :: Python :: 3.9
-Classifier: Programming Language :: Python :: 3.10
-Classifier: Operating System :: POSIX :: Linux
-Classifier: Operating System :: MacOS
-Classifier: Development Status :: 4 - Beta
-Description-Content-Type: text/markdown
-
-# Hydra Ax Sweeper
-Provides a [`Ax Sweeper`](https://ax.dev/) based Hydra Sweeper supporting parallel execution.
-
-See [website](https://hydra.cc/docs/plugins/ax_sweeper/) for more information

hydra-ax-sweeper-1.3.0.dev0/hydra_ax_sweeper.egg-info/requires.txt

@@ -1,4 +0,0 @@
-hydra-core>=1.1.0.dev7
-ax-platform<0.2.1,>=0.1.20
-torch
-gpytorch<=1.8.1