chore: import upstream snapshot with attribution

python/ray/tune/search/ax/__init__.py

@@ -0,0 +1,3 @@
+from ray.tune.search.ax.ax_search import AxSearch
+
+__all__ = ["AxSearch"]

python/ray/tune/search/ax/ax_search.py

@@ -0,0 +1,476 @@
+import copy
+import logging
+from typing import Dict, List, Optional, Union
+
+import numpy as np
+
+from ray import cloudpickle
+from ray.tune.result import DEFAULT_METRIC
+from ray.tune.search import (
+    UNDEFINED_METRIC_MODE,
+    UNDEFINED_SEARCH_SPACE,
+    UNRESOLVED_SEARCH_SPACE,
+    Searcher,
+)
+from ray.tune.search.sample import (
+    Categorical,
+    Float,
+    Integer,
+    LogUniform,
+    Quantized,
+    Uniform,
+)
+from ray.tune.search.variant_generator import parse_spec_vars
+from ray.tune.utils.util import flatten_dict, unflatten_list_dict
+
+try:
+    import ax
+    from ax.service.ax_client import AxClient
+
+    try:
+        # Newer Ax versions
+        from ax.service.ax_client import ObjectiveProperties
+    except Exception:
+        ObjectiveProperties = None
+except ImportError:
+    ax = AxClient = ObjectiveProperties = None
+
+# This exception only exists in newer Ax releases for python 3.7
+try:
+    from ax.exceptions.core import DataRequiredError
+    from ax.exceptions.generation_strategy import MaxParallelismReachedException
+except ImportError:
+    MaxParallelismReachedException = DataRequiredError = Exception
+
+
+logger = logging.getLogger(__name__)
+
+
+class AxSearch(Searcher):
+    """Uses `Ax <https://ax.dev/>`_ to optimize hyperparameters.
+
+    Ax is a platform for understanding, managing, deploying, and
+    automating adaptive experiments. Ax provides an easy to use
+    interface with BoTorch, a flexible, modern library for Bayesian
+    optimization in PyTorch. More information can be found in https://ax.dev/.
+
+    To use this search algorithm, you must install Ax:
+
+    .. code-block:: bash
+
+        $ pip install ax-platform
+
+    Parameters:
+        space: Parameters in the experiment search space.
+            Required elements in the dictionaries are: "name" (name of
+            this parameter, string), "type" (type of the parameter: "range",
+            "fixed", or "choice", string), "bounds" for range parameters
+            (list of two values, lower bound first), "values" for choice
+            parameters (list of values), and "value" for fixed parameters
+            (single value).
+        metric: Name of the metric used as objective in this
+            experiment. This metric must be present in `raw_data` argument
+            to `log_data`. This metric must also be present in the dict
+            reported/returned by the Trainable. If None but a mode was passed,
+            the `ray.tune.result.DEFAULT_METRIC` will be used per default.
+        mode: One of {min, max}. Determines whether objective is
+            minimizing or maximizing the metric attribute. Defaults to "max".
+        points_to_evaluate: Initial parameter suggestions to be run
+            first. This is for when you already have some good parameters
+            you want to run first to help the algorithm make better suggestions
+            for future parameters. Needs to be a list of dicts containing the
+            configurations.
+        parameter_constraints: Parameter constraints, such as
+            "x3 >= x4" or "x3 + x4 >= 2".
+        outcome_constraints: Outcome constraints of form
+            "metric_name >= bound", like "m1 <= 3."
+        ax_client: Optional AxClient instance. If this is set, do
+            not pass any values to these parameters: `space`, `metric`,
+            `parameter_constraints`, `outcome_constraints`.
+        **ax_kwargs: Passed to AxClient instance. Ignored if `AxClient` is not
+            None.
+
+    Tune automatically converts search spaces to Ax's format:
+
+    .. code-block:: python
+
+        from ray import tune
+        from ray.tune.search.ax import AxSearch
+
+        config = {
+            "x1": tune.uniform(0.0, 1.0),
+            "x2": tune.uniform(0.0, 1.0)
+        }
+
+        def easy_objective(config):
+            for i in range(100):
+                intermediate_result = config["x1"] + config["x2"] * i
+                tune.report({"score": intermediate_result})
+
+        ax_search = AxSearch()
+        tuner = tune.Tuner(
+            easy_objective,
+            tune_config=tune.TuneConfig(
+                search_alg=ax_search,
+                metric="score",
+                mode="max",
+            ),
+            param_space=config,
+        )
+        tuner.fit()
+
+    If you would like to pass the search space manually, the code would
+    look like this:
+
+    .. code-block:: python
+
+        from ray import tune
+        from ray.tune.search.ax import AxSearch
+
+        parameters = [
+            {"name": "x1", "type": "range", "bounds": [0.0, 1.0]},
+            {"name": "x2", "type": "range", "bounds": [0.0, 1.0]},
+        ]
+
+        def easy_objective(config):
+            for i in range(100):
+                intermediate_result = config["x1"] + config["x2"] * i
+                tune.report({"score": intermediate_result})
+
+        ax_search = AxSearch(space=parameters, metric="score", mode="max")
+        tuner = tune.Tuner(
+            easy_objective,
+            tune_config=tune.TuneConfig(
+                search_alg=ax_search,
+            ),
+        )
+        tuner.fit()
+
+    """
+
+    def __init__(
+        self,
+        space: Optional[Union[Dict, List[Dict]]] = None,
+        metric: Optional[str] = None,
+        mode: Optional[str] = None,
+        points_to_evaluate: Optional[List[Dict]] = None,
+        parameter_constraints: Optional[List] = None,
+        outcome_constraints: Optional[List] = None,
+        ax_client: Optional[AxClient] = None,
+        **ax_kwargs,
+    ):
+        assert (
+            ax is not None
+        ), """Ax must be installed!
+            You can install AxSearch with the command:
+            `pip install ax-platform`."""
+
+        if mode:
+            assert mode in ["min", "max"], "`mode` must be 'min' or 'max'."
+
+        super(AxSearch, self).__init__(
+            metric=metric,
+            mode=mode,
+        )
+
+        self._ax = ax_client
+        self._ax_kwargs = ax_kwargs or {}
+
+        if isinstance(space, dict) and space:
+            resolved_vars, domain_vars, grid_vars = parse_spec_vars(space)
+            if domain_vars or grid_vars:
+                logger.warning(
+                    UNRESOLVED_SEARCH_SPACE.format(par="space", cls=type(self))
+                )
+                space = self.convert_search_space(space)
+
+        self._space = space
+        self._parameter_constraints = parameter_constraints
+        self._outcome_constraints = outcome_constraints
+
+        self._points_to_evaluate = copy.deepcopy(points_to_evaluate)
+
+        self._parameters = []
+        self._live_trial_mapping = {}
+
+        if self._ax or self._space:
+            self._setup_experiment()
+
+    def _setup_experiment(self):
+        if self._metric is None and self._mode:
+            self._metric = DEFAULT_METRIC
+
+        if not self._ax:
+            self._ax = AxClient(**self._ax_kwargs)
+
+        # Detect whether the AxClient already has an experiment attached.
+        # ax 1.0+ uses AssertionError ("Experiment not set"); older Ax used ValueError.
+        try:
+            _ = self._ax.experiment
+            has_experiment = True
+        except AssertionError as e:
+            if "Experiment not set" not in str(e):
+                raise
+            has_experiment = False
+        except ValueError:
+            has_experiment = False
+
+        if has_experiment:
+            # User provided an already-initialized client; they must not also pass
+            # experiment-defining args to AxSearch.
+            if any(
+                [
+                    self._space,
+                    self._parameter_constraints,
+                    self._outcome_constraints,
+                    self._mode,
+                    self._metric,
+                ]
+            ):
+                raise ValueError(
+                    "If you create the Ax experiment yourself (by passing an AxClient "
+                    "with an experiment), do not pass values for these parameters to "
+                    "AxSearch: space, parameter_constraints, outcome_constraints, mode, metric."
+                )
+        else:
+            # AxSearch must create the experiment.
+            if not self._space:
+                raise ValueError(
+                    "You have to create an Ax experiment by calling "
+                    "`AxClient.create_experiment()`, or pass an Ax search space as "
+                    "`space` to AxSearch, or pass a `param_space` dict to `tune.Tuner()`."
+                )
+            if self._mode not in ["min", "max"]:
+                raise ValueError(
+                    "Please specify the `mode` argument when initializing the AxSearch "
+                    "object or pass it to `tune.TuneConfig()`."
+                )
+
+            minimize = self._mode != "max"
+
+            # New Ax API (objectives=...) vs old (objective_name/minimize)
+            try:
+                if ObjectiveProperties is None:
+                    raise TypeError(
+                        "ObjectiveProperties not available in this Ax version"
+                    )
+
+                self._ax.create_experiment(
+                    parameters=self._space,
+                    objectives={self._metric: ObjectiveProperties(minimize=minimize)},
+                    parameter_constraints=self._parameter_constraints,
+                    outcome_constraints=self._outcome_constraints,
+                )
+            except TypeError:
+                self._ax.create_experiment(
+                    parameters=self._space,
+                    objective_name=self._metric,
+                    parameter_constraints=self._parameter_constraints,
+                    outcome_constraints=self._outcome_constraints,
+                    minimize=minimize,
+                )
+
+        # Access experiment - should exist now (either created above or already existed)
+        try:
+            exp = self._ax.experiment
+        except AssertionError as e:
+            # Re-raise if the message is unexpected to avoid masking real Ax failures.
+            if "Experiment not set" not in str(e):
+                raise
+            raise RuntimeError(
+                "Failed to access Ax experiment after setup. "
+                "This may indicate an issue with the Ax client setup."
+            ) from e
+        except ValueError as e:
+            raise RuntimeError(
+                "Failed to access Ax experiment after setup. "
+                "This may indicate an issue with the Ax client setup."
+            ) from e
+
+        # Populate mode/metric from experiment (keep your existing logic here)
+        try:
+            obj = exp.optimization_config.objective
+            self._mode = "min" if obj.minimize else "max"
+            self._metric = obj.metric.name
+        except Exception:
+            objectives = exp.optimization_config.objective.objectives
+            first = objectives[0]
+            self._mode = "min" if first.minimize else "max"
+            self._metric = first.metric.name
+
+        self._parameters = list(exp.parameters)
+
+        if getattr(self._ax, "_enforce_sequential_optimization", False):
+            logger.warning(
+                "Detected sequential enforcement. Be sure to use a ConcurrencyLimiter."
+            )
+
+    def set_search_properties(
+        self, metric: Optional[str], mode: Optional[str], config: Dict, **spec
+    ):
+        if self._ax:
+            return False
+        space = self.convert_search_space(config)
+        self._space = space
+        if metric:
+            self._metric = metric
+        if mode:
+            self._mode = mode
+
+        self._setup_experiment()
+        return True
+
+    def suggest(self, trial_id: str) -> Optional[Dict]:
+        if not self._ax:
+            raise RuntimeError(
+                UNDEFINED_SEARCH_SPACE.format(
+                    cls=self.__class__.__name__, space="space"
+                )
+            )
+
+        if not self._metric or not self._mode:
+            raise RuntimeError(
+                UNDEFINED_METRIC_MODE.format(
+                    cls=self.__class__.__name__, metric=self._metric, mode=self._mode
+                )
+            )
+
+        if self._points_to_evaluate:
+            config = self._points_to_evaluate.pop(0)
+            parameters, trial_index = self._ax.attach_trial(config)
+        else:
+            try:
+                parameters, trial_index = self._ax.get_next_trial()
+            except (MaxParallelismReachedException, DataRequiredError):
+                return None
+
+        self._live_trial_mapping[trial_id] = trial_index
+        try:
+            suggested_config = unflatten_list_dict(parameters)
+        except AssertionError:
+            # Fails to unflatten if keys are out of order, which only happens
+            # if search space includes a list with both constants and
+            # tunable hyperparameters:
+            # Ex: "a": [1, tune.uniform(2, 3), 4]
+            suggested_config = unflatten_list_dict(
+                {k: parameters[k] for k in sorted(parameters.keys())}
+            )
+        return suggested_config
+
+    def on_trial_complete(self, trial_id, result=None, error=False):
+        """Notification for the completion of trial.
+
+        Data of form key value dictionary of metric names and values.
+        """
+        if result:
+            self._process_result(trial_id, result)
+        self._live_trial_mapping.pop(trial_id)
+
+    def _process_result(self, trial_id, result):
+        ax_trial_index = self._live_trial_mapping[trial_id]
+        metrics_to_include = [self._metric] + [
+            oc.metric.name
+            for oc in self._ax.experiment.optimization_config.outcome_constraints
+        ]
+        metric_dict = {}
+        for key in metrics_to_include:
+            val = result[key]
+            if np.isnan(val) or np.isinf(val):
+                # Don't report trials with NaN metrics to Ax
+                self._ax.abandon_trial(
+                    trial_index=ax_trial_index,
+                    reason=f"nan/inf metrics reported by {trial_id}",
+                )
+                return
+            metric_dict[key] = (val, None)
+        self._ax.complete_trial(trial_index=ax_trial_index, raw_data=metric_dict)
+
+    @staticmethod
+    def convert_search_space(spec: Dict):
+        resolved_vars, domain_vars, grid_vars = parse_spec_vars(spec)
+
+        if grid_vars:
+            raise ValueError(
+                "Grid search parameters cannot be automatically converted "
+                "to an Ax search space."
+            )
+
+        # Flatten and resolve again after checking for grid search.
+        spec = flatten_dict(spec, prevent_delimiter=True)
+        resolved_vars, domain_vars, grid_vars = parse_spec_vars(spec)
+
+        def resolve_value(par, domain):
+            sampler = domain.get_sampler()
+            if isinstance(sampler, Quantized):
+                logger.warning(
+                    "AxSearch does not support quantization. Dropped quantization."
+                )
+                sampler = sampler.sampler
+
+            if isinstance(domain, Float):
+                if isinstance(sampler, LogUniform):
+                    return {
+                        "name": par,
+                        "type": "range",
+                        "bounds": [domain.lower, domain.upper],
+                        "value_type": "float",
+                        "log_scale": True,
+                    }
+                elif isinstance(sampler, Uniform):
+                    return {
+                        "name": par,
+                        "type": "range",
+                        "bounds": [domain.lower, domain.upper],
+                        "value_type": "float",
+                        "log_scale": False,
+                    }
+            elif isinstance(domain, Integer):
+                if isinstance(sampler, LogUniform):
+                    return {
+                        "name": par,
+                        "type": "range",
+                        "bounds": [domain.lower, domain.upper - 1],
+                        "value_type": "int",
+                        "log_scale": True,
+                    }
+                elif isinstance(sampler, Uniform):
+                    return {
+                        "name": par,
+                        "type": "range",
+                        "bounds": [domain.lower, domain.upper - 1],
+                        "value_type": "int",
+                        "log_scale": False,
+                    }
+            elif isinstance(domain, Categorical):
+                if isinstance(sampler, Uniform):
+                    return {"name": par, "type": "choice", "values": domain.categories}
+
+            raise ValueError(
+                "AxSearch does not support parameters of type "
+                "`{}` with samplers of type `{}`".format(
+                    type(domain).__name__, type(domain.sampler).__name__
+                )
+            )
+
+        # Parameter name is e.g. "a/b/c" for nested dicts,
+        # "a/d/0", "a/d/1" for nested lists (using the index in the list)
+        fixed_values = [
+            {"name": "/".join(str(p) for p in path), "type": "fixed", "value": val}
+            for path, val in resolved_vars
+        ]
+        resolved_values = [
+            resolve_value("/".join(str(p) for p in path), domain)
+            for path, domain in domain_vars
+        ]
+
+        return fixed_values + resolved_values
+
+    def save(self, checkpoint_path: str):
+        save_object = self.__dict__
+        with open(checkpoint_path, "wb") as outputFile:
+            cloudpickle.dump(save_object, outputFile)
+
+    def restore(self, checkpoint_path: str):
+        with open(checkpoint_path, "rb") as inputFile:
+            save_object = cloudpickle.load(inputFile)
+        self.__dict__.update(save_object)