Add files using upload-large-folder tool

.venv/lib/python3.11/site-packages/ray/tune/analysis/__init__.py

@@ -0,0 +1,3 @@
+from ray.tune.analysis.experiment_analysis import ExperimentAnalysis
+
+__all__ = ["ExperimentAnalysis"]

.venv/lib/python3.11/site-packages/ray/tune/analysis/experiment_analysis.py

@@ -0,0 +1,678 @@
+import copy
+import io
+import json
+import logging
+import os
+from numbers import Number
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import pyarrow.fs
+
+from ray.air.constants import EXPR_PROGRESS_FILE, EXPR_RESULT_FILE, TRAINING_ITERATION
+from ray.train import Checkpoint
+from ray.train._internal.storage import _exists_at_fs_path, get_fs_and_path
+from ray.tune.execution.experiment_state import _find_newest_experiment_checkpoint
+from ray.tune.execution.tune_controller import TuneController
+from ray.tune.experiment import Trial
+from ray.tune.result import CONFIG_PREFIX, DEFAULT_METRIC
+from ray.tune.utils import flatten_dict
+from ray.tune.utils.serialization import TuneFunctionDecoder
+from ray.tune.utils.util import is_nan, is_nan_or_inf, unflattened_lookup
+from ray.util.annotations import PublicAPI
+
+try:
+    import pandas as pd
+    from pandas import DataFrame
+except ImportError:
+    pd = None
+    DataFrame = None
+
+
+logger = logging.getLogger(__name__)
+
+
+@PublicAPI(stability="beta")
+class ExperimentAnalysis:
+    """Analyze results from a Ray Train/Tune experiment.
+
+    To use this class, the run must store the history of reported metrics
+    in log files (e.g., `result.json` and `progress.csv`).
+    This is the default behavior, unless default loggers are explicitly excluded
+    with the `TUNE_DISABLE_AUTO_CALLBACK_LOGGERS=1` environment variable.
+
+    Parameters:
+        experiment_checkpoint_path: Path to an `experiment_state.json` file,
+            or a directory that contains an `experiment_state.json` file.
+        default_metric: Default metric for comparing results. Can be
+            overwritten with the ``metric`` parameter in the respective
+            functions.
+        default_mode: Default mode for comparing results. Has to be one
+            of [min, max]. Can be overwritten with the ``mode`` parameter
+            in the respective functions.
+        trials: List of trials that can be accessed via `analysis.trials`.
+    """
+
+    def __init__(
+        self,
+        experiment_checkpoint_path: Union[str, os.PathLike],
+        *,
+        storage_filesystem: Optional[pyarrow.fs.FileSystem] = None,
+        trials: Optional[List[Trial]] = None,
+        default_metric: Optional[str] = None,
+        default_mode: Optional[str] = None,
+    ):
+        self.default_metric = default_metric
+        if default_mode and default_mode not in ["min", "max"]:
+            raise ValueError("`default_mode` has to be None or one of [min, max]")
+        self.default_mode = default_mode
+        if self.default_metric is None and self.default_mode is not None:
+            # If only a mode was passed, use anonymous metric
+            self.default_metric = DEFAULT_METRIC
+
+        # Resolve the filesystem if not specified.
+        if storage_filesystem:
+            self._fs = storage_filesystem
+        else:
+            self._fs, experiment_checkpoint_path = get_fs_and_path(
+                experiment_checkpoint_path
+            )
+
+        # Find the json state file.
+        experiment_checkpoint_path = str(experiment_checkpoint_path)
+        if experiment_checkpoint_path.endswith(".json"):
+            self._experiment_fs_path = os.path.dirname(experiment_checkpoint_path)
+            self._experiment_json_fs_path = experiment_checkpoint_path
+        else:
+            self._experiment_fs_path = experiment_checkpoint_path
+
+            experiment_json_fs_path = _find_newest_experiment_checkpoint(
+                experiment_path=self._experiment_fs_path, fs=self._fs
+            )
+            if experiment_json_fs_path is None:
+                pattern = TuneController.CKPT_FILE_TMPL.format("*")
+                raise ValueError(
+                    f"No experiment snapshot file of form '{pattern}' was found at: "
+                    f"({self._fs.type_name}, {self._experiment_fs_path})\n"
+                    "Please check if you specified the correct experiment path, "
+                    "which should be a combination of the `storage_path` and `name` "
+                    "specified in your run."
+                )
+
+            self._experiment_json_fs_path = experiment_json_fs_path
+
+        self.trials = trials or self._load_trials()
+        self._trial_dataframes = self._fetch_trial_dataframes()
+        self._configs = self.get_all_configs()
+
+    def _load_trials(self) -> List[Trial]:
+        with self._fs.open_input_stream(self._experiment_json_fs_path) as f:
+            experiment_state = json.loads(f.readall(), cls=TuneFunctionDecoder)
+
+        experiment_fs_path = Path(self._experiment_fs_path)
+
+        trials = []
+        trial_states = experiment_state["trial_data"]
+        for trial_json_state, trial_runtime_metadata in trial_states:
+            trial = Trial.from_json_state(trial_json_state, stub=True)
+            trial.restore_run_metadata(trial_runtime_metadata)
+
+            new_storage = copy.copy(trial.storage)
+            new_storage.storage_fs_path = experiment_fs_path.parent.as_posix()
+            new_storage.storage_filesystem = self._fs
+            new_storage.experiment_dir_name = experiment_fs_path.name
+            trial.set_storage(new_storage)
+
+            trials.append(trial)
+        return trials
+
+    def _fetch_trial_dataframe(self, trial: Trial) -> DataFrame:
+        force_dtype = {"trial_id": str}  # Never convert trial_id to float.
+
+        # If there were no reported results, there will be no files into a DataFrame
+        if trial.last_result is None:
+            return DataFrame()
+
+        json_fs_path = Path(trial.storage.trial_fs_path, EXPR_RESULT_FILE).as_posix()
+        csv_fs_path = Path(trial.storage.trial_fs_path, EXPR_PROGRESS_FILE).as_posix()
+        # Prefer reading the JSON if it exists.
+        if _exists_at_fs_path(trial.storage.storage_filesystem, json_fs_path):
+            with trial.storage.storage_filesystem.open_input_stream(json_fs_path) as f:
+                content = f.readall().decode("utf-8").rstrip("\n")
+                if not content:
+                    return DataFrame()
+                json_list = [json.loads(row) for row in content.split("\n")]
+            df = pd.json_normalize(json_list, sep="/")
+        # Fallback to reading the CSV.
+        elif _exists_at_fs_path(trial.storage.storage_filesystem, csv_fs_path):
+            with trial.storage.storage_filesystem.open_input_stream(csv_fs_path) as f:
+                csv_str = f.readall().decode("utf-8")
+            df = pd.read_csv(io.StringIO(csv_str), dtype=force_dtype)
+        else:
+            raise FileNotFoundError(
+                f"Could not fetch metrics for {trial}: both {EXPR_RESULT_FILE} and "
+                f"{EXPR_PROGRESS_FILE} were not found at {trial.storage.trial_fs_path}"
+            )
+
+        return df
+
+    def _fetch_trial_dataframes(self) -> Dict[str, DataFrame]:
+        """Fetches trial dataframes from files.
+
+        Returns:
+            A dictionary mapping trial_id -> pd.DataFrame
+        """
+        failures = []
+
+        trial_dfs = {}
+        for trial in self.trials:
+            try:
+                trial_dfs[trial.trial_id] = self._fetch_trial_dataframe(trial)
+            except Exception as e:
+                failures.append((trial, e))
+                trial_dfs[trial.trial_id] = DataFrame()
+                continue
+
+        if failures:
+            fail_str = "\n".join(
+                [f"- {trial}: {repr(error)}" for trial, error in failures]
+            )
+            logger.warning(
+                f"Failed to fetch metrics for {len(failures)} trial(s):\n{fail_str}"
+            )
+        return trial_dfs
+
+    def get_all_configs(self, prefix: bool = False) -> Dict[str, Dict]:
+        """Returns all trial hyperparameter configurations.
+
+        Args:
+            prefix: If True, flattens the config dict
+                and prepends `config/`.
+
+        Returns:
+            Dict[str, Dict]: Mapping trial_id -> config dict
+        """
+        return {
+            trial.trial_id: (
+                flatten_dict({CONFIG_PREFIX: trial.config}) if prefix else trial.config
+            )
+            for trial in self.trials
+        }
+
+    @property
+    def experiment_path(self) -> str:
+        """Path pointing to the experiment directory on persistent storage.
+
+        This can point to a remote storage location (e.g. S3) or to a local
+        location (path on the head node)."""
+        return self._experiment_fs_path
+
+    @property
+    def best_trial(self) -> Trial:
+        """Get the best trial of the experiment
+
+        The best trial is determined by comparing the last trial results
+        using the `metric` and `mode` parameters passed to `tune.run()`.
+
+        If you didn't pass these parameters, use
+        `get_best_trial(metric, mode, scope)` instead.
+        """
+        if not self.default_metric or not self.default_mode:
+            raise ValueError(
+                "To fetch the `best_trial`, pass a `metric` and `mode` "
+                "parameter to `tune.run()`. Alternatively, use the "
+                "`get_best_trial(metric, mode)` method to set the metric "
+                "and mode explicitly."
+            )
+        return self.get_best_trial(self.default_metric, self.default_mode)
+
+    @property
+    def best_config(self) -> Dict:
+        """Get the config of the best trial of the experiment
+
+        The best trial is determined by comparing the last trial results
+        using the `metric` and `mode` parameters passed to `tune.run()`.
+
+        If you didn't pass these parameters, use
+        `get_best_config(metric, mode, scope)` instead.
+        """
+        if not self.default_metric or not self.default_mode:
+            raise ValueError(
+                "To fetch the `best_config`, pass a `metric` and `mode` "
+                "parameter to `tune.run()`. Alternatively, use the "
+                "`get_best_config(metric, mode)` method to set the metric "
+                "and mode explicitly."
+            )
+        return self.get_best_config(self.default_metric, self.default_mode)
+
+    @property
+    def best_checkpoint(self) -> Checkpoint:
+        """Get the checkpoint path of the best trial of the experiment
+
+        The best trial is determined by comparing the last trial results
+        using the `metric` and `mode` parameters passed to `tune.run()`.
+
+        If you didn't pass these parameters, use
+        `get_best_checkpoint(trial, metric, mode)` instead.
+
+        Returns:
+            :class:`Checkpoint <ray.train.Checkpoint>` object.
+        """
+        if not self.default_metric or not self.default_mode:
+            raise ValueError(
+                "To fetch the `best_checkpoint`, pass a `metric` and `mode` "
+                "parameter to `tune.run()`. Alternatively, use the "
+                "`get_best_checkpoint(trial, metric, mode)` method to set the "
+                "metric and mode explicitly."
+            )
+        best_trial = self.best_trial
+        if not best_trial:
+            raise ValueError(
+                f"No best trial found. Please check if you specified the "
+                f"correct default metric ({self.default_metric}) and mode "
+                f"({self.default_mode})."
+            )
+        return self.get_best_checkpoint(
+            best_trial, self.default_metric, self.default_mode
+        )
+
+    @property
+    def best_dataframe(self) -> DataFrame:
+        """Get the full result dataframe of the best trial of the experiment
+
+        The best trial is determined by comparing the last trial results
+        using the `metric` and `mode` parameters passed to `tune.run()`.
+
+        If you didn't pass these parameters, use
+        `get_best_trial(metric, mode)` and use it to look for the dataframe
+        in the `self.trial_dataframes` dict.
+        """
+        if not self.default_metric or not self.default_mode:
+            raise ValueError(
+                "To fetch the `best_result`, pass a `metric` and `mode` "
+                "parameter to `tune.run()`."
+            )
+        return self.trial_dataframes[self.best_trial.trial_id]
+
+    @property
+    def best_result(self) -> Dict:
+        """Get the last result of the best trial of the experiment
+
+        The best trial is determined by comparing the last trial results
+        using the `metric` and `mode` parameters passed to `tune.run()`.
+
+        If you didn't pass these parameters, use
+        `get_best_trial(metric, mode, scope).last_result` instead.
+        """
+        if not self.default_metric or not self.default_mode:
+            raise ValueError(
+                "To fetch the `best_result`, pass a `metric` and `mode` "
+                "parameter to `tune.run()`. Alternatively, use "
+                "`get_best_trial(metric, mode).last_result` to set "
+                "the metric and mode explicitly and fetch the last result."
+            )
+        return self.best_trial.last_result
+
+    def _delimiter(self):
+        return os.environ.get("TUNE_RESULT_DELIM", "/")
+
+    @property
+    def best_result_df(self) -> DataFrame:
+        """Get the best result of the experiment as a pandas dataframe.
+
+        The best trial is determined by comparing the last trial results
+        using the `metric` and `mode` parameters passed to `tune.run()`.
+
+        If you didn't pass these parameters, use
+        `get_best_trial(metric, mode, scope).last_result` instead.
+        """
+        if not pd:
+            raise ValueError(
+                "`best_result_df` requires pandas. Install with "
+                "`pip install pandas`."
+            )
+
+        best_result = flatten_dict(self.best_result, delimiter=self._delimiter())
+        return pd.DataFrame.from_records([best_result], index="trial_id")
+
+    @property
+    def results(self) -> Dict[str, Dict]:
+        """Get the last result of the all trials of the experiment"""
+        return {trial.trial_id: trial.last_result for trial in self.trials}
+
+    @property
+    def results_df(self) -> DataFrame:
+        """Get all the last results as a pandas dataframe."""
+        if not pd:
+            raise ValueError(
+                "`results_df` requires pandas. Install with `pip install pandas`."
+            )
+        return pd.DataFrame.from_records(
+            [
+                flatten_dict(trial.last_result, delimiter=self._delimiter())
+                for trial in self.trials
+            ],
+            index="trial_id",
+        )
+
+    @property
+    def trial_dataframes(self) -> Dict[str, DataFrame]:
+        """List of all dataframes of the trials.
+
+        Each dataframe is indexed by iterations and contains reported
+        metrics.
+        """
+        return self._trial_dataframes
+
+    def dataframe(
+        self, metric: Optional[str] = None, mode: Optional[str] = None
+    ) -> DataFrame:
+        """Returns a pandas.DataFrame object constructed from the trials.
+
+        This function will look through all observed results of each trial
+        and return the one corresponding to the passed ``metric`` and
+        ``mode``: If ``mode=min``, it returns the result with the lowest
+        *ever* observed ``metric`` for this trial (this is not necessarily
+        the last)! For ``mode=max``, it's the highest, respectively. If
+        ``metric=None`` or ``mode=None``, the last result will be returned.
+
+        Args:
+            metric: Key for trial info to order on. If None, uses last result.
+            mode: One of [None, "min", "max"].
+
+        Returns:
+            pd.DataFrame: Constructed from a result dict of each trial.
+        """
+        # Do not validate metric/mode here or set from default metric/mode!
+        # Otherwise we will get confusing results as the lowest ever observed
+        # result may not be the last result.
+        if mode and mode not in ["min", "max"]:
+            raise ValueError("If set, `mode` has to be one of [min, max]")
+
+        if mode and not metric:
+            raise ValueError(
+                "If a `mode` is passed to `ExperimentAnalysis.dataframe(),"
+                " you'll also have to pass a `metric`!"
+            )
+
+        rows = self._retrieve_rows(metric=metric, mode=mode)
+        all_configs = self.get_all_configs(prefix=True)
+        for path, config in all_configs.items():
+            if path in rows:
+                rows[path].update(config)
+                rows[path].update(logdir=path)
+        return pd.DataFrame(list(rows.values()))
+
+    def _get_trial_checkpoints_with_metric(
+        self, trial: Trial, metric: Optional[str] = None
+    ) -> List[Tuple[Checkpoint, Number]]:
+        """Get all checkpoints and a specified metric of a trial.
+
+        Args:
+            trial: The log directory of a trial, or a trial instance.
+            metric: key for trial info to return, e.g. "mean_accuracy".
+                "training_iteration" is used by default if no value was
+                passed to ``self.default_metric``.
+
+        Returns:
+            List of [Checkpoint, metric] for all checkpoints of the trial.
+        """
+        metric = metric or self.default_metric or TRAINING_ITERATION
+
+        best_checkpoint_results = (
+            trial.run_metadata.checkpoint_manager.best_checkpoint_results
+        )
+        best_checkpoints = [
+            (checkpoint_result.checkpoint, checkpoint_result.metrics)
+            for checkpoint_result in best_checkpoint_results
+        ]
+        # Support nested metrics given as flattened strings, e.g.
+        # "info/learner/default_policy/policy_loss".
+        return [
+            (checkpoint, unflattened_lookup(metric, metrics))
+            for checkpoint, metrics in best_checkpoints
+        ]
+
+    def get_best_checkpoint(
+        self,
+        trial: Trial,
+        metric: Optional[str] = None,
+        mode: Optional[str] = None,
+    ) -> Optional[Checkpoint]:
+        """Gets best persistent checkpoint path of provided trial.
+
+        Any checkpoints with an associated metric value of ``nan`` will be filtered out.
+
+        Args:
+            trial: The log directory of a trial, or a trial instance.
+            metric: key of trial info to return, e.g. "mean_accuracy".
+                "training_iteration" is used by default if no value was
+                passed to ``self.default_metric``.
+            mode: One of [min, max]. Defaults to ``self.default_mode``.
+
+        Returns:
+            A :class:`Checkpoint <ray.train.Checkpoint>` object
+        """
+        metric = metric or self.default_metric or TRAINING_ITERATION
+        mode = self._validate_mode(mode)
+
+        checkpoints_and_metrics = self._get_trial_checkpoints_with_metric(trial, metric)
+
+        # Filter out nan. Sorting nan values leads to undefined behavior.
+        checkpoints_and_metrics = list(
+            filter(lambda x: not is_nan(x[1]), checkpoints_and_metrics)
+        )
+
+        if not checkpoints_and_metrics:
+            logger.error(f"No checkpoints have been found for trial {trial}.")
+            return None
+
+        score_order_factor = -1 if mode == "min" else 1
+        best_checkpoint, _ = max(
+            checkpoints_and_metrics, key=lambda x: score_order_factor * x[1]
+        )
+        return best_checkpoint
+
+    def get_best_trial(
+        self,
+        metric: Optional[str] = None,
+        mode: Optional[str] = None,
+        scope: str = "last",
+        filter_nan_and_inf: bool = True,
+    ) -> Optional[Trial]:
+        """Retrieve the best trial object.
+
+        Compares all trials' scores on ``metric``.
+        If ``metric`` is not specified, ``self.default_metric`` will be used.
+        If `mode` is not specified, ``self.default_mode`` will be used.
+        These values are usually initialized by passing the ``metric`` and
+        ``mode`` parameters to ``tune.run()``.
+
+        Args:
+            metric: Key for trial info to order on. Defaults to
+                ``self.default_metric``.
+            mode: One of [min, max]. Defaults to ``self.default_mode``.
+            scope: One of [all, last, avg, last-5-avg, last-10-avg].
+                If `scope=last`, only look at each trial's final step for
+                `metric`, and compare across trials based on `mode=[min,max]`.
+                If `scope=avg`, consider the simple average over all steps
+                for `metric` and compare across trials based on
+                `mode=[min,max]`. If `scope=last-5-avg` or `scope=last-10-avg`,
+                consider the simple average over the last 5 or 10 steps for
+                `metric` and compare across trials based on `mode=[min,max]`.
+                If `scope=all`, find each trial's min/max score for `metric`
+                based on `mode`, and compare trials based on `mode=[min,max]`.
+            filter_nan_and_inf: If True (default), NaN or infinite
+                values are disregarded and these trials are never selected as
+                the best trial.
+
+        Returns:
+            The best trial for the provided metric. If no trials contain the provided
+                metric, or if the value for the metric is NaN for all trials,
+                then returns None.
+        """
+        if len(self.trials) == 1:
+            return self.trials[0]
+
+        metric = self._validate_metric(metric)
+        mode = self._validate_mode(mode)
+
+        if scope not in ["all", "last", "avg", "last-5-avg", "last-10-avg"]:
+            raise ValueError(
+                "ExperimentAnalysis: attempting to get best trial for "
+                'metric {} for scope {} not in ["all", "last", "avg", '
+                '"last-5-avg", "last-10-avg"]. '
+                "If you didn't pass a `metric` parameter to `tune.run()`, "
+                "you have to pass one when fetching the best trial.".format(
+                    metric, scope
+                )
+            )
+        best_trial = None
+        best_metric_score = None
+
+        for trial in self.trials:
+            if metric not in trial.metric_analysis:
+                continue
+
+            if scope in ["last", "avg", "last-5-avg", "last-10-avg"]:
+                metric_score = trial.metric_analysis[metric][scope]
+            else:
+                metric_score = trial.metric_analysis[metric][mode]
+
+            if filter_nan_and_inf and is_nan_or_inf(metric_score):
+                continue
+
+            if best_metric_score is None:
+                best_metric_score = metric_score
+                best_trial = trial
+                continue
+
+            if (mode == "max") and (best_metric_score < metric_score):
+                best_metric_score = metric_score
+                best_trial = trial
+            elif (mode == "min") and (best_metric_score > metric_score):
+                best_metric_score = metric_score
+                best_trial = trial
+
+        if not best_trial:
+            logger.warning(
+                "Could not find best trial. Did you pass the correct `metric` "
+                "parameter?"
+            )
+        return best_trial
+
+    def get_best_config(
+        self,
+        metric: Optional[str] = None,
+        mode: Optional[str] = None,
+        scope: str = "last",
+    ) -> Optional[Dict]:
+        """Retrieve the best config corresponding to the trial.
+
+        Compares all trials' scores on `metric`.
+        If ``metric`` is not specified, ``self.default_metric`` will be used.
+        If `mode` is not specified, ``self.default_mode`` will be used.
+        These values are usually initialized by passing the ``metric`` and
+        ``mode`` parameters to ``tune.run()``.
+
+        Args:
+            metric: Key for trial info to order on. Defaults to
+                ``self.default_metric``.
+            mode: One of [min, max]. Defaults to ``self.default_mode``.
+            scope: One of [all, last, avg, last-5-avg, last-10-avg].
+                If `scope=last`, only look at each trial's final step for
+                `metric`, and compare across trials based on `mode=[min,max]`.
+                If `scope=avg`, consider the simple average over all steps
+                for `metric` and compare across trials based on
+                `mode=[min,max]`. If `scope=last-5-avg` or `scope=last-10-avg`,
+                consider the simple average over the last 5 or 10 steps for
+                `metric` and compare across trials based on `mode=[min,max]`.
+                If `scope=all`, find each trial's min/max score for `metric`
+                based on `mode`, and compare trials based on `mode=[min,max]`.
+        """
+        best_trial = self.get_best_trial(metric, mode, scope)
+        return best_trial.config if best_trial else None
+
+    def get_last_checkpoint(
+        self, trial=None, metric="training_iteration", mode="max"
+    ) -> Optional[Checkpoint]:
+        """Gets the last checkpoint of the provided trial,
+        i.e., with the highest "training_iteration".
+
+        If no trial is specified, it loads the best trial according to the
+        provided metric and mode (defaults to max. training iteration).
+
+        Args:
+            trial: If None, load the best trial automatically.
+            metric: If no trial is specified, use this metric to identify
+                the best trial and load the last checkpoint from this trial.
+            mode: If no trial is specified, use the metric and this mode
+                to identify the best trial and load the last checkpoint from it.
+
+        Returns:
+            Path for last checkpoint of trial
+        """
+        trial = trial or self.get_best_trial(metric, mode)
+        return self.get_best_checkpoint(trial, TRAINING_ITERATION, "max")
+
+    def _validate_metric(self, metric: str) -> str:
+        if not metric and not self.default_metric:
+            raise ValueError(
+                "No `metric` has been passed and  `default_metric` has "
+                "not been set. Please specify the `metric` parameter."
+            )
+        return metric or self.default_metric
+
+    def _validate_mode(self, mode: str) -> str:
+        if not mode and not self.default_mode:
+            raise ValueError(
+                "No `mode` has been passed and  `default_mode` has "
+                "not been set. Please specify the `mode` parameter."
+            )
+        if mode and mode not in ["min", "max"]:
+            raise ValueError("If set, `mode` has to be one of [min, max]")
+        return mode or self.default_mode
+
+    def _retrieve_rows(
+        self, metric: Optional[str] = None, mode: Optional[str] = None
+    ) -> Dict[str, Any]:
+        assert mode is None or mode in ["max", "min"]
+        assert not mode or metric
+        rows = {}
+        for path, df in self.trial_dataframes.items():
+            if df.empty:
+                continue
+            if metric not in df:
+                idx = -1
+            elif mode == "max":
+                idx = df[metric].idxmax()
+            elif mode == "min":
+                idx = df[metric].idxmin()
+            else:
+                idx = -1
+            try:
+                rows[path] = df.iloc[idx].to_dict()
+            except TypeError:
+                # idx is nan
+                logger.warning(
+                    "Warning: Non-numerical value(s) encountered for {}".format(path)
+                )
+
+        return rows
+
+    def __getstate__(self) -> Dict[str, Any]:
+        """Ensure that trials are marked as stubs when pickling,
+        so that they can be loaded later without the trainable
+        being registered.
+        """
+        state = self.__dict__.copy()
+
+        def make_stub_if_needed(trial: Trial) -> Trial:
+            if trial.stub:
+                return trial
+            trial_copy = Trial(trial.trainable_name, stub=True)
+            trial_copy.__setstate__(trial.__getstate__())
+            return trial_copy
+
+        state["trials"] = [make_stub_if_needed(t) for t in state["trials"]]
+        return state

.venv/lib/python3.11/site-packages/ray/tune/impl/config.py

@@ -0,0 +1,46 @@
+from dataclasses import dataclass
+
+from ray.air.config import CheckpointConfig as _CheckpointConfig
+from ray.air.config import FailureConfig as _FailureConfig
+from ray.air.config import RunConfig as _RunConfig
+from ray.train.constants import _v2_migration_warnings_enabled
+from ray.train.utils import _copy_doc, _log_deprecation_warning
+
+# NOTE: This is just a pass-through wrapper around `ray.train.RunConfig`
+# in order to detect whether the import module was correct (e.g. `ray.tune.RunConfig`).
+
+
+@dataclass
+@_copy_doc(_CheckpointConfig)
+class CheckpointConfig(_CheckpointConfig):
+    pass
+
+
+@dataclass
+@_copy_doc(_FailureConfig)
+class FailureConfig(_FailureConfig):
+    pass
+
+
+@dataclass
+@_copy_doc(_RunConfig)
+class RunConfig(_RunConfig):
+    def __post_init__(self):
+        self.checkpoint_config = self.checkpoint_config or CheckpointConfig()
+        self.failure_config = self.failure_config or FailureConfig()
+
+        super().__post_init__()
+
+        if not isinstance(self.checkpoint_config, CheckpointConfig):
+            if _v2_migration_warnings_enabled():
+                _log_deprecation_warning(
+                    "The `CheckpointConfig` class should be imported from `ray.tune` "
+                    "when passing it to the Tuner. Please update your imports."
+                )
+
+        if not isinstance(self.failure_config, FailureConfig):
+            if _v2_migration_warnings_enabled():
+                _log_deprecation_warning(
+                    "The `FailureConfig` class should be imported from `ray.tune` "
+                    "when passing it to the Tuner. Please update your imports."
+                )

.venv/lib/python3.11/site-packages/ray/tune/impl/out_of_band_serialize_dataset.py

@@ -0,0 +1,33 @@
+import contextlib
+import traceback
+
+import ray
+
+
+def _deserialize_and_fully_execute_if_needed(serialized_ds: bytes):
+    ds = ray.data.Dataset.deserialize_lineage(serialized_ds)
+    return ds
+
+
+def _reduce(ds: ray.data.Dataset):
+    tb_list = traceback.format_list(traceback.extract_stack())
+    _already_in_out_of_band_serialization = False
+    for tb in tb_list:
+        # TODO(xwjiang): Let's make this less hacky.
+        if "serialize_lineage" in tb:
+            _already_in_out_of_band_serialization = True
+            break
+    if not _already_in_out_of_band_serialization and ds.has_serializable_lineage():
+        return _deserialize_and_fully_execute_if_needed, (ds.serialize_lineage(),)
+    else:
+        return ds.__reduce__()
+
+
+@contextlib.contextmanager
+def out_of_band_serialize_dataset():
+    context = ray._private.worker.global_worker.get_serialization_context()
+    try:
+        context._register_cloudpickle_reducer(ray.data.Dataset, _reduce)
+        yield
+    finally:
+        context._unregister_cloudpickle_reducer(ray.data.Dataset)

.venv/lib/python3.11/site-packages/ray/tune/impl/placeholder.py

@@ -0,0 +1,244 @@
+import hashlib
+from collections import defaultdict
+from typing import Any, Dict, Tuple
+
+from ray.tune.search.sample import Categorical, Domain, Function
+from ray.tune.search.variant_generator import assign_value
+from ray.util.annotations import DeveloperAPI
+
+ID_HASH_LENGTH = 8
+
+
+def create_resolvers_map():
+    return defaultdict(list)
+
+
+def _id_hash(path_tuple):
+    """Compute a hash for the specific placeholder based on its path."""
+    return hashlib.sha1(str(path_tuple).encode("utf-8")).hexdigest()[:ID_HASH_LENGTH]
+
+
+class _FunctionResolver:
+    """Replaced value for function typed objects."""
+
+    TOKEN = "__fn_ph"
+
+    def __init__(self, hash, fn):
+        self.hash = hash
+        self._fn = fn
+
+    def resolve(self, config: Dict):
+        """Some functions take a resolved spec dict as input.
+
+        Note: Function placeholders are independently sampled during
+        resolution. Therefore their random states are not restored.
+        """
+        return self._fn.sample(config=config)
+
+    def get_placeholder(self) -> str:
+        return (self.TOKEN, self.hash)
+
+
+class _RefResolver:
+    """Replaced value for all other non-primitive objects."""
+
+    TOKEN = "__ref_ph"
+
+    def __init__(self, hash, value):
+        self.hash = hash
+        self._value = value
+
+    def resolve(self):
+        return self._value
+
+    def get_placeholder(self) -> str:
+        return (self.TOKEN, self.hash)
+
+
+def _is_primitive(x):
+    """Returns True if x is a primitive type.
+
+    Primitive types are int, float, str, bool, and None.
+    """
+    return isinstance(x, (int, float, str, bool)) or x is None
+
+
+@DeveloperAPI
+def inject_placeholders(
+    config: Any,
+    resolvers: defaultdict,
+    id_prefix: Tuple = (),
+    path_prefix: Tuple = (),
+) -> Dict:
+    """Replaces reference objects contained by a config dict with placeholders.
+
+    Given a config dict, this function replaces all reference objects contained
+    by this dict with placeholder strings. It recursively expands nested dicts
+    and lists, and properly handles Tune native search objects such as Categorical
+    and Function.
+    This makes sure the config dict only contains primitive typed values, which
+    can then be handled by different search algorithms.
+
+    A few details about id_prefix and path_prefix. Consider the following config,
+    where "param1" is a simple grid search of 3 tuples.
+
+    config = {
+        "param1": tune.grid_search([
+            (Cat, None, None),
+            (None, Dog, None),
+            (None, None, Fish),
+        ]),
+    }
+
+    We will replace the 3 objects contained with placeholders. And after trial
+    expansion, the config may look like this:
+
+    config = {
+        "param1": (None, (placeholder, hash), None)
+    }
+
+    Now you need 2 pieces of information to resolve the placeholder. One is the
+    path of ("param1", 1), which tells you that the first element of the tuple
+    under "param1" key is a placeholder that needs to be resolved.
+    The other is the mapping from the placeholder to the actual object. In this
+    case hash -> Dog.
+
+    id and path prefixes serve exactly this purpose here. The difference between
+    these two is that id_prefix is the location of the value in the pre-injected
+    config tree. So if a value is the second option in a grid_search, it gets an
+    id part of 1. Injected placeholders all get unique id prefixes. path prefix
+    identifies a placeholder in the expanded config tree. So for example, all
+    options of a single grid_search will get the same path prefix. This is how
+    we know which location has a placeholder to be resolved in the post-expansion
+    tree.
+
+    Args:
+        config: The config dict to replace references in.
+        resolvers: A dict from path to replaced objects.
+        id_prefix: The prefix to prepend to id every single placeholders.
+        path_prefix: The prefix to prepend to every path identifying
+            potential locations of placeholders in an expanded tree.
+
+    Returns:
+        The config with all references replaced.
+    """
+    if isinstance(config, dict) and "grid_search" in config and len(config) == 1:
+        config["grid_search"] = [
+            # Different options gets different id prefixes.
+            # But we should omit appending to path_prefix because after expansion,
+            # this level will not be there.
+            inject_placeholders(choice, resolvers, id_prefix + (i,), path_prefix)
+            for i, choice in enumerate(config["grid_search"])
+        ]
+        return config
+    elif isinstance(config, dict):
+        return {
+            k: inject_placeholders(v, resolvers, id_prefix + (k,), path_prefix + (k,))
+            for k, v in config.items()
+        }
+    elif isinstance(config, list):
+        return [
+            inject_placeholders(elem, resolvers, id_prefix + (i,), path_prefix + (i,))
+            for i, elem in enumerate(config)
+        ]
+    elif isinstance(config, tuple):
+        return tuple(
+            inject_placeholders(elem, resolvers, id_prefix + (i,), path_prefix + (i,))
+            for i, elem in enumerate(config)
+        )
+    elif _is_primitive(config):
+        # Primitive types.
+        return config
+    elif isinstance(config, Categorical):
+        config.categories = [
+            # Different options gets different id prefixes.
+            # But we should omit appending to path_prefix because after expansion,
+            # this level will not be there.
+            inject_placeholders(choice, resolvers, id_prefix + (i,), path_prefix)
+            for i, choice in enumerate(config.categories)
+        ]
+        return config
+    elif isinstance(config, Function):
+        # Function type.
+        id_hash = _id_hash(id_prefix)
+        v = _FunctionResolver(id_hash, config)
+        resolvers[path_prefix].append(v)
+        return v.get_placeholder()
+    elif not isinstance(config, Domain):
+        # Other non-search space reference objects, dataset, actor handle, etc.
+        id_hash = _id_hash(id_prefix)
+        v = _RefResolver(id_hash, config)
+        resolvers[path_prefix].append(v)
+        return v.get_placeholder()
+    else:
+        # All the other cases, do nothing.
+        return config
+
+
+def _get_placeholder(config: Any, prefix: Tuple, path: Tuple):
+    if not path:
+        return prefix, config
+
+    key = path[0]
+    if isinstance(config, tuple):
+        if config[0] in (_FunctionResolver.TOKEN, _RefResolver.TOKEN):
+            # Found a matching placeholder.
+            # Note that we do not require that the full path are consumed before
+            # declaring a match. Because this placeholder may be part of a nested
+            # search space. For example, the following config:
+            #   config = {
+            #       "param1": tune.grid_search([
+            #           tune.grid_search([Object1, 2, 3]),
+            #           tune.grid_search([Object2, 5, 6]),
+            #       ]),
+            #   }
+            # will result in placeholders under path ("param1", 0, 0).
+            # After expansion though, the choosen placeholder will live under path
+            # ("param1", 0) like this: config = {"param1": (Placeholder1, 2, 3)}
+            return prefix, config
+        elif key < len(config):
+            return _get_placeholder(
+                config[key], prefix=prefix + (path[0],), path=path[1:]
+            )
+    elif (isinstance(config, dict) and key in config) or (
+        isinstance(config, list) and key < len(config)
+    ):
+        # Expand config tree recursively.
+        return _get_placeholder(config[key], prefix=prefix + (path[0],), path=path[1:])
+
+    # Can not find a matching placeholder.
+    return None, None
+
+
+@DeveloperAPI
+def resolve_placeholders(config: Any, replaced: defaultdict):
+    """Replaces placeholders contained by a config dict with the original values.
+
+    Args:
+        config: The config to replace placeholders in.
+        replaced: A dict from path to replaced objects.
+    """
+
+    def __resolve(resolver_type, args):
+        for path, resolvers in replaced.items():
+            assert resolvers
+
+            if not isinstance(resolvers[0], resolver_type):
+                continue
+
+            prefix, ph = _get_placeholder(config, (), path)
+            if not ph:
+                # Represents an unchosen value. Just skip.
+                continue
+
+            for resolver in resolvers:
+                if resolver.hash != ph[1]:
+                    continue
+                # Found the matching resolver.
+                assign_value(config, prefix, resolver.resolve(*args))
+
+    # RefResolvers first.
+    __resolve(_RefResolver, args=())
+    # Functions need to be resolved after RefResolvers, in case they are
+    # referencing values from the RefResolvers.
+    __resolve(_FunctionResolver, args=(config,))

.venv/lib/python3.11/site-packages/ray/tune/impl/test_utils.py

@@ -0,0 +1,66 @@
+from sklearn.datasets import load_breast_cancer
+
+from ray import tune
+from ray.data import Dataset, Datasource, ReadTask, read_datasource
+from ray.data.block import BlockMetadata
+from ray.tune.impl.utils import execute_dataset
+
+
+# TODO(xwjiang): Enable this when Clark's out-of-band-serialization is landed.
+class TestDatasource(Datasource):
+    def prepare_read(self, parallelism: int, **read_args):
+        import pyarrow as pa
+
+        def load_data():
+            data_raw = load_breast_cancer(as_frame=True)
+            dataset_df = data_raw["data"]
+            dataset_df["target"] = data_raw["target"]
+            return [pa.Table.from_pandas(dataset_df)]
+
+        meta = BlockMetadata(
+            num_rows=None,
+            size_bytes=None,
+            schema=None,
+            input_files=None,
+            exec_stats=None,
+        )
+        return [ReadTask(load_data, meta)]
+
+
+def gen_dataset_func() -> Dataset:
+    test_datasource = TestDatasource()
+    return read_datasource(test_datasource)
+
+
+def test_grid_search():
+    ds1 = gen_dataset_func().lazy().map(lambda x: x)
+    ds2 = gen_dataset_func().lazy().map(lambda x: x)
+    assert not ds1._plan._has_final_stage_snapshot()
+    assert not ds2._plan._has_final_stage_snapshot()
+    param_space = {"train_dataset": tune.grid_search([ds1, ds2])}
+    execute_dataset(param_space)
+    executed_ds = param_space["train_dataset"]["grid_search"]
+    assert len(executed_ds) == 2
+    assert executed_ds[0]._plan._has_final_stage_snapshot()
+    assert executed_ds[1]._plan._has_final_stage_snapshot()
+
+
+def test_choice():
+    ds1 = gen_dataset_func().lazy().map(lambda x: x)
+    ds2 = gen_dataset_func().lazy().map(lambda x: x)
+    assert not ds1._plan._has_final_stage_snapshot()
+    assert not ds2._plan._has_final_stage_snapshot()
+    param_space = {"train_dataset": tune.choice([ds1, ds2])}
+    execute_dataset(param_space)
+    executed_ds = param_space["train_dataset"].categories
+    assert len(executed_ds) == 2
+    assert executed_ds[0]._plan._has_final_stage_snapshot()
+    assert executed_ds[1]._plan._has_final_stage_snapshot()
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

.venv/lib/python3.11/site-packages/ray/tune/impl/tuner_internal.py

@@ -0,0 +1,669 @@
+import copy
+import io
+import logging
+import math
+from pathlib import Path
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    List,
+    Optional,
+    Tuple,
+    Type,
+    Union,
+)
+
+import pyarrow.fs
+
+import ray.cloudpickle as pickle
+import ray.train
+from ray.air._internal.uri_utils import URI
+from ray.air._internal.usage import AirEntrypoint
+from ray.train import ScalingConfig
+from ray.train._internal.storage import StorageContext, get_fs_and_path
+from ray.train.constants import _v2_migration_warnings_enabled
+from ray.train.utils import _log_deprecation_warning
+from ray.tune import (
+    Experiment,
+    ExperimentAnalysis,
+    ResumeConfig,
+    RunConfig,
+    TuneConfig,
+    TuneError,
+)
+from ray.tune.registry import is_function_trainable
+from ray.tune.result_grid import ResultGrid
+from ray.tune.trainable import Trainable
+from ray.tune.tune import _Config, run
+from ray.tune.utils import flatten_dict
+from ray.util import inspect_serializability
+
+if TYPE_CHECKING:
+    from ray.train.trainer import BaseTrainer
+    from ray.util.queue import Queue
+
+
+_TUNER_PKL = "tuner.pkl"
+_TRAINABLE_KEY = "_trainable"
+_CONVERTED_TRAINABLE_KEY = "_converted_trainable"
+_PARAM_SPACE_KEY = "_param_space"
+_EXPERIMENT_ANALYSIS_KEY = "_experiment_analysis"
+
+logger = logging.getLogger(__name__)
+
+TrainableType = Union[str, Callable, Type[Trainable]]
+TrainableTypeOrTrainer = Union[TrainableType, "BaseTrainer"]
+
+
+class TunerInternal:
+    """The real implementation behind external facing ``Tuner``.
+
+    The external facing ``Tuner`` multiplexes between local Tuner and remote Tuner
+    depending on whether in Ray client mode.
+
+    In Ray client mode, external ``Tuner`` wraps ``TunerInternal`` into a remote actor,
+    which is guaranteed to be placed on head node.
+
+    ``TunerInternal`` can be constructed from fresh, in which case, ``trainable`` needs
+    to be provided, together with optional ``param_space``, ``tune_config`` and
+    ``run_config``.
+
+    It can also be restored from a previous failed run (given ``restore_path``).
+
+    Args:
+        restore_path: The path from where the Tuner can be restored. If provided, None
+            of the rest args are needed.
+        resume_config: Resume config to configure which trials to continue.
+        trainable: The trainable to be tuned.
+        param_space: Search space of the tuning job.
+            One thing to note is that both preprocessor and dataset can be tuned here.
+        tune_config: Tuning algorithm specific configs.
+            Refer to ray.tune.tune_config.TuneConfig for more info.
+        run_config: Runtime configuration that is specific to individual trials.
+            If passed, this will overwrite the run config passed to the Trainer,
+            if applicable. Refer to ray.tune.RunConfig for more info.
+    """
+
+    def __init__(
+        self,
+        restore_path: str = None,
+        storage_filesystem: Optional[pyarrow.fs.FileSystem] = None,
+        resume_config: Optional[ResumeConfig] = None,
+        trainable: Optional[TrainableTypeOrTrainer] = None,
+        param_space: Optional[Dict[str, Any]] = None,
+        tune_config: Optional[TuneConfig] = None,
+        run_config: Optional[RunConfig] = None,
+        _tuner_kwargs: Optional[Dict] = None,
+        _entrypoint: AirEntrypoint = AirEntrypoint.TUNER,
+    ):
+        from ray.train.trainer import BaseTrainer
+
+        if isinstance(trainable, BaseTrainer):
+            if _v2_migration_warnings_enabled():
+                _log_deprecation_warning(
+                    "Passing a Trainer to the Tuner is deprecated. "
+                    "See the section on hyperparameter optimization in this "
+                    "REP for more information: "
+                    "https://github.com/ray-project/enhancements/pull/57"
+                )
+
+            run_config = self._choose_run_config(
+                tuner_run_config=run_config,
+                trainer=trainable,
+                param_space=param_space,
+            )
+
+        self._tune_config = tune_config or TuneConfig()
+        self._run_config = copy.copy(run_config) or RunConfig()
+
+        if not isinstance(self._run_config, RunConfig):
+            if _v2_migration_warnings_enabled():
+                _log_deprecation_warning(
+                    "The `RunConfig` class should be imported from `ray.tune` "
+                    "when passing it to the Tuner. Please update your imports."
+                )
+
+        self._entrypoint = _entrypoint
+
+        # Restore from Tuner checkpoint.
+        if restore_path:
+            self._restore_from_path_or_uri(
+                path_or_uri=restore_path,
+                trainable=trainable,
+                overwrite_param_space=param_space,
+                resume_config=resume_config,
+                storage_filesystem=storage_filesystem,
+            )
+            return
+
+        # Start from fresh
+        if not trainable:
+            raise TuneError("You need to provide a trainable to tune.")
+
+        self.trainable = trainable
+        assert self.converted_trainable
+        self._validate_trainable(self.converted_trainable)
+
+        self.param_space = param_space
+
+        self._resume_config = None
+        self._is_restored = False
+        self._tuner_kwargs = copy.deepcopy(_tuner_kwargs) or {}
+        self._experiment_analysis = None
+
+        self._run_config.name = (
+            self._run_config.name
+            or StorageContext.get_experiment_dir_name(self.converted_trainable)
+        )
+        # The storage context here is only used to access the resolved
+        # storage fs and experiment path, in order to avoid duplicating that logic.
+        # This is NOT the storage context object that gets passed to remote workers.
+        storage = StorageContext(
+            storage_path=self._run_config.storage_path,
+            experiment_dir_name=self._run_config.name,
+            storage_filesystem=self._run_config.storage_filesystem,
+        )
+
+        fs = storage.storage_filesystem
+        fs.create_dir(storage.experiment_fs_path)
+        with fs.open_output_stream(
+            Path(storage.experiment_fs_path, _TUNER_PKL).as_posix()
+        ) as f:
+            f.write(pickle.dumps(self.__getstate__()))
+
+    def get_run_config(self) -> RunConfig:
+        return self._run_config
+
+    # For Jupyter output with Ray Client
+    def set_run_config_and_remote_string_queue(
+        self, run_config: RunConfig, string_queue: "Queue"
+    ):
+        self._run_config = run_config
+        self._tuner_kwargs["_remote_string_queue"] = string_queue
+
+    def clear_remote_string_queue(self):
+        self._tuner_kwargs.pop("_remote_string_queue", None)
+
+    def _expected_utilization(self, cpus_per_trial, cpus_total):
+        num_samples = self._tune_config.num_samples
+        if num_samples < 0:  # TODO: simplify this in Tune
+            num_samples = math.inf
+        concurrent_trials = self._tune_config.max_concurrent_trials or 0
+        if concurrent_trials < 1:  # TODO: simplify this in Tune
+            concurrent_trials = math.inf
+
+        actual_concurrency = min(
+            (
+                (cpus_total // cpus_per_trial) if cpus_per_trial else 0,
+                num_samples,
+                concurrent_trials,
+            )
+        )
+        return (actual_concurrency * cpus_per_trial) / (cpus_total + 0.001)
+
+    def _validate_trainable(
+        self, trainable: TrainableType, required_trainable_name: Optional[str] = None
+    ):
+        """Determines whether or not the trainable is valid.
+
+        This includes checks on the serializability of the trainable, as well
+        asserting that the trainable name is as expected on restoration.
+
+        This trainable name validation is needed due to an implementation detail
+        where the trainable name (which is differently generated depending on
+        the trainable type) is saved in the Trial metadata and needs to match
+        upon restoration. This does not affect the typical path, since `Tuner.restore`
+        expects the exact same trainable (which will have the same name).
+
+        Raises:
+            ValueError: if the trainable name does not match or if the trainable
+                is not serializable.
+        """
+        try:
+            pickle.dumps(trainable)
+        except TypeError as e:
+            sio = io.StringIO()
+            inspect_serializability(trainable, print_file=sio)
+            msg = (
+                "The provided trainable is not serializable, which is a requirement "
+                "since the trainable is serialized and deserialized when transferred "
+                "to remote workers. See below for a trace of the non-serializable "
+                "objects that were found in your trainable:\n"
+                f"{sio.getvalue()}"
+            )
+            raise TypeError(msg) from e
+
+        if not required_trainable_name:
+            return
+
+        trainable_name = Experiment.get_trainable_name(trainable)
+
+        if trainable_name != required_trainable_name:
+            raise ValueError(
+                "Invalid `trainable` input to `Tuner.restore()`. To fix this error, "
+                "pass in the same trainable that was used to initialize the Tuner. "
+                "Got a trainable with identifier "
+                f"'{trainable_name}' but expected '{required_trainable_name}'."
+            )
+
+    def _set_trainable_on_restore(
+        self, trainable: TrainableType, old_trainable_name: Optional[str]
+    ):
+        from ray.train.base_trainer import BaseTrainer
+
+        self.trainable = trainable
+        assert self.converted_trainable
+        self._validate_trainable(
+            trainable=self.converted_trainable,
+            required_trainable_name=old_trainable_name,
+        )
+
+        if isinstance(self.trainable, BaseTrainer):
+            # Log a warning in case the user tries to modify the
+            # `RunConfig` from the Trainer
+            trainer: BaseTrainer = self.trainable
+
+            # Only log if the Trainer has a non-default RunConfig
+            if trainer.run_config != RunConfig():
+                logger.warning(
+                    "The Tune experiment will restore using the original run's "
+                    "`RunConfig`. If you made any changes to the `RunConfig` "
+                    "within the Trainer you passed into `Tuner.restore`, "
+                    "they will be ignored in the resumed run."
+                )
+
+            trainer.run_config = self._run_config
+
+    def _validate_param_space_on_restore(
+        self,
+        new_param_space: Dict[str, Any],
+        flattened_param_space_keys: Optional[List[str]],
+    ):
+        """Determines whether the (optionally) re-specified `param_space` is valid.
+
+        This method performs very loose validation on the new param_space to
+        prevent users from trying to specify new hyperparameters to tune over.
+
+        Raises:
+            ValueError: if not all keys match the original param_space.
+        """
+        if flattened_param_space_keys is None:
+            # Backwards compatibility: skip validation
+            return
+
+        keys = sorted(flatten_dict(new_param_space).keys())
+        if keys != flattened_param_space_keys:
+            raise ValueError(
+                "Invalid `param_space` input to `Tuner.restore()`. To fix this error, "
+                "pass in the same `param_space` that was used to initialize the Tuner. "
+                "Only re-specify the `param_space` to refresh Ray object references "
+                "that no longer exist due to restoring from a new Ray cluster session. "
+                "It should not be used to introduce new hyperparameters to tune."
+                f"\n\nGot: {keys}\nExpected: {flattened_param_space_keys}"
+            )
+
+    def _set_param_space_on_restore(
+        self,
+        param_space: Optional[Dict[str, Any]],
+        flattened_param_space_keys: Optional[List[str]],
+    ):
+        self.param_space = param_space
+
+        if self.param_space is not None:
+            # param_space = None -> use the original param_space
+            self._validate_param_space_on_restore(
+                new_param_space=self.param_space,
+                flattened_param_space_keys=flattened_param_space_keys,
+            )
+
+    def _load_tuner_state(
+        self, tuner_state: Dict[str, Any]
+    ) -> Tuple[Optional[str], Optional[List[str]]]:
+        """Loads Tuner state from the previously saved `tuner.pkl`.
+
+        Args:
+            tuner_pkl_path: pathlib.Path of the `tuner.pkl` file saved during the
+                original Tuner initialization.
+
+        Returns:
+            tuple: of `(old_trainable_name, flattened_param_space_keys)` used for
+                validating the re-specified `trainable` and `param_space`.
+        """
+        # NOTE: These are magic keys used for validating restore args.
+        old_trainable_name = tuner_state.pop("__trainable_name", None)
+        flattened_param_space_keys = tuner_state.pop(
+            "__flattened_param_space_keys", None
+        )
+
+        self.__setstate__(tuner_state)
+
+        return old_trainable_name, flattened_param_space_keys
+
+    def _restore_from_path_or_uri(
+        self,
+        path_or_uri: str,
+        trainable: TrainableTypeOrTrainer,
+        overwrite_param_space: Optional[Dict[str, Any]],
+        resume_config: ResumeConfig,
+        storage_filesystem: Optional[pyarrow.fs.FileSystem],
+    ):
+        fs, fs_path = get_fs_and_path(path_or_uri, storage_filesystem)
+        with fs.open_input_file(Path(fs_path, _TUNER_PKL).as_posix()) as f:
+            tuner_state = pickle.loads(f.readall())
+
+        old_trainable_name, flattened_param_space_keys = self._load_tuner_state(
+            tuner_state
+        )
+
+        # Perform validation and set the re-specified `trainable` and `param_space`
+        self._set_trainable_on_restore(
+            trainable=trainable, old_trainable_name=old_trainable_name
+        )
+        self._set_param_space_on_restore(
+            param_space=overwrite_param_space,
+            flattened_param_space_keys=flattened_param_space_keys,
+        )
+
+        # Update RunConfig to reflect changes in the experiment directory
+        path_or_uri_obj = URI(path_or_uri)
+
+        # Infer the `storage_path` and run `name` of the restored run using the
+        # experiment directory.
+        # Ex: ~/ray_results/exp_name -> ~/ray_results, exp_name
+        # Ex: s3://bucket/exp_name -> s3://bucket, exp_name
+        self._run_config.name = path_or_uri_obj.name
+        self._run_config.storage_path = str(path_or_uri_obj.parent)
+        # Update the storage_filesystem with the one passed in on restoration, if any.
+        self._run_config.storage_filesystem = storage_filesystem
+
+        # Load the experiment results at the point where it left off.
+        try:
+            self._experiment_analysis = ExperimentAnalysis(
+                experiment_checkpoint_path=path_or_uri,
+                default_metric=self._tune_config.metric,
+                default_mode=self._tune_config.mode,
+                storage_filesystem=storage_filesystem,
+            )
+        except Exception:
+            self._experiment_analysis = None
+
+        self._resume_config = resume_config
+        self._is_restored = True
+
+    def _choose_run_config(
+        self,
+        tuner_run_config: Optional[RunConfig],
+        trainer: "BaseTrainer",
+        param_space: Optional[Dict[str, Any]],
+    ) -> RunConfig:
+        """Chooses which `RunConfig` to use when multiple can be passed in
+        through a Trainer or the Tuner itself.
+
+        Args:
+            tuner_run_config: The run config passed into the Tuner constructor.
+            trainer: The Trainer instance to use with Tune, which may have
+                a RunConfig specified by the user.
+            param_space: The param space passed to the Tuner.
+
+        Raises:
+            ValueError: if the `run_config` is specified as a hyperparameter.
+        """
+        if param_space and "run_config" in param_space:
+            raise ValueError(
+                "`RunConfig` cannot be tuned as part of the `param_space`! "
+                "Move the run config to be a parameter of the `Tuner`: "
+                "Tuner(..., run_config=RunConfig(...))"
+            )
+
+        # Both Tuner RunConfig + Trainer RunConfig --> prefer Tuner RunConfig
+        if tuner_run_config and trainer.run_config != ray.train.RunConfig():
+            logger.info(
+                "A `RunConfig` was passed to both the `Tuner` and the "
+                f"`{trainer.__class__.__name__}`. The run config passed to "
+                "the `Tuner` is the one that will be used."
+            )
+            return tuner_run_config
+
+        # No Tuner RunConfig -> pass the Trainer config through
+        # This returns either a user-specified config, or the default RunConfig
+        # if nothing was provided to both the Trainer or Tuner.
+        if not tuner_run_config:
+            return trainer.run_config
+
+        # Tuner RunConfig + No Trainer RunConfig --> Use the Tuner config
+        return tuner_run_config
+
+    def _process_scaling_config(self) -> None:
+        """Converts ``self._param_space["scaling_config"]`` to a dict.
+
+        The dict is converted back to a dataclass by the Trainer, after the
+        Tune search specification is resolved.
+        """
+        # TODO: introduce `ray.tune.sample.TuneableDataclass` and allow Tune to
+        # natively resolve specs with dataclasses.
+        scaling_config = self._param_space.get("scaling_config")
+        if not isinstance(scaling_config, ScalingConfig):
+            return
+        self._param_space["scaling_config"] = scaling_config.__dict__.copy()
+
+    @property
+    def trainable(self) -> TrainableTypeOrTrainer:
+        return self._trainable
+
+    @property
+    def converted_trainable(self) -> TrainableType:
+        return self._converted_trainable
+
+    @trainable.setter
+    def trainable(self, trainable: TrainableTypeOrTrainer):
+        self._trainable = trainable
+        self._converted_trainable = self._convert_trainable(trainable)
+
+    @property
+    def param_space(self) -> Optional[Dict[str, Any]]:
+        return self._param_space
+
+    @param_space.setter
+    def param_space(self, param_space: Optional[Dict[str, Any]]):
+        # Handle any configs that adhere to the `to_dict` interface.
+        # Ex: AlgorithmConfig from RLlib
+        if isinstance(param_space, _Config):
+            param_space = param_space.to_dict()
+
+        if not isinstance(param_space, dict) and param_space is not None:
+            raise ValueError(
+                "The `param_space` passed to the `Tuner` must be a dict. "
+                f"Got '{type(param_space)}' instead."
+            )
+
+        self._param_space = param_space
+
+        if param_space:
+            self._process_scaling_config()
+
+    def _convert_trainable(self, trainable: TrainableTypeOrTrainer) -> TrainableType:
+        """Converts a Trainer to a Tune trainable and saves the converted
+        trainable. If not using a Trainer, this leaves the trainable as is."""
+        from ray.train.trainer import BaseTrainer
+
+        return (
+            trainable.as_trainable()
+            if isinstance(trainable, BaseTrainer)
+            else trainable
+        )
+
+    def fit(self) -> ResultGrid:
+        trainable = self.converted_trainable
+        param_space = copy.deepcopy(self.param_space)
+        if not self._is_restored:
+            analysis = self._fit_internal(trainable, param_space)
+        else:
+            analysis = self._fit_resume(trainable, param_space)
+
+        self._experiment_analysis = analysis
+
+        return ResultGrid(self._experiment_analysis)
+
+    def get_results(self) -> ResultGrid:
+        if not self._experiment_analysis:
+            raise RuntimeError(
+                "Can't return results as experiment has not been run, yet. "
+                "Call `Tuner.fit()` to run the experiment first."
+            )
+        return ResultGrid(self._experiment_analysis)
+
+    def _get_tune_run_arguments(self, trainable: TrainableType) -> Dict[str, Any]:
+        """Get tune.run arguments common for both new and resumed runs."""
+        # Avoid overwriting the originally configured checkpoint config.
+        checkpoint_config = copy.deepcopy(self._run_config.checkpoint_config)
+
+        if checkpoint_config.checkpoint_frequency:
+            # Function trainables (and thus most of our trainers) usually don't handle
+            # this argument.
+            handle_checkpoint_freq = getattr(
+                trainable, "_handles_checkpoint_freq", None
+            )
+            if handle_checkpoint_freq is False:
+                # If we specifically know this trainable doesn't support the
+                # argument, raise an error
+                raise ValueError(
+                    "You passed `checkpoint_frequency="
+                    f"{checkpoint_config.checkpoint_frequency}` to your "
+                    "CheckpointConfig, but this trainer does not support "
+                    "this argument. If you passed in a Trainer that takes in a "
+                    "custom training loop, you will need to "
+                    "report a checkpoint every `checkpoint_frequency` iterations "
+                    "within your training loop using "
+                    "`ray.train.report(metrics=..., checkpoint=...)` "
+                    "to get this behavior."
+                )
+            elif handle_checkpoint_freq is True:
+                # If we specifically support it, it's handled in the training loop,
+                # so we disable tune's bookkeeping.
+                checkpoint_config.checkpoint_frequency = 0
+            # Otherwise, the trainable is not a Trainer and we just keep the
+            # user-supplied value.
+            # Function trainables will raise a runtime error later if set > 0
+        if checkpoint_config.checkpoint_at_end is not None:
+            # Again, function trainables usually don't handle this argument.
+            handle_cp_at_end = getattr(trainable, "_handles_checkpoint_at_end", None)
+            if handle_cp_at_end is False:
+                # If we specifically know we don't support it, raise an error.
+                raise ValueError(
+                    "You passed `checkpoint_at_end="
+                    f"{checkpoint_config.checkpoint_at_end}` "
+                    "to your CheckpointConfig, but this trainer does not support "
+                    "this argument. If you passed in a Trainer that takes in a "
+                    "custom training loop, you should include one last call to "
+                    "`ray.train.report(metrics=..., checkpoint=...)` "
+                    "at the end of your training loop to get this behavior."
+                )
+            elif handle_cp_at_end is True:
+                # If we specifically support it, it's handled in the training loop,
+                # so we disable tune's internal bookkeeping.
+                checkpoint_config.checkpoint_at_end = False
+            # If this is a user-defined trainable, just keep the value
+            # Function trainables will raise a runtime error later if set to True
+        else:
+            # Set default to False for function trainables and True for everything else
+            if is_function_trainable(trainable):
+                checkpoint_config.checkpoint_at_end = False
+            else:
+                checkpoint_config.checkpoint_at_end = True
+
+        return dict(
+            storage_path=self._run_config.storage_path,
+            storage_filesystem=self._run_config.storage_filesystem,
+            name=self._run_config.name,
+            mode=self._tune_config.mode,
+            metric=self._tune_config.metric,
+            callbacks=self._run_config.callbacks,
+            sync_config=self._run_config.sync_config,
+            stop=self._run_config.stop,
+            max_failures=self._run_config.failure_config.max_failures,
+            checkpoint_config=checkpoint_config,
+            raise_on_failed_trial=False,
+            fail_fast=(self._run_config.failure_config.fail_fast),
+            progress_reporter=self._run_config.progress_reporter,
+            verbose=self._run_config.verbose,
+            reuse_actors=self._tune_config.reuse_actors,
+            max_concurrent_trials=self._tune_config.max_concurrent_trials,
+            time_budget_s=self._tune_config.time_budget_s,
+            trial_name_creator=self._tune_config.trial_name_creator,
+            trial_dirname_creator=self._tune_config.trial_dirname_creator,
+            _entrypoint=self._entrypoint,
+            # Deprecated
+            chdir_to_trial_dir=self._tune_config.chdir_to_trial_dir,
+        )
+
+    def _fit_internal(
+        self, trainable: TrainableType, param_space: Optional[Dict[str, Any]]
+    ) -> ExperimentAnalysis:
+        """Fitting for a fresh Tuner."""
+        args = {
+            **self._get_tune_run_arguments(trainable),
+            **dict(
+                run_or_experiment=trainable,
+                config=param_space,
+                num_samples=self._tune_config.num_samples,
+                search_alg=self._tune_config.search_alg,
+                scheduler=self._tune_config.scheduler,
+                log_to_file=self._run_config.log_to_file,
+            ),
+            **self._tuner_kwargs,
+        }
+        analysis = run(
+            **args,
+        )
+        self.clear_remote_string_queue()
+        return analysis
+
+    def _fit_resume(
+        self, trainable: TrainableType, param_space: Optional[Dict[str, Any]]
+    ) -> ExperimentAnalysis:
+        """Fitting for a restored Tuner."""
+        assert self._resume_config
+
+        args = {
+            **self._get_tune_run_arguments(trainable),
+            **dict(
+                run_or_experiment=trainable,
+                config=param_space,
+                resume_config=self._resume_config,
+                search_alg=self._tune_config.search_alg,
+                scheduler=self._tune_config.scheduler,
+            ),
+            **self._tuner_kwargs,
+        }
+        analysis = run(**args)
+        self.clear_remote_string_queue()
+        return analysis
+
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        state["_tuner_kwargs"] = state["_tuner_kwargs"].copy()
+        state["_tuner_kwargs"].pop("_remote_string_queue", None)
+        state.pop(_TRAINABLE_KEY, None)
+        trainable = state.pop(_CONVERTED_TRAINABLE_KEY, None)
+        param_space = state.pop(_PARAM_SPACE_KEY, None)
+        state.pop(_EXPERIMENT_ANALYSIS_KEY, None)
+
+        state["__trainable_name"] = (
+            Experiment.get_trainable_name(trainable) if trainable else None
+        )
+        state["__flattened_param_space_keys"] = (
+            sorted(flatten_dict(param_space).keys())
+            if param_space is not None
+            else None
+        )
+
+        return state
+
+    def __setstate__(self, state):
+        # Make sure the magic metadata gets removed first.
+        state.pop("__flattened_param_space_keys", None)
+        state.pop("__trainable_name", None)
+
+        self.__dict__.update(state)

.venv/lib/python3.11/site-packages/ray/tune/search/__init__.py

@@ -0,0 +1,153 @@
+from ray._private.utils import get_function_args
+from ray.tune.search.basic_variant import BasicVariantGenerator
+from ray.tune.search.concurrency_limiter import ConcurrencyLimiter
+from ray.tune.search.repeater import Repeater
+from ray.tune.search.search_algorithm import SearchAlgorithm
+from ray.tune.search.search_generator import SearchGenerator
+from ray.tune.search.searcher import Searcher
+from ray.tune.search.variant_generator import grid_search
+from ray.util import PublicAPI
+
+
+def _import_variant_generator():
+    return BasicVariantGenerator
+
+
+def _import_ax_search():
+    from ray.tune.search.ax.ax_search import AxSearch
+
+    return AxSearch
+
+
+def _import_hyperopt_search():
+    from ray.tune.search.hyperopt.hyperopt_search import HyperOptSearch
+
+    return HyperOptSearch
+
+
+def _import_bayesopt_search():
+    from ray.tune.search.bayesopt.bayesopt_search import BayesOptSearch
+
+    return BayesOptSearch
+
+
+def _import_bohb_search():
+    from ray.tune.search.bohb.bohb_search import TuneBOHB
+
+    return TuneBOHB
+
+
+def _import_nevergrad_search():
+    from ray.tune.search.nevergrad.nevergrad_search import NevergradSearch
+
+    return NevergradSearch
+
+
+def _import_optuna_search():
+    from ray.tune.search.optuna.optuna_search import OptunaSearch
+
+    return OptunaSearch
+
+
+def _import_zoopt_search():
+    from ray.tune.search.zoopt.zoopt_search import ZOOptSearch
+
+    return ZOOptSearch
+
+
+def _import_hebo_search():
+    from ray.tune.search.hebo.hebo_search import HEBOSearch
+
+    return HEBOSearch
+
+
+SEARCH_ALG_IMPORT = {
+    "variant_generator": _import_variant_generator,
+    "random": _import_variant_generator,
+    "ax": _import_ax_search,
+    "hyperopt": _import_hyperopt_search,
+    "bayesopt": _import_bayesopt_search,
+    "bohb": _import_bohb_search,
+    "nevergrad": _import_nevergrad_search,
+    "optuna": _import_optuna_search,
+    "zoopt": _import_zoopt_search,
+    "hebo": _import_hebo_search,
+}
+
+
+@PublicAPI(stability="beta")
+def create_searcher(
+    search_alg,
+    **kwargs,
+):
+    """Instantiate a search algorithm based on the given string.
+
+    This is useful for swapping between different search algorithms.
+
+    Args:
+        search_alg: The search algorithm to use.
+        metric: The training result objective value attribute. Stopping
+            procedures will use this attribute.
+        mode: One of {min, max}. Determines whether objective is
+            minimizing or maximizing the metric attribute.
+        **kwargs: Additional parameters.
+            These keyword arguments will be passed to the initialization
+            function of the chosen class.
+    Returns:
+        ray.tune.search.Searcher: The search algorithm.
+    Example:
+        >>> from ray import tune # doctest: +SKIP
+        >>> search_alg = tune.create_searcher('ax') # doctest: +SKIP
+    """
+
+    search_alg = search_alg.lower()
+    if search_alg not in SEARCH_ALG_IMPORT:
+        raise ValueError(
+            f"The `search_alg` argument must be one of "
+            f"{list(SEARCH_ALG_IMPORT)}. "
+            f"Got: {search_alg}"
+        )
+
+    SearcherClass = SEARCH_ALG_IMPORT[search_alg]()
+
+    search_alg_args = get_function_args(SearcherClass)
+    trimmed_kwargs = {k: v for k, v in kwargs.items() if k in search_alg_args}
+
+    return SearcherClass(**trimmed_kwargs)
+
+
+UNRESOLVED_SEARCH_SPACE = str(
+    "You passed a `{par}` parameter to {cls} that contained unresolved search "
+    "space definitions. {cls} should however be instantiated with fully "
+    "configured search spaces only. To use Ray Tune's automatic search space "
+    "conversion, pass the space definition as part of the `param_space` argument "
+    "to `tune.Tuner()` instead."
+)
+
+UNDEFINED_SEARCH_SPACE = str(
+    "Trying to sample a configuration from {cls}, but no search "
+    "space has been defined. Either pass the `{space}` argument when "
+    "instantiating the search algorithm, or pass a `param_space` to "
+    "`tune.Tuner()`."
+)
+
+UNDEFINED_METRIC_MODE = str(
+    "Trying to sample a configuration from {cls}, but the `metric` "
+    "({metric}) or `mode` ({mode}) parameters have not been set. "
+    "Either pass these arguments when instantiating the search algorithm, "
+    "or pass them to `tune.TuneConfig()`."
+)
+
+
+__all__ = [
+    "SearchAlgorithm",
+    "Searcher",
+    "ConcurrencyLimiter",
+    "Repeater",
+    "BasicVariantGenerator",
+    "grid_search",
+    "SearchGenerator",
+    "UNRESOLVED_SEARCH_SPACE",
+    "UNDEFINED_SEARCH_SPACE",
+    "UNDEFINED_METRIC_MODE",
+]

.venv/lib/python3.11/site-packages/ray/tune/search/_mock.py

@@ -0,0 +1,55 @@
+from typing import Dict, List, Optional
+
+from ray.tune.experiment import Trial
+from ray.tune.search import ConcurrencyLimiter, Searcher
+from ray.tune.search.search_generator import SearchGenerator
+
+
+class _MockSearcher(Searcher):
+    def __init__(self, **kwargs):
+        self.live_trials = {}
+        self.counter = {"result": 0, "complete": 0}
+        self.final_results = []
+        self.stall = False
+        self.results = []
+        super(_MockSearcher, self).__init__(**kwargs)
+
+    def suggest(self, trial_id: str):
+        if not self.stall:
+            self.live_trials[trial_id] = 1
+            return {"test_variable": 2}
+        return None
+
+    def on_trial_result(self, trial_id: str, result: Dict):
+        self.counter["result"] += 1
+        self.results += [result]
+
+    def on_trial_complete(
+        self, trial_id: str, result: Optional[Dict] = None, error: bool = False
+    ):
+        self.counter["complete"] += 1
+        if result:
+            self._process_result(result)
+        if trial_id in self.live_trials:
+            del self.live_trials[trial_id]
+
+    def _process_result(self, result: Dict):
+        self.final_results += [result]
+
+
+class _MockSuggestionAlgorithm(SearchGenerator):
+    def __init__(self, max_concurrent: Optional[int] = None, **kwargs):
+        self.searcher = _MockSearcher(**kwargs)
+        if max_concurrent:
+            self.searcher = ConcurrencyLimiter(
+                self.searcher, max_concurrent=max_concurrent
+            )
+        super(_MockSuggestionAlgorithm, self).__init__(self.searcher)
+
+    @property
+    def live_trials(self) -> List[Trial]:
+        return self.searcher.live_trials
+
+    @property
+    def results(self) -> List[Dict]:
+        return self.searcher.results

.venv/lib/python3.11/site-packages/ray/tune/search/basic_variant.py

@@ -0,0 +1,421 @@
+import copy
+import itertools
+import os
+import uuid
+import warnings
+from pathlib import Path
+from typing import TYPE_CHECKING, Dict, List, Optional, Union
+
+import numpy as np
+
+from ray.air._internal.usage import tag_searcher
+from ray.tune.error import TuneError
+from ray.tune.experiment.config_parser import _create_trial_from_spec, _make_parser
+from ray.tune.search.sample import _BackwardsCompatibleNumpyRng, np_random_generator
+from ray.tune.search.search_algorithm import SearchAlgorithm
+from ray.tune.search.variant_generator import (
+    _count_spec_samples,
+    _count_variants,
+    _flatten_resolved_vars,
+    _get_preset_variants,
+    format_vars,
+    generate_variants,
+)
+from ray.tune.utils.util import _atomic_save, _load_newest_checkpoint
+from ray.util import PublicAPI
+
+if TYPE_CHECKING:
+    from ray.tune.experiment import Experiment
+
+SERIALIZATION_THRESHOLD = 1e6
+
+
+class _VariantIterator:
+    """Iterates over generated variants from the search space.
+
+    This object also toggles between lazy evaluation and
+    eager evaluation of samples. If lazy evaluation is enabled,
+    this object cannot be serialized.
+    """
+
+    def __init__(self, iterable, lazy_eval=False):
+        self.lazy_eval = lazy_eval
+        self.iterable = iterable
+        self._has_next = True
+        if lazy_eval:
+            self._load_value()
+        else:
+            self.iterable = list(iterable)
+            self._has_next = bool(self.iterable)
+
+    def _load_value(self):
+        try:
+            self.next_value = next(self.iterable)
+        except StopIteration:
+            self._has_next = False
+
+    def has_next(self):
+        return self._has_next
+
+    def __next__(self):
+        if self.lazy_eval:
+            current_value = self.next_value
+            self._load_value()
+            return current_value
+        current_value = self.iterable.pop(0)
+        self._has_next = bool(self.iterable)
+        return current_value
+
+
+class _TrialIterator:
+    """Generates trials from the spec.
+
+    Args:
+        uuid_prefix: Used in creating the trial name.
+        num_samples: Number of samples from distribution
+             (same as tune.TuneConfig).
+        unresolved_spec: Experiment specification
+            that might have unresolved distributions.
+        constant_grid_search: Should random variables be sampled
+            first before iterating over grid variants (True) or not (False).
+        points_to_evaluate: Configurations that will be tried out without sampling.
+        lazy_eval: Whether variants should be generated
+            lazily or eagerly. This is toggled depending
+            on the size of the grid search.
+        start: index at which to start counting trials.
+        random_state (int | np.random.Generator | np.random.RandomState):
+            Seed or numpy random generator to use for reproducible results.
+            If None (default), will use the global numpy random generator
+            (``np.random``). Please note that full reproducibility cannot
+            be guaranteed in a distributed enviroment.
+    """
+
+    def __init__(
+        self,
+        uuid_prefix: str,
+        num_samples: int,
+        unresolved_spec: dict,
+        constant_grid_search: bool = False,
+        points_to_evaluate: Optional[List] = None,
+        lazy_eval: bool = False,
+        start: int = 0,
+        random_state: Optional[
+            Union[int, "np_random_generator", np.random.RandomState]
+        ] = None,
+    ):
+        self.parser = _make_parser()
+        self.num_samples = num_samples
+        self.uuid_prefix = uuid_prefix
+        self.num_samples_left = num_samples
+        self.unresolved_spec = unresolved_spec
+        self.constant_grid_search = constant_grid_search
+        self.points_to_evaluate = points_to_evaluate or []
+        self.num_points_to_evaluate = len(self.points_to_evaluate)
+        self.counter = start
+        self.lazy_eval = lazy_eval
+        self.variants = None
+        self.random_state = random_state
+
+    def create_trial(self, resolved_vars, spec):
+        trial_id = self.uuid_prefix + ("%05d" % self.counter)
+        experiment_tag = str(self.counter)
+        # Always append resolved vars to experiment tag?
+        if resolved_vars:
+            experiment_tag += "_{}".format(format_vars(resolved_vars))
+        self.counter += 1
+        return _create_trial_from_spec(
+            spec,
+            self.parser,
+            evaluated_params=_flatten_resolved_vars(resolved_vars),
+            trial_id=trial_id,
+            experiment_tag=experiment_tag,
+        )
+
+    def __next__(self):
+        """Generates Trial objects with the variant generation process.
+
+        Uses a fixed point iteration to resolve variants. All trials
+        should be able to be generated at once.
+
+        See also: `ray.tune.search.variant_generator`.
+
+        Returns:
+            Trial object
+        """
+
+        if "run" not in self.unresolved_spec:
+            raise TuneError("Must specify `run` in {}".format(self.unresolved_spec))
+
+        if self.variants and self.variants.has_next():
+            # This block will be skipped upon instantiation.
+            # `variants` will be set later after the first loop.
+            resolved_vars, spec = next(self.variants)
+            return self.create_trial(resolved_vars, spec)
+
+        if self.points_to_evaluate:
+            config = self.points_to_evaluate.pop(0)
+            self.num_samples_left -= 1
+            self.variants = _VariantIterator(
+                _get_preset_variants(
+                    self.unresolved_spec,
+                    config,
+                    constant_grid_search=self.constant_grid_search,
+                    random_state=self.random_state,
+                ),
+                lazy_eval=self.lazy_eval,
+            )
+            resolved_vars, spec = next(self.variants)
+            return self.create_trial(resolved_vars, spec)
+        elif self.num_samples_left > 0:
+            self.variants = _VariantIterator(
+                generate_variants(
+                    self.unresolved_spec,
+                    constant_grid_search=self.constant_grid_search,
+                    random_state=self.random_state,
+                ),
+                lazy_eval=self.lazy_eval,
+            )
+            self.num_samples_left -= 1
+            resolved_vars, spec = next(self.variants)
+            return self.create_trial(resolved_vars, spec)
+        else:
+            raise StopIteration
+
+    def __iter__(self):
+        return self
+
+
+@PublicAPI
+class BasicVariantGenerator(SearchAlgorithm):
+    """Uses Tune's variant generation for resolving variables.
+
+    This is the default search algorithm used if no other search algorithm
+    is specified.
+
+
+    Args:
+        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.
+        max_concurrent: Maximum number of concurrently running trials.
+            If 0 (default), no maximum is enforced.
+        constant_grid_search: If this is set to ``True``, Ray Tune will
+            *first* try to sample random values and keep them constant over
+            grid search parameters. If this is set to ``False`` (default),
+            Ray Tune will sample new random parameters in each grid search
+            condition.
+        random_state:
+            Seed or numpy random generator to use for reproducible results.
+            If None (default), will use the global numpy random generator
+            (``np.random``). Please note that full reproducibility cannot
+            be guaranteed in a distributed environment.
+
+
+    Example:
+
+    .. code-block:: python
+
+        from ray import tune
+
+        # This will automatically use the `BasicVariantGenerator`
+        tuner = tune.Tuner(
+            lambda config: config["a"] + config["b"],
+            tune_config=tune.TuneConfig(
+                num_samples=4
+            ),
+            param_space={
+                "a": tune.grid_search([1, 2]),
+                "b": tune.randint(0, 3)
+            },
+        )
+        tuner.fit()
+
+    In the example above, 8 trials will be generated: For each sample
+    (``4``), each of the grid search variants for ``a`` will be sampled
+    once. The ``b`` parameter will be sampled randomly.
+
+    The generator accepts a pre-set list of points that should be evaluated.
+    The points will replace the first samples of each experiment passed to
+    the ``BasicVariantGenerator``.
+
+    Each point will replace one sample of the specified ``num_samples``. If
+    grid search variables are overwritten with the values specified in the
+    presets, the number of samples will thus be reduced.
+
+    Example:
+
+    .. code-block:: python
+
+        from ray import tune
+        from ray.tune.search.basic_variant import BasicVariantGenerator
+
+        tuner = tune.Tuner(
+            lambda config: config["a"] + config["b"],
+            tune_config=tune.TuneConfig(
+                search_alg=BasicVariantGenerator(points_to_evaluate=[
+                    {"a": 2, "b": 2},
+                    {"a": 1},
+                    {"b": 2}
+                ]),
+                num_samples=4
+            ),
+            param_space={
+                "a": tune.grid_search([1, 2]),
+                "b": tune.randint(0, 3)
+            },
+        )
+        tuner.fit()
+
+    The example above will produce six trials via four samples:
+
+    - The first sample will produce one trial with ``a=2`` and ``b=2``.
+    - The second sample will produce one trial with ``a=1`` and ``b`` sampled
+      randomly
+    - The third sample will produce two trials, one for each grid search
+      value of ``a``. It will be ``b=2`` for both of these trials.
+    - The fourth sample will produce two trials, one for each grid search
+      value of ``a``. ``b`` will be sampled randomly and independently for
+      both of these trials.
+
+    """
+
+    CKPT_FILE_TMPL = "basic-variant-state-{}.json"
+
+    def __init__(
+        self,
+        points_to_evaluate: Optional[List[Dict]] = None,
+        max_concurrent: int = 0,
+        constant_grid_search: bool = False,
+        random_state: Optional[
+            Union[int, "np_random_generator", np.random.RandomState]
+        ] = None,
+    ):
+        tag_searcher(self)
+        self._trial_generator = []
+        self._iterators = []
+        self._trial_iter = None
+        self._finished = False
+        self._random_state = _BackwardsCompatibleNumpyRng(random_state)
+
+        self._points_to_evaluate = points_to_evaluate or []
+
+        # Unique prefix for all trials generated, e.g., trial ids start as
+        # 2f1e_00001, 2f1ef_00002, 2f1ef_0003, etc. Overridable for testing.
+        force_test_uuid = os.environ.get("_TEST_TUNE_TRIAL_UUID")
+        if force_test_uuid:
+            self._uuid_prefix = force_test_uuid + "_"
+        else:
+            self._uuid_prefix = str(uuid.uuid1().hex)[:5] + "_"
+
+        self._total_samples = 0
+        self.max_concurrent = max_concurrent
+        self._constant_grid_search = constant_grid_search
+        self._live_trials = set()
+
+    @property
+    def total_samples(self):
+        return self._total_samples
+
+    def add_configurations(
+        self, experiments: Union["Experiment", List["Experiment"], Dict[str, Dict]]
+    ):
+        """Chains generator given experiment specifications.
+
+        Arguments:
+            experiments: Experiments to run.
+        """
+        from ray.tune.experiment import _convert_to_experiment_list
+
+        experiment_list = _convert_to_experiment_list(experiments)
+
+        for experiment in experiment_list:
+            grid_vals = _count_spec_samples(experiment.spec, num_samples=1)
+            lazy_eval = grid_vals > SERIALIZATION_THRESHOLD
+            if lazy_eval:
+                warnings.warn(
+                    f"The number of pre-generated samples ({grid_vals}) "
+                    "exceeds the serialization threshold "
+                    f"({int(SERIALIZATION_THRESHOLD)}). Resume ability is "
+                    "disabled. To fix this, reduce the number of "
+                    "dimensions/size of the provided grid search."
+                )
+
+            previous_samples = self._total_samples
+            points_to_evaluate = copy.deepcopy(self._points_to_evaluate)
+            self._total_samples += _count_variants(experiment.spec, points_to_evaluate)
+            iterator = _TrialIterator(
+                uuid_prefix=self._uuid_prefix,
+                num_samples=experiment.spec.get("num_samples", 1),
+                unresolved_spec=experiment.spec,
+                constant_grid_search=self._constant_grid_search,
+                points_to_evaluate=points_to_evaluate,
+                lazy_eval=lazy_eval,
+                start=previous_samples,
+                random_state=self._random_state,
+            )
+            self._iterators.append(iterator)
+            self._trial_generator = itertools.chain(self._trial_generator, iterator)
+
+    def next_trial(self):
+        """Provides one Trial object to be queued into the TrialRunner.
+
+        Returns:
+            Trial: Returns a single trial.
+        """
+        if self.is_finished():
+            return None
+        if self.max_concurrent > 0 and len(self._live_trials) >= self.max_concurrent:
+            return None
+        if not self._trial_iter:
+            self._trial_iter = iter(self._trial_generator)
+        try:
+            trial = next(self._trial_iter)
+            self._live_trials.add(trial.trial_id)
+            return trial
+        except StopIteration:
+            self._trial_generator = []
+            self._trial_iter = None
+            self.set_finished()
+            return None
+
+    def on_trial_complete(
+        self, trial_id: str, result: Optional[Dict] = None, error: bool = False
+    ):
+        if trial_id in self._live_trials:
+            self._live_trials.remove(trial_id)
+
+    def get_state(self):
+        if any(iterator.lazy_eval for iterator in self._iterators):
+            return False
+        state = self.__dict__.copy()
+        del state["_trial_generator"]
+        return state
+
+    def set_state(self, state):
+        self.__dict__.update(state)
+        for iterator in self._iterators:
+            self._trial_generator = itertools.chain(self._trial_generator, iterator)
+
+    def save_to_dir(self, dirpath, session_str):
+        if any(iterator.lazy_eval for iterator in self._iterators):
+            return False
+        state_dict = self.get_state()
+        _atomic_save(
+            state=state_dict,
+            checkpoint_dir=dirpath,
+            file_name=self.CKPT_FILE_TMPL.format(session_str),
+            tmp_file_name=".tmp_generator",
+        )
+
+    def has_checkpoint(self, dirpath: str):
+        """Whether a checkpoint file exists within dirpath."""
+        return any(Path(dirpath).glob(self.CKPT_FILE_TMPL.format("*")))
+
+    def restore_from_dir(self, dirpath: str):
+        """Restores self + searcher + search wrappers from dirpath."""
+        state_dict = _load_newest_checkpoint(dirpath, self.CKPT_FILE_TMPL.format("*"))
+        if not state_dict:
+            raise RuntimeError("Unable to find checkpoint in {}.".format(dirpath))
+        self.set_state(state_dict)

.venv/lib/python3.11/site-packages/ray/tune/search/concurrency_limiter.py

@@ -0,0 +1,176 @@
+import copy
+import logging
+from typing import Dict, List, Optional
+
+from ray.tune.search.searcher import Searcher
+from ray.tune.search.util import _set_search_properties_backwards_compatible
+from ray.util.annotations import PublicAPI
+
+logger = logging.getLogger(__name__)
+
+
+@PublicAPI
+class ConcurrencyLimiter(Searcher):
+    """A wrapper algorithm for limiting the number of concurrent trials.
+
+    Certain Searchers have their own internal logic for limiting
+    the number of concurrent trials. If such a Searcher is passed to a
+    ``ConcurrencyLimiter``, the ``max_concurrent`` of the
+    ``ConcurrencyLimiter`` will override the ``max_concurrent`` value
+    of the Searcher. The ``ConcurrencyLimiter`` will then let the
+    Searcher's internal logic take over.
+
+    Args:
+        searcher: Searcher object that the
+            ConcurrencyLimiter will manage.
+        max_concurrent: Maximum concurrent samples from the underlying
+            searcher.
+        batch: Whether to wait for all concurrent samples
+            to finish before updating the underlying searcher.
+
+    Example:
+
+    .. code-block:: python
+
+        from ray.tune.search import ConcurrencyLimiter
+        search_alg = HyperOptSearch(metric="accuracy")
+        search_alg = ConcurrencyLimiter(search_alg, max_concurrent=2)
+        tuner = tune.Tuner(
+            trainable,
+            tune_config=tune.TuneConfig(
+                search_alg=search_alg
+            ),
+        )
+        tuner.fit()
+    """
+
+    def __init__(self, searcher: Searcher, max_concurrent: int, batch: bool = False):
+        assert type(max_concurrent) is int and max_concurrent > 0
+        self.searcher = searcher
+        self.max_concurrent = max_concurrent
+        self.batch = batch
+        self.live_trials = set()
+        self.num_unfinished_live_trials = 0
+        self.cached_results = {}
+        self._limit_concurrency = True
+
+        if not isinstance(searcher, Searcher):
+            raise RuntimeError(
+                f"The `ConcurrencyLimiter` only works with `Searcher` "
+                f"objects (got {type(searcher)}). Please try to pass "
+                f"`max_concurrent` to the search generator directly."
+            )
+
+        self._set_searcher_max_concurrency()
+
+        super(ConcurrencyLimiter, self).__init__(
+            metric=self.searcher.metric, mode=self.searcher.mode
+        )
+
+    def _set_searcher_max_concurrency(self):
+        # If the searcher has special logic for handling max concurrency,
+        # we do not do anything inside the ConcurrencyLimiter
+        self._limit_concurrency = not self.searcher.set_max_concurrency(
+            self.max_concurrent
+        )
+
+    def set_max_concurrency(self, max_concurrent: int) -> bool:
+        # Determine if this behavior is acceptable, or if it should
+        # raise an exception.
+        self.max_concurrent = max_concurrent
+        return True
+
+    def set_search_properties(
+        self, metric: Optional[str], mode: Optional[str], config: Dict, **spec
+    ) -> bool:
+        self._set_searcher_max_concurrency()
+        return _set_search_properties_backwards_compatible(
+            self.searcher.set_search_properties, metric, mode, config, **spec
+        )
+
+    def suggest(self, trial_id: str) -> Optional[Dict]:
+        if not self._limit_concurrency:
+            return self.searcher.suggest(trial_id)
+
+        assert (
+            trial_id not in self.live_trials
+        ), f"Trial ID {trial_id} must be unique: already found in set."
+        if len(self.live_trials) >= self.max_concurrent:
+            logger.debug(
+                f"Not providing a suggestion for {trial_id} due to "
+                "concurrency limit: %s/%s.",
+                len(self.live_trials),
+                self.max_concurrent,
+            )
+            return
+
+        suggestion = self.searcher.suggest(trial_id)
+        if suggestion not in (None, Searcher.FINISHED):
+            self.live_trials.add(trial_id)
+            self.num_unfinished_live_trials += 1
+        return suggestion
+
+    def on_trial_complete(
+        self, trial_id: str, result: Optional[Dict] = None, error: bool = False
+    ):
+        if not self._limit_concurrency:
+            return self.searcher.on_trial_complete(trial_id, result=result, error=error)
+
+        if trial_id not in self.live_trials:
+            return
+        elif self.batch:
+            self.cached_results[trial_id] = (result, error)
+            self.num_unfinished_live_trials -= 1
+            if self.num_unfinished_live_trials <= 0:
+                # Update the underlying searcher once the
+                # full batch is completed.
+                for trial_id, (result, error) in self.cached_results.items():
+                    self.searcher.on_trial_complete(
+                        trial_id, result=result, error=error
+                    )
+                    self.live_trials.remove(trial_id)
+                self.cached_results = {}
+                self.num_unfinished_live_trials = 0
+            else:
+                return
+        else:
+            self.searcher.on_trial_complete(trial_id, result=result, error=error)
+            self.live_trials.remove(trial_id)
+            self.num_unfinished_live_trials -= 1
+
+    def on_trial_result(self, trial_id: str, result: Dict) -> None:
+        self.searcher.on_trial_result(trial_id, result)
+
+    def add_evaluated_point(
+        self,
+        parameters: Dict,
+        value: float,
+        error: bool = False,
+        pruned: bool = False,
+        intermediate_values: Optional[List[float]] = None,
+    ):
+        return self.searcher.add_evaluated_point(
+            parameters, value, error, pruned, intermediate_values
+        )
+
+    def get_state(self) -> Dict:
+        state = self.__dict__.copy()
+        del state["searcher"]
+        return copy.deepcopy(state)
+
+    def set_state(self, state: Dict):
+        self.__dict__.update(state)
+
+    def save(self, checkpoint_path: str):
+        self.searcher.save(checkpoint_path)
+
+    def restore(self, checkpoint_path: str):
+        self.searcher.restore(checkpoint_path)
+
+    # BOHB Specific.
+    # TODO(team-ml): Refactor alongside HyperBandForBOHB
+    def on_pause(self, trial_id: str):
+        self.searcher.on_pause(trial_id)
+
+    def on_unpause(self, trial_id: str):
+        self.searcher.on_unpause(trial_id)

.venv/lib/python3.11/site-packages/ray/tune/search/hebo/__init__.py

@@ -0,0 +1,3 @@
+from ray.tune.search.hebo.hebo_search import HEBOSearch
+
+__all__ = ["HEBOSearch"]

.venv/lib/python3.11/site-packages/ray/tune/search/nevergrad/__init__.py

@@ -0,0 +1,3 @@
+from ray.tune.search.nevergrad.nevergrad_search import NevergradSearch
+
+__all__ = ["NevergradSearch"]

.venv/lib/python3.11/site-packages/ray/tune/search/nevergrad/nevergrad_search.py

@@ -0,0 +1,373 @@
+import inspect
+import logging
+import pickle
+from typing import Dict, List, Optional, Sequence, Type, Union
+
+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,
+    Domain,
+    Float,
+    Integer,
+    LogUniform,
+    Quantized,
+)
+from ray.tune.search.variant_generator import parse_spec_vars
+from ray.tune.utils.util import flatten_dict, unflatten_dict
+
+try:
+    import nevergrad as ng
+    from nevergrad.optimization import Optimizer
+    from nevergrad.optimization.base import ConfiguredOptimizer
+
+    Parameter = ng.p.Parameter
+except ImportError:
+    ng = None
+    Optimizer = None
+    ConfiguredOptimizer = None
+    Parameter = None
+
+logger = logging.getLogger(__name__)
+
+
+class NevergradSearch(Searcher):
+    """Uses Nevergrad to optimize hyperparameters.
+
+    Nevergrad is an open source tool from Facebook for derivative free
+    optimization.  More info can be found at:
+    https://github.com/facebookresearch/nevergrad.
+
+    You will need to install Nevergrad via the following command:
+
+    .. code-block:: bash
+
+        $ pip install nevergrad
+
+    Parameters:
+        optimizer: Optimizer class provided from Nevergrad.
+            See here for available optimizers:
+            https://facebookresearch.github.io/nevergrad/optimizers_ref.html#optimizers
+            This can also be an instance of a `ConfiguredOptimizer`. See the
+            section on configured optimizers in the above link.
+        optimizer_kwargs: Kwargs passed in when instantiating the `optimizer`
+        space: Nevergrad parametrization
+            to be passed to optimizer on instantiation, or list of parameter
+            names if you passed an optimizer object.
+        metric: The training result objective value attribute. If None
+            but a mode was passed, the anonymous metric `_metric` will be used
+            per default.
+        mode: One of {min, max}. Determines whether objective is
+            minimizing or maximizing the metric attribute.
+        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.
+
+    Tune automatically converts search spaces to Nevergrad's format:
+
+    .. code-block:: python
+
+        import nevergrad as ng
+
+        config = {
+            "width": tune.uniform(0, 20),
+            "height": tune.uniform(-100, 100),
+            "activation": tune.choice(["relu", "tanh"])
+        }
+
+        current_best_params = [{
+            "width": 10,
+            "height": 0,
+            "activation": relu",
+        }]
+
+        ng_search = NevergradSearch(
+            optimizer=ng.optimizers.OnePlusOne,
+            metric="mean_loss",
+            mode="min",
+            points_to_evaluate=current_best_params)
+
+        run(my_trainable, config=config, search_alg=ng_search)
+
+    If you would like to pass the search space manually, the code would
+    look like this:
+
+    .. code-block:: python
+
+        import nevergrad as ng
+
+        space = ng.p.Dict(
+            width=ng.p.Scalar(lower=0, upper=20),
+            height=ng.p.Scalar(lower=-100, upper=100),
+            activation=ng.p.Choice(choices=["relu", "tanh"])
+        )
+
+        ng_search = NevergradSearch(
+            optimizer=ng.optimizers.OnePlusOne,
+            space=space,
+            metric="mean_loss",
+            mode="min")
+
+        run(my_trainable, search_alg=ng_search)
+
+    """
+
+    def __init__(
+        self,
+        optimizer: Optional[
+            Union[Optimizer, Type[Optimizer], ConfiguredOptimizer]
+        ] = None,
+        optimizer_kwargs: Optional[Dict] = None,
+        space: Optional[Union[Dict, Parameter]] = None,
+        metric: Optional[str] = None,
+        mode: Optional[str] = None,
+        points_to_evaluate: Optional[List[Dict]] = None,
+    ):
+        assert (
+            ng is not None
+        ), """Nevergrad must be installed!
+            You can install Nevergrad with the command:
+            `pip install nevergrad`."""
+        if mode:
+            assert mode in ["min", "max"], "`mode` must be 'min' or 'max'."
+
+        super(NevergradSearch, self).__init__(metric=metric, mode=mode)
+
+        self._space = None
+        self._opt_factory = None
+        self._nevergrad_opt = None
+        self._optimizer_kwargs = optimizer_kwargs or {}
+
+        if points_to_evaluate is None:
+            self._points_to_evaluate = None
+        elif not isinstance(points_to_evaluate, Sequence):
+            raise ValueError(
+                "Invalid object type passed for `points_to_evaluate`: "
+                f"{type(points_to_evaluate)}. "
+                "Please pass a list of points (dictionaries) instead."
+            )
+        else:
+            self._points_to_evaluate = list(points_to_evaluate)
+
+        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)
+
+        if isinstance(optimizer, Optimizer):
+            if space is not None and not isinstance(space, list):
+                raise ValueError(
+                    "If you pass a configured optimizer to Nevergrad, either "
+                    "pass a list of parameter names or None as the `space` "
+                    "parameter."
+                )
+            if self._optimizer_kwargs:
+                raise ValueError(
+                    "If you pass in optimizer kwargs, either pass "
+                    "an `Optimizer` subclass or an instance of "
+                    "`ConfiguredOptimizer`."
+                )
+
+            self._parameters = space
+            self._nevergrad_opt = optimizer
+        elif (
+            inspect.isclass(optimizer) and issubclass(optimizer, Optimizer)
+        ) or isinstance(optimizer, ConfiguredOptimizer):
+            self._opt_factory = optimizer
+            self._parameters = None
+            self._space = space
+        else:
+            raise ValueError(
+                "The `optimizer` argument passed to NevergradSearch must be "
+                "either an `Optimizer` or a `ConfiguredOptimizer`."
+            )
+
+        self._live_trial_mapping = {}
+
+        if self._nevergrad_opt or self._space:
+            self._setup_nevergrad()
+
+    def _setup_nevergrad(self):
+        if self._opt_factory:
+            self._nevergrad_opt = self._opt_factory(
+                self._space, **self._optimizer_kwargs
+            )
+
+        # nevergrad.tell internally minimizes, so "max" => -1
+        if self._mode == "max":
+            self._metric_op = -1.0
+        elif self._mode == "min":
+            self._metric_op = 1.0
+
+        if self._metric is None and self._mode:
+            # If only a mode was passed, use anonymous metric
+            self._metric = DEFAULT_METRIC
+
+        if hasattr(self._nevergrad_opt, "instrumentation"):  # added in v0.2.0
+            if self._nevergrad_opt.instrumentation.kwargs:
+                if self._nevergrad_opt.instrumentation.args:
+                    raise ValueError("Instrumented optimizers should use kwargs only")
+                if self._parameters is not None:
+                    raise ValueError(
+                        "Instrumented optimizers should provide "
+                        "None as parameter_names"
+                    )
+            else:
+                if self._parameters is None:
+                    raise ValueError(
+                        "Non-instrumented optimizers should have "
+                        "a list of parameter_names"
+                    )
+                if len(self._nevergrad_opt.instrumentation.args) != 1:
+                    raise ValueError("Instrumented optimizers should use kwargs only")
+        if self._parameters is not None and self._nevergrad_opt.dimension != len(
+            self._parameters
+        ):
+            raise ValueError(
+                "len(parameters_names) must match optimizer "
+                "dimension for non-instrumented optimizers"
+            )
+
+        if self._points_to_evaluate:
+            # Nevergrad is LIFO, so we add the points to evaluate in reverse
+            # order.
+            for i in range(len(self._points_to_evaluate) - 1, -1, -1):
+                self._nevergrad_opt.suggest(self._points_to_evaluate[i])
+
+    def set_search_properties(
+        self, metric: Optional[str], mode: Optional[str], config: Dict, **spec
+    ) -> bool:
+        if self._nevergrad_opt or self._space:
+            return False
+        space = self.convert_search_space(config)
+        self._space = space
+
+        if metric:
+            self._metric = metric
+        if mode:
+            self._mode = mode
+
+        self._setup_nevergrad()
+        return True
+
+    def suggest(self, trial_id: str) -> Optional[Dict]:
+        if not self._nevergrad_opt:
+            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
+                )
+            )
+
+        suggested_config = self._nevergrad_opt.ask()
+
+        self._live_trial_mapping[trial_id] = suggested_config
+        # in v0.2.0+, output of ask() is a Candidate,
+        # with fields args and kwargs
+        if not suggested_config.kwargs:
+            if self._parameters:
+                return unflatten_dict(
+                    dict(zip(self._parameters, suggested_config.args[0]))
+                )
+            return unflatten_dict(suggested_config.value)
+        else:
+            return unflatten_dict(suggested_config.kwargs)
+
+    def on_trial_complete(
+        self, trial_id: str, result: Optional[Dict] = None, error: bool = False
+    ):
+        """Notification for the completion of trial.
+
+        The result is internally negated when interacting with Nevergrad
+        so that Nevergrad Optimizers can "maximize" this value,
+        as it minimizes on default.
+        """
+        if result:
+            self._process_result(trial_id, result)
+
+        self._live_trial_mapping.pop(trial_id)
+
+    def _process_result(self, trial_id: str, result: Dict):
+        ng_trial_info = self._live_trial_mapping[trial_id]
+        self._nevergrad_opt.tell(ng_trial_info, self._metric_op * result[self._metric])
+
+    def save(self, checkpoint_path: str):
+        save_object = self.__dict__
+        with open(checkpoint_path, "wb") as outputFile:
+            pickle.dump(save_object, outputFile)
+
+    def restore(self, checkpoint_path: str):
+        with open(checkpoint_path, "rb") as inputFile:
+            save_object = pickle.load(inputFile)
+        self.__dict__.update(save_object)
+
+    @staticmethod
+    def convert_search_space(spec: Dict) -> Parameter:
+        resolved_vars, domain_vars, grid_vars = parse_spec_vars(spec)
+
+        if grid_vars:
+            raise ValueError(
+                "Grid search parameters cannot be automatically converted "
+                "to a Nevergrad 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(domain: Domain) -> Parameter:
+            sampler = domain.get_sampler()
+            if isinstance(sampler, Quantized):
+                logger.warning(
+                    "Nevergrad does not support quantization. Dropped quantization."
+                )
+                sampler = sampler.get_sampler()
+
+            if isinstance(domain, Float):
+                if isinstance(sampler, LogUniform):
+                    return ng.p.Log(
+                        lower=domain.lower, upper=domain.upper, exponent=sampler.base
+                    )
+                return ng.p.Scalar(lower=domain.lower, upper=domain.upper)
+
+            elif isinstance(domain, Integer):
+                if isinstance(sampler, LogUniform):
+                    return ng.p.Log(
+                        lower=domain.lower,
+                        upper=domain.upper - 1,  # Upper bound exclusive
+                        exponent=sampler.base,
+                    ).set_integer_casting()
+                return ng.p.Scalar(
+                    lower=domain.lower,
+                    upper=domain.upper - 1,  # Upper bound exclusive
+                ).set_integer_casting()
+
+            elif isinstance(domain, Categorical):
+                return ng.p.Choice(choices=domain.categories)
+
+            raise ValueError(
+                "Nevergrad 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
+        space = {"/".join(path): resolve_value(domain) for path, domain in domain_vars}
+
+        return ng.p.Dict(**space)

.venv/lib/python3.11/site-packages/ray/tune/search/optuna/__init__.py

@@ -0,0 +1,3 @@
+from ray.tune.search.optuna.optuna_search import OptunaSearch
+
+__all__ = ["OptunaSearch"]

.venv/lib/python3.11/site-packages/ray/tune/search/optuna/optuna_search.py

@@ -0,0 +1,733 @@
+import functools
+import logging
+import pickle
+import time
+import warnings
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+from packaging import version
+
+from ray.air.constants import TRAINING_ITERATION
+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,
+    Domain,
+    Float,
+    Integer,
+    LogUniform,
+    Quantized,
+    Uniform,
+)
+from ray.tune.search.variant_generator import parse_spec_vars
+from ray.tune.utils.util import flatten_dict, unflatten_dict, validate_warmstart
+
+try:
+    import optuna as ot
+    from optuna.distributions import BaseDistribution as OptunaDistribution
+    from optuna.samplers import BaseSampler
+    from optuna.storages import BaseStorage
+    from optuna.trial import Trial as OptunaTrial
+    from optuna.trial import TrialState as OptunaTrialState
+except ImportError:
+    ot = None
+    OptunaDistribution = None
+    BaseSampler = None
+    BaseStorage = None
+    OptunaTrialState = None
+    OptunaTrial = None
+
+logger = logging.getLogger(__name__)
+
+# print a warning if define by run function takes longer than this to execute
+DEFINE_BY_RUN_WARN_THRESHOLD_S = 1  # 1 is arbitrary
+
+
+class _OptunaTrialSuggestCaptor:
+    """Utility to capture returned values from Optuna's suggest_ methods.
+
+    This will wrap around the ``optuna.Trial` object and decorate all
+    `suggest_` callables with a function capturing the returned value,
+    which will be saved in the ``captured_values`` dict.
+    """
+
+    def __init__(self, ot_trial: OptunaTrial) -> None:
+        self.ot_trial = ot_trial
+        self.captured_values: Dict[str, Any] = {}
+
+    def _get_wrapper(self, func: Callable) -> Callable:
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            # name is always the first arg for suggest_ methods
+            name = kwargs.get("name", args[0])
+            ret = func(*args, **kwargs)
+            self.captured_values[name] = ret
+            return ret
+
+        return wrapper
+
+    def __getattr__(self, item_name: str) -> Any:
+        item = getattr(self.ot_trial, item_name)
+        if item_name.startswith("suggest_") and callable(item):
+            return self._get_wrapper(item)
+        return item
+
+
+class OptunaSearch(Searcher):
+    """A wrapper around Optuna to provide trial suggestions.
+
+    `Optuna <https://optuna.org/>`_ is a hyperparameter optimization library.
+    In contrast to other libraries, it employs define-by-run style
+    hyperparameter definitions.
+
+    This Searcher is a thin wrapper around Optuna's search algorithms.
+    You can pass any Optuna sampler, which will be used to generate
+    hyperparameter suggestions.
+
+    Multi-objective optimization is supported.
+
+    Args:
+        space: Hyperparameter search space definition for
+            Optuna's sampler. This can be either a :class:`dict` with
+            parameter names as keys and ``optuna.distributions`` as values,
+            or a Callable - in which case, it should be a define-by-run
+            function using ``optuna.trial`` to obtain the hyperparameter
+            values. The function should return either a :class:`dict` of
+            constant values with names as keys, or None.
+            For more information, see https://optuna.readthedocs.io\
+/en/stable/tutorial/10_key_features/002_configurations.html.
+
+            .. warning::
+                No actual computation should take place in the define-by-run
+                function. Instead, put the training logic inside the function
+                or class trainable passed to ``tune.Tuner()``.
+
+        metric: The training result objective value attribute. If
+            None but a mode was passed, the anonymous metric ``_metric``
+            will be used per default. Can be a list of metrics for
+            multi-objective optimization.
+        mode: One of {min, max}. Determines whether objective is
+            minimizing or maximizing the metric attribute. Can be a list of
+            modes for multi-objective optimization (corresponding to
+            ``metric``).
+        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.
+        sampler: Optuna sampler used to
+            draw hyperparameter configurations. Defaults to ``MOTPESampler``
+            for multi-objective optimization with Optuna<2.9.0, and
+            ``TPESampler`` in every other case.
+            See https://optuna.readthedocs.io/en/stable/reference/samplers/index.html
+            for available Optuna samplers.
+
+            .. warning::
+                Please note that with Optuna 2.10.0 and earlier
+                default ``MOTPESampler``/``TPESampler`` suffer
+                from performance issues when dealing with a large number of
+                completed trials (approx. >100). This will manifest as
+                a delay when suggesting new configurations.
+                This is an Optuna issue and may be fixed in a future
+                Optuna release.
+        study_name: Optuna study name that uniquely identifies the trial
+            results. Defaults to ``"optuna"``.
+        storage: Optuna storage used for storing trial results to
+            storages other than in-memory storage,
+            for instance optuna.storages.RDBStorage.
+        seed: Seed to initialize sampler with. This parameter is only
+            used when ``sampler=None``. In all other cases, the sampler
+            you pass should be initialized with the seed already.
+        evaluated_rewards: If you have previously evaluated the
+            parameters passed in as points_to_evaluate you can avoid
+            re-running those trials by passing in the reward attributes
+            as a list so the optimiser can be told the results without
+            needing to re-compute the trial. Must be the same length as
+            points_to_evaluate.
+
+            .. warning::
+                When using ``evaluated_rewards``, the search space ``space``
+                must be provided as a :class:`dict` with parameter names as
+                keys and ``optuna.distributions`` instances as values. The
+                define-by-run search space definition is not yet supported with
+                this functionality.
+
+    Tune automatically converts search spaces to Optuna's format:
+
+    .. code-block:: python
+
+        from ray.tune.search.optuna import OptunaSearch
+
+        config = {
+            "a": tune.uniform(6, 8)
+            "b": tune.loguniform(1e-4, 1e-2)
+        }
+
+        optuna_search = OptunaSearch(
+            metric="loss",
+            mode="min")
+
+        tuner = tune.Tuner(
+            trainable,
+            tune_config=tune.TuneConfig(
+                search_alg=optuna_search,
+            ),
+            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.tune.search.optuna import OptunaSearch
+        import optuna
+
+        space = {
+            "a": optuna.distributions.FloatDistribution(6, 8),
+            "b": optuna.distributions.FloatDistribution(1e-4, 1e-2, log=True),
+        }
+
+        optuna_search = OptunaSearch(
+            space,
+            metric="loss",
+            mode="min")
+
+        tuner = tune.Tuner(
+            trainable,
+            tune_config=tune.TuneConfig(
+                search_alg=optuna_search,
+            ),
+        )
+        tuner.fit()
+
+        # Equivalent Optuna define-by-run function approach:
+
+        def define_search_space(trial: optuna.Trial):
+            trial.suggest_float("a", 6, 8)
+            trial.suggest_float("b", 1e-4, 1e-2, log=True)
+            # training logic goes into trainable, this is just
+            # for search space definition
+
+        optuna_search = OptunaSearch(
+            define_search_space,
+            metric="loss",
+            mode="min")
+
+        tuner = tune.Tuner(
+            trainable,
+            tune_config=tune.TuneConfig(
+                search_alg=optuna_search,
+            ),
+        )
+        tuner.fit()
+
+    Multi-objective optimization is supported:
+
+    .. code-block:: python
+
+        from ray.tune.search.optuna import OptunaSearch
+        import optuna
+
+        space = {
+            "a": optuna.distributions.FloatDistribution(6, 8),
+            "b": optuna.distributions.FloatDistribution(1e-4, 1e-2, log=True),
+        }
+
+        # Note you have to specify metric and mode here instead of
+        # in tune.TuneConfig
+        optuna_search = OptunaSearch(
+            space,
+            metric=["loss1", "loss2"],
+            mode=["min", "max"])
+
+        # Do not specify metric and mode here!
+        tuner = tune.Tuner(
+            trainable,
+            tune_config=tune.TuneConfig(
+                search_alg=optuna_search,
+            ),
+        )
+        tuner.fit()
+
+    You can pass configs that will be evaluated first using
+    ``points_to_evaluate``:
+
+    .. code-block:: python
+
+        from ray.tune.search.optuna import OptunaSearch
+        import optuna
+
+        space = {
+            "a": optuna.distributions.FloatDistribution(6, 8),
+            "b": optuna.distributions.FloatDistribution(1e-4, 1e-2, log=True),
+        }
+
+        optuna_search = OptunaSearch(
+            space,
+            points_to_evaluate=[{"a": 6.5, "b": 5e-4}, {"a": 7.5, "b": 1e-3}]
+            metric="loss",
+            mode="min")
+
+        tuner = tune.Tuner(
+            trainable,
+            tune_config=tune.TuneConfig(
+                search_alg=optuna_search,
+            ),
+        )
+        tuner.fit()
+
+    Avoid re-running evaluated trials by passing the rewards together with
+    `points_to_evaluate`:
+
+    .. code-block:: python
+
+        from ray.tune.search.optuna import OptunaSearch
+        import optuna
+
+        space = {
+            "a": optuna.distributions.FloatDistribution(6, 8),
+            "b": optuna.distributions.FloatDistribution(1e-4, 1e-2, log=True),
+        }
+
+        optuna_search = OptunaSearch(
+            space,
+            points_to_evaluate=[{"a": 6.5, "b": 5e-4}, {"a": 7.5, "b": 1e-3}]
+            evaluated_rewards=[0.89, 0.42]
+            metric="loss",
+            mode="min")
+
+        tuner = tune.Tuner(
+            trainable,
+            tune_config=tune.TuneConfig(
+                search_alg=optuna_search,
+            ),
+        )
+        tuner.fit()
+
+    .. versionadded:: 0.8.8
+
+    """
+
+    def __init__(
+        self,
+        space: Optional[
+            Union[
+                Dict[str, "OptunaDistribution"],
+                List[Tuple],
+                Callable[["OptunaTrial"], Optional[Dict[str, Any]]],
+            ]
+        ] = None,
+        metric: Optional[Union[str, List[str]]] = None,
+        mode: Optional[Union[str, List[str]]] = None,
+        points_to_evaluate: Optional[List[Dict]] = None,
+        sampler: Optional["BaseSampler"] = None,
+        study_name: Optional[str] = None,
+        storage: Optional["BaseStorage"] = None,
+        seed: Optional[int] = None,
+        evaluated_rewards: Optional[List] = None,
+    ):
+        assert ot is not None, "Optuna must be installed! Run `pip install optuna`."
+        super(OptunaSearch, self).__init__(metric=metric, mode=mode)
+
+        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).__name__)
+                )
+                space = self.convert_search_space(space)
+            else:
+                # Flatten to support nested dicts
+                space = flatten_dict(space, "/")
+
+        self._space = space
+
+        self._points_to_evaluate = points_to_evaluate or []
+        self._evaluated_rewards = evaluated_rewards
+        if study_name:
+            self._study_name = study_name
+        else:
+            self._study_name = "optuna"  # Fixed study name for in-memory storage
+
+        if sampler and seed:
+            logger.warning(
+                "You passed an initialized sampler to `OptunaSearch`. The "
+                "`seed` parameter has to be passed to the sampler directly "
+                "and will be ignored."
+            )
+        elif sampler:
+            assert isinstance(sampler, BaseSampler), (
+                "You can only pass an instance of "
+                "`optuna.samplers.BaseSampler` "
+                "as a sampler to `OptunaSearcher`."
+            )
+
+        self._sampler = sampler
+        self._seed = seed
+
+        if storage:
+            assert isinstance(storage, BaseStorage), (
+                "The `storage` parameter in `OptunaSearcher` must be an instance "
+                "of `optuna.storages.BaseStorage`."
+            )
+        # If storage is not provided, just set self._storage to None
+        # so that the default in-memory storage is used.
+        self._storage = storage
+
+        self._completed_trials = set()
+
+        self._ot_trials = {}
+        self._ot_study = None
+        if self._space:
+            self._setup_study(mode)
+
+    def _setup_study(self, mode: Union[str, list]):
+        if self._metric is None and self._mode:
+            if isinstance(self._mode, list):
+                raise ValueError(
+                    "If ``mode`` is a list (multi-objective optimization "
+                    "case), ``metric`` must be defined."
+                )
+            # If only a mode was passed, use anonymous metric
+            self._metric = DEFAULT_METRIC
+
+        pruner = ot.pruners.NopPruner()
+
+        if self._sampler:
+            sampler = self._sampler
+        elif isinstance(mode, list) and version.parse(ot.__version__) < version.parse(
+            "2.9.0"
+        ):
+            # MOTPESampler deprecated in Optuna>=2.9.0
+            sampler = ot.samplers.MOTPESampler(seed=self._seed)
+        else:
+            sampler = ot.samplers.TPESampler(seed=self._seed)
+
+        if isinstance(mode, list):
+            study_direction_args = dict(
+                directions=["minimize" if m == "min" else "maximize" for m in mode],
+            )
+        else:
+            study_direction_args = dict(
+                direction="minimize" if mode == "min" else "maximize",
+            )
+
+        self._ot_study = ot.study.create_study(
+            storage=self._storage,
+            sampler=sampler,
+            pruner=pruner,
+            study_name=self._study_name,
+            load_if_exists=True,
+            **study_direction_args,
+        )
+
+        if self._points_to_evaluate:
+            validate_warmstart(
+                self._space,
+                self._points_to_evaluate,
+                self._evaluated_rewards,
+                validate_point_name_lengths=not callable(self._space),
+            )
+            if self._evaluated_rewards:
+                for point, reward in zip(
+                    self._points_to_evaluate, self._evaluated_rewards
+                ):
+                    self.add_evaluated_point(point, reward)
+            else:
+                for point in self._points_to_evaluate:
+                    self._ot_study.enqueue_trial(point)
+
+    def set_search_properties(
+        self, metric: Optional[str], mode: Optional[str], config: Dict, **spec
+    ) -> bool:
+        if self._space:
+            return False
+        space = self.convert_search_space(config)
+        self._space = space
+        if metric:
+            self._metric = metric
+        if mode:
+            self._mode = mode
+
+        self._setup_study(self._mode)
+        return True
+
+    def _suggest_from_define_by_run_func(
+        self,
+        func: Callable[["OptunaTrial"], Optional[Dict[str, Any]]],
+        ot_trial: "OptunaTrial",
+    ) -> Dict:
+        captor = _OptunaTrialSuggestCaptor(ot_trial)
+        time_start = time.time()
+        ret = func(captor)
+        time_taken = time.time() - time_start
+        if time_taken > DEFINE_BY_RUN_WARN_THRESHOLD_S:
+            warnings.warn(
+                "Define-by-run function passed in the `space` argument "
+                f"took {time_taken} seconds to "
+                "run. Ensure that actual computation, training takes "
+                "place inside Tune's train functions or Trainables "
+                "passed to `tune.Tuner()`."
+            )
+        if ret is not None:
+            if not isinstance(ret, dict):
+                raise TypeError(
+                    "The return value of the define-by-run function "
+                    "passed in the `space` argument should be "
+                    "either None or a `dict` with `str` keys. "
+                    f"Got {type(ret)}."
+                )
+            if not all(isinstance(k, str) for k in ret.keys()):
+                raise TypeError(
+                    "At least one of the keys in the dict returned by the "
+                    "define-by-run function passed in the `space` argument "
+                    "was not a `str`."
+                )
+        return {**captor.captured_values, **ret} if ret else captor.captured_values
+
+    def suggest(self, trial_id: str) -> Optional[Dict]:
+        if not self._space:
+            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 callable(self._space):
+            # Define-by-run case
+            if trial_id not in self._ot_trials:
+                self._ot_trials[trial_id] = self._ot_study.ask()
+
+            ot_trial = self._ot_trials[trial_id]
+
+            params = self._suggest_from_define_by_run_func(self._space, ot_trial)
+        else:
+            # Use Optuna ask interface (since version 2.6.0)
+            if trial_id not in self._ot_trials:
+                self._ot_trials[trial_id] = self._ot_study.ask(
+                    fixed_distributions=self._space
+                )
+            ot_trial = self._ot_trials[trial_id]
+            params = ot_trial.params
+
+        return unflatten_dict(params)
+
+    def on_trial_result(self, trial_id: str, result: Dict):
+        if isinstance(self.metric, list):
+            # Optuna doesn't support incremental results
+            # for multi-objective optimization
+            return
+        if trial_id in self._completed_trials:
+            logger.warning(
+                f"Received additional result for trial {trial_id}, but "
+                f"it already finished. Result: {result}"
+            )
+            return
+        metric = result[self.metric]
+        step = result[TRAINING_ITERATION]
+        ot_trial = self._ot_trials[trial_id]
+        ot_trial.report(metric, step)
+
+    def on_trial_complete(
+        self, trial_id: str, result: Optional[Dict] = None, error: bool = False
+    ):
+        if trial_id in self._completed_trials:
+            logger.warning(
+                f"Received additional completion for trial {trial_id}, but "
+                f"it already finished. Result: {result}"
+            )
+            return
+
+        ot_trial = self._ot_trials[trial_id]
+
+        if result:
+            if isinstance(self.metric, list):
+                val = [result.get(metric, None) for metric in self.metric]
+            else:
+                val = result.get(self.metric, None)
+        else:
+            val = None
+        ot_trial_state = OptunaTrialState.COMPLETE
+        if val is None:
+            if error:
+                ot_trial_state = OptunaTrialState.FAIL
+            else:
+                ot_trial_state = OptunaTrialState.PRUNED
+        try:
+            self._ot_study.tell(ot_trial, val, state=ot_trial_state)
+        except Exception as exc:
+            logger.warning(exc)  # E.g. if NaN was reported
+
+        self._completed_trials.add(trial_id)
+
+    def add_evaluated_point(
+        self,
+        parameters: Dict,
+        value: float,
+        error: bool = False,
+        pruned: bool = False,
+        intermediate_values: Optional[List[float]] = None,
+    ):
+        if not self._space:
+            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 callable(self._space):
+            raise TypeError(
+                "Define-by-run function passed in `space` argument is not "
+                "yet supported when using `evaluated_rewards`. Please provide "
+                "an `OptunaDistribution` dict or pass a Ray Tune "
+                "search space to `tune.Tuner()`."
+            )
+
+        ot_trial_state = OptunaTrialState.COMPLETE
+        if error:
+            ot_trial_state = OptunaTrialState.FAIL
+        elif pruned:
+            ot_trial_state = OptunaTrialState.PRUNED
+
+        if intermediate_values:
+            intermediate_values_dict = {
+                i: value for i, value in enumerate(intermediate_values)
+            }
+        else:
+            intermediate_values_dict = None
+
+        # If the trial state is FAILED, the value must be `None` in Optuna==4.1.0
+        # Reference: https://github.com/optuna/optuna/pull/5211
+        # This is a temporary fix for the issue that Optuna enforces the value
+        # to be `None` if the trial state is FAILED.
+        # TODO (hpguo): A better solution may requires us to update the base class
+        # to allow the `value` arg in `add_evaluated_point` being `Optional[float]`.
+        if ot_trial_state == OptunaTrialState.FAIL:
+            value = None
+
+        trial = ot.trial.create_trial(
+            state=ot_trial_state,
+            value=value,
+            params=parameters,
+            distributions=self._space,
+            intermediate_values=intermediate_values_dict,
+        )
+
+        self._ot_study.add_trial(trial)
+
+    def save(self, checkpoint_path: str):
+        save_object = self.__dict__.copy()
+        with open(checkpoint_path, "wb") as outputFile:
+            pickle.dump(save_object, outputFile)
+
+    def restore(self, checkpoint_path: str):
+        with open(checkpoint_path, "rb") as inputFile:
+            save_object = pickle.load(inputFile)
+        if isinstance(save_object, dict):
+            self.__dict__.update(save_object)
+        else:
+            # Backwards compatibility
+            (
+                self._sampler,
+                self._ot_trials,
+                self._ot_study,
+                self._points_to_evaluate,
+                self._evaluated_rewards,
+            ) = save_object
+
+    @staticmethod
+    def convert_search_space(spec: Dict) -> Dict[str, Any]:
+        resolved_vars, domain_vars, grid_vars = parse_spec_vars(spec)
+
+        if not domain_vars and not grid_vars:
+            return {}
+
+        if grid_vars:
+            raise ValueError(
+                "Grid search parameters cannot be automatically converted "
+                "to an Optuna 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(domain: Domain) -> ot.distributions.BaseDistribution:
+            quantize = None
+
+            sampler = domain.get_sampler()
+            if isinstance(sampler, Quantized):
+                quantize = sampler.q
+                sampler = sampler.sampler
+                if isinstance(sampler, LogUniform):
+                    logger.warning(
+                        "Optuna does not handle quantization in loguniform "
+                        "sampling. The parameter will be passed but it will "
+                        "probably be ignored."
+                    )
+
+            if isinstance(domain, Float):
+                if isinstance(sampler, LogUniform):
+                    if quantize:
+                        logger.warning(
+                            "Optuna does not support both quantization and "
+                            "sampling from LogUniform. Dropped quantization."
+                        )
+                    return ot.distributions.FloatDistribution(
+                        domain.lower, domain.upper, log=True
+                    )
+
+                elif isinstance(sampler, Uniform):
+                    if quantize:
+                        return ot.distributions.FloatDistribution(
+                            domain.lower, domain.upper, step=quantize
+                        )
+                    return ot.distributions.FloatDistribution(
+                        domain.lower, domain.upper
+                    )
+
+            elif isinstance(domain, Integer):
+                if isinstance(sampler, LogUniform):
+                    return ot.distributions.IntDistribution(
+                        domain.lower, domain.upper - 1, step=quantize or 1, log=True
+                    )
+                elif isinstance(sampler, Uniform):
+                    # Upper bound should be inclusive for quantization and
+                    # exclusive otherwise
+                    return ot.distributions.IntDistribution(
+                        domain.lower,
+                        domain.upper - int(bool(not quantize)),
+                        step=quantize or 1,
+                    )
+            elif isinstance(domain, Categorical):
+                if isinstance(sampler, Uniform):
+                    return ot.distributions.CategoricalDistribution(domain.categories)
+
+            raise ValueError(
+                "Optuna search 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
+        values = {"/".join(path): resolve_value(domain) for path, domain in domain_vars}
+
+        return values

.venv/lib/python3.11/site-packages/ray/tune/search/repeater.py

@@ -0,0 +1,199 @@
+import copy
+import logging
+from typing import Dict, List, Optional
+
+import numpy as np
+
+from ray.tune.search.searcher import Searcher
+from ray.tune.search.util import _set_search_properties_backwards_compatible
+from ray.util import PublicAPI
+
+logger = logging.getLogger(__name__)
+
+TRIAL_INDEX = "__trial_index__"
+"""str: A constant value representing the repeat index of the trial."""
+
+
+def _warn_num_samples(searcher: Searcher, num_samples: int):
+    if isinstance(searcher, Repeater) and num_samples % searcher.repeat:
+        logger.warning(
+            "`num_samples` is now expected to be the total number of trials, "
+            "including the repeat trials. For example, set num_samples=15 if "
+            "you intend to obtain 3 search algorithm suggestions and repeat "
+            "each suggestion 5 times. Any leftover trials "
+            "(num_samples mod repeat) will be ignored."
+        )
+
+
+class _TrialGroup:
+    """Internal class for grouping trials of same parameters.
+
+    This is used when repeating trials for reducing training variance.
+
+    Args:
+        primary_trial_id: Trial ID of the "primary trial".
+            This trial is the one that the Searcher is aware of.
+        config: Suggested configuration shared across all trials
+            in the trial group.
+        max_trials: Max number of trials to execute within this group.
+
+    """
+
+    def __init__(self, primary_trial_id: str, config: Dict, max_trials: int = 1):
+        assert type(config) is dict, "config is not a dict, got {}".format(config)
+        self.primary_trial_id = primary_trial_id
+        self.config = config
+        self._trials = {primary_trial_id: None}
+        self.max_trials = max_trials
+
+    def add(self, trial_id: str):
+        assert len(self._trials) < self.max_trials
+        self._trials.setdefault(trial_id, None)
+
+    def full(self) -> bool:
+        return len(self._trials) == self.max_trials
+
+    def report(self, trial_id: str, score: float):
+        assert trial_id in self._trials
+        if score is None:
+            raise ValueError("Internal Error: Score cannot be None.")
+        self._trials[trial_id] = score
+
+    def finished_reporting(self) -> bool:
+        return (
+            None not in self._trials.values() and len(self._trials) == self.max_trials
+        )
+
+    def scores(self) -> List[Optional[float]]:
+        return list(self._trials.values())
+
+    def count(self) -> int:
+        return len(self._trials)
+
+
+@PublicAPI
+class Repeater(Searcher):
+    """A wrapper algorithm for repeating trials of same parameters.
+
+    Set tune.TuneConfig(num_samples=...) to be a multiple of `repeat`. For example,
+    set num_samples=15 if you intend to obtain 3 search algorithm suggestions
+    and repeat each suggestion 5 times. Any leftover trials
+    (num_samples mod repeat) will be ignored.
+
+    It is recommended that you do not run an early-stopping TrialScheduler
+    simultaneously.
+
+    Args:
+        searcher: Searcher object that the
+            Repeater will optimize. Note that the Searcher
+            will only see 1 trial among multiple repeated trials.
+            The result/metric passed to the Searcher upon
+            trial completion will be averaged among all repeats.
+        repeat: Number of times to generate a trial with a repeated
+            configuration. Defaults to 1.
+        set_index: Sets a tune.search.repeater.TRIAL_INDEX in
+            Trainable/Function config which corresponds to the index of the
+            repeated trial. This can be used for seeds. Defaults to True.
+
+    Example:
+
+    .. code-block:: python
+
+        from ray.tune.search import Repeater
+
+        search_alg = BayesOptSearch(...)
+        re_search_alg = Repeater(search_alg, repeat=10)
+
+        # Repeat 2 samples 10 times each.
+        tuner = tune.Tuner(
+            trainable,
+            tune_config=tune.TuneConfig(
+                search_alg=re_search_alg,
+                num_samples=20,
+            ),
+        )
+        tuner.fit()
+
+    """
+
+    def __init__(self, searcher: Searcher, repeat: int = 1, set_index: bool = True):
+        self.searcher = searcher
+        self.repeat = repeat
+        self._set_index = set_index
+        self._groups = []
+        self._trial_id_to_group = {}
+        self._current_group = None
+        super(Repeater, self).__init__(
+            metric=self.searcher.metric, mode=self.searcher.mode
+        )
+
+    def suggest(self, trial_id: str) -> Optional[Dict]:
+        if self._current_group is None or self._current_group.full():
+            config = self.searcher.suggest(trial_id)
+            if config is None:
+                return config
+            self._current_group = _TrialGroup(
+                trial_id, copy.deepcopy(config), max_trials=self.repeat
+            )
+            self._groups.append(self._current_group)
+            index_in_group = 0
+        else:
+            index_in_group = self._current_group.count()
+            self._current_group.add(trial_id)
+
+        config = self._current_group.config.copy()
+        if self._set_index:
+            config[TRIAL_INDEX] = index_in_group
+        self._trial_id_to_group[trial_id] = self._current_group
+        return config
+
+    def on_trial_complete(self, trial_id: str, result: Optional[Dict] = None, **kwargs):
+        """Stores the score for and keeps track of a completed trial.
+
+        Stores the metric of a trial as nan if any of the following conditions
+        are met:
+
+        1. ``result`` is empty or not provided.
+        2. ``result`` is provided but no metric was provided.
+
+        """
+        if trial_id not in self._trial_id_to_group:
+            logger.error(
+                "Trial {} not in group; cannot report score. "
+                "Seen trials: {}".format(trial_id, list(self._trial_id_to_group))
+            )
+        trial_group = self._trial_id_to_group[trial_id]
+        if not result or self.searcher.metric not in result:
+            score = np.nan
+        else:
+            score = result[self.searcher.metric]
+        trial_group.report(trial_id, score)
+
+        if trial_group.finished_reporting():
+            scores = trial_group.scores()
+            self.searcher.on_trial_complete(
+                trial_group.primary_trial_id,
+                result={self.searcher.metric: np.nanmean(scores)},
+                **kwargs
+            )
+
+    def get_state(self) -> Dict:
+        self_state = self.__dict__.copy()
+        del self_state["searcher"]
+        return self_state
+
+    def set_state(self, state: Dict):
+        self.__dict__.update(state)
+
+    def save(self, checkpoint_path: str):
+        self.searcher.save(checkpoint_path)
+
+    def restore(self, checkpoint_path: str):
+        self.searcher.restore(checkpoint_path)
+
+    def set_search_properties(
+        self, metric: Optional[str], mode: Optional[str], config: Dict, **spec
+    ) -> bool:
+        return _set_search_properties_backwards_compatible(
+            self.searcher.set_search_properties, metric, mode, config, **spec
+        )

.venv/lib/python3.11/site-packages/ray/tune/search/sample.py

@@ -0,0 +1,742 @@
+import logging
+from copy import copy
+from inspect import signature
+from math import isclose
+from typing import Any, Callable, Dict, List, Optional, Sequence, Union
+
+import numpy as np
+
+# Backwards compatibility
+from ray.util.annotations import DeveloperAPI, PublicAPI
+
+try:
+    # Added in numpy>=1.17 but we require numpy>=1.16
+    np_random_generator = np.random.Generator
+    LEGACY_RNG = False
+except AttributeError:
+
+    class np_random_generator:
+        pass
+
+    LEGACY_RNG = True
+
+logger = logging.getLogger(__name__)
+
+
+class _BackwardsCompatibleNumpyRng:
+    """Thin wrapper to ensure backwards compatibility between
+    new and old numpy randomness generators.
+    """
+
+    _rng = None
+
+    def __init__(
+        self,
+        generator_or_seed: Optional[
+            Union["np_random_generator", np.random.RandomState, int]
+        ] = None,
+    ):
+        if generator_or_seed is None or isinstance(
+            generator_or_seed, (np.random.RandomState, np_random_generator)
+        ):
+            self._rng = generator_or_seed
+        elif LEGACY_RNG:
+            self._rng = np.random.RandomState(generator_or_seed)
+        else:
+            self._rng = np.random.default_rng(generator_or_seed)
+
+    @property
+    def legacy_rng(self) -> bool:
+        return not isinstance(self._rng, np_random_generator)
+
+    @property
+    def rng(self):
+        # don't set self._rng to np.random to avoid picking issues
+        return self._rng if self._rng is not None else np.random
+
+    def __getattr__(self, name: str) -> Any:
+        # https://numpy.org/doc/stable/reference/random/new-or-different.html
+        if self.legacy_rng:
+            if name == "integers":
+                name = "randint"
+            elif name == "random":
+                name = "rand"
+        return getattr(self.rng, name)
+
+
+RandomState = Union[
+    None, _BackwardsCompatibleNumpyRng, np_random_generator, np.random.RandomState, int
+]
+
+
+@DeveloperAPI
+class Domain:
+    """Base class to specify a type and valid range to sample parameters from.
+
+    This base class is implemented by parameter spaces, like float ranges
+    (``Float``), integer ranges (``Integer``), or categorical variables
+    (``Categorical``). The ``Domain`` object contains information about
+    valid values (e.g. minimum and maximum values), and exposes methods that
+    allow specification of specific samplers (e.g. ``uniform()`` or
+    ``loguniform()``).
+
+    """
+
+    sampler = None
+    default_sampler_cls = None
+
+    def cast(self, value):
+        """Cast value to domain type"""
+        return value
+
+    def set_sampler(self, sampler, allow_override=False):
+        if self.sampler and not allow_override:
+            raise ValueError(
+                "You can only choose one sampler for parameter "
+                "domains. Existing sampler for parameter {}: "
+                "{}. Tried to add {}".format(
+                    self.__class__.__name__, self.sampler, sampler
+                )
+            )
+        self.sampler = sampler
+
+    def get_sampler(self):
+        sampler = self.sampler
+        if not sampler:
+            sampler = self.default_sampler_cls()
+        return sampler
+
+    def sample(
+        self,
+        config: Optional[Union[List[Dict], Dict]] = None,
+        size: int = 1,
+        random_state: "RandomState" = None,
+    ):
+        if not isinstance(random_state, _BackwardsCompatibleNumpyRng):
+            random_state = _BackwardsCompatibleNumpyRng(random_state)
+        sampler = self.get_sampler()
+        return sampler.sample(self, config=config, size=size, random_state=random_state)
+
+    def is_grid(self):
+        return isinstance(self.sampler, Grid)
+
+    def is_function(self):
+        return False
+
+    def is_valid(self, value: Any):
+        """Returns True if `value` is a valid value in this domain."""
+        raise NotImplementedError
+
+    @property
+    def domain_str(self):
+        return "(unknown)"
+
+
+@DeveloperAPI
+class Sampler:
+    def sample(
+        self,
+        domain: Domain,
+        config: Optional[Union[List[Dict], Dict]] = None,
+        size: int = 1,
+        random_state: "RandomState" = None,
+    ):
+        raise NotImplementedError
+
+
+@DeveloperAPI
+class BaseSampler(Sampler):
+    def __str__(self):
+        return "Base"
+
+
+@DeveloperAPI
+class Uniform(Sampler):
+    def __str__(self):
+        return "Uniform"
+
+
+@DeveloperAPI
+class LogUniform(Sampler):
+    def __init__(self, base: float = 10):
+        self.base = base
+        assert self.base > 0, "Base has to be strictly greater than 0"
+
+    def __str__(self):
+        return "LogUniform"
+
+
+@DeveloperAPI
+class Normal(Sampler):
+    def __init__(self, mean: float = 0.0, sd: float = 0.0):
+        self.mean = mean
+        self.sd = sd
+
+        assert self.sd > 0, "SD has to be strictly greater than 0"
+
+    def __str__(self):
+        return "Normal"
+
+
+@DeveloperAPI
+class Grid(Sampler):
+    """Dummy sampler used for grid search"""
+
+    def sample(
+        self,
+        domain: Domain,
+        config: Optional[Union[List[Dict], Dict]] = None,
+        size: int = 1,
+        random_state: "RandomState" = None,
+    ):
+        return RuntimeError("Do not call `sample()` on grid.")
+
+
+@DeveloperAPI
+class Float(Domain):
+    class _Uniform(Uniform):
+        def sample(
+            self,
+            domain: "Float",
+            config: Optional[Union[List[Dict], Dict]] = None,
+            size: int = 1,
+            random_state: "RandomState" = None,
+        ):
+            if not isinstance(random_state, _BackwardsCompatibleNumpyRng):
+                random_state = _BackwardsCompatibleNumpyRng(random_state)
+            assert domain.lower > float("-inf"), "Uniform needs a lower bound"
+            assert domain.upper < float("inf"), "Uniform needs a upper bound"
+            items = random_state.uniform(domain.lower, domain.upper, size=size)
+            return items if len(items) > 1 else domain.cast(items[0])
+
+    class _LogUniform(LogUniform):
+        def sample(
+            self,
+            domain: "Float",
+            config: Optional[Union[List[Dict], Dict]] = None,
+            size: int = 1,
+            random_state: "RandomState" = None,
+        ):
+            if not isinstance(random_state, _BackwardsCompatibleNumpyRng):
+                random_state = _BackwardsCompatibleNumpyRng(random_state)
+            assert domain.lower > 0, "LogUniform needs a lower bound greater than 0"
+            assert (
+                0 < domain.upper < float("inf")
+            ), "LogUniform needs a upper bound greater than 0"
+            logmin = np.log(domain.lower) / np.log(self.base)
+            logmax = np.log(domain.upper) / np.log(self.base)
+
+            items = self.base ** (random_state.uniform(logmin, logmax, size=size))
+            return items if len(items) > 1 else domain.cast(items[0])
+
+    class _Normal(Normal):
+        def sample(
+            self,
+            domain: "Float",
+            config: Optional[Union[List[Dict], Dict]] = None,
+            size: int = 1,
+            random_state: "RandomState" = None,
+        ):
+            if not isinstance(random_state, _BackwardsCompatibleNumpyRng):
+                random_state = _BackwardsCompatibleNumpyRng(random_state)
+            assert not domain.lower or domain.lower == float(
+                "-inf"
+            ), "Normal sampling does not allow a lower value bound."
+            assert not domain.upper or domain.upper == float(
+                "inf"
+            ), "Normal sampling does not allow a upper value bound."
+            items = random_state.normal(self.mean, self.sd, size=size)
+            return items if len(items) > 1 else domain.cast(items[0])
+
+    default_sampler_cls = _Uniform
+
+    def __init__(self, lower: Optional[float], upper: Optional[float]):
+        # Need to explicitly check for None
+        self.lower = lower if lower is not None else float("-inf")
+        self.upper = upper if upper is not None else float("inf")
+
+    def cast(self, value):
+        return float(value)
+
+    def uniform(self):
+        if not self.lower > float("-inf"):
+            raise ValueError(
+                "Uniform requires a lower bound. Make sure to set the "
+                "`lower` parameter of `Float()`."
+            )
+        if not self.upper < float("inf"):
+            raise ValueError(
+                "Uniform requires a upper bound. Make sure to set the "
+                "`upper` parameter of `Float()`."
+            )
+        new = copy(self)
+        new.set_sampler(self._Uniform())
+        return new
+
+    def loguniform(self, base: float = 10):
+        if not self.lower > 0:
+            raise ValueError(
+                "LogUniform requires a lower bound greater than 0."
+                f"Got: {self.lower}. Did you pass a variable that has "
+                "been log-transformed? If so, pass the non-transformed value "
+                "instead."
+            )
+        if not 0 < self.upper < float("inf"):
+            raise ValueError(
+                "LogUniform requires a upper bound greater than 0. "
+                f"Got: {self.lower}. Did you pass a variable that has "
+                "been log-transformed? If so, pass the non-transformed value "
+                "instead."
+            )
+        new = copy(self)
+        new.set_sampler(self._LogUniform(base))
+        return new
+
+    def normal(self, mean=0.0, sd=1.0):
+        new = copy(self)
+        new.set_sampler(self._Normal(mean, sd))
+        return new
+
+    def quantized(self, q: float):
+        if self.lower > float("-inf") and not isclose(
+            self.lower / q, round(self.lower / q)
+        ):
+            raise ValueError(
+                f"Your lower variable bound {self.lower} is not divisible by "
+                f"quantization factor {q}."
+            )
+        if self.upper < float("inf") and not isclose(
+            self.upper / q, round(self.upper / q)
+        ):
+            raise ValueError(
+                f"Your upper variable bound {self.upper} is not divisible by "
+                f"quantization factor {q}."
+            )
+
+        new = copy(self)
+        new.set_sampler(Quantized(new.get_sampler(), q), allow_override=True)
+        return new
+
+    def is_valid(self, value: float):
+        return self.lower <= value <= self.upper
+
+    @property
+    def domain_str(self):
+        return f"({self.lower}, {self.upper})"
+
+
+@DeveloperAPI
+class Integer(Domain):
+    class _Uniform(Uniform):
+        def sample(
+            self,
+            domain: "Integer",
+            config: Optional[Union[List[Dict], Dict]] = None,
+            size: int = 1,
+            random_state: "RandomState" = None,
+        ):
+            if not isinstance(random_state, _BackwardsCompatibleNumpyRng):
+                random_state = _BackwardsCompatibleNumpyRng(random_state)
+            items = random_state.integers(domain.lower, domain.upper, size=size)
+            return items if len(items) > 1 else domain.cast(items[0])
+
+    class _LogUniform(LogUniform):
+        def sample(
+            self,
+            domain: "Integer",
+            config: Optional[Union[List[Dict], Dict]] = None,
+            size: int = 1,
+            random_state: "RandomState" = None,
+        ):
+            if not isinstance(random_state, _BackwardsCompatibleNumpyRng):
+                random_state = _BackwardsCompatibleNumpyRng(random_state)
+            assert domain.lower > 0, "LogUniform needs a lower bound greater than 0"
+            assert (
+                0 < domain.upper < float("inf")
+            ), "LogUniform needs a upper bound greater than 0"
+            logmin = np.log(domain.lower) / np.log(self.base)
+            logmax = np.log(domain.upper) / np.log(self.base)
+
+            items = self.base ** (random_state.uniform(logmin, logmax, size=size))
+            items = np.floor(items).astype(int)
+            return items if len(items) > 1 else domain.cast(items[0])
+
+    default_sampler_cls = _Uniform
+
+    def __init__(self, lower, upper):
+        self.lower = lower
+        self.upper = upper
+
+    def cast(self, value):
+        return int(value)
+
+    def quantized(self, q: int):
+        new = copy(self)
+        new.set_sampler(Quantized(new.get_sampler(), q), allow_override=True)
+        return new
+
+    def uniform(self):
+        new = copy(self)
+        new.set_sampler(self._Uniform())
+        return new
+
+    def loguniform(self, base: float = 10):
+        if not self.lower > 0:
+            raise ValueError(
+                "LogUniform requires a lower bound greater than 0."
+                f"Got: {self.lower}. Did you pass a variable that has "
+                "been log-transformed? If so, pass the non-transformed value "
+                "instead."
+            )
+        if not 0 < self.upper < float("inf"):
+            raise ValueError(
+                "LogUniform requires a upper bound greater than 0. "
+                f"Got: {self.lower}. Did you pass a variable that has "
+                "been log-transformed? If so, pass the non-transformed value "
+                "instead."
+            )
+        new = copy(self)
+        new.set_sampler(self._LogUniform(base))
+        return new
+
+    def is_valid(self, value: int):
+        return self.lower <= value <= self.upper
+
+    @property
+    def domain_str(self):
+        return f"({self.lower}, {self.upper})"
+
+
+@DeveloperAPI
+class Categorical(Domain):
+    class _Uniform(Uniform):
+        def sample(
+            self,
+            domain: "Categorical",
+            config: Optional[Union[List[Dict], Dict]] = None,
+            size: int = 1,
+            random_state: "RandomState" = None,
+        ):
+            if not isinstance(random_state, _BackwardsCompatibleNumpyRng):
+                random_state = _BackwardsCompatibleNumpyRng(random_state)
+            # do not use .choice() directly on domain.categories
+            # as that will coerce them to a single dtype
+            indices = random_state.choice(
+                np.arange(0, len(domain.categories)), size=size
+            )
+            items = [domain.categories[index] for index in indices]
+            return items if len(items) > 1 else domain.cast(items[0])
+
+    default_sampler_cls = _Uniform
+
+    def __init__(self, categories: Sequence):
+        self.categories = list(categories)
+
+    def uniform(self):
+        new = copy(self)
+        new.set_sampler(self._Uniform())
+        return new
+
+    def grid(self):
+        new = copy(self)
+        new.set_sampler(Grid())
+        return new
+
+    def __len__(self):
+        return len(self.categories)
+
+    def __getitem__(self, item):
+        return self.categories[item]
+
+    def is_valid(self, value: Any):
+        return value in self.categories
+
+    @property
+    def domain_str(self):
+        return f"{self.categories}"
+
+
+@DeveloperAPI
+class Function(Domain):
+    class _CallSampler(BaseSampler):
+        def __try_fn(self, domain: "Function", config: Dict[str, Any]):
+            try:
+                return domain.func(config)
+            except (AttributeError, KeyError):
+                from ray.tune.search.variant_generator import _UnresolvedAccessGuard
+
+                r = domain.func(_UnresolvedAccessGuard({"config": config}))
+                logger.warning(
+                    "sample_from functions that take a spec dict are "
+                    "deprecated. Please update your function to work with "
+                    "the config dict directly."
+                )
+                return r
+
+        def sample(
+            self,
+            domain: "Function",
+            config: Optional[Union[List[Dict], Dict]] = None,
+            size: int = 1,
+            random_state: "RandomState" = None,
+        ):
+            if not isinstance(random_state, _BackwardsCompatibleNumpyRng):
+                random_state = _BackwardsCompatibleNumpyRng(random_state)
+            if domain.pass_config:
+                items = [
+                    self.__try_fn(domain, config[i])
+                    if isinstance(config, list)
+                    else self.__try_fn(domain, config)
+                    for i in range(size)
+                ]
+            else:
+                items = [domain.func() for i in range(size)]
+
+            return items if len(items) > 1 else domain.cast(items[0])
+
+    default_sampler_cls = _CallSampler
+
+    def __init__(self, func: Callable):
+        sig = signature(func)
+
+        pass_config = True  # whether we should pass `config` when calling `func`
+        try:
+            sig.bind({})
+        except TypeError:
+            pass_config = False
+
+        if not pass_config:
+            try:
+                sig.bind()
+            except TypeError as exc:
+                raise ValueError(
+                    "The function passed to a `Function` parameter must be "
+                    "callable with either 0 or 1 parameters."
+                ) from exc
+
+        self.pass_config = pass_config
+        self.func = func
+
+    def is_function(self):
+        return True
+
+    def is_valid(self, value: Any):
+        return True  # This is user-defined, so lets not assume anything
+
+    @property
+    def domain_str(self):
+        return f"{self.func}()"
+
+
+@DeveloperAPI
+class Quantized(Sampler):
+    def __init__(self, sampler: Sampler, q: Union[float, int]):
+        self.sampler = sampler
+        self.q = q
+
+        assert self.sampler, "Quantized() expects a sampler instance"
+
+    def get_sampler(self):
+        return self.sampler
+
+    def sample(
+        self,
+        domain: Domain,
+        config: Optional[Union[List[Dict], Dict]] = None,
+        size: int = 1,
+        random_state: "RandomState" = None,
+    ):
+        if not isinstance(random_state, _BackwardsCompatibleNumpyRng):
+            random_state = _BackwardsCompatibleNumpyRng(random_state)
+
+        if self.q == 1:
+            return self.sampler.sample(domain, config, size, random_state=random_state)
+
+        quantized_domain = copy(domain)
+        quantized_domain.lower = np.ceil(domain.lower / self.q) * self.q
+        quantized_domain.upper = np.floor(domain.upper / self.q) * self.q
+        values = self.sampler.sample(
+            quantized_domain, config, size, random_state=random_state
+        )
+        quantized = np.round(np.divide(values, self.q)) * self.q
+
+        if not isinstance(quantized, np.ndarray):
+            return domain.cast(quantized)
+        return list(quantized)
+
+
+@PublicAPI
+def sample_from(func: Callable[[Dict], Any]):
+    """Specify that tune should sample configuration values from this function.
+
+    Arguments:
+        func: An callable function to draw a sample from.
+    """
+    return Function(func)
+
+
+@PublicAPI
+def uniform(lower: float, upper: float):
+    """Sample a float value uniformly between ``lower`` and ``upper``.
+
+    Sampling from ``tune.uniform(1, 10)`` is equivalent to sampling from
+    ``np.random.uniform(1, 10))``
+
+    """
+    return Float(lower, upper).uniform()
+
+
+@PublicAPI
+def quniform(lower: float, upper: float, q: float):
+    """Sample a quantized float value uniformly between ``lower`` and ``upper``.
+
+    Sampling from ``tune.uniform(1, 10)`` is equivalent to sampling from
+    ``np.random.uniform(1, 10))``
+
+    The value will be quantized, i.e. rounded to an integer increment of ``q``.
+    Quantization makes the upper bound inclusive.
+
+    """
+    return Float(lower, upper).uniform().quantized(q)
+
+
+@PublicAPI
+def loguniform(lower: float, upper: float, base: float = 10):
+    """Sugar for sampling in different orders of magnitude.
+
+    Args:
+        lower: Lower boundary of the output interval (e.g. 1e-4)
+        upper: Upper boundary of the output interval (e.g. 1e-2)
+        base: Base of the log. Defaults to 10.
+
+    """
+    return Float(lower, upper).loguniform(base)
+
+
+@PublicAPI
+def qloguniform(lower: float, upper: float, q: float, base: float = 10):
+    """Sugar for sampling in different orders of magnitude.
+
+    The value will be quantized, i.e. rounded to an integer increment of ``q``.
+
+    Quantization makes the upper bound inclusive.
+
+    Args:
+        lower: Lower boundary of the output interval (e.g. 1e-4)
+        upper: Upper boundary of the output interval (e.g. 1e-2)
+        q: Quantization number. The result will be rounded to an
+            integer increment of this value.
+        base: Base of the log. Defaults to 10.
+
+    """
+    return Float(lower, upper).loguniform(base).quantized(q)
+
+
+@PublicAPI
+def choice(categories: Sequence):
+    """Sample a categorical value.
+
+    Sampling from ``tune.choice([1, 2])`` is equivalent to sampling from
+    ``np.random.choice([1, 2])``
+
+    """
+    return Categorical(categories).uniform()
+
+
+@PublicAPI
+def randint(lower: int, upper: int):
+    """Sample an integer value uniformly between ``lower`` and ``upper``.
+
+    ``lower`` is inclusive, ``upper`` is exclusive.
+
+    Sampling from ``tune.randint(10)`` is equivalent to sampling from
+    ``np.random.randint(10)``
+
+    .. versionchanged:: 1.5.0
+        When converting Ray Tune configs to searcher-specific search spaces,
+        the lower and upper limits are adjusted to keep compatibility with
+        the bounds stated in the docstring above.
+
+    """
+    return Integer(lower, upper).uniform()
+
+
+@PublicAPI
+def lograndint(lower: int, upper: int, base: float = 10):
+    """Sample an integer value log-uniformly between ``lower`` and ``upper``,
+    with ``base`` being the base of logarithm.
+
+    ``lower`` is inclusive, ``upper`` is exclusive.
+
+    .. versionchanged:: 1.5.0
+        When converting Ray Tune configs to searcher-specific search spaces,
+        the lower and upper limits are adjusted to keep compatibility with
+        the bounds stated in the docstring above.
+
+    """
+    return Integer(lower, upper).loguniform(base)
+
+
+@PublicAPI
+def qrandint(lower: int, upper: int, q: int = 1):
+    """Sample an integer value uniformly between ``lower`` and ``upper``.
+
+    ``lower`` is inclusive, ``upper`` is also inclusive (!).
+
+    The value will be quantized, i.e. rounded to an integer increment of ``q``.
+    Quantization makes the upper bound inclusive.
+
+    .. versionchanged:: 1.5.0
+        When converting Ray Tune configs to searcher-specific search spaces,
+        the lower and upper limits are adjusted to keep compatibility with
+        the bounds stated in the docstring above.
+
+    """
+    return Integer(lower, upper).uniform().quantized(q)
+
+
+@PublicAPI
+def qlograndint(lower: int, upper: int, q: int, base: float = 10):
+    """Sample an integer value log-uniformly between ``lower`` and ``upper``,
+    with ``base`` being the base of logarithm.
+
+    ``lower`` is inclusive, ``upper`` is also inclusive (!).
+
+    The value will be quantized, i.e. rounded to an integer increment of ``q``.
+    Quantization makes the upper bound inclusive.
+
+    .. versionchanged:: 1.5.0
+        When converting Ray Tune configs to searcher-specific search spaces,
+        the lower and upper limits are adjusted to keep compatibility with
+        the bounds stated in the docstring above.
+
+    """
+    return Integer(lower, upper).loguniform(base).quantized(q)
+
+
+@PublicAPI
+def randn(mean: float = 0.0, sd: float = 1.0):
+    """Sample a float value normally with ``mean`` and ``sd``.
+
+    Args:
+        mean: Mean of the normal distribution. Defaults to 0.
+        sd: SD of the normal distribution. Defaults to 1.
+
+    """
+    return Float(None, None).normal(mean, sd)
+
+
+@PublicAPI
+def qrandn(mean: float, sd: float, q: float):
+    """Sample a float value normally with ``mean`` and ``sd``.
+
+    The value will be quantized, i.e. rounded to an integer increment of ``q``.
+
+    Args:
+        mean: Mean of the normal distribution.
+        sd: SD of the normal distribution.
+        q: Quantization number. The result will be rounded to an
+            integer increment of this value.
+
+    """
+    return Float(None, None).normal(mean, sd).quantized(q)

.venv/lib/python3.11/site-packages/ray/tune/search/search_algorithm.py

@@ -0,0 +1,127 @@
+from typing import TYPE_CHECKING, Dict, List, Optional, Union
+
+from ray.util.annotations import DeveloperAPI
+
+if TYPE_CHECKING:
+    from ray.tune.experiment import Experiment
+
+
+@DeveloperAPI
+class SearchAlgorithm:
+    """Interface of an event handler API for hyperparameter search.
+
+    Unlike TrialSchedulers, SearchAlgorithms will not have the ability
+    to modify the execution (i.e., stop and pause trials).
+
+    Trials added manually (i.e., via the Client API) will also notify
+    this class upon new events, so custom search algorithms should
+    maintain a list of trials ID generated from this class.
+
+    See also: `ray.tune.search.BasicVariantGenerator`.
+    """
+
+    _finished = False
+
+    _metric = None
+
+    @property
+    def metric(self):
+        return self._metric
+
+    def set_search_properties(
+        self, metric: Optional[str], mode: Optional[str], config: Dict, **spec
+    ) -> bool:
+        """Pass search properties to search algorithm.
+
+        This method acts as an alternative to instantiating search algorithms
+        with their own specific search spaces. Instead they can accept a
+        Tune config through this method.
+
+        The search algorithm will usually pass this method to their
+        ``Searcher`` instance.
+
+        Args:
+            metric: Metric to optimize
+            mode: One of ["min", "max"]. Direction to optimize.
+            config: Tune config dict.
+            **spec: Any kwargs for forward compatiblity.
+                Info like Experiment.PUBLIC_KEYS is provided through here.
+        """
+        if self._metric and metric:
+            return False
+        if metric:
+            self._metric = metric
+        return True
+
+    @property
+    def total_samples(self):
+        """Get number of total trials to be generated"""
+        return 0
+
+    def add_configurations(
+        self, experiments: Union["Experiment", List["Experiment"], Dict[str, Dict]]
+    ):
+        """Tracks given experiment specifications.
+
+        Arguments:
+            experiments: Experiments to run.
+        """
+        raise NotImplementedError
+
+    def next_trial(self):
+        """Returns single Trial object to be queued into the TrialRunner.
+
+        Returns:
+            trial: Returns a Trial object.
+        """
+        raise NotImplementedError
+
+    def on_trial_result(self, trial_id: str, result: Dict):
+        """Called on each intermediate result returned by a trial.
+
+        This will only be called when the trial is in the RUNNING state.
+
+        Arguments:
+            trial_id: Identifier for the trial.
+            result: Result dictionary.
+        """
+        pass
+
+    def on_trial_complete(
+        self, trial_id: str, result: Optional[Dict] = None, error: bool = False
+    ):
+        """Notification for the completion of trial.
+
+        Arguments:
+            trial_id: Identifier for the trial.
+            result: Defaults to None. A dict will
+                be provided with this notification when the trial is in
+                the RUNNING state AND either completes naturally or
+                by manual termination.
+            error: Defaults to False. True if the trial is in
+                the RUNNING state and errors.
+        """
+        pass
+
+    def is_finished(self) -> bool:
+        """Returns True if no trials left to be queued into TrialRunner.
+
+        Can return True before all trials have finished executing.
+        """
+        return self._finished
+
+    def set_finished(self):
+        """Marks the search algorithm as finished."""
+        self._finished = True
+
+    def has_checkpoint(self, dirpath: str) -> bool:
+        """Should return False if restoring is not implemented."""
+        return False
+
+    def save_to_dir(self, dirpath: str, **kwargs):
+        """Saves a search algorithm."""
+        pass
+
+    def restore_from_dir(self, dirpath: str):
+        """Restores a search algorithm along with its wrapped state."""
+        pass

.venv/lib/python3.11/site-packages/ray/tune/search/search_generator.py

@@ -0,0 +1,222 @@
+import copy
+import logging
+from typing import Dict, List, Optional, Union
+
+from ray.tune.error import TuneError
+from ray.tune.experiment import Experiment, Trial, _convert_to_experiment_list
+from ray.tune.experiment.config_parser import _create_trial_from_spec, _make_parser
+from ray.tune.search.search_algorithm import SearchAlgorithm
+from ray.tune.search.searcher import Searcher
+from ray.tune.search.util import _set_search_properties_backwards_compatible
+from ray.tune.search.variant_generator import _resolve_nested_dict, format_vars
+from ray.tune.utils.util import (
+    _atomic_save,
+    _load_newest_checkpoint,
+    flatten_dict,
+    merge_dicts,
+)
+from ray.util.annotations import DeveloperAPI
+
+logger = logging.getLogger(__name__)
+
+
+def _warn_on_repeater(searcher, total_samples):
+    from ray.tune.search.repeater import _warn_num_samples
+
+    _warn_num_samples(searcher, total_samples)
+
+
+@DeveloperAPI
+class SearchGenerator(SearchAlgorithm):
+    """Generates trials to be passed to the TrialRunner.
+
+    Uses the provided ``searcher`` object to generate trials. This class
+    transparently handles repeating trials with score aggregation
+    without embedding logic into the Searcher.
+
+    Args:
+        searcher: Search object that subclasses the Searcher base class. This
+            is then used for generating new hyperparameter samples.
+    """
+
+    CKPT_FILE_TMPL = "search_gen_state-{}.json"
+
+    def __init__(self, searcher: Searcher):
+        assert issubclass(
+            type(searcher), Searcher
+        ), "Searcher should be subclassing Searcher."
+        self.searcher = searcher
+        self._parser = _make_parser()
+        self._experiment = None
+        self._counter = 0  # Keeps track of number of trials created.
+        self._total_samples = 0  # int: total samples to evaluate.
+        self._finished = False
+
+    @property
+    def metric(self):
+        return self.searcher.metric
+
+    def set_search_properties(
+        self, metric: Optional[str], mode: Optional[str], config: Dict, **spec
+    ) -> bool:
+        return _set_search_properties_backwards_compatible(
+            self.searcher.set_search_properties, metric, mode, config, **spec
+        )
+
+    @property
+    def total_samples(self):
+        return self._total_samples
+
+    def add_configurations(
+        self, experiments: Union[Experiment, List[Experiment], Dict[str, Dict]]
+    ):
+        """Registers experiment specifications.
+
+        Arguments:
+            experiments: Experiments to run.
+        """
+        assert not self._experiment
+        logger.debug("added configurations")
+        experiment_list = _convert_to_experiment_list(experiments)
+        assert (
+            len(experiment_list) == 1
+        ), "SearchAlgorithms can only support 1 experiment at a time."
+        self._experiment = experiment_list[0]
+        experiment_spec = self._experiment.spec
+        self._total_samples = self._experiment.spec.get("num_samples", 1)
+
+        _warn_on_repeater(self.searcher, self._total_samples)
+        if "run" not in experiment_spec:
+            raise TuneError("Must specify `run` in {}".format(experiment_spec))
+
+    def next_trial(self):
+        """Provides one Trial object to be queued into the TrialRunner.
+
+        Returns:
+            Trial: Returns a single trial.
+        """
+        if not self.is_finished():
+            return self.create_trial_if_possible(self._experiment.spec)
+        return None
+
+    def create_trial_if_possible(self, experiment_spec: Dict) -> Optional[Trial]:
+        logger.debug("creating trial")
+        trial_id = Trial.generate_id()
+        suggested_config = self.searcher.suggest(trial_id)
+        if suggested_config == Searcher.FINISHED:
+            self._finished = True
+            logger.debug("Searcher has finished.")
+            return
+
+        if suggested_config is None:
+            return
+        spec = copy.deepcopy(experiment_spec)
+        spec["config"] = merge_dicts(spec["config"], copy.deepcopy(suggested_config))
+
+        # Create a new trial_id if duplicate trial is created
+        flattened_config = _resolve_nested_dict(spec["config"])
+        self._counter += 1
+        tag = "{0}_{1}".format(str(self._counter), format_vars(flattened_config))
+        trial = _create_trial_from_spec(
+            spec,
+            self._parser,
+            evaluated_params=flatten_dict(suggested_config),
+            experiment_tag=tag,
+            trial_id=trial_id,
+        )
+        return trial
+
+    def on_trial_result(self, trial_id: str, result: Dict):
+        """Notifies the underlying searcher."""
+        self.searcher.on_trial_result(trial_id, result)
+
+    def on_trial_complete(
+        self, trial_id: str, result: Optional[Dict] = None, error: bool = False
+    ):
+        self.searcher.on_trial_complete(trial_id=trial_id, result=result, error=error)
+
+    def is_finished(self) -> bool:
+        return self._counter >= self._total_samples or self._finished
+
+    def get_state(self) -> Dict:
+        return {
+            "counter": self._counter,
+            "total_samples": self._total_samples,
+            "finished": self._finished,
+            "experiment": self._experiment,
+        }
+
+    def set_state(self, state: Dict):
+        self._counter = state["counter"]
+        self._total_samples = state["total_samples"]
+        self._finished = state["finished"]
+        self._experiment = state["experiment"]
+
+    def has_checkpoint(self, dirpath: str):
+        return bool(_load_newest_checkpoint(dirpath, self.CKPT_FILE_TMPL.format("*")))
+
+    def save_to_dir(self, dirpath: str, session_str: str):
+        """Saves self + searcher to dir.
+
+        Separates the "searcher" from its wrappers (concurrency, repeating).
+        This allows the user to easily restore a given searcher.
+
+        The save operation is atomic (write/swap).
+
+        Args:
+            dirpath: Filepath to experiment dir.
+            session_str: Unique identifier of the current run
+                session.
+        """
+        searcher = self.searcher
+        search_alg_state = self.get_state()
+        while hasattr(searcher, "searcher"):
+            searcher_name = type(searcher).__name__
+            if searcher_name in search_alg_state:
+                logger.warning(
+                    "There was a duplicate when saving {}. "
+                    "Restore may not work properly.".format(searcher_name)
+                )
+            else:
+                search_alg_state["name:" + searcher_name] = searcher.get_state()
+            searcher = searcher.searcher
+        base_searcher = searcher
+        # We save the base searcher separately for users to easily
+        # separate the searcher.
+        base_searcher.save_to_dir(dirpath, session_str)
+        _atomic_save(
+            state=search_alg_state,
+            checkpoint_dir=dirpath,
+            file_name=self.CKPT_FILE_TMPL.format(session_str),
+            tmp_file_name=".tmp_search_generator_ckpt",
+        )
+
+    def restore_from_dir(self, dirpath: str):
+        """Restores self + searcher + search wrappers from dirpath."""
+
+        searcher = self.searcher
+        search_alg_state = _load_newest_checkpoint(
+            dirpath, self.CKPT_FILE_TMPL.format("*")
+        )
+        if not search_alg_state:
+            raise RuntimeError("Unable to find checkpoint in {}.".format(dirpath))
+        while hasattr(searcher, "searcher"):
+            searcher_name = "name:" + type(searcher).__name__
+            if searcher_name not in search_alg_state:
+                names = [
+                    key.split("name:")[1]
+                    for key in search_alg_state
+                    if key.startswith("name:")
+                ]
+                logger.warning(
+                    "{} was not found in the experiment "
+                    "state when restoring. Found {}.".format(searcher_name, names)
+                )
+            else:
+                searcher.set_state(search_alg_state.pop(searcher_name))
+            searcher = searcher.searcher
+        base_searcher = searcher
+
+        logger.debug(f"searching base {base_searcher}")
+        base_searcher.restore_from_dir(dirpath)
+        self.set_state(search_alg_state)

.venv/lib/python3.11/site-packages/ray/tune/search/searcher.py

@@ -0,0 +1,597 @@
+import copy
+import glob
+import logging
+import os
+import warnings
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+
+from ray.air._internal.usage import tag_searcher
+from ray.tune.search.util import _set_search_properties_backwards_compatible
+from ray.util.annotations import DeveloperAPI, PublicAPI
+from ray.util.debug import log_once
+
+if TYPE_CHECKING:
+    from ray.tune.analysis import ExperimentAnalysis
+    from ray.tune.experiment import Trial
+
+logger = logging.getLogger(__name__)
+
+
+@DeveloperAPI
+class Searcher:
+    """Abstract class for wrapping suggesting algorithms.
+
+    Custom algorithms can extend this class easily by overriding the
+    `suggest` method provide generated parameters for the trials.
+
+    Any subclass that implements ``__init__`` must also call the
+    constructor of this class: ``super(Subclass, self).__init__(...)``.
+
+    To track suggestions and their corresponding evaluations, the method
+    `suggest` will be passed a trial_id, which will be used in
+    subsequent notifications.
+
+    Not all implementations support multi objectives.
+
+    Note to Tune developers: If a new searcher is added, please update
+    `air/_internal/usage.py`.
+
+    Args:
+        metric: The training result objective value attribute. If
+            list then list of training result objective value attributes
+        mode: If string One of {min, max}. If list then
+            list of max and min, determines whether objective is minimizing
+            or maximizing the metric attribute. Must match type of metric.
+
+    .. code-block:: python
+
+        class ExampleSearch(Searcher):
+            def __init__(self, metric="mean_loss", mode="min", **kwargs):
+                super(ExampleSearch, self).__init__(
+                    metric=metric, mode=mode, **kwargs)
+                self.optimizer = Optimizer()
+                self.configurations = {}
+
+            def suggest(self, trial_id):
+                configuration = self.optimizer.query()
+                self.configurations[trial_id] = configuration
+
+            def on_trial_complete(self, trial_id, result, **kwargs):
+                configuration = self.configurations[trial_id]
+                if result and self.metric in result:
+                    self.optimizer.update(configuration, result[self.metric])
+
+        tuner = tune.Tuner(
+            trainable_function,
+            tune_config=tune.TuneConfig(
+                search_alg=ExampleSearch()
+            )
+        )
+        tuner.fit()
+
+
+    """
+
+    FINISHED = "FINISHED"
+    CKPT_FILE_TMPL = "searcher-state-{}.pkl"
+
+    def __init__(
+        self,
+        metric: Optional[str] = None,
+        mode: Optional[str] = None,
+    ):
+        tag_searcher(self)
+        self._metric = metric
+        self._mode = mode
+
+        if not mode or not metric:
+            # Early return to avoid assertions
+            return
+
+        assert isinstance(
+            metric, type(mode)
+        ), "metric and mode must be of the same type"
+        if isinstance(mode, str):
+            assert mode in ["min", "max"], "if `mode` is a str must be 'min' or 'max'!"
+        elif isinstance(mode, list):
+            assert len(mode) == len(metric), "Metric and mode must be the same length"
+            assert all(
+                mod in ["min", "max", "obs"] for mod in mode
+            ), "All of mode must be 'min' or 'max' or 'obs'!"
+        else:
+            raise ValueError("Mode most either be a list or string")
+
+    def set_search_properties(
+        self, metric: Optional[str], mode: Optional[str], config: Dict, **spec
+    ) -> bool:
+        """Pass search properties to searcher.
+
+        This method acts as an alternative to instantiating search algorithms
+        with their own specific search spaces. Instead they can accept a
+        Tune config through this method. A searcher should return ``True``
+        if setting the config was successful, or ``False`` if it was
+        unsuccessful, e.g. when the search space has already been set.
+
+        Args:
+            metric: Metric to optimize
+            mode: One of ["min", "max"]. Direction to optimize.
+            config: Tune config dict.
+            **spec: Any kwargs for forward compatiblity.
+                Info like Experiment.PUBLIC_KEYS is provided through here.
+        """
+        return False
+
+    def on_trial_result(self, trial_id: str, result: Dict) -> None:
+        """Optional notification for result during training.
+
+        Note that by default, the result dict may include NaNs or
+        may not include the optimization metric. It is up to the
+        subclass implementation to preprocess the result to
+        avoid breaking the optimization process.
+
+        Args:
+            trial_id: A unique string ID for the trial.
+            result: Dictionary of metrics for current training progress.
+                Note that the result dict may include NaNs or
+                may not include the optimization metric. It is up to the
+                subclass implementation to preprocess the result to
+                avoid breaking the optimization process.
+        """
+        pass
+
+    def on_trial_complete(
+        self, trial_id: str, result: Optional[Dict] = None, error: bool = False
+    ) -> None:
+        """Notification for the completion of trial.
+
+        Typically, this method is used for notifying the underlying
+        optimizer of the result.
+
+        Args:
+            trial_id: A unique string ID for the trial.
+            result: Dictionary of metrics for current training progress.
+                Note that the result dict may include NaNs or
+                may not include the optimization metric. It is up to the
+                subclass implementation to preprocess the result to
+                avoid breaking the optimization process. Upon errors, this
+                may also be None.
+            error: True if the training process raised an error.
+
+        """
+        raise NotImplementedError
+
+    def suggest(self, trial_id: str) -> Optional[Dict]:
+        """Queries the algorithm to retrieve the next set of parameters.
+
+        Arguments:
+            trial_id: Trial ID used for subsequent notifications.
+
+        Returns:
+            dict | FINISHED | None: Configuration for a trial, if possible.
+                If FINISHED is returned, Tune will be notified that
+                no more suggestions/configurations will be provided.
+                If None is returned, Tune will skip the querying of the
+                searcher for this step.
+
+        """
+        raise NotImplementedError
+
+    def add_evaluated_point(
+        self,
+        parameters: Dict,
+        value: float,
+        error: bool = False,
+        pruned: bool = False,
+        intermediate_values: Optional[List[float]] = None,
+    ):
+        """Pass results from a point that has been evaluated separately.
+
+        This method allows for information from outside the
+        suggest - on_trial_complete loop to be passed to the search
+        algorithm.
+        This functionality depends on the underlying search algorithm
+        and may not be always available.
+
+        Args:
+            parameters: Parameters used for the trial.
+            value: Metric value obtained in the trial.
+            error: True if the training process raised an error.
+            pruned: True if trial was pruned.
+            intermediate_values: List of metric values for
+                intermediate iterations of the result. None if not
+                applicable.
+
+        """
+        raise NotImplementedError
+
+    def add_evaluated_trials(
+        self,
+        trials_or_analysis: Union["Trial", List["Trial"], "ExperimentAnalysis"],
+        metric: str,
+    ):
+        """Pass results from trials that have been evaluated separately.
+
+        This method allows for information from outside the
+        suggest - on_trial_complete loop to be passed to the search
+        algorithm.
+        This functionality depends on the underlying search algorithm
+        and may not be always available (same as ``add_evaluated_point``.)
+
+        Args:
+            trials_or_analysis: Trials to pass results form to the searcher.
+            metric: Metric name reported by trials used for
+                determining the objective value.
+
+        """
+        if self.add_evaluated_point == Searcher.add_evaluated_point:
+            raise NotImplementedError
+
+        # lazy imports to avoid circular dependencies
+        from ray.tune.analysis import ExperimentAnalysis
+        from ray.tune.experiment import Trial
+        from ray.tune.result import DONE
+
+        if isinstance(trials_or_analysis, (list, tuple)):
+            trials = trials_or_analysis
+        elif isinstance(trials_or_analysis, Trial):
+            trials = [trials_or_analysis]
+        elif isinstance(trials_or_analysis, ExperimentAnalysis):
+            trials = trials_or_analysis.trials
+        else:
+            raise NotImplementedError(
+                "Expected input to be a `Trial`, a list of `Trial`s, or "
+                f"`ExperimentAnalysis`, got: {trials_or_analysis}"
+            )
+
+        any_trial_had_metric = False
+
+        def trial_to_points(trial: Trial) -> Dict[str, Any]:
+            nonlocal any_trial_had_metric
+            has_trial_been_pruned = (
+                trial.status == Trial.TERMINATED
+                and not trial.last_result.get(DONE, False)
+            )
+            has_trial_finished = (
+                trial.status == Trial.TERMINATED and trial.last_result.get(DONE, False)
+            )
+            if not any_trial_had_metric:
+                any_trial_had_metric = (
+                    metric in trial.last_result and has_trial_finished
+                )
+            if Trial.TERMINATED and metric not in trial.last_result:
+                return None
+            return dict(
+                parameters=trial.config,
+                value=trial.last_result.get(metric, None),
+                error=trial.status == Trial.ERROR,
+                pruned=has_trial_been_pruned,
+                intermediate_values=None,  # we do not save those
+            )
+
+        for trial in trials:
+            kwargs = trial_to_points(trial)
+            if kwargs:
+                self.add_evaluated_point(**kwargs)
+
+        if not any_trial_had_metric:
+            warnings.warn(
+                "No completed trial returned the specified metric. "
+                "Make sure the name you have passed is correct. "
+            )
+
+    def save(self, checkpoint_path: str):
+        """Save state to path for this search algorithm.
+
+        Args:
+            checkpoint_path: File where the search algorithm
+                state is saved. This path should be used later when
+                restoring from file.
+
+        Example:
+
+        .. code-block:: python
+
+            search_alg = Searcher(...)
+
+            tuner = tune.Tuner(
+                cost,
+                tune_config=tune.TuneConfig(
+                    search_alg=search_alg,
+                    num_samples=5
+                ),
+                param_space=config
+            )
+            results = tuner.fit()
+
+            search_alg.save("./my_favorite_path.pkl")
+
+        .. versionchanged:: 0.8.7
+            Save is automatically called by `Tuner().fit()`. You can use
+            `Tuner().restore()` to restore from an experiment directory
+            such as `~/ray_results/trainable`.
+
+        """
+        raise NotImplementedError
+
+    def restore(self, checkpoint_path: str):
+        """Restore state for this search algorithm
+
+
+        Args:
+            checkpoint_path: File where the search algorithm
+                state is saved. This path should be the same
+                as the one provided to "save".
+
+        Example:
+
+        .. code-block:: python
+
+            search_alg.save("./my_favorite_path.pkl")
+
+            search_alg2 = Searcher(...)
+            search_alg2 = ConcurrencyLimiter(search_alg2, 1)
+            search_alg2.restore(checkpoint_path)
+            tuner = tune.Tuner(
+                cost,
+                tune_config=tune.TuneConfig(
+                    search_alg=search_alg2,
+                    num_samples=5
+                ),
+            )
+            tuner.fit()
+
+        """
+        raise NotImplementedError
+
+    def set_max_concurrency(self, max_concurrent: int) -> bool:
+        """Set max concurrent trials this searcher can run.
+
+        This method will be called on the wrapped searcher by the
+        ``ConcurrencyLimiter``. It is intended to allow for searchers
+        which have custom, internal logic handling max concurrent trials
+        to inherit the value passed to ``ConcurrencyLimiter``.
+
+        If this method returns False, it signifies that no special
+        logic for handling this case is present in the searcher.
+
+        Args:
+            max_concurrent: Number of maximum concurrent trials.
+        """
+        return False
+
+    def get_state(self) -> Dict:
+        raise NotImplementedError
+
+    def set_state(self, state: Dict):
+        raise NotImplementedError
+
+    def save_to_dir(self, checkpoint_dir: str, session_str: str = "default"):
+        """Automatically saves the given searcher to the checkpoint_dir.
+
+        This is automatically used by Tuner().fit() during a Tune job.
+
+        Args:
+            checkpoint_dir: Filepath to experiment dir.
+            session_str: Unique identifier of the current run
+                session.
+        """
+        tmp_search_ckpt_path = os.path.join(checkpoint_dir, ".tmp_searcher_ckpt")
+        success = True
+        try:
+            self.save(tmp_search_ckpt_path)
+        except NotImplementedError:
+            if log_once("suggest:save_to_dir"):
+                logger.warning("save not implemented for Searcher. Skipping save.")
+            success = False
+
+        if success and os.path.exists(tmp_search_ckpt_path):
+            os.replace(
+                tmp_search_ckpt_path,
+                os.path.join(checkpoint_dir, self.CKPT_FILE_TMPL.format(session_str)),
+            )
+
+    def restore_from_dir(self, checkpoint_dir: str):
+        """Restores the state of a searcher from a given checkpoint_dir.
+
+        Typically, you should use this function to restore from an
+        experiment directory such as `~/ray_results/trainable`.
+
+        .. code-block:: python
+
+            tuner = tune.Tuner(
+                cost,
+                run_config=train.RunConfig(
+                    name=self.experiment_name,
+                    storage_path="~/my_results",
+                ),
+                tune_config=tune.TuneConfig(
+                    search_alg=search_alg,
+                    num_samples=5
+                ),
+                param_space=config
+            )
+            tuner.fit()
+
+            search_alg2 = Searcher()
+            search_alg2.restore_from_dir(
+                os.path.join("~/my_results", self.experiment_name)
+        """
+
+        pattern = self.CKPT_FILE_TMPL.format("*")
+        full_paths = glob.glob(os.path.join(checkpoint_dir, pattern))
+        if not full_paths:
+            raise RuntimeError(
+                "Searcher unable to find checkpoint in {}".format(checkpoint_dir)
+            )  # TODO
+        most_recent_checkpoint = max(full_paths)
+        self.restore(most_recent_checkpoint)
+
+    @property
+    def metric(self) -> str:
+        """The training result objective value attribute."""
+        return self._metric
+
+    @property
+    def mode(self) -> str:
+        """Specifies if minimizing or maximizing the metric."""
+        return self._mode
+
+
+@PublicAPI
+class ConcurrencyLimiter(Searcher):
+    """A wrapper algorithm for limiting the number of concurrent trials.
+
+    Certain Searchers have their own internal logic for limiting
+    the number of concurrent trials. If such a Searcher is passed to a
+    ``ConcurrencyLimiter``, the ``max_concurrent`` of the
+    ``ConcurrencyLimiter`` will override the ``max_concurrent`` value
+    of the Searcher. The ``ConcurrencyLimiter`` will then let the
+    Searcher's internal logic take over.
+
+    Args:
+        searcher: Searcher object that the
+            ConcurrencyLimiter will manage.
+        max_concurrent: Maximum concurrent samples from the underlying
+            searcher.
+        batch: Whether to wait for all concurrent samples
+            to finish before updating the underlying searcher.
+
+    Example:
+
+    .. code-block:: python
+
+        from ray.tune.search import ConcurrencyLimiter
+        search_alg = HyperOptSearch(metric="accuracy")
+        search_alg = ConcurrencyLimiter(search_alg, max_concurrent=2)
+        tuner = tune.Tuner(
+            trainable_function,
+            tune_config=tune.TuneConfig(
+                search_alg=search_alg
+            ),
+        )
+        tuner.fit()
+
+    """
+
+    def __init__(self, searcher: Searcher, max_concurrent: int, batch: bool = False):
+        assert type(max_concurrent) is int and max_concurrent > 0
+        self.searcher = searcher
+        self.max_concurrent = max_concurrent
+        self.batch = batch
+        self.live_trials = set()
+        self.num_unfinished_live_trials = 0
+        self.cached_results = {}
+        self._limit_concurrency = True
+
+        if not isinstance(searcher, Searcher):
+            raise RuntimeError(
+                f"The `ConcurrencyLimiter` only works with `Searcher` "
+                f"objects (got {type(searcher)}). Please try to pass "
+                f"`max_concurrent` to the search generator directly."
+            )
+
+        self._set_searcher_max_concurrency()
+
+        super(ConcurrencyLimiter, self).__init__(
+            metric=self.searcher.metric, mode=self.searcher.mode
+        )
+
+    def _set_searcher_max_concurrency(self):
+        # If the searcher has special logic for handling max concurrency,
+        # we do not do anything inside the ConcurrencyLimiter
+        self._limit_concurrency = not self.searcher.set_max_concurrency(
+            self.max_concurrent
+        )
+
+    def set_max_concurrency(self, max_concurrent: int) -> bool:
+        # Determine if this behavior is acceptable, or if it should
+        # raise an exception.
+        self.max_concurrent = max_concurrent
+        return True
+
+    def set_search_properties(
+        self, metric: Optional[str], mode: Optional[str], config: Dict, **spec
+    ) -> bool:
+        self._set_searcher_max_concurrency()
+        return _set_search_properties_backwards_compatible(
+            self.searcher.set_search_properties, metric, mode, config, **spec
+        )
+
+    def suggest(self, trial_id: str) -> Optional[Dict]:
+        if not self._limit_concurrency:
+            return self.searcher.suggest(trial_id)
+
+        assert (
+            trial_id not in self.live_trials
+        ), f"Trial ID {trial_id} must be unique: already found in set."
+        if len(self.live_trials) >= self.max_concurrent:
+            logger.debug(
+                f"Not providing a suggestion for {trial_id} due to "
+                "concurrency limit: %s/%s.",
+                len(self.live_trials),
+                self.max_concurrent,
+            )
+            return
+
+        suggestion = self.searcher.suggest(trial_id)
+        if suggestion not in (None, Searcher.FINISHED):
+            self.live_trials.add(trial_id)
+            self.num_unfinished_live_trials += 1
+        return suggestion
+
+    def on_trial_complete(
+        self, trial_id: str, result: Optional[Dict] = None, error: bool = False
+    ):
+        if not self._limit_concurrency:
+            return self.searcher.on_trial_complete(trial_id, result=result, error=error)
+
+        if trial_id not in self.live_trials:
+            return
+        elif self.batch:
+            self.cached_results[trial_id] = (result, error)
+            self.num_unfinished_live_trials -= 1
+            if self.num_unfinished_live_trials <= 0:
+                # Update the underlying searcher once the
+                # full batch is completed.
+                for trial_id, (result, error) in self.cached_results.items():
+                    self.searcher.on_trial_complete(
+                        trial_id, result=result, error=error
+                    )
+                    self.live_trials.remove(trial_id)
+                self.cached_results = {}
+                self.num_unfinished_live_trials = 0
+            else:
+                return
+        else:
+            self.searcher.on_trial_complete(trial_id, result=result, error=error)
+            self.live_trials.remove(trial_id)
+            self.num_unfinished_live_trials -= 1
+
+    def on_trial_result(self, trial_id: str, result: Dict) -> None:
+        self.searcher.on_trial_result(trial_id, result)
+
+    def add_evaluated_point(
+        self,
+        parameters: Dict,
+        value: float,
+        error: bool = False,
+        pruned: bool = False,
+        intermediate_values: Optional[List[float]] = None,
+    ):
+        return self.searcher.add_evaluated_point(
+            parameters, value, error, pruned, intermediate_values
+        )
+
+    def get_state(self) -> Dict:
+        state = self.__dict__.copy()
+        del state["searcher"]
+        return copy.deepcopy(state)
+
+    def set_state(self, state: Dict):
+        self.__dict__.update(state)
+
+    def save(self, checkpoint_path: str):
+        self.searcher.save(checkpoint_path)
+
+    def restore(self, checkpoint_path: str):
+        self.searcher.restore(checkpoint_path)

.venv/lib/python3.11/site-packages/ray/tune/search/util.py

@@ -0,0 +1,31 @@
+import logging
+from typing import Dict, Optional
+
+logger = logging.getLogger(__name__)
+
+
+def _set_search_properties_backwards_compatible(
+    set_search_properties_func,
+    metric: Optional[str],
+    mode: Optional[str],
+    config: Dict,
+    **spec
+) -> bool:
+    """Wraps around set_search_properties() so that it is backward compatible.
+
+    Also outputs a warning to encourage custom searchers to be updated.
+    """
+    try:
+        return set_search_properties_func(metric, mode, config, **spec)
+    except TypeError as e:
+        if str(e).startswith(
+            "set_search_properties() got an unexpected keyword argument"
+        ):
+            logger.warning(
+                "Please update custom Searcher to take in function signature "
+                "as ``def set_search_properties(metric, mode, config, "
+                "**spec) -> bool``."
+            )
+            return set_search_properties_func(metric, mode, config)
+        else:
+            raise e

.venv/lib/python3.11/site-packages/ray/tune/search/variant_generator.py

@@ -0,0 +1,523 @@
+import copy
+import logging
+import random
+import re
+from collections.abc import Mapping
+from typing import Any, Dict, Generator, Iterable, List, Optional, Tuple
+
+import numpy
+
+from ray.tune.search.sample import Categorical, Domain, Function, RandomState
+from ray.util.annotations import DeveloperAPI, PublicAPI
+
+logger = logging.getLogger(__name__)
+
+
+@DeveloperAPI
+def generate_variants(
+    unresolved_spec: Dict,
+    constant_grid_search: bool = False,
+    random_state: "RandomState" = None,
+) -> Generator[Tuple[Dict, Dict], None, None]:
+    """Generates variants from a spec (dict) with unresolved values.
+
+    There are two types of unresolved values:
+
+        Grid search: These define a grid search over values. For example, the
+        following grid search values in a spec will produce six distinct
+        variants in combination:
+
+            "activation": grid_search(["relu", "tanh"])
+            "learning_rate": grid_search([1e-3, 1e-4, 1e-5])
+
+        Lambda functions: These are evaluated to produce a concrete value, and
+        can express dependencies or conditional distributions between values.
+        They can also be used to express random search (e.g., by calling
+        into the `random` or `np` module).
+
+            "cpu": lambda spec: spec.config.num_workers
+            "batch_size": lambda spec: random.uniform(1, 1000)
+
+    Finally, to support defining specs in plain JSON / YAML, grid search
+    and lambda functions can also be defined alternatively as follows:
+
+        "activation": {"grid_search": ["relu", "tanh"]}
+        "cpu": {"eval": "spec.config.num_workers"}
+
+    Use `format_vars` to format the returned dict of hyperparameters.
+
+    Yields:
+        (Dict of resolved variables, Spec object)
+    """
+    for resolved_vars, spec in _generate_variants_internal(
+        unresolved_spec,
+        constant_grid_search=constant_grid_search,
+        random_state=random_state,
+    ):
+        assert not _unresolved_values(spec)
+        yield resolved_vars, spec
+
+
+@PublicAPI(stability="beta")
+def grid_search(values: Iterable) -> Dict[str, Iterable]:
+    """Specify a grid of values to search over.
+
+    Values specified in a grid search are guaranteed to be sampled.
+
+    If multiple grid search variables are defined, they are combined with the
+    combinatorial product. This means every possible combination of values will
+    be sampled.
+
+    Example:
+
+        >>> from ray import tune
+        >>> param_space={
+        ...   "x": tune.grid_search([10, 20]),
+        ...   "y": tune.grid_search(["a", "b", "c"])
+        ... }
+
+    This will create a grid of 6 samples:
+    ``{"x": 10, "y": "a"}``, ``{"x": 10, "y": "b"}``, etc.
+
+    When specifying ``num_samples`` in the
+    :class:`TuneConfig <ray.tune.tune_config.TuneConfig>`, this will specify
+    the number of random samples per grid search combination.
+
+    For instance, in the example above, if ``num_samples=4``,
+    a total of 24 trials will be started -
+    4 trials for each of the 6 grid search combinations.
+
+    Args:
+        values: An iterable whose parameters will be used for creating a trial grid.
+
+    """
+    return {"grid_search": values}
+
+
+_STANDARD_IMPORTS = {
+    "random": random,
+    "np": numpy,
+}
+
+_MAX_RESOLUTION_PASSES = 20
+
+
+def _resolve_nested_dict(nested_dict: Dict) -> Dict[Tuple, Any]:
+    """Flattens a nested dict by joining keys into tuple of paths.
+
+    Can then be passed into `format_vars`.
+    """
+    res = {}
+    for k, v in nested_dict.items():
+        if isinstance(v, dict):
+            for k_, v_ in _resolve_nested_dict(v).items():
+                res[(k,) + k_] = v_
+        else:
+            res[(k,)] = v
+    return res
+
+
+@DeveloperAPI
+def format_vars(resolved_vars: Dict) -> str:
+    """Format variables to be used as experiment tags.
+
+    Experiment tags are used in directory names, so this method makes sure
+    the resulting tags can be legally used in directory names on all systems.
+
+    The input to this function is a dict of the form
+    ``{("nested", "config", "path"): "value"}``. The output will be a comma
+    separated string of the form ``last_key=value``, so in this example
+    ``path=value``.
+
+    Note that the sanitizing implies that empty strings are possible return
+    values. This is expected and acceptable, as it is not a common case and
+    the resulting directory names will still be valid.
+
+    Args:
+        resolved_vars: Dictionary mapping from config path tuples to a value.
+
+    Returns:
+        Comma-separated key=value string.
+    """
+    vars = resolved_vars.copy()
+    # TrialRunner already has these in the experiment_tag
+    for v in ["run", "env", "resources_per_trial"]:
+        vars.pop(v, None)
+
+    return ",".join(
+        f"{_clean_value(k[-1])}={_clean_value(v)}" for k, v in sorted(vars.items())
+    )
+
+
+def _flatten_resolved_vars(resolved_vars: Dict) -> Dict:
+    """Formats the resolved variable dict into a mapping of (str -> value)."""
+    flattened_resolved_vars_dict = {}
+    for pieces, value in resolved_vars.items():
+        if pieces[0] == "config":
+            pieces = pieces[1:]
+        pieces = [str(piece) for piece in pieces]
+        flattened_resolved_vars_dict["/".join(pieces)] = value
+    return flattened_resolved_vars_dict
+
+
+def _clean_value(value: Any) -> str:
+    """Format floats and replace invalid string characters with ``_``."""
+    if isinstance(value, float):
+        return f"{value:.4f}"
+    else:
+        # Define an invalid alphabet, which is the inverse of the
+        # stated regex characters
+        invalid_alphabet = r"[^a-zA-Z0-9_-]+"
+        return re.sub(invalid_alphabet, "_", str(value)).strip("_")
+
+
+@DeveloperAPI
+def parse_spec_vars(
+    spec: Dict,
+) -> Tuple[List[Tuple[Tuple, Any]], List[Tuple[Tuple, Any]], List[Tuple[Tuple, Any]]]:
+    resolved, unresolved = _split_resolved_unresolved_values(spec)
+    resolved_vars = list(resolved.items())
+
+    if not unresolved:
+        return resolved_vars, [], []
+
+    grid_vars = []
+    domain_vars = []
+    for path, value in unresolved.items():
+        if value.is_grid():
+            grid_vars.append((path, value))
+        else:
+            domain_vars.append((path, value))
+    grid_vars.sort()
+
+    return resolved_vars, domain_vars, grid_vars
+
+
+def _count_spec_samples(spec: Dict, num_samples=1) -> int:
+    """Count samples for a specific spec"""
+    _, domain_vars, grid_vars = parse_spec_vars(spec)
+    grid_count = 1
+    for path, domain in grid_vars:
+        grid_count *= len(domain.categories)
+    return num_samples * grid_count
+
+
+def _count_variants(spec: Dict, presets: Optional[List[Dict]] = None) -> int:
+    # Helper function: Deep update dictionary
+    def deep_update(d, u):
+        for k, v in u.items():
+            if isinstance(v, Mapping):
+                d[k] = deep_update(d.get(k, {}), v)
+            else:
+                d[k] = v
+        return d
+
+    total_samples = 0
+    total_num_samples = spec.get("num_samples", 1)
+    # For each preset, overwrite the spec and count the samples generated
+    # for this preset
+    for preset in presets:
+        preset_spec = copy.deepcopy(spec)
+        deep_update(preset_spec["config"], preset)
+        total_samples += _count_spec_samples(preset_spec, 1)
+        total_num_samples -= 1
+
+    # Add the remaining samples
+    if total_num_samples > 0:
+        total_samples += _count_spec_samples(spec, total_num_samples)
+    return total_samples
+
+
+def _generate_variants_internal(
+    spec: Dict, constant_grid_search: bool = False, random_state: "RandomState" = None
+) -> Tuple[Dict, Dict]:
+    spec = copy.deepcopy(spec)
+    _, domain_vars, grid_vars = parse_spec_vars(spec)
+
+    if not domain_vars and not grid_vars:
+        yield {}, spec
+        return
+
+    # Variables to resolve
+    to_resolve = domain_vars
+
+    all_resolved = True
+    if constant_grid_search:
+        # In this path, we first sample random variables and keep them constant
+        # for grid search.
+        # `_resolve_domain_vars` will alter `spec` directly
+        all_resolved, resolved_vars = _resolve_domain_vars(
+            spec, domain_vars, allow_fail=True, random_state=random_state
+        )
+        if not all_resolved:
+            # Not all variables have been resolved, but remove those that have
+            # from the `to_resolve` list.
+            to_resolve = [(r, d) for r, d in to_resolve if r not in resolved_vars]
+    grid_search = _grid_search_generator(spec, grid_vars)
+    for resolved_spec in grid_search:
+        if not constant_grid_search or not all_resolved:
+            # In this path, we sample the remaining random variables
+            _, resolved_vars = _resolve_domain_vars(
+                resolved_spec, to_resolve, random_state=random_state
+            )
+
+        for resolved, spec in _generate_variants_internal(
+            resolved_spec,
+            constant_grid_search=constant_grid_search,
+            random_state=random_state,
+        ):
+            for path, value in grid_vars:
+                resolved_vars[path] = _get_value(spec, path)
+            for k, v in resolved.items():
+                if (
+                    k in resolved_vars
+                    and v != resolved_vars[k]
+                    and _is_resolved(resolved_vars[k])
+                ):
+                    raise ValueError(
+                        "The variable `{}` could not be unambiguously "
+                        "resolved to a single value. Consider simplifying "
+                        "your configuration.".format(k)
+                    )
+                resolved_vars[k] = v
+            yield resolved_vars, spec
+
+
+def _get_preset_variants(
+    spec: Dict,
+    config: Dict,
+    constant_grid_search: bool = False,
+    random_state: "RandomState" = None,
+):
+    """Get variants according to a spec, initialized with a config.
+
+    Variables from the spec are overwritten by the variables in the config.
+    Thus, we may end up with less sampled parameters.
+
+    This function also checks if values used to overwrite search space
+    parameters are valid, and logs a warning if not.
+    """
+    spec = copy.deepcopy(spec)
+
+    resolved, _, _ = parse_spec_vars(config)
+
+    for path, val in resolved:
+        try:
+            domain = _get_value(spec["config"], path)
+            if isinstance(domain, dict):
+                if "grid_search" in domain:
+                    domain = Categorical(domain["grid_search"])
+                else:
+                    # If users want to overwrite an entire subdict,
+                    # let them do it.
+                    domain = None
+        except IndexError as exc:
+            raise ValueError(
+                f"Pre-set config key `{'/'.join(path)}` does not correspond "
+                f"to a valid key in the search space definition. Please add "
+                f"this path to the `param_space` variable passed to `tune.Tuner()`."
+            ) from exc
+
+        if domain:
+            if isinstance(domain, Domain):
+                if not domain.is_valid(val):
+                    logger.warning(
+                        f"Pre-set value `{val}` is not within valid values of "
+                        f"parameter `{'/'.join(path)}`: {domain.domain_str}"
+                    )
+            else:
+                # domain is actually a fixed value
+                if domain != val:
+                    logger.warning(
+                        f"Pre-set value `{val}` is not equal to the value of "
+                        f"parameter `{'/'.join(path)}`: {domain}"
+                    )
+        assign_value(spec["config"], path, val)
+
+    return _generate_variants_internal(
+        spec, constant_grid_search=constant_grid_search, random_state=random_state
+    )
+
+
+@DeveloperAPI
+def assign_value(spec: Dict, path: Tuple, value: Any):
+    """Assigns a value to a nested dictionary.
+
+    Handles the special case of tuples, in which case the tuples
+    will be re-constructed to accomodate the updated value.
+    """
+    parent_spec = None
+    parent_key = None
+    for k in path[:-1]:
+        parent_spec = spec
+        parent_key = k
+        spec = spec[k]
+    key = path[-1]
+    if not isinstance(spec, tuple):
+        # spec is mutable. Just assign the value.
+        spec[key] = value
+    else:
+        if parent_spec is None:
+            raise ValueError("Cannot assign value to a tuple.")
+        assert isinstance(key, int), "Tuple key must be an int."
+        # Special handling since tuples are immutable.
+        parent_spec[parent_key] = spec[:key] + (value,) + spec[key + 1 :]
+
+
+def _get_value(spec: Dict, path: Tuple) -> Any:
+    for k in path:
+        spec = spec[k]
+    return spec
+
+
+def _resolve_domain_vars(
+    spec: Dict,
+    domain_vars: List[Tuple[Tuple, Domain]],
+    allow_fail: bool = False,
+    random_state: "RandomState" = None,
+) -> Tuple[bool, Dict]:
+    resolved = {}
+    error = True
+    num_passes = 0
+    while error and num_passes < _MAX_RESOLUTION_PASSES:
+        num_passes += 1
+        error = False
+        for path, domain in domain_vars:
+            if path in resolved:
+                continue
+            try:
+                value = domain.sample(
+                    _UnresolvedAccessGuard(spec), random_state=random_state
+                )
+            except RecursiveDependencyError as e:
+                error = e
+            except Exception:
+                raise ValueError(
+                    "Failed to evaluate expression: {}: {}".format(path, domain)
+                )
+            else:
+                assign_value(spec, path, value)
+                resolved[path] = value
+    if error:
+        if not allow_fail:
+            raise error
+        else:
+            return False, resolved
+    return True, resolved
+
+
+def _grid_search_generator(
+    unresolved_spec: Dict, grid_vars: List
+) -> Generator[Dict, None, None]:
+    value_indices = [0] * len(grid_vars)
+
+    def increment(i):
+        value_indices[i] += 1
+        if value_indices[i] >= len(grid_vars[i][1]):
+            value_indices[i] = 0
+            if i + 1 < len(value_indices):
+                return increment(i + 1)
+            else:
+                return True
+        return False
+
+    if not grid_vars:
+        yield unresolved_spec
+        return
+
+    while value_indices[-1] < len(grid_vars[-1][1]):
+        spec = copy.deepcopy(unresolved_spec)
+        for i, (path, values) in enumerate(grid_vars):
+            assign_value(spec, path, values[value_indices[i]])
+        yield spec
+        if grid_vars:
+            done = increment(0)
+            if done:
+                break
+
+
+def _is_resolved(v) -> bool:
+    resolved, _ = _try_resolve(v)
+    return resolved
+
+
+def _try_resolve(v) -> Tuple[bool, Any]:
+    if isinstance(v, Domain):
+        # Domain to sample from
+        return False, v
+    elif isinstance(v, dict) and len(v) == 1 and "eval" in v:
+        # Lambda function in eval syntax
+        return False, Function(
+            lambda spec: eval(v["eval"], _STANDARD_IMPORTS, {"spec": spec})
+        )
+    elif isinstance(v, dict) and len(v) == 1 and "grid_search" in v:
+        # Grid search values
+        grid_values = v["grid_search"]
+        return False, Categorical(grid_values).grid()
+    return True, v
+
+
+def _split_resolved_unresolved_values(
+    spec: Dict,
+) -> Tuple[Dict[Tuple, Any], Dict[Tuple, Any]]:
+    resolved_vars = {}
+    unresolved_vars = {}
+    for k, v in spec.items():
+        resolved, v = _try_resolve(v)
+        if not resolved:
+            unresolved_vars[(k,)] = v
+        elif isinstance(v, dict):
+            # Recurse into a dict
+            (
+                _resolved_children,
+                _unresolved_children,
+            ) = _split_resolved_unresolved_values(v)
+            for path, value in _resolved_children.items():
+                resolved_vars[(k,) + path] = value
+            for path, value in _unresolved_children.items():
+                unresolved_vars[(k,) + path] = value
+        elif isinstance(v, (list, tuple)):
+            # Recurse into a list
+            for i, elem in enumerate(v):
+                (
+                    _resolved_children,
+                    _unresolved_children,
+                ) = _split_resolved_unresolved_values({i: elem})
+                for path, value in _resolved_children.items():
+                    resolved_vars[(k,) + path] = value
+                for path, value in _unresolved_children.items():
+                    unresolved_vars[(k,) + path] = value
+        else:
+            resolved_vars[(k,)] = v
+    return resolved_vars, unresolved_vars
+
+
+def _unresolved_values(spec: Dict) -> Dict[Tuple, Any]:
+    return _split_resolved_unresolved_values(spec)[1]
+
+
+def _has_unresolved_values(spec: Dict) -> bool:
+    return True if _unresolved_values(spec) else False
+
+
+class _UnresolvedAccessGuard(dict):
+    def __init__(self, *args, **kwds):
+        super(_UnresolvedAccessGuard, self).__init__(*args, **kwds)
+        self.__dict__ = self
+
+    def __getattribute__(self, item):
+        value = dict.__getattribute__(self, item)
+        if not _is_resolved(value):
+            raise RecursiveDependencyError(
+                "`{}` recursively depends on {}".format(item, value)
+            )
+        elif isinstance(value, dict):
+            return _UnresolvedAccessGuard(value)
+        else:
+            return value
+
+
+@DeveloperAPI
+class RecursiveDependencyError(Exception):
+    def __init__(self, msg: str):
+        Exception.__init__(self, msg)

.venv/lib/python3.11/site-packages/ray/tune/stopper/__init__.py

@@ -0,0 +1,18 @@
+from ray.tune.stopper.experiment_plateau import ExperimentPlateauStopper
+from ray.tune.stopper.function_stopper import FunctionStopper
+from ray.tune.stopper.maximum_iteration import MaximumIterationStopper
+from ray.tune.stopper.noop import NoopStopper
+from ray.tune.stopper.stopper import CombinedStopper, Stopper
+from ray.tune.stopper.timeout import TimeoutStopper
+from ray.tune.stopper.trial_plateau import TrialPlateauStopper
+
+__all__ = [
+    "Stopper",
+    "CombinedStopper",
+    "ExperimentPlateauStopper",
+    "FunctionStopper",
+    "MaximumIterationStopper",
+    "NoopStopper",
+    "TimeoutStopper",
+    "TrialPlateauStopper",
+]