chore: import upstream snapshot with attribution

python/ray/tune/.claude/CLAUDE.md

@@ -0,0 +1,10 @@
+<!-- Loaded on-demand when Claude works on Ray Tune files. -->
+<!-- Keep under 50 lines. Multi-step procedures → skills. Code style → rules/. -->
+
+# Ray Tune
+
+## Key Modules
+<!-- Entry points, important abstractions, non-obvious dependencies -->
+
+## Gotchas
+<!-- Non-obvious behaviors, common mistakes, things that break silently -->

python/ray/tune/.claude/rules/README.md

@@ -0,0 +1,9 @@
+<!-- Add Ray Tune team-specific rules here as .md files. -->
+<!-- Rules with paths: frontmatter only load when matching files are edited. -->
+<!-- Example:
+---
+paths:
+  - "python/ray/tune/**/*.py"
+---
+- Use the Tuner API for new search algorithm integrations
+-->

python/ray/tune/README.rst

@@ -0,0 +1,30 @@
+Tune: Scalable Hyperparameter Tuning
+====================================
+
+Tune is a scalable framework for hyperparameter search with a focus on deep learning and deep reinforcement learning.
+
+User documentation can be `found here <http://docs.ray.io/en/master/tune.html>`__.
+
+
+Tutorial
+--------
+
+To get started with Tune, try going through `our tutorial of using Tune with Keras <https://github.com/ray-project/tutorial/blob/master/tune_exercises/exercise_1_basics.ipynb>`__.
+
+(Experimental): You can try out `the above tutorial on a free hosted server via Binder <https://mybinder.org/v2/gh/ray-project/tutorial/master?filepath=tune_exercises%2Fexercise_1_basics.ipynb>`__.
+
+
+Citing Tune
+-----------
+
+If Tune helps you in your academic research, you are encouraged to cite `our paper <https://arxiv.org/abs/1807.05118>`__. Here is an example bibtex:
+
+.. code-block:: tex
+
+    @article{liaw2018tune,
+        title={Tune: A Research Platform for Distributed Model Selection and Training},
+        author={Liaw, Richard and Liang, Eric and Nishihara, Robert and
+                Moritz, Philipp and Gonzalez, Joseph E and Stoica, Ion},
+        journal={arXiv preprint arXiv:1807.05118},
+        year={2018}
+    }

python/ray/tune/__init__.py

@@ -0,0 +1,114 @@
+# isort: off
+# Try import ray[tune] core requirements (defined in setup.py)
+try:
+    import fsspec  # noqa: F401
+    import pandas  # noqa: F401
+    import pyarrow  # noqa: F401
+    import requests  # noqa: F401
+except ImportError as exc:
+    raise ImportError(
+        "Can't import ray.tune as some dependencies are missing. "
+        'Run `pip install "ray[tune]"` to fix.'
+    ) from exc
+# isort: on
+
+from ray.tune.trainable.trainable_fn_utils import Checkpoint, get_checkpoint, report
+from ray.tune.impl.config import CheckpointConfig, FailureConfig, RunConfig
+from ray.tune.syncer import SyncConfig
+from ray.air.result import Result
+from ray.tune.analysis import ExperimentAnalysis
+from ray.tune.callback import Callback
+from ray.tune.context import TuneContext, get_context
+from ray.tune.error import TuneError
+from ray.tune.execution.placement_groups import PlacementGroupFactory
+from ray.tune.experiment import Experiment
+from ray.tune.progress_reporter import (
+    CLIReporter,
+    JupyterNotebookReporter,
+    ProgressReporter,
+)
+from ray.tune.registry import register_env, register_trainable
+from ray.tune.result_grid import ResultGrid
+from ray.tune.schedulers import create_scheduler
+from ray.tune.search import create_searcher, grid_search
+from ray.tune.search.sample import (
+    choice,
+    lograndint,
+    loguniform,
+    qlograndint,
+    qloguniform,
+    qrandint,
+    qrandn,
+    quniform,
+    randint,
+    randn,
+    sample_from,
+    uniform,
+)
+from ray.tune.stopper import Stopper
+from ray.tune.trainable import Trainable
+from ray.tune.trainable.util import with_parameters, with_resources
+from ray.tune.tune import run, run_experiments
+from ray.tune.tune_config import ResumeConfig, TuneConfig
+from ray.tune.tuner import Tuner
+
+__all__ = [
+    "Trainable",
+    "Callback",
+    "TuneError",
+    "grid_search",
+    "register_env",
+    "register_trainable",
+    "run",
+    "run_experiments",
+    "with_parameters",
+    "with_resources",
+    "Stopper",
+    "Experiment",
+    "sample_from",
+    "uniform",
+    "quniform",
+    "choice",
+    "randint",
+    "lograndint",
+    "qrandint",
+    "qlograndint",
+    "randn",
+    "qrandn",
+    "loguniform",
+    "qloguniform",
+    "ExperimentAnalysis",
+    "CLIReporter",
+    "JupyterNotebookReporter",
+    "ProgressReporter",
+    "ResultGrid",
+    "create_searcher",
+    "create_scheduler",
+    "PlacementGroupFactory",
+    "Tuner",
+    "TuneConfig",
+    "ResumeConfig",
+    "RunConfig",
+    "CheckpointConfig",
+    "FailureConfig",
+    "Result",
+    "Checkpoint",
+    "get_checkpoint",
+    "report",
+    "get_context",
+    "TuneContext",
+    "SyncConfig",
+]
+
+report.__module__ = "ray.tune"
+get_checkpoint.__module__ = "ray.tune"
+get_context.__module__ = "ray.tune"
+TuneContext.__module__ = "ray.tune"
+Checkpoint.__module__ = "ray.tune"
+Result.__module__ = "ray.tune"
+RunConfig.__module__ = "ray.tune"
+CheckpointConfig.__module__ = "ray.tune"
+FailureConfig.__module__ = "ray.tune"
+
+
+# DO NOT ADD ANYTHING AFTER THIS LINE.

python/ray/tune/analysis/__init__.py

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

python/ray/tune/analysis/experiment_analysis.py

@@ -0,0 +1,691 @@
+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._internal.storage import _exists_at_fs_path, get_fs_and_path
+from ray.tune import Checkpoint
+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 _loads_with_cloudpickle
+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.
+
+    """
+
+    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,
+    ):
+        """Initialize an ``ExperimentAnalysis``.
+
+        Args:
+            experiment_checkpoint_path: Path to an `experiment_state.json` file,
+                or a directory that contains an `experiment_state.json` file.
+            storage_filesystem: A custom ``pyarrow.fs.FileSystem`` corresponding
+                to ``experiment_checkpoint_path``. This may be necessary if the
+                original experiment used a custom filesystem.
+            trials: List of trials that can be accessed via `analysis.trials`.
+            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.
+        """
+        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 = _loads_with_cloudpickle(f.readall())
+
+        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.tune.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.tune.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]`.
+
+        Returns:
+            The hyperparameter configuration of the best trial, or ``None`` if
+            no best trial could be identified.
+        """
+        best_trial = self.get_best_trial(metric, mode, scope)
+        return best_trial.config if best_trial else None
+
+    def get_last_checkpoint(
+        self,
+        trial: Optional[Trial] = None,
+        metric: str = "training_iteration",
+        mode: str = "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

python/ray/tune/automl/__init__.py

@@ -0,0 +1 @@
+raise DeprecationWarning("`ray.tune.automl` is deprecated in Ray 2.6.")

python/ray/tune/callback.py

@@ -0,0 +1,515 @@
+import glob
+import warnings
+from abc import ABCMeta
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
+
+import ray.tune
+from ray.tune.utils.util import _atomic_save, _load_newest_checkpoint
+from ray.util.annotations import DeveloperAPI, PublicAPI
+
+if TYPE_CHECKING:
+    from ray.tune.experiment import Trial
+    from ray.tune.stopper import Stopper
+
+
+class _CallbackMeta(ABCMeta):
+    """A helper metaclass to ensure container classes (e.g. CallbackList) have
+    implemented all the callback methods (e.g. `on_*`).
+    """
+
+    def __new__(mcs, name: str, bases: Tuple[type], attrs: Dict[str, Any]) -> type:
+        cls = super().__new__(mcs, name, bases, attrs)
+
+        if mcs.need_check(cls, name, bases, attrs):
+            mcs.check(cls, name, bases, attrs)
+
+        return cls
+
+    @classmethod
+    def need_check(
+        mcs, cls: type, name: str, bases: Tuple[type], attrs: Dict[str, Any]
+    ) -> bool:
+        return attrs.get("IS_CALLBACK_CONTAINER", False)
+
+    @classmethod
+    def check(
+        mcs, cls: type, name: str, bases: Tuple[type], attrs: Dict[str, Any]
+    ) -> None:
+        methods = set()
+        for base in bases:
+            methods.update(
+                attr_name
+                for attr_name, attr in vars(base).items()
+                if mcs.need_override_by_subclass(attr_name, attr)
+            )
+        overridden = {
+            attr_name
+            for attr_name, attr in attrs.items()
+            if mcs.need_override_by_subclass(attr_name, attr)
+        }
+        missing = methods.difference(overridden)
+        if missing:
+            raise TypeError(
+                f"Found missing callback method: {missing} "
+                f"in class {cls.__module__}.{cls.__qualname__}."
+            )
+
+    @classmethod
+    def need_override_by_subclass(mcs, attr_name: str, attr: Any) -> bool:
+        return (
+            (
+                attr_name.startswith("on_")
+                and not attr_name.startswith("on_trainer_init")
+            )
+            or attr_name == "setup"
+        ) and callable(attr)
+
+
+@PublicAPI(stability="beta")
+class Callback(metaclass=_CallbackMeta):
+    """Tune base callback that can be extended and passed to a ``TrialRunner``
+
+    Tune callbacks are called from within the ``TrialRunner`` class. There are
+    several hooks that can be used, all of which are found in the submethod
+    definitions of this base class.
+
+    The parameters passed to the ``**info`` dict vary between hooks. The
+    parameters passed are described in the docstrings of the methods.
+
+    This example will print a metric each time a result is received:
+
+    .. testcode::
+
+        from ray import tune
+        from ray.tune import Callback
+
+
+        class MyCallback(Callback):
+            def on_trial_result(self, iteration, trials, trial, result,
+                                **info):
+                print(f"Got result: {result['metric']}")
+
+
+        def train_func(config):
+            for i in range(10):
+                tune.report(metric=i)
+
+        tuner = tune.Tuner(
+            train_func,
+            run_config=tune.RunConfig(
+                callbacks=[MyCallback()]
+            )
+        )
+        tuner.fit()
+
+    .. testoutput::
+        :hide:
+
+        ...
+    """
+
+    # File templates for any artifacts written by this callback
+    # These files should live in the `trial.local_path` for each trial.
+    # TODO(ml-team): Make this more visible to users to override. Internal use for now.
+    _SAVED_FILE_TEMPLATES = []
+
+    # arguments here match Experiment.public_spec
+    def setup(
+        self,
+        stop: Optional["Stopper"] = None,
+        num_samples: Optional[int] = None,
+        total_num_samples: Optional[int] = None,
+        **info,
+    ):
+        """Called once at the very beginning of training.
+
+        Any Callback setup should be added here (setting environment
+        variables, etc.)
+
+        Arguments:
+            stop: Stopping criteria.
+                If ``time_budget_s`` was passed to ``tune.RunConfig``, a
+                ``TimeoutStopper`` will be passed here, either by itself
+                or as a part of a ``CombinedStopper``.
+            num_samples: Number of times to sample from the
+                hyperparameter space. Defaults to 1. If `grid_search` is
+                provided as an argument, the grid will be repeated
+                `num_samples` of times. If this is -1, (virtually) infinite
+                samples are generated until a stopping condition is met.
+            total_num_samples: Total number of samples factoring
+                in grid search samplers.
+            **info: Kwargs dict for forward compatibility.
+        """
+        pass
+
+    def on_step_begin(self, iteration: int, trials: List["Trial"], **info):
+        """Called at the start of each tuning loop step.
+
+        Arguments:
+            iteration: Number of iterations of the tuning loop.
+            trials: List of trials.
+            **info: Kwargs dict for forward compatibility.
+        """
+        pass
+
+    def on_step_end(self, iteration: int, trials: List["Trial"], **info):
+        """Called at the end of each tuning loop step.
+
+        The iteration counter is increased before this hook is called.
+
+        Arguments:
+            iteration: Number of iterations of the tuning loop.
+            trials: List of trials.
+            **info: Kwargs dict for forward compatibility.
+        """
+        pass
+
+    def on_trial_start(
+        self, iteration: int, trials: List["Trial"], trial: "Trial", **info
+    ):
+        """Called after starting a trial instance.
+
+        Arguments:
+            iteration: Number of iterations of the tuning loop.
+            trials: List of trials.
+            trial: Trial that just has been started.
+            **info: Kwargs dict for forward compatibility.
+
+        """
+        pass
+
+    def on_trial_restore(
+        self, iteration: int, trials: List["Trial"], trial: "Trial", **info
+    ):
+        """Called after restoring a trial instance.
+
+        Arguments:
+            iteration: Number of iterations of the tuning loop.
+            trials: List of trials.
+            trial: Trial that just has been restored.
+            **info: Kwargs dict for forward compatibility.
+        """
+        pass
+
+    def on_trial_save(
+        self, iteration: int, trials: List["Trial"], trial: "Trial", **info
+    ):
+        """Called after receiving a checkpoint from a trial.
+
+        Arguments:
+            iteration: Number of iterations of the tuning loop.
+            trials: List of trials.
+            trial: Trial that just saved a checkpoint.
+            **info: Kwargs dict for forward compatibility.
+        """
+        pass
+
+    def on_trial_result(
+        self,
+        iteration: int,
+        trials: List["Trial"],
+        trial: "Trial",
+        result: Dict,
+        **info,
+    ):
+        """Called after receiving a result from a trial.
+
+        The search algorithm and scheduler are notified before this
+        hook is called.
+
+        Arguments:
+            iteration: Number of iterations of the tuning loop.
+            trials: List of trials.
+            trial: Trial that just sent a result.
+            result: Result that the trial sent.
+            **info: Kwargs dict for forward compatibility.
+        """
+        pass
+
+    def on_trial_complete(
+        self, iteration: int, trials: List["Trial"], trial: "Trial", **info
+    ):
+        """Called after a trial instance completed.
+
+        The search algorithm and scheduler are notified before this
+        hook is called.
+
+        Arguments:
+            iteration: Number of iterations of the tuning loop.
+            trials: List of trials.
+            trial: Trial that just has been completed.
+            **info: Kwargs dict for forward compatibility.
+        """
+        pass
+
+    def on_trial_recover(
+        self, iteration: int, trials: List["Trial"], trial: "Trial", **info
+    ):
+        """Called after a trial instance failed (errored) but the trial is scheduled
+        for retry.
+
+        The search algorithm and scheduler are not notified.
+
+        Arguments:
+            iteration: Number of iterations of the tuning loop.
+            trials: List of trials.
+            trial: Trial that just has errored.
+            **info: Kwargs dict for forward compatibility.
+        """
+        pass
+
+    def on_trial_error(
+        self, iteration: int, trials: List["Trial"], trial: "Trial", **info
+    ):
+        """Called after a trial instance failed (errored).
+
+        The search algorithm and scheduler are notified before this
+        hook is called.
+
+        Arguments:
+            iteration: Number of iterations of the tuning loop.
+            trials: List of trials.
+            trial: Trial that just has errored.
+            **info: Kwargs dict for forward compatibility.
+        """
+        pass
+
+    def on_checkpoint(
+        self,
+        iteration: int,
+        trials: List["Trial"],
+        trial: "Trial",
+        checkpoint: "ray.tune.Checkpoint",
+        **info,
+    ):
+        """Called after a trial saved a checkpoint with Tune.
+
+        Arguments:
+            iteration: Number of iterations of the tuning loop.
+            trials: List of trials.
+            trial: Trial that just has errored.
+            checkpoint: Checkpoint object that has been saved
+                by the trial.
+            **info: Kwargs dict for forward compatibility.
+        """
+        pass
+
+    def on_experiment_end(self, trials: List["Trial"], **info):
+        """Called after experiment is over and all trials have concluded.
+
+        Arguments:
+            trials: List of trials.
+            **info: Kwargs dict for forward compatibility.
+        """
+        pass
+
+    def get_state(self) -> Optional[Dict]:
+        """Get the state of the callback.
+
+        This method should be implemented by subclasses to return a dictionary
+        representation of the object's current state.
+
+        This is called automatically by Tune to periodically checkpoint callback state.
+        Upon :ref:`Tune experiment restoration <tune-experiment-level-fault-tolerance>`,
+        callback state will be restored via :meth:`~ray.tune.Callback.set_state`.
+
+        .. testcode::
+
+            from typing import Dict, List, Optional
+
+            from ray.tune import Callback
+            from ray.tune.experiment import Trial
+
+            class MyCallback(Callback):
+                def __init__(self):
+                    self._trial_ids = set()
+
+                def on_trial_start(
+                    self, iteration: int, trials: List["Trial"], trial: "Trial", **info
+                ):
+                    self._trial_ids.add(trial.trial_id)
+
+                def get_state(self) -> Optional[Dict]:
+                    return {"trial_ids": self._trial_ids.copy()}
+
+                def set_state(self, state: Dict) -> Optional[Dict]:
+                    self._trial_ids = state["trial_ids"]
+
+        Returns:
+            dict: State of the callback. Should be `None` if the callback does not
+            have any state to save (this is the default).
+        """
+        return None
+
+    def set_state(self, state: Dict):
+        """Set the state of the callback.
+
+        This method should be implemented by subclasses to restore the callback's
+        state based on the given dict state.
+
+        This is used automatically by Tune to restore checkpoint callback state
+        on :ref:`Tune experiment restoration <tune-experiment-level-fault-tolerance>`.
+
+        See :meth:`~ray.tune.Callback.get_state` for an example implementation.
+
+        Args:
+            state: State of the callback.
+        """
+        pass
+
+
+@DeveloperAPI
+class CallbackList(Callback):
+    """Call multiple callbacks at once."""
+
+    IS_CALLBACK_CONTAINER = True
+    CKPT_FILE_TMPL = "callback-states-{}.pkl"
+
+    def __init__(self, callbacks: List[Callback]):
+        self._callbacks = callbacks
+
+    def setup(self, **info):
+        for callback in self._callbacks:
+            try:
+                callback.setup(**info)
+            except TypeError as e:
+                if "argument" in str(e):
+                    warnings.warn(
+                        "Please update `setup` method in callback "
+                        f"`{callback.__class__}` to match the method signature"
+                        " in `ray.tune.callback.Callback`.",
+                        FutureWarning,
+                    )
+                    callback.setup()
+                else:
+                    raise e
+
+    def on_step_begin(self, **info):
+        for callback in self._callbacks:
+            callback.on_step_begin(**info)
+
+    def on_step_end(self, **info):
+        for callback in self._callbacks:
+            callback.on_step_end(**info)
+
+    def on_trial_start(self, **info):
+        for callback in self._callbacks:
+            callback.on_trial_start(**info)
+
+    def on_trial_restore(self, **info):
+        for callback in self._callbacks:
+            callback.on_trial_restore(**info)
+
+    def on_trial_save(self, **info):
+        for callback in self._callbacks:
+            callback.on_trial_save(**info)
+
+    def on_trial_result(self, **info):
+        for callback in self._callbacks:
+            callback.on_trial_result(**info)
+
+    def on_trial_complete(self, **info):
+        for callback in self._callbacks:
+            callback.on_trial_complete(**info)
+
+    def on_trial_recover(self, **info):
+        for callback in self._callbacks:
+            callback.on_trial_recover(**info)
+
+    def on_trial_error(self, **info):
+        for callback in self._callbacks:
+            callback.on_trial_error(**info)
+
+    def on_checkpoint(self, **info):
+        for callback in self._callbacks:
+            callback.on_checkpoint(**info)
+
+    def on_experiment_end(self, **info):
+        for callback in self._callbacks:
+            callback.on_experiment_end(**info)
+
+    def get_state(self) -> Optional[Dict]:
+        """Gets the state of all callbacks contained within this list.
+        If there are no stateful callbacks, then None will be returned in order
+        to avoid saving an unnecessary callback checkpoint file."""
+        state = {}
+        any_stateful_callbacks = False
+        for i, callback in enumerate(self._callbacks):
+            callback_state = callback.get_state()
+            if callback_state:
+                any_stateful_callbacks = True
+            state[i] = callback_state
+        if not any_stateful_callbacks:
+            return None
+        return state
+
+    def set_state(self, state: Dict):
+        """Sets the state for all callbacks contained within this list.
+        Skips setting state for all stateless callbacks where `get_state`
+        returned None."""
+        for i, callback in enumerate(self._callbacks):
+            callback_state = state.get(i, None)
+            if callback_state:
+                callback.set_state(callback_state)
+
+    def save_to_dir(self, checkpoint_dir: str, session_str: str = "default"):
+        """Save the state of the callback list to the checkpoint_dir.
+
+        Args:
+            checkpoint_dir: directory where the checkpoint is stored.
+            session_str: Unique identifier of the current run session (ex: timestamp).
+        """
+        state_dict = self.get_state()
+
+        if state_dict:
+            file_name = self.CKPT_FILE_TMPL.format(session_str)
+            tmp_file_name = f"tmp-{file_name}"
+            _atomic_save(
+                state=state_dict,
+                checkpoint_dir=checkpoint_dir,
+                file_name=file_name,
+                tmp_file_name=tmp_file_name,
+            )
+
+    def restore_from_dir(self, checkpoint_dir: str):
+        """Restore the state of the list of callbacks from the checkpoint_dir.
+
+        You should check if it's possible to restore with `can_restore`
+        before calling this method.
+
+        Args:
+            checkpoint_dir: directory where the checkpoint is stored.
+
+        Raises:
+            RuntimeError: if unable to find checkpoint.
+            NotImplementedError: if the `set_state` method is not implemented.
+        """
+        state_dict = _load_newest_checkpoint(
+            checkpoint_dir, self.CKPT_FILE_TMPL.format("*")
+        )
+        if not state_dict:
+            raise RuntimeError(
+                "Unable to find checkpoint in {}.".format(checkpoint_dir)
+            )
+        self.set_state(state_dict)
+
+    def can_restore(self, checkpoint_dir: str) -> bool:
+        """Check if the checkpoint_dir contains the saved state for this callback list.
+
+        Args:
+            checkpoint_dir: Directory to look for a saved state file in.
+
+        Returns:
+            can_restore: True if the checkpoint_dir contains a file of the
+                format `CKPT_FILE_TMPL`. False otherwise.
+        """
+        return any(
+            glob.iglob(Path(checkpoint_dir, self.CKPT_FILE_TMPL.format("*")).as_posix())
+        )
+
+    def __len__(self) -> int:
+        return len(self._callbacks)
+
+    def __getitem__(self, i: int) -> "Callback":
+        return self._callbacks[i]

python/ray/tune/cli/commands.py

@@ -0,0 +1,309 @@
+import logging
+import operator
+import os
+import shutil
+import subprocess
+from datetime import datetime
+from pathlib import Path
+from typing import List, Optional
+
+import click
+import pandas as pd
+from pandas.api.types import is_numeric_dtype, is_string_dtype
+
+from ray._private.thirdparty.tabulate.tabulate import tabulate
+from ray.air.constants import EXPR_RESULT_FILE
+from ray.tune import TuneError
+from ray.tune.analysis import ExperimentAnalysis
+from ray.tune.result import (
+    CONFIG_PREFIX,
+    DEFAULT_EXPERIMENT_INFO_KEYS,
+    DEFAULT_RESULT_KEYS,
+)
+
+logger = logging.getLogger(__name__)
+
+EDITOR = os.getenv("EDITOR", "vim")
+
+TIMESTAMP_FORMAT = "%Y-%m-%d %H:%M:%S (%A)"
+
+DEFAULT_CLI_KEYS = DEFAULT_EXPERIMENT_INFO_KEYS + DEFAULT_RESULT_KEYS
+
+DEFAULT_PROJECT_INFO_KEYS = (
+    "name",
+    "total_trials",
+    "last_updated",
+)
+
+TERM_WIDTH, TERM_HEIGHT = shutil.get_terminal_size(fallback=(100, 100))
+
+OPERATORS = {
+    "<": operator.lt,
+    "<=": operator.le,
+    "==": operator.eq,
+    "!=": operator.ne,
+    ">=": operator.ge,
+    ">": operator.gt,
+}
+
+
+def _check_tabulate():
+    """Checks whether tabulate is installed."""
+    if tabulate is None:
+        raise ImportError("Tabulate not installed. Please run `pip install tabulate`.")
+
+
+def print_format_output(dataframe: pd.DataFrame):
+    """Prints output of given dataframe to fit into terminal.
+
+    Args:
+        dataframe: The dataframe to print to the terminal.
+
+    Returns:
+        table: Final outputted dataframe.
+        dropped_cols: Columns dropped due to terminal size.
+        empty_cols: Empty columns (dropped on default).
+    """
+    print_df = pd.DataFrame()
+    dropped_cols = []
+    empty_cols = []
+    # column display priority is based on the info_keys passed in
+    for i, col in enumerate(dataframe):
+        if dataframe[col].isnull().all():
+            # Don't add col to print_df if is fully empty
+            empty_cols += [col]
+            continue
+
+        print_df[col] = dataframe[col]
+        test_table = tabulate(print_df, headers="keys", tablefmt="psql")
+        if str(test_table).index("\n") > TERM_WIDTH:
+            # Drop all columns beyond terminal width
+            print_df.drop(col, axis=1, inplace=True)
+            dropped_cols += list(dataframe.columns)[i:]
+            break
+
+    table = tabulate(print_df, headers="keys", tablefmt="psql", showindex="never")
+
+    print(table)
+    if dropped_cols:
+        click.secho("Dropped columns: {}".format(dropped_cols), fg="yellow")
+        click.secho("Please increase your terminal size to view remaining columns.")
+    if empty_cols:
+        click.secho("Empty columns: {}".format(empty_cols), fg="yellow")
+
+    return table, dropped_cols, empty_cols
+
+
+def list_trials(
+    experiment_path: str,
+    sort: Optional[List[str]] = None,
+    output: Optional[str] = None,
+    filter_op: Optional[str] = None,
+    info_keys: Optional[List[str]] = None,
+    limit: int = None,
+    desc: bool = False,
+):
+    """Lists trials in the directory subtree starting at the given path.
+
+    Args:
+        experiment_path: Directory where trials are located.
+            Like Experiment.local_dir/Experiment.name/experiment*.json.
+        sort: Keys to sort by.
+        output: Name of file where output is saved.
+        filter_op: Filter operation in the format
+            "<column> <operator> <value>".
+        info_keys: Keys that are displayed.
+        limit: Number of rows to display.
+        desc: Sort ascending vs. descending.
+    """
+    _check_tabulate()
+
+    try:
+        checkpoints_df = ExperimentAnalysis(experiment_path).dataframe()  # last result
+    except TuneError as e:
+        raise click.ClickException("No trial data found!") from e
+
+    config_prefix = CONFIG_PREFIX + "/"
+
+    def key_filter(k):
+        return k in DEFAULT_CLI_KEYS or k.startswith(config_prefix)
+
+    col_keys = [k for k in checkpoints_df.columns if key_filter(k)]
+
+    if info_keys:
+        for k in info_keys:
+            if k not in checkpoints_df.columns:
+                raise click.ClickException(
+                    "Provided key invalid: {}. "
+                    "Available keys: {}.".format(k, checkpoints_df.columns)
+                )
+        col_keys = [k for k in checkpoints_df.columns if k in info_keys]
+
+    if not col_keys:
+        raise click.ClickException("No columns to output.")
+
+    checkpoints_df = checkpoints_df[col_keys]
+    if "last_update_time" in checkpoints_df:
+        with pd.option_context("mode.use_inf_as_null", True):
+            datetime_series = checkpoints_df["last_update_time"].dropna()
+
+        datetime_series = datetime_series.apply(
+            lambda t: datetime.fromtimestamp(t).strftime(TIMESTAMP_FORMAT)
+        )
+        checkpoints_df["last_update_time"] = datetime_series
+
+    if "logdir" in checkpoints_df:
+        # logdir often too long to view in table, so drop experiment_path
+        checkpoints_df["logdir"] = checkpoints_df["logdir"].str.replace(
+            experiment_path, ""
+        )
+
+    if filter_op:
+        col, op, val = filter_op.split(" ")
+        col_type = checkpoints_df[col].dtype
+        if is_numeric_dtype(col_type):
+            val = float(val)
+        elif is_string_dtype(col_type):
+            val = str(val)
+        # TODO(Andrew): add support for datetime and boolean
+        else:
+            raise click.ClickException(
+                "Unsupported dtype for {}: {}".format(val, col_type)
+            )
+        op = OPERATORS[op]
+        filtered_index = op(checkpoints_df[col], val)
+        checkpoints_df = checkpoints_df[filtered_index]
+
+    if sort:
+        for key in sort:
+            if key not in checkpoints_df:
+                raise click.ClickException(
+                    "{} not in: {}".format(key, list(checkpoints_df))
+                )
+        ascending = not desc
+        checkpoints_df = checkpoints_df.sort_values(by=sort, ascending=ascending)
+
+    if limit:
+        checkpoints_df = checkpoints_df[:limit]
+
+    print_format_output(checkpoints_df)
+
+    if output:
+        file_extension = os.path.splitext(output)[1].lower()
+        if file_extension in (".p", ".pkl", ".pickle"):
+            checkpoints_df.to_pickle(output)
+        elif file_extension == ".csv":
+            checkpoints_df.to_csv(output, index=False)
+        else:
+            raise click.ClickException("Unsupported filetype: {}".format(output))
+        click.secho("Output saved at {}".format(output), fg="green")
+
+
+def list_experiments(
+    project_path: str,
+    sort: Optional[List[str]] = None,
+    output: str = None,
+    filter_op: str = None,
+    info_keys: Optional[List[str]] = None,
+    limit: int = None,
+    desc: bool = False,
+):
+    """Lists experiments in the directory subtree.
+
+    Args:
+        project_path: Directory where experiments are located.
+            Corresponds to Experiment.local_dir.
+        sort: Keys to sort by.
+        output: Name of file where output is saved.
+        filter_op: Filter operation in the format
+            "<column> <operator> <value>".
+        info_keys: Keys that are displayed.
+        limit: Number of rows to display.
+        desc: Sort ascending vs. descending.
+    """
+    _check_tabulate()
+    base, experiment_folders, _ = next(os.walk(project_path))
+
+    experiment_data_collection = []
+
+    for experiment_dir in experiment_folders:
+        num_trials = sum(
+            EXPR_RESULT_FILE in files
+            for _, _, files in os.walk(os.path.join(base, experiment_dir))
+        )
+
+        experiment_data = {"name": experiment_dir, "total_trials": num_trials}
+        experiment_data_collection.append(experiment_data)
+
+    if not experiment_data_collection:
+        raise click.ClickException("No experiments found!")
+
+    info_df = pd.DataFrame(experiment_data_collection)
+    if not info_keys:
+        info_keys = DEFAULT_PROJECT_INFO_KEYS
+    col_keys = [k for k in list(info_keys) if k in info_df]
+    if not col_keys:
+        raise click.ClickException(
+            "None of keys {} in experiment data!".format(info_keys)
+        )
+    info_df = info_df[col_keys]
+
+    if filter_op:
+        col, op, val = filter_op.split(" ")
+        col_type = info_df[col].dtype
+        if is_numeric_dtype(col_type):
+            val = float(val)
+        elif is_string_dtype(col_type):
+            val = str(val)
+        # TODO(Andrew): add support for datetime and boolean
+        else:
+            raise click.ClickException(
+                "Unsupported dtype for {}: {}".format(val, col_type)
+            )
+        op = OPERATORS[op]
+        filtered_index = op(info_df[col], val)
+        info_df = info_df[filtered_index]
+
+    if sort:
+        for key in sort:
+            if key not in info_df:
+                raise click.ClickException("{} not in: {}".format(key, list(info_df)))
+        ascending = not desc
+        info_df = info_df.sort_values(by=sort, ascending=ascending)
+
+    if limit:
+        info_df = info_df[:limit]
+
+    print_format_output(info_df)
+
+    if output:
+        file_extension = os.path.splitext(output)[1].lower()
+        if file_extension in (".p", ".pkl", ".pickle"):
+            info_df.to_pickle(output)
+        elif file_extension == ".csv":
+            info_df.to_csv(output, index=False)
+        else:
+            raise click.ClickException("Unsupported filetype: {}".format(output))
+        click.secho("Output saved at {}".format(output), fg="green")
+
+
+def add_note(path: str, filename: str = "note.txt"):
+    """Opens a txt file at the given path where user can add and save notes.
+
+    Args:
+        path: Directory where note will be saved.
+        filename: Name of note. Defaults to "note.txt"
+    """
+    path = Path(path).expanduser()
+    assert path.is_dir(), "{} is not a valid directory.".format(path)
+
+    filepath = path / filename
+
+    try:
+        subprocess.call([EDITOR, filepath.as_posix()])
+    except Exception as exc:
+        click.secho("Editing note failed: {}".format(str(exc)), fg="red")
+    if filepath.exists():
+        print("Note updated at:", filepath.as_posix())
+    else:
+        print("Note created at:", filepath.as_posix())

python/ray/tune/cli/scripts.py

@@ -0,0 +1,101 @@
+import click
+
+import ray.tune.cli.commands as commands
+
+
+@click.group()
+def cli():
+    pass
+
+
+@cli.command()
+@click.argument("experiment_path", required=True, type=str)
+@click.option("--sort", default=None, type=str, help="Select which column to sort on.")
+@click.option(
+    "--output",
+    "-o",
+    default=None,
+    type=str,
+    help="Select file to output information to.",
+)
+@click.option(
+    "--filter",
+    "filter_op",
+    default=None,
+    type=str,
+    help="Select filter in the format '<column> <operator> <value>'.",
+)
+@click.option(
+    "--columns", default=None, type=str, help="Select columns to be displayed."
+)
+@click.option(
+    "--limit", default=None, type=int, help="Select number of rows to display."
+)
+@click.option("--desc", default=False, type=bool, help="Sort ascending vs. descending.")
+def list_trials(experiment_path, sort, output, filter_op, columns, limit, desc):
+    """Lists trials in the directory subtree starting at the given path."""
+    if sort:
+        sort = sort.split(",")
+    if columns:
+        columns = columns.split(",")
+    commands.list_trials(experiment_path, sort, output, filter_op, columns, limit, desc)
+
+
+@cli.command()
+@click.argument("project_path", required=True, type=str)
+@click.option("--sort", default=None, type=str, help="Select which column to sort on.")
+@click.option(
+    "--output",
+    "-o",
+    default=None,
+    type=str,
+    help="Select file to output information to.",
+)
+@click.option(
+    "--filter",
+    "filter_op",
+    default=None,
+    type=str,
+    help="Select filter in the format '<column> <operator> <value>'.",
+)
+@click.option(
+    "--columns", default=None, type=str, help="Select columns to be displayed."
+)
+@click.option(
+    "--limit", default=None, type=int, help="Select number of rows to display."
+)
+@click.option("--desc", default=False, type=bool, help="Sort ascending vs. descending.")
+def list_experiments(project_path, sort, output, filter_op, columns, limit, desc):
+    """Lists experiments in the directory subtree."""
+    if sort:
+        sort = sort.split(",")
+    if columns:
+        columns = columns.split(",")
+    commands.list_experiments(
+        project_path, sort, output, filter_op, columns, limit, desc
+    )
+
+
+@cli.command()
+@click.argument("path", required=True, type=str)
+@click.option(
+    "--filename", default="note.txt", type=str, help="Specify filename for note."
+)
+def add_note(path, filename):
+    """Adds user notes as a text file at the given path."""
+    commands.add_note(path, filename)
+
+
+cli.add_command(list_trials, name="ls")
+cli.add_command(list_trials, name="list-trials")
+cli.add_command(list_experiments, name="lsx")
+cli.add_command(list_experiments, name="list-experiments")
+cli.add_command(add_note, name="add-note")
+
+
+def main():
+    return cli()
+
+
+if __name__ == "__main__":
+    main()

python/ray/tune/constants.py

@@ -0,0 +1,36 @@
+# ==================================================
+#               Environment Variables
+# ==================================================
+
+# Environment variable for Tune execution callbacks
+RAY_TUNE_CALLBACKS_ENV_VAR = "RAY_TUNE_CALLBACKS"
+
+# NOTE: When adding a new environment variable, please track it in this list.
+TUNE_ENV_VARS = {
+    "RAY_AIR_LOCAL_CACHE_DIR",
+    "TUNE_DISABLE_AUTO_CALLBACK_LOGGERS",
+    "TUNE_DISABLE_AUTO_INIT",
+    "TUNE_DISABLE_DATED_SUBDIR",
+    "TUNE_DISABLE_STRICT_METRIC_CHECKING",
+    "TUNE_DISABLE_SIGINT_HANDLER",
+    "TUNE_FORCE_TRIAL_CLEANUP_S",
+    "TUNE_FUNCTION_THREAD_TIMEOUT_S",
+    "TUNE_GLOBAL_CHECKPOINT_S",
+    "TUNE_MAX_LEN_IDENTIFIER",
+    "TUNE_MAX_PENDING_TRIALS_PG",
+    "TUNE_PLACEMENT_GROUP_PREFIX",
+    "TUNE_PLACEMENT_GROUP_RECON_INTERVAL",
+    "TUNE_PRINT_ALL_TRIAL_ERRORS",
+    "TUNE_RESULT_DIR",
+    "TUNE_RESULT_BUFFER_LENGTH",
+    "TUNE_RESULT_DELIM",
+    "TUNE_RESULT_BUFFER_MAX_TIME_S",
+    "TUNE_RESULT_BUFFER_MIN_TIME_S",
+    "TUNE_WARN_THRESHOLD_S",
+    "TUNE_WARN_INSUFFICENT_RESOURCE_THRESHOLD_S",
+    "TUNE_WARN_INSUFFICENT_RESOURCE_THRESHOLD_S_AUTOSCALER",
+    "TUNE_WARN_EXCESSIVE_EXPERIMENT_CHECKPOINT_SYNC_THRESHOLD_S",
+    "TUNE_STATE_REFRESH_PERIOD",
+    "TUNE_RESTORE_RETRY_NUM",
+    RAY_TUNE_CALLBACKS_ENV_VAR,
+}

python/ray/tune/context.py

@@ -0,0 +1,116 @@
+import threading
+from typing import Any, Dict, Optional
+
+from ray.train._internal import session
+from ray.train.constants import (
+    V2_MIGRATION_GUIDE_MESSAGE,
+    _v2_migration_warnings_enabled,
+)
+from ray.train.context import TrainContext as TrainV1Context
+from ray.train.utils import _copy_doc
+from ray.tune.execution.placement_groups import PlacementGroupFactory
+from ray.util.annotations import Deprecated, PublicAPI
+
+# The context singleton on this process.
+_tune_context: Optional["TuneContext"] = None
+_tune_context_lock = threading.Lock()
+
+
+_TRAIN_SPECIFIC_CONTEXT_DEPRECATION_MESSAGE = (
+    "`{}` is deprecated for Ray Tune because there is no concept of worker ranks "
+    "for Ray Tune, so these methods only make sense to use in the context of "
+    f"a Ray Train worker. {V2_MIGRATION_GUIDE_MESSAGE}"
+)
+
+
+@PublicAPI(stability="beta")
+class TuneContext(TrainV1Context):
+    """Context to access metadata within Ray Tune functions."""
+
+    # NOTE: These methods are deprecated on the TrainContext, but are still
+    # available on the TuneContext. Re-defining them here to avoid the
+    # deprecation warnings.
+
+    @_copy_doc(session.get_trial_name)
+    def get_trial_name(self) -> str:
+        return session.get_trial_name()
+
+    @_copy_doc(session.get_trial_id)
+    def get_trial_id(self) -> str:
+        return session.get_trial_id()
+
+    @_copy_doc(session.get_trial_resources)
+    def get_trial_resources(self) -> PlacementGroupFactory:
+        return session.get_trial_resources()
+
+    @_copy_doc(session.get_trial_dir)
+    def get_trial_dir(self) -> str:
+        return session.get_trial_dir()
+
+    # Deprecated APIs
+
+    @Deprecated
+    def get_metadata(self) -> Dict[str, Any]:
+        raise DeprecationWarning(
+            "`get_metadata` is deprecated for Ray Tune, as it has never been usable."
+        )
+
+    @Deprecated(
+        message=_TRAIN_SPECIFIC_CONTEXT_DEPRECATION_MESSAGE.format("get_world_size"),
+        warning=_v2_migration_warnings_enabled(),
+    )
+    @_copy_doc(TrainV1Context.get_world_size)
+    def get_world_size(self) -> int:
+        return session.get_world_size()
+
+    @Deprecated(
+        message=_TRAIN_SPECIFIC_CONTEXT_DEPRECATION_MESSAGE.format("get_world_rank"),
+        warning=_v2_migration_warnings_enabled(),
+    )
+    @_copy_doc(TrainV1Context.get_world_rank)
+    def get_world_rank(self) -> int:
+        return session.get_world_rank()
+
+    @Deprecated(
+        message=_TRAIN_SPECIFIC_CONTEXT_DEPRECATION_MESSAGE.format("get_local_rank"),
+        warning=_v2_migration_warnings_enabled(),
+    )
+    @_copy_doc(TrainV1Context.get_local_rank)
+    def get_local_rank(self) -> int:
+        return session.get_local_rank()
+
+    @Deprecated(
+        message=_TRAIN_SPECIFIC_CONTEXT_DEPRECATION_MESSAGE.format(
+            "get_local_world_size"
+        ),
+        warning=_v2_migration_warnings_enabled(),
+    )
+    @_copy_doc(TrainV1Context.get_local_world_size)
+    def get_local_world_size(self) -> int:
+        return session.get_local_world_size()
+
+    @Deprecated(
+        message=_TRAIN_SPECIFIC_CONTEXT_DEPRECATION_MESSAGE.format("get_node_rank"),
+        warning=_v2_migration_warnings_enabled(),
+    )
+    @_copy_doc(TrainV1Context.get_node_rank)
+    def get_node_rank(self) -> int:
+        return session.get_node_rank()
+
+
+@PublicAPI(stability="beta")
+def get_context() -> TuneContext:
+    """Get or create a singleton Ray Tune context.
+
+    The context is only available in a tune function passed to the `ray.tune.Tuner`.
+
+    See the :class:`~ray.tune.TuneContext` API reference to see available methods.
+    """
+    global _tune_context
+
+    with _tune_context_lock:
+        if _tune_context is None:
+            # TODO(justinvyu): This default should be a dummy context
+            # that is only used for testing / running outside of Tune.
+            _tune_context = TuneContext()
+        return _tune_context

python/ray/tune/error.py

@@ -0,0 +1,48 @@
+from ray.util.annotations import PublicAPI
+
+
+@PublicAPI
+class TuneError(Exception):
+    """General error class raised by ray.tune."""
+
+    pass
+
+
+class _AbortTrialExecution(TuneError):
+    """Error that indicates a trial should not be retried."""
+
+    pass
+
+
+class _SubCategoryTuneError(TuneError):
+    """The more specific TuneError that happens for a certain Tune
+    subroutine. For example starting/stopping a trial.
+    """
+
+    def __init__(self, traceback_str: str):
+        self.traceback_str = traceback_str
+
+    def __str__(self):
+        return self.traceback_str
+
+
+class _TuneStopTrialError(_SubCategoryTuneError):
+    """Error that happens when stopping a tune trial."""
+
+    pass
+
+
+class _TuneStartTrialError(_SubCategoryTuneError):
+    """Error that happens when starting a tune trial."""
+
+    pass
+
+
+class _TuneNoNextExecutorEventError(_SubCategoryTuneError):
+    """Error that happens when waiting to get the next event to
+    handle from RayTrialExecutor.
+
+    Note: RayTaskError will be raised by itself and will not be using
+    this category. This category is for everything else."""
+
+    pass

python/ray/tune/examples/README.rst

@@ -0,0 +1,81 @@
+Tune Examples
+=============
+
+.. Keep this in sync with ray/doc/tune-examples.rst
+
+In our repository, we provide a variety of examples for the various use cases and features of Tune.
+
+If any example is broken, or if you'd like to add an example to this page, feel free to raise an issue on our Github repository.
+
+
+General Examples
+----------------
+
+- `async_hyperband_example <https://github.com/ray-project/ray/blob/master/python/ray/tune/examples/async_hyperband_example.py>`__: Example of using a Trainable class with AsyncHyperBandScheduler.
+- `hyperband_example <https://github.com/ray-project/ray/blob/master/python/ray/tune/examples/hyperband_example.py>`__: Example of using a Trainable class with HyperBandScheduler. Also uses the Experiment class API for specifying the experiment configuration. Also uses the AsyncHyperBandScheduler.
+- `pbt_example <https://github.com/ray-project/ray/blob/master/python/ray/tune/examples/pbt_example.py>`__: Example of using a Trainable class with PopulationBasedTraining scheduler.
+- `PBT with Function API <https://github.com/ray-project/ray/blob/master/python/ray/tune/examples/pbt_function.py>`__: Example of using the function API with a PopulationBasedTraining scheduler.
+- `pbt_ppo_example <https://github.com/ray-project/ray/blob/master/python/ray/tune/examples/pbt_ppo_example.py>`__: Example of optimizing a distributed RLlib algorithm (PPO) with the PopulationBasedTraining scheduler.
+- `logging_example <https://github.com/ray-project/ray/blob/master/python/ray/tune/examples/logging_example.py>`__: Example of custom loggers and custom trial directory naming.
+- `custom_func_checkpointing <https://github.com/ray-project/ray/blob/master/python/ray/tune/examples/logging_example.py>`__: Example of custom checkpointing logic using the function API.
+
+Search Algorithm Examples
+-------------------------
+
+- `Ax example <https://github.com/ray-project/ray/blob/master/python/ray/tune/examples/ax_example.py>`__: Optimize a Hartmann function with `Ax <https://ax.dev>`_ with 4 parallel workers.
+- `Nevergrad example <https://github.com/ray-project/ray/blob/master/python/ray/tune/examples/nevergrad_example.py>`__: Optimize a simple toy function with the gradient-free optimization package `Nevergrad <https://github.com/facebookresearch/nevergrad>`_ with 4 parallel workers.
+- `Bayesian Optimization example <https://github.com/ray-project/ray/blob/master/python/ray/tune/examples/bayesopt_example.py>`__: Optimize a simple toy function using `Bayesian Optimization <https://github.com/fmfn/BayesianOptimization>`_ with 4 parallel workers.
+
+Tensorflow/Keras Examples
+-------------------------
+
+- `tune_mnist_keras <https://github.com/ray-project/ray/blob/master/python/ray/tune/examples/tune_mnist_keras.py>`__: Converts the Keras MNIST example to use Tune with the function-based API and a Keras callback. Also shows how to easily convert something relying on argparse to use Tune.
+- `pbt_memnn_example <https://github.com/ray-project/ray/blob/master/python/ray/tune/examples/pbt_memnn_example.py>`__: Example of training a Memory NN on bAbI with Keras using PBT.
+- `Tensorflow 2 Example <https://github.com/ray-project/ray/blob/master/python/ray/tune/examples/tf_mnist_example.py>`__: Converts the Advanced TF2.0 MNIST example to use Tune with the Trainable. This uses `tf.function`. Original code from tensorflow: https://www.tensorflow.org/tutorials/quickstart/advanced
+
+
+PyTorch Examples
+----------------
+
+- `mnist_pytorch <https://github.com/ray-project/ray/blob/master/python/ray/tune/examples/mnist_pytorch.py>`__: Converts the PyTorch MNIST example to use Tune with the function-based API. Also shows how to easily convert something relying on argparse to use Tune.
+- `mnist_pytorch_trainable <https://github.com/ray-project/ray/blob/master/python/ray/tune/examples/mnist_pytorch_trainable.py>`__: Converts the PyTorch MNIST example to use Tune with Trainable API. Also uses the HyperBandScheduler and checkpoints the model at the end.
+
+
+PyTorch Lightning Examples
+--------------------------
+
+For a full walkthrough of tuning a PyTorch Lightning model with Ray Tune, see the
+`Using PyTorch Lightning with Tune <https://docs.ray.io/en/latest/tune/examples/tune-pytorch-lightning.html>`__ tutorial.
+
+- `mnist_ptl_mini <https://github.com/ray-project/ray/blob/master/python/ray/tune/examples/mnist_ptl_mini.py>`__: A minimal example of tuning a PyTorch Lightning MNIST classifier with Ray Tune.
+- `mlflow_ptl <https://github.com/ray-project/ray/blob/master/python/ray/tune/examples/mlflow_ptl.py>`__: Example for using `MLflow <https://github.com/mlflow/mlflow/>`__ and PyTorch Lightning with Ray Tune.
+
+
+XGBoost Example
+---------------
+
+- `xgboost_example <https://github.com/ray-project/ray/blob/master/python/ray/tune/examples/xgboost_example.py>`__: Trains a basic XGBoost model with Tune with the function-based API and a XGBoost callback.
+
+
+XGBoost with Dynamic Resources Example
+--------------------------------------
+
+- `xgboost_dynamic_resources_example <https://github.com/ray-project/ray/blob/master/python/ray/tune/examples/xgboost_dynamic_resources_example.py>`__: Trains a basic XGBoost model with Tune with the class-based API and a ResourceChangingScheduler, ensuring all resources are being used at all time.
+
+
+LightGBM Example
+----------------
+
+- `lightgbm_example <https://github.com/ray-project/ray/blob/master/python/ray/tune/examples/lightgbm_example.py>`__: Trains a basic LightGBM model with Tune with the function-based API and a LightGBM callback.
+
+Huggingface Transformers Example
+--------------------------------
+
+- `pbt_transformers <https://github.com/ray-project/ray/blob/master/python/ray/tune/examples/pbt_transformers/pbt_transformers.py>`__: Fine-tunes a Huggingface transformer with Tune Population Based Training.
+
+
+Contributed Examples
+--------------------
+
+- `pbt_tune_cifar10_with_keras <https://github.com/ray-project/ray/blob/master/python/ray/tune/examples/pbt_tune_cifar10_with_keras.py>`__: A contributed example of tuning a Keras model on CIFAR10 with the PopulationBasedTraining scheduler.
+- `hyperopt_conditional_search_space_example <https://github.com/ray-project/ray/blob/master/python/ray/tune/examples/hyperopt_conditional_search_space_example.py>`__: Conditional search space example using HyperOpt.

python/ray/tune/examples/async_hyperband_example.py

@@ -0,0 +1,63 @@
+#!/usr/bin/env python
+
+import argparse
+import time
+from typing import Any, Dict
+
+from ray import tune
+from ray.tune.schedulers import AsyncHyperBandScheduler
+
+
+def evaluation_fn(step, width, height) -> float:
+    # simulate model evaluation
+    time.sleep(0.1)
+    return (0.1 + width * step / 100) ** (-1) + height * 0.1
+
+
+def easy_objective(config: Dict[str, Any]) -> None:
+    # Config contains the hyperparameters to tune
+    width, height = config["width"], config["height"]
+
+    for step in range(config["steps"]):
+        # Iterative training function - can be an arbitrary training procedure
+        intermediate_score = evaluation_fn(step, width, height)
+        # Feed the score back back to Tune.
+        tune.report({"iterations": step, "mean_loss": intermediate_score})
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="AsyncHyperBand optimization example")
+    parser.add_argument(
+        "--smoke-test", action="store_true", help="Finish quickly for testing"
+    )
+    args, _ = parser.parse_known_args()
+
+    # AsyncHyperBand enables aggressive early stopping of poorly performing trials
+    scheduler = AsyncHyperBandScheduler(
+        grace_period=5,  # Minimum training iterations before stopping
+        max_t=100,  # Maximum training iterations
+    )
+
+    tuner = tune.Tuner(
+        tune.with_resources(easy_objective, {"cpu": 1, "gpu": 0}),
+        run_config=tune.RunConfig(
+            name="asynchyperband_test",
+            stop={"training_iteration": 1 if args.smoke_test else 9999},
+            verbose=1,
+        ),
+        tune_config=tune.TuneConfig(
+            metric="mean_loss",
+            mode="min",
+            scheduler=scheduler,
+            num_samples=20,  # Number of trials to run
+        ),
+        param_space={
+            "steps": 100,
+            "width": tune.uniform(10, 100),
+            "height": tune.uniform(0, 100),
+        },
+    )
+
+    # Run the hyperparameter optimization
+    results = tuner.fit()
+    print(f"Best hyperparameters found: {results.get_best_result().config}")

python/ray/tune/examples/ax_example.py

@@ -0,0 +1,97 @@
+"""This example demonstrates the usage of AxSearch with Ray Tune.
+
+It also checks that it is usable with a separate scheduler.
+
+Requires the Ax library to be installed (`pip install ax-platform`).
+"""
+
+import time
+
+import numpy as np
+
+from ray import tune
+from ray.tune.schedulers import AsyncHyperBandScheduler
+from ray.tune.search.ax import AxSearch
+
+
+def hartmann6(x):
+    alpha = np.array([1.0, 1.2, 3.0, 3.2])
+    A = np.array(
+        [
+            [10, 3, 17, 3.5, 1.7, 8],
+            [0.05, 10, 17, 0.1, 8, 14],
+            [3, 3.5, 1.7, 10, 17, 8],
+            [17, 8, 0.05, 10, 0.1, 14],
+        ]
+    )
+    P = 10 ** (-4) * np.array(
+        [
+            [1312, 1696, 5569, 124, 8283, 5886],
+            [2329, 4135, 8307, 3736, 1004, 9991],
+            [2348, 1451, 3522, 2883, 3047, 6650],
+            [4047, 8828, 8732, 5743, 1091, 381],
+        ]
+    )
+    y = 0.0
+    for j, alpha_j in enumerate(alpha):
+        t = 0
+        for k in range(6):
+            t += A[j, k] * ((x[k] - P[j, k]) ** 2)
+        y -= alpha_j * np.exp(-t)
+    return y
+
+
+def easy_objective(config):
+    for i in range(config["iterations"]):
+        x = np.array([config.get("x{}".format(i + 1)) for i in range(6)])
+        tune.report(
+            {
+                "timesteps_total": i,
+                "hartmann6": hartmann6(x),
+                "l2norm": np.sqrt((x**2).sum()),
+            }
+        )
+        time.sleep(0.02)
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--smoke-test", action="store_true", help="Finish quickly for testing"
+    )
+    args, _ = parser.parse_known_args()
+
+    algo = AxSearch(
+        parameter_constraints=["x1 + x2 <= 2.0"],  # Optional.
+        outcome_constraints=["l2norm <= 1.25"],  # Optional.
+    )
+    # Limit to 4 concurrent trials
+    algo = tune.search.ConcurrencyLimiter(algo, max_concurrent=4)
+    scheduler = AsyncHyperBandScheduler()
+    tuner = tune.Tuner(
+        easy_objective,
+        run_config=tune.RunConfig(
+            name="ax",
+            stop={"timesteps_total": 100},
+        ),
+        tune_config=tune.TuneConfig(
+            metric="hartmann6",  # provided in the 'easy_objective' function
+            mode="min",
+            search_alg=algo,
+            scheduler=scheduler,
+            num_samples=10 if args.smoke_test else 50,
+        ),
+        param_space={
+            "iterations": 100,
+            "x1": tune.uniform(0.0, 1.0),
+            "x2": tune.uniform(0.0, 1.0),
+            "x3": tune.uniform(0.0, 1.0),
+            "x4": tune.uniform(0.0, 1.0),
+            "x5": tune.uniform(0.0, 1.0),
+            "x6": tune.uniform(0.0, 1.0),
+        },
+    )
+    results = tuner.fit()
+    print("Best hyperparameters found were: ", results.get_best_result().config)

python/ray/tune/examples/bayesopt_example.py

@@ -0,0 +1,61 @@
+"""This example demonstrates the usage of BayesOpt with Ray Tune.
+
+It also checks that it is usable with a separate scheduler.
+
+Requires the BayesOpt library to be installed (`pip install bayesian-optimization`).
+"""
+import time
+
+from ray import tune
+from ray.tune.schedulers import AsyncHyperBandScheduler
+from ray.tune.search import ConcurrencyLimiter
+from ray.tune.search.bayesopt import BayesOptSearch
+
+
+def evaluation_fn(step, width, height):
+    return (0.1 + width * step / 100) ** (-1) + height * 0.1
+
+
+def easy_objective(config):
+    # Hyperparameters
+    width, height = config["width"], config["height"]
+
+    for step in range(config["steps"]):
+        # Iterative training function - can be any arbitrary training procedure
+        intermediate_score = evaluation_fn(step, width, height)
+        # Feed the score back back to Tune.
+        tune.report({"iterations": step, "mean_loss": intermediate_score})
+        time.sleep(0.1)
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--smoke-test", action="store_true", help="Finish quickly for testing"
+    )
+    args, _ = parser.parse_known_args()
+
+    algo = BayesOptSearch(utility_kwargs={"kind": "ucb", "kappa": 2.5, "xi": 0.0})
+    algo = ConcurrencyLimiter(algo, max_concurrent=4)
+    scheduler = AsyncHyperBandScheduler()
+    tuner = tune.Tuner(
+        easy_objective,
+        tune_config=tune.TuneConfig(
+            metric="mean_loss",
+            mode="min",
+            search_alg=algo,
+            scheduler=scheduler,
+            num_samples=10 if args.smoke_test else 1000,
+        ),
+        run_config=tune.RunConfig(name="my_exp"),
+        param_space={
+            "steps": 100,
+            "width": tune.uniform(0, 20),
+            "height": tune.uniform(-100, 100),
+        },
+    )
+    results = tuner.fit()
+
+    print("Best hyperparameters found were: ", results.get_best_result().config)

python/ray/tune/examples/bohb_example.py

@@ -0,0 +1,108 @@
+#!/usr/bin/env python
+
+"""This example demonstrates the usage of BOHB with Ray Tune.
+
+Requires the HpBandSter and ConfigSpace libraries to be installed
+(`pip install hpbandster ConfigSpace`).
+"""
+
+import json
+import os
+import time
+
+import numpy as np
+
+import ray
+from ray import tune
+from ray.tune import Trainable
+from ray.tune.schedulers.hb_bohb import HyperBandForBOHB
+from ray.tune.search.bohb import TuneBOHB
+
+
+class MyTrainableClass(Trainable):
+    """Example agent whose learning curve is a random sigmoid.
+
+    The dummy hyperparameters "width" and "height" determine the slope and
+    maximum reward value reached.
+    """
+
+    def setup(self, config):
+        self.timestep = 0
+
+    def step(self):
+        self.timestep += 1
+        v = np.tanh(float(self.timestep) / self.config.get("width", 1))
+        v *= self.config.get("height", 1)
+        time.sleep(0.1)
+        # Here we use `episode_reward_mean`, but you can also report other
+        # objectives such as loss or accuracy.
+        return {"episode_reward_mean": v}
+
+    def save_checkpoint(self, checkpoint_dir):
+        path = os.path.join(checkpoint_dir, "checkpoint")
+        with open(path, "w") as f:
+            f.write(json.dumps({"timestep": self.timestep}))
+
+    def load_checkpoint(self, checkpoint_dir):
+        path = os.path.join(checkpoint_dir, "checkpoint")
+        with open(path, "r") as f:
+            self.timestep = json.loads(f.read())["timestep"]
+
+
+if __name__ == "__main__":
+    import sys
+
+    if sys.version_info >= (3, 12):
+        # TuneBOHB is not compatible with Python 3.12
+        sys.exit(0)
+
+    ray.init(num_cpus=8)
+
+    config = {
+        "iterations": 100,
+        "width": tune.uniform(0, 20),
+        "height": tune.uniform(-100, 100),
+        "activation": tune.choice(["relu", "tanh"]),
+    }
+
+    # Optional: Pass the parameter space yourself
+    # import ConfigSpace as CS
+    # config_space = CS.ConfigurationSpace()
+    # config_space.add_hyperparameter(
+    #     CS.UniformFloatHyperparameter("width", lower=0, upper=20))
+    # config_space.add_hyperparameter(
+    #     CS.UniformFloatHyperparameter("height", lower=-100, upper=100))
+    # config_space.add_hyperparameter(
+    #     CS.CategoricalHyperparameter(
+    #         "activation", choices=["relu", "tanh"]))
+
+    max_iterations = 10
+    bohb_hyperband = HyperBandForBOHB(
+        time_attr="training_iteration",
+        max_t=max_iterations,
+        reduction_factor=2,
+        stop_last_trials=False,
+    )
+
+    bohb_search = TuneBOHB(
+        # space=config_space,  # If you want to set the space manually
+    )
+    bohb_search = tune.search.ConcurrencyLimiter(bohb_search, max_concurrent=4)
+
+    tuner = tune.Tuner(
+        MyTrainableClass,
+        run_config=tune.RunConfig(
+            name="bohb_test", stop={"training_iteration": max_iterations}
+        ),
+        tune_config=tune.TuneConfig(
+            metric="episode_reward_mean",
+            mode="max",
+            scheduler=bohb_hyperband,
+            search_alg=bohb_search,
+            num_samples=32,
+        ),
+        param_space=config,
+    )
+    results = tuner.fit()
+
+    print("Best hyperparameters found were: ", results.get_best_result().config)

python/ray/tune/examples/cifar10_pytorch.py

@@ -0,0 +1,285 @@
+# ruff: noqa
+# fmt: off
+
+# __import_begin__
+import os
+import tempfile
+from typing import Dict
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+import torchvision
+import torchvision.transforms as transforms
+from filelock import FileLock
+from torch.utils.data import random_split
+
+import ray
+from ray import tune
+from ray.tune import Checkpoint
+from ray.tune.schedulers import ASHAScheduler
+
+# __import_end__
+
+
+# __load_data_begin__
+DATA_DIR = tempfile.mkdtemp()
+
+def load_data(data_dir):
+    transform = transforms.Compose([
+        transforms.ToTensor(),
+        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
+    ])
+
+    # We add FileLock here because multiple workers will want to
+    # download data, and this may cause overwrites since
+    # DataLoader is not threadsafe.
+    with FileLock(os.path.expanduser("~/.data.lock")):
+        trainset = torchvision.datasets.CIFAR10(
+            root=data_dir, train=True, download=True, transform=transform)
+
+        testset = torchvision.datasets.CIFAR10(
+            root=data_dir, train=False, download=True, transform=transform)
+
+    return trainset, testset
+# __load_data_end__
+
+def load_test_data():
+    # Loads a fake dataset for testing so it doesn't rely on external download.
+    trainset = torchvision.datasets.FakeData(
+        128, (3, 32, 32), num_classes=10, transform=transforms.ToTensor()
+    )
+    testset = torchvision.datasets.FakeData(
+        16, (3, 32, 32), num_classes=10, transform=transforms.ToTensor()
+    )
+    return trainset, testset
+
+
+# __net_begin__
+class Net(nn.Module):
+    def __init__(self, l1=120, l2=84):
+        super(Net, self).__init__()
+        self.conv1 = nn.Conv2d(3, 6, 5)
+        self.pool = nn.MaxPool2d(2, 2)
+        self.conv2 = nn.Conv2d(6, 16, 5)
+        self.fc1 = nn.Linear(16 * 5 * 5, l1)
+        self.fc2 = nn.Linear(l1, l2)
+        self.fc3 = nn.Linear(l2, 10)
+
+    def forward(self, x):
+        x = self.pool(F.relu(self.conv1(x)))
+        x = self.pool(F.relu(self.conv2(x)))
+        x = x.view(-1, 16 * 5 * 5)
+        x = F.relu(self.fc1(x))
+        x = F.relu(self.fc2(x))
+        x = self.fc3(x)
+        return x
+# __net_end__
+
+
+# __train_begin__
+def train_cifar(config):
+    net = Net(config["l1"], config["l2"])
+
+    device = "cpu"
+    if torch.cuda.is_available():
+        device = "cuda:0"
+        if torch.cuda.device_count() > 1:
+            net = nn.DataParallel(net)
+    net.to(device)
+
+    criterion = nn.CrossEntropyLoss()
+    optimizer = optim.SGD(net.parameters(), lr=config["lr"], momentum=0.9)
+
+    # Load existing checkpoint through `get_checkpoint()` API.
+    if tune.get_checkpoint():
+        loaded_checkpoint = tune.get_checkpoint()
+        with loaded_checkpoint.as_directory() as loaded_checkpoint_dir:
+            model_state, optimizer_state = torch.load(
+                os.path.join(loaded_checkpoint_dir, "checkpoint.pt")
+            )
+            net.load_state_dict(model_state)
+            optimizer.load_state_dict(optimizer_state)
+
+    if config["smoke_test"]:
+        trainset, testset = load_test_data()
+    else:
+        trainset, testset = load_data(DATA_DIR)
+
+    test_abs = int(len(trainset) * 0.8)
+    train_subset, val_subset = random_split(
+        trainset, [test_abs, len(trainset) - test_abs])
+
+    trainloader = torch.utils.data.DataLoader(
+        train_subset,
+        batch_size=int(config["batch_size"]),
+        shuffle=True,
+        num_workers=0 if config["smoke_test"] else 8,
+    )
+    valloader = torch.utils.data.DataLoader(
+        val_subset,
+        batch_size=int(config["batch_size"]),
+        shuffle=True,
+        num_workers=0 if config["smoke_test"] else 8,
+    )
+
+    for epoch in range(10):  # loop over the dataset multiple times
+        running_loss = 0.0
+        epoch_steps = 0
+        for i, data in enumerate(trainloader):
+            # get the inputs; data is a list of [inputs, labels]
+            inputs, labels = data
+            inputs, labels = inputs.to(device), labels.to(device)
+
+            # zero the parameter gradients
+            optimizer.zero_grad()
+
+            # forward + backward + optimize
+            outputs = net(inputs)
+            loss = criterion(outputs, labels)
+            loss.backward()
+            optimizer.step()
+
+            # print statistics
+            running_loss += loss.item()
+            epoch_steps += 1
+            if i % 2000 == 1999:  # print every 2000 mini-batches
+                print("[%d, %5d] loss: %.3f" % (epoch + 1, i + 1,
+                                                running_loss / epoch_steps))
+                running_loss = 0.0
+
+        # Validation loss
+        val_loss = 0.0
+        val_steps = 0
+        total = 0
+        correct = 0
+        for i, data in enumerate(valloader, 0):
+            with torch.no_grad():
+                inputs, labels = data
+                inputs, labels = inputs.to(device), labels.to(device)
+
+                outputs = net(inputs)
+                _, predicted = torch.max(outputs.data, 1)
+                total += labels.size(0)
+                correct += (predicted == labels).sum().item()
+
+                loss = criterion(outputs, labels)
+                val_loss += loss.cpu().numpy()
+                val_steps += 1
+
+        # Here we save a checkpoint. It is automatically registered with
+        # Ray Tune and will potentially be accessed through in ``get_checkpoint()``
+        # in future iterations.
+        # Note to save a file like checkpoint, you still need to put it under a directory
+        # to construct a checkpoint.
+        with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
+            path = os.path.join(temp_checkpoint_dir, "checkpoint.pt")
+            torch.save(
+                (net.state_dict(), optimizer.state_dict()), path
+            )
+            checkpoint = Checkpoint.from_directory(temp_checkpoint_dir)
+            tune.report(
+                {"loss": (val_loss / val_steps), "accuracy": correct / total},
+                checkpoint=checkpoint,
+            )
+    print("Finished Training")
+# __train_end__
+
+
+# __test_acc_begin__
+def test_best_model(config: Dict, checkpoint: "Checkpoint", smoke_test=False):
+    best_trained_model = Net(config["l1"], config["l2"])
+    device = "cuda:0" if torch.cuda.is_available() else "cpu"
+    best_trained_model.to(device)
+
+    with checkpoint.as_directory() as checkpoint_dir:
+        checkpoint_path = os.path.join(checkpoint_dir, "checkpoint.pt")
+        model_state, optimizer_state = torch.load(checkpoint_path)
+        best_trained_model.load_state_dict(model_state)
+
+    if smoke_test:
+        _, testset = load_test_data()
+    else:
+        _, testset = load_data(DATA_DIR)
+
+    testloader = torch.utils.data.DataLoader(
+        testset, batch_size=4, shuffle=False, num_workers=2)
+
+    correct = 0
+    total = 0
+    with torch.no_grad():
+        for data in testloader:
+            images, labels = data
+            images, labels = images.to(device), labels.to(device)
+            outputs = best_trained_model(images)
+            _, predicted = torch.max(outputs.data, 1)
+            total += labels.size(0)
+            correct += (predicted == labels).sum().item()
+
+
+    print("Best trial test set accuracy: {}".format(correct / total))
+
+# __test_acc_end__
+
+# __main_begin__
+def main(num_samples=10, max_num_epochs=10, gpus_per_trial=2, smoke_test=False):
+    config = {
+        "l1": tune.sample_from(lambda _: 2 ** np.random.randint(2, 9)),
+        "l2": tune.sample_from(lambda _: 2 ** np.random.randint(2, 9)),
+        "lr": tune.loguniform(1e-4, 1e-1),
+        "batch_size": tune.choice([2, 4, 8, 16]),
+        "smoke_test": smoke_test,
+    }
+    scheduler = ASHAScheduler(
+        max_t=max_num_epochs,
+        grace_period=1,
+        reduction_factor=2)
+
+    tuner = tune.Tuner(
+        tune.with_resources(
+            tune.with_parameters(train_cifar),
+            resources={"cpu": 2, "gpu": gpus_per_trial},
+        ),
+        tune_config=tune.TuneConfig(
+            metric="loss",
+            mode="min",
+            num_samples=num_samples,
+            scheduler=scheduler
+        ),
+        param_space=config,
+    )
+    results = tuner.fit()
+    best_result = results.get_best_result("loss", "min")
+    print("Best trial config: {}".format(best_result.config))
+    print("Best trial final validation loss: {}".format(
+        best_result.metrics["loss"]))
+    print("Best trial final validation accuracy: {}".format(
+        best_result.metrics["accuracy"]))
+
+    test_best_model(best_result.config, best_result.checkpoint, smoke_test=smoke_test)
+
+
+# __main_end__
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--smoke-test", action="store_true", help="Finish quickly for testing")
+    parser.add_argument(
+        "--ray-address",
+        help="Address of Ray cluster for seamless distributed execution.",
+        required=False)
+    args, _ = parser.parse_known_args()
+
+    if args.smoke_test:
+        ray.init(num_cpus=2)
+        main(num_samples=1, max_num_epochs=1, gpus_per_trial=0, smoke_test=True)
+    else:
+        ray.init(args.ray_address)
+        # Change this to activate training on GPUs
+        main(num_samples=10, max_num_epochs=10, gpus_per_trial=0)

python/ray/tune/examples/custom_checkpointing_with_callback.py

@@ -0,0 +1,221 @@
+# Example demonstrating how to use SHOULD_CHECKPOINT in a tuner callback
+# for smart checkpointing logic. This shows how to trigger checkpointing from
+# callbacks based on training progress rather than fixed intervals.
+
+import argparse
+import json
+import os
+import time
+
+from ray import tune
+from ray.tune import Callback
+from ray.tune.result import SHOULD_CHECKPOINT
+
+# Hint: SHOULD_CHECKPOINT is an alias of the string "should_checkpoint"
+
+
+# Some dummy function
+def evaluation_fn(step, width, height):
+    time.sleep(0.1)
+    return (0.1 + width * step / 100) ** (-1) + height * 0.1
+
+
+class SmartCheckpointCallback(Callback):
+    """Custom callback that triggers checkpointing by updating the result dict.
+
+    This callback demonstrates checkpointing logic beyond
+    simple periodic checkpointing. It checkpoints based on performance improvements
+    or when the loss becomes unstable.
+
+    Args:
+        checkpoint_on_improvement: Checkpoint when loss improves significantly
+        checkpoint_on_instability: Checkpoint when loss becomes unstable
+    """
+
+    def __init__(
+        self,
+        *,
+        checkpoint_on_improvement: bool = True,
+        checkpoint_on_instability: bool = True,
+    ):
+        self.checkpoint_on_improvement = checkpoint_on_improvement
+        self.checkpoint_on_instability = checkpoint_on_instability
+        self.best_loss_per_trial = {}
+        self.recent_losses_per_trial = {}
+
+    def on_trial_result(self, iteration, trials, trial, result, **info):
+        """Called after receiving a result from the trainable.
+
+        This hook implements intelligent checkpointing logic:
+        1. Checkpoint when we see significant improvement
+        2. Checkpoint when loss becomes unstable (variance increases)
+        3. Always checkpoint at specific milestones (every 10 steps)
+        """
+        trial_id = trial.trial_id
+        current_loss = result.get("mean_loss", float("inf"))
+        current_step = result.get("iterations", 0)
+
+        # Initialize tracking for this trial
+        if trial_id not in self.best_loss_per_trial:
+            self.best_loss_per_trial[trial_id] = float("inf")
+            self.recent_losses_per_trial[trial_id] = []
+
+        should_checkpoint = False
+        reason = ""
+
+        # 1. Checkpoint every 10 steps as a baseline
+        if current_step > 0 and current_step % 10 == 0:
+            should_checkpoint = True
+            reason = f"milestone at step {current_step}"
+
+        # 2. Checkpoint on significant improvement
+        if self.checkpoint_on_improvement:
+            if (
+                current_loss < self.best_loss_per_trial[trial_id] * 0.9
+            ):  # 10% improvement
+                should_checkpoint = True
+                reason = f"significant improvement: {current_loss:.4f} < {self.best_loss_per_trial[trial_id]:.4f}"
+                self.best_loss_per_trial[trial_id] = current_loss
+
+        # 3. Checkpoint on instability (high variance in recent losses)
+        if self.checkpoint_on_instability and current_step > 5:
+            recent_losses = self.recent_losses_per_trial[trial_id]
+            recent_losses.append(current_loss)
+            if len(recent_losses) > 5:
+                recent_losses.pop(0)  # Keep only last 5 losses
+
+            if len(recent_losses) == 5:
+                variance = (
+                    sum((x - sum(recent_losses) / 5) ** 2 for x in recent_losses) / 5
+                )
+                if variance > 0.1:  # High variance threshold
+                    should_checkpoint = True
+                    reason = f"instability detected: variance={variance:.4f}"
+        else:
+            # Track recent losses
+            recent_losses = self.recent_losses_per_trial[trial_id]
+            recent_losses.append(current_loss)
+            if len(recent_losses) > 5:
+                recent_losses.pop(0)
+
+        if should_checkpoint:
+            print(
+                f"Callback requesting checkpoint for trial {trial_id} at step {current_step}: {reason}"
+            )
+            result[SHOULD_CHECKPOINT] = True
+
+
+class OptimizationTrainable(tune.Trainable):
+    """A simple trainable that demonstrates automatic checkpointing with callbacks"""
+
+    def setup(self, config):
+        """Initialize the trainable"""
+        self.current_step = 0
+        self.width = config["width"]
+        self.height = config["height"]
+
+    def step(self):
+        """Perform one step of training"""
+        intermediate_score = evaluation_fn(self.current_step, self.width, self.height)
+        self.current_step += 1
+
+        return {
+            "iterations": self.current_step,
+            "mean_loss": intermediate_score,
+            "step": self.current_step,  # For tracking
+        }
+
+    def save_checkpoint(self, checkpoint_dir):
+        """Save checkpoint
+
+        Called automatically by Tune when SHOULD_CHECKPOINT is in the result
+        """
+        checkpoint_path = os.path.join(checkpoint_dir, "checkpoint.json")
+        with open(checkpoint_path, "w") as f:
+            json.dump(
+                {"step": self.current_step, "width": self.width, "height": self.height},
+                f,
+            )
+        print(f"Checkpoint saved at step {self.current_step}")
+
+    def load_checkpoint(self, checkpoint):
+        """Load checkpoint - called automatically by Tune during restoration"""
+        checkpoint_path = os.path.join(checkpoint, "checkpoint.json")
+        with open(checkpoint_path, "r") as f:
+            state = json.load(f)
+        self.current_step = state["step"]
+        self.width = state["width"]
+        self.height = state["height"]
+        print(f"Checkpoint loaded from step {self.current_step}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--smoke-test", action="store_true", help="Finish quickly for testing"
+    )
+    args, _ = parser.parse_known_args()
+
+    print(
+        "=" * 60,
+        "Ray Tune Example: Smart Checkpointing with custom SHOULD_CHECKPOINT key",
+        "=" * 60,
+        "",
+        "This example demonstrates how to set the SHOULD_CHECKPOINT key in a callback",
+        "to implement intelligent checkpointing based on training progress.",
+        "",
+        "Key features:",
+        "- Callback-driven checkpointing by setting result[SHOULD_CHECKPOINT] = True",
+        "- Checkpoints triggered by performance improvements",
+        "- Milestone-based checkpointing every 10 steps",
+        "- Instability detection (high variance in recent losses)",
+        "- Automatic checkpoint save/load via class trainable",
+        sep="\n",
+    )
+
+    # Create the smart checkpoint callback
+    checkpoint_callback = SmartCheckpointCallback(
+        checkpoint_on_improvement=True, checkpoint_on_instability=True
+    )
+
+    tuner = tune.Tuner(
+        OptimizationTrainable,
+        run_config=tune.RunConfig(
+            name="smart_checkpoint_test",
+            stop={"training_iteration": 1 if args.smoke_test else 20},
+            callbacks=[checkpoint_callback],  # Add our custom callback
+            # Disable automatic periodic checkpointing to show callback control
+            checkpoint_config=tune.CheckpointConfig(
+                checkpoint_frequency=0,  # Disable periodic checkpointing
+                checkpoint_at_end=True,  # Still checkpoint at the end
+            ),
+        ),
+        tune_config=tune.TuneConfig(
+            metric="mean_loss",
+            mode="min",
+            num_samples=3,
+        ),
+        param_space={
+            "width": tune.randint(10, 100),
+            "height": tune.loguniform(10, 100),
+        },
+    )
+
+    print(
+        "Starting hyperparameter tuning with smart checkpointing...",
+        "Watch for checkpoint messages triggered by the callback!",
+        sep="\n",
+    )
+
+    results = tuner.fit()
+    best_result = results.get_best_result()
+    print(
+        "\n" + "=" * 60,
+        "RESULTS",
+        "=" * 60,
+        f"Best hyperparameters: {best_result.config}",
+        f"Best checkpoint: {best_result.checkpoint}",
+        "",
+        "The checkpoints were triggered by the SmartCheckpointCallback",
+        sep="\n",
+    )

python/ray/tune/examples/custom_func_checkpointing.py

@@ -0,0 +1,70 @@
+# If want to use checkpointing with a custom training function (not a Ray
+# integration like PyTorch or Tensorflow), your function can read/write
+# checkpoint through the ``ray.tune.report(metrics, checkpoint=...)`` API.
+import argparse
+import json
+import os
+import tempfile
+import time
+
+from ray import tune
+from ray.tune import Checkpoint
+
+
+def evaluation_fn(step, width, height):
+    time.sleep(0.1)
+    return (0.1 + width * step / 100) ** (-1) + height * 0.1
+
+
+def train_func(config):
+    step = 0
+    width, height = config["width"], config["height"]
+
+    checkpoint = tune.get_checkpoint()
+    if checkpoint:
+        with checkpoint.as_directory() as checkpoint_dir:
+            with open(os.path.join(checkpoint_dir, "checkpoint.json")) as f:
+                state = json.load(f)
+            step = state["step"] + 1
+
+    for current_step in range(step, 100):
+        intermediate_score = evaluation_fn(current_step, width, height)
+
+        with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
+            with open(os.path.join(temp_checkpoint_dir, "checkpoint.json"), "w") as f:
+                json.dump({"step": current_step}, f)
+            tune.report(
+                {"iterations": current_step, "mean_loss": intermediate_score},
+                checkpoint=Checkpoint.from_directory(temp_checkpoint_dir),
+            )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--smoke-test", action="store_true", help="Finish quickly for testing"
+    )
+    args, _ = parser.parse_known_args()
+
+    tuner = tune.Tuner(
+        train_func,
+        run_config=tune.RunConfig(
+            name="hyperband_test",
+            stop={"training_iteration": 1 if args.smoke_test else 10},
+        ),
+        tune_config=tune.TuneConfig(
+            metric="mean_loss",
+            mode="min",
+            num_samples=5,
+        ),
+        param_space={
+            "steps": 10,
+            "width": tune.randint(10, 100),
+            "height": tune.loguniform(10, 100),
+        },
+    )
+    results = tuner.fit()
+    best_result = results.get_best_result()
+    print("Best hyperparameters: ", best_result.config)
+    best_checkpoint = best_result.checkpoint
+    print("Best checkpoint: ", best_checkpoint)

python/ray/tune/examples/hyperband_example.py

@@ -0,0 +1,44 @@
+#!/usr/bin/env python
+
+import argparse
+
+import ray
+from ray import tune
+from ray.tune.schedulers import HyperBandScheduler
+from ray.tune.utils.mock_trainable import MyTrainableClass
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--smoke-test", action="store_true", help="Finish quickly for testing"
+    )
+    args, _ = parser.parse_known_args()
+
+    ray.init(num_cpus=4 if args.smoke_test else None)
+
+    # Hyperband early stopping, configured with `episode_reward_mean` as the
+    # objective and `training_iteration` as the time unit,
+    # which is automatically filled by Tune.
+    hyperband = HyperBandScheduler(time_attr="training_iteration", max_t=200)
+
+    tuner = tune.Tuner(
+        MyTrainableClass,
+        run_config=tune.RunConfig(
+            name="hyperband_test",
+            stop={"training_iteration": 1 if args.smoke_test else 200},
+            verbose=1,
+            failure_config=tune.FailureConfig(
+                fail_fast=True,
+            ),
+        ),
+        tune_config=tune.TuneConfig(
+            num_samples=20 if args.smoke_test else 200,
+            metric="episode_reward_mean",
+            mode="max",
+            scheduler=hyperband,
+        ),
+        param_space={"width": tune.randint(10, 90), "height": tune.randint(0, 100)},
+    )
+    results = tuner.fit()
+
+    print("Best hyperparameters found were: ", results.get_best_result().config)

python/ray/tune/examples/hyperband_function_example.py

@@ -0,0 +1,76 @@
+#!/usr/bin/env python
+
+import argparse
+import json
+import os
+import tempfile
+
+import numpy as np
+
+import ray
+from ray import tune
+from ray.tune import Checkpoint
+from ray.tune.schedulers import HyperBandScheduler
+
+
+def train_func(config):
+    step = 0
+    checkpoint = tune.get_checkpoint()
+    if checkpoint:
+        with checkpoint.as_directory() as checkpoint_dir:
+            with open(os.path.join(checkpoint_dir, "checkpoint.json")) as f:
+                step = json.load(f)["timestep"] + 1
+
+    for timestep in range(step, 100):
+        v = np.tanh(float(timestep) / config.get("width", 1))
+        v *= config.get("height", 1)
+
+        # Checkpoint the state of the training every 3 steps
+        # Note that this is only required for certain schedulers
+        with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
+            checkpoint = None
+            if timestep % 3 == 0:
+                with open(
+                    os.path.join(temp_checkpoint_dir, "checkpoint.json"), "w"
+                ) as f:
+                    json.dump({"timestep": timestep}, f)
+                checkpoint = Checkpoint.from_directory(temp_checkpoint_dir)
+
+            # Here we use `episode_reward_mean`, but you can also report other
+            # objectives such as loss or accuracy.
+            tune.report({"episode_reward_mean": v}, checkpoint=checkpoint)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--smoke-test", action="store_true", help="Finish quickly for testing"
+    )
+    args, _ = parser.parse_known_args()
+
+    ray.init(num_cpus=4 if args.smoke_test else None)
+
+    # Hyperband early stopping, configured with `episode_reward_mean` as the
+    # objective and `training_iteration` as the time unit,
+    # which is automatically filled by Tune.
+    hyperband = HyperBandScheduler(max_t=200)
+
+    tuner = tune.Tuner(
+        train_func,
+        run_config=tune.RunConfig(
+            name="hyperband_test",
+            stop={"training_iteration": 10 if args.smoke_test else 99999},
+            failure_config=tune.FailureConfig(
+                fail_fast=True,
+            ),
+        ),
+        tune_config=tune.TuneConfig(
+            num_samples=20,
+            metric="episode_reward_mean",
+            mode="max",
+            scheduler=hyperband,
+        ),
+        param_space={"height": tune.uniform(0, 100)},
+    )
+    results = tuner.fit()
+    print("Best hyperparameters found were: ", results.get_best_result().config)

python/ray/tune/examples/hyperopt_conditional_search_space_example.py

@@ -0,0 +1,107 @@
+"""This example demonstrates the usage of conditional search spaces with Tune.
+
+It also checks that it is usable with a separate scheduler.
+
+Requires the HyperOpt library to be installed (`pip install hyperopt`).
+
+For an example of using a Tune search space, see
+:doc:`/tune/examples/hyperopt_example`.
+"""
+
+import time
+
+from hyperopt import hp
+
+import ray
+from ray import tune
+from ray.tune.schedulers import AsyncHyperBandScheduler
+from ray.tune.search import ConcurrencyLimiter
+from ray.tune.search.hyperopt import HyperOptSearch
+
+
+def f_unpack_dict(dct: dict) -> dict:
+    """Unpacks all sub-dictionaries in given dictionary recursively.
+    There should be no duplicated keys across all nested
+    subdictionaries, or some instances will be lost without warning
+
+    Source: https://www.kaggle.com/fanvacoolt/tutorial-on-hyperopt
+
+    Args:
+        dct: dictionary to unpack
+
+    Returns:
+        dict: unpacked dictionary
+    """
+
+    res = {}
+    for k, v in dct.items():
+        if isinstance(v, dict):
+            res = {**res, **f_unpack_dict(v)}
+        else:
+            res[k] = v
+
+    return res
+
+
+def evaluation_fn(step, width, height, mult=1):
+    return (0.1 + width * step / 100) ** (-1) + height * 0.1 * mult
+
+
+def easy_objective(config_in):
+    # Hyperparameters
+    config = f_unpack_dict(config_in)
+    width, height, mult = config["width"], config["height"], config.get("mult", 1)
+    print(config)
+
+    for step in range(config["steps"]):
+        # Iterative training function - can be any arbitrary training procedure
+        intermediate_score = evaluation_fn(step, width, height, mult)
+        # Feed the score back back to Tune.
+        tune.report({"iterations": step, "mean_loss": intermediate_score})
+        time.sleep(0.1)
+
+
+config_space = {
+    "activation": hp.choice(
+        "activation",
+        [
+            {"activation": "relu", "mult": hp.uniform("mult", 1, 2)},
+            {"activation": "tanh"},
+        ],
+    ),
+    "width": hp.uniform("width", 0, 20),
+    "height": hp.uniform("heright", -100, 100),
+    "steps": 100,
+}
+
+
+def run_hyperopt_tune(config_dict=config_space, smoke_test=False):
+    algo = HyperOptSearch(space=config_dict, metric="mean_loss", mode="min")
+    algo = ConcurrencyLimiter(algo, max_concurrent=4)
+    scheduler = AsyncHyperBandScheduler()
+    tuner = tune.Tuner(
+        easy_objective,
+        tune_config=tune.TuneConfig(
+            metric="mean_loss",
+            mode="min",
+            search_alg=algo,
+            scheduler=scheduler,
+            num_samples=10 if smoke_test else 100,
+        ),
+    )
+    results = tuner.fit()
+    print("Best hyperparameters found were: ", results.get_best_result().config)
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--smoke-test", action="store_true", help="Finish quickly for testing"
+    )
+    args, _ = parser.parse_known_args()
+
+    ray.init(configure_logging=False)
+
+    run_hyperopt_tune(smoke_test=args.smoke_test)

python/ray/tune/examples/lightgbm_example.py

@@ -0,0 +1,103 @@
+import lightgbm as lgb
+import sklearn.datasets
+import sklearn.metrics
+from sklearn.model_selection import train_test_split
+
+from ray import tune
+from ray.tune.integration.lightgbm import TuneReportCheckpointCallback
+from ray.tune.schedulers import ASHAScheduler
+
+
+def train_breast_cancer(config: dict):
+    # This is a simple training function to be passed into Tune
+
+    # Load dataset
+    data, target = sklearn.datasets.load_breast_cancer(return_X_y=True)
+
+    # Split into train and test set
+    train_x, test_x, train_y, test_y = train_test_split(data, target, test_size=0.25)
+
+    # Build input Datasets for LightGBM
+    train_set = lgb.Dataset(train_x, label=train_y)
+    test_set = lgb.Dataset(test_x, label=test_y)
+
+    # Train the classifier, using the Tune callback
+    lgb.train(
+        config,
+        train_set,
+        valid_sets=[test_set],
+        valid_names=["eval"],
+        callbacks=[
+            TuneReportCheckpointCallback(
+                {
+                    "binary_error": "eval-binary_error",
+                    "binary_logloss": "eval-binary_logloss",
+                }
+            )
+        ],
+    )
+
+
+def train_breast_cancer_cv(config: dict):
+    # This is a simple training function to be passed into Tune, using
+    # lightgbm's cross validation functionality
+
+    # Load dataset
+    data, target = sklearn.datasets.load_breast_cancer(return_X_y=True)
+
+    train_set = lgb.Dataset(data, label=target)
+
+    # Run CV, using the Tune callback
+    lgb.cv(
+        config,
+        train_set,
+        stratified=True,
+        # Checkpointing is not supported for CV
+        # LightGBM aggregates metrics over folds automatically
+        # with the cv_agg key. Both mean and standard deviation
+        # are provided.
+        callbacks=[
+            TuneReportCheckpointCallback(
+                {
+                    "binary_error": "valid-binary_error-mean",
+                    "binary_logloss": "valid-binary_logloss-mean",
+                    "binary_error_stdv": "valid-binary_error-stdv",
+                    "binary_logloss_stdv": "valid-binary_logloss-stdv",
+                },
+                frequency=0,
+            )
+        ],
+    )
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--use-cv", action="store_true", help="Use `lgb.cv` instead of `lgb.train`."
+    )
+    args, _ = parser.parse_known_args()
+
+    config = {
+        "objective": "binary",
+        "metric": ["binary_error", "binary_logloss"],
+        "verbose": -1,
+        "boosting_type": tune.grid_search(["gbdt", "dart"]),
+        "num_leaves": tune.randint(10, 1000),
+        "learning_rate": tune.loguniform(1e-8, 1e-1),
+    }
+
+    tuner = tune.Tuner(
+        train_breast_cancer if not args.use_cv else train_breast_cancer_cv,
+        tune_config=tune.TuneConfig(
+            metric="binary_error",
+            mode="min",
+            num_samples=2,
+            scheduler=ASHAScheduler(),
+        ),
+        param_space=config,
+    )
+    results = tuner.fit()
+
+    print("Best hyperparameters found were: ", results.get_best_result().config)

python/ray/tune/examples/logging_example.py

@@ -0,0 +1,64 @@
+#!/usr/bin/env python
+
+import argparse
+import time
+
+from ray import tune
+from ray.tune.logger import LoggerCallback
+
+
+class TestLoggerCallback(LoggerCallback):
+    def on_trial_result(self, iteration, trials, trial, result, **info):
+        print(f"TestLogger for trial {trial}: {result}")
+
+
+def trial_str_creator(trial):
+    return "{}_{}_123".format(trial.trainable_name, trial.trial_id)
+
+
+def evaluation_fn(step, width, height):
+    time.sleep(0.1)
+    return (0.1 + width * step / 100) ** (-1) + height * 0.1
+
+
+def easy_objective(config):
+    # Hyperparameters
+    width, height = config["width"], config["height"]
+
+    for step in range(config["steps"]):
+        # Iterative training function - can be any arbitrary training procedure
+        intermediate_score = evaluation_fn(step, width, height)
+        # Feed the score back back to Tune.
+        tune.report({"iterations": step, "mean_loss": intermediate_score})
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--smoke-test", action="store_true", help="Finish quickly for testing"
+    )
+    args, _ = parser.parse_known_args()
+
+    tuner = tune.Tuner(
+        easy_objective,
+        run_config=tune.RunConfig(
+            name="hyperband_test",
+            callbacks=[TestLoggerCallback()],
+            stop={"training_iteration": 1 if args.smoke_test else 100},
+        ),
+        tune_config=tune.TuneConfig(
+            metric="mean_loss",
+            mode="min",
+            num_samples=5,
+            trial_name_creator=trial_str_creator,
+            trial_dirname_creator=trial_str_creator,
+        ),
+        param_space={
+            "steps": 100,
+            "width": tune.randint(10, 100),
+            "height": tune.loguniform(10, 100),
+        },
+    )
+    results = tuner.fit()
+
+    print("Best hyperparameters: ", results.get_best_result().config)

python/ray/tune/examples/mlflow_example.py

@@ -0,0 +1,128 @@
+#!/usr/bin/env python
+"""Examples using MLfowLoggerCallback and setup_mlflow.
+"""
+import os
+import tempfile
+import time
+
+import mlflow
+
+from ray import tune
+from ray.air.integrations.mlflow import MLflowLoggerCallback, setup_mlflow
+
+
+def evaluation_fn(step, width, height):
+    return (0.1 + width * step / 100) ** (-1) + height * 0.1
+
+
+def train_function(config):
+    # Hyperparameters
+    width, height = config["width"], config["height"]
+
+    for step in range(config.get("steps", 100)):
+        # Iterative training function - can be any arbitrary training procedure
+        intermediate_score = evaluation_fn(step, width, height)
+        # Feed the score back to Tune.
+        tune.report({"iterations": step, "mean_loss": intermediate_score})
+        time.sleep(0.1)
+
+
+def tune_with_callback(mlflow_tracking_uri, finish_fast=False):
+
+    tuner = tune.Tuner(
+        train_function,
+        run_config=tune.RunConfig(
+            name="mlflow",
+            callbacks=[
+                MLflowLoggerCallback(
+                    tracking_uri=mlflow_tracking_uri,
+                    experiment_name="example",
+                    save_artifact=True,
+                )
+            ],
+        ),
+        tune_config=tune.TuneConfig(
+            num_samples=5,
+        ),
+        param_space={
+            "width": tune.randint(10, 100),
+            "height": tune.randint(0, 100),
+            "steps": 5 if finish_fast else 100,
+        },
+    )
+    tuner.fit()
+
+
+def train_function_mlflow(config):
+    setup_mlflow(config)
+
+    # Hyperparameters
+    width, height = config["width"], config["height"]
+
+    for step in range(config.get("steps", 100)):
+        # Iterative training function - can be any arbitrary training procedure
+        intermediate_score = evaluation_fn(step, width, height)
+        # Log the metrics to mlflow
+        mlflow.log_metrics(dict(mean_loss=intermediate_score), step=step)
+        # Feed the score back to Tune.
+        tune.report({"iterations": step, "mean_loss": intermediate_score})
+        time.sleep(0.1)
+
+
+def tune_with_setup(mlflow_tracking_uri, finish_fast=False):
+    # Set the experiment, or create a new one if does not exist yet.
+    mlflow.set_tracking_uri(mlflow_tracking_uri)
+    mlflow.set_experiment(experiment_name="mixin_example")
+    tuner = tune.Tuner(
+        train_function_mlflow,
+        run_config=tune.RunConfig(
+            name="mlflow",
+        ),
+        tune_config=tune.TuneConfig(
+            num_samples=5,
+        ),
+        param_space={
+            "width": tune.randint(10, 100),
+            "height": tune.randint(0, 100),
+            "steps": 5 if finish_fast else 100,
+            "mlflow": {
+                "experiment_name": "mixin_example",
+                "tracking_uri": mlflow.get_tracking_uri(),
+            },
+        },
+    )
+    tuner.fit()
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--smoke-test", action="store_true", help="Finish quickly for testing"
+    )
+    parser.add_argument(
+        "--tracking-uri",
+        type=str,
+        help="The tracking URI for the MLflow tracking server.",
+    )
+    args, _ = parser.parse_known_args()
+
+    if args.smoke_test:
+        mlflow_tracking_uri = os.path.join(tempfile.gettempdir(), "mlruns")
+    else:
+        mlflow_tracking_uri = args.tracking_uri
+
+    tune_with_callback(mlflow_tracking_uri, finish_fast=args.smoke_test)
+    if not args.smoke_test:
+        df = mlflow.search_runs(
+            [mlflow.get_experiment_by_name("example").experiment_id]
+        )
+        print(df)
+
+    tune_with_setup(mlflow_tracking_uri, finish_fast=args.smoke_test)
+    if not args.smoke_test:
+        df = mlflow.search_runs(
+            [mlflow.get_experiment_by_name("mixin_example").experiment_id]
+        )
+        print(df)

python/ray/tune/examples/mlflow_ptl.py

@@ -0,0 +1,105 @@
+"""An example showing how to use Pytorch Lightning training, Ray Tune
+HPO, and MLflow autologging all together."""
+
+import os
+import tempfile
+
+import lightning.pytorch as pl
+import mlflow
+
+from ray import tune
+from ray.air.integrations.mlflow import setup_mlflow
+from ray.tune.examples.mnist_ptl_mini import LightningMNISTClassifier, MNISTDataModule
+from ray.tune.integration.pytorch_lightning import TuneReportCallback
+
+
+def train_mnist_tune(config, data_dir=None, num_epochs=10, num_gpus=0):
+    setup_mlflow(
+        config,
+        experiment_name=config.get("experiment_name", None),
+        tracking_uri=config.get("tracking_uri", None),
+    )
+
+    model = LightningMNISTClassifier(config, data_dir)
+    dm = MNISTDataModule(
+        data_dir=data_dir, num_workers=1, batch_size=config["batch_size"]
+    )
+    metrics = {"loss": "ptl/val_loss", "acc": "ptl/val_accuracy"}
+    mlflow.pytorch.autolog()
+    trainer = pl.Trainer(
+        max_epochs=num_epochs,
+        gpus=num_gpus,
+        progress_bar_refresh_rate=0,
+        callbacks=[TuneReportCallback(metrics, on="validation_end")],
+    )
+    trainer.fit(model, dm)
+
+
+def tune_mnist(
+    num_samples=10,
+    num_epochs=10,
+    gpus_per_trial=0,
+    tracking_uri=None,
+    experiment_name="ptl_autologging_example",
+):
+    data_dir = os.path.join(tempfile.gettempdir(), "mnist_data_")
+    # Download data
+    MNISTDataModule(data_dir=data_dir, batch_size=32).prepare_data()
+
+    # Set the MLflow experiment, or create it if it does not exist.
+    mlflow.set_tracking_uri(tracking_uri)
+    mlflow.set_experiment(experiment_name)
+
+    config = {
+        "layer_1": tune.choice([32, 64, 128]),
+        "layer_2": tune.choice([64, 128, 256]),
+        "lr": tune.loguniform(1e-4, 1e-1),
+        "batch_size": tune.choice([32, 64, 128]),
+        "experiment_name": experiment_name,
+        "tracking_uri": mlflow.get_tracking_uri(),
+        "data_dir": os.path.join(tempfile.gettempdir(), "mnist_data_"),
+        "num_epochs": num_epochs,
+    }
+
+    trainable = tune.with_parameters(
+        train_mnist_tune,
+        data_dir=data_dir,
+        num_epochs=num_epochs,
+        num_gpus=gpus_per_trial,
+    )
+
+    tuner = tune.Tuner(
+        tune.with_resources(trainable, resources={"cpu": 1, "gpu": gpus_per_trial}),
+        tune_config=tune.TuneConfig(
+            metric="loss",
+            mode="min",
+            num_samples=num_samples,
+        ),
+        run_config=tune.RunConfig(
+            name="tune_mnist",
+        ),
+        param_space=config,
+    )
+    results = tuner.fit()
+
+    print("Best hyperparameters found were: ", results.get_best_result().config)
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--smoke-test", action="store_true", help="Finish quickly for testing"
+    )
+    args, _ = parser.parse_known_args()
+
+    if args.smoke_test:
+        tune_mnist(
+            num_samples=1,
+            num_epochs=1,
+            gpus_per_trial=0,
+            tracking_uri=os.path.join(tempfile.gettempdir(), "mlruns"),
+        )
+    else:
+        tune_mnist(num_samples=10, num_epochs=10, gpus_per_trial=0)

python/ray/tune/examples/mnist_ptl_mini.py

@@ -0,0 +1,166 @@
+import math
+import os
+
+import lightning.pytorch as pl
+import torch
+from datasets import load_dataset
+from filelock import FileLock
+from torch.nn import functional as F
+from torch.utils.data import DataLoader
+from torchmetrics import Accuracy
+from torchvision import transforms
+
+from ray import tune
+from ray.tune.integration.pytorch_lightning import TuneReportCheckpointCallback
+
+PATH_DATASETS = os.environ.get("PATH_DATASETS", ".")
+
+
+class MNISTDataModule(pl.LightningDataModule):
+    def __init__(self, batch_size: int, data_dir: str = PATH_DATASETS):
+        super().__init__()
+        self.data_dir = data_dir
+        self.transform = transforms.Compose(
+            [
+                transforms.ToTensor(),
+                transforms.Normalize((0.1307,), (0.3081,)),
+            ]
+        )
+        self.batch_size = batch_size
+        self.dims = (1, 28, 28)
+        self.num_classes = 10
+
+    def prepare_data(self):
+        # download
+        with FileLock(os.path.expanduser("~/.data.lock")):
+            load_dataset("ylecun/mnist", cache_dir=self.data_dir)
+
+    def setup(self, stage=None):
+        dataset = load_dataset("ylecun/mnist", cache_dir=self.data_dir)
+
+        def transform_fn(sample):
+            return (self.transform(sample["image"]), sample["label"])
+
+        self.mnist_train = [transform_fn(sample) for sample in dataset["train"]]
+        self.mnist_val = [transform_fn(sample) for sample in dataset["test"]]
+
+    def train_dataloader(self):
+        return DataLoader(self.mnist_train, batch_size=self.batch_size)
+
+    def val_dataloader(self):
+        return DataLoader(self.mnist_val, batch_size=self.batch_size)
+
+
+class LightningMNISTClassifier(pl.LightningModule):
+    def __init__(self, config, data_dir=None):
+        super(LightningMNISTClassifier, self).__init__()
+
+        self.data_dir = data_dir or os.getcwd()
+        self.lr = config["lr"]
+        layer_1, layer_2 = config["layer_1"], config["layer_2"]
+        self.batch_size = config["batch_size"]
+
+        # mnist images are (1, 28, 28) (channels, width, height)
+        self.layer_1 = torch.nn.Linear(28 * 28, layer_1)
+        self.layer_2 = torch.nn.Linear(layer_1, layer_2)
+        self.layer_3 = torch.nn.Linear(layer_2, 10)
+        self.accuracy = Accuracy(task="multiclass", num_classes=10, top_k=1)
+
+    def forward(self, x):
+        batch_size, channels, width, height = x.size()
+        x = x.view(batch_size, -1)
+        x = self.layer_1(x)
+        x = torch.relu(x)
+        x = self.layer_2(x)
+        x = torch.relu(x)
+        x = self.layer_3(x)
+        x = torch.log_softmax(x, dim=1)
+        return x
+
+    def configure_optimizers(self):
+        return torch.optim.Adam(self.parameters(), lr=self.lr)
+
+    def training_step(self, train_batch, batch_idx):
+        x, y = train_batch
+        logits = self.forward(x)
+        loss = F.nll_loss(logits, y)
+        acc = self.accuracy(logits, y)
+        self.log("ptl/train_loss", loss)
+        self.log("ptl/train_accuracy", acc)
+        return loss
+
+    def validation_step(self, val_batch, batch_idx):
+        x, y = val_batch
+        logits = self.forward(x)
+        loss = F.nll_loss(logits, y)
+        acc = self.accuracy(logits, y)
+        return {"val_loss": loss, "val_accuracy": acc}
+
+    def validation_epoch_end(self, outputs):
+        avg_loss = torch.stack([x["val_loss"] for x in outputs]).mean()
+        avg_acc = torch.stack([x["val_accuracy"] for x in outputs]).mean()
+        self.log("ptl/val_loss", avg_loss)
+        self.log("ptl/val_accuracy", avg_acc)
+
+
+def train_mnist_tune(config, num_epochs=10, num_gpus=0):
+    data_dir = os.path.abspath("./data")
+    model = LightningMNISTClassifier(config, data_dir)
+    with FileLock(os.path.expanduser("~/.data.lock")):
+        dm = MNISTDataModule(data_dir=data_dir, batch_size=config["batch_size"])
+    metrics = {"loss": "ptl/val_loss", "acc": "ptl/val_accuracy"}
+    trainer = pl.Trainer(
+        max_epochs=num_epochs,
+        # If fractional GPUs passed in, convert to int.
+        gpus=math.ceil(num_gpus),
+        enable_progress_bar=False,
+        callbacks=[
+            TuneReportCheckpointCallback(
+                metrics, on="validation_end", save_checkpoints=False
+            )
+        ],
+    )
+    trainer.fit(model, dm)
+
+
+def tune_mnist(num_samples=10, num_epochs=10, gpus_per_trial=0):
+    config = {
+        "layer_1": tune.choice([32, 64, 128]),
+        "layer_2": tune.choice([64, 128, 256]),
+        "lr": tune.loguniform(1e-4, 1e-1),
+        "batch_size": tune.choice([32, 64, 128]),
+    }
+
+    trainable = tune.with_parameters(
+        train_mnist_tune, num_epochs=num_epochs, num_gpus=gpus_per_trial
+    )
+    tuner = tune.Tuner(
+        tune.with_resources(trainable, resources={"cpu": 1, "gpu": gpus_per_trial}),
+        tune_config=tune.TuneConfig(
+            metric="loss",
+            mode="min",
+            num_samples=num_samples,
+        ),
+        run_config=tune.RunConfig(
+            name="tune_mnist",
+        ),
+        param_space=config,
+    )
+    results = tuner.fit()
+
+    print("Best hyperparameters found were: ", results.get_best_result().config)
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--smoke-test", action="store_true", help="Finish quickly for testing"
+    )
+    args, _ = parser.parse_known_args()
+
+    if args.smoke_test:
+        tune_mnist(num_samples=1, num_epochs=1, gpus_per_trial=0)
+    else:
+        tune_mnist(num_samples=10, num_epochs=10, gpus_per_trial=0)

python/ray/tune/examples/mnist_pytorch.py

@@ -0,0 +1,161 @@
+# Original Code here:
+# https://github.com/pytorch/examples/blob/master/mnist/main.py
+
+import argparse
+import os
+import tempfile
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from filelock import FileLock
+from torchvision import datasets, transforms
+
+import ray
+from ray import tune
+from ray.tune import Checkpoint
+from ray.tune.schedulers import AsyncHyperBandScheduler
+
+# Change these values if you want the training to run quicker or slower.
+EPOCH_SIZE = 512
+TEST_SIZE = 256
+
+
+class ConvNet(nn.Module):
+    def __init__(self):
+        super(ConvNet, self).__init__()
+        self.conv1 = nn.Conv2d(1, 3, kernel_size=3)
+        self.fc = nn.Linear(192, 10)
+
+    def forward(self, x):
+        x = F.relu(F.max_pool2d(self.conv1(x), 3))
+        x = x.view(-1, 192)
+        x = self.fc(x)
+        return F.log_softmax(x, dim=1)
+
+
+def train_func(model, optimizer, train_loader, device=None):
+    device = device or torch.device("cpu")
+    model.train()
+    for batch_idx, (data, target) in enumerate(train_loader):
+        if batch_idx * len(data) > EPOCH_SIZE:
+            return
+        data, target = data.to(device), target.to(device)
+        optimizer.zero_grad()
+        output = model(data)
+        loss = F.nll_loss(output, target)
+        loss.backward()
+        optimizer.step()
+
+
+def test_func(model, data_loader, device=None):
+    device = device or torch.device("cpu")
+    model.eval()
+    correct = 0
+    total = 0
+    with torch.no_grad():
+        for batch_idx, (data, target) in enumerate(data_loader):
+            if batch_idx * len(data) > TEST_SIZE:
+                break
+            data, target = data.to(device), target.to(device)
+            outputs = model(data)
+            _, predicted = torch.max(outputs.data, 1)
+            total += target.size(0)
+            correct += (predicted == target).sum().item()
+
+    return correct / total
+
+
+def get_data_loaders(batch_size=64):
+    mnist_transforms = transforms.Compose(
+        [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
+    )
+
+    # We add FileLock here because multiple workers will want to
+    # download data, and this may cause overwrites since
+    # DataLoader is not threadsafe.
+    with FileLock(os.path.expanduser("~/data.lock")):
+        train_loader = torch.utils.data.DataLoader(
+            datasets.MNIST(
+                "~/data", train=True, download=True, transform=mnist_transforms
+            ),
+            batch_size=batch_size,
+            shuffle=True,
+        )
+        test_loader = torch.utils.data.DataLoader(
+            datasets.MNIST(
+                "~/data", train=False, download=True, transform=mnist_transforms
+            ),
+            batch_size=batch_size,
+            shuffle=True,
+        )
+    return train_loader, test_loader
+
+
+def train_mnist(config):
+    should_checkpoint = config.get("should_checkpoint", False)
+    use_cuda = torch.cuda.is_available()
+    device = torch.device("cuda" if use_cuda else "cpu")
+    train_loader, test_loader = get_data_loaders()
+    model = ConvNet().to(device)
+
+    optimizer = optim.SGD(
+        model.parameters(), lr=config["lr"], momentum=config["momentum"]
+    )
+
+    while True:
+        train_func(model, optimizer, train_loader, device)
+        acc = test_func(model, test_loader, device)
+        metrics = {"mean_accuracy": acc}
+
+        # Report metrics (and possibly a checkpoint)
+        if should_checkpoint:
+            with tempfile.TemporaryDirectory() as tempdir:
+                torch.save(model.state_dict(), os.path.join(tempdir, "model.pt"))
+                tune.report(metrics, checkpoint=Checkpoint.from_directory(tempdir))
+        else:
+            tune.report(metrics)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="PyTorch MNIST Example")
+    parser.add_argument(
+        "--cuda", action="store_true", default=False, help="Enables GPU training"
+    )
+    parser.add_argument(
+        "--smoke-test", action="store_true", help="Finish quickly for testing"
+    )
+    args, _ = parser.parse_known_args()
+
+    ray.init(num_cpus=2 if args.smoke_test else None)
+
+    # for early stopping
+    sched = AsyncHyperBandScheduler()
+
+    resources_per_trial = {"cpu": 2, "gpu": int(args.cuda)}  # set this for GPUs
+    tuner = tune.Tuner(
+        tune.with_resources(train_mnist, resources=resources_per_trial),
+        tune_config=tune.TuneConfig(
+            metric="mean_accuracy",
+            mode="max",
+            scheduler=sched,
+            num_samples=1 if args.smoke_test else 50,
+        ),
+        run_config=tune.RunConfig(
+            name="exp",
+            stop={
+                "mean_accuracy": 0.98,
+                "training_iteration": 5 if args.smoke_test else 100,
+            },
+        ),
+        param_space={
+            "lr": tune.loguniform(1e-4, 1e-2),
+            "momentum": tune.uniform(0.1, 0.9),
+        },
+    )
+    results = tuner.fit()
+
+    print("Best config is:", results.get_best_result().config)
+
+    assert not results.errors

python/ray/tune/examples/mnist_pytorch_trainable.py

@@ -0,0 +1,98 @@
+# Original Code here:
+# https://github.com/pytorch/examples/blob/master/mnist/main.py
+from __future__ import print_function
+
+import argparse
+import os
+
+import torch
+import torch.optim as optim
+
+import ray
+from ray import tune
+from ray.tune.examples.mnist_pytorch import (
+    ConvNet,
+    get_data_loaders,
+    test_func,
+    train_func,
+)
+from ray.tune.schedulers import ASHAScheduler
+
+# Change these values if you want the training to run quicker or slower.
+EPOCH_SIZE = 512
+TEST_SIZE = 256
+
+# Training settings
+parser = argparse.ArgumentParser(description="PyTorch MNIST Example")
+parser.add_argument(
+    "--use-gpu", action="store_true", default=False, help="enables CUDA training"
+)
+parser.add_argument("--ray-address", type=str, help="The Redis address of the cluster.")
+parser.add_argument(
+    "--smoke-test", action="store_true", help="Finish quickly for testing"
+)
+
+
+# Below comments are for documentation purposes only.
+# fmt: off
+# __trainable_example_begin__
+class TrainMNIST(tune.Trainable):
+    def setup(self, config):
+        use_cuda = config.get("use_gpu") and torch.cuda.is_available()
+        self.device = torch.device("cuda" if use_cuda else "cpu")
+        self.train_loader, self.test_loader = get_data_loaders()
+        self.model = ConvNet().to(self.device)
+        self.optimizer = optim.SGD(
+            self.model.parameters(),
+            lr=config.get("lr", 0.01),
+            momentum=config.get("momentum", 0.9))
+
+    def step(self):
+        train_func(
+            self.model, self.optimizer, self.train_loader, device=self.device)
+        acc = test_func(self.model, self.test_loader, self.device)
+        return {"mean_accuracy": acc}
+
+    def save_checkpoint(self, checkpoint_dir):
+        checkpoint_path = os.path.join(checkpoint_dir, "model.pth")
+        torch.save(self.model.state_dict(), checkpoint_path)
+
+    def load_checkpoint(self, checkpoint_dir):
+        checkpoint_path = os.path.join(checkpoint_dir, "model.pth")
+        self.model.load_state_dict(torch.load(checkpoint_path))
+
+
+# __trainable_example_end__
+# fmt: on
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+    ray.init(address=args.ray_address, num_cpus=6 if args.smoke_test else None)
+    sched = ASHAScheduler()
+
+    tuner = tune.Tuner(
+        tune.with_resources(TrainMNIST, resources={"cpu": 3, "gpu": int(args.use_gpu)}),
+        run_config=tune.RunConfig(
+            stop={
+                "mean_accuracy": 0.95,
+                "training_iteration": 3 if args.smoke_test else 20,
+            },
+            checkpoint_config=tune.CheckpointConfig(
+                checkpoint_at_end=True, checkpoint_frequency=3
+            ),
+        ),
+        tune_config=tune.TuneConfig(
+            metric="mean_accuracy",
+            mode="max",
+            scheduler=sched,
+            num_samples=1 if args.smoke_test else 20,
+        ),
+        param_space={
+            "args": args,
+            "lr": tune.uniform(0.001, 0.1),
+            "momentum": tune.uniform(0.1, 0.9),
+        },
+    )
+    results = tuner.fit()
+
+    print("Best config is:", results.get_best_result().config)

python/ray/tune/examples/nevergrad_example.py

@@ -0,0 +1,77 @@
+"""This example demonstrates the usage of Nevergrad with Ray Tune.
+
+It also checks that it is usable with a separate scheduler.
+
+Requires the Nevergrad library to be installed (`pip install nevergrad`).
+"""
+
+import time
+
+from ray import tune
+from ray.tune.schedulers import AsyncHyperBandScheduler
+from ray.tune.search import ConcurrencyLimiter
+from ray.tune.search.nevergrad import NevergradSearch
+
+
+def evaluation_fn(step, width, height):
+    return (0.1 + width * step / 100) ** (-1) + height * 0.1
+
+
+def easy_objective(config):
+    # Hyperparameters
+    width, height = config["width"], config["height"]
+
+    for step in range(config["steps"]):
+        # Iterative training function - can be any arbitrary training procedure
+        intermediate_score = evaluation_fn(step, width, height)
+        # Feed the score back back to Tune.
+        tune.report({"iterations": step, "mean_loss": intermediate_score})
+        time.sleep(0.1)
+
+
+if __name__ == "__main__":
+    import argparse
+
+    import nevergrad as ng
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--smoke-test", action="store_true", help="Finish quickly for testing"
+    )
+    args, _ = parser.parse_known_args()
+
+    # Optional: Pass the parameter space yourself
+    # 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"])
+    # )
+
+    algo = NevergradSearch(
+        optimizer=ng.optimizers.OnePlusOne,
+        # space=space,  # If you want to set the space manually
+    )
+    algo = ConcurrencyLimiter(algo, max_concurrent=4)
+
+    scheduler = AsyncHyperBandScheduler()
+
+    tuner = tune.Tuner(
+        easy_objective,
+        tune_config=tune.TuneConfig(
+            metric="mean_loss",
+            mode="min",
+            search_alg=algo,
+            scheduler=scheduler,
+            num_samples=10 if args.smoke_test else 50,
+        ),
+        run_config=tune.RunConfig(name="nevergrad"),
+        param_space={
+            "steps": 100,
+            "width": tune.uniform(0, 20),
+            "height": tune.uniform(-100, 100),
+            "activation": tune.choice(["relu", "tanh"]),
+        },
+    )
+    results = tuner.fit()
+
+    print("Best hyperparameters found were: ", results.get_best_result().config)

python/ray/tune/examples/optuna_define_by_run_example.py

@@ -0,0 +1,93 @@
+"""This example demonstrates the usage of Optuna define-by-run with Ray Tune.
+
+It also checks that it is usable with a separate scheduler.
+
+Requires the Optuna library to be installed (`pip install optuna`).
+
+For an example of using a Tune search space, see
+:doc:`/tune/examples/optuna_example`.
+"""
+
+import time
+from typing import Any, Dict, Optional
+
+import ray
+from ray import tune
+from ray.tune.schedulers import AsyncHyperBandScheduler
+from ray.tune.search import ConcurrencyLimiter
+from ray.tune.search.optuna import OptunaSearch
+
+
+def evaluation_fn(step, width, height, mult=1):
+    return (0.1 + width * step / 100) ** (-1) + height * 0.1 * mult
+
+
+def easy_objective(config):
+    # Hyperparameters
+    width, height, mult = config["width"], config["height"], config.get("mult", 1)
+    print(config)
+
+    for step in range(config["steps"]):
+        # Iterative training function - can be any arbitrary training procedure
+        intermediate_score = evaluation_fn(step, width, height, mult)
+        # Feed the score back back to Tune.
+        tune.report({"iterations": step, "mean_loss": intermediate_score})
+        time.sleep(0.1)
+
+
+def define_by_run_func(trial) -> Optional[Dict[str, Any]]:
+    """Define-by-run function to create the search space.
+
+    Ensure no actual computation takes place here. That should go into
+    the trainable passed to ``Tuner`` (in this example, that's
+    ``easy_objective``).
+
+    For more information, see https://optuna.readthedocs.io/en/stable\
+/tutorial/10_key_features/002_configurations.html
+
+    This function should either return None or a dict with constant values.
+    """
+    # This param is not used in the objective function.
+    activation = trial.suggest_categorical("activation", ["relu", "tanh"])
+    trial.suggest_float("width", 0, 20)
+    trial.suggest_float("height", -100, 100)
+
+    # Define-by-run allows for conditional search spaces.
+    if activation == "relu":
+        trial.suggest_float("mult", 1, 2)
+
+    # Return all constants in a dictionary.
+    return {"steps": 100}
+
+
+def run_optuna_tune(smoke_test=False):
+    algo = OptunaSearch(space=define_by_run_func, metric="mean_loss", mode="min")
+    algo = ConcurrencyLimiter(algo, max_concurrent=4)
+    scheduler = AsyncHyperBandScheduler()
+    tuner = tune.Tuner(
+        easy_objective,
+        tune_config=tune.TuneConfig(
+            metric="mean_loss",
+            mode="min",
+            search_alg=algo,
+            scheduler=scheduler,
+            num_samples=10 if smoke_test else 100,
+        ),
+    )
+    results = tuner.fit()
+
+    print("Best hyperparameters found were: ", results.get_best_result().config)
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--smoke-test", action="store_true", help="Finish quickly for testing"
+    )
+    args, _ = parser.parse_known_args()
+
+    ray.init(configure_logging=False)
+
+    run_optuna_tune(smoke_test=args.smoke_test)

python/ray/tune/examples/optuna_example.py

@@ -0,0 +1,73 @@
+"""This example demonstrates the usage of Optuna with Ray Tune.
+
+It also checks that it is usable with a separate scheduler.
+
+Requires the Optuna library to be installed (`pip install optuna`).
+
+For an example of using an Optuna define-by-run function, see
+:doc:`/tune/examples/optuna_define_by_run_example`.
+"""
+
+import time
+
+import ray
+from ray import tune
+from ray.tune.schedulers import AsyncHyperBandScheduler
+from ray.tune.search import ConcurrencyLimiter
+from ray.tune.search.optuna import OptunaSearch
+
+
+def evaluation_fn(step, width, height):
+    return (0.1 + width * step / 100) ** (-1) + height * 0.1
+
+
+def easy_objective(config):
+    # Hyperparameters
+    width, height = config["width"], config["height"]
+
+    for step in range(config["steps"]):
+        # Iterative training function - can be any arbitrary training procedure
+        intermediate_score = evaluation_fn(step, width, height)
+        # Feed the score back back to Tune.
+        tune.report({"iterations": step, "mean_loss": intermediate_score})
+        time.sleep(0.1)
+
+
+def run_optuna_tune(smoke_test=False):
+    algo = OptunaSearch()
+    algo = ConcurrencyLimiter(algo, max_concurrent=4)
+    scheduler = AsyncHyperBandScheduler()
+    tuner = tune.Tuner(
+        easy_objective,
+        tune_config=tune.TuneConfig(
+            metric="mean_loss",
+            mode="min",
+            search_alg=algo,
+            scheduler=scheduler,
+            num_samples=10 if smoke_test else 100,
+        ),
+        param_space={
+            "steps": 100,
+            "width": tune.uniform(0, 20),
+            "height": tune.uniform(-100, 100),
+            # This is an ignored parameter.
+            "activation": tune.choice(["relu", "tanh"]),
+        },
+    )
+    results = tuner.fit()
+
+    print("Best hyperparameters found were: ", results.get_best_result().config)
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--smoke-test", action="store_true", help="Finish quickly for testing"
+    )
+    args, _ = parser.parse_known_args()
+
+    ray.init(configure_logging=False)
+
+    run_optuna_tune(smoke_test=args.smoke_test)

python/ray/tune/examples/optuna_multiobjective_example.py

@@ -0,0 +1,79 @@
+"""This example demonstrates the usage of Optuna with Ray Tune for
+multi-objective optimization.
+
+Please note that schedulers may not work correctly with multi-objective
+optimization.
+
+Requires the Optuna library to be installed (`pip install optuna`).
+"""
+import time
+
+import ray
+from ray import tune
+from ray.tune.search import ConcurrencyLimiter
+from ray.tune.search.optuna import OptunaSearch
+
+
+def evaluation_fn(step, width, height):
+    return (0.1 + width * step / 100) ** (-1) + height * 0.1
+
+
+def easy_objective(config):
+    # Hyperparameters
+    width, height = config["width"], config["height"]
+
+    for step in range(config["steps"]):
+        # Iterative training function - can be any arbitrary training procedure
+        intermediate_score = evaluation_fn(step, width, height)
+        # Feed the score back back to Tune.
+        tune.report(
+            {
+                "iterations": step,
+                "loss": intermediate_score,
+                "gain": intermediate_score * width,
+            }
+        )
+        time.sleep(0.1)
+
+
+def run_optuna_tune(smoke_test=False):
+    algo = OptunaSearch(metric=["loss", "gain"], mode=["min", "max"])
+    algo = ConcurrencyLimiter(algo, max_concurrent=4)
+    tuner = tune.Tuner(
+        easy_objective,
+        tune_config=tune.TuneConfig(
+            search_alg=algo,
+            num_samples=10 if smoke_test else 100,
+        ),
+        param_space={
+            "steps": 100,
+            "width": tune.uniform(0, 20),
+            "height": tune.uniform(-100, 100),
+            # This is an ignored parameter.
+            "activation": tune.choice(["relu", "tanh"]),
+        },
+    )
+    results = tuner.fit()
+
+    print(
+        "Best hyperparameters for loss found were: ",
+        results.get_best_result("loss", "min").config,
+    )
+    print(
+        "Best hyperparameters for gain found were: ",
+        results.get_best_result("gain", "max").config,
+    )
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--smoke-test", action="store_true", help="Finish quickly for testing"
+    )
+    args, _ = parser.parse_known_args()
+
+    ray.init(configure_logging=False)
+
+    run_optuna_tune(smoke_test=args.smoke_test)

python/ray/tune/examples/pb2_example.py

@@ -0,0 +1,62 @@
+#!/usr/bin/env python
+
+import argparse
+
+import ray
+from ray import tune
+from ray.tune.examples.pbt_function import pbt_function
+from ray.tune.schedulers.pb2 import PB2
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--smoke-test", action="store_true", help="Finish quickly for testing"
+    )
+    args, _ = parser.parse_known_args()
+
+    if args.smoke_test:
+        ray.init(num_cpus=2)  # force pausing to happen for test
+
+    perturbation_interval = 5
+    pbt = PB2(
+        time_attr="training_iteration",
+        perturbation_interval=perturbation_interval,
+        hyperparam_bounds={
+            # hyperparameter bounds.
+            "lr": [0.0001, 0.02],
+        },
+    )
+
+    tuner = tune.Tuner(
+        pbt_function,
+        run_config=tune.RunConfig(
+            name="pbt_test",
+            verbose=False,
+            stop={
+                "training_iteration": 30,
+            },
+            failure_config=tune.FailureConfig(
+                fail_fast=True,
+            ),
+        ),
+        tune_config=tune.TuneConfig(
+            scheduler=pbt,
+            metric="mean_accuracy",
+            mode="max",
+            num_samples=8,
+            reuse_actors=True,
+        ),
+        param_space={
+            "lr": 0.0001,
+            # note: this parameter is perturbed but has no effect on
+            # the model training in this example
+            "some_other_factor": 1,
+            # This parameter is not perturbed and is used to determine
+            # checkpoint frequency. We set checkpoints and perturbations
+            # to happen at the same frequency.
+            "checkpoint_interval": perturbation_interval,
+        },
+    )
+    results = tuner.fit()
+
+    print("Best hyperparameters found were: ", results.get_best_result().config)

python/ray/tune/examples/pb2_ppo_example.py

@@ -0,0 +1,157 @@
+import argparse
+import os
+import random
+from datetime import datetime
+
+import pandas as pd
+
+from ray.tune import run, sample_from
+from ray.tune.schedulers import PopulationBasedTraining
+from ray.tune.schedulers.pb2 import PB2
+
+
+# Postprocess the perturbed config to ensure it's still valid used if PBT.
+def explore(config):
+    # Ensure we collect enough timesteps to do sgd.
+    if config["train_batch_size"] < config["sgd_minibatch_size"] * 2:
+        config["train_batch_size"] = config["sgd_minibatch_size"] * 2
+    # Ensure we run at least one sgd iter.
+    if config["lambda"] > 1:
+        config["lambda"] = 1
+    config["train_batch_size"] = int(config["train_batch_size"])
+    return config
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--max", type=int, default=1000000)
+    parser.add_argument("--algo", type=str, default="PPO")
+    parser.add_argument("--num_workers", type=int, default=4)
+    parser.add_argument("--num_samples", type=int, default=4)
+    parser.add_argument("--t_ready", type=int, default=50000)
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument(
+        "--horizon", type=int, default=1600
+    )  # make this 1000 for other envs
+    parser.add_argument("--perturb", type=float, default=0.25)  # if using PBT
+    parser.add_argument("--env_name", type=str, default="BipedalWalker-v2")
+    parser.add_argument(
+        "--criteria", type=str, default="timesteps_total"
+    )  # "training_iteration", "time_total_s"
+    parser.add_argument(
+        "--net", type=str, default="32_32"
+    )  # May be important to use a larger network for bigger tasks.
+    parser.add_argument("--filename", type=str, default="")
+    parser.add_argument("--method", type=str, default="pb2")  # ['pbt', 'pb2']
+    parser.add_argument("--save_csv", type=bool, default=False)
+
+    args = parser.parse_args()
+
+    # bipedalwalker needs 1600
+    if args.env_name in ["BipedalWalker-v2", "BipedalWalker-v3"]:
+        horizon = 1600
+    else:
+        horizon = 1000
+
+    pbt = PopulationBasedTraining(
+        time_attr=args.criteria,
+        metric="episode_reward_mean",
+        mode="max",
+        perturbation_interval=args.t_ready,
+        resample_probability=args.perturb,
+        quantile_fraction=args.perturb,  # copy bottom % with top %
+        # Specifies the search space for these hyperparams
+        hyperparam_mutations={
+            "lambda": lambda: random.uniform(0.9, 1.0),
+            "clip_param": lambda: random.uniform(0.1, 0.5),
+            "lr": lambda: random.uniform(1e-3, 1e-5),
+            "train_batch_size": lambda: random.randint(1000, 60000),
+        },
+        custom_explore_fn=explore,
+    )
+
+    pb2 = PB2(
+        time_attr=args.criteria,
+        metric="episode_reward_mean",
+        mode="max",
+        perturbation_interval=args.t_ready,
+        quantile_fraction=args.perturb,  # copy bottom % with top %
+        # Specifies the hyperparam search space
+        hyperparam_bounds={
+            "lambda": [0.9, 1.0],
+            "clip_param": [0.1, 0.5],
+            "lr": [1e-5, 1e-3],
+            "train_batch_size": [1000, 60000],
+        },
+    )
+
+    methods = {"pbt": pbt, "pb2": pb2}
+
+    timelog = (
+        str(datetime.date(datetime.now())) + "_" + str(datetime.time(datetime.now()))
+    )
+
+    args.dir = "{}_{}_{}_Size{}_{}_{}".format(
+        args.algo,
+        args.filename,
+        args.method,
+        str(args.num_samples),
+        args.env_name,
+        args.criteria,
+    )
+
+    analysis = run(
+        args.algo,
+        name="{}_{}_{}_seed{}_{}".format(
+            timelog, args.method, args.env_name, str(args.seed), args.filename
+        ),
+        scheduler=methods[args.method],
+        verbose=1,
+        num_samples=args.num_samples,
+        reuse_actors=True,
+        stop={args.criteria: args.max},
+        config={
+            "env": args.env_name,
+            "log_level": "INFO",
+            "seed": args.seed,
+            "kl_coeff": 1.0,
+            "num_gpus": 0,
+            "horizon": horizon,
+            "observation_filter": "MeanStdFilter",
+            "model": {
+                "fcnet_hiddens": [
+                    int(args.net.split("_")[0]),
+                    int(args.net.split("_")[1]),
+                ],
+                "free_log_std": True,
+            },
+            "num_sgd_iter": 10,
+            "sgd_minibatch_size": 128,
+            "lambda": sample_from(lambda spec: random.uniform(0.9, 1.0)),
+            "clip_param": sample_from(lambda spec: random.uniform(0.1, 0.5)),
+            "lr": sample_from(lambda spec: random.uniform(1e-3, 1e-5)),
+            "train_batch_size": sample_from(lambda spec: random.randint(1000, 60000)),
+        },
+    )
+
+    all_dfs = list(analysis.trial_dataframes.values())
+
+    results = pd.DataFrame()
+    for i in range(args.num_samples):
+        df = all_dfs[i]
+        df = df[
+            [
+                "timesteps_total",
+                "episodes_total",
+                "episode_reward_mean",
+                "info/learner/default_policy/cur_kl_coeff",
+            ]
+        ]
+        df["Agent"] = i
+        results = pd.concat([results, df]).reset_index(drop=True)
+
+    if args.save_csv:
+        if not (os.path.exists("data/" + args.dir)):
+            os.makedirs("data/" + args.dir)
+
+        results.to_csv("data/{}/seed{}.csv".format(args.dir, str(args.seed)))

python/ray/tune/examples/pbt_convnet_example.py

@@ -0,0 +1,138 @@
+#!/usr/bin/env python
+
+# ruff: noqa
+# fmt: off
+
+# __tutorial_imports_begin__
+import argparse
+import os
+
+import numpy as np
+import torch
+import torch.optim as optim
+from torchvision import datasets
+
+import ray
+from ray import tune
+from ray.tune.examples.mnist_pytorch import (
+    ConvNet,
+    get_data_loaders,
+    test_func,
+    train_func,
+)
+from ray.tune.schedulers import PopulationBasedTraining
+from ray.tune.utils import validate_save_restore
+
+# __tutorial_imports_end__
+
+
+# __trainable_begin__
+class PytorchTrainable(tune.Trainable):
+    """Train a Pytorch ConvNet with Trainable and PopulationBasedTraining
+       scheduler. The example reuse some of the functions in mnist_pytorch,
+       and is a good demo for how to add the tuning function without
+       changing the original training code.
+    """
+
+    def setup(self, config):
+        self.train_loader, self.test_loader = get_data_loaders()
+        self.model = ConvNet()
+        self.optimizer = optim.SGD(
+            self.model.parameters(),
+            lr=config.get("lr", 0.01),
+            momentum=config.get("momentum", 0.9))
+
+    def step(self):
+        train_func(self.model, self.optimizer, self.train_loader)
+        acc = test_func(self.model, self.test_loader)
+        return {"mean_accuracy": acc}
+
+    def save_checkpoint(self, checkpoint_dir):
+        checkpoint_path = os.path.join(checkpoint_dir, "model.pth")
+        torch.save(self.model.state_dict(), checkpoint_path)
+
+    def load_checkpoint(self, checkpoint_dir):
+        checkpoint_path = os.path.join(checkpoint_dir, "model.pth")
+        self.model.load_state_dict(torch.load(checkpoint_path))
+
+    def reset_config(self, new_config):
+        for param_group in self.optimizer.param_groups:
+            if "lr" in new_config:
+                param_group["lr"] = new_config["lr"]
+            if "momentum" in new_config:
+                param_group["momentum"] = new_config["momentum"]
+
+        self.config = new_config
+        return True
+# __trainable_end__
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--smoke-test", action="store_true", help="Finish quickly for testing")
+    args, _ = parser.parse_known_args()
+
+    ray.init(num_cpus=2)
+    datasets.MNIST("~/data", train=True, download=True)
+
+    # check if PytorchTrainble will save/restore correctly before execution
+    validate_save_restore(PytorchTrainable)
+
+    # __pbt_begin__
+    scheduler = PopulationBasedTraining(
+        time_attr="training_iteration",
+        perturbation_interval=5,
+        hyperparam_mutations={
+            # distribution for resampling
+            "lr": lambda: np.random.uniform(0.0001, 1),
+            # allow perturbations within this set of categorical values
+            "momentum": [0.8, 0.9, 0.99],
+        })
+    # __pbt_end__
+
+    # __tune_begin__
+    class CustomStopper(tune.Stopper):
+        def __init__(self):
+            self.should_stop = False
+
+        def __call__(self, trial_id, result):
+            max_iter = 5 if args.smoke_test else 100
+            if not self.should_stop and result["mean_accuracy"] > 0.96:
+                self.should_stop = True
+            return self.should_stop or result["training_iteration"] >= max_iter
+
+        def stop_all(self):
+            return self.should_stop
+
+    stopper = CustomStopper()
+
+    tuner = tune.Tuner(
+        PytorchTrainable,
+        run_config=tune.RunConfig(
+            name="pbt_test",
+            stop=stopper,
+            verbose=1,
+            checkpoint_config=tune.CheckpointConfig(
+                checkpoint_score_attribute="mean_accuracy",
+                checkpoint_frequency=5,
+                num_to_keep=4,
+            ),
+        ),
+        tune_config=tune.TuneConfig(
+            scheduler=scheduler,
+            metric="mean_accuracy",
+            mode="max",
+            num_samples=4,
+            reuse_actors=True,
+        ),
+        param_space={
+            "lr": tune.uniform(0.001, 1),
+            "momentum": tune.uniform(0.001, 1),
+        },
+    )
+    results = tuner.fit()
+    # __tune_end__
+
+    best_result = results.get_best_result()
+    best_checkpoint = best_result.checkpoint

python/ray/tune/examples/pbt_convnet_function_example.py

@@ -0,0 +1,146 @@
+#!/usr/bin/env python
+
+# __tutorial_imports_begin__
+import argparse
+import os
+
+import numpy as np
+import torch
+import torch.optim as optim
+
+import ray
+from ray import tune
+from ray.tune import Checkpoint
+from ray.tune.examples.mnist_pytorch import ConvNet, get_data_loaders, test_func
+from ray.tune.schedulers import PopulationBasedTraining
+
+# __tutorial_imports_end__
+
+
+# __train_begin__
+def train_convnet(config):
+    # Create our data loaders, model, and optmizer.
+    step = 0
+    train_loader, test_loader = get_data_loaders()
+    model = ConvNet()
+    optimizer = optim.SGD(
+        model.parameters(),
+        lr=config.get("lr", 0.01),
+        momentum=config.get("momentum", 0.9),
+    )
+
+    # If `get_checkpoint()` is not None, then we are resuming from a checkpoint.
+    # Load model state and iteration step from checkpoint.
+    if tune.get_checkpoint():
+        print("Loading from checkpoint.")
+        loaded_checkpoint = tune.get_checkpoint()
+        with loaded_checkpoint.as_directory() as loaded_checkpoint_dir:
+            path = os.path.join(loaded_checkpoint_dir, "checkpoint.pt")
+            checkpoint = torch.load(path)
+            model.load_state_dict(checkpoint["model"])
+            step = checkpoint["step"]
+
+    while True:
+        ray.tune.examples.mnist_pytorch.train_func(model, optimizer, train_loader)
+        acc = test_func(model, test_loader)
+        checkpoint = None
+        if step % 5 == 0:
+            # Every 5 steps, checkpoint our current state.
+            # First get the checkpoint directory from tune.
+            # Need to create a directory under current working directory
+            # to construct checkpoint object from.
+            os.makedirs("my_model", exist_ok=True)
+            torch.save(
+                {
+                    "step": step,
+                    "model": model.state_dict(),
+                },
+                "my_model/checkpoint.pt",
+            )
+            checkpoint = Checkpoint.from_directory("my_model")
+
+        step += 1
+        tune.report({"mean_accuracy": acc}, checkpoint=checkpoint)
+
+
+# __train_end__
+
+
+def eval_best_model(results: tune.ResultGrid):
+    """Test the best model given output of tuner.fit()."""
+    with results.get_best_result().checkpoint.as_directory() as best_checkpoint_path:
+        best_model = ConvNet()
+        best_checkpoint = torch.load(
+            os.path.join(best_checkpoint_path, "checkpoint.pt")
+        )
+        best_model.load_state_dict(best_checkpoint["model"])
+        # Note that test only runs on a small random set of the test data, thus the
+        # accuracy may be different from metrics shown in tuning process.
+        test_acc = test_func(best_model, get_data_loaders()[1])
+        print("best model accuracy: ", test_acc)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--smoke-test", action="store_true", help="Finish quickly for testing"
+    )
+    args, _ = parser.parse_known_args()
+
+    # __pbt_begin__
+    scheduler = PopulationBasedTraining(
+        time_attr="training_iteration",
+        perturbation_interval=5,
+        hyperparam_mutations={
+            # distribution for resampling
+            "lr": lambda: np.random.uniform(0.0001, 1),
+            # allow perturbations within this set of categorical values
+            "momentum": [0.8, 0.9, 0.99],
+        },
+    )
+
+    # __pbt_end__
+
+    # __tune_begin__
+    class CustomStopper(tune.Stopper):
+        def __init__(self):
+            self.should_stop = False
+
+        def __call__(self, trial_id, result):
+            max_iter = 5 if args.smoke_test else 100
+            if not self.should_stop and result["mean_accuracy"] > 0.96:
+                self.should_stop = True
+            return self.should_stop or result["training_iteration"] >= max_iter
+
+        def stop_all(self):
+            return self.should_stop
+
+    stopper = CustomStopper()
+
+    tuner = tune.Tuner(
+        train_convnet,
+        run_config=tune.RunConfig(
+            name="pbt_test",
+            stop=stopper,
+            verbose=1,
+            checkpoint_config=tune.CheckpointConfig(
+                checkpoint_score_attribute="mean_accuracy",
+                num_to_keep=4,
+            ),
+        ),
+        tune_config=tune.TuneConfig(
+            scheduler=scheduler,
+            metric="mean_accuracy",
+            mode="max",
+            num_samples=4,
+            reuse_actors=True,
+        ),
+        param_space={
+            "lr": tune.uniform(0.001, 1),
+            "momentum": tune.uniform(0.001, 1),
+        },
+    )
+    results = tuner.fit()
+    # __tune_end__
+
+    eval_best_model(results)

python/ray/tune/examples/pbt_dcgan_mnist/common.py

@@ -0,0 +1,285 @@
+import os
+
+import matplotlib.animation as animation
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.parallel
+import torch.utils.data
+import torchvision.datasets as dset
+import torchvision.transforms as transforms
+import torchvision.utils as vutils
+from scipy.stats import entropy
+from torch.autograd import Variable
+from torch.nn import functional as F
+
+import ray
+
+# Training parameters
+workers = 2
+batch_size = 64
+image_size = 32
+
+# Number of channels in the training images. For color images this is 3
+nc = 1
+
+# Size of z latent vector (i.e. size of generator input)
+nz = 100
+
+# Size of feature maps in generator
+ngf = 32
+
+# Size of feature maps in discriminator
+ndf = 32
+
+# Beta1 hyperparam for Adam optimizers
+beta1 = 0.5
+
+# iterations of actual training in each Trainable _train
+train_iterations_per_step = 5
+
+MODEL_PATH = os.path.expanduser("~/.ray/models/mnist_cnn.pt")
+
+
+def get_data_loader(data_dir="~/data"):
+    dataset = dset.MNIST(
+        root=data_dir,
+        download=True,
+        transform=transforms.Compose(
+            [
+                transforms.Resize(image_size),
+                transforms.ToTensor(),
+                transforms.Normalize((0.13066,), (0.30131,)),
+            ]
+        ),
+    )
+
+    # Create the dataloader
+    dataloader = torch.utils.data.DataLoader(
+        dataset, batch_size=batch_size, shuffle=True, num_workers=workers
+    )
+
+    return dataloader
+
+
+# __GANmodel_begin__
+# custom weights initialization called on netG and netD
+def weights_init(m):
+    classname = m.__class__.__name__
+    if classname.find("Conv") != -1:
+        nn.init.normal_(m.weight.data, 0.0, 0.02)
+    elif classname.find("BatchNorm") != -1:
+        nn.init.normal_(m.weight.data, 1.0, 0.02)
+        nn.init.constant_(m.bias.data, 0)
+
+
+# Generator Code
+class Generator(nn.Module):
+    def __init__(self):
+        super(Generator, self).__init__()
+        self.main = nn.Sequential(
+            # input is Z, going into a convolution
+            nn.ConvTranspose2d(nz, ngf * 4, 4, 1, 0, bias=False),
+            nn.BatchNorm2d(ngf * 4),
+            nn.ReLU(True),
+            nn.ConvTranspose2d(ngf * 4, ngf * 2, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(ngf * 2),
+            nn.ReLU(True),
+            nn.ConvTranspose2d(ngf * 2, ngf, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(ngf),
+            nn.ReLU(True),
+            nn.ConvTranspose2d(ngf, nc, 4, 2, 1, bias=False),
+            nn.Tanh(),
+        )
+
+    def forward(self, input):
+        return self.main(input)
+
+
+class Discriminator(nn.Module):
+    def __init__(self):
+        super(Discriminator, self).__init__()
+        self.main = nn.Sequential(
+            nn.Conv2d(nc, ndf, 4, 2, 1, bias=False),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(ndf * 2),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(ndf * 4),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Conv2d(ndf * 4, 1, 4, 1, 0, bias=False),
+            nn.Sigmoid(),
+        )
+
+    def forward(self, input):
+        return self.main(input)
+
+
+# __GANmodel_end__
+
+
+# __INCEPTION_SCORE_begin__
+class Net(nn.Module):
+    """
+    LeNet for MNist classification, used for inception_score
+    """
+
+    def __init__(self):
+        super(Net, self).__init__()
+        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
+        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
+        self.conv2_drop = nn.Dropout2d()
+        self.fc1 = nn.Linear(320, 50)
+        self.fc2 = nn.Linear(50, 10)
+
+    def forward(self, x):
+        x = F.relu(F.max_pool2d(self.conv1(x), 2))
+        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
+        x = x.view(-1, 320)
+        x = F.relu(self.fc1(x))
+        x = F.dropout(x, training=self.training)
+        x = self.fc2(x)
+        return F.log_softmax(x, dim=1)
+
+
+def inception_score(imgs, mnist_model_ref, batch_size=32, splits=1):
+    N = len(imgs)
+    dtype = torch.FloatTensor
+    dataloader = torch.utils.data.DataLoader(imgs, batch_size=batch_size)
+    cm = ray.get(mnist_model_ref)  # Get the mnist model from Ray object store.
+    up = nn.Upsample(size=(28, 28), mode="bilinear").type(dtype)
+
+    def get_pred(x):
+        x = up(x)
+        x = cm(x)
+        return F.softmax(x).data.cpu().numpy()
+
+    preds = np.zeros((N, 10))
+    for i, batch in enumerate(dataloader, 0):
+        batch = batch.type(dtype)
+        batchv = Variable(batch)
+        batch_size_i = batch.size()[0]
+        preds[i * batch_size : i * batch_size + batch_size_i] = get_pred(batchv)
+
+    # Now compute the mean kl-div
+    split_scores = []
+    for k in range(splits):
+        part = preds[k * (N // splits) : (k + 1) * (N // splits), :]
+        py = np.mean(part, axis=0)
+        scores = []
+        for i in range(part.shape[0]):
+            pyx = part[i, :]
+            scores.append(entropy(pyx, py))
+        split_scores.append(np.exp(np.mean(scores)))
+
+    return np.mean(split_scores), np.std(split_scores)
+
+
+# __INCEPTION_SCORE_end__
+
+
+def train_func(
+    netD,
+    netG,
+    optimG,
+    optimD,
+    criterion,
+    dataloader,
+    iteration,
+    device,
+    mnist_model_ref,
+):
+    real_label = 1
+    fake_label = 0
+
+    for i, data in enumerate(dataloader, 0):
+        if i >= train_iterations_per_step:
+            break
+
+        netD.zero_grad()
+        real_cpu = data[0].to(device)
+        b_size = real_cpu.size(0)
+        label = torch.full((b_size,), real_label, dtype=torch.float, device=device)
+        output = netD(real_cpu).view(-1)
+        errD_real = criterion(output, label)
+        errD_real.backward()
+        D_x = output.mean().item()
+
+        noise = torch.randn(b_size, nz, 1, 1, device=device)
+        fake = netG(noise)
+        label.fill_(fake_label)
+        output = netD(fake.detach()).view(-1)
+        errD_fake = criterion(output, label)
+        errD_fake.backward()
+        D_G_z1 = output.mean().item()
+        errD = errD_real + errD_fake
+        optimD.step()
+
+        netG.zero_grad()
+        label.fill_(real_label)
+        output = netD(fake).view(-1)
+        errG = criterion(output, label)
+        errG.backward()
+        D_G_z2 = output.mean().item()
+        optimG.step()
+
+        is_score, is_std = inception_score(fake, mnist_model_ref)
+
+        # Output training stats
+        if iteration % 10 == 0:
+            print(
+                "[%d/%d]\tLoss_D: %.4f\tLoss_G: %.4f\tD(x): %.4f\tD(G(z))"
+                ": %.4f / %.4f \tInception score: %.4f"
+                % (
+                    iteration,
+                    len(dataloader),
+                    errD.item(),
+                    errG.item(),
+                    D_x,
+                    D_G_z1,
+                    D_G_z2,
+                    is_score,
+                )
+            )
+
+    return errG.item(), errD.item(), is_score
+
+
+def plot_images(dataloader):
+    # Plot some training images
+    real_batch = next(iter(dataloader))
+    plt.figure(figsize=(8, 8))
+    plt.axis("off")
+    plt.title("Original Images")
+    plt.imshow(
+        np.transpose(
+            vutils.make_grid(real_batch[0][:64], padding=2, normalize=True).cpu(),
+            (1, 2, 0),
+        )
+    )
+
+    plt.show()
+
+
+def demo_gan(checkpoint_paths):
+    img_list = []
+    fixed_noise = torch.randn(64, nz, 1, 1)
+    for path in checkpoint_paths:
+        checkpoint_dict = torch.load(os.path.join(path, "checkpoint.pt"))
+
+        loadedG = Generator()
+        loadedG.load_state_dict(checkpoint_dict["netGmodel"])
+        with torch.no_grad():
+            fake = loadedG(fixed_noise).detach().cpu()
+        img_list.append(vutils.make_grid(fake, padding=2, normalize=True))
+
+    fig = plt.figure(figsize=(8, 8))
+    plt.axis("off")
+    ims = [[plt.imshow(np.transpose(i, (1, 2, 0)), animated=True)] for i in img_list]
+    ani = animation.ArtistAnimation(
+        fig, ims, interval=1000, repeat_delay=1000, blit=True
+    )
+    ani.save("./generated.gif", writer="imagemagick", dpi=72)
+    plt.show()

python/ray/tune/examples/pbt_dcgan_mnist/pbt_dcgan_mnist_func.py

@@ -0,0 +1,191 @@
+#!/usr/bin/env python
+"""
+Example of training DCGAN on MNIST using PBT with Tune's function API.
+"""
+import argparse
+import os
+import tempfile
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.parallel
+import torch.optim as optim
+import torch.utils.data
+from filelock import FileLock
+
+import ray
+from ray import tune
+from ray.tune import Checkpoint
+from ray.tune.examples.pbt_dcgan_mnist.common import (
+    MODEL_PATH,
+    Discriminator,
+    Generator,
+    Net,
+    beta1,
+    demo_gan,
+    get_data_loader,
+    plot_images,
+    train_func,
+    weights_init,
+)
+from ray.tune.schedulers import PopulationBasedTraining
+
+
+# __Train_begin__
+def dcgan_train(config):
+    use_cuda = config.get("use_gpu") and torch.cuda.is_available()
+    device = torch.device("cuda" if use_cuda else "cpu")
+    netD = Discriminator().to(device)
+    netD.apply(weights_init)
+    netG = Generator().to(device)
+    netG.apply(weights_init)
+    criterion = nn.BCELoss()
+    optimizerD = optim.Adam(
+        netD.parameters(), lr=config.get("lr", 0.01), betas=(beta1, 0.999)
+    )
+    optimizerG = optim.Adam(
+        netG.parameters(), lr=config.get("lr", 0.01), betas=(beta1, 0.999)
+    )
+    with FileLock(os.path.expanduser("~/ray_results/.data.lock")):
+        dataloader = get_data_loader()
+
+    step = 1
+    checkpoint = tune.get_checkpoint()
+    if checkpoint:
+        with checkpoint.as_directory() as checkpoint_dir:
+            checkpoint_dict = torch.load(os.path.join(checkpoint_dir, "checkpoint.pt"))
+        netD.load_state_dict(checkpoint_dict["netDmodel"])
+        netG.load_state_dict(checkpoint_dict["netGmodel"])
+        optimizerD.load_state_dict(checkpoint_dict["optimD"])
+        optimizerG.load_state_dict(checkpoint_dict["optimG"])
+        # Note: Make sure to increment the loaded step by 1 to get the
+        # current step.
+        last_step = checkpoint_dict["step"]
+        step = last_step + 1
+
+        # NOTE: It's important to set the optimizer learning rates
+        # again, since we want to explore the parameters passed in by PBT.
+        # Without this, we would continue using the exact same
+        # configuration as the trial whose checkpoint we are exploiting.
+        if "netD_lr" in config:
+            for param_group in optimizerD.param_groups:
+                param_group["lr"] = config["netD_lr"]
+        if "netG_lr" in config:
+            for param_group in optimizerG.param_groups:
+                param_group["lr"] = config["netG_lr"]
+
+    while True:
+        lossG, lossD, is_score = train_func(
+            netD,
+            netG,
+            optimizerG,
+            optimizerD,
+            criterion,
+            dataloader,
+            step,
+            device,
+            config["mnist_model_ref"],
+        )
+        metrics = {"lossg": lossG, "lossd": lossD, "is_score": is_score}
+
+        if step % config["checkpoint_interval"] == 0:
+            with tempfile.TemporaryDirectory() as tmpdir:
+                torch.save(
+                    {
+                        "netDmodel": netD.state_dict(),
+                        "netGmodel": netG.state_dict(),
+                        "optimD": optimizerD.state_dict(),
+                        "optimG": optimizerG.state_dict(),
+                        "step": step,
+                    },
+                    os.path.join(tmpdir, "checkpoint.pt"),
+                )
+                tune.report(metrics, checkpoint=Checkpoint.from_directory(tmpdir))
+        else:
+            tune.report(metrics)
+
+        step += 1
+
+
+# __Train_end__
+
+
+def download_mnist_cnn():
+    import urllib.request
+
+    # Download a pre-trained MNIST model for inception score calculation.
+    # This is a tiny model (<100kb).
+    if not os.path.exists(MODEL_PATH):
+        print("downloading model")
+        os.makedirs(os.path.dirname(MODEL_PATH), exist_ok=True)
+        urllib.request.urlretrieve(
+            "https://github.com/ray-project/ray/raw/master/python/ray/tune/"
+            "examples/pbt_dcgan_mnist/mnist_cnn.pt",
+            MODEL_PATH,
+        )
+    return MODEL_PATH
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--smoke-test", action="store_true", help="Finish quickly for testing"
+    )
+    parser.add_argument(
+        "--data-dir", type=str, default="~/data/", help="Set the path of the dataset."
+    )
+    args, _ = parser.parse_known_args()
+    ray.init()
+
+    download_mnist_cnn()
+
+    dataloader = get_data_loader(args.data_dir)
+    if not args.smoke_test:
+        plot_images(dataloader)
+
+    # __tune_begin__
+
+    # load the pretrained mnist classification model for inception_score
+    mnist_cnn = Net()
+    mnist_cnn.load_state_dict(torch.load(MODEL_PATH))
+    mnist_cnn.eval()
+    # Put the model in Ray object store.
+    mnist_model_ref = ray.put(mnist_cnn)
+
+    scheduler = PopulationBasedTraining(
+        perturbation_interval=5,
+        hyperparam_mutations={
+            # distribution for resampling
+            "netG_lr": lambda: np.random.uniform(1e-2, 1e-5),
+            "netD_lr": lambda: np.random.uniform(1e-2, 1e-5),
+        },
+    )
+
+    tune_iter = 5 if args.smoke_test else 300
+    tuner = tune.Tuner(
+        dcgan_train,
+        run_config=tune.RunConfig(
+            name="pbt_dcgan_mnist",
+            stop={"training_iteration": tune_iter},
+            verbose=1,
+        ),
+        tune_config=tune.TuneConfig(
+            metric="is_score",
+            mode="max",
+            num_samples=8,
+            scheduler=scheduler,
+        ),
+        param_space={
+            "netG_lr": tune.choice([0.0001, 0.0002, 0.0005]),
+            "netD_lr": tune.choice([0.0001, 0.0002, 0.0005]),
+            "mnist_model_ref": mnist_model_ref,
+        },
+    )
+    results = tuner.fit()
+    # __tune_end__
+
+    # demo of the trained Generators
+    if not args.smoke_test:
+        checkpoint_paths = [result.checkpoint.to_directory() for result in results]
+        demo_gan(checkpoint_paths)

python/ray/tune/examples/pbt_dcgan_mnist/pbt_dcgan_mnist_trainable.py

@@ -0,0 +1,185 @@
+#!/usr/bin/env python
+"""
+Example of training DCGAN on MNIST using PBT with Tune's Trainable Class
+API.
+"""
+import argparse
+import os
+import random
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.parallel
+import torch.optim as optim
+import torch.utils.data
+from filelock import FileLock
+
+import ray
+from ray import tune
+from ray.tune.examples.pbt_dcgan_mnist.common import (
+    MODEL_PATH,
+    Discriminator,
+    Generator,
+    Net,
+    beta1,
+    demo_gan,
+    get_data_loader,
+    plot_images,
+    train_func,
+    weights_init,
+)
+from ray.tune.schedulers import PopulationBasedTraining
+
+
+# __Trainable_begin__
+class PytorchTrainable(tune.Trainable):
+    def setup(self, config):
+        use_cuda = config.get("use_gpu") and torch.cuda.is_available()
+        self.device = torch.device("cuda" if use_cuda else "cpu")
+        self.netD = Discriminator().to(self.device)
+        self.netD.apply(weights_init)
+        self.netG = Generator().to(self.device)
+        self.netG.apply(weights_init)
+        self.criterion = nn.BCELoss()
+        self.optimizerD = optim.Adam(
+            self.netD.parameters(), lr=config.get("lr", 0.01), betas=(beta1, 0.999)
+        )
+        self.optimizerG = optim.Adam(
+            self.netG.parameters(), lr=config.get("lr", 0.01), betas=(beta1, 0.999)
+        )
+        with FileLock(os.path.expanduser("~/.data.lock")):
+            self.dataloader = get_data_loader(config.get("data_dir", "~/data"))
+        self.mnist_model_ref = config["mnist_model_ref"]
+
+    def step(self):
+        lossG, lossD, is_score = train_func(
+            self.netD,
+            self.netG,
+            self.optimizerG,
+            self.optimizerD,
+            self.criterion,
+            self.dataloader,
+            self._iteration,
+            self.device,
+            self.mnist_model_ref,
+        )
+        return {"lossg": lossG, "lossd": lossD, "is_score": is_score}
+
+    def save_checkpoint(self, checkpoint_dir):
+        path = os.path.join(checkpoint_dir, "checkpoint.pt")
+        torch.save(
+            {
+                "netDmodel": self.netD.state_dict(),
+                "netGmodel": self.netG.state_dict(),
+                "optimD": self.optimizerD.state_dict(),
+                "optimG": self.optimizerG.state_dict(),
+            },
+            path,
+        )
+
+        return checkpoint_dir
+
+    def load_checkpoint(self, checkpoint_dir):
+        path = os.path.join(checkpoint_dir, "checkpoint.pt")
+        checkpoint = torch.load(path)
+        self.netD.load_state_dict(checkpoint["netDmodel"])
+        self.netG.load_state_dict(checkpoint["netGmodel"])
+        self.optimizerD.load_state_dict(checkpoint["optimD"])
+        self.optimizerG.load_state_dict(checkpoint["optimG"])
+
+    def reset_config(self, new_config):
+        if "netD_lr" in new_config:
+            for param_group in self.optimizerD.param_groups:
+                param_group["lr"] = new_config["netD_lr"]
+        if "netG_lr" in new_config:
+            for param_group in self.optimizerG.param_groups:
+                param_group["lr"] = new_config["netG_lr"]
+
+        self.config = new_config
+        return True
+
+
+# __Trainable_end__
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--smoke-test", action="store_true", help="Finish quickly for testing"
+    )
+    parser.add_argument(
+        "--data-dir", type=str, default="~/data/", help="Set the path of the dataset."
+    )
+    args, _ = parser.parse_known_args()
+    ray.init()
+
+    import urllib.request
+
+    # Download a pre-trained MNIST model for inception score calculation.
+    # This is a tiny model (<100kb).
+    if not os.path.exists(MODEL_PATH):
+        print("downloading model")
+        os.makedirs(os.path.dirname(MODEL_PATH), exist_ok=True)
+        urllib.request.urlretrieve(
+            "https://github.com/ray-project/ray/raw/master/python/ray/tune/"
+            "examples/pbt_dcgan_mnist/mnist_cnn.pt",
+            MODEL_PATH,
+        )
+
+    dataloader = get_data_loader()
+    if not args.smoke_test:
+        plot_images(dataloader)
+
+    # load the pretrained mnist classification model for inception_score
+    mnist_cnn = Net()
+    mnist_cnn.load_state_dict(torch.load(MODEL_PATH))
+    mnist_cnn.eval()
+    mnist_model_ref = ray.put(mnist_cnn)
+
+    # __tune_begin__
+    scheduler = PopulationBasedTraining(
+        time_attr="training_iteration",
+        perturbation_interval=5,
+        hyperparam_mutations={
+            # distribution for resampling
+            "netG_lr": lambda: np.random.uniform(1e-2, 1e-5),
+            "netD_lr": lambda: np.random.uniform(1e-2, 1e-5),
+        },
+    )
+
+    tune_iter = 10 if args.smoke_test else 300
+    tuner = tune.Tuner(
+        PytorchTrainable,
+        run_config=tune.RunConfig(
+            name="pbt_dcgan_mnist",
+            stop={"training_iteration": tune_iter},
+            verbose=1,
+            checkpoint_config=tune.CheckpointConfig(checkpoint_at_end=True),
+        ),
+        tune_config=tune.TuneConfig(
+            metric="is_score",
+            mode="max",
+            num_samples=8,
+            scheduler=scheduler,
+            reuse_actors=True,
+        ),
+        param_space={
+            "netG_lr": tune.sample_from(
+                lambda spec: random.choice([0.0001, 0.0002, 0.0005])
+            ),
+            "netD_lr": tune.sample_from(
+                lambda spec: random.choice([0.0001, 0.0002, 0.0005])
+            ),
+            "mnist_model_ref": mnist_model_ref,
+            "data_dir": args.data_dir,
+        },
+    )
+    results = tuner.fit()
+
+    # export_formats=[ExportFormat.MODEL]
+    # __tune_end__
+
+    # demo of the trained Generators
+    if not args.smoke_test:
+        checkpoint_paths = [result.checkpoint.to_directory() for result in results]
+        demo_gan(checkpoint_paths)

python/ray/tune/examples/pbt_example.py

@@ -0,0 +1,146 @@
+#!/usr/bin/env python
+
+import argparse
+import random
+
+import numpy as np
+
+import ray
+from ray import tune
+from ray.tune.schedulers import PopulationBasedTraining
+
+
+class PBTBenchmarkExample(tune.Trainable):
+    """Toy PBT problem for benchmarking adaptive learning rate.
+
+    The goal is to optimize this trainable's accuracy. The accuracy increases
+    fastest at the optimal lr, which is a function of the current accuracy.
+
+    The optimal lr schedule for this problem is the triangle wave as follows.
+    Note that many lr schedules for real models also follow this shape:
+
+     best lr
+      ^
+      |    /\
+      |   /  \
+      |  /    \
+      | /      \
+      ------------> accuracy
+
+    In this problem, using PBT with a population of 2-4 is sufficient to
+    roughly approximate this lr schedule. Higher population sizes will yield
+    faster convergence. Training will not converge without PBT.
+    """
+
+    def setup(self, config):
+        self.lr = config["lr"]
+        self.accuracy = 0.0  # end = 1000
+
+    def step(self):
+        midpoint = 100  # lr starts decreasing after acc > midpoint
+        q_tolerance = 3  # penalize exceeding lr by more than this multiple
+        noise_level = 2  # add gaussian noise to the acc increase
+        # triangle wave:
+        #  - start at 0.001 @ t=0,
+        #  - peak at 0.01 @ t=midpoint,
+        #  - end at 0.001 @ t=midpoint * 2,
+        if self.accuracy < midpoint:
+            optimal_lr = 0.01 * self.accuracy / midpoint
+        else:
+            optimal_lr = 0.01 - 0.01 * (self.accuracy - midpoint) / midpoint
+        optimal_lr = min(0.01, max(0.001, optimal_lr))
+
+        # compute accuracy increase
+        q_err = max(self.lr, optimal_lr) / min(self.lr, optimal_lr)
+        if q_err < q_tolerance:
+            self.accuracy += (1.0 / q_err) * random.random()
+        elif self.lr > optimal_lr:
+            self.accuracy -= (q_err - q_tolerance) * random.random()
+        self.accuracy += noise_level * np.random.normal()
+        self.accuracy = max(0, self.accuracy)
+
+        return {
+            "mean_accuracy": self.accuracy,
+            "cur_lr": self.lr,
+            "optimal_lr": optimal_lr,  # for debugging
+            "q_err": q_err,  # for debugging
+            "done": self.accuracy > midpoint * 2,
+        }
+
+    def save_checkpoint(self, checkpoint_dir):
+        return {
+            "accuracy": self.accuracy,
+            "lr": self.lr,
+        }
+
+    def load_checkpoint(self, checkpoint):
+        self.accuracy = checkpoint["accuracy"]
+
+    def reset_config(self, new_config):
+        self.lr = new_config["lr"]
+        self.config = new_config
+        return True
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--smoke-test", action="store_true", help="Finish quickly for testing"
+    )
+    args, _ = parser.parse_known_args()
+
+    if args.smoke_test:
+        ray.init(num_cpus=2)  # force pausing to happen for test
+
+    perturbation_interval = 5
+    pbt = PopulationBasedTraining(
+        time_attr="training_iteration",
+        perturbation_interval=perturbation_interval,
+        hyperparam_mutations={
+            # distribution for resampling
+            "lr": lambda: random.uniform(0.0001, 0.02),
+            # allow perturbations within this set of categorical values
+            "some_other_factor": [1, 2],
+        },
+    )
+
+    tuner = tune.Tuner(
+        PBTBenchmarkExample,
+        run_config=tune.RunConfig(
+            name="pbt_class_api_example",
+            # Stop when done = True or at some # of train steps (whichever comes first)
+            stop={
+                "done": True,
+                "training_iteration": 10 if args.smoke_test else 1000,
+            },
+            verbose=0,
+            # We recommend matching `perturbation_interval` and `checkpoint_interval`
+            # (e.g. checkpoint every 4 steps, and perturb on those same steps)
+            # or making `perturbation_interval` a multiple of `checkpoint_interval`
+            # (e.g. checkpoint every 2 steps, and perturb every 4 steps).
+            # This is to ensure that the lastest checkpoints are being used by PBT
+            # when trials decide to exploit. If checkpointing and perturbing are not
+            # aligned, then PBT may use a stale checkpoint to resume from.
+            checkpoint_config=tune.CheckpointConfig(
+                checkpoint_frequency=perturbation_interval,
+                checkpoint_score_attribute="mean_accuracy",
+                num_to_keep=4,
+            ),
+        ),
+        tune_config=tune.TuneConfig(
+            scheduler=pbt,
+            metric="mean_accuracy",
+            mode="max",
+            reuse_actors=True,
+            num_samples=8,
+        ),
+        param_space={
+            "lr": 0.0001,
+            # note: this parameter is perturbed but has no effect on
+            # the model training in this example
+            "some_other_factor": 1,
+        },
+    )
+    results = tuner.fit()
+
+    print("Best hyperparameters found were: ", results.get_best_result().config)

python/ray/tune/examples/pbt_function.py

@@ -0,0 +1,181 @@
+#!/usr/bin/env python
+
+import argparse
+import json
+import os
+import random
+import tempfile
+
+import numpy as np
+
+import ray
+from ray import tune
+from ray.tune import Checkpoint
+from ray.tune.schedulers import PopulationBasedTraining
+
+
+def pbt_function(config):
+    """Toy PBT problem for benchmarking adaptive learning rate.
+
+    The goal is to optimize this trainable's accuracy. The accuracy increases
+    fastest at the optimal lr, which is a function of the current accuracy.
+
+    The optimal lr schedule for this problem is the triangle wave as follows.
+    Note that many lr schedules for real models also follow this shape:
+
+     best lr
+      ^
+      |    /\
+      |   /  \
+      |  /    \
+      | /      \
+      ------------> accuracy
+
+    In this problem, using PBT with a population of 2-4 is sufficient to
+    roughly approximate this lr schedule. Higher population sizes will yield
+    faster convergence. Training will not converge without PBT.
+    """
+    lr = config["lr"]
+    checkpoint_interval = config.get("checkpoint_interval", 1)
+
+    accuracy = 0.0  # end = 1000
+
+    # NOTE: See below why step is initialized to 1
+    step = 1
+    checkpoint = tune.get_checkpoint()
+    if checkpoint:
+        with checkpoint.as_directory() as checkpoint_dir:
+            with open(os.path.join(checkpoint_dir, "checkpoint.json"), "r") as f:
+                checkpoint_dict = json.load(f)
+
+        accuracy = checkpoint_dict["acc"]
+        last_step = checkpoint_dict["step"]
+        # Current step should be 1 more than the last checkpoint step
+        step = last_step + 1
+
+    # triangle wave:
+    #  - start at 0.001 @ t=0,
+    #  - peak at 0.01 @ t=midpoint,
+    #  - end at 0.001 @ t=midpoint * 2,
+    midpoint = 100  # lr starts decreasing after acc > midpoint
+    q_tolerance = 3  # penalize exceeding lr by more than this multiple
+    noise_level = 2  # add gaussian noise to the acc increase
+
+    # Let `stop={"done": True}` in the configs below handle trial stopping
+    while True:
+        if accuracy < midpoint:
+            optimal_lr = 0.01 * accuracy / midpoint
+        else:
+            optimal_lr = 0.01 - 0.01 * (accuracy - midpoint) / midpoint
+        optimal_lr = min(0.01, max(0.001, optimal_lr))
+
+        # compute accuracy increase
+        q_err = max(lr, optimal_lr) / min(lr, optimal_lr)
+        if q_err < q_tolerance:
+            accuracy += (1.0 / q_err) * random.random()
+        elif lr > optimal_lr:
+            accuracy -= (q_err - q_tolerance) * random.random()
+        accuracy += noise_level * np.random.normal()
+        accuracy = max(0, accuracy)
+
+        metrics = {
+            "mean_accuracy": accuracy,
+            "cur_lr": lr,
+            "optimal_lr": optimal_lr,  # for debugging
+            "q_err": q_err,  # for debugging
+            "done": accuracy > midpoint * 2,  # this stops the training process
+        }
+
+        if step % checkpoint_interval == 0:
+            # Checkpoint every `checkpoint_interval` steps
+            # NOTE: if we initialized `step=0` above, our checkpointing and perturbing
+            # would be out of sync by 1 step.
+            # Ex: if `checkpoint_interval` = `perturbation_interval` = 3
+            # step:                0 (checkpoint)  1     2            3 (checkpoint)
+            # training_iteration:  1               2     3 (perturb)  4
+            with tempfile.TemporaryDirectory() as tempdir:
+                with open(os.path.join(tempdir, "checkpoint.json"), "w") as f:
+                    checkpoint_dict = {"acc": accuracy, "step": step}
+                    json.dump(checkpoint_dict, f)
+                tune.report(metrics, checkpoint=Checkpoint.from_directory(tempdir))
+        else:
+            tune.report(metrics)
+        step += 1
+
+
+def run_tune_pbt(smoke_test=False):
+    perturbation_interval = 5
+    pbt = PopulationBasedTraining(
+        time_attr="training_iteration",
+        perturbation_interval=perturbation_interval,
+        hyperparam_mutations={
+            # distribution for resampling
+            "lr": tune.uniform(0.0001, 0.02),
+            # allow perturbations within this set of categorical values
+            "some_other_factor": [1, 2],
+        },
+    )
+
+    tuner = tune.Tuner(
+        pbt_function,
+        run_config=tune.RunConfig(
+            name="pbt_function_api_example",
+            verbose=False,
+            stop={
+                # Stop when done = True or at some # of train steps
+                # (whichever comes first)
+                "done": True,
+                "training_iteration": 10 if smoke_test else 1000,
+            },
+            failure_config=tune.FailureConfig(
+                fail_fast=True,
+            ),
+            checkpoint_config=tune.CheckpointConfig(
+                checkpoint_score_attribute="mean_accuracy",
+                num_to_keep=2,
+            ),
+        ),
+        tune_config=tune.TuneConfig(
+            scheduler=pbt,
+            metric="mean_accuracy",
+            mode="max",
+            num_samples=8,
+            reuse_actors=True,
+        ),
+        param_space={
+            "lr": 0.0001,
+            # Note: `some_other_factor` is perturbed because it is specified under
+            # the PBT scheduler's `hyperparam_mutations` argument, but has no effect on
+            # the model training in this example
+            "some_other_factor": 1,
+            # Note: `checkpoint_interval` will not be perturbed (since it's not
+            # included above), and it will be used to determine how many steps to take
+            # between each checkpoint.
+            # We recommend matching `perturbation_interval` and `checkpoint_interval`
+            # (e.g. checkpoint every 4 steps, and perturb on those same steps)
+            # or making `perturbation_interval` a multiple of `checkpoint_interval`
+            # (e.g. checkpoint every 2 steps, and perturb every 4 steps).
+            # This is to ensure that the lastest checkpoints are being used by PBT
+            # when trials decide to exploit. If checkpointing and perturbing are not
+            # aligned, then PBT may use a stale checkpoint to resume from.
+            "checkpoint_interval": perturbation_interval,
+        },
+    )
+    results = tuner.fit()
+
+    print("Best hyperparameters found were: ", results.get_best_result().config)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--smoke-test",
+        action="store_true",
+        default=False,
+        help="Finish quickly for testing",
+    )
+    args, _ = parser.parse_known_args()
+    if args.smoke_test:
+        ray.init(num_cpus=2)  # force pausing to happen for test
+
+    run_tune_pbt(smoke_test=args.smoke_test)

python/ray/tune/examples/pbt_memnn_example.py

@@ -0,0 +1,325 @@
+"""Example training a memory neural net on the bAbI dataset.
+
+References Keras and is based off of https://keras.io/examples/babi_memnn/.
+"""
+
+from __future__ import print_function
+
+import argparse
+import os
+import re
+import sys
+import tarfile
+
+import numpy as np
+from filelock import FileLock
+
+from ray import tune
+
+if sys.version_info >= (3, 12):
+    # Skip this test in Python 3.12+ because TensorFlow is not supported.
+    sys.exit(0)
+else:
+    from tensorflow.keras.layers import (
+        LSTM,
+        Activation,
+        Dense,
+        Dropout,
+        Embedding,
+        Input,
+        Permute,
+        add,
+        concatenate,
+        dot,
+    )
+    from tensorflow.keras.models import Model, Sequential, load_model
+    from tensorflow.keras.optimizers import RMSprop
+    from tensorflow.keras.preprocessing.sequence import pad_sequences
+    from tensorflow.keras.utils import get_file
+
+
+def tokenize(sent):
+    """Return the tokens of a sentence including punctuation.
+
+    >>> tokenize("Bob dropped the apple. Where is the apple?")
+    ["Bob", "dropped", "the", "apple", ".", "Where", "is", "the", "apple", "?"]
+    """
+    return [x.strip() for x in re.split(r"(\W+)?", sent) if x and x.strip()]
+
+
+def parse_stories(lines, only_supporting=False):
+    """Parse stories provided in the bAbi tasks format
+
+    If only_supporting is true, only the sentences
+    that support the answer are kept.
+    """
+    data = []
+    story = []
+    for line in lines:
+        line = line.decode("utf-8").strip()
+        nid, line = line.split(" ", 1)
+        nid = int(nid)
+        if nid == 1:
+            story = []
+        if "\t" in line:
+            q, a, supporting = line.split("\t")
+            q = tokenize(q)
+            if only_supporting:
+                # Only select the related substory
+                supporting = map(int, supporting.split())
+                substory = [story[i - 1] for i in supporting]
+            else:
+                # Provide all the substories
+                substory = [x for x in story if x]
+            data.append((substory, q, a))
+            story.append("")
+        else:
+            sent = tokenize(line)
+            story.append(sent)
+    return data
+
+
+def get_stories(f, only_supporting=False, max_length=None):
+    """Given a file name, read the file,
+    retrieve the stories,
+    and then convert the sentences into a single story.
+
+    If max_length is supplied,
+    any stories longer than max_length tokens will be discarded.
+    """
+
+    def flatten(data):
+        return sum(data, [])
+
+    data = parse_stories(f.readlines(), only_supporting=only_supporting)
+    data = [
+        (flatten(story), q, answer)
+        for story, q, answer in data
+        if not max_length or len(flatten(story)) < max_length
+    ]
+    return data
+
+
+def vectorize_stories(word_idx, story_maxlen, query_maxlen, data):
+    inputs, queries, answers = [], [], []
+    for story, query, answer in data:
+        inputs.append([word_idx[w] for w in story])
+        queries.append([word_idx[w] for w in query])
+        answers.append(word_idx[answer])
+    return (
+        pad_sequences(inputs, maxlen=story_maxlen),
+        pad_sequences(queries, maxlen=query_maxlen),
+        np.array(answers),
+    )
+
+
+def read_data(finish_fast=False):
+    # Get the file
+    try:
+        path = get_file(
+            "babi-tasks-v1-2.tar.gz",
+            origin="https://s3.amazonaws.com/text-datasets/"
+            "babi_tasks_1-20_v1-2.tar.gz",
+        )
+    except Exception:
+        print(
+            "Error downloading dataset, please download it manually:\n"
+            "$ wget http://www.thespermwhale.com/jaseweston/babi/tasks_1-20_v1-2"  # noqa: E501
+            ".tar.gz\n"
+            "$ mv tasks_1-20_v1-2.tar.gz ~/.keras/datasets/babi-tasks-v1-2.tar.gz"  # noqa: E501
+        )
+        raise
+
+    # Choose challenge
+    challenges = {
+        # QA1 with 10,000 samples
+        "single_supporting_fact_10k": "tasks_1-20_v1-2/en-10k/qa1_"
+        "single-supporting-fact_{}.txt",
+        # QA2 with 10,000 samples
+        "two_supporting_facts_10k": "tasks_1-20_v1-2/en-10k/qa2_"
+        "two-supporting-facts_{}.txt",
+    }
+    challenge_type = "single_supporting_fact_10k"
+    challenge = challenges[challenge_type]
+
+    with tarfile.open(path) as tar:
+        train_stories = get_stories(tar.extractfile(challenge.format("train")))
+        test_stories = get_stories(tar.extractfile(challenge.format("test")))
+    if finish_fast:
+        train_stories = train_stories[:64]
+        test_stories = test_stories[:64]
+    return train_stories, test_stories
+
+
+class MemNNModel(tune.Trainable):
+    def build_model(self):
+        """Helper method for creating the model"""
+        vocab = set()
+        for story, q, answer in self.train_stories + self.test_stories:
+            vocab |= set(story + q + [answer])
+        vocab = sorted(vocab)
+
+        # Reserve 0 for masking via pad_sequences
+        vocab_size = len(vocab) + 1
+        story_maxlen = max(len(x) for x, _, _ in self.train_stories + self.test_stories)
+        query_maxlen = max(len(x) for _, x, _ in self.train_stories + self.test_stories)
+
+        word_idx = {c: i + 1 for i, c in enumerate(vocab)}
+        self.inputs_train, self.queries_train, self.answers_train = vectorize_stories(
+            word_idx, story_maxlen, query_maxlen, self.train_stories
+        )
+        self.inputs_test, self.queries_test, self.answers_test = vectorize_stories(
+            word_idx, story_maxlen, query_maxlen, self.test_stories
+        )
+
+        # placeholders
+        input_sequence = Input((story_maxlen,))
+        question = Input((query_maxlen,))
+
+        # encoders
+        # embed the input sequence into a sequence of vectors
+        input_encoder_m = Sequential()
+        input_encoder_m.add(Embedding(input_dim=vocab_size, output_dim=64))
+        input_encoder_m.add(Dropout(self.config.get("dropout", 0.3)))
+        # output: (samples, story_maxlen, embedding_dim)
+
+        # embed the input into a sequence of vectors of size query_maxlen
+        input_encoder_c = Sequential()
+        input_encoder_c.add(Embedding(input_dim=vocab_size, output_dim=query_maxlen))
+        input_encoder_c.add(Dropout(self.config.get("dropout", 0.3)))
+        # output: (samples, story_maxlen, query_maxlen)
+
+        # embed the question into a sequence of vectors
+        question_encoder = Sequential()
+        question_encoder.add(
+            Embedding(input_dim=vocab_size, output_dim=64, input_length=query_maxlen)
+        )
+        question_encoder.add(Dropout(self.config.get("dropout", 0.3)))
+        # output: (samples, query_maxlen, embedding_dim)
+
+        # encode input sequence and questions (which are indices)
+        # to sequences of dense vectors
+        input_encoded_m = input_encoder_m(input_sequence)
+        input_encoded_c = input_encoder_c(input_sequence)
+        question_encoded = question_encoder(question)
+
+        # compute a "match" between the first input vector sequence
+        # and the question vector sequence
+        # shape: `(samples, story_maxlen, query_maxlen)`
+        match = dot([input_encoded_m, question_encoded], axes=(2, 2))
+        match = Activation("softmax")(match)
+
+        # add the match matrix with the second input vector sequence
+        response = add(
+            [match, input_encoded_c]
+        )  # (samples, story_maxlen, query_maxlen)
+        response = Permute((2, 1))(response)  # (samples, query_maxlen, story_maxlen)
+
+        # concatenate the match matrix with the question vector sequence
+        answer = concatenate([response, question_encoded])
+
+        # the original paper uses a matrix multiplication.
+        # we choose to use a RNN instead.
+        answer = LSTM(32)(answer)  # (samples, 32)
+
+        # one regularization layer -- more would probably be needed.
+        answer = Dropout(self.config.get("dropout", 0.3))(answer)
+        answer = Dense(vocab_size)(answer)  # (samples, vocab_size)
+        # we output a probability distribution over the vocabulary
+        answer = Activation("softmax")(answer)
+
+        # build the final model
+        model = Model([input_sequence, question], answer)
+        return model
+
+    def setup(self, config):
+        with FileLock(os.path.expanduser("~/.tune.lock")):
+            self.train_stories, self.test_stories = read_data(config["finish_fast"])
+        model = self.build_model()
+        rmsprop = RMSprop(
+            lr=self.config.get("lr", 1e-3), rho=self.config.get("rho", 0.9)
+        )
+        model.compile(
+            optimizer=rmsprop,
+            loss="sparse_categorical_crossentropy",
+            metrics=["accuracy"],
+        )
+        self.model = model
+
+    def step(self):
+        # train
+        self.model.fit(
+            [self.inputs_train, self.queries_train],
+            self.answers_train,
+            batch_size=self.config.get("batch_size", 32),
+            epochs=self.config.get("epochs", 1),
+            validation_data=([self.inputs_test, self.queries_test], self.answers_test),
+            verbose=0,
+        )
+        _, accuracy = self.model.evaluate(
+            [self.inputs_train, self.queries_train], self.answers_train, verbose=0
+        )
+        return {"mean_accuracy": accuracy}
+
+    def save_checkpoint(self, checkpoint_dir):
+        file_path = checkpoint_dir + "/model"
+        self.model.save(file_path)
+
+    def load_checkpoint(self, checkpoint_dir):
+        # See https://stackoverflow.com/a/42763323
+        del self.model
+        file_path = checkpoint_dir + "/model"
+        self.model = load_model(file_path)
+
+
+if __name__ == "__main__":
+    import ray
+    from ray.tune.schedulers import PopulationBasedTraining
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--smoke-test", action="store_true", help="Finish quickly for testing"
+    )
+    args, _ = parser.parse_known_args()
+
+    if args.smoke_test:
+        ray.init(num_cpus=2)
+
+    perturbation_interval = 2
+    pbt = PopulationBasedTraining(
+        perturbation_interval=perturbation_interval,
+        hyperparam_mutations={
+            "dropout": lambda: np.random.uniform(0, 1),
+            "lr": lambda: 10 ** np.random.randint(-10, 0),
+            "rho": lambda: np.random.uniform(0, 1),
+        },
+    )
+
+    tuner = tune.Tuner(
+        MemNNModel,
+        run_config=tune.RunConfig(
+            name="pbt_babi_memnn",
+            stop={"training_iteration": 4 if args.smoke_test else 100},
+            checkpoint_config=tune.CheckpointConfig(
+                checkpoint_frequency=perturbation_interval,
+                checkpoint_score_attribute="mean_accuracy",
+                num_to_keep=2,
+            ),
+        ),
+        tune_config=tune.TuneConfig(
+            scheduler=pbt,
+            metric="mean_accuracy",
+            mode="max",
+            num_samples=2,
+            reuse_actors=True,
+        ),
+        param_space={
+            "finish_fast": args.smoke_test,
+            "batch_size": 32,
+            "epochs": 1,
+            "dropout": 0.3,
+            "lr": 0.01,
+            "rho": 0.9,
+        },
+    )
+    tuner.fit()

python/ray/tune/examples/pbt_ppo_example.py

@@ -0,0 +1,75 @@
+#!/usr/bin/env python
+"""Example of using PBT with RLlib.
+
+Note that this requires a cluster with at least 8 GPUs in order for all trials
+to run concurrently, otherwise PBT will round-robin train the trials which
+is less efficient (or you can set {"gpu": 0} to use CPUs for SGD instead).
+
+Note that Tune in general does not need 8 GPUs, and this is just a more
+computationally demanding example.
+"""
+
+import random
+
+from ray import tune
+from ray.rllib.algorithms.ppo import PPO
+from ray.tune.schedulers import PopulationBasedTraining
+
+if __name__ == "__main__":
+    # Postprocess the perturbed config to ensure it's still valid
+    def explore(config):
+        # ensure we collect enough timesteps to do sgd
+        if config["train_batch_size"] < config["sgd_minibatch_size"] * 2:
+            config["train_batch_size"] = config["sgd_minibatch_size"] * 2
+        # ensure we run at least one sgd iter
+        if config["num_sgd_iter"] < 1:
+            config["num_sgd_iter"] = 1
+        return config
+
+    pbt = PopulationBasedTraining(
+        time_attr="time_total_s",
+        perturbation_interval=120,
+        resample_probability=0.25,
+        # Specifies the mutations of these hyperparams
+        hyperparam_mutations={
+            "lambda": lambda: random.uniform(0.9, 1.0),
+            "clip_param": lambda: random.uniform(0.01, 0.5),
+            "lr": [1e-3, 5e-4, 1e-4, 5e-5, 1e-5],
+            "num_sgd_iter": lambda: random.randint(1, 30),
+            "sgd_minibatch_size": lambda: random.randint(128, 16384),
+            "train_batch_size": lambda: random.randint(2000, 160000),
+        },
+        custom_explore_fn=explore,
+    )
+
+    tuner = tune.Tuner(
+        PPO,
+        run_config=tune.RunConfig(
+            name="pbt_humanoid_test",
+        ),
+        tune_config=tune.TuneConfig(
+            scheduler=pbt,
+            num_samples=8,
+            metric="episode_reward_mean",
+            mode="max",
+            reuse_actors=True,
+        ),
+        param_space={
+            "env": "Humanoid-v1",
+            "kl_coeff": 1.0,
+            "num_workers": 8,
+            "num_gpus": 1,
+            "model": {"free_log_std": True},
+            # These params are tuned from a fixed starting value.
+            "lambda": 0.95,
+            "clip_param": 0.2,
+            "lr": 1e-4,
+            # These params start off randomly drawn from a set.
+            "num_sgd_iter": tune.choice([10, 20, 30]),
+            "sgd_minibatch_size": tune.choice([128, 512, 2048]),
+            "train_batch_size": tune.choice([10000, 20000, 40000]),
+        },
+    )
+    results = tuner.fit()
+
+    print("best hyperparameters: ", results.get_best_result().config)

python/ray/tune/examples/pbt_transformers/pbt_transformers.py

@@ -0,0 +1,165 @@
+"""
+This example is uses the official
+huggingface transformers `hyperparameter_search` API.
+"""
+
+import os
+
+from transformers import (
+    AutoConfig,
+    AutoModelForSequenceClassification,
+    AutoTokenizer,
+    GlueDataset,
+    GlueDataTrainingArguments,
+    Trainer,
+    TrainingArguments,
+    glue_tasks_num_labels,
+)
+
+from ray import tune
+from ray.tune import CheckpointConfig, CLIReporter
+from ray.tune.examples.pbt_transformers.utils import (
+    build_compute_metrics_fn,
+    download_data,
+)
+from ray.tune.schedulers import PopulationBasedTraining
+
+
+def tune_transformer(num_samples=8, gpus_per_trial=0, smoke_test=False):
+    data_dir_name = "./data" if not smoke_test else "./test_data"
+    data_dir = os.path.abspath(os.path.join(os.getcwd(), data_dir_name))
+    if not os.path.exists(data_dir):
+        os.mkdir(data_dir, 0o755)
+
+    # Change these as needed.
+    model_name = (
+        "bert-base-uncased" if not smoke_test else "sshleifer/tiny-distilroberta-base"
+    )
+    task_name = "rte"
+
+    task_data_dir = os.path.join(data_dir, task_name.upper())
+
+    num_labels = glue_tasks_num_labels[task_name]
+
+    config = AutoConfig.from_pretrained(
+        model_name, num_labels=num_labels, finetuning_task=task_name
+    )
+
+    # Download and cache tokenizer, model, and features
+    print("Downloading and caching Tokenizer")
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+    # Triggers tokenizer download to cache
+    print("Downloading and caching pre-trained model")
+    AutoModelForSequenceClassification.from_pretrained(
+        model_name,
+        config=config,
+    )
+
+    def get_model():
+        return AutoModelForSequenceClassification.from_pretrained(
+            model_name,
+            config=config,
+        )
+
+    # Download data.
+    download_data(task_name, data_dir)
+
+    data_args = GlueDataTrainingArguments(task_name=task_name, data_dir=task_data_dir)
+
+    train_dataset = GlueDataset(
+        data_args, tokenizer=tokenizer, mode="train", cache_dir=task_data_dir
+    )
+    eval_dataset = GlueDataset(
+        data_args, tokenizer=tokenizer, mode="dev", cache_dir=task_data_dir
+    )
+
+    training_args = TrainingArguments(
+        output_dir=".",
+        learning_rate=1e-5,  # config
+        do_train=True,
+        do_eval=True,
+        use_cpu=gpus_per_trial <= 0,
+        eval_strategy="epoch",
+        save_strategy="epoch",
+        load_best_model_at_end=True,
+        num_train_epochs=2,  # config
+        max_steps=-1,
+        per_device_train_batch_size=16,  # config
+        per_device_eval_batch_size=16,  # config
+        warmup_steps=0,
+        weight_decay=0.1,  # config
+        logging_dir="./logs",
+        skip_memory_metrics=True,
+        report_to="none",
+    )
+
+    trainer = Trainer(
+        model_init=get_model,
+        args=training_args,
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+        compute_metrics=build_compute_metrics_fn(task_name),
+    )
+
+    tune_config = {
+        "per_device_train_batch_size": 32,
+        "per_device_eval_batch_size": 32,
+        "num_train_epochs": tune.choice([2, 3, 4, 5]),
+        "max_steps": 1 if smoke_test else -1,  # Used for smoke test.
+    }
+
+    scheduler = PopulationBasedTraining(
+        time_attr="training_iteration",
+        metric="eval_acc",
+        mode="max",
+        perturbation_interval=1,
+        hyperparam_mutations={
+            "weight_decay": tune.uniform(0.0, 0.3),
+            "learning_rate": tune.uniform(1e-5, 5e-5),
+            "per_device_train_batch_size": [16, 32, 64],
+        },
+    )
+
+    reporter = CLIReporter(
+        parameter_columns={
+            "weight_decay": "w_decay",
+            "learning_rate": "lr",
+            "per_device_train_batch_size": "train_bs/gpu",
+            "num_train_epochs": "num_epochs",
+        },
+        metric_columns=["eval_acc", "eval_loss", "epoch", "training_iteration"],
+    )
+
+    trainer.hyperparameter_search(
+        hp_space=lambda _: tune_config,
+        backend="ray",
+        n_trials=num_samples,
+        resources_per_trial={"cpu": 1, "gpu": gpus_per_trial},
+        scheduler=scheduler,
+        checkpoint_config=CheckpointConfig(
+            num_to_keep=1,
+            checkpoint_score_attribute="training_iteration",
+        ),
+        stop={"training_iteration": 1} if smoke_test else None,
+        progress_reporter=reporter,
+        local_dir="~/ray_results/",
+        name="tune_transformer_pbt",
+        log_to_file=True,
+    )
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--smoke-test", action="store_true", help="Finish quickly for testing"
+    )
+    args, _ = parser.parse_known_args()
+
+    if args.smoke_test:
+        tune_transformer(num_samples=1, gpus_per_trial=0, smoke_test=True)
+    else:
+        # You can change the number of GPUs here:
+        tune_transformer(num_samples=8, gpus_per_trial=1)

python/ray/tune/examples/pbt_transformers/test_data/RTE/dev.tsv

@@ -0,0 +1,10 @@
+index	sentence1	sentence2	label
+0	Dana Reeve, the widow of the actor Christopher Reeve, has died of lung cancer at age 44, according to the Christopher Reeve Foundation.	Christopher Reeve had an accident.	not_entailment
+1	Yet, we now are discovering that antibiotics are losing their effectiveness against illness. Disease-causing bacteria are mutating faster than we can come up with new antibiotics to fight the new variations.	Bacteria is winning the war against antibiotics.	entailment
+2	Cairo is now home to some 15 million people - a burgeoning population that produces approximately 10,000 tonnes of rubbish per day, putting an enormous strain on public services. In the past 10 years, the government has tried hard to encourage private investment in the refuse sector, but some estimate 4,000 tonnes of waste is left behind every day, festering in the heat as it waits for someone to clear it up. It is often the people in the poorest neighbourhoods that are worst affected. But in some areas they are fighting back. In Shubra, one of the northern districts of the city, the residents have taken to the streets armed with dustpans and brushes to clean up public areas which have been used as public dumps.	15 million tonnes of rubbish are produced daily in Cairo.	not_entailment
+3	The Amish community in Pennsylvania, which numbers about 55,000, lives an agrarian lifestyle, shunning technological advances like electricity and automobiles. And many say their insular lifestyle gives them a sense that they are protected from the violence of American society. But as residents gathered near the school, some wearing traditional garb and arriving in horse-drawn buggies, they said that sense of safety had been shattered. "If someone snaps and wants to do something stupid, there's no distance that's going to stop them," said Jake King, 56, an Amish lantern maker who knew several families whose children had been shot.	Pennsylvania has the biggest Amish community in the U.S.	not_entailment
+4	Security forces were on high alert after an election campaign in which more than 1,000 people, including seven election candidates, have been killed.	Security forces were on high alert after a campaign marred by violence.	entailment
+5	In 1979, the leaders signed the Egypt-Israel peace treaty on the White House lawn. Both President Begin and Sadat received the Nobel Peace Prize for their work. The two nations have enjoyed peaceful relations to this day.	The Israel-Egypt Peace Agreement was signed in 1979.	entailment
+6	singer and actress Britney Spears, 24, has filled papers in Los Angeles County Superior Court to divorce her husband Kevin Federline, 28. A spokeswoman for the court, Kathy Roberts stated that the papers cited irreconcilable differences" as the reason for the divorce and have, according to the courts, been legally separated as of Monday, November 6, the same day that Spears appeared on Late Night with David Letterman.	Spears is to divorce from Kevin Federline.	entailment
+7	Following the successful bid to bring the 2010 Ryder Cup to Wales, the Wales Tourist Board has wasted little time in commissioning work to ensure that the benefits accruing from the event are felt throughout the country.	Wales to host 2010 Ryder Cup.	entailment
+8	Steve Jobs was attacked by Sculley and other Apple executives for not delivering enough hot new products and resigned from the company a few weeks later.	Steve Jobs worked for Apple.	entailment

python/ray/tune/examples/pbt_transformers/test_data/RTE/train.tsv

@@ -0,0 +1,10 @@
+index	sentence1	sentence2	label
+0	No Weapons of Mass Destruction Found in Iraq Yet.	Weapons of Mass Destruction Found in Iraq.	not_entailment
+1	A place of sorrow, after Pope John Paul II died, became a place of celebration, as Roman Catholic faithful gathered in downtown Chicago to mark the installation of new Pope Benedict XVI.	Pope Benedict XVI is the new leader of the Roman Catholic Church.	entailment
+2	Herceptin was already approved to treat the sickest breast cancer patients, and the company said, Monday, it will discuss with federal regulators the possibility of prescribing the drug for more breast cancer patients.	Herceptin can be used to treat breast cancer.	entailment
+3	Judie Vivian, chief executive at ProMedica, a medical service company that helps sustain the 2-year-old Vietnam Heart Institute in Ho Chi Minh City (formerly Saigon), said that so far about 1,500 children have received treatment.	The previous name of Ho Chi Minh City was Saigon.	entailment
+4	A man is due in court later charged with the murder 26 years ago of a teenager whose case was the first to be featured on BBC One's Crimewatch. Colette Aram, 16, was walking to her boyfriend's house in Keyworth, Nottinghamshire, on 30 October 1983 when she disappeared. Her body was later found in a field close to her home. Paul Stewart Hutchinson, 50, has been charged with murder and is due before Nottingham magistrates later.	Paul Stewart Hutchinson is accused of having stabbed a girl.	not_entailment
+5	Britain said, Friday, that it has barred cleric, Omar Bakri, from returning to the country from Lebanon, where he was released by police after being detained for 24 hours.	Bakri was briefly detained, but was released.	entailment
+6	Nearly 4 million children who have at least one parent who entered the U.S. illegally were born in the United States and are U.S. citizens as a result, according to the study conducted by the Pew Hispanic Center. That's about three quarters of the estimated 5.5 million children of illegal immigrants inside the United States, according to the study. About 1.8 million children of undocumented immigrants live in poverty, the study found.	Three quarters of U.S. illegal immigrants have children.	not_entailment
+7	Like the United States, U.N. officials are also dismayed that Aristide killed a conference called by Prime Minister Robert Malval in Port-au-Prince in hopes of bringing all the feuding parties together.	Aristide had Prime Minister Robert Malval  murdered in Port-au-Prince.	not_entailment
+8	WASHINGTON --  A newly declassified narrative of the Bush administration's advice to the CIA on harsh interrogations shows that the small group of Justice Department lawyers who wrote memos authorizing controversial interrogation techniques were operating not on their own but with direction from top administration officials, including then-Vice President Dick Cheney and national security adviser Condoleezza Rice. At the same time, the narrative suggests that then-Defense Secretary Donald H. Rumsfeld and then-Secretary of State Colin Powell were largely left out of the decision-making process.	Dick Cheney was the Vice President of Bush.	entailment

python/ray/tune/examples/pbt_transformers/utils.py

@@ -0,0 +1,46 @@
+"""Utilities to load and cache data."""
+
+import os
+from typing import Callable, Dict
+
+import numpy as np
+from transformers import EvalPrediction, glue_compute_metrics, glue_output_modes
+
+
+def build_compute_metrics_fn(task_name: str) -> Callable[[EvalPrediction], Dict]:
+    """Function from transformers/examples/text-classification/run_glue.py"""
+    output_mode = glue_output_modes[task_name]
+
+    def compute_metrics_fn(p: EvalPrediction):
+        if output_mode == "classification":
+            preds = np.argmax(p.predictions, axis=1)
+        elif output_mode == "regression":
+            preds = np.squeeze(p.predictions)
+        metrics = glue_compute_metrics(task_name, preds, p.label_ids)
+        return metrics
+
+    return compute_metrics_fn
+
+
+def download_data(task_name, data_dir="./data"):
+    # Download RTE training data
+    print("Downloading dataset.")
+    import urllib
+    import zipfile
+
+    if task_name == "rte":
+        url = "https://dl.fbaipublicfiles.com/glue/data/RTE.zip"
+    else:
+        raise ValueError("Unknown task: {}".format(task_name))
+    data_file = os.path.join(data_dir, "{}.zip".format(task_name))
+    if not os.path.exists(data_file):
+        urllib.request.urlretrieve(url, data_file)
+        with zipfile.ZipFile(data_file) as zip_ref:
+            zip_ref.extractall(data_dir)
+        print("Downloaded data for task {} to {}".format(task_name, data_dir))
+    else:
+        print(
+            "Data already exists. Using downloaded data for task {} from {}".format(
+                task_name, data_dir
+            )
+        )

python/ray/tune/examples/pbt_tune_cifar10_with_keras.py

@@ -0,0 +1,243 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""Train keras CNN on the CIFAR10 small images dataset.
+
+The model comes from: https://zhuanlan.zhihu.com/p/29214791,
+and it gets to about 87% validation accuracy in 100 epochs.
+
+Note that the script requires a machine with 4 GPUs. You
+can set {"gpu": 0} to use CPUs for training, although
+it is less efficient.
+"""
+
+from __future__ import print_function
+
+import argparse
+
+import numpy as np
+import tensorflow as tf
+from tensorflow.keras.datasets import cifar10
+from tensorflow.keras.layers import (
+    Convolution2D,
+    Dense,
+    Dropout,
+    Flatten,
+    Input,
+    MaxPooling2D,
+)
+from tensorflow.keras.models import Model, load_model
+from tensorflow.keras.preprocessing.image import ImageDataGenerator
+
+from ray import tune
+from ray.tune import Trainable
+from ray.tune.schedulers import PopulationBasedTraining
+
+num_classes = 10
+NUM_SAMPLES = 128
+
+
+class Cifar10Model(Trainable):
+    def _read_data(self):
+        # The data, split between train and test sets:
+        (x_train, y_train), (x_test, y_test) = cifar10.load_data()
+
+        # Convert class vectors to binary class matrices.
+        y_train = tf.keras.utils.to_categorical(y_train, num_classes)
+        y_test = tf.keras.utils.to_categorical(y_test, num_classes)
+
+        x_train = x_train.astype("float32")
+        x_train /= 255
+        x_test = x_test.astype("float32")
+        x_test /= 255
+
+        return (x_train, y_train), (x_test, y_test)
+
+    def _build_model(self, input_shape):
+        x = Input(shape=(32, 32, 3))
+        y = x
+        y = Convolution2D(
+            filters=64,
+            kernel_size=3,
+            strides=1,
+            padding="same",
+            activation="relu",
+            kernel_initializer="he_normal",
+        )(y)
+        y = Convolution2D(
+            filters=64,
+            kernel_size=3,
+            strides=1,
+            padding="same",
+            activation="relu",
+            kernel_initializer="he_normal",
+        )(y)
+        y = MaxPooling2D(pool_size=2, strides=2, padding="same")(y)
+
+        y = Convolution2D(
+            filters=128,
+            kernel_size=3,
+            strides=1,
+            padding="same",
+            activation="relu",
+            kernel_initializer="he_normal",
+        )(y)
+        y = Convolution2D(
+            filters=128,
+            kernel_size=3,
+            strides=1,
+            padding="same",
+            activation="relu",
+            kernel_initializer="he_normal",
+        )(y)
+        y = MaxPooling2D(pool_size=2, strides=2, padding="same")(y)
+
+        y = Convolution2D(
+            filters=256,
+            kernel_size=3,
+            strides=1,
+            padding="same",
+            activation="relu",
+            kernel_initializer="he_normal",
+        )(y)
+        y = Convolution2D(
+            filters=256,
+            kernel_size=3,
+            strides=1,
+            padding="same",
+            activation="relu",
+            kernel_initializer="he_normal",
+        )(y)
+        y = MaxPooling2D(pool_size=2, strides=2, padding="same")(y)
+
+        y = Flatten()(y)
+        y = Dropout(self.config.get("dropout", 0.5))(y)
+        y = Dense(units=10, activation="softmax", kernel_initializer="he_normal")(y)
+
+        model = Model(inputs=x, outputs=y, name="model1")
+        return model
+
+    def setup(self, config):
+        self.train_data, self.test_data = self._read_data()
+        x_train = self.train_data[0]
+        model = self._build_model(x_train.shape[1:])
+
+        opt = tf.keras.optimizers.Adadelta(
+            lr=self.config.get("lr", 1e-4), weight_decay=self.config.get("decay", 1e-4)
+        )
+        model.compile(
+            loss="categorical_crossentropy", optimizer=opt, metrics=["accuracy"]
+        )
+        self.model = model
+
+    def step(self):
+        x_train, y_train = self.train_data
+        x_train, y_train = x_train[:NUM_SAMPLES], y_train[:NUM_SAMPLES]
+        x_test, y_test = self.test_data
+        x_test, y_test = x_test[:NUM_SAMPLES], y_test[:NUM_SAMPLES]
+
+        aug_gen = ImageDataGenerator(
+            # set input mean to 0 over the dataset
+            featurewise_center=False,
+            # set each sample mean to 0
+            samplewise_center=False,
+            # divide inputs by dataset std
+            featurewise_std_normalization=False,
+            # divide each input by its std
+            samplewise_std_normalization=False,
+            # apply ZCA whitening
+            zca_whitening=False,
+            # randomly rotate images in the range (degrees, 0 to 180)
+            rotation_range=0,
+            # randomly shift images horizontally (fraction of total width)
+            width_shift_range=0.1,
+            # randomly shift images vertically (fraction of total height)
+            height_shift_range=0.1,
+            # randomly flip images
+            horizontal_flip=True,
+            # randomly flip images
+            vertical_flip=False,
+        )
+
+        aug_gen.fit(x_train)
+        batch_size = self.config.get("batch_size", 64)
+        gen = aug_gen.flow(x_train, y_train, batch_size=batch_size)
+        self.model.fit_generator(
+            generator=gen, epochs=self.config.get("epochs", 1), validation_data=None
+        )
+
+        # loss, accuracy
+        _, accuracy = self.model.evaluate(x_test, y_test, verbose=0)
+        return {"mean_accuracy": accuracy}
+
+    def save_checkpoint(self, checkpoint_dir):
+        file_path = checkpoint_dir + "/model"
+        self.model.save(file_path)
+
+    def load_checkpoint(self, checkpoint_dir):
+        # See https://stackoverflow.com/a/42763323
+        del self.model
+        file_path = checkpoint_dir + "/model"
+        self.model = load_model(file_path)
+
+    def cleanup(self):
+        # If need, save your model when exit.
+        # saved_path = self.model.save(self.logdir)
+        # print("save model at: ", saved_path)
+        pass
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--smoke-test", action="store_true", help="Finish quickly for testing"
+    )
+    args, _ = parser.parse_known_args()
+
+    space = {
+        "epochs": 1,
+        "batch_size": 64,
+        "lr": tune.grid_search([10**-4, 10**-5]),
+        "decay": tune.sample_from(lambda config: config["lr"] / 100.0),
+        "dropout": tune.grid_search([0.25, 0.5]),
+    }
+    if args.smoke_test:
+        space["lr"] = 10**-4
+        space["dropout"] = 0.5
+
+    perturbation_interval = 10
+    pbt = PopulationBasedTraining(
+        time_attr="training_iteration",
+        perturbation_interval=perturbation_interval,
+        hyperparam_mutations={
+            "dropout": lambda _: np.random.uniform(0, 1),
+        },
+    )
+
+    tuner = tune.Tuner(
+        tune.with_resources(
+            Cifar10Model,
+            resources={"cpu": 1, "gpu": 1},
+        ),
+        run_config=tune.RunConfig(
+            name="pbt_cifar10",
+            stop={
+                "mean_accuracy": 0.80,
+                "training_iteration": 30,
+            },
+            checkpoint_config=tune.CheckpointConfig(
+                checkpoint_frequency=perturbation_interval,
+                checkpoint_score_attribute="mean_accuracy",
+                num_to_keep=2,
+            ),
+        ),
+        tune_config=tune.TuneConfig(
+            scheduler=pbt,
+            num_samples=4,
+            metric="mean_accuracy",
+            mode="max",
+            reuse_actors=True,
+        ),
+        param_space=space,
+    )
+    results = tuner.fit()
+    print("Best hyperparameters found were: ", results.get_best_result().config)

python/ray/tune/examples/tf_mnist_example.py

@@ -0,0 +1,152 @@
+#!/usr/bin/env python
+# coding: utf-8
+#
+# This example showcases how to use TF2.0 APIs with Tune.
+# Original code: https://www.tensorflow.org/tutorials/quickstart/advanced
+#
+# As of 10/12/2019: One caveat of using TF2.0 is that TF AutoGraph
+# functionality does not interact nicely with Ray actors. One way to get around
+# this is to `import tensorflow` inside the Tune Trainable.
+#
+
+import argparse
+import os
+import sys
+
+from filelock import FileLock
+
+from ray import tune
+
+MAX_TRAIN_BATCH = 10
+
+if sys.version_info >= (3, 12):
+    # Tensorflow is not installed for Python 3.12 because of keras compatibility.
+    sys.exit(0)
+else:
+    from tensorflow.keras import Model
+    from tensorflow.keras.datasets.mnist import load_data
+    from tensorflow.keras.layers import Conv2D, Dense, Flatten
+
+
+class MyModel(Model):
+    def __init__(self, hiddens=128):
+        super(MyModel, self).__init__()
+        self.conv1 = Conv2D(32, 3, activation="relu")
+        self.flatten = Flatten()
+        self.d1 = Dense(hiddens, activation="relu")
+        self.d2 = Dense(10, activation="softmax")
+
+    def call(self, x):
+        x = self.conv1(x)
+        x = self.flatten(x)
+        x = self.d1(x)
+        return self.d2(x)
+
+
+class MNISTTrainable(tune.Trainable):
+    def setup(self, config):
+        # IMPORTANT: See the above note.
+        import tensorflow as tf
+
+        # Use FileLock to avoid race conditions.
+        with FileLock(os.path.expanduser("~/.tune.lock")):
+            (x_train, y_train), (x_test, y_test) = load_data()
+        x_train, x_test = x_train / 255.0, x_test / 255.0
+
+        # Add a channels dimension
+        x_train = x_train[..., tf.newaxis]
+        x_test = x_test[..., tf.newaxis]
+        self.train_ds = tf.data.Dataset.from_tensor_slices((x_train, y_train))
+        self.train_ds = self.train_ds.shuffle(10000).batch(config.get("batch", 32))
+
+        self.test_ds = tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(32)
+
+        self.model = MyModel(hiddens=config.get("hiddens", 128))
+        self.loss_object = tf.keras.losses.SparseCategoricalCrossentropy()
+        self.optimizer = tf.keras.optimizers.Adam()
+        self.train_loss = tf.keras.metrics.Mean(name="train_loss")
+        self.train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(
+            name="train_accuracy"
+        )
+
+        self.test_loss = tf.keras.metrics.Mean(name="test_loss")
+        self.test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(
+            name="test_accuracy"
+        )
+
+        @tf.function
+        def train_step(images, labels):
+            with tf.GradientTape() as tape:
+                predictions = self.model(images)
+                loss = self.loss_object(labels, predictions)
+            gradients = tape.gradient(loss, self.model.trainable_variables)
+            self.optimizer.apply_gradients(
+                zip(gradients, self.model.trainable_variables)
+            )
+
+            self.train_loss(loss)
+            self.train_accuracy(labels, predictions)
+
+        @tf.function
+        def test_step(images, labels):
+            predictions = self.model(images)
+            t_loss = self.loss_object(labels, predictions)
+
+            self.test_loss(t_loss)
+            self.test_accuracy(labels, predictions)
+
+        self.tf_train_step = train_step
+        self.tf_test_step = test_step
+
+    def save_checkpoint(self, checkpoint_dir: str):
+        return None
+
+    def load_checkpoint(self, checkpoint):
+        return None
+
+    def step(self):
+        self.train_loss.reset_states()
+        self.train_accuracy.reset_states()
+        self.test_loss.reset_states()
+        self.test_accuracy.reset_states()
+
+        for idx, (images, labels) in enumerate(self.train_ds):
+            if idx > MAX_TRAIN_BATCH:  # This is optional and can be removed.
+                break
+            self.tf_train_step(images, labels)
+
+        for test_images, test_labels in self.test_ds:
+            self.tf_test_step(test_images, test_labels)
+
+        # It is important to return tf.Tensors as numpy objects.
+        return {
+            "epoch": self.iteration,
+            "loss": self.train_loss.result().numpy(),
+            "accuracy": self.train_accuracy.result().numpy() * 100,
+            "test_loss": self.test_loss.result().numpy(),
+            "mean_accuracy": self.test_accuracy.result().numpy() * 100,
+        }
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--smoke-test", action="store_true", help="Finish quickly for testing"
+    )
+    args, _ = parser.parse_known_args()
+
+    tuner = tune.Tuner(
+        MNISTTrainable,
+        tune_config=tune.TuneConfig(
+            metric="test_loss",
+            mode="min",
+        ),
+        run_config=tune.RunConfig(
+            stop={"training_iteration": 5 if args.smoke_test else 50},
+            verbose=1,
+        ),
+        param_space={"hiddens": tune.grid_search([32, 64, 128])},
+    )
+    results = tuner.fit()
+
+    print("Best hyperparameters found were: ", results.get_best_result().config)

python/ray/tune/examples/tune-default.yaml

@@ -0,0 +1,14 @@
+cluster_name: tune-default
+provider: {type: aws, region: us-west-2}
+auth: {ssh_user: ubuntu}
+min_workers: 3
+max_workers: 3
+# Deep Learning AMI (Ubuntu) Version 21.0
+available_node_types:
+  head_node:
+    node_config: {InstanceType: c5.xlarge, ImageId: ami-0b294f219d14e6a82}
+  worker_nodes:
+    node_config: {InstanceType: c5.xlarge, ImageId: ami-0b294f219d14e6a82}
+head_node_type: head_node
+setup_commands: # Set up each node.
+    - pip install ray torch torchvision tensorboard

python/ray/tune/examples/tune-local-default.yaml

@@ -0,0 +1,11 @@
+cluster_name: local-default
+provider:
+    type: local
+    head_ip: YOUR_HEAD_NODE_HOSTNAME
+    worker_ips: [WORKER_NODE_1_HOSTNAME, WORKER_NODE_2_HOSTNAME, ... ]
+auth: {ssh_user: YOUR_USERNAME, ssh_private_key: ~/.ssh/id_rsa}
+## Typically for local clusters, min_workers == max_workers.
+min_workers: 3
+max_workers: 3
+setup_commands:  # Set up each node.
+    - pip install ray torch torchvision tensorboard

python/ray/tune/examples/tune_basic_example.py

@@ -0,0 +1,57 @@
+"""This example demonstrates basic Ray Tune random search and grid search."""
+import time
+
+import ray
+from ray import tune
+
+
+def evaluation_fn(step, width, height):
+    time.sleep(0.1)
+    return (0.1 + width * step / 100) ** (-1) + height * 0.1
+
+
+def easy_objective(config):
+    # Hyperparameters
+    width, height = config["width"], config["height"]
+
+    for step in range(config["steps"]):
+        # Iterative training function - can be any arbitrary training procedure
+        intermediate_score = evaluation_fn(step, width, height)
+        # Feed the score back back to Tune.
+        tune.report({"iterations": step, "mean_loss": intermediate_score})
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--smoke-test", action="store_true", help="Finish quickly for testing"
+    )
+    args, _ = parser.parse_known_args()
+
+    ray.init(configure_logging=False)
+
+    # This will do a grid search over the `activation` parameter. This means
+    # that each of the two values (`relu` and `tanh`) will be sampled once
+    # for each sample (`num_samples`). We end up with 2 * 50 = 100 samples.
+    # The `width` and `height` parameters are sampled randomly.
+    # `steps` is a constant parameter.
+
+    tuner = tune.Tuner(
+        easy_objective,
+        tune_config=tune.TuneConfig(
+            metric="mean_loss",
+            mode="min",
+            num_samples=5 if args.smoke_test else 50,
+        ),
+        param_space={
+            "steps": 5 if args.smoke_test else 100,
+            "width": tune.uniform(0, 20),
+            "height": tune.uniform(-100, 100),
+            "activation": tune.grid_search(["relu", "tanh"]),
+        },
+    )
+    results = tuner.fit()
+
+    print("Best hyperparameters found were: ", results.get_best_result().config)

python/ray/tune/examples/tune_mnist_keras.py

@@ -0,0 +1,99 @@
+import argparse
+import os
+import sys
+
+from filelock import FileLock
+
+import ray
+from ray import tune
+from ray.tune.schedulers import AsyncHyperBandScheduler
+
+if sys.version_info >= (3, 12):
+    # Tensorflow is not installed for Python 3.12 because of keras compatibility.
+    sys.exit(0)
+else:
+    from tensorflow.keras.datasets import mnist
+
+    from ray.tune.integration.keras import TuneReportCheckpointCallback
+
+
+def train_mnist(config):
+    # https://github.com/tensorflow/tensorflow/issues/32159
+    import tensorflow as tf
+
+    batch_size = 128
+    num_classes = 10
+    epochs = 12
+
+    with FileLock(os.path.expanduser("~/.data.lock")):
+        (x_train, y_train), (x_test, y_test) = mnist.load_data()
+    x_train, x_test = x_train / 255.0, x_test / 255.0
+    model = tf.keras.models.Sequential(
+        [
+            tf.keras.layers.Flatten(input_shape=(28, 28)),
+            tf.keras.layers.Dense(config["hidden"], activation="relu"),
+            tf.keras.layers.Dropout(0.2),
+            tf.keras.layers.Dense(num_classes, activation="softmax"),
+        ]
+    )
+
+    model.compile(
+        loss="sparse_categorical_crossentropy",
+        optimizer=tf.keras.optimizers.SGD(lr=config["lr"], momentum=config["momentum"]),
+        metrics=["accuracy"],
+    )
+
+    model.fit(
+        x_train,
+        y_train,
+        batch_size=batch_size,
+        epochs=epochs,
+        verbose=0,
+        validation_data=(x_test, y_test),
+        callbacks=[
+            TuneReportCheckpointCallback(
+                checkpoint_on=[], metrics={"mean_accuracy": "accuracy"}
+            )
+        ],
+    )
+
+
+def tune_mnist(num_training_iterations):
+    sched = AsyncHyperBandScheduler(
+        time_attr="training_iteration", max_t=400, grace_period=20
+    )
+
+    tuner = tune.Tuner(
+        tune.with_resources(train_mnist, resources={"cpu": 2, "gpu": 0}),
+        run_config=tune.RunConfig(
+            name="exp",
+            stop={"mean_accuracy": 0.99, "training_iteration": num_training_iterations},
+        ),
+        tune_config=tune.TuneConfig(
+            scheduler=sched,
+            metric="mean_accuracy",
+            mode="max",
+            num_samples=10,
+        ),
+        param_space={
+            "threads": 2,
+            "lr": tune.uniform(0.001, 0.1),
+            "momentum": tune.uniform(0.1, 0.9),
+            "hidden": tune.randint(32, 512),
+        },
+    )
+    results = tuner.fit()
+    print("Best hyperparameters found were: ", results.get_best_result().config)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--smoke-test", action="store_true", help="Finish quickly for testing"
+    )
+    args, _ = parser.parse_known_args()
+
+    if args.smoke_test:
+        ray.init(num_cpus=4)
+
+    tune_mnist(num_training_iterations=2 if args.smoke_test else 300)

python/ray/tune/examples/utils.py

@@ -0,0 +1,21 @@
+import tensorflow as tf
+from sklearn.datasets import load_iris
+from sklearn.model_selection import train_test_split
+from sklearn.preprocessing import OneHotEncoder
+
+
+def get_iris_data(test_size=0.2):
+    iris_data = load_iris()
+    x = iris_data.data
+    y = iris_data.target.reshape(-1, 1)
+    encoder = OneHotEncoder(sparse=False)
+    y = encoder.fit_transform(y)
+    train_x, test_x, train_y, test_y = train_test_split(x, y)
+    return train_x, train_y, test_x, test_y
+
+
+def set_keras_threads(threads):
+    # We set threads here to avoid contention, as Keras
+    # is heavily parallelized across multiple cores.
+    tf.config.threading.set_inter_op_parallelism_threads(threads)
+    tf.config.threading.set_intra_op_parallelism_threads(threads)

python/ray/tune/examples/xgboost_dynamic_resources_example.py

@@ -0,0 +1,187 @@
+from typing import TYPE_CHECKING, Any, Dict, Optional
+
+import sklearn.datasets
+import sklearn.metrics
+import xgboost as xgb
+from sklearn.model_selection import train_test_split
+
+import ray
+from ray import tune
+from ray.tune.execution.placement_groups import PlacementGroupFactory
+from ray.tune.experiment import Trial
+from ray.tune.integration.xgboost import TuneReportCheckpointCallback
+from ray.tune.schedulers import ASHAScheduler, ResourceChangingScheduler
+
+if TYPE_CHECKING:
+    from ray.tune.execution.tune_controller import TuneController
+
+CHECKPOINT_FILENAME = "booster-checkpoint.json"
+
+
+def get_best_model_checkpoint(best_result: "ray.tune.Result"):
+    best_bst = TuneReportCheckpointCallback.get_model(
+        best_result.checkpoint, filename=CHECKPOINT_FILENAME
+    )
+
+    accuracy = 1.0 - best_result.metrics["eval-logloss"]
+    print(f"Best model parameters: {best_result.config}")
+    print(f"Best model total accuracy: {accuracy:.4f}")
+    return best_bst
+
+
+# our train function needs to be able to checkpoint
+# to work with ResourceChangingScheduler
+def train_breast_cancer(config: dict):
+    # This is a simple training function to be passed into Tune
+    # Load dataset
+    data, labels = sklearn.datasets.load_breast_cancer(return_X_y=True)
+    # Split into train and test set
+    train_x, test_x, train_y, test_y = train_test_split(data, labels, test_size=0.25)
+    # Build input matrices for XGBoost
+    train_set = xgb.DMatrix(train_x, label=train_y)
+    test_set = xgb.DMatrix(test_x, label=test_y)
+
+    # Checkpointing needs to be set up in order for dynamic
+    # resource allocation to work as intended
+    xgb_model = None
+    checkpoint = tune.get_checkpoint()
+    if checkpoint:
+        xgb_model = TuneReportCheckpointCallback.get_model(
+            checkpoint, filename=CHECKPOINT_FILENAME
+        )
+
+    # Set `nthread` to the number of CPUs available to the trial,
+    # which is assigned by the scheduler.
+    config["nthread"] = int(tune.get_context().get_trial_resources().head_cpus)
+    print(f"nthreads: {config['nthread']} xgb_model: {xgb_model}")
+    # Train the classifier, using the Tune callback
+    xgb.train(
+        config,
+        train_set,
+        evals=[(test_set, "eval")],
+        verbose_eval=False,
+        xgb_model=xgb_model,
+        callbacks=[
+            TuneReportCheckpointCallback(
+                # checkpointing should happen every iteration
+                # with dynamic resource allocation
+                frequency=1,
+                filename=CHECKPOINT_FILENAME,
+            )
+        ],
+    )
+
+
+def tune_xgboost():
+    search_space = {
+        # You can mix constants with search space objects.
+        "objective": "binary:logistic",
+        "eval_metric": ["logloss", "error"],
+        "max_depth": 9,
+        "learning_rate": 1,
+        "min_child_weight": tune.grid_search([2, 3]),
+        "subsample": tune.grid_search([0.8, 0.9]),
+        "colsample_bynode": tune.grid_search([0.8, 0.9]),
+        "random_state": 1,
+        "num_parallel_tree": 2000,
+    }
+    # This will enable aggressive early stopping of bad trials.
+    base_scheduler = ASHAScheduler(
+        max_t=16, grace_period=1, reduction_factor=2  # 16 training iterations
+    )
+
+    def example_resources_allocation_function(
+        tune_controller: "TuneController",
+        trial: Trial,
+        result: Dict[str, Any],
+        scheduler: "ResourceChangingScheduler",
+    ) -> Optional[PlacementGroupFactory]:
+        """This is a basic example of a resource allocating function.
+
+        The function naively balances available CPUs over live trials.
+
+        This function returns a new ``PlacementGroupFactory`` with updated
+        resource requirements, or None. If the returned
+        ``PlacementGroupFactory`` is equal by value to the one the
+        trial has currently, the scheduler will skip the update process
+        internally (same with None).
+
+        See :class:`DistributeResources` for a more complex,
+        robust approach.
+
+        Args:
+            tune_controller: Trial runner for this Tune run.
+                Can be used to obtain information about other trials.
+            trial: The trial to allocate new resources to.
+            result: The latest results of trial.
+            scheduler: The scheduler calling the function.
+
+        Returns:
+            A new ``PlacementGroupFactory`` with the updated resource
+            requirements, or ``None`` to leave the trial's resources unchanged.
+        """
+
+        # Get base trial resources as defined in
+        # ``tune.with_resources``
+        base_trial_resource = scheduler._base_trial_resources
+
+        # Don't bother if this is just the first iteration
+        if result["training_iteration"] < 1:
+            return None
+
+        # default values if resources_per_trial is unspecified
+        if base_trial_resource is None:
+            base_trial_resource = PlacementGroupFactory([{"CPU": 1, "GPU": 0}])
+
+        # Assume that the number of CPUs cannot go below what was
+        # specified in ``Tuner.fit()``.
+        min_cpu = base_trial_resource.required_resources.get("CPU", 0)
+
+        # Get the number of CPUs available in total (not just free)
+        total_available_cpus = tune_controller._resource_updater.get_num_cpus()
+
+        # Divide the free CPUs among all live trials
+        cpu_to_use = max(
+            min_cpu, total_available_cpus // len(tune_controller.get_live_trials())
+        )
+
+        # Assign new CPUs to the trial in a PlacementGroupFactory
+        return PlacementGroupFactory([{"CPU": cpu_to_use, "GPU": 0}])
+
+    # You can either define your own resources_allocation_function, or
+    # use the default one - DistributeResources
+
+    # from ray.tune.schedulers.resource_changing_scheduler import \
+    #    DistributeResources
+
+    scheduler = ResourceChangingScheduler(
+        base_scheduler=base_scheduler,
+        resources_allocation_function=example_resources_allocation_function,
+        # resources_allocation_function=DistributeResources()  # default
+    )
+
+    tuner = tune.Tuner(
+        tune.with_resources(
+            train_breast_cancer, resources=PlacementGroupFactory([{"CPU": 1, "GPU": 0}])
+        ),
+        tune_config=tune.TuneConfig(
+            metric="eval-logloss",
+            mode="min",
+            num_samples=1,
+            scheduler=scheduler,
+        ),
+        param_space=search_space,
+    )
+    results = tuner.fit()
+
+    return results.get_best_result()
+
+
+if __name__ == "__main__":
+    ray.init(num_cpus=8)
+
+    best_result = tune_xgboost()
+    best_bst = get_best_model_checkpoint(best_result)
+
+    # You could now do further predictions with
+    # best_bst.predict(...)

python/ray/tune/examples/xgboost_example.py

@@ -0,0 +1,130 @@
+from typing import Dict, List
+
+import numpy as np
+import sklearn.datasets
+import sklearn.metrics
+import xgboost as xgb
+from sklearn.model_selection import train_test_split
+
+import ray
+from ray import tune
+from ray.tune.integration.xgboost import TuneReportCheckpointCallback
+from ray.tune.schedulers import ASHAScheduler
+
+CHECKPOINT_FILENAME = "booster-checkpoint.json"
+
+
+def train_breast_cancer(config: dict):
+    # This is a simple training function to be passed into Tune
+
+    # Load dataset
+    data, labels = sklearn.datasets.load_breast_cancer(return_X_y=True)
+
+    # Split into train and test set
+    train_x, test_x, train_y, test_y = train_test_split(data, labels, test_size=0.25)
+    # Build input matrices for XGBoost
+    train_set = xgb.DMatrix(train_x, label=train_y)
+    test_set = xgb.DMatrix(test_x, label=test_y)
+
+    # Train the classifier, using the Tune callback
+    xgb.train(
+        config,
+        train_set,
+        evals=[(test_set, "test")],
+        verbose_eval=False,
+        callbacks=[
+            TuneReportCheckpointCallback(frequency=1, filename=CHECKPOINT_FILENAME)
+        ],
+    )
+
+
+def train_breast_cancer_cv(config: dict):
+    # This is a simple training function to be passed into Tune
+    # using xgboost's cross validation functionality
+
+    # Load dataset
+    data, labels = sklearn.datasets.load_breast_cancer(return_X_y=True)
+
+    # For CV, we need to average over a list of results form folds
+    def average_cv_folds(results_dict: Dict[str, List[float]]) -> Dict[str, float]:
+        return {k: np.mean(v) for k, v in results_dict.items()}
+
+    train_set = xgb.DMatrix(data, label=labels)
+
+    # Run CV, using the Tune callback
+    xgb.cv(
+        config,
+        train_set,
+        verbose_eval=False,
+        stratified=True,
+        # Checkpointing is not supported for CV
+        callbacks=[
+            TuneReportCheckpointCallback(
+                results_postprocessing_fn=average_cv_folds, frequency=0
+            )
+        ],
+    )
+
+
+def get_best_model_checkpoint(best_result: "ray.tune.Result"):
+    best_bst = TuneReportCheckpointCallback.get_model(
+        best_result.checkpoint, filename=CHECKPOINT_FILENAME
+    )
+    accuracy = 1.0 - best_result.metrics["test-error"]
+    print(f"Best model parameters: {best_result.config}")
+    print(f"Best model total accuracy: {accuracy:.4f}")
+    return best_bst
+
+
+def tune_xgboost(use_cv: bool = False):
+    search_space = {
+        # You can mix constants with search space objects.
+        "objective": "binary:logistic",
+        "eval_metric": ["logloss", "error"],
+        "max_depth": tune.randint(1, 9),
+        "min_child_weight": tune.choice([1, 2, 3]),
+        "subsample": tune.uniform(0.5, 1.0),
+        "eta": tune.loguniform(1e-4, 1e-1),
+    }
+    # This will enable aggressive early stopping of bad trials.
+    scheduler = ASHAScheduler(
+        max_t=10, grace_period=1, reduction_factor=2  # 10 training iterations
+    )
+
+    tuner = tune.Tuner(
+        tune.with_resources(
+            train_breast_cancer if not use_cv else train_breast_cancer_cv,
+            # You can add "gpu": 0.1 to allocate GPUs
+            resources={"cpu": 1},
+        ),
+        tune_config=tune.TuneConfig(
+            metric="test-logloss",
+            mode="min",
+            num_samples=10,
+            scheduler=scheduler,
+        ),
+        param_space=search_space,
+    )
+    results = tuner.fit()
+
+    return results.get_best_result()
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--use-cv", action="store_true", help="Use `xgb.cv` instead of `xgb.train`."
+    )
+    args, _ = parser.parse_known_args()
+
+    best_result = tune_xgboost(args.use_cv)
+
+    # Load the best model checkpoint.
+    # Checkpointing is not supported when using `xgb.cv`
+    if not args.use_cv:
+        best_bst = get_best_model_checkpoint(best_result)
+
+        # You could now do further predictions with
+        # best_bst.predict(...)

python/ray/tune/execution/class_cache.py

@@ -0,0 +1,68 @@
+import os
+
+import ray
+from ray.air.constants import COPY_DIRECTORY_CHECKPOINTS_INSTEAD_OF_MOVING_ENV
+from ray.train.constants import (
+    ENABLE_V2_MIGRATION_WARNINGS_ENV_VAR,
+    RAY_CHDIR_TO_TRIAL_DIR,
+)
+from ray.train.v2._internal.constants import (
+    ENV_VARS_TO_PROPAGATE as TRAIN_ENV_VARS_TO_PROPAGATE,
+)
+
+DEFAULT_ENV_VARS = {
+    # https://github.com/ray-project/ray/issues/28197
+    "PL_DISABLE_FORK": "1"
+}
+ENV_VARS_TO_PROPAGATE = (
+    {
+        COPY_DIRECTORY_CHECKPOINTS_INSTEAD_OF_MOVING_ENV,
+        RAY_CHDIR_TO_TRIAL_DIR,
+        ENABLE_V2_MIGRATION_WARNINGS_ENV_VAR,
+        "AWS_ACCESS_KEY_ID",
+        "AWS_SECRET_ACCESS_KEY",
+        "AWS_SECURITY_TOKEN",
+        "AWS_SESSION_TOKEN",
+    }
+    # Propagate the Ray Train environment variables from the driver process
+    # to the trainable process so that Tune + Train v2 can be used together.
+    | TRAIN_ENV_VARS_TO_PROPAGATE
+)
+
+
+class _ActorClassCache:
+    """Caches actor classes.
+
+    ray.remote is a registration call. It sends the serialized object to the
+    key value store (redis), and will be fetched at an arbitrary worker
+    later. Registration does not use any Ray scheduling resources.
+
+    Later, class.remote() actually creates the remote actor. The
+    actor will be instantiated on some arbitrary machine,
+    according to the underlying Ray scheduler.
+
+    Without this cache, you would register the same serialized object
+    over and over again. Naturally, since redis doesn’t spill to disk,
+    this can easily nuke the redis instance (and basically blow up Ray).
+    This cache instead allows us to register once and only once.
+
+    Note that we assume there can be multiple trainables in the
+    system at once.
+    """
+
+    def __init__(self):
+        self._cache = {}
+
+    def get(self, trainable_cls):
+        """Gets the wrapped trainable_cls, otherwise calls ray.remote."""
+        env_vars = DEFAULT_ENV_VARS.copy()
+
+        for env_var_to_propagate in ENV_VARS_TO_PROPAGATE:
+            if env_var_to_propagate in os.environ:
+                env_vars[env_var_to_propagate] = os.environ[env_var_to_propagate]
+
+        runtime_env = {"env_vars": env_vars}
+        if trainable_cls not in self._cache:
+            remote_cls = ray.remote(runtime_env=runtime_env)(trainable_cls)
+            self._cache[trainable_cls] = remote_cls
+        return self._cache[trainable_cls]

python/ray/tune/execution/cluster_info.py

@@ -0,0 +1,12 @@
+from functools import lru_cache
+from pathlib import Path
+
+
+@lru_cache()
+def _is_ray_cluster():
+    """Checks if the bootstrap config file exists.
+
+    This will always exist if using an autoscaling cluster/started
+    with the ray cluster launcher.
+    """
+    return Path("~/ray_bootstrap_config.yaml").expanduser().exists()

python/ray/tune/execution/experiment_state.py

@@ -0,0 +1,290 @@
+import fnmatch
+import logging
+import os
+import time
+from collections import Counter
+from pathlib import Path
+from typing import Callable, Dict, Optional, Union
+
+import pyarrow.fs
+
+from ray.train._internal.storage import (
+    StorageContext,
+    _download_from_fs_path,
+    _list_at_fs_path,
+    get_fs_and_path,
+)
+from ray.tune.experiment.trial import Trial
+from ray.tune.impl.out_of_band_serialize_dataset import out_of_band_serialize_dataset
+
+logger = logging.getLogger(__name__)
+
+
+_SLOW_SYNC_WARNING = (
+    "This could be due to a large number of trials, "
+    "large logfiles from lots of reported metrics, or throttling from the "
+    "remote storage if uploading too frequently.\n"
+    "You may want to consider switching the `RunConfig(storage_filesystem)`"
+    " to a more performant storage backend such as s3fs for a "
+    "S3 storage path.\n"
+    "You can suppress this error by setting the environment variable "
+    "TUNE_WARN_SLOW_EXPERIMENT_CHECKPOINT_SYNC_THRESHOLD_S to a higher "
+    "value than the current threshold ({threshold})."
+)
+
+
+def _find_newest_experiment_checkpoint(
+    experiment_path: str, fs: Optional[pyarrow.fs.FileSystem] = None
+) -> Optional[str]:
+    """Returns file name of most recently created experiment checkpoint.
+
+    Args:
+        experiment_path: Local or remote path to the experiment directory
+            containing at least one experiment checkpoint file.
+        fs: Optional custom ``pyarrow.fs.FileSystem`` corresponding to
+            ``experiment_path``. If not provided, one is inferred from the
+            path.
+
+    Returns:
+        str: The local or remote path to the latest experiment checkpoint file
+            based on timestamp. None if no experiment checkpoints were found.
+    """
+    from ray.tune.execution.tune_controller import TuneController
+
+    fs, experiment_fs_path = get_fs_and_path(experiment_path, storage_filesystem=fs)
+    filenames = _list_at_fs_path(fs=fs, fs_path=experiment_fs_path)
+    pattern = TuneController.CKPT_FILE_TMPL.format("*")
+    matching = fnmatch.filter(filenames, pattern)
+    if not matching:
+        return None
+    filename = max(matching)
+    return Path(experiment_fs_path, filename).as_posix()
+
+
+class _ExperimentCheckpointManager:
+    """Helper class for managing experiment-level checkpoints.
+
+    This class implements the ``checkpoint()`` method used to checkpoint
+    experiment state. When called, this will serialize and write to disk
+    the state of the trial runner, trial executor, and search algorithm, to
+    a specified checkpoint file.
+
+    The checkpoint period is automatically adjusted to
+    ``max(10, time_per_checkpoint * 19)``. This means that at most 5% of the
+    time (1/20) will be used for writing checkpoints, while 95% of the time
+    (19/20) will be used to handle the rest of the training loop.
+    """
+
+    def __init__(
+        self,
+        *,
+        storage: Optional[StorageContext],
+        checkpoint_period: Union[int, float, str],
+        sync_every_n_trial_checkpoints: Optional[int] = None,
+    ):
+        self._storage = storage
+
+        self._last_save_time = float("-inf")
+        self._last_sync_time = None
+
+        # Dynamic checkpointing period
+        self._auto_checkpoint_enabled = checkpoint_period == "auto"
+        if self._auto_checkpoint_enabled:
+            self._checkpoint_period = 10.0  # Initial value
+        else:
+            self._checkpoint_period = float(checkpoint_period)
+
+        # TODO(justinvyu): This is a non-performant workaround to force sync
+        # every num_to_keep checkpoints in order to maintain consistency
+        # between the experiment state's view of the latest checkpoint,
+        # and the actual latest checkpoint that was uploaded.
+        self._sync_every_n_trial_checkpoints = sync_every_n_trial_checkpoints
+        self._trial_num_checkpoints_since_last_sync: Dict[Trial, int] = Counter()
+        self._should_force_sync_up: bool = False
+
+        self._excessive_sync_threshold = float(
+            os.environ.get(
+                "TUNE_WARN_EXCESSIVE_EXPERIMENT_CHECKPOINT_SYNC_THRESHOLD_S", "5"
+            )
+        )
+        self._slow_sync_threshold = float(
+            os.environ.get(
+                "TUNE_WARN_SLOW_EXPERIMENT_CHECKPOINT_SYNC_THRESHOLD_S", "30"
+            )
+        )
+
+    @property
+    def auto_checkpoint_enabled(self):
+        return self._auto_checkpoint_enabled
+
+    def _update_auto_checkpoint_time(self, time_taken: float):
+        if self._auto_checkpoint_enabled:
+            # Multiplying this time by 19 means we spend ~5% of the time
+            # writing global checkpoints and 95% of the time processing trials
+            self._checkpoint_period = max(10.0, time_taken * 19)
+            logger.debug(
+                f"Experiment state snapshotting took "
+                f"{time_taken:.2f} seconds. "
+                f"Adjusting snapshotting period to "
+                f"{self._checkpoint_period:.2f} seconds."
+            )
+
+    def sync_up_experiment_state(
+        self,
+        save_fn: Callable[[], None],
+        force: bool = False,
+        wait: bool = False,
+    ) -> None:
+        """Saves execution state to the experiment directory on the storage path.
+        This includes an experiment checkpoint file that contains trial statuses
+        and the searcher state.
+
+        Overwrites the current session checkpoint, which starts when self
+        is instantiated. Throttle depends on self._checkpoint_period.
+
+        Args:
+            save_fn: Function to call to actually save data to the driver
+                staging path. The files in the driver staging path will be
+                uploaded to the storage path.
+            force: Forces an experiment checkpoint and launches a sync to storage.
+                This happens regardless of checkpoint_period
+            wait: Waits for the sync up to complete before returning.
+        """
+        driver_staging_path = self._storage.experiment_driver_staging_path
+
+        force = force or self._should_force_sync_up
+
+        now = time.monotonic()
+        if now - self._last_save_time < self._checkpoint_period and not force:
+            return
+
+        # Checkpoint
+        checkpoint_time_start = time.monotonic()
+
+        # NOTE: This context manager is for Datasets captured in a trial config.
+        # This is the case when *tuning over datasets*.
+        # If the datasets have already been full executed, then serializing
+        # block refs means that this checkpoint is not usable in a new Ray cluster.
+        # This context will serialize the dataset execution plan instead, if available.
+        with out_of_band_serialize_dataset():
+            save_fn()
+
+        def wait_for_sync():
+            try:
+                self._storage.syncer.wait()
+            except Exception:
+                logger.error(
+                    "Saving experiment state to storage at "
+                    f"'{self._storage.experiment_fs_path}' failed with exception: ",
+                    exc_info=True,
+                )
+
+        if force:
+            start_time = time.monotonic()
+            wait_for_sync()
+            wait_time = time.monotonic() - start_time
+            if wait_time > self._slow_sync_threshold:
+                logger.warning(
+                    "Saving the experiment state (which holds a global view "
+                    "of trial statuses and is used to restore the experiment) "
+                    f"took ~{wait_time:.2f} seconds, which may be a performance "
+                    "bottleneck.\n"
+                    f"{_SLOW_SYNC_WARNING.format(threshold=self._slow_sync_threshold)}"
+                )
+
+        time_since_last_sync = (
+            time.monotonic() - self._last_sync_time
+            if self._last_sync_time is not None
+            else None
+        )
+        launched_sync = self._storage.syncer.sync_up(
+            driver_staging_path, self._storage.experiment_fs_path
+        )
+        if launched_sync:
+            if (
+                time_since_last_sync is not None
+                and time_since_last_sync < self._excessive_sync_threshold
+                and self._should_force_sync_up
+            ):
+                logger.warning(
+                    "Experiment state snapshotting has been triggered multiple "
+                    f"times in the last {self._excessive_sync_threshold} seconds "
+                    "and may become a bottleneck. "
+                    "A snapshot is forced if `CheckpointConfig(num_to_keep)` is set, "
+                    "and a trial has checkpointed >= `num_to_keep` times "
+                    "since the last snapshot.\n"
+                    "You may want to consider increasing the "
+                    "`CheckpointConfig(num_to_keep)` or decreasing the frequency of "
+                    "saving checkpoints.\n"
+                    "You can suppress this warning by setting the environment variable "
+                    "TUNE_WARN_EXCESSIVE_EXPERIMENT_CHECKPOINT_SYNC_THRESHOLD_S "
+                    "to a smaller value than the current threshold "
+                    f"({self._excessive_sync_threshold}). "
+                    "Set it to 0 to completely suppress this warning."
+                )
+
+            self._last_sync_time = time.monotonic()
+
+            # We just synced, so reset the force flag
+            self._trial_num_checkpoints_since_last_sync.clear()
+            self._should_force_sync_up = False
+        else:
+            if (
+                time_since_last_sync is not None
+                and time_since_last_sync > self._slow_sync_threshold
+            ):
+                logger.warning(
+                    "Saving the experiment state (which holds a global view "
+                    "of trial statuses and is used to restore the experiment) "
+                    f"has already taken {time_since_last_sync:.2f} seconds, "
+                    "which may cause consistency issues upon restoration if your "
+                    "driver script ungracefully exits.\n"
+                    f"{_SLOW_SYNC_WARNING.format(threshold=self._slow_sync_threshold)}"
+                )
+
+        if wait:
+            wait_for_sync()
+
+        checkpoint_time_taken = time.monotonic() - checkpoint_time_start
+
+        # Adjust dynamic checkpointing
+        self._update_auto_checkpoint_time(time_taken=checkpoint_time_taken)
+
+        # Finish
+        self._last_save_time = time.monotonic()
+
+    def sync_down_experiment_state(self) -> None:
+        fs = self._storage.storage_filesystem
+        filepaths = _list_at_fs_path(fs=fs, fs_path=self._storage.experiment_fs_path)
+        # TODO(ekl) we should refactor our restore code to read the necessary data
+        # directly from the storage context. As a temporary hack, restore all the
+        # serialized files from the root dir where other modules expect them to be.
+        matches = [
+            path
+            for path in filepaths
+            if path.endswith(".json") or path.endswith(".pkl")
+        ]
+        for relpath in matches:
+            fs_path = Path(self._storage.experiment_fs_path, relpath).as_posix()
+            local_path = Path(
+                self._storage.experiment_driver_staging_path, relpath
+            ).as_posix()
+            _download_from_fs_path(fs=fs, fs_path=fs_path, local_path=local_path)
+        logger.debug(
+            f"Copied {matches} from:\n(fs, path) = "
+            f"({self._storage.storage_filesystem.type_name}, "
+            f"{self._storage.experiment_fs_path})\n"
+            f"-> {self._storage.experiment_driver_staging_path}"
+        )
+
+    def on_trial_checkpoint(self, trial: Trial):
+        if not self._sync_every_n_trial_checkpoints:
+            return
+
+        self._trial_num_checkpoints_since_last_sync[trial] += 1
+
+        if (
+            self._trial_num_checkpoints_since_last_sync[trial]
+            >= self._sync_every_n_trial_checkpoints
+        ):
+            self._should_force_sync_up = True

python/ray/tune/execution/insufficient_resources_manager.py

@@ -0,0 +1,167 @@
+import logging
+import os
+import time
+from functools import lru_cache
+from typing import Dict, Optional, Tuple
+
+import ray
+from ray.tune.execution.cluster_info import _is_ray_cluster
+from ray.tune.experiment import Trial
+
+logger = logging.getLogger(__name__)
+
+
+# Ideally we want to use @cache; but it's only available for python 3.9.
+# Caching is only helpful/correct for no autoscaler case.
+@lru_cache()
+def _get_cluster_resources_no_autoscaler() -> Dict:
+    return ray.cluster_resources()
+
+
+def _get_trial_cpu_and_gpu(trial: Trial) -> Tuple[int, int]:
+    cpu = trial.placement_group_factory.required_resources.get("CPU", 0)
+    gpu = trial.placement_group_factory.required_resources.get("GPU", 0)
+    return cpu, gpu
+
+
+def _can_fulfill_no_autoscaler(trial: Trial) -> bool:
+    """Calculates if there is enough resources for a PENDING trial.
+
+    For no autoscaler case.
+    """
+    assert trial.status == Trial.PENDING
+    asked_cpus, asked_gpus = _get_trial_cpu_and_gpu(trial)
+
+    return asked_cpus <= _get_cluster_resources_no_autoscaler().get(
+        "CPU", 0
+    ) and asked_gpus <= _get_cluster_resources_no_autoscaler().get("GPU", 0)
+
+
+@lru_cache()
+def _get_insufficient_resources_warning_threshold() -> float:
+    if _is_ray_cluster():
+        return float(
+            os.environ.get(
+                "TUNE_WARN_INSUFFICENT_RESOURCE_THRESHOLD_S_AUTOSCALER", "60"
+            )
+        )
+    else:
+        # Set the default to 10s so that we don't prematurely determine that
+        # a cluster cannot fulfill the resources requirements.
+        # TODO(xwjiang): Change it back once #18608 is resolved.
+        return float(os.environ.get("TUNE_WARN_INSUFFICENT_RESOURCE_THRESHOLD_S", "60"))
+
+
+MSG_TRAIN_START = (
+    "Training has not started in the last {wait_time:.0f} seconds. "
+    "This could be due to the cluster not having enough resources available. "
+)
+MSG_TRAIN_INSUFFICIENT = (
+    "You asked for {asked_cpus} CPUs and {asked_gpus} GPUs, but the cluster only "
+    "has {cluster_cpus} CPUs and {cluster_gpus} GPUs available. "
+)
+MSG_TRAIN_END = (
+    "Stop the training and adjust the required resources (e.g. via the "
+    "`ScalingConfig` or `resources_per_trial`, or `num_workers` for rllib), "
+    "or add more resources to your cluster."
+)
+
+MSG_TUNE_START = (
+    "No trial is running and no new trial has been started within "
+    "the last {wait_time:.0f} seconds. "
+    "This could be due to the cluster not having enough resources available. "
+)
+MSG_TUNE_INSUFFICIENT = (
+    "You asked for {asked_cpus} CPUs and {asked_gpus} GPUs per trial, "
+    "but the cluster only has {cluster_cpus} CPUs and {cluster_gpus} GPUs available. "
+)
+MSG_TUNE_END = (
+    "Stop the tuning and adjust the required resources (e.g. via the "
+    "`ScalingConfig` or `resources_per_trial`, or `num_workers` for rllib), "
+    "or add more resources to your cluster."
+)
+
+
+# TODO(xwjiang): Consider having a help page with more detailed instructions.
+@lru_cache()
+def _get_insufficient_resources_warning_msg(
+    for_train: bool = False, trial: Optional[Trial] = None
+) -> str:
+    msg = "Ignore this message if the cluster is autoscaling. "
+
+    if for_train:
+        start = MSG_TRAIN_START
+        insufficient = MSG_TRAIN_INSUFFICIENT
+        end = MSG_TRAIN_END
+    else:
+        start = MSG_TUNE_START
+        insufficient = MSG_TUNE_INSUFFICIENT
+        end = MSG_TUNE_END
+
+    msg += start.format(wait_time=_get_insufficient_resources_warning_threshold())
+
+    if trial:
+        asked_cpus, asked_gpus = _get_trial_cpu_and_gpu(trial)
+        cluster_resources = _get_cluster_resources_no_autoscaler()
+
+        msg += insufficient.format(
+            asked_cpus=asked_cpus,
+            asked_gpus=asked_gpus,
+            cluster_cpus=cluster_resources.get("CPU", 0),
+            cluster_gpus=cluster_resources.get("GPU", 0),
+        )
+
+    msg += end
+
+    return msg
+
+
+class _InsufficientResourcesManager:
+    """Insufficient resources manager.
+
+    Makes best effort, conservative guesses about if Tune loop is stuck due to
+    infeasible resources. If so, outputs usability messages for users to
+    act upon.
+    """
+
+    def __init__(self, for_train: bool = False):
+        # The information tracked across the life time of Tune loop.
+        self._no_running_trials_since = -1
+        self._last_trial_num = -1
+        self._for_train = for_train
+
+    def on_no_available_trials(self, all_trials):
+        """Tracks information across the life of Tune loop and makes guesses
+        about if Tune loop is stuck due to infeasible resources.
+        If so, outputs certain warning messages.
+        The logic should be conservative, non-intrusive and informative.
+        For example, rate limiting is applied so that the message is not
+        spammy.
+        """
+        # This is approximately saying we are not making progress.
+        if len(all_trials) == self._last_trial_num:
+            if self._no_running_trials_since == -1:
+                self._no_running_trials_since = time.monotonic()
+            elif (
+                time.monotonic() - self._no_running_trials_since
+                > _get_insufficient_resources_warning_threshold()
+            ):
+                can_fulfill_any = any(
+                    trial.status == Trial.PENDING and _can_fulfill_no_autoscaler(trial)
+                    for trial in all_trials
+                )
+
+                if can_fulfill_any:
+                    # If one trial can be fulfilled, it will be fulfilled eventually
+                    self._no_running_trials_since = -1
+                    return
+
+                # Otherwise, can fulfill none
+                msg = _get_insufficient_resources_warning_msg(
+                    for_train=self._for_train, trial=all_trials[0]
+                )
+                logger.warning(msg)
+                self._no_running_trials_since = time.monotonic()
+        else:
+            self._no_running_trials_since = -1
+        self._last_trial_num = len(all_trials)

python/ray/tune/execution/placement_groups.py

@@ -0,0 +1,131 @@
+import warnings
+from typing import Dict, Optional
+
+from ray.air.execution.resources.request import ResourceRequest
+from ray.util.annotations import DeveloperAPI, PublicAPI
+from ray.util.placement_group import placement_group
+
+
+@PublicAPI(stability="beta")
+class PlacementGroupFactory(ResourceRequest):
+    """Wrapper class that creates placement groups for trials.
+
+    This function should be used to define resource requests for Ray Tune
+    trials. It holds the parameters to create
+    :ref:`placement groups <ray-placement-group-doc-ref>`.
+    At a minimum, this will hold at least one bundle specifying the
+    resource requirements for each trial:
+
+    .. code-block:: python
+
+        from ray import tune
+
+        tuner = tune.Tuner(
+            tune.with_resources(
+                train,
+                resources=tune.PlacementGroupFactory([
+                    {"CPU": 1, "GPU": 0.5, "custom_resource": 2}
+                ])
+            )
+        )
+        tuner.fit()
+
+    If the trial itself schedules further remote workers, the resource
+    requirements should be specified in additional bundles. You can also
+    pass the placement strategy for these bundles, e.g. to enforce
+    co-located placement:
+
+    .. code-block:: python
+
+        from ray import tune
+
+        tuner = tune.Tuner(
+            tune.with_resources(
+                train,
+                resources=tune.PlacementGroupFactory([
+                    {"CPU": 1, "GPU": 0.5, "custom_resource": 2},
+                    {"CPU": 2},
+                    {"CPU": 2},
+                ], strategy="PACK")
+            )
+        )
+        tuner.fit()
+
+    The example above will reserve 1 CPU, 0.5 GPUs and 2 custom_resources
+    for the trainable itself, and reserve another 2 bundles of 2 CPUs each.
+    The trial will only start when all these resources are available. This
+    could be used e.g. if you had one learner running in the main trainable
+    that schedules two remote workers that need access to 2 CPUs each.
+
+    If the trainable itself doesn't require resources.
+    You can specify it as:
+
+    .. code-block:: python
+
+        from ray import tune
+
+        tuner = tune.Tuner(
+            tune.with_resources(
+                train,
+                resources=tune.PlacementGroupFactory([
+                    {},
+                    {"CPU": 2},
+                    {"CPU": 2},
+                ], strategy="PACK")
+            )
+        )
+        tuner.fit()
+
+    Args:
+        bundles: A list of bundles which
+            represent the resources requirements.
+        strategy: The strategy to create the placement group.
+
+         - "PACK": Packs Bundles into as few nodes as possible.
+         - "SPREAD": Places Bundles across distinct nodes as even as possible.
+         - "STRICT_PACK": Packs Bundles into one node. The group is
+           not allowed to span multiple nodes.
+         - "STRICT_SPREAD": Packs Bundles across distinct nodes.
+        *args: Passed to the call of ``placement_group()``
+        **kwargs: Passed to the call of ``placement_group()``
+
+    """
+
+    def __call__(self, *args, **kwargs):
+        warnings.warn(
+            "Calling PlacementGroupFactory objects is deprecated. Use "
+            "`to_placement_group()` instead.",
+            DeprecationWarning,
+        )
+        kwargs.update(self._bound.kwargs)
+        # Call with bounded *args and **kwargs
+        return placement_group(*self._bound.args, **kwargs)
+
+
+@DeveloperAPI
+def resource_dict_to_pg_factory(spec: Optional[Dict[str, float]] = None):
+    """Translates resource dict into PlacementGroupFactory."""
+    spec = spec or {"cpu": 1}
+
+    spec = spec.copy()
+
+    cpus = spec.pop("cpu", spec.pop("CPU", 0.0))
+    gpus = spec.pop("gpu", spec.pop("GPU", 0.0))
+    memory = spec.pop("memory", 0.0)
+
+    # If there is a custom_resources key, use as base for bundle
+    bundle = dict(spec.pop("custom_resources", {}))
+
+    # Otherwise, consider all other keys as custom resources
+    if not bundle:
+        bundle = spec
+
+    bundle.update(
+        {
+            "CPU": cpus,
+            "GPU": gpus,
+            "memory": memory,
+        }
+    )
+
+    return PlacementGroupFactory([bundle])

python/ray/tune/experiment/__init__.py

@@ -0,0 +1,4 @@
+from ray.tune.experiment.experiment import Experiment, _convert_to_experiment_list
+from ray.tune.experiment.trial import Trial
+
+__all__ = ["Experiment", "_convert_to_experiment_list", "Trial"]

python/ray/tune/experiment/config_parser.py

@@ -0,0 +1,218 @@
+import argparse
+import json
+from typing import Any, Callable, Optional
+
+from ray.tune import CheckpointConfig
+from ray.tune.error import TuneError
+from ray.tune.experiment import Trial
+from ray.tune.resources import json_to_resources
+
+# For compatibility under py2 to consider unicode as str
+from ray.tune.utils.serialization import TuneFunctionEncoder
+from ray.tune.utils.util import SafeFallbackEncoder
+
+
+def _make_parser(
+    parser_creator: Optional[Callable[..., argparse.ArgumentParser]] = None,
+    **kwargs: Any,
+) -> argparse.ArgumentParser:
+    """Returns a base argument parser for the ray.tune tool.
+
+    Args:
+        parser_creator: A constructor for the parser class.
+        **kwargs: Non-positional args to be passed into the
+            parser class constructor.
+
+    Returns:
+        An ``argparse.ArgumentParser`` configured with the standard Tune
+        command-line flags.
+    """
+
+    if parser_creator:
+        parser = parser_creator(**kwargs)
+    else:
+        parser = argparse.ArgumentParser(**kwargs)
+
+    # Note: keep this in sync with rllib/train.py
+    parser.add_argument(
+        "--run",
+        default=None,
+        type=str,
+        help="The algorithm or model to train. This may refer to the name "
+        "of a built-on algorithm (e.g. RLlib's DQN or PPO), or a "
+        "user-defined trainable function or class registered in the "
+        "tune registry.",
+    )
+    parser.add_argument(
+        "--stop",
+        default="{}",
+        type=json.loads,
+        help="The stopping criteria, specified in JSON. The keys may be any "
+        "field returned by 'train()' e.g. "
+        '\'{"time_total_s": 600, "training_iteration": 100000}\' to stop '
+        "after 600 seconds or 100k iterations, whichever is reached first.",
+    )
+    parser.add_argument(
+        "--config",
+        default="{}",
+        type=json.loads,
+        help="Algorithm-specific configuration (e.g. env, hyperparams), "
+        "specified in JSON.",
+    )
+    parser.add_argument(
+        "--resources-per-trial",
+        default=None,
+        type=json_to_resources,
+        help="Override the machine resources to allocate per trial, e.g. "
+        '\'{"cpu": 64, "gpu": 8}\'. Note that GPUs will not be assigned '
+        "unless you specify them here. For RLlib, you probably want to "
+        "leave this alone and use RLlib configs to control parallelism.",
+    )
+    parser.add_argument(
+        "--num-samples",
+        default=1,
+        type=int,
+        help="Number of times to repeat each trial.",
+    )
+    parser.add_argument(
+        "--checkpoint-freq",
+        default=0,
+        type=int,
+        help="How many training iterations between checkpoints. "
+        "A value of 0 (default) disables checkpointing.",
+    )
+    parser.add_argument(
+        "--checkpoint-at-end",
+        action="store_true",
+        help="Whether to checkpoint at the end of the experiment. Default is False.",
+    )
+    parser.add_argument(
+        "--keep-checkpoints-num",
+        default=None,
+        type=int,
+        help="Number of best checkpoints to keep. Others get "
+        "deleted. Default (None) keeps all checkpoints.",
+    )
+    parser.add_argument(
+        "--checkpoint-score-attr",
+        default="training_iteration",
+        type=str,
+        help="Specifies by which attribute to rank the best checkpoint. "
+        "Default is increasing order. If attribute starts with min- it "
+        "will rank attribute in decreasing order. Example: "
+        "min-validation_loss",
+    )
+    parser.add_argument(
+        "--export-formats",
+        default=None,
+        help="List of formats that exported at the end of the experiment. "
+        "Default is None. For RLlib, 'checkpoint' and 'model' are "
+        "supported for TensorFlow policy graphs.",
+    )
+    parser.add_argument(
+        "--max-failures",
+        default=3,
+        type=int,
+        help="Try to recover a trial from its last checkpoint at least this "
+        "many times. Only applies if checkpointing is enabled.",
+    )
+    parser.add_argument(
+        "--scheduler",
+        default="FIFO",
+        type=str,
+        help="FIFO (default), MedianStopping, AsyncHyperBand, "
+        "HyperBand, or HyperOpt.",
+    )
+    parser.add_argument(
+        "--scheduler-config",
+        default="{}",
+        type=json.loads,
+        help="Config options to pass to the scheduler.",
+    )
+
+    # Note: this currently only makes sense when running a single trial
+    parser.add_argument(
+        "--restore",
+        default=None,
+        type=str,
+        help="If specified, restore from this checkpoint.",
+    )
+
+    return parser
+
+
+def _to_argv(config):
+    """Converts configuration to a command line argument format."""
+    argv = []
+    for k, v in config.items():
+        if "-" in k:
+            raise ValueError("Use '_' instead of '-' in `{}`".format(k))
+        if v is None:
+            continue
+        if not isinstance(v, bool) or v:  # for argparse flags
+            argv.append("--{}".format(k.replace("_", "-")))
+        if isinstance(v, str):
+            argv.append(v)
+        elif isinstance(v, bool):
+            pass
+        elif callable(v):
+            argv.append(json.dumps(v, cls=TuneFunctionEncoder))
+        else:
+            argv.append(json.dumps(v, cls=SafeFallbackEncoder))
+    return argv
+
+
+_cached_pgf = {}
+
+
+def _create_trial_from_spec(
+    spec: dict, parser: argparse.ArgumentParser, **trial_kwargs
+):
+    """Creates a Trial object from parsing the spec.
+
+    Args:
+        spec: A resolved experiment specification. Arguments should
+            The args here should correspond to the command line flags
+            in ray.tune.experiment.config_parser.
+        parser: An argument parser object from
+            make_parser.
+        **trial_kwargs: Extra keyword arguments used in instantiating the Trial.
+
+    Returns:
+        A trial object with corresponding parameters to the specification.
+    """
+    global _cached_pgf
+
+    spec = spec.copy()
+    resources = spec.pop("resources_per_trial", None)
+
+    try:
+        args, _ = parser.parse_known_args(_to_argv(spec))
+    except SystemExit:
+        raise TuneError("Error parsing args, see above message", spec)
+
+    if resources:
+        trial_kwargs["placement_group_factory"] = resources
+
+    checkpoint_config = spec.get("checkpoint_config", CheckpointConfig())
+
+    return Trial(
+        # Submitting trial via server in py2.7 creates Unicode, which does not
+        # convert to string in a straightforward manner.
+        trainable_name=spec["run"],
+        # json.load leads to str -> unicode in py2.7
+        config=spec.get("config", {}),
+        # json.load leads to str -> unicode in py2.7
+        stopping_criterion=spec.get("stop", {}),
+        checkpoint_config=checkpoint_config,
+        export_formats=spec.get("export_formats", []),
+        # str(None) doesn't create None
+        restore_path=spec.get("restore"),
+        trial_name_creator=spec.get("trial_name_creator"),
+        trial_dirname_creator=spec.get("trial_dirname_creator"),
+        log_to_file=spec.get("log_to_file"),
+        # str(None) doesn't create None
+        max_failures=args.max_failures,
+        storage=spec.get("storage"),
+        **trial_kwargs,
+    )

python/ray/tune/experiment/experiment.py

@@ -0,0 +1,450 @@
+import copy
+import datetime
+import logging
+import pprint as pp
+import traceback
+from functools import partial
+from pathlib import Path
+from pickle import PicklingError
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    List,
+    Mapping,
+    Optional,
+    Sequence,
+    Type,
+    Union,
+)
+
+import ray
+from ray.exceptions import RpcError
+from ray.train._internal.storage import StorageContext
+from ray.train.constants import DEFAULT_STORAGE_PATH
+from ray.tune import CheckpointConfig, SyncConfig
+from ray.tune.error import TuneError
+from ray.tune.registry import is_function_trainable, register_trainable
+from ray.tune.stopper import CombinedStopper, FunctionStopper, Stopper, TimeoutStopper
+from ray.util.annotations import Deprecated, DeveloperAPI
+
+if TYPE_CHECKING:
+    import pyarrow.fs
+
+    from ray.tune import PlacementGroupFactory
+    from ray.tune.experiment import Trial
+
+
+logger = logging.getLogger(__name__)
+
+
+def _validate_log_to_file(log_to_file):
+    """Validate ``tune.RunConfig``'s ``log_to_file`` parameter. Return
+    validated relative stdout and stderr filenames."""
+    if not log_to_file:
+        stdout_file = stderr_file = None
+    elif isinstance(log_to_file, bool) and log_to_file:
+        stdout_file = "stdout"
+        stderr_file = "stderr"
+    elif isinstance(log_to_file, str):
+        stdout_file = stderr_file = log_to_file
+    elif isinstance(log_to_file, Sequence):
+        if len(log_to_file) != 2:
+            raise ValueError(
+                "If you pass a Sequence to `log_to_file` it has to have "
+                "a length of 2 (for stdout and stderr, respectively). The "
+                "Sequence you passed has length {}.".format(len(log_to_file))
+            )
+        stdout_file, stderr_file = log_to_file
+    else:
+        raise ValueError(
+            "You can pass a boolean, a string, or a Sequence of length 2 to "
+            "`log_to_file`, but you passed something else ({}).".format(
+                type(log_to_file)
+            )
+        )
+    return stdout_file, stderr_file
+
+
+@DeveloperAPI
+class Experiment:
+    """Tracks experiment specifications.
+
+    Implicitly registers the Trainable if needed. The args here take
+    the same meaning as the arguments defined `tune.py:run`.
+
+    .. code-block:: python
+
+        experiment_spec = Experiment(
+            "my_experiment_name",
+            my_func,
+            stop={"mean_accuracy": 100},
+            config={
+                "alpha": tune.grid_search([0.2, 0.4, 0.6]),
+                "beta": tune.grid_search([1, 2]),
+            },
+            resources_per_trial={
+                "cpu": 1,
+                "gpu": 0
+            },
+            num_samples=10,
+            local_dir="~/ray_results",
+            checkpoint_freq=10,
+            max_failures=2)
+
+    """
+
+    # Keys that will be present in `public_spec` dict.
+    PUBLIC_KEYS = {"stop", "num_samples", "time_budget_s"}
+    _storage_context_cls = StorageContext
+
+    def __init__(
+        self,
+        name: str,
+        run: Union[str, Callable, Type],
+        *,
+        stop: Optional[Union[Mapping, Stopper, Callable[[str, Mapping], bool]]] = None,
+        time_budget_s: Optional[Union[int, float, datetime.timedelta]] = None,
+        config: Optional[Dict[str, Any]] = None,
+        resources_per_trial: Union[
+            None, Mapping[str, Union[float, int, Mapping]], "PlacementGroupFactory"
+        ] = None,
+        num_samples: int = 1,
+        storage_path: Optional[str] = None,
+        storage_filesystem: Optional["pyarrow.fs.FileSystem"] = None,
+        sync_config: Optional[Union[SyncConfig, dict]] = None,
+        checkpoint_config: Optional[Union[CheckpointConfig, dict]] = None,
+        trial_name_creator: Optional[Callable[["Trial"], str]] = None,
+        trial_dirname_creator: Optional[Callable[["Trial"], str]] = None,
+        log_to_file: bool = False,
+        export_formats: Optional[Sequence] = None,
+        max_failures: int = 0,
+        restore: Optional[str] = None,
+        # Deprecated
+        local_dir: Optional[str] = None,
+    ):
+        if isinstance(checkpoint_config, dict):
+            checkpoint_config = CheckpointConfig(**checkpoint_config)
+        else:
+            checkpoint_config = checkpoint_config or CheckpointConfig()
+
+        if is_function_trainable(run):
+            if checkpoint_config.checkpoint_at_end:
+                raise ValueError(
+                    "'checkpoint_at_end' cannot be used with a function trainable. "
+                    "You should include one last call to "
+                    "`ray.tune.report(metrics=..., checkpoint=...)` "
+                    "at the end of your training loop to get this behavior."
+                )
+            if checkpoint_config.checkpoint_frequency:
+                raise ValueError(
+                    "'checkpoint_frequency' cannot be set for a function trainable. "
+                    "You will need to report a checkpoint every "
+                    "`checkpoint_frequency` iterations within your training loop using "
+                    "`ray.tune.report(metrics=..., checkpoint=...)` "
+                    "to get this behavior."
+                )
+        try:
+            self._run_identifier = Experiment.register_if_needed(run)
+        except RpcError as e:
+            if e.rpc_code == ray._raylet.GRPC_STATUS_CODE_RESOURCE_EXHAUSTED:
+                raise TuneError(
+                    f"The Trainable/training function is too large for grpc resource "
+                    f"limit. Check that its definition is not implicitly capturing a "
+                    f"large array or other object in scope. "
+                    f"Tip: use tune.with_parameters() to put large objects "
+                    f"in the Ray object store. \n"
+                    f"Original exception: {traceback.format_exc()}"
+                )
+            else:
+                raise e
+
+        if not name:
+            name = StorageContext.get_experiment_dir_name(run)
+
+        storage_path = storage_path or DEFAULT_STORAGE_PATH
+        self.storage = self._storage_context_cls(
+            storage_path=storage_path,
+            storage_filesystem=storage_filesystem,
+            sync_config=sync_config,
+            experiment_dir_name=name,
+        )
+        logger.debug(f"StorageContext on the DRIVER:\n{self.storage}")
+
+        config = config or {}
+        if not isinstance(config, dict):
+            raise ValueError(
+                f"`Experiment(config)` must be a dict, got: {type(config)}. "
+                "Please convert your search space to a dict before passing it in."
+            )
+
+        self._stopper = None
+        stopping_criteria = {}
+        if not stop:
+            pass
+        elif isinstance(stop, list):
+            bad_stoppers = [s for s in stop if not isinstance(s, Stopper)]
+            if bad_stoppers:
+                stopper_types = [type(s) for s in stop]
+                raise ValueError(
+                    "If you pass a list as the `stop` argument to "
+                    "`tune.RunConfig()`, each element must be an instance of "
+                    f"`tune.stopper.Stopper`. Got {stopper_types}."
+                )
+            self._stopper = CombinedStopper(*stop)
+        elif isinstance(stop, dict):
+            stopping_criteria = stop
+        elif callable(stop):
+            if FunctionStopper.is_valid_function(stop):
+                self._stopper = FunctionStopper(stop)
+            elif isinstance(stop, Stopper):
+                self._stopper = stop
+            else:
+                raise ValueError(
+                    "Provided stop object must be either a dict, "
+                    "a function, or a subclass of "
+                    f"`ray.tune.Stopper`. Got {type(stop)}."
+                )
+        else:
+            raise ValueError(
+                f"Invalid stop criteria: {stop}. Must be a "
+                f"callable or dict. Got {type(stop)}."
+            )
+
+        if time_budget_s:
+            if self._stopper:
+                self._stopper = CombinedStopper(
+                    self._stopper, TimeoutStopper(time_budget_s)
+                )
+            else:
+                self._stopper = TimeoutStopper(time_budget_s)
+
+        stdout_file, stderr_file = _validate_log_to_file(log_to_file)
+
+        spec = {
+            "run": self._run_identifier,
+            "stop": stopping_criteria,
+            "time_budget_s": time_budget_s,
+            "config": config,
+            "resources_per_trial": resources_per_trial,
+            "num_samples": num_samples,
+            "checkpoint_config": checkpoint_config,
+            "trial_name_creator": trial_name_creator,
+            "trial_dirname_creator": trial_dirname_creator,
+            "log_to_file": (stdout_file, stderr_file),
+            "export_formats": export_formats or [],
+            "max_failures": max_failures,
+            "restore": (
+                Path(restore).expanduser().absolute().as_posix() if restore else None
+            ),
+            "storage": self.storage,
+        }
+        self.spec = spec
+
+    @classmethod
+    def from_json(cls, name: str, spec: dict):
+        """Generates an Experiment object from JSON.
+
+        Args:
+            name: Name of Experiment.
+            spec: JSON configuration of experiment.
+
+        Returns:
+            An ``Experiment`` constructed from the provided ``spec``.
+        """
+        if "run" not in spec:
+            raise TuneError("No trainable specified!")
+
+        # Special case the `env` param for RLlib by automatically
+        # moving it into the `config` section.
+        if "env" in spec:
+            spec["config"] = spec.get("config", {})
+            spec["config"]["env"] = spec["env"]
+            del spec["env"]
+
+        if "sync_config" in spec and isinstance(spec["sync_config"], dict):
+            spec["sync_config"] = SyncConfig(**spec["sync_config"])
+
+        if "checkpoint_config" in spec and isinstance(spec["checkpoint_config"], dict):
+            spec["checkpoint_config"] = CheckpointConfig(**spec["checkpoint_config"])
+
+        spec = copy.deepcopy(spec)
+
+        run_value = spec.pop("run")
+        try:
+            exp = cls(name, run_value, **spec)
+        except TypeError as e:
+            raise TuneError(
+                f"Failed to load the following Tune experiment "
+                f"specification:\n\n {pp.pformat(spec)}.\n\n"
+                f"Please check that the arguments are valid. "
+                f"Experiment creation failed with the following "
+                f"error:\n {e}"
+            )
+        return exp
+
+    @classmethod
+    def get_trainable_name(cls, run_object: Union[str, Callable, Type]):
+        """Get Trainable name.
+
+        Args:
+            run_object: Trainable to run. If string,
+                assumes it is an ID and does not modify it. Otherwise,
+                returns a string corresponding to the run_object name.
+
+        Returns:
+            A string representing the trainable identifier.
+
+        Raises:
+            TuneError: if ``run_object`` passed in is invalid.
+        """
+        from ray.tune.search.sample import Domain
+
+        if isinstance(run_object, str) or isinstance(run_object, Domain):
+            return run_object
+        elif isinstance(run_object, type) or callable(run_object):
+            name = "DEFAULT"
+            if hasattr(run_object, "_name"):
+                name = run_object._name
+            elif hasattr(run_object, "__name__"):
+                fn_name = run_object.__name__
+                if fn_name == "<lambda>":
+                    name = "lambda"
+                elif fn_name.startswith("<"):
+                    name = "DEFAULT"
+                else:
+                    name = fn_name
+            elif (
+                isinstance(run_object, partial)
+                and hasattr(run_object, "func")
+                and hasattr(run_object.func, "__name__")
+            ):
+                name = run_object.func.__name__
+            else:
+                logger.warning("No name detected on trainable. Using {}.".format(name))
+            return name
+        else:
+            raise TuneError("Improper 'run' - not string nor trainable.")
+
+    @classmethod
+    def register_if_needed(cls, run_object: Union[str, Callable, Type]):
+        """Registers Trainable or Function at runtime.
+
+        Assumes already registered if run_object is a string.
+        Also, does not inspect interface of given run_object.
+
+        Args:
+            run_object: Trainable to run. If string,
+                assumes it is an ID and does not modify it. Otherwise,
+                returns a string corresponding to the run_object name.
+
+        Returns:
+            A string representing the trainable identifier.
+        """
+        from ray.tune.search.sample import Domain
+
+        if isinstance(run_object, str):
+            return run_object
+        elif isinstance(run_object, Domain):
+            logger.warning("Not registering trainable. Resolving as variant.")
+            return run_object
+        name = cls.get_trainable_name(run_object)
+        try:
+            register_trainable(name, run_object)
+        except (TypeError, PicklingError) as e:
+            extra_msg = (
+                "Other options: "
+                "\n-Try reproducing the issue by calling "
+                "`pickle.dumps(trainable)`. "
+                "\n-If the error is typing-related, try removing "
+                "the type annotations and try again."
+            )
+            raise type(e)(str(e) + " " + extra_msg) from None
+        return name
+
+    @property
+    def stopper(self):
+        return self._stopper
+
+    @property
+    def local_path(self) -> Optional[str]:
+        return self.storage.experiment_driver_staging_path
+
+    @property
+    @Deprecated("Replaced by `local_path`")
+    def local_dir(self):
+        # TODO(justinvyu): [Deprecated] Remove in 2.11.
+        raise DeprecationWarning("Use `local_path` instead of `local_dir`.")
+
+    @property
+    def remote_path(self) -> Optional[str]:
+        return self.storage.experiment_fs_path
+
+    @property
+    def path(self) -> Optional[str]:
+        return self.remote_path or self.local_path
+
+    @property
+    def checkpoint_config(self):
+        return self.spec.get("checkpoint_config")
+
+    @property
+    @Deprecated("Replaced by `local_path`")
+    def checkpoint_dir(self):
+        # TODO(justinvyu): [Deprecated] Remove in 2.11.
+        raise DeprecationWarning("Use `local_path` instead of `checkpoint_dir`.")
+
+    @property
+    def run_identifier(self):
+        """Returns a string representing the trainable identifier."""
+        return self._run_identifier
+
+    @property
+    def public_spec(self) -> Dict[str, Any]:
+        """Returns the spec dict with only the public-facing keys.
+
+        Intended to be used for passing information to callbacks,
+        Searchers and Schedulers.
+        """
+        return {k: v for k, v in self.spec.items() if k in self.PUBLIC_KEYS}
+
+
+def _convert_to_experiment_list(experiments: Union[Experiment, List[Experiment], Dict]):
+    """Produces a list of Experiment objects.
+
+    Converts input from dict, single experiment, or list of
+    experiments to list of experiments. If input is None,
+    will return an empty list.
+
+    Arguments:
+        experiments: Experiments to run.
+
+    Returns:
+        List of experiments.
+    """
+    exp_list = experiments
+
+    # Transform list if necessary
+    if experiments is None:
+        exp_list = []
+    elif isinstance(experiments, Experiment):
+        exp_list = [experiments]
+    elif isinstance(experiments, dict):
+        exp_list = [
+            Experiment.from_json(name, spec) for name, spec in experiments.items()
+        ]
+
+    # Validate exp_list
+    if isinstance(exp_list, list) and all(
+        isinstance(exp, Experiment) for exp in exp_list
+    ):
+        if len(exp_list) > 1:
+            logger.info(
+                "Running with multiple concurrent experiments. "
+                "All experiments will be using the same SearchAlgorithm."
+            )
+    else:
+        raise TuneError("Invalid argument: {}".format(experiments))
+
+    return exp_list

python/ray/tune/impl/config.py

@@ -0,0 +1,53 @@
+from dataclasses import dataclass
+
+from ray.air.config import (
+    CheckpointConfig as _CheckpointConfig,
+    FailureConfig as _FailureConfig,
+    RunConfig as _RunConfig,
+)
+from ray.train.constants import (
+    V2_MIGRATION_GUIDE_MESSAGE,
+    _v2_migration_warnings_enabled,
+)
+from ray.train.utils import _copy_doc, _log_deprecation_warning
+
+# NOTE: This is just a pass-through wrapper around `ray.tune.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."
+                    f"{V2_MIGRATION_GUIDE_MESSAGE}"
+                )
+
+        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."
+                    f"{V2_MIGRATION_GUIDE_MESSAGE}"
+                )

python/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)

python/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.sha256(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,))

python/ray/tune/impl/test_utils.py

@@ -0,0 +1,65 @@
+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,
+            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._has_computed_output()
+    assert not ds2._has_computed_output()
+    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]._has_computed_output()
+    assert executed_ds[1]._has_computed_output()
+
+
+def test_choice():
+    ds1 = gen_dataset_func().lazy().map(lambda x: x)
+    ds2 = gen_dataset_func().lazy().map(lambda x: x)
+    assert not ds1._has_computed_output()
+    assert not ds2._has_computed_output()
+    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]._has_computed_output()
+    assert executed_ds[1]._has_computed_output()
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/tune/impl/tuner_internal.py

@@ -0,0 +1,698 @@
+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
+import ray.tune
+from ray.air._internal.uri_utils import URI
+from ray.air._internal.usage import AirEntrypoint
+from ray.train._internal.storage import StorageContext, get_fs_and_path
+from ray.train.constants import (
+    V2_MIGRATION_GUIDE_MESSAGE,
+    _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.
+        storage_filesystem: A custom ``pyarrow.fs.FileSystem`` corresponding
+            to ``restore_path``. This may be necessary if the original
+            experiment used a custom filesystem.
+        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.
+        _tuner_kwargs: Internal. Extra kwargs forwarded to ``tune.run`` when
+            this Tuner is fit.
+        _entrypoint: Internal. Tracks which user-facing entrypoint constructed
+            this Tuner so that warnings and errors can be specialized.
+    """
+
+    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(
+                    "The Ray Train + Ray Tune integration has been reworked. "
+                    "Passing a Trainer to the Tuner is deprecated and will be removed "
+                    "in a future release. "
+                    f"{V2_MIGRATION_GUIDE_MESSAGE}"
+                )
+
+            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()
+        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.")
+
+        if self._entrypoint == AirEntrypoint.TUNER and not isinstance(
+            self._run_config, ray.tune.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. "
+                    f"{V2_MIGRATION_GUIDE_MESSAGE}"
+                )
+
+        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,
+    ) -> 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).
+
+        Args:
+            trainable: The trainable to validate.
+            required_trainable_name: If provided, the trainable's generated
+                name must match this value; used on restoration to detect a
+                trainable swap.
+
+        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]],
+    ) -> None:
+        """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.
+
+        Args:
+            new_param_space: The newly provided search space to validate.
+            flattened_param_space_keys: Sorted flat keys of the original
+                ``param_space``. ``None`` skips validation for backwards
+                compatibility.
+
+        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_state: Deserialized contents of the `tuner.pkl` 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.
+
+        Returns:
+            The resolved ``RunConfig`` to use for the Tune experiment.
+
+        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, ray.train.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.tune.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.tune.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)

python/ray/tune/integration/keras.py

@@ -0,0 +1,72 @@
+from typing import Dict
+
+import ray.tune
+from ray.train.tensorflow import TensorflowCheckpoint
+from ray.train.tensorflow.keras import RayReportCallback
+from ray.util.annotations import PublicAPI
+
+_DEPRECATION_MESSAGE = (
+    "The `ray.tune.integration.keras` module is deprecated in favor of "
+    "`ray.train.tensorflow.keras.ReportCheckpointCallback`."
+)
+
+
+class TuneReportCallback:
+    """Deprecated.
+    Use :class:`ray.train.tensorflow.keras.ReportCheckpointCallback` instead."""
+
+    def __new__(cls, *args, **kwargs):
+        raise DeprecationWarning(_DEPRECATION_MESSAGE)
+
+
+class _TuneCheckpointCallback:
+    """Deprecated.
+    Use :class:`ray.train.tensorflow.keras.ReportCheckpointCallback` instead."""
+
+    def __new__(cls, *args, **kwargs):
+        raise DeprecationWarning(_DEPRECATION_MESSAGE)
+
+
+@PublicAPI(stability="alpha")
+class TuneReportCheckpointCallback(RayReportCallback):
+    """Keras callback for Ray Tune reporting and checkpointing.
+
+    .. note::
+        Metrics are always reported with checkpoints, even if the event isn't specified
+        in ``report_metrics_on``.
+
+    Example:
+        .. code-block:: python
+
+            ############# Using it in Ray Tune ###############
+            from ray.tune.integrations.keras import TuneReportCheckpointCallback
+
+            def train_fn():
+                model = build_model()
+                model.fit(dataset_shard, callbacks=[TuneReportCheckpointCallback()])
+
+            tuner = tune.Tuner(train_fn)
+            results = tuner.fit()
+
+    Args:
+        metrics: Metrics to report. If this is a list, each item describes
+            the metric key reported to Keras, and it's reported under the
+            same name. If this is a dict, each key is the name reported
+            and the respective value is the metric key reported to Keras.
+            If this is None, all Keras logs are reported.
+        report_metrics_on: When to report metrics. Must be one of
+            the Keras event hooks (less the ``on_``), e.g.
+            "train_start" or "predict_end". Defaults to "epoch_end".
+        checkpoint_on: When to save checkpoints. Must be one of the Keras event hooks
+            (less the ``on_``), e.g. "train_start" or "predict_end". Defaults to
+            "epoch_end".
+
+    """
+
+    def _save_and_report_checkpoint(
+        self, metrics: Dict, checkpoint: TensorflowCheckpoint
+    ):
+        ray.tune.report(metrics, checkpoint=checkpoint)
+
+    def _report_metrics(self, metrics: Dict):
+        ray.tune.report(metrics)

python/ray/tune/integration/lightgbm.py

@@ -0,0 +1,84 @@
+import tempfile
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Dict, Optional
+
+from lightgbm import Booster
+
+import ray.tune
+from ray.train.lightgbm._lightgbm_utils import RayReportCallback
+from ray.tune import Checkpoint
+from ray.util.annotations import Deprecated, PublicAPI
+
+
+@PublicAPI(stability="beta")
+class TuneReportCheckpointCallback(RayReportCallback):
+    """Creates a callback that reports metrics and checkpoints model.
+
+    Args:
+        metrics: Metrics to report. If this is a list,
+            each item should be a metric key reported by LightGBM,
+            and it will be reported to Ray Train/Tune under the same name.
+            This can also be a dict of {<key-to-report>: <lightgbm-metric-key>},
+            which can be used to rename LightGBM default metrics.
+        filename: Customize the saved checkpoint file type by passing
+            a filename. Defaults to "model.txt".
+        frequency: How often to save checkpoints, in terms of iterations.
+            Defaults to 0 (no checkpoints are saved during training).
+        checkpoint_at_end: Whether or not to save a checkpoint at the end of training.
+        results_postprocessing_fn: An optional Callable that takes in
+            the metrics dict that will be reported (after it has been flattened)
+            and returns a modified dict.
+
+    Examples
+    --------
+
+    Reporting checkpoints and metrics to Ray Tune when running many
+    independent LightGBM trials (without data parallelism within a trial).
+
+    .. testcode::
+        :skipif: True
+
+        import lightgbm
+
+        from ray.tune.integration.lightgbm import TuneReportCheckpointCallback
+
+        config = {
+            # ...
+            "metric": ["binary_logloss", "binary_error"],
+        }
+
+        # Report only log loss to Tune after each validation epoch.
+        bst = lightgbm.train(
+            ...,
+            callbacks=[
+                TuneReportCheckpointCallback(
+                    metrics={"loss": "eval-binary_logloss"}, frequency=1
+                )
+            ],
+        )
+
+    """
+
+    @contextmanager
+    def _get_checkpoint(self, model: Booster) -> Optional[Checkpoint]:
+        with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
+            model.save_model(Path(temp_checkpoint_dir, self._filename).as_posix())
+            yield Checkpoint.from_directory(temp_checkpoint_dir)
+
+    def _save_and_report_checkpoint(self, report_dict: Dict, model: Booster):
+        with self._get_checkpoint(model=model) as checkpoint:
+            ray.tune.report(report_dict, checkpoint=checkpoint)
+
+    def _report_metrics(self, report_dict: Dict):
+        ray.tune.report(report_dict)
+
+
+@Deprecated
+class TuneReportCallback:
+    def __new__(cls: type, *args, **kwargs):
+        # TODO(justinvyu): [code_removal] Remove in 2.11.
+        raise DeprecationWarning(
+            "`TuneReportCallback` is deprecated. "
+            "Use `ray.tune.integration.lightgbm.TuneReportCheckpointCallback` instead."
+        )

python/ray/tune/integration/pytorch_lightning.py

@@ -0,0 +1,206 @@
+import inspect
+import logging
+import os
+import tempfile
+import warnings
+from contextlib import contextmanager
+from typing import Dict, List, Optional, Type, Union
+
+import ray.tune
+from ray.tune import Checkpoint
+from ray.util import log_once
+from ray.util.annotations import Deprecated, PublicAPI
+
+try:
+    from lightning.pytorch import Callback, LightningModule, Trainer
+except ModuleNotFoundError:
+    from pytorch_lightning import Callback, LightningModule, Trainer
+
+
+logger = logging.getLogger(__name__)
+
+# Get all Pytorch Lightning Callback hooks based on whatever PTL version is being used.
+_allowed_hooks = {
+    name
+    for name, fn in inspect.getmembers(Callback, predicate=inspect.isfunction)
+    if name.startswith("on_")
+}
+
+
+def _override_ptl_hooks(callback_cls: Type["TuneCallback"]) -> Type["TuneCallback"]:
+    """Overrides all allowed PTL Callback hooks with our custom handle logic."""
+
+    def generate_overridden_hook(fn_name):
+        def overridden_hook(
+            self,
+            trainer: Trainer,
+            *args,
+            pl_module: Optional[LightningModule] = None,
+            **kwargs,
+        ):
+            if fn_name in self._on:
+                self._handle(trainer=trainer, pl_module=pl_module)
+
+        return overridden_hook
+
+    # Set the overridden hook to all the allowed hooks in TuneCallback.
+    for fn_name in _allowed_hooks:
+        setattr(callback_cls, fn_name, generate_overridden_hook(fn_name))
+
+    return callback_cls
+
+
+@_override_ptl_hooks
+class TuneCallback(Callback):
+    """Base class for Tune's PyTorch Lightning callbacks.
+
+    Args:
+        on: When to trigger checkpoint creations. Must be one of
+            the PyTorch Lightning event hooks (less the ``on_``), e.g.
+            "train_batch_start", or "train_end". Defaults to "validation_end"
+    """
+
+    def __init__(self, on: Union[str, List[str]] = "validation_end"):
+        if not isinstance(on, list):
+            on = [on]
+
+        for hook in on:
+            if f"on_{hook}" not in _allowed_hooks:
+                raise ValueError(
+                    f"Invalid hook selected: {hook}. Must be one of "
+                    f"{_allowed_hooks}"
+                )
+
+        # Add back the "on_" prefix for internal consistency.
+        on = [f"on_{hook}" for hook in on]
+
+        self._on = on
+
+    def _handle(self, trainer: Trainer, pl_module: Optional[LightningModule]):
+        raise NotImplementedError
+
+
+@PublicAPI
+class TuneReportCheckpointCallback(TuneCallback):
+    """PyTorch Lightning report and checkpoint callback
+
+    Saves checkpoints after each validation step. Also reports metrics to Tune,
+    which is needed for checkpoint registration.
+
+    Args:
+        metrics: Metrics to report to Tune. If this is a list,
+            each item describes the metric key reported to PyTorch Lightning,
+            and it will reported under the same name to Tune. If this is a
+            dict, each key will be the name reported to Tune and the respective
+            value will be the metric key reported to PyTorch Lightning.
+        filename: Filename of the checkpoint within the checkpoint
+            directory. Defaults to "checkpoint".
+        save_checkpoints: If True (default), checkpoints will be saved and
+            reported to Ray. If False, only metrics will be reported.
+        on: When to trigger checkpoint creations and metric reports. Must be one of
+            the PyTorch Lightning event hooks (less the ``on_``), e.g.
+            "train_batch_start", or "train_end". Defaults to "validation_end".
+
+
+    Example:
+
+    .. code-block:: python
+
+        import lightning.pytorch as pl
+        from ray.tune.integration.pytorch_lightning import (
+            TuneReportCheckpointCallback)
+
+        # Save checkpoint after each training batch and after each
+        # validation epoch.
+        trainer = pl.Trainer(callbacks=[TuneReportCheckpointCallback(
+            metrics={"loss": "val_loss", "mean_accuracy": "val_acc"},
+            filename="trainer.ckpt", on="validation_end")])
+
+
+    """
+
+    def __init__(
+        self,
+        metrics: Optional[Union[str, List[str], Dict[str, str]]] = None,
+        filename: str = "checkpoint",
+        save_checkpoints: bool = True,
+        on: Union[str, List[str]] = "validation_end",
+    ):
+        super(TuneReportCheckpointCallback, self).__init__(on=on)
+        if isinstance(metrics, str):
+            metrics = [metrics]
+        self._save_checkpoints = save_checkpoints
+        self._filename = filename
+        self._metrics = metrics
+
+    def _get_report_dict(self, trainer: Trainer, pl_module: LightningModule):
+        # Don't report if just doing initial validation sanity checks.
+        if trainer.sanity_checking:
+            return
+        if not self._metrics:
+            report_dict = {k: v.item() for k, v in trainer.callback_metrics.items()}
+        else:
+            report_dict = {}
+            for key in self._metrics:
+                if isinstance(self._metrics, dict):
+                    metric = self._metrics[key]
+                else:
+                    metric = key
+                if metric in trainer.callback_metrics:
+                    report_dict[key] = trainer.callback_metrics[metric].item()
+                else:
+                    logger.warning(
+                        f"Metric {metric} does not exist in "
+                        "`trainer.callback_metrics."
+                    )
+
+        return report_dict
+
+    @contextmanager
+    def _get_checkpoint(self, trainer: Trainer) -> Optional[Checkpoint]:
+        if not self._save_checkpoints:
+            yield None
+            return
+
+        with tempfile.TemporaryDirectory() as checkpoint_dir:
+            trainer.save_checkpoint(os.path.join(checkpoint_dir, self._filename))
+            checkpoint = Checkpoint.from_directory(checkpoint_dir)
+            yield checkpoint
+
+    def _handle(self, trainer: Trainer, pl_module: LightningModule):
+        if trainer.sanity_checking:
+            return
+
+        report_dict = self._get_report_dict(trainer, pl_module)
+        if not report_dict:
+            return
+
+        with self._get_checkpoint(trainer) as checkpoint:
+            ray.tune.report(report_dict, checkpoint=checkpoint)
+
+
+class _TuneCheckpointCallback(TuneCallback):
+    def __init__(self, *args, **kwargs):
+        raise DeprecationWarning(
+            "`ray.tune.integration.pytorch_lightning._TuneCheckpointCallback` "
+            "is deprecated."
+        )
+
+
+@Deprecated
+class TuneReportCallback(TuneReportCheckpointCallback):
+    def __init__(
+        self,
+        metrics: Optional[Union[str, List[str], Dict[str, str]]] = None,
+        on: Union[str, List[str]] = "validation_end",
+    ):
+        if log_once("tune_ptl_report_deprecated"):
+            warnings.warn(
+                "`ray.tune.integration.pytorch_lightning.TuneReportCallback` "
+                "is deprecated. Use "
+                "`ray.tune.integration.pytorch_lightning.TuneReportCheckpointCallback`"
+                " instead."
+            )
+        super(TuneReportCallback, self).__init__(
+            metrics=metrics, save_checkpoints=False, on=on
+        )

python/ray/tune/integration/ray_train.py

@@ -0,0 +1,40 @@
+from typing import Any, Dict, List, Optional
+
+from ray.train import Checkpoint as RayTrainCheckpoint
+from ray.train._internal.session import get_session
+from ray.train.v2._internal.execution.context import TrainRunContext
+from ray.train.v2.api.callback import UserCallback
+from ray.tune.trainable.trainable_fn_utils import _in_tune_session
+from ray.util.annotations import DeveloperAPI
+
+CHECKPOINT_PATH_KEY = "checkpoint_path"
+
+
+@DeveloperAPI
+class TuneReportCallback(UserCallback):
+    """Propagate metrics and checkpoint paths from Ray Train workers to Ray Tune."""
+
+    def __init__(self):
+        if not _in_tune_session():
+            raise RuntimeError("TuneReportCallback must be used in a Tune session.")
+        self._training_actor_item_queue = (
+            get_session()._get_or_create_inter_actor_queue()
+        )
+
+    def after_report(
+        self,
+        run_context: TrainRunContext,
+        metrics: List[Dict[str, Any]],
+        checkpoint: Optional[RayTrainCheckpoint],
+    ):
+        # TODO: This can be changed to aggregate the metrics from all workers.
+        # For now, just achieve feature parity with the old Tune+Train integration.
+        metrics = metrics[0].copy()
+
+        # If a checkpoint is provided, add the checkpoint path to the metrics.
+        # Don't report the checkpoint again since it's already been uploaded
+        # to storage.
+        if checkpoint:
+            metrics[CHECKPOINT_PATH_KEY] = checkpoint.path
+
+        self._training_actor_item_queue.put(metrics)

python/ray/tune/integration/xgboost.py

@@ -0,0 +1,101 @@
+import tempfile
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Callable, Dict, List, Optional, Union
+
+from xgboost.core import Booster
+
+import ray.tune
+from ray.train.xgboost._xgboost_utils import RayReportCallback
+from ray.tune import Checkpoint
+from ray.util.annotations import Deprecated, PublicAPI
+
+
+@PublicAPI(stability="beta")
+class TuneReportCheckpointCallback(RayReportCallback):
+    """XGBoost callback to save checkpoints and report metrics for Ray Tune.
+
+    Args:
+        metrics: Metrics to report. If this is a list,
+            each item describes the metric key reported to XGBoost,
+            and it will be reported under the same name.
+            This can also be a dict of {<key-to-report>: <xgboost-metric-key>},
+            which can be used to rename xgboost default metrics.
+        filename: Customize the saved checkpoint file type by passing
+            a filename. Defaults to "model.ubj".
+        frequency: How often to save checkpoints, in terms of iterations.
+            Defaults to 0 (no checkpoints are saved during training).
+        checkpoint_at_end: Whether or not to save a checkpoint at the end of training.
+        results_postprocessing_fn: An optional Callable that takes in
+            the metrics dict that will be reported (after it has been flattened)
+            and returns a modified dict. For example, this can be used to
+            average results across CV fold when using ``xgboost.cv``.
+
+    Examples:
+
+        Reporting checkpoints and metrics to Ray Tune when running many
+        independent xgboost trials (without data parallelism within a trial).
+
+        .. testcode::
+            :skipif: True
+
+            import xgboost
+
+            from ray.tune import Tuner
+            from ray.tune.integration.xgboost import TuneReportCheckpointCallback
+
+            def train_fn(config):
+                # Report log loss to Ray Tune after each validation epoch.
+                bst = xgboost.train(
+                    ...,
+                    callbacks=[
+                        TuneReportCheckpointCallback(
+                            metrics={"loss": "eval-logloss"}, frequency=1
+                        )
+                    ],
+                )
+
+            tuner = Tuner(train_fn)
+            results = tuner.fit()
+    """
+
+    def __init__(
+        self,
+        metrics: Optional[Union[str, List[str], Dict[str, str]]] = None,
+        filename: str = RayReportCallback.CHECKPOINT_NAME,
+        frequency: int = 0,
+        checkpoint_at_end: bool = True,
+        results_postprocessing_fn: Optional[
+            Callable[[Dict[str, Union[float, List[float]]]], Dict[str, float]]
+        ] = None,
+    ):
+        super().__init__(
+            metrics=metrics,
+            filename=filename,
+            frequency=frequency,
+            checkpoint_at_end=checkpoint_at_end,
+            results_postprocessing_fn=results_postprocessing_fn,
+        )
+
+    @contextmanager
+    def _get_checkpoint(self, model: Booster) -> Optional[Checkpoint]:
+        with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
+            model.save_model(Path(temp_checkpoint_dir, self._filename).as_posix())
+            yield Checkpoint(temp_checkpoint_dir)
+
+    def _save_and_report_checkpoint(self, report_dict: Dict, model: Booster):
+        with self._get_checkpoint(model=model) as checkpoint:
+            ray.tune.report(report_dict, checkpoint=checkpoint)
+
+    def _report_metrics(self, report_dict: Dict):
+        ray.tune.report(report_dict)
+
+
+@Deprecated
+class TuneReportCallback:
+    def __new__(cls: type, *args, **kwargs):
+        # TODO(justinvyu): [code_removal] Remove in 2.11.
+        raise DeprecationWarning(
+            "`TuneReportCallback` is deprecated. "
+            "Use `ray.tune.integration.xgboost.TuneReportCheckpointCallback` instead."
+        )

python/ray/tune/logger/__init__.py

@@ -0,0 +1,15 @@
+from ray.tune.logger.csv import CSVLoggerCallback
+from ray.tune.logger.json import JsonLoggerCallback
+from ray.tune.logger.logger import (
+    LoggerCallback,
+    pretty_print,
+)
+from ray.tune.logger.tensorboardx import TBXLoggerCallback
+
+__all__ = [
+    "LoggerCallback",
+    "pretty_print",
+    "CSVLoggerCallback",
+    "JsonLoggerCallback",
+    "TBXLoggerCallback",
+]

python/ray/tune/logger/aim.py

@@ -0,0 +1,185 @@
+import logging
+from typing import TYPE_CHECKING, Dict, List, Optional, Union
+
+import numpy as np
+
+from ray.air.constants import TRAINING_ITERATION
+from ray.tune.logger.logger import LoggerCallback
+from ray.tune.result import TIME_TOTAL_S, TIMESTEPS_TOTAL
+from ray.tune.utils import flatten_dict
+from ray.util.annotations import PublicAPI
+
+if TYPE_CHECKING:
+    from ray.tune.experiment.trial import Trial
+
+try:
+    from aim.sdk import Repo, Run
+except ImportError:
+    Repo, Run = None, None
+
+logger = logging.getLogger(__name__)
+
+VALID_SUMMARY_TYPES = [int, float, np.float32, np.float64, np.int32, np.int64]
+
+
+@PublicAPI
+class AimLoggerCallback(LoggerCallback):
+    """Aim Logger: logs metrics in Aim format.
+
+    Aim is an open-source, self-hosted ML experiment tracking tool.
+    It's good at tracking lots (thousands) of training runs, and it allows you to
+    compare them with a performant and well-designed UI.
+
+    Source: https://github.com/aimhubio/aim
+    """
+
+    VALID_HPARAMS = (str, bool, int, float, list, type(None))
+    VALID_NP_HPARAMS = (np.bool_, np.float32, np.float64, np.int32, np.int64)
+
+    def __init__(
+        self,
+        repo: Optional[Union[str, "Repo"]] = None,
+        experiment_name: Optional[str] = None,
+        metrics: Optional[List[str]] = None,
+        **aim_run_kwargs,
+    ):
+        """Initialize the Aim logger callback.
+
+        Args:
+            repo: Aim repository directory or a `Repo` object that the Run object
+                will log results to. If not provided, a default repo will be set
+                up in the experiment directory (one level above trial directories).
+            experiment_name: Sets the `experiment` property of each Run object,
+                which is the experiment name associated with it. Can be used
+                later to query runs/sequences. If not provided, the default
+                will be the Tune experiment name set by `RunConfig(name=...)`.
+            metrics: List of metric names (out of the metrics reported by Tune)
+                to track in Aim. If no metric are specified, log everything
+                that is reported.
+            **aim_run_kwargs: Additional arguments that will be passed when
+                creating the individual `Run` objects for each trial. For the
+                full list of arguments, please see the Aim documentation:
+                https://aimstack.readthedocs.io/en/latest/refs/sdk.html
+        """
+        assert Run is not None, (
+            "aim must be installed!. You can install aim with"
+            " the command: `pip install aim`."
+        )
+        self._repo_path = repo
+        self._experiment_name = experiment_name
+        if not (bool(metrics) or metrics is None):
+            raise ValueError(
+                "`metrics` must either contain at least one metric name, or be None, "
+                "in which case all reported metrics will be logged to the aim repo."
+            )
+        self._metrics = metrics
+        self._aim_run_kwargs = aim_run_kwargs
+        self._trial_to_run: Dict["Trial", Run] = {}
+
+    def _create_run(self, trial: "Trial") -> Run:
+        """Initializes an Aim Run object for a given trial.
+
+        Args:
+            trial: The Tune trial that aim will track as a Run.
+
+        Returns:
+            Run: The created aim run for a specific trial.
+        """
+        experiment_dir = trial.local_experiment_path
+        run = Run(
+            repo=self._repo_path or experiment_dir,
+            experiment=self._experiment_name or trial.experiment_dir_name,
+            **self._aim_run_kwargs,
+        )
+        # Attach a few useful trial properties
+        run["trial_id"] = trial.trial_id
+        run["trial_log_dir"] = trial.path
+        trial_ip = trial.get_ray_actor_ip()
+        if trial_ip:
+            run["trial_ip"] = trial_ip
+        return run
+
+    def log_trial_start(self, trial: "Trial"):
+        if trial in self._trial_to_run:
+            # Cleanup an existing run if the trial has been restarted
+            self._trial_to_run[trial].close()
+
+        trial.init_local_path()
+        self._trial_to_run[trial] = self._create_run(trial)
+
+        if trial.evaluated_params:
+            self._log_trial_hparams(trial)
+
+    def log_trial_result(self, iteration: int, trial: "Trial", result: Dict):
+        tmp_result = result.copy()
+
+        step = result.get(TIMESTEPS_TOTAL, None) or result[TRAINING_ITERATION]
+
+        for k in ["config", "pid", "timestamp", TIME_TOTAL_S, TRAINING_ITERATION]:
+            tmp_result.pop(k, None)  # not useful to log these
+
+        # `context` and `epoch` are special keys that users can report,
+        # which are treated as special aim metrics/configurations.
+        context = tmp_result.pop("context", None)
+        epoch = tmp_result.pop("epoch", None)
+
+        trial_run = self._trial_to_run[trial]
+        path = ["ray", "tune"]
+
+        flat_result = flatten_dict(tmp_result, delimiter="/")
+        valid_result = {}
+
+        for attr, value in flat_result.items():
+            if self._metrics and attr not in self._metrics:
+                continue
+
+            full_attr = "/".join(path + [attr])
+            if isinstance(value, tuple(VALID_SUMMARY_TYPES)) and not (
+                np.isnan(value) or np.isinf(value)
+            ):
+                valid_result[attr] = value
+                trial_run.track(
+                    value=value,
+                    name=full_attr,
+                    epoch=epoch,
+                    step=step,
+                    context=context,
+                )
+            elif (isinstance(value, (list, tuple, set)) and len(value) > 0) or (
+                isinstance(value, np.ndarray) and value.size > 0
+            ):
+                valid_result[attr] = value
+
+    def log_trial_end(self, trial: "Trial", failed: bool = False):
+        trial_run = self._trial_to_run.pop(trial)
+        trial_run.close()
+
+    def _log_trial_hparams(self, trial: "Trial"):
+        params = flatten_dict(trial.evaluated_params, delimiter="/")
+        flat_params = flatten_dict(params)
+
+        scrubbed_params = {
+            k: v for k, v in flat_params.items() if isinstance(v, self.VALID_HPARAMS)
+        }
+
+        np_params = {
+            k: v.tolist()
+            for k, v in flat_params.items()
+            if isinstance(v, self.VALID_NP_HPARAMS)
+        }
+
+        scrubbed_params.update(np_params)
+        removed = {
+            k: v
+            for k, v in flat_params.items()
+            if not isinstance(v, self.VALID_HPARAMS + self.VALID_NP_HPARAMS)
+        }
+        if removed:
+            logger.info(
+                "Removed the following hyperparameter values when "
+                "logging to aim: %s",
+                str(removed),
+            )
+
+        run = self._trial_to_run[trial]
+        run["hparams"] = scrubbed_params

python/ray/tune/logger/comet.py

@@ -0,0 +1,3 @@
+from ray.air.integrations.comet import CometLoggerCallback
+
+CometLoggerCallback.__module__ = "ray.tune.logger.comet"

python/ray/tune/logger/csv.py

@@ -0,0 +1,79 @@
+import csv
+import logging
+from pathlib import Path
+from typing import TYPE_CHECKING, Dict, TextIO
+
+from ray.air.constants import EXPR_PROGRESS_FILE
+from ray.tune.logger.logger import LoggerCallback
+from ray.tune.utils import flatten_dict
+from ray.util.annotations import PublicAPI
+
+if TYPE_CHECKING:
+    from ray.tune.experiment.trial import Trial  # noqa: F401
+
+logger = logging.getLogger(__name__)
+
+
+@PublicAPI
+class CSVLoggerCallback(LoggerCallback):
+    """Logs results to progress.csv under the trial directory.
+
+    Automatically flattens nested dicts in the result dict before writing
+    to csv:
+
+        {"a": {"b": 1, "c": 2}} -> {"a/b": 1, "a/c": 2}
+
+    """
+
+    _SAVED_FILE_TEMPLATES = [EXPR_PROGRESS_FILE]
+
+    def __init__(self):
+        self._trial_continue: Dict["Trial", bool] = {}
+        self._trial_files: Dict["Trial", TextIO] = {}
+        self._trial_csv: Dict["Trial", csv.DictWriter] = {}
+
+    def _setup_trial(self, trial: "Trial"):
+        if trial in self._trial_files:
+            self._trial_files[trial].close()
+
+        # Make sure logdir exists
+        trial.init_local_path()
+        local_file_path = Path(trial.local_path, EXPR_PROGRESS_FILE)
+
+        # Resume the file from remote storage.
+        self._restore_from_remote(EXPR_PROGRESS_FILE, trial)
+
+        self._trial_continue[trial] = (
+            local_file_path.exists() and local_file_path.stat().st_size > 0
+        )
+
+        self._trial_files[trial] = local_file_path.open("at")
+        self._trial_csv[trial] = None
+
+    def log_trial_result(self, iteration: int, trial: "Trial", result: Dict):
+        if trial not in self._trial_files:
+            self._setup_trial(trial)
+
+        tmp = result.copy()
+        tmp.pop("config", None)
+        result = flatten_dict(tmp, delimiter="/")
+
+        if not self._trial_csv[trial]:
+            self._trial_csv[trial] = csv.DictWriter(
+                self._trial_files[trial], result.keys()
+            )
+            if not self._trial_continue[trial]:
+                self._trial_csv[trial].writeheader()
+
+        self._trial_csv[trial].writerow(
+            {k: v for k, v in result.items() if k in self._trial_csv[trial].fieldnames}
+        )
+        self._trial_files[trial].flush()
+
+    def log_trial_end(self, trial: "Trial", failed: bool = False):
+        if trial not in self._trial_files:
+            return
+
+        del self._trial_csv[trial]
+        self._trial_files[trial].close()
+        del self._trial_files[trial]

python/ray/tune/logger/json.py

@@ -0,0 +1,78 @@
+import json
+import logging
+from pathlib import Path
+from typing import TYPE_CHECKING, Dict, TextIO
+
+import ray.cloudpickle as cloudpickle
+from ray.air.constants import EXPR_PARAM_FILE, EXPR_PARAM_PICKLE_FILE, EXPR_RESULT_FILE
+from ray.tune.logger.logger import LoggerCallback
+from ray.tune.utils.util import SafeFallbackEncoder
+from ray.util.annotations import PublicAPI
+
+if TYPE_CHECKING:
+    from ray.tune.experiment.trial import Trial  # noqa: F401
+
+logger = logging.getLogger(__name__)
+
+
+@PublicAPI
+class JsonLoggerCallback(LoggerCallback):
+    """Logs trial results in json format.
+
+    Also writes to a results file and param.json file when results or
+    configurations are updated. Experiments must be executed with the
+    JsonLoggerCallback to be compatible with the ExperimentAnalysis tool.
+    """
+
+    _SAVED_FILE_TEMPLATES = [EXPR_RESULT_FILE, EXPR_PARAM_FILE, EXPR_PARAM_PICKLE_FILE]
+
+    def __init__(self):
+        self._trial_configs: Dict["Trial", Dict] = {}
+        self._trial_files: Dict["Trial", TextIO] = {}
+
+    def log_trial_start(self, trial: "Trial"):
+        if trial in self._trial_files:
+            self._trial_files[trial].close()
+
+        # Update config
+        self.update_config(trial, trial.config)
+
+        # Make sure logdir exists
+        trial.init_local_path()
+        local_file = Path(trial.local_path, EXPR_RESULT_FILE)
+
+        # Resume the file from remote storage.
+        self._restore_from_remote(EXPR_RESULT_FILE, trial)
+
+        self._trial_files[trial] = local_file.open("at")
+
+    def log_trial_result(self, iteration: int, trial: "Trial", result: Dict):
+        if trial not in self._trial_files:
+            self.log_trial_start(trial)
+        json.dump(result, self._trial_files[trial], cls=SafeFallbackEncoder)
+        self._trial_files[trial].write("\n")
+        self._trial_files[trial].flush()
+
+    def log_trial_end(self, trial: "Trial", failed: bool = False):
+        if trial not in self._trial_files:
+            return
+
+        self._trial_files[trial].close()
+        del self._trial_files[trial]
+
+    def update_config(self, trial: "Trial", config: Dict):
+        self._trial_configs[trial] = config
+
+        config_out = Path(trial.local_path, EXPR_PARAM_FILE)
+        with config_out.open("w") as f:
+            json.dump(
+                self._trial_configs[trial],
+                f,
+                indent=2,
+                sort_keys=True,
+                cls=SafeFallbackEncoder,
+            )
+
+        config_pkl = Path(trial.local_path, EXPR_PARAM_PICKLE_FILE)
+        with config_pkl.open("wb") as f:
+            cloudpickle.dump(self._trial_configs[trial], f)

python/ray/tune/logger/logger.py

@@ -0,0 +1,155 @@
+import json
+import logging
+from pathlib import Path
+from typing import TYPE_CHECKING, Dict, List, Optional, Set
+
+import pyarrow
+import yaml
+
+from ray.air._internal.json import SafeFallbackEncoder
+from ray.tune.callback import Callback
+from ray.util.annotations import DeveloperAPI, PublicAPI
+
+if TYPE_CHECKING:
+    from ray.tune.experiment.trial import Trial  # noqa: F401
+
+logger = logging.getLogger(__name__)
+
+
+# Apply flow style for sequences of this length
+_SEQUENCE_LEN_FLOW_STYLE = 3
+
+
+@PublicAPI
+class LoggerCallback(Callback):
+    """Base class for experiment-level logger callbacks
+
+    This base class defines a general interface for logging events,
+    like trial starts, restores, ends, checkpoint saves, and receiving
+    trial results.
+
+    Callbacks implementing this interface should make sure that logging
+    utilities are cleaned up properly on trial termination, i.e. when
+    ``log_trial_end`` is received. This includes e.g. closing files.
+    """
+
+    def log_trial_start(self, trial: "Trial"):
+        """Handle logging when a trial starts.
+
+        Args:
+            trial: Trial object.
+        """
+        pass
+
+    def log_trial_restore(self, trial: "Trial"):
+        """Handle logging when a trial restores.
+
+        Args:
+            trial: Trial object.
+        """
+        pass
+
+    def log_trial_save(self, trial: "Trial"):
+        """Handle logging when a trial saves a checkpoint.
+
+        Args:
+            trial: Trial object.
+        """
+        pass
+
+    def log_trial_result(self, iteration: int, trial: "Trial", result: Dict):
+        """Handle logging when a trial reports a result.
+
+        Args:
+            iteration: Iteration of the experiment that this result belongs to.
+            trial: Trial object.
+            result: Result dictionary.
+        """
+        pass
+
+    def log_trial_end(self, trial: "Trial", failed: bool = False):
+        """Handle logging when a trial ends.
+
+        Args:
+            trial: Trial object.
+            failed: True if the Trial finished gracefully, False if
+                it failed (e.g. when it raised an exception).
+        """
+        pass
+
+    def on_trial_result(
+        self,
+        iteration: int,
+        trials: List["Trial"],
+        trial: "Trial",
+        result: Dict,
+        **info,
+    ):
+        self.log_trial_result(iteration, trial, result)
+
+    def on_trial_start(
+        self, iteration: int, trials: List["Trial"], trial: "Trial", **info
+    ):
+        self.log_trial_start(trial)
+
+    def on_trial_restore(
+        self, iteration: int, trials: List["Trial"], trial: "Trial", **info
+    ):
+        self.log_trial_restore(trial)
+
+    def on_trial_save(
+        self, iteration: int, trials: List["Trial"], trial: "Trial", **info
+    ):
+        self.log_trial_save(trial)
+
+    def on_trial_complete(
+        self, iteration: int, trials: List["Trial"], trial: "Trial", **info
+    ):
+        self.log_trial_end(trial, failed=False)
+
+    def on_trial_error(
+        self, iteration: int, trials: List["Trial"], trial: "Trial", **info
+    ):
+        self.log_trial_end(trial, failed=True)
+
+    def _restore_from_remote(self, file_name: str, trial: "Trial") -> None:
+        if not trial.checkpoint:
+            # If there's no checkpoint, there's no logging artifacts to restore
+            # since we're starting from scratch.
+            return
+
+        local_file = Path(trial.local_path, file_name).as_posix()
+        remote_file = Path(trial.storage.trial_fs_path, file_name).as_posix()
+
+        try:
+            pyarrow.fs.copy_files(
+                remote_file,
+                local_file,
+                source_filesystem=trial.storage.storage_filesystem,
+            )
+            logger.debug(f"Copied {remote_file} to {local_file}")
+        except FileNotFoundError:
+            logger.warning(f"Remote file not found: {remote_file}")
+        except Exception:
+            logger.exception(f"Error downloading {remote_file}")
+
+
+class _RayDumper(yaml.SafeDumper):
+    def represent_sequence(self, tag, sequence, flow_style=None):
+        if len(sequence) > _SEQUENCE_LEN_FLOW_STYLE:
+            return super().represent_sequence(tag, sequence, flow_style=True)
+        return super().represent_sequence(tag, sequence, flow_style=flow_style)
+
+
+@DeveloperAPI
+def pretty_print(result, exclude: Optional[Set[str]] = None):
+    result = result.copy()
+    result.update(config=None)  # drop config from pretty print
+    result.update(hist_stats=None)  # drop hist_stats from pretty print
+    out = {}
+    for k, v in result.items():
+        if v is not None and (exclude is None or k not in exclude):
+            out[k] = v
+
+    cleaned = json.dumps(out, cls=SafeFallbackEncoder)
+    return yaml.dump(json.loads(cleaned), Dumper=_RayDumper, default_flow_style=False)

python/ray/tune/logger/mlflow.py

@@ -0,0 +1,3 @@
+from ray.air.integrations.mlflow import MLflowLoggerCallback
+
+MLflowLoggerCallback.__module__ = "ray.tune.logger.mlflow"

python/ray/tune/logger/tensorboardx.py

@@ -0,0 +1,178 @@
+import logging
+from typing import TYPE_CHECKING, Dict
+
+import numpy as np
+
+from ray.air.constants import TRAINING_ITERATION
+from ray.tune.logger.logger import LoggerCallback
+from ray.tune.result import TIME_TOTAL_S, TIMESTEPS_TOTAL
+from ray.tune.utils import flatten_dict
+from ray.util.annotations import PublicAPI
+from ray.util.debug import log_once
+
+if TYPE_CHECKING:
+    from ray.tune.experiment.trial import Trial  # noqa: F401
+
+logger = logging.getLogger(__name__)
+
+VALID_SUMMARY_TYPES = [int, float, np.float32, np.float64, np.int32, np.int64]
+
+
+@PublicAPI
+class TBXLoggerCallback(LoggerCallback):
+    """TensorBoardX Logger.
+
+    Note that hparams will be written only after a trial has terminated.
+    This logger automatically flattens nested dicts to show on TensorBoard:
+
+        {"a": {"b": 1, "c": 2}} -> {"a/b": 1, "a/c": 2}
+    """
+
+    _SAVED_FILE_TEMPLATES = ["events.out.tfevents.*"]
+
+    VALID_HPARAMS = (str, bool, int, float, list, type(None))
+    VALID_NP_HPARAMS = (np.bool_, np.float32, np.float64, np.int32, np.int64)
+
+    def __init__(self):
+        try:
+            from tensorboardX import SummaryWriter
+
+            self._summary_writer_cls = SummaryWriter
+        except ImportError:
+            if log_once("tbx-install"):
+                logger.info('pip install "ray[tune]" to see TensorBoard files.')
+            raise
+        self._trial_writer: Dict["Trial", SummaryWriter] = {}
+        self._trial_result: Dict["Trial", Dict] = {}
+
+    def log_trial_start(self, trial: "Trial"):
+        if trial in self._trial_writer:
+            self._trial_writer[trial].close()
+        trial.init_local_path()
+        self._trial_writer[trial] = self._summary_writer_cls(
+            trial.local_path, flush_secs=30
+        )
+        self._trial_result[trial] = {}
+
+    def log_trial_result(self, iteration: int, trial: "Trial", result: Dict):
+        if trial not in self._trial_writer:
+            self.log_trial_start(trial)
+
+        step = result.get(TIMESTEPS_TOTAL) or result[TRAINING_ITERATION]
+
+        tmp = result.copy()
+        for k in ["config", "pid", "timestamp", TIME_TOTAL_S, TRAINING_ITERATION]:
+            if k in tmp:
+                del tmp[k]  # not useful to log these
+
+        flat_result = flatten_dict(tmp, delimiter="/")
+        path = ["ray", "tune"]
+        valid_result = {}
+
+        for attr, value in flat_result.items():
+            full_attr = "/".join(path + [attr])
+            if isinstance(value, tuple(VALID_SUMMARY_TYPES)) and not np.isnan(value):
+                valid_result[full_attr] = value
+                self._trial_writer[trial].add_scalar(full_attr, value, global_step=step)
+            elif (isinstance(value, list) and len(value) > 0) or (
+                isinstance(value, np.ndarray) and value.size > 0
+            ):
+                valid_result[full_attr] = value
+
+                # Must be a single image.
+                if isinstance(value, np.ndarray) and value.ndim == 3:
+                    self._trial_writer[trial].add_image(
+                        full_attr,
+                        value,
+                        global_step=step,
+                    )
+                    continue
+
+                # Must be a batch of images.
+                if isinstance(value, np.ndarray) and value.ndim == 4:
+                    self._trial_writer[trial].add_images(
+                        full_attr,
+                        value,
+                        global_step=step,
+                    )
+                    continue
+
+                # Must be video
+                if isinstance(value, np.ndarray) and value.ndim == 5:
+                    self._trial_writer[trial].add_video(
+                        full_attr, value, global_step=step, fps=20
+                    )
+                    continue
+
+                try:
+                    self._trial_writer[trial].add_histogram(
+                        full_attr, value, global_step=step
+                    )
+                # In case TensorboardX still doesn't think it's a valid value
+                # (e.g. `[[]]`), warn and move on.
+                except (ValueError, TypeError):
+                    if log_once("invalid_tbx_value"):
+                        logger.warning(
+                            "You are trying to log an invalid value ({}={}) "
+                            "via {}!".format(full_attr, value, type(self).__name__)
+                        )
+
+        self._trial_result[trial] = valid_result
+        self._trial_writer[trial].flush()
+
+    def log_trial_end(self, trial: "Trial", failed: bool = False):
+        if trial in self._trial_writer:
+            if trial and trial.evaluated_params and self._trial_result[trial]:
+                flat_result = flatten_dict(self._trial_result[trial], delimiter="/")
+                scrubbed_result = {
+                    k: value
+                    for k, value in flat_result.items()
+                    if isinstance(value, tuple(VALID_SUMMARY_TYPES))
+                }
+                self._try_log_hparams(trial, scrubbed_result)
+            self._trial_writer[trial].close()
+            del self._trial_writer[trial]
+            del self._trial_result[trial]
+
+    def _try_log_hparams(self, trial: "Trial", result: Dict):
+        # TBX currently errors if the hparams value is None.
+        flat_params = flatten_dict(trial.evaluated_params)
+        scrubbed_params = {
+            k: v for k, v in flat_params.items() if isinstance(v, self.VALID_HPARAMS)
+        }
+
+        np_params = {
+            k: v.tolist()
+            for k, v in flat_params.items()
+            if isinstance(v, self.VALID_NP_HPARAMS)
+        }
+
+        scrubbed_params.update(np_params)
+
+        removed = {
+            k: v
+            for k, v in flat_params.items()
+            if not isinstance(v, self.VALID_HPARAMS + self.VALID_NP_HPARAMS)
+        }
+        if removed:
+            logger.info(
+                "Removed the following hyperparameter values when "
+                "logging to tensorboard: %s",
+                str(removed),
+            )
+
+        from tensorboardX.summary import hparams
+
+        try:
+            experiment_tag, session_start_tag, session_end_tag = hparams(
+                hparam_dict=scrubbed_params, metric_dict=result
+            )
+            self._trial_writer[trial].file_writer.add_summary(experiment_tag)
+            self._trial_writer[trial].file_writer.add_summary(session_start_tag)
+            self._trial_writer[trial].file_writer.add_summary(session_end_tag)
+        except Exception:
+            logger.exception(
+                "TensorboardX failed to log hparams. "
+                "This may be due to an unsupported type "
+                "in the hyperparameter values."
+            )

python/ray/tune/logger/wandb.py

@@ -0,0 +1,3 @@
+from ray.air.integrations.wandb import WandbLoggerCallback
+
+WandbLoggerCallback.__module__ = "ray.tune.logger.wandb"