chore: import upstream snapshot with attribution

python/ray/tune/schedulers/__init__.py

@@ -0,0 +1,96 @@
+import inspect
+from typing import Any
+
+from ray._common.utils import get_function_args
+from ray.tune.schedulers.async_hyperband import ASHAScheduler, AsyncHyperBandScheduler
+from ray.tune.schedulers.hb_bohb import HyperBandForBOHB
+from ray.tune.schedulers.hyperband import HyperBandScheduler
+from ray.tune.schedulers.median_stopping_rule import MedianStoppingRule
+from ray.tune.schedulers.pbt import (
+    PopulationBasedTraining,
+    PopulationBasedTrainingReplay,
+)
+from ray.tune.schedulers.resource_changing_scheduler import ResourceChangingScheduler
+from ray.tune.schedulers.trial_scheduler import FIFOScheduler, TrialScheduler
+from ray.util import PublicAPI
+
+
+def _pb2_importer():
+    # PB2 is imported lazily since it has additional dependencies.
+    from ray.tune.schedulers.pb2 import PB2
+
+    return PB2
+
+
+# Values in this dictionary will be one two kinds:
+#    class of the scheduler object to create
+#    wrapper function to support a lazy import of the scheduler class
+SCHEDULER_IMPORT = {
+    "fifo": FIFOScheduler,
+    "async_hyperband": AsyncHyperBandScheduler,
+    "asynchyperband": AsyncHyperBandScheduler,
+    "median_stopping_rule": MedianStoppingRule,
+    "medianstopping": MedianStoppingRule,
+    "hyperband": HyperBandScheduler,
+    "hb_bohb": HyperBandForBOHB,
+    "pbt": PopulationBasedTraining,
+    "pbt_replay": PopulationBasedTrainingReplay,
+    "pb2": _pb2_importer,
+    "resource_changing": ResourceChangingScheduler,
+}
+
+
+@PublicAPI(stability="beta")
+def create_scheduler(
+    scheduler: str,
+    **kwargs: Any,
+) -> TrialScheduler:
+    """Instantiate a scheduler based on the given string.
+
+    This is useful for swapping between different schedulers.
+
+    Args:
+        scheduler: The scheduler to use.
+        **kwargs: Scheduler parameters.
+            These keyword arguments will be passed to the initialization
+            function of the chosen scheduler.
+    Returns:
+        ray.tune.schedulers.trial_scheduler.TrialScheduler: The scheduler.
+    Example:
+        >>> from ray import tune
+        >>> pbt_kwargs = {}
+        >>> scheduler = tune.create_scheduler('pbt', **pbt_kwargs) # doctest: +SKIP
+    """
+
+    scheduler = scheduler.lower()
+    if scheduler not in SCHEDULER_IMPORT:
+        raise ValueError(
+            f"The `scheduler` argument must be one of "
+            f"{list(SCHEDULER_IMPORT)}. "
+            f"Got: {scheduler}"
+        )
+
+    SchedulerClass = SCHEDULER_IMPORT[scheduler]
+
+    if inspect.isfunction(SchedulerClass):
+        # invoke the wrapper function to retrieve class
+        SchedulerClass = SchedulerClass()
+
+    scheduler_args = get_function_args(SchedulerClass)
+    trimmed_kwargs = {k: v for k, v in kwargs.items() if k in scheduler_args}
+
+    return SchedulerClass(**trimmed_kwargs)
+
+
+__all__ = [
+    "TrialScheduler",
+    "HyperBandScheduler",
+    "AsyncHyperBandScheduler",
+    "ASHAScheduler",
+    "MedianStoppingRule",
+    "FIFOScheduler",
+    "PopulationBasedTraining",
+    "PopulationBasedTrainingReplay",
+    "HyperBandForBOHB",
+    "ResourceChangingScheduler",
+]

python/ray/tune/schedulers/async_hyperband.py

@@ -0,0 +1,291 @@
+import logging
+import pickle
+from typing import TYPE_CHECKING, Dict, Optional, Union
+
+import numpy as np
+
+from ray.tune.experiment import Trial
+from ray.tune.result import DEFAULT_METRIC
+from ray.tune.schedulers.trial_scheduler import FIFOScheduler, TrialScheduler
+from ray.util import PublicAPI
+
+if TYPE_CHECKING:
+    from ray.tune.execution.tune_controller import TuneController
+
+logger = logging.getLogger(__name__)
+
+
+@PublicAPI
+class AsyncHyperBandScheduler(FIFOScheduler):
+    """Implements the Async Successive Halving.
+
+    This should provide similar theoretical performance as HyperBand but
+    avoid straggler issues that HyperBand faces. One implementation detail
+    is when using multiple brackets, trial allocation to bracket is done
+    randomly with over a softmax probability.
+
+    See https://arxiv.org/abs/1810.05934
+
+    Args:
+        time_attr: A training result attr to use for comparing time.
+            Note that you can pass in something non-temporal such as
+            `training_iteration` as a measure of progress, the only requirement
+            is that the attribute should increase monotonically.
+            Valid values are any key reported in the result dict by your
+            trainable. The auto-filled keys ``"training_iteration"`` (the
+            iteration count) and ``"time_total_s"`` (wall-clock seconds since
+            the trial started) always work; any additional numeric, monotonic
+            key your trainable reports via ``tune.report({...})`` is also valid
+            (for example ``"timesteps_total"`` or a custom progress counter).
+            Passing a key that is not present in the reported result causes
+            the scheduler to skip its decision for that step.
+        metric: The training result objective value attribute. Stopping
+            procedures will use this attribute. If None but a mode was passed,
+            the `ray.tune.result.DEFAULT_METRIC` will be used per default.
+        mode: One of {min, max}. Determines whether objective is
+            minimizing or maximizing the metric attribute.
+        max_t: max time units per trial. Trials will be stopped after
+            max_t time units (determined by time_attr) have passed.
+        grace_period: Only stop trials at least this old in time.
+            The units are the same as the attribute named by `time_attr`.
+        reduction_factor: Used to set halving rate and amount. This
+            is simply a unit-less scalar.
+        brackets: Number of brackets. Each bracket has a different
+            halving rate, specified by the reduction factor.
+        stop_last_trials: Whether to terminate the trials after
+            reaching max_t. Defaults to True.
+    """
+
+    def __init__(
+        self,
+        time_attr: str = "training_iteration",
+        metric: Optional[str] = None,
+        mode: Optional[str] = None,
+        max_t: int = 100,
+        grace_period: int = 1,
+        reduction_factor: float = 4,
+        brackets: int = 1,
+        stop_last_trials: bool = True,
+    ):
+        assert max_t > 0, "Max (time_attr) not valid!"
+        assert max_t >= grace_period, "grace_period must be <= max_t!"
+        assert grace_period > 0, "grace_period must be positive!"
+        assert reduction_factor > 1, "Reduction Factor not valid!"
+        assert brackets > 0, "brackets must be positive!"
+        if mode:
+            assert mode in ["min", "max"], "`mode` must be 'min' or 'max'!"
+
+        super().__init__()
+        self._reduction_factor = reduction_factor
+        self._max_t = max_t
+
+        self._trial_info = {}  # Stores Trial -> Bracket
+
+        # Tracks state for new trial add
+        self._brackets = [
+            _Bracket(
+                grace_period,
+                max_t,
+                reduction_factor,
+                s,
+                stop_last_trials=stop_last_trials,
+            )
+            for s in range(brackets)
+        ]
+        self._counter = 0  # for
+        self._num_stopped = 0
+        self._metric = metric
+        self._mode = mode
+        self._metric_op = None
+        if self._mode == "max":
+            self._metric_op = 1.0
+        elif self._mode == "min":
+            self._metric_op = -1.0
+        self._time_attr = time_attr
+        self._stop_last_trials = stop_last_trials
+
+    def set_search_properties(
+        self, metric: Optional[str], mode: Optional[str], **spec
+    ) -> bool:
+        if self._metric and metric:
+            return False
+        if self._mode and mode:
+            return False
+
+        if metric:
+            self._metric = metric
+        if mode:
+            self._mode = mode
+
+        if self._mode == "max":
+            self._metric_op = 1.0
+        elif self._mode == "min":
+            self._metric_op = -1.0
+
+        if self._metric is None and self._mode:
+            # If only a mode was passed, use anonymous metric
+            self._metric = DEFAULT_METRIC
+
+        return True
+
+    def on_trial_add(self, tune_controller: "TuneController", trial: Trial):
+        if not self._metric or not self._metric_op:
+            raise ValueError(
+                "{} has been instantiated without a valid `metric` ({}) or "
+                "`mode` ({}) parameter. Either pass these parameters when "
+                "instantiating the scheduler, or pass them as parameters "
+                "to `tune.TuneConfig()`".format(
+                    self.__class__.__name__, self._metric, self._mode
+                )
+            )
+
+        sizes = np.array([len(b._rungs) for b in self._brackets])
+        probs = np.e ** (sizes - sizes.max())
+        normalized = probs / probs.sum()
+        idx = np.random.choice(len(self._brackets), p=normalized)
+        self._trial_info[trial.trial_id] = self._brackets[idx]
+
+    def on_trial_result(
+        self, tune_controller: "TuneController", trial: Trial, result: Dict
+    ) -> str:
+        action = TrialScheduler.CONTINUE
+        if self._time_attr not in result or self._metric not in result:
+            return action
+        if result[self._time_attr] >= self._max_t and self._stop_last_trials:
+            action = TrialScheduler.STOP
+        else:
+            bracket = self._trial_info[trial.trial_id]
+            action = bracket.on_result(
+                trial, result[self._time_attr], self._metric_op * result[self._metric]
+            )
+        if action == TrialScheduler.STOP:
+            self._num_stopped += 1
+        return action
+
+    def on_trial_complete(
+        self, tune_controller: "TuneController", trial: Trial, result: Dict
+    ):
+        if self._time_attr not in result or self._metric not in result:
+            return
+        bracket = self._trial_info[trial.trial_id]
+        bracket.on_result(
+            trial, result[self._time_attr], self._metric_op * result[self._metric]
+        )
+        del self._trial_info[trial.trial_id]
+
+    def on_trial_remove(self, tune_controller: "TuneController", trial: Trial):
+        del self._trial_info[trial.trial_id]
+
+    def debug_string(self) -> str:
+        out = "Using AsyncHyperBand: num_stopped={}".format(self._num_stopped)
+        out += "\n" + "\n".join([b.debug_str() for b in self._brackets])
+        return out
+
+    def save(self, checkpoint_path: str):
+        save_object = self.__dict__
+        with open(checkpoint_path, "wb") as outputFile:
+            pickle.dump(save_object, outputFile)
+
+    def restore(self, checkpoint_path: str):
+        with open(checkpoint_path, "rb") as inputFile:
+            save_object = pickle.load(inputFile)
+        self.__dict__.update(save_object)
+
+
+class _Bracket:
+    """Bookkeeping system to track the cutoffs.
+
+    Rungs are created in reversed order so that we can more easily find
+    the correct rung corresponding to the current iteration of the result.
+
+    Example:
+        >>> trial1, trial2, trial3 = ... # doctest: +SKIP
+        >>> b = _Bracket(1, 10, 2, 0) # doctest: +SKIP
+        >>> # CONTINUE
+        >>> b.on_result(trial1, 1, 2) # doctest: +SKIP
+        >>> # CONTINUE
+        >>> b.on_result(trial2, 1, 4) # doctest: +SKIP
+        >>> # rungs are reversed
+        >>> b.cutoff(b._rungs[-1][1]) == 3.0 # doctest: +SKIP
+         # STOP
+        >>> b.on_result(trial3, 1, 1) # doctest: +SKIP
+        >>> b.cutoff(b._rungs[3][1]) == 2.0 # doctest: +SKIP
+    """
+
+    def __init__(
+        self,
+        min_t: int,
+        max_t: int,
+        reduction_factor: float,
+        s: int,
+        stop_last_trials: bool = True,
+    ):
+        """Initialize a bracket of the asynchronous successive halving algorithm.
+
+        Args:
+            min_t: Minimum number of iterations before a trial can be stopped.
+            max_t: Maximum number of iterations per trial.
+            reduction_factor: Halving rate used to derive rung spacing.
+            s: Bracket index, used to offset the first rung.
+            stop_last_trials: If True, allow trials that survive the final rung
+                to still be stopped by the bracket.
+        """
+        self.rf = reduction_factor
+        MAX_RUNGS = int(np.log(max_t / min_t) / np.log(self.rf) - s + 1)
+        self._rungs = [
+            (min_t * self.rf ** (k + s), {}) for k in reversed(range(MAX_RUNGS))
+        ]
+        self._stop_last_trials = stop_last_trials
+
+    def cutoff(self, recorded) -> Optional[Union[int, float, complex, np.ndarray]]:
+        if not recorded:
+            return None
+        return np.nanpercentile(list(recorded.values()), (1 - 1 / self.rf) * 100)
+
+    def on_result(self, trial: Trial, cur_iter: int, cur_rew: Optional[float]) -> str:
+        action = TrialScheduler.CONTINUE
+        for milestone, recorded in self._rungs:
+            if (
+                cur_iter >= milestone
+                and trial.trial_id in recorded
+                and not self._stop_last_trials
+            ):
+                # If our result has been recorded for this trial already, the
+                # decision to continue training has already been made. Thus we can
+                # skip new cutoff calculation and just continue training.
+                # We can also break as milestones are descending.
+                break
+            if cur_iter < milestone or trial.trial_id in recorded:
+                continue
+            else:
+                cutoff = self.cutoff(recorded)
+                if cutoff is not None and cur_rew < cutoff:
+                    action = TrialScheduler.STOP
+                if cur_rew is None:
+                    logger.warning(
+                        "Reward attribute is None! Consider"
+                        " reporting using a different field."
+                    )
+                else:
+                    recorded[trial.trial_id] = cur_rew
+                break
+        return action
+
+    def debug_str(self) -> str:
+        # TODO: fix up the output for this
+        iters = " | ".join(
+            [
+                "Iter {:.3f}: {}".format(milestone, self.cutoff(recorded))
+                for milestone, recorded in self._rungs
+            ]
+        )
+        return "Bracket: " + iters
+
+
+ASHAScheduler = AsyncHyperBandScheduler
+
+if __name__ == "__main__":
+    sched = AsyncHyperBandScheduler(grace_period=1, max_t=10, reduction_factor=2)
+    print(sched.debug_string())
+    bracket = sched._brackets[0]
+    print(bracket.cutoff({str(i): i for i in range(20)}))

python/ray/tune/schedulers/hb_bohb.py

@@ -0,0 +1,176 @@
+import logging
+from typing import TYPE_CHECKING, Dict, Optional
+
+from ray.tune.experiment import Trial
+from ray.tune.schedulers.hyperband import HyperBandScheduler
+from ray.tune.schedulers.trial_scheduler import TrialScheduler
+from ray.util import PublicAPI
+
+if TYPE_CHECKING:
+    from ray.tune.execution.tune_controller import TuneController
+
+logger = logging.getLogger(__name__)
+
+
+@PublicAPI
+class HyperBandForBOHB(HyperBandScheduler):
+    """Extends HyperBand early stopping algorithm for BOHB.
+
+    This implementation removes the ``HyperBandScheduler`` pipelining. This
+    class introduces key changes:
+
+    1. Trials are now placed so that the bracket with the largest size is
+    filled first.
+
+    2. Trials will be paused even if the bracket is not filled. This allows
+    BOHB to insert new trials into the training.
+
+    See ray.tune.schedulers.HyperBandScheduler for parameter docstring.
+    """
+
+    def on_trial_add(self, tune_controller: "TuneController", trial: Trial):
+        """Adds new trial.
+
+        On a new trial add, if current bracket is not filled, add to current
+        bracket. Else, if current band is not filled, create new bracket, add
+        to current bracket. Else, create new iteration, create new bracket,
+        add to bracket.
+        """
+        if not self._metric or not self._metric_op:
+            raise ValueError(
+                "{} has been instantiated without a valid `metric` ({}) or "
+                "`mode` ({}) parameter. Either pass these parameters when "
+                "instantiating the scheduler, or pass them as parameters "
+                "to `tune.TuneConfig()`".format(
+                    self.__class__.__name__, self._metric, self._mode
+                )
+            )
+
+        cur_bracket = self._state["bracket"]
+        cur_band = self._hyperbands[self._state["band_idx"]]
+        if cur_bracket is None or cur_bracket.filled():
+            retry = True
+            while retry:
+                # if current iteration is filled, create new iteration
+                if self._cur_band_filled():
+                    cur_band = []
+                    self._hyperbands.append(cur_band)
+                    self._state["band_idx"] += 1
+
+                # MAIN CHANGE HERE - largest bracket first!
+                # cur_band will always be less than s_max_1 or else filled
+                s = self._s_max_1 - len(cur_band) - 1
+                assert s >= 0, "Current band is filled!"
+                if self._get_r0(s) == 0:
+                    logger.debug("BOHB: Bracket too small - Retrying...")
+                    cur_bracket = None
+                else:
+                    retry = False
+                    cur_bracket = self._create_bracket(s)
+                cur_band.append(cur_bracket)
+                self._state["bracket"] = cur_bracket
+
+        self._state["bracket"].add_trial(trial)
+        self._trial_info[trial] = cur_bracket, self._state["band_idx"]
+
+    def on_trial_result(
+        self, tune_controller: "TuneController", trial: Trial, result: Dict
+    ) -> str:
+        """If bracket is finished, all trials will be stopped.
+
+        If a given trial finishes and bracket iteration is not done,
+        the trial will be paused and resources will be given up.
+
+        This scheduler will not start trials but will stop trials.
+        The current running trial will not be handled,
+        as the trialrunner will be given control to handle it."""
+
+        result["hyperband_info"] = {}
+        bracket, _ = self._trial_info[trial]
+        bracket.update_trial_stats(trial, result)
+
+        if bracket.continue_trial(trial):
+            return TrialScheduler.CONTINUE
+
+        result["hyperband_info"]["budget"] = bracket._cumul_r
+
+        # MAIN CHANGE HERE!
+        statuses = [(t, t.status) for t in bracket._live_trials]
+        if not bracket.filled() or any(
+            status != Trial.PAUSED for t, status in statuses if t is not trial
+        ):
+            # BOHB Specific. This hack existed in old Ray versions
+            # and was removed, but it needs to be brought back
+            # as otherwise the BOHB doesn't behave as intended.
+            # The default concurrency limiter works by discarding
+            # new suggestions if there are more running trials
+            # than the limit. That doesn't take into account paused
+            # trials. With BOHB, this leads to N trials finishing
+            # completely and then another N trials starting,
+            # instead of trials being paused and resumed in brackets
+            # as intended.
+            # There should be a better API for this.
+            # TODO(team-ml): Refactor alongside HyperBandForBOHB
+            tune_controller.search_alg.searcher.on_pause(trial.trial_id)
+            return TrialScheduler.PAUSE
+
+        logger.debug(f"Processing bracket after trial {trial} result")
+        action = self._process_bracket(tune_controller, bracket)
+        if action == TrialScheduler.PAUSE:
+            tune_controller.search_alg.searcher.on_pause(trial.trial_id)
+        return action
+
+    def _unpause_trial(self, tune_controller: "TuneController", trial: Trial):
+        # Hack. See comment in on_trial_result
+        tune_controller.search_alg.searcher.on_unpause(trial.trial_id)
+
+    def choose_trial_to_run(
+        self, tune_controller: "TuneController", allow_recurse: bool = True
+    ) -> Optional[Trial]:
+        """Fair scheduling within iteration by completion percentage.
+
+        List of trials not used since all trials are tracked as state
+        of scheduler. If iteration is occupied (ie, no trials to run),
+        then look into next iteration.
+        """
+
+        for hyperband in self._hyperbands:
+            # band will have None entries if no resources
+            # are to be allocated to that bracket.
+            scrubbed = [b for b in hyperband if b is not None]
+            for bracket in scrubbed:
+                for trial in bracket.current_trials():
+                    if (
+                        trial.status == Trial.PAUSED
+                        and trial in bracket.trials_to_unpause
+                    ) or trial.status == Trial.PENDING:
+                        return trial
+        # MAIN CHANGE HERE!
+        if not any(t.status == Trial.RUNNING for t in tune_controller.get_trials()):
+            for hyperband in self._hyperbands:
+                for bracket in hyperband:
+                    if bracket and any(
+                        trial.status == Trial.PAUSED
+                        for trial in bracket.current_trials()
+                    ):
+                        # This will change the trial state
+                        logger.debug("Processing bracket since no trial is running.")
+                        self._process_bracket(tune_controller, bracket)
+
+                        # If there are pending trials now, suggest one.
+                        # This is because there might be both PENDING and
+                        # PAUSED trials now, and PAUSED trials will raise
+                        # an error before the trial runner tries again.
+                        if allow_recurse and any(
+                            (
+                                trial.status == Trial.PAUSED
+                                and trial in bracket.trials_to_unpause
+                            )
+                            or trial.status == Trial.PENDING
+                            for trial in bracket.current_trials()
+                        ):
+                            return self.choose_trial_to_run(
+                                tune_controller, allow_recurse=False
+                            )
+        # MAIN CHANGE HERE!
+        return None

python/ray/tune/schedulers/hyperband.py

@@ -0,0 +1,614 @@
+import collections
+import logging
+from typing import TYPE_CHECKING, Dict, List, Optional, Tuple
+
+import numpy as np
+
+from ray.tune.error import TuneError
+from ray.tune.experiment import Trial
+from ray.tune.result import DEFAULT_METRIC
+from ray.tune.schedulers.trial_scheduler import FIFOScheduler, TrialScheduler
+from ray.util.annotations import PublicAPI
+
+if TYPE_CHECKING:
+    from ray.tune.execution.tune_controller import TuneController
+
+logger = logging.getLogger(__name__)
+
+
+# Implementation notes:
+#    This implementation contains 3 logical levels.
+#    Each HyperBand iteration is a "band". There can be multiple
+#    bands running at once, and there can be 1 band that is incomplete.
+#
+#    In each band, there are at most `s` + 1 brackets.
+#    `s` is a value determined by given parameters, and assigned on
+#    a cyclic basis.
+#
+#    In each bracket, there are at most `n(s)` trials, indicating that
+#    `n` is a function of `s`. These trials go through a series of
+#    halving procedures, dropping lowest performers. Multiple
+#    brackets are running at once.
+#
+#    Trials added will be inserted into the most recent bracket
+#    and band and will spill over to new brackets/bands accordingly.
+#
+#    This maintains the bracket size and max trial count per band
+#    to 5 and 117 respectively, which correspond to that of
+#    `max_attr=81, eta=3` from the blog post. Trials will fill up
+#    from smallest bracket to largest, with largest
+#    having the most rounds of successive halving.
+@PublicAPI
+class HyperBandScheduler(FIFOScheduler):
+    """Implements the HyperBand early stopping algorithm.
+
+    HyperBandScheduler early stops trials using the HyperBand optimization
+    algorithm. It divides trials into brackets of varying sizes, and
+    periodically early stops low-performing trials within each bracket.
+
+    To use this implementation of HyperBand with Tune, all you need
+    to do is specify the max length of time a trial can run `max_t`, the time
+    units `time_attr`, the name of the reported objective value `metric`,
+    and if `metric` is to be maximized or minimized (`mode`).
+    We automatically determine reasonable values for the other
+    HyperBand parameters based on the given values.
+
+    For example, to limit trials to 10 minutes and early stop based on the
+    `episode_mean_reward` attr, construct:
+
+    ``HyperBand('time_total_s', 'episode_reward_mean', max_t=600)``
+
+    Note that Tune's stopping criteria will be applied in conjunction with
+    HyperBand's early stopping mechanisms.
+
+    See also: https://blog.ml.cmu.edu/2018/12/12/massively-parallel-hyperparameter-optimization/
+
+    Args:
+        time_attr: The training result attr to use for comparing time.
+            Note that you can pass in something non-temporal such as
+            `training_iteration` as a measure of progress, the only requirement
+            is that the attribute should increase monotonically.
+            Valid values are any key reported in the result dict by your
+            trainable. The auto-filled keys ``"training_iteration"`` (the
+            iteration count) and ``"time_total_s"`` (wall-clock seconds since
+            the trial started) always work; any additional numeric, monotonic
+            key your trainable reports via ``tune.report({...})`` is also valid
+            (for example ``"timesteps_total"`` or a custom progress counter).
+            Passing a key that is not present in the reported result causes
+            the scheduler to skip its decision for that step.
+        metric: The training result objective value attribute. Stopping
+            procedures will use this attribute. If None but a mode was passed,
+            the `ray.tune.result.DEFAULT_METRIC` will be used per default.
+        mode: One of {min, max}. Determines whether objective is
+            minimizing or maximizing the metric attribute.
+        max_t: max time units per trial. Trials will be stopped after
+            max_t time units (determined by time_attr) have passed.
+            The scheduler will terminate trials after this time has passed.
+            Note that this is different from the semantics of `max_t` as
+            mentioned in the original HyperBand paper.
+        reduction_factor: Same as `eta`. Determines how sharp
+            the difference is between bracket space-time allocation ratios.
+        stop_last_trials: Whether to terminate the trials after
+            reaching max_t. Defaults to True.
+    """  # noqa: E501
+
+    _supports_buffered_results = False
+
+    def __init__(
+        self,
+        time_attr: str = "training_iteration",
+        metric: Optional[str] = None,
+        mode: Optional[str] = None,
+        max_t: int = 81,
+        reduction_factor: float = 3,
+        stop_last_trials: bool = True,
+    ):
+        assert max_t > 0, "Max (time_attr) not valid!"
+        if mode:
+            assert mode in ["min", "max"], "`mode` must be 'min' or 'max'!"
+
+        super().__init__()
+        self._eta = reduction_factor
+        self._s_max_1 = int(np.round(np.log(max_t) / np.log(reduction_factor))) + 1
+        self._max_t_attr = max_t
+        # bracket max trials
+        self._get_n0 = lambda s: int(np.ceil(self._s_max_1 / (s + 1) * self._eta**s))
+        # bracket initial iterations
+        self._get_r0 = lambda s: int((max_t * self._eta ** (-s)))
+        self._hyperbands = [[]]  # list of hyperband iterations
+        self._trial_info = {}  # Stores Trial -> Bracket, Band Iteration
+
+        # Tracks state for new trial add
+        self._state = {"bracket": None, "band_idx": 0}
+        self._num_stopped = 0
+        self._metric = metric
+        self._mode = mode
+        self._metric_op = None
+
+        if self._mode == "max":
+            self._metric_op = 1.0
+        elif self._mode == "min":
+            self._metric_op = -1.0
+        self._time_attr = time_attr
+        self._stop_last_trials = stop_last_trials
+
+    def set_search_properties(
+        self, metric: Optional[str], mode: Optional[str], **spec
+    ) -> bool:
+        if self._metric and metric:
+            return False
+        if self._mode and mode:
+            return False
+
+        if metric:
+            self._metric = metric
+        if mode:
+            self._mode = mode
+
+        if self._mode == "max":
+            self._metric_op = 1.0
+        elif self._mode == "min":
+            self._metric_op = -1.0
+
+        if self._metric is None and self._mode:
+            # If only a mode was passed, use anonymous metric
+            self._metric = DEFAULT_METRIC
+
+        return True
+
+    def on_trial_add(self, tune_controller: "TuneController", trial: Trial):
+        """Adds new trial.
+
+        On a new trial add, if current bracket is not filled,
+        add to current bracket. Else, if current band is not filled,
+        create new bracket, add to current bracket.
+        Else, create new iteration, create new bracket, add to bracket."""
+        if not self._metric or not self._metric_op:
+            raise ValueError(
+                "{} has been instantiated without a valid `metric` ({}) or "
+                "`mode` ({}) parameter. Either pass these parameters when "
+                "instantiating the scheduler, or pass them as parameters "
+                "to `tune.TuneConfig()`".format(
+                    self.__class__.__name__, self._metric, self._mode
+                )
+            )
+
+        cur_bracket = self._state["bracket"]
+        cur_band = self._hyperbands[self._state["band_idx"]]
+        if cur_bracket is None or cur_bracket.filled():
+            retry = True
+            while retry:
+                # if current iteration is filled, create new iteration
+                if self._cur_band_filled():
+                    cur_band = []
+                    self._hyperbands.append(cur_band)
+                    self._state["band_idx"] += 1
+
+                # cur_band will always be less than s_max_1 or else filled
+                s = len(cur_band)
+                assert s < self._s_max_1, "Current band is filled!"
+                if self._get_r0(s) == 0:
+                    logger.info("Bracket too small - Retrying...")
+                    cur_bracket = None
+                else:
+                    retry = False
+                    cur_bracket = self._create_bracket(s)
+                cur_band.append(cur_bracket)
+                self._state["bracket"] = cur_bracket
+
+        self._state["bracket"].add_trial(trial)
+        self._trial_info[trial] = cur_bracket, self._state["band_idx"]
+
+    def _create_bracket(self, s):
+        return _Bracket(
+            time_attr=self._time_attr,
+            max_trials=self._get_n0(s),
+            init_t_attr=self._get_r0(s),
+            max_t_attr=self._max_t_attr,
+            eta=self._eta,
+            s=s,
+            stop_last_trials=self._stop_last_trials,
+        )
+
+    def _cur_band_filled(self) -> bool:
+        """Checks if the current band is filled.
+
+        The size of the current band should be equal to s_max_1"""
+
+        cur_band = self._hyperbands[self._state["band_idx"]]
+        return len(cur_band) == self._s_max_1
+
+    def on_trial_result(
+        self, tune_controller: "TuneController", trial: Trial, result: Dict
+    ):
+        """If bracket is finished, all trials will be stopped.
+
+        If a given trial finishes and bracket iteration is not done,
+        the trial will be paused and resources will be given up.
+
+        This scheduler will not start trials but will stop trials.
+        The current running trial will not be handled,
+        as the trialrunner will be given control to handle it."""
+
+        bracket, _ = self._trial_info[trial]
+        bracket.update_trial_stats(trial, result)
+
+        if bracket.continue_trial(trial):
+            return TrialScheduler.CONTINUE
+
+        logger.debug(f"Processing bracket after trial {trial} result")
+        action = self._process_bracket(tune_controller, bracket)
+        logger.debug(
+            f"{action} for {trial} on "
+            f"{self._time_attr}={result.get(self._time_attr)}"
+        )
+        return action
+
+    def _process_bracket(
+        self, tune_controller: "TuneController", bracket: "_Bracket"
+    ) -> str:
+        """This is called whenever a trial makes progress.
+
+        When all live trials in the bracket have no more iterations left,
+        Trials will be successively halved. If bracket is done, all
+        non-running trials will be stopped and cleaned up,
+        and during each halving phase, bad trials will be stopped while good
+        trials will return to "PENDING".
+
+        Note some implicit conditions here: In ``on_trial_result`` a trial is
+        either continued (e.g. if it didn't reach the time threshold for the bracket)
+        or this method (``_process_bracket``) is called. If there are other trials left
+        that still haven't reached the threshold, the trial is PAUSED. This means
+        that when the bracket is actually processed (``bracket.cur_iter_done``), there
+        is at most one RUNNING trial (which is the trial that is currently processed)
+        and the rest are either PAUSED (as explained above) or TERMINATED/ERRORED
+        (if they finish separately).
+        """
+
+        action = TrialScheduler.PAUSE
+        if bracket.cur_iter_done():
+            if bracket.finished():
+                bracket.cleanup_full(tune_controller)
+                return TrialScheduler.STOP
+
+            bracket.is_being_processed = True
+
+            good, bad = bracket.successive_halving(self._metric, self._metric_op)
+
+            logger.debug(
+                f"Processing {len(good)} good and {len(bad)} bad trials in "
+                f"bracket {bracket}.\n"
+                f"Good: {good}\nBad: {bad}"
+            )
+
+            # kill bad trials
+            self._num_stopped += len(bad)
+            for t in bad:
+                if t.status == Trial.PAUSED or t.is_saving:
+                    logger.debug(f"Stopping other trial {str(t)}")
+                    tune_controller.stop_trial(t)
+                elif t.status == Trial.RUNNING:
+                    # See the docstring: There can only be at most one RUNNING
+                    # trial, which is the current trial.
+                    logger.debug(f"Stopping current trial {str(t)}")
+                    bracket.cleanup_trial(t)
+                    action = TrialScheduler.STOP
+                else:
+                    # Trials cannot be ERROR/TERMINATED, as then they would have
+                    # been removed from the bracket (in `bracket.cleanup_trial`).
+                    # Trials cannot be PENDING, as then they wouldn't have reported
+                    # enough results to finish the bracket, and it wouldn't be
+                    # processed.
+                    raise TuneError(
+                        f"Trial with unexpected bad status encountered: "
+                        f"{str(t)} is {t.status}"
+                    )
+
+            # ready the good trials - if trial is too far ahead, don't continue
+            for t in good:
+                if bracket.continue_trial(t):
+                    # The scheduler should have cleaned up this trial already.
+                    assert t.status not in (Trial.ERROR, Trial.TERMINATED), (
+                        f"Good trial {t.trial_id} is in an invalid state: {t.status}\n"
+                        "Expected trial to be either PAUSED, PENDING, or RUNNING.\n"
+                        "If you encounter this, please file an issue on the Ray Github."
+                    )
+                    if t.status == Trial.PAUSED or t.is_saving:
+                        logger.debug(f"Unpausing trial {str(t)}")
+                        self._unpause_trial(tune_controller, t)
+                        bracket.trials_to_unpause.add(t)
+                    elif t.status == Trial.RUNNING:
+                        # See the docstring: There can only be at most one RUNNING
+                        # trial, which is the current trial.
+                        logger.debug(f"Continuing current trial {str(t)}")
+                        action = TrialScheduler.CONTINUE
+                    # else: PENDING trial (from a previous unpause) should stay as is.
+                elif bracket.finished() and bracket.stop_last_trials:
+                    # Scheduler decides to not continue trial because the bracket
+                    # reached max_t. In this case, stop the trials
+                    if t.status == Trial.PAUSED or t.is_saving:
+                        logger.debug(f"Bracket finished. Stopping other trial {str(t)}")
+                        tune_controller.stop_trial(t)
+                    elif t.status == Trial.RUNNING:
+                        # See the docstring: There can only be at most one RUNNING
+                        # trial, which is the current trial.
+                        logger.debug(
+                            f"Bracket finished. Stopping current trial {str(t)}"
+                        )
+                        bracket.cleanup_trial(t)
+                        action = TrialScheduler.STOP
+        return action
+
+    def _unpause_trial(self, tune_controller: "TuneController", trial: Trial):
+        """No-op by default."""
+        return
+
+    def on_trial_remove(self, tune_controller: "TuneController", trial: Trial):
+        """Notification when trial terminates.
+
+        Trial info is removed from bracket. Triggers halving if bracket is
+        not finished."""
+        bracket, _ = self._trial_info[trial]
+        bracket.cleanup_trial(trial)
+        if not bracket.finished() and not bracket.is_being_processed:
+            logger.debug(f"Processing bracket after trial {trial} removed")
+            self._process_bracket(tune_controller, bracket)
+
+    def on_trial_complete(
+        self, tune_controller: "TuneController", trial: Trial, result: Dict
+    ):
+        """Cleans up trial info from bracket if trial completed early."""
+        self.on_trial_remove(tune_controller, trial)
+
+    def on_trial_error(self, tune_controller: "TuneController", trial: Trial):
+        """Cleans up trial info from bracket if trial errored early."""
+        self.on_trial_remove(tune_controller, trial)
+
+    def choose_trial_to_run(self, tune_controller: "TuneController") -> Optional[Trial]:
+        """Fair scheduling within iteration by completion percentage.
+
+        List of trials not used since all trials are tracked as state
+        of scheduler. If iteration is occupied (ie, no trials to run),
+        then look into next iteration.
+        """
+
+        for hyperband in self._hyperbands:
+            # band will have None entries if no resources
+            # are to be allocated to that bracket.
+            scrubbed = [b for b in hyperband if b is not None]
+            for bracket in sorted(scrubbed, key=lambda b: b.completion_percentage()):
+                for trial in bracket.current_trials():
+                    if (
+                        trial.status == Trial.PAUSED
+                        and trial in bracket.trials_to_unpause
+                    ) or trial.status == Trial.PENDING:
+                        return trial
+        return None
+
+    def debug_string(self) -> str:
+        """This provides a progress notification for the algorithm.
+
+        For each bracket, the algorithm will output a string as follows:
+
+            Bracket(Max Size (n)=5, Milestone (r)=33, completed=14.6%):
+            {PENDING: 2, RUNNING: 3, TERMINATED: 2}
+
+        "Max Size" indicates the max number of pending/running experiments
+        set according to the Hyperband algorithm.
+
+        "Milestone" indicates the iterations a trial will run for before
+        the next halving will occur.
+
+        "Completed" indicates an approximate progress metric. Some brackets,
+        like ones that are unfilled, will not reach 100%.
+        """
+        out = "Using HyperBand: "
+        out += "num_stopped={} total_brackets={}".format(
+            self._num_stopped, sum(len(band) for band in self._hyperbands)
+        )
+        for i, band in enumerate(self._hyperbands):
+            out += "\nRound #{}:".format(i)
+            for bracket in band:
+                if bracket:
+                    out += "\n  {}".format(bracket)
+        return out
+
+    def state(self) -> Dict[str, int]:
+        return {
+            "num_brackets": sum(len(band) for band in self._hyperbands),
+            "num_stopped": self._num_stopped,
+        }
+
+
+class _Bracket:
+    """Logical object for tracking Hyperband bracket progress. Keeps track
+    of proper parameters as designated by HyperBand.
+
+    Also keeps track of progress to ensure good scheduling.
+    """
+
+    def __init__(
+        self,
+        time_attr: str,
+        max_trials: int,
+        init_t_attr: int,
+        max_t_attr: int,
+        eta: float,
+        s: int,
+        stop_last_trials: bool = True,
+    ):
+        self._live_trials = {}  # maps trial -> current result
+        self._all_trials = []
+        self._time_attr = time_attr  # attribute to
+
+        self._n = self._n0 = max_trials
+        self._r = self._r0 = init_t_attr
+        self._max_t_attr = max_t_attr
+        self._cumul_r = self._r0
+
+        self._eta = eta
+        self._halves = s
+
+        self._total_work = self._calculate_total_work(self._n0, self._r0, s)
+        self._completed_progress = 0
+        self.stop_last_trials = stop_last_trials
+        self.is_being_processed = False
+
+        self.trials_to_unpause = set()
+
+    def add_trial(self, trial: Trial):
+        """Add trial to bracket assuming bracket is not filled.
+
+        At a later iteration, a newly added trial will be given equal
+        opportunity to catch up."""
+        assert not self.filled(), "Cannot add trial to filled bracket!"
+        self._live_trials[trial] = None
+        self._all_trials.append(trial)
+
+    def cur_iter_done(self) -> bool:
+        """Checks if all iterations have completed.
+
+        TODO(rliaw): also check that `t.iterations == self._r`"""
+        return all(
+            self._get_result_time(result) >= self._cumul_r
+            for result in self._live_trials.values()
+        )
+
+    def finished(self) -> bool:
+        if not self.stop_last_trials:
+            return False
+        return self._halves == 0 and self.cur_iter_done()
+
+    def current_trials(self) -> List[Trial]:
+        return list(self._live_trials)
+
+    def continue_trial(self, trial: Trial) -> bool:
+        result = self._live_trials[trial]
+        if not self.stop_last_trials and self._halves == 0:
+            return True
+        elif self._get_result_time(result) < self._cumul_r:
+            logger.debug(
+                f"Continuing trial {trial} as it hasn't reached the time threshold "
+                f"{self._cumul_r}, yet."
+            )
+            return True
+        return False
+
+    def filled(self) -> bool:
+        """Checks if bracket is filled.
+
+        Only let new trials be added at current level minimizing the need
+        to backtrack and bookkeep previous medians."""
+
+        return len(self._live_trials) == self._n
+
+    def successive_halving(
+        self, metric: str, metric_op: float
+    ) -> Tuple[List[Trial], List[Trial]]:
+        if self._halves == 0 and not self.stop_last_trials:
+            return self._live_trials, []
+        assert self._halves > 0
+
+        # "Halving" is a misnomer. We're actually reducing by factor `eta`.
+        self._halves -= 1
+
+        # If we had 8 trials in the bracket and eta=2, we will keep 4.
+        # If we had 9 trials in the bracket and eta=3, we will keep 3.
+        self._n = int(np.ceil(self._n / self._eta))
+
+        # Likewise, we increase the number of iterations until we process the bracket
+        # again.
+        # Remember r0 = max_t * self._eta ** (-s)
+        # Let max_t=16, eta=2, s=1. Then r0=8, and we calculate r1=16.
+        # Let max_t=16, eta=2, s=2. Then r0=4, and we calculate r1=8, r2=16.
+
+        # Let max_t=81, eta=3, s=1. Then r0=27, and we calculate r1=81.
+        # Let max_t=81, eta=3, s=2. Then r0=9, and we calculate r1=27, r2=81.
+        self._r *= self._eta
+        self._r = int(min(self._r, self._max_t_attr))
+        self._cumul_r = self._r
+        sorted_trials = sorted(
+            self._live_trials, key=lambda t: metric_op * self._live_trials[t][metric]
+        )
+
+        good, bad = sorted_trials[-self._n :], sorted_trials[: -self._n]
+        return good, bad
+
+    def update_trial_stats(self, trial: Trial, result: Dict):
+        """Update result for trial. Called after trial has finished
+        an iteration - will decrement iteration count.
+
+        TODO(rliaw): The other alternative is to keep the trials
+        in and make sure they're not set as pending later."""
+
+        assert trial in self._live_trials
+        assert self._get_result_time(result) >= 0
+        observed_time = self._get_result_time(result)
+        last_observed = self._get_result_time(self._live_trials[trial])
+
+        delta = observed_time - last_observed
+        if delta <= 0:
+            logger.info(
+                "Restoring from a previous point in time. "
+                "Previous={}; Now={}".format(last_observed, observed_time)
+            )
+        self._completed_progress += delta
+        self._live_trials[trial] = result
+        self.trials_to_unpause.discard(trial)
+
+    def cleanup_trial(self, trial: Trial):
+        """Clean up statistics tracking for terminated trials (either by force
+        or otherwise).
+
+        This may cause bad trials to continue for a long time, in the case
+        where all the good trials finish early and there are only bad trials
+        left in a bracket with a large max-iteration."""
+        self._live_trials.pop(trial, None)
+
+    def cleanup_full(self, tune_controller: "TuneController"):
+        """Cleans up bracket after bracket is completely finished.
+
+        Lets the last trial continue to run until termination condition
+        kicks in."""
+        for trial in self.current_trials():
+            if trial.status == Trial.PAUSED:
+                tune_controller.stop_trial(trial)
+
+    def completion_percentage(self) -> float:
+        """Returns a progress metric.
+
+        This will not be always finish with 100 since dead trials
+        are dropped."""
+        if self.finished():
+            return 1.0
+        return min(self._completed_progress / self._total_work, 1.0)
+
+    def _get_result_time(self, result: Dict) -> float:
+        if result is None:
+            return 0
+        return result[self._time_attr]
+
+    def _calculate_total_work(self, n: int, r: float, s: int):
+        work = 0
+        cumulative_r = r
+        for _ in range(s + 1):
+            work += int(n) * int(r)
+            n /= self._eta
+            n = int(np.ceil(n))
+            r *= self._eta
+            r = int(min(r, self._max_t_attr - cumulative_r))
+        return work
+
+    def __repr__(self) -> str:
+        status = ", ".join(
+            [
+                "Max Size (n)={}".format(self._n),
+                "Milestone (r)={}".format(self._cumul_r),
+                "completed={:.1%}".format(self.completion_percentage()),
+            ]
+        )
+        counts = collections.Counter([t.status for t in self._all_trials])
+        trial_statuses = ", ".join(
+            sorted("{}: {}".format(k, v) for k, v in counts.items())
+        )
+        return "Bracket({}): {{{}}} ".format(status, trial_statuses)

python/ray/tune/schedulers/median_stopping_rule.py

@@ -0,0 +1,225 @@
+import collections
+import logging
+from typing import TYPE_CHECKING, Dict, List, Optional
+
+import numpy as np
+
+from ray.tune.experiment import Trial
+from ray.tune.result import DEFAULT_METRIC
+from ray.tune.schedulers.trial_scheduler import FIFOScheduler, TrialScheduler
+from ray.util.annotations import PublicAPI
+
+if TYPE_CHECKING:
+    from ray.tune.execution.tune_controller import TuneController
+
+logger = logging.getLogger(__name__)
+
+
+@PublicAPI
+class MedianStoppingRule(FIFOScheduler):
+    """Implements the median stopping rule as described in the Vizier paper:
+
+    https://research.google.com/pubs/pub46180.html
+
+    Args:
+        time_attr: The training result attr to use for comparing time.
+            Note that you can pass in something non-temporal such as
+            `training_iteration` as a measure of progress, the only requirement
+            is that the attribute should increase monotonically.
+            Valid values are any key reported in the result dict by your
+            trainable. The auto-filled keys ``"training_iteration"`` (the
+            iteration count) and ``"time_total_s"`` (wall-clock seconds since
+            the trial started) always work; any additional numeric, monotonic
+            key your trainable reports via ``tune.report({...})`` is also valid
+            (for example ``"timesteps_total"`` or a custom progress counter).
+            Passing a key that is not present in the reported result causes
+            the scheduler to skip its decision for that step.
+        metric: The training result objective value attribute. Stopping
+            procedures will use this attribute. If None but a mode was passed,
+            the `ray.tune.result.DEFAULT_METRIC` will be used per default.
+        mode: One of {min, max}. Determines whether objective is
+            minimizing or maximizing the metric attribute.
+        grace_period: Only stop trials at least this old in time.
+            The mean will only be computed from this time onwards. The units
+            are the same as the attribute named by `time_attr`.
+        min_samples_required: Minimum number of trials to compute median
+            over.
+        min_time_slice: Each trial runs at least this long before
+            yielding (assuming it isn't stopped). Note: trials ONLY yield if
+            there are not enough samples to evaluate performance for the
+            current result AND there are other trials waiting to run.
+            The units are the same as the attribute named by `time_attr`.
+        hard_stop: If False, pauses trials instead of stopping
+            them. When all other trials are complete, paused trials will be
+            resumed and allowed to run FIFO.
+    """
+
+    def __init__(
+        self,
+        time_attr: str = "time_total_s",
+        metric: Optional[str] = None,
+        mode: Optional[str] = None,
+        grace_period: float = 60.0,
+        min_samples_required: int = 3,
+        min_time_slice: int = 0,
+        hard_stop: bool = True,
+    ):
+        super().__init__()
+        self._stopped_trials = set()
+        self._grace_period = grace_period
+        self._min_samples_required = min_samples_required
+        self._min_time_slice = min_time_slice
+        self._metric = metric
+        self._worst = None
+        self._compare_op = None
+
+        self._mode = mode
+        if mode:
+            assert mode in ["min", "max"], "`mode` must be 'min' or 'max'."
+            self._worst = float("-inf") if self._mode == "max" else float("inf")
+            self._compare_op = max if self._mode == "max" else min
+
+        self._time_attr = time_attr
+        self._hard_stop = hard_stop
+        self._trial_state = {}
+        self._last_pause = collections.defaultdict(lambda: float("-inf"))
+        self._results = collections.defaultdict(list)
+
+    def set_search_properties(
+        self, metric: Optional[str], mode: Optional[str], **spec
+    ) -> bool:
+        if self._metric and metric:
+            return False
+        if self._mode and mode:
+            return False
+
+        if metric:
+            self._metric = metric
+        if mode:
+            self._mode = mode
+
+        self._worst = float("-inf") if self._mode == "max" else float("inf")
+        self._compare_op = max if self._mode == "max" else min
+
+        if self._metric is None and self._mode:
+            # If only a mode was passed, use anonymous metric
+            self._metric = DEFAULT_METRIC
+
+        return True
+
+    def on_trial_add(self, tune_controller: "TuneController", trial: Trial):
+        if not self._metric or not self._worst or not self._compare_op:
+            raise ValueError(
+                "{} has been instantiated without a valid `metric` ({}) or "
+                "`mode` ({}) parameter. Either pass these parameters when "
+                "instantiating the scheduler, or pass them as parameters "
+                "to `tune.TuneConfig()`".format(
+                    self.__class__.__name__, self._metric, self._mode
+                )
+            )
+
+        super(MedianStoppingRule, self).on_trial_add(tune_controller, trial)
+
+    def on_trial_result(
+        self, tune_controller: "TuneController", trial: Trial, result: Dict
+    ) -> str:
+        """Callback for early stopping.
+
+        This stopping rule stops a running trial if the trial's best objective
+        value by step `t` is strictly worse than the median of the running
+        averages of all completed trials' objectives reported up to step `t`.
+        """
+        if self._time_attr not in result or self._metric not in result:
+            return TrialScheduler.CONTINUE
+
+        if trial in self._stopped_trials:
+            assert not self._hard_stop
+            # Fall back to FIFO
+            return TrialScheduler.CONTINUE
+
+        time = result[self._time_attr]
+        self._results[trial].append(result)
+
+        if time < self._grace_period:
+            return TrialScheduler.CONTINUE
+
+        trials = self._trials_beyond_time(time)
+        trials.remove(trial)
+
+        if len(trials) < self._min_samples_required:
+            action = self._on_insufficient_samples(tune_controller, trial, time)
+            if action == TrialScheduler.PAUSE:
+                self._last_pause[trial] = time
+                action_str = "Yielding time to other trials."
+            else:
+                action_str = "Continuing anyways."
+            logger.debug(
+                "MedianStoppingRule: insufficient samples={} to evaluate "
+                "trial {} at t={}. {}".format(
+                    len(trials), trial.trial_id, time, action_str
+                )
+            )
+            return action
+
+        median_result = self._median_result(trials, time)
+        best_result = self._best_result(trial)
+        logger.debug(
+            "Trial {} best res={} vs median res={} at t={}".format(
+                trial, best_result, median_result, time
+            )
+        )
+
+        if self._compare_op(median_result, best_result) != best_result:
+            logger.debug("MedianStoppingRule: early stopping {}".format(trial))
+            self._stopped_trials.add(trial)
+            if self._hard_stop:
+                return TrialScheduler.STOP
+            else:
+                return TrialScheduler.PAUSE
+        else:
+            return TrialScheduler.CONTINUE
+
+    def on_trial_complete(
+        self, tune_controller: "TuneController", trial: Trial, result: Dict
+    ):
+        self._results[trial].append(result)
+
+    def debug_string(self) -> str:
+        return "Using MedianStoppingRule: num_stopped={}.".format(
+            len(self._stopped_trials)
+        )
+
+    def _on_insufficient_samples(
+        self, tune_controller: "TuneController", trial: Trial, time: float
+    ) -> str:
+        pause = time - self._last_pause[trial] > self._min_time_slice
+        pause = pause and [
+            t
+            for t in tune_controller.get_live_trials()
+            if t.status in (Trial.PENDING, Trial.PAUSED)
+        ]
+        return TrialScheduler.PAUSE if pause else TrialScheduler.CONTINUE
+
+    def _trials_beyond_time(self, time: float) -> List[Trial]:
+        trials = [
+            trial
+            for trial in self._results
+            if self._results[trial][-1][self._time_attr] >= time
+        ]
+        return trials
+
+    def _median_result(self, trials: List[Trial], time: float):
+        return np.median([self._running_mean(trial, time) for trial in trials])
+
+    def _running_mean(self, trial: Trial, time: float) -> np.ndarray:
+        results = self._results[trial]
+        # TODO(ekl) we could do interpolation to be more precise, but for now
+        # assume len(results) is large and the time diffs are roughly equal
+        scoped_results = [
+            r for r in results if self._grace_period <= r[self._time_attr] <= time
+        ]
+        return np.mean([r[self._metric] for r in scoped_results])
+
+    def _best_result(self, trial):
+        results = self._results[trial]
+        return self._compare_op([r[self._metric] for r in results])

python/ray/tune/schedulers/pb2.py

@@ -0,0 +1,505 @@
+import logging
+from copy import deepcopy
+from typing import TYPE_CHECKING, Callable, Dict, Optional, Tuple, Union
+
+import numpy as np
+import pandas as pd
+
+from ray.tune import TuneError
+from ray.tune.experiment import Trial
+from ray.tune.schedulers import PopulationBasedTraining
+from ray.tune.schedulers.pbt import _PBTTrialState
+from ray.tune.utils.util import flatten_dict, unflatten_dict
+from ray.util.debug import log_once
+
+if TYPE_CHECKING:
+    from ray.tune.execution.tune_controller import TuneController
+
+
+def import_pb2_dependencies():
+    try:
+        import sklearn
+    except ImportError:
+        sklearn = None
+    return sklearn
+
+
+has_sklearn = import_pb2_dependencies()
+
+if has_sklearn:
+    from sklearn.gaussian_process import GaussianProcessRegressor
+
+    from ray.tune.schedulers.pb2_utils import (
+        UCB,
+        TV_SquaredExp,
+        normalize,
+        optimize_acq,
+        select_length,
+        standardize,
+    )
+
+logger = logging.getLogger(__name__)
+
+
+def _fill_config(
+    config: Dict, hyperparam_bounds: Dict[str, Union[dict, list, tuple]]
+) -> Dict:
+    """Fills missing hyperparameters in config by sampling uniformly from the
+    specified `hyperparam_bounds`.
+    Recursively fills the config if `hyperparam_bounds` is a nested dict.
+
+    This is a helper used to set initial hyperparameter values if the user doesn't
+    specify them in the Tuner `param_space`.
+
+    Returns the dict of filled hyperparameters.
+    """
+    filled_hyperparams = {}
+    for param_name, bounds in hyperparam_bounds.items():
+        if isinstance(bounds, dict):
+            if param_name not in config:
+                config[param_name] = {}
+            filled_hyperparams[param_name] = _fill_config(config[param_name], bounds)
+        elif isinstance(bounds, (list, tuple)) and param_name not in config:
+            if log_once(param_name + "-missing"):
+                logger.debug(
+                    f"Cannot find {param_name} in config. Initializing by "
+                    "sampling uniformly from the provided `hyperparam_bounds`."
+                )
+            assert len(bounds) == 2
+            low, high = bounds
+            config[param_name] = filled_hyperparams[param_name] = np.random.uniform(
+                low, high
+            )
+    return filled_hyperparams
+
+
+def _select_config(
+    Xraw: np.array,
+    yraw: np.array,
+    current: list,
+    newpoint: np.array,
+    bounds: dict,
+    num_f: int,
+) -> np.ndarray:
+    """Selects the next hyperparameter config to try.
+
+    This function takes the formatted data, fits the GP model and optimizes the
+    UCB acquisition function to select the next point.
+
+    Args:
+        Xraw: The un-normalized array of hyperparams, Time and
+            Reward
+        yraw: The un-normalized vector of reward changes.
+        current: The hyperparams of trials currently running. This is
+            important so we do not select the same config twice. If there is
+            data here then we fit a second GP including it
+            (with fake y labels). The GP variance doesn't depend on the y
+            labels so it is ok.
+        newpoint: The Reward and Time for the new point.
+            We cannot change these as they are based on the *new weights*.
+        bounds: Bounds for the hyperparameters. Used to normalize.
+        num_f: The number of fixed params. Almost always 2 (reward+time)
+
+    Returns:
+        xt: A vector of new hyperparameters.
+    """
+    length = select_length(Xraw, yraw, bounds, num_f)
+
+    Xraw = Xraw[-length:, :]
+    yraw = yraw[-length:]
+
+    base_vals = np.array(list(bounds.values())).T
+    oldpoints = Xraw[:, :num_f]
+    old_lims = np.concatenate(
+        (np.max(oldpoints, axis=0), np.min(oldpoints, axis=0))
+    ).reshape(2, oldpoints.shape[1])
+    limits = np.concatenate((old_lims, base_vals), axis=1)
+
+    X = normalize(Xraw, limits)
+    y = standardize(yraw).reshape(yraw.size, 1)
+
+    fixed = normalize(newpoint, oldpoints)
+
+    kernel = TV_SquaredExp(variance=1.0, lengthscale=1.0, epsilon=0.1)
+
+    try:
+        m = GaussianProcessRegressor(
+            kernel=kernel, optimizer="fmin_l_bfgs_b", alpha=1e-10
+        )
+        m.fit(X, y)
+    except np.linalg.LinAlgError:
+        # add diagonal ** we would ideally make this something more robust...
+        X += np.eye(X.shape[0]) * 1e-3
+        m = GaussianProcessRegressor(
+            kernel=kernel, optimizer="fmin_l_bfgs_b", alpha=1e-10
+        )
+        m.fit(X, y)
+
+    if current is None:
+        m1 = deepcopy(m)
+    else:
+        # add the current trials to the dataset
+        padding = np.array([fixed for _ in range(current.shape[0])])
+        current = normalize(current, base_vals)
+        current = np.hstack((padding, current))
+
+        Xnew = np.vstack((X, current))
+        ypad = np.zeros(current.shape[0])
+        ypad = ypad.reshape(-1, 1)
+        ynew = np.vstack((y, ypad))
+
+        kernel1 = TV_SquaredExp(variance=1.0, lengthscale=1.0, epsilon=0.1)
+        m1 = GaussianProcessRegressor(
+            kernel=kernel1, optimizer="fmin_l_bfgs_b", alpha=1e-10
+        )
+        m1.fit(Xnew, ynew)
+
+    xt = optimize_acq(UCB, m, m1, fixed, num_f)
+
+    # convert back...
+    xt = xt * (np.max(base_vals, axis=0) - np.min(base_vals, axis=0)) + np.min(
+        base_vals, axis=0
+    )
+
+    xt = xt.astype(np.float32)
+    return xt
+
+
+def _explore(
+    data: pd.DataFrame,
+    bounds: Dict[str, Tuple[float, float]],
+    current: list,
+    base: Trial,
+    old: Trial,
+    config: Dict[str, Tuple[float, float]],
+) -> Tuple[Dict, pd.DataFrame]:
+    """Returns next hyperparameter configuration to use.
+
+    This function primarily processes the data from completed trials
+    and then requests the next config from the select_config function.
+    It then adds the new trial to the dataframe, so that the reward change
+    can be computed using the new weights.
+    It returns the new point and the dataframe with the new entry.
+    """
+
+    df = data.sort_values(by="Time").reset_index(drop=True)
+
+    # Group by trial ID and hyperparams.
+    # Compute change in timesteps and reward.
+    df["y"] = df.groupby(["Trial"] + list(bounds.keys()))["Reward"].diff()
+    df["t_change"] = df.groupby(["Trial"] + list(bounds.keys()))["Time"].diff()
+
+    # Delete entries without positive change in t.
+    df = df[df["t_change"] > 0].reset_index(drop=True)
+    df["R_before"] = df.Reward - df.y
+
+    # Normalize the reward change by the update size.
+    # For example if trials took diff lengths of time.
+    df["y"] = df.y / df.t_change
+    df = df[~df.y.isna()].reset_index(drop=True)
+    df = df.sort_values(by="Time").reset_index(drop=True)
+
+    # Only use the last 1k datapoints, so the GP is not too slow.
+    df = df.iloc[-1000:, :].reset_index(drop=True)
+
+    # We need this to know the T and Reward for the weights.
+    dfnewpoint = df[df["Trial"] == str(base)]
+
+    if not dfnewpoint.empty:
+        # N ow specify the dataset for the GP.
+        y = np.array(df.y.values)
+        # Meta data we keep -> episodes and reward.
+        # (TODO: convert to curve)
+        t_r = df[["Time", "R_before"]]
+        hparams = df[bounds.keys()]
+        X = pd.concat([t_r, hparams], axis=1).values
+        newpoint = df[df["Trial"] == str(base)].iloc[-1, :][["Time", "R_before"]].values
+        new = _select_config(X, y, current, newpoint, bounds, num_f=len(t_r.columns))
+
+        new_config = config.copy()
+        values = []
+        # Cast types for new hyperparameters.
+        for i, col in enumerate(hparams.columns):
+            # Use the type from the old config. Like this types
+            # should be passed on from the first config downwards.
+            type_ = type(config[col])
+            new_config[col] = type_(new[i])
+            values.append(type_(new[i]))
+
+        new_T = df[df["Trial"] == str(base)].iloc[-1, :]["Time"]
+        new_Reward = df[df["Trial"] == str(base)].iloc[-1, :].Reward
+
+        lst = [[str(old)] + [new_T] + values + [new_Reward]]
+        cols = ["Trial", "Time"] + list(bounds) + ["Reward"]
+        new_entry = pd.DataFrame(lst, columns=cols)
+
+        # Create an entry for the new config, with the reward from the
+        # copied agent.
+        data = pd.concat([data, new_entry]).reset_index(drop=True)
+
+    else:
+        new_config = config.copy()
+
+    return new_config, data
+
+
+class PB2(PopulationBasedTraining):
+    """Implements the Population Based Bandit (PB2) algorithm.
+
+    PB2 trains a group of models (or agents) in parallel. Periodically, poorly
+    performing models clone the state of the top performers, and the hyper-
+    parameters are re-selected using GP-bandit optimization. The GP model is
+    trained to predict the improvement in the next training period.
+
+    Like PBT, PB2 adapts hyperparameters during training time. This enables
+    very fast hyperparameter discovery and also automatically discovers
+    schedules.
+
+    This Tune PB2 implementation is built on top of Tune's PBT implementation.
+    It considers all trials added as part of the PB2 population. If the number
+    of trials exceeds the cluster capacity, they will be time-multiplexed as to
+    balance training progress across the population. To run multiple trials,
+    use `tune.TuneConfig(num_samples=<int>)`.
+
+    In {LOG_DIR}/{MY_EXPERIMENT_NAME}/, all mutations are logged in
+    `pb2_global.txt` and individual policy perturbations are recorded
+    in pb2_policy_{i}.txt. Tune logs: [target trial tag, clone trial tag,
+    target trial iteration, clone trial iteration, old config, new config]
+    on each perturbation step.
+
+    Args:
+        time_attr: The training result attr to use for comparing time.
+            Note that you can pass in something non-temporal such as
+            `training_iteration` as a measure of progress, the only requirement
+            is that the attribute should increase monotonically.
+            Valid values are any key reported in the result dict by your
+            trainable. The auto-filled keys ``"training_iteration"`` (the
+            iteration count) and ``"time_total_s"`` (wall-clock seconds since
+            the trial started) always work; any additional numeric, monotonic
+            key your trainable reports via ``tune.report({...})`` is also valid
+            (for example ``"timesteps_total"`` or a custom progress counter).
+            Passing a key that is not present in the reported result causes
+            the scheduler to skip its decision for that step.
+        metric: The training result objective value attribute. Stopping
+            procedures will use this attribute.
+        mode: One of {min, max}. Determines whether objective is
+            minimizing or maximizing the metric attribute.
+        perturbation_interval: Models will be considered for
+            perturbation at this interval of `time_attr`. Note that
+            perturbation incurs checkpoint overhead, so you shouldn't set this
+            to be too frequent.
+        hyperparam_bounds: Hyperparameters to mutate. The format is
+            as follows: for each key, enter a list of the form [min, max]
+            representing the minimum and maximum possible hyperparam values.
+            A key can also hold a dict for nested hyperparameters.
+            Tune will sample uniformly between the bounds provided by
+            `hyperparam_bounds` for the initial hyperparameter values if the
+            corresponding hyperparameters are not present in a trial's initial `config`.
+        quantile_fraction: Parameters are transferred from the top
+            `quantile_fraction` fraction of trials to the bottom
+            `quantile_fraction` fraction. Needs to be between 0 and 0.5.
+            Setting it to 0 essentially implies doing no exploitation at all.
+        log_config: Whether to log the ray config of each model to
+            local_dir at each exploit. Allows config schedule to be
+            reconstructed.
+        require_attrs: Whether to require time_attr and metric to appear
+            in result for every iteration. If True, error will be raised
+            if these values are not present in trial result.
+        synch: If False, will use asynchronous implementation of
+            PBT. Trial perturbations occur every perturbation_interval for each
+            trial independently. If True, will use synchronous implementation
+            of PBT. Perturbations will occur only after all trials are
+            synced at the same time_attr every perturbation_interval.
+            Defaults to False. See Appendix A.1 here
+            https://arxiv.org/pdf/1711.09846.pdf.
+        custom_explore_fn: You can also specify a custom exploration
+            function. This function is invoked as `f(config)`, where the input
+            is the new config generated by Bayesian Optimization. This function
+            should return the `config` updated as needed.
+
+    Example:
+
+        .. code-block:: python
+
+            from ray import tune
+            from ray.tune.schedulers.pb2 import PB2
+            from ray.tune.examples.pbt_function import pbt_function
+
+            pb2 = PB2(
+                metric="mean_accuracy",
+                mode="max",
+                perturbation_interval=20,
+                hyperparam_bounds={"lr": [0.0001, 0.1]},
+            )
+            tuner = tune.Tuner(
+                pbt_function,
+                tune_config=tune.TuneConfig(
+                    scheduler=pb2,
+                    num_samples=8,
+                ),
+                param_space={"lr": 0.0001},
+            )
+            tuner.fit()
+
+    """
+
+    def __init__(
+        self,
+        time_attr: str = "time_total_s",
+        metric: Optional[str] = None,
+        mode: Optional[str] = None,
+        perturbation_interval: float = 60.0,
+        hyperparam_bounds: Dict[str, Union[dict, list, tuple]] = None,
+        quantile_fraction: float = 0.25,
+        log_config: bool = True,
+        require_attrs: bool = True,
+        synch: bool = False,
+        custom_explore_fn: Optional[Callable[[dict], dict]] = None,
+    ):
+
+        sklearn_available = import_pb2_dependencies()
+        if not sklearn_available:
+            raise RuntimeError("Please install scikit-learn to use PB2.")
+
+        hyperparam_bounds = hyperparam_bounds or {}
+
+        if not hyperparam_bounds:
+            raise TuneError(
+                "`hyperparam_bounds` must be specified to use PB2 scheduler."
+            )
+
+        super(PB2, self).__init__(
+            time_attr=time_attr,
+            metric=metric,
+            mode=mode,
+            perturbation_interval=perturbation_interval,
+            hyperparam_mutations=hyperparam_bounds,
+            quantile_fraction=quantile_fraction,
+            resample_probability=0,
+            custom_explore_fn=custom_explore_fn,
+            log_config=log_config,
+            require_attrs=require_attrs,
+            synch=synch,
+        )
+
+        self.last_exploration_time = 0  # when we last explored
+        self.data = pd.DataFrame()
+
+        self._hyperparam_bounds = hyperparam_bounds
+        self._hyperparam_bounds_flat = flatten_dict(
+            hyperparam_bounds, prevent_delimiter=True
+        )
+        self._validate_hyperparam_bounds(self._hyperparam_bounds_flat)
+
+        # Current = trials running that have already re-started after reaching
+        #           the checkpoint. When exploring we care if these trials
+        #           are already in or scheduled to be in the next round.
+        self.current = None
+
+    def on_trial_add(self, tune_controller: "TuneController", trial: Trial):
+        filled_hyperparams = _fill_config(trial.config, self._hyperparam_bounds)
+        # Make sure that the params we sampled show up in the CLI output
+        trial.evaluated_params.update(flatten_dict(filled_hyperparams))
+        super().on_trial_add(tune_controller, trial)
+
+    def _validate_hyperparam_bounds(self, hyperparam_bounds: dict):
+        """Check that each hyperparam bound is of the form [low, high].
+
+        Args:
+            hyperparam_bounds: Flattened mapping of hyperparameter name to a
+                ``[low, high]`` pair (or 2-tuple) describing the allowed range.
+
+        Raises:
+            ValueError: if any of the hyperparam bounds are of an invalid format.
+        """
+        for key, value in hyperparam_bounds.items():
+            if not isinstance(value, (list, tuple)) or len(value) != 2:
+                raise ValueError(
+                    "`hyperparam_bounds` values must either be "
+                    f"a list or tuple of size 2, but got {value} "
+                    f"instead for the param '{key}'"
+                )
+            low, high = value
+            if low > high:
+                raise ValueError(
+                    "`hyperparam_bounds` values must be of the form [low, high] "
+                    f"where low <= high, but got {value} instead for param '{key}'."
+                )
+
+    def _save_trial_state(
+        self, state: _PBTTrialState, time: int, result: Dict, trial: Trial
+    ):
+        score = super(PB2, self)._save_trial_state(state, time, result, trial)
+
+        # Data logging for PB2.
+
+        # Collect hyperparams names and current values for this trial.
+        names = list(self._hyperparam_bounds_flat.keys())
+        flattened_config = flatten_dict(trial.config)
+        values = [flattened_config[key] for key in names]
+
+        # Store trial state and hyperparams in dataframe.
+        # this needs to be made more general.
+        lst = [[trial, result[self._time_attr]] + values + [score]]
+        cols = ["Trial", "Time"] + names + ["Reward"]
+        entry = pd.DataFrame(lst, columns=cols)
+
+        self.data = pd.concat([self.data, entry]).reset_index(drop=True)
+        self.data.Trial = self.data.Trial.astype("str")
+
+    def _get_new_config(self, trial: Trial, trial_to_clone: Trial) -> Tuple[Dict, Dict]:
+        """Gets new config for trial by exploring trial_to_clone's config using
+        Bayesian Optimization (BO) to choose the hyperparameter values to explore.
+
+        Overrides `PopulationBasedTraining._get_new_config`.
+
+        Args:
+            trial: The current trial that decided to exploit trial_to_clone.
+            trial_to_clone: The top-performing trial with a hyperparameter config
+                that the current trial will explore.
+
+        Returns:
+            new_config: New hyperparameter configuration (after BO).
+            operations: Empty dict since PB2 doesn't explore in easily labeled ways
+                like PBT does.
+        """
+        # If we are at a new timestep, we dont want to penalise for trials
+        # still going.
+        if self.data["Time"].max() > self.last_exploration_time:
+            self.current = None
+
+        new_config_flat, data = _explore(
+            self.data,
+            self._hyperparam_bounds_flat,
+            self.current,
+            trial_to_clone,
+            trial,
+            flatten_dict(trial_to_clone.config),
+        )
+
+        # Important to replace the old values, since we are copying across
+        self.data = data.copy()
+
+        # If the current guy being selecting is at a point that is already
+        # done, then append the data to the "current" which contains the
+        # points in the current batch.
+        new = [new_config_flat[key] for key in self._hyperparam_bounds_flat]
+
+        new = np.array(new)
+        new = new.reshape(1, new.size)
+        if self.data["Time"].max() > self.last_exploration_time:
+            self.last_exploration_time = self.data["Time"].max()
+            self.current = new.copy()
+        else:
+            self.current = np.concatenate((self.current, new), axis=0)
+            logger.debug(self.current)
+
+        new_config = unflatten_dict(new_config_flat)
+
+        if self._custom_explore_fn:
+            new_config = self._custom_explore_fn(new_config)
+            assert (
+                new_config is not None
+            ), "Custom explore function failed to return a new config"
+
+        return new_config, {}

python/ray/tune/schedulers/pb2_utils.py

@@ -0,0 +1,212 @@
+import numpy as np
+from scipy.optimize import minimize
+from sklearn.gaussian_process import GaussianProcessRegressor
+from sklearn.gaussian_process.kernels import Hyperparameter, Kernel
+from sklearn.metrics import pairwise_distances
+from sklearn.metrics.pairwise import euclidean_distances
+
+
+class TV_SquaredExp(Kernel):
+    """Time varying squared exponential kernel.
+    For more info see the TV-GP-UCB paper:
+    http://proceedings.mlr.press/v51/bogunovic16.pdf
+    """
+
+    def __init__(
+        self,
+        variance=1.0,
+        lengthscale=1.0,
+        epsilon=0.1,
+        variance_bounds=(1e-5, 1e5),
+        lengthscale_bounds=(1e-5, 1e5),
+        epsilon_bounds=(1e-5, 0.5),
+    ):
+        self.variance = variance
+        self.lengthscale = lengthscale
+        self.epsilon = epsilon
+        self.variance_bounds = variance_bounds
+        self.lengthscale_bounds = lengthscale_bounds
+        self.epsilon_bounds = epsilon_bounds
+
+    @property
+    def hyperparameter_variance(self):
+        return Hyperparameter("variance", "numeric", self.variance_bounds)
+
+    @property
+    def hyperparameter_lengthscale(self):
+        return Hyperparameter("lengthscale", "numeric", self.lengthscale_bounds)
+
+    @property
+    def hyperparameter_epsilon(self):
+        return Hyperparameter("epsilon", "numeric", self.epsilon_bounds)
+
+    def __call__(self, X, Y=None, eval_gradient=False):
+        X = np.atleast_2d(X)
+        if Y is None:
+            Y = X
+
+        epsilon = np.clip(self.epsilon, 1e-5, 0.5)
+
+        # Time must be in the first column
+        T1 = X[:, 0].reshape(-1, 1)
+        T2 = Y[:, 0].reshape(-1, 1)
+        dists = pairwise_distances(T1, T2, "cityblock")
+        timekernel = (1 - epsilon) ** (0.5 * dists)
+
+        # RBF kernel on remaining features
+        X_spatial = X[:, 1:]
+        Y_spatial = Y[:, 1:]
+        rbf = self.variance * np.exp(
+            -np.square(euclidean_distances(X_spatial, Y_spatial)) / self.lengthscale
+        )
+
+        K = rbf * timekernel
+
+        if eval_gradient:
+            K_gradient_variance = K
+            dist2 = np.square(euclidean_distances(X_spatial, Y_spatial))
+            K_gradient_lengthscale = K * dist2 / self.lengthscale
+            n = dists / 2
+            K_gradient_epsilon = -K * n * epsilon / (1 - epsilon)
+            return K, np.dstack(
+                [K_gradient_variance, K_gradient_lengthscale, K_gradient_epsilon]
+            )
+
+        return K
+
+    def diag(self, X):
+        return np.full(X.shape[0], self.variance, dtype=np.float64)
+
+    def is_stationary(self):
+        return False
+
+    @property
+    def theta(self):
+        return np.log([self.variance, self.lengthscale, self.epsilon])
+
+    @theta.setter
+    def theta(self, theta):
+        self.variance = np.exp(theta[0])
+        self.lengthscale = np.exp(theta[1])
+        self.epsilon = np.exp(theta[2])
+
+    @property
+    def bounds(self):
+        return np.log(
+            [
+                list(self.variance_bounds),
+                list(self.lengthscale_bounds),
+                list(self.epsilon_bounds),
+            ]
+        )
+
+
+def normalize(data, wrt):
+    """Normalize data to be in range (0,1), with respect to (wrt) boundaries,
+    which can be specified.
+    """
+    return (data - np.min(wrt, axis=0)) / (
+        np.max(wrt, axis=0) - np.min(wrt, axis=0) + 1e-8
+    )
+
+
+def standardize(data):
+    """Standardize to be Gaussian N(0,1). Clip final values."""
+    data = (data - np.mean(data, axis=0)) / (np.std(data, axis=0) + 1e-8)
+    return np.clip(data, -2, 2)
+
+
+def UCB(m, m1, x, fixed, kappa=None):
+    """UCB acquisition function. Interesting points to note:
+    1) We concat with the fixed points, because we are not optimizing wrt
+       these. This is the Reward and Time, which we can't change. We want
+       to find the best hyperparameters *given* the reward and time.
+    2) We use m to get the mean and m1 to get the variance. If we already
+       have trials running, then m1 contains this information. This reduces
+       the variance at points currently running, even if we don't have
+       their label.
+       Ref: https://jmlr.org/papers/volume15/desautels14a/desautels14a.pdf
+
+    """
+    c1 = 0.2
+    c2 = 0.4
+    beta_t = c1 + max(0, np.log(c2 * m.X_train_.shape[0]))
+    kappa = np.sqrt(beta_t) if kappa is None else kappa
+
+    xtest = np.concatenate((fixed.reshape(-1, 1), np.array(x).reshape(-1, 1))).T
+
+    try:
+        mean = m.predict(xtest)[0]
+    except ValueError:
+        mean = -9999
+
+    try:
+        _, std = m1.predict(xtest, return_std=True)
+        var = std[0] ** 2
+    except ValueError:
+        var = 0
+    return mean + kappa * var
+
+
+def optimize_acq(func, m, m1, fixed, num_f):
+    """Optimize acquisition function."""
+
+    opts = {"maxiter": 200, "maxfun": 200, "disp": False}
+
+    T = 10
+    best_value = -999
+    best_theta = m1.X_train_[0, :]
+
+    bounds = [(0, 1) for _ in range(m.X_train_.shape[1] - num_f)]
+
+    for ii in range(T):
+        x0 = np.random.uniform(0, 1, m.X_train_.shape[1] - num_f)
+
+        res = minimize(
+            lambda x: -func(m, m1, x, fixed),
+            x0,
+            bounds=bounds,
+            method="L-BFGS-B",
+            options=opts,
+        )
+
+        val = func(m, m1, res.x, fixed)
+        if val > best_value:
+            best_value = val
+            best_theta = res.x
+
+    return np.clip(best_theta, 0, 1)
+
+
+def select_length(Xraw, yraw, bounds, num_f):
+    """Select the number of datapoints to keep, using cross validation"""
+    min_len = 200
+
+    if Xraw.shape[0] < min_len:
+        return Xraw.shape[0]
+    else:
+        length = min_len - 10
+        scores = []
+        while length + 10 <= Xraw.shape[0]:
+            length += 10
+
+            base_vals = np.array(list(bounds.values())).T
+            X_len = Xraw[-length:, :]
+            y_len = yraw[-length:]
+            oldpoints = X_len[:, :num_f]
+            old_lims = np.concatenate(
+                (np.max(oldpoints, axis=0), np.min(oldpoints, axis=0))
+            ).reshape(2, oldpoints.shape[1])
+            limits = np.concatenate((old_lims, base_vals), axis=1)
+
+            X = normalize(X_len, limits)
+            y = standardize(y_len).reshape(y_len.size, 1)
+
+            kernel = TV_SquaredExp(variance=1.0, lengthscale=1.0, epsilon=0.1)
+            m = GaussianProcessRegressor(kernel=kernel, optimizer="fmin_l_bfgs_b")
+            m.fit(X, y)
+
+            scores.append(m.log_marginal_likelihood_value_)
+        idx = np.argmax(scores)
+        length = (idx + int((min_len / 10))) * 10
+        return length

python/ray/tune/schedulers/resource_changing_scheduler.py

@@ -0,0 +1,875 @@
+import logging
+import pickle
+import warnings
+from copy import deepcopy
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Set, Tuple, Union
+
+import numpy as np
+
+from ray.air.execution.resources.request import _sum_bundles
+from ray.tune.execution.placement_groups import PlacementGroupFactory
+from ray.tune.experiment import Trial
+from ray.tune.schedulers.trial_scheduler import FIFOScheduler, TrialScheduler
+from ray.util.annotations import PublicAPI
+
+if TYPE_CHECKING:
+    from ray.tune.execution.tune_controller import TuneController
+
+logger = logging.getLogger(__name__)
+
+
+@PublicAPI(stability="beta")
+class DistributeResources:
+    """This class creates a basic uniform resource allocation function.
+
+    The function naively balances free resources (CPUs and GPUs) between
+    trials, giving them all equal priority, ensuring that all resources
+    are always being used. The free resources will be placed in new bundles.
+    The function assumes that all bundles are equal (there is no "head"
+    bundle).
+
+    If for some reason a trial ends up with
+    more resources than there are free ones, it will adjust downwards.
+    It will also ensure that trial as at least as many resources as
+    it started with (``base_trial_resource``).
+
+    The 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).
+
+    If you wish to implement your own resource distribution logic,
+    you can do so by extending this class, as it provides several
+    generic methods. You can also implement a function instead.
+
+    Args:
+        add_bundles: If True, create new bundles from free resources.
+            Otherwise, spread them among base_trial_resource bundles.
+        increase_by: A dict with key-value
+            pairs representing an atomic unit of resources (name-amount)
+            the trial will be increased by. If not set, the trial will
+            increase by 1 CPU/GPU.
+        increase_by_times: If set to >=1 and ``increase_by`` is set,
+            the trial will increase by maximum of
+            ``increase_by_times * increase_by`` resources. If set to <1,
+            no upper limit is set. Ignored if ``increase_by`` is not set.
+        reserve_resources: A dict of
+            resource_name-amount pairs representing the resources
+            that will not be allocated to resized trials.
+    """
+
+    def __init__(
+        self,
+        add_bundles: bool = False,
+        increase_by: Optional[Dict[str, float]] = None,
+        increase_by_times: int = -1,
+        reserve_resources: Optional[Dict[str, float]] = None,
+    ):
+        self.add_bundles = add_bundles
+        self.increase_by = increase_by or {}
+        self.increase_by_times = increase_by_times
+        self.reserve_resources = reserve_resources or {}
+
+    def _validate(
+        self, base_trial_resource: PlacementGroupFactory, result: Dict[str, Any]
+    ) -> bool:
+        """Return False if we should keep the current resources outright."""
+        if not isinstance(base_trial_resource, PlacementGroupFactory):
+            raise ValueError(
+                f"{self.__class__.__name__} only supports PlacementGroupFactories."
+            )
+
+        if not self.add_bundles and len(base_trial_resource.bundles) > 1:
+            raise ValueError(
+                "If `add_bundles` is False, the number of bundles in "
+                "`resources_per_trial` must be 1 "
+                f"(got {len(base_trial_resource.bundles)})."
+            )
+
+        # Don't bother if this is just the first iteration
+        if result["training_iteration"] < 1:
+            return False
+        return True
+
+    def _get_total_available_resources(
+        self, tune_controller: "TuneController"
+    ) -> Tuple[float, float]:
+        """Get the number of CPUs and GPUs avaialble in total (not just free)"""
+        total_available_cpus = (
+            tune_controller._resource_updater.get_num_cpus()
+            - self.reserve_resources.get("CPU", 0)
+        )
+        total_available_gpus = (
+            tune_controller._resource_updater.get_num_gpus()
+            - self.reserve_resources.get("GPU", 0)
+        )
+        return total_available_cpus, total_available_gpus
+
+    def _get_used_cpus_and_gpus(self, t: Trial) -> Tuple[float, float]:
+        """Check how many CPUs and GPUs a trial is using currently"""
+        return (
+            t.placement_group_factory.required_resources.get("CPU", 0),
+            t.placement_group_factory.required_resources.get("GPU", 0),
+        )
+
+    def _get_resources_from_bundles(
+        self, bundles: List[Dict[str, float]]
+    ) -> Dict[str, float]:
+        """Get total sums of resources in bundles"""
+        if not bundles:
+            return {"CPU": 0, "GPU": 0}
+        return _sum_bundles(bundles)
+
+    def _is_bundle_empty(self, bundle: Dict[str, float]) -> bool:
+        return not (bundle.get("CPU", 0) or bundle.get("GPU", 0))
+
+    def _add_two_bundles(
+        self,
+        bundles_a: List[Dict[str, float]],
+        bundles_b: List[Dict[str, float]],
+        increase_by: Dict[str, float],
+        limit_to_increase_by_times: bool,
+        max_increase_by_times: int = -1,
+    ):
+        """Add two bundles together.
+
+        If ``limit_to_increase_by_times`` is True, ``self.increase_by_times`` > 0
+        and ``max_increase_by_times`` > 0, ensure that the resulting number of
+        bundles is not above ``min(max_increase_by_times, self.increase_by_times)``.
+
+        If ``limit_to_increase_by_times`` is True and ``self.increase_by_times`` > 0,
+        ensure that the resulting number of bundles is not above
+        `self.increase_by_times``.
+        """
+        if limit_to_increase_by_times:
+            if max_increase_by_times > 0 and self.increase_by_times > 0:
+                max_increase_by_times = min(
+                    max_increase_by_times, self.increase_by_times
+                )
+            elif self.increase_by_times > 0:
+                max_increase_by_times = self.increase_by_times
+
+        if self.add_bundles:
+            bundles = [b for b in bundles_a if not self._is_bundle_empty(b)] + [
+                b for b in bundles_b if not self._is_bundle_empty(b)
+            ]
+            if max_increase_by_times > 0:
+                bundles = bundles[:max_increase_by_times]
+        else:
+            bundles_a = bundles_a or [{}]
+            bundles_b = bundles_b or [{}]
+            bundles = [
+                {
+                    "CPU": bundles_a[0].get("CPU", 0) + bundles_b[0].get("CPU", 0),
+                    "GPU": bundles_a[0].get("GPU", 0) + bundles_b[0].get("GPU", 0),
+                }
+            ]
+            if max_increase_by_times > 0:
+                bundles[0]["CPU"] = min(
+                    bundles[0]["CPU"],
+                    increase_by.get("CPU", 0) * max_increase_by_times,
+                )
+                bundles[0]["GPU"] = min(
+                    bundles[0]["GPU"],
+                    increase_by.get("GPU", 0) * max_increase_by_times,
+                )
+
+        return bundles
+
+    def _get_multiplier(
+        self,
+        increase_by: Dict[str, float],
+        cpus: float = 0,
+        gpus: float = 0,
+        max_multiplier: int = -1,
+    ) -> int:
+        """Get how many times ``increase_by`` bundles
+        occur in ``cpus`` and ``gpus``."""
+        if increase_by.get("CPU", 0) and increase_by.get("GPU", 0):
+            multiplier = min(
+                cpus // increase_by.get("CPU", 0),
+                gpus // increase_by.get("GPU", 0),
+            )
+        elif increase_by.get("GPU", 0):
+            multiplier = gpus // increase_by.get("GPU", 0)
+        else:
+            multiplier = cpus // increase_by.get("CPU", 0)
+
+        if max_multiplier > 0 and multiplier > 0:
+            multiplier = min(max_multiplier, multiplier)
+        return int(multiplier)
+
+    def _remove_bundles(
+        self,
+        bundles: List[Dict[str, float]],
+        increase_by: Dict[str, float],
+        multiplier: int,
+    ) -> List[Dict[str, float]]:
+        """Remove ``multiplier`` ``increase_by`` bundles from ``bundles``."""
+        multiplier = -abs(multiplier)
+        if self.add_bundles:
+            bundles = bundles[:multiplier]
+        else:
+            bundles = deepcopy(bundles)
+            bundles[0]["CPU"] += increase_by.get("CPU", 0) * multiplier
+            bundles[0]["GPU"] += increase_by.get("GPU", 0) * multiplier
+            bundles[0]["CPU"] = max(bundles[0]["CPU"], 0)
+            bundles[0]["GPU"] = max(bundles[0]["GPU"], 0)
+        return bundles
+
+    def _create_new_bundles(
+        self,
+        increase_by: Dict[str, float],
+        multiplier: int,
+    ) -> List[Dict[str, float]]:
+        """Create a list of new bundles containing ``increase_by`` * ``multiplier``."""
+        multiplier = abs(multiplier)
+
+        if self.add_bundles:
+            bundles = [increase_by] * int(multiplier)
+        else:
+            bundles = [{}]
+            bundles[0]["CPU"] = increase_by.get("CPU", 0) * multiplier
+            bundles[0]["GPU"] = increase_by.get("GPU", 0) * multiplier
+
+        return bundles
+
+    def _modify_bundles_with_free_resources(
+        self,
+        bundles: List[Dict[str, float]],
+        increase_by: Dict[str, float],
+        free_cpus: float,
+        free_gpus: float,
+        *,
+        max_multiplier: int = -1,
+        max_increase_by_times: int = -1,
+    ):
+        """Given free resources, increase/decrease the number of bundles in
+        ``bundles``."""
+        multiplier = self._get_multiplier(
+            increase_by, free_cpus, free_gpus, max_multiplier
+        )
+        if multiplier < 0:
+            bundles = self._remove_bundles(bundles, increase_by, multiplier)
+        elif multiplier > 0:
+            bundles_to_add = self._create_new_bundles(increase_by, multiplier)
+            bundles = self._add_two_bundles(
+                bundles, bundles_to_add, increase_by, True, max_increase_by_times
+            )
+        return bundles
+
+    def _get_added_bundles(
+        self, bundles: List[Dict[str, float]], base_bundles: List[Dict[str, float]]
+    ) -> List[Dict[str, float]]:
+        """Return the difference between bundles and base_bundles"""
+        if self.add_bundles:
+            added_bundles = bundles[len(base_bundles) :]
+        else:
+            if not bundles:
+                bundles = [{"CPU": 0, "GPU": 0}]
+            if not base_bundles:
+                base_bundles = [{"CPU": 0, "GPU": 0}]
+            added_bundles = [
+                {
+                    "CPU": bundles[0].get("CPU", 0) - base_bundles[0].get("CPU", 0),
+                    "GPU": bundles[0].get("GPU", 0) - base_bundles[0].get("GPU", 0),
+                }
+            ]
+        return added_bundles
+
+    def _are_bundles_below_limit(
+        self,
+        bundles: List[Dict[str, float]],
+        base_bundles: Optional[List[Dict[str, float]]] = None,
+        max_added_cpus: Optional[float] = None,
+        max_added_gpus: Optional[float] = None,
+    ):
+        if not max_added_cpus:
+            if self.increase_by_times > 0:
+                max_added_cpus = self.increase_by.get("CPU", 0) * self.increase_by_times
+            else:
+                max_added_cpus = np.inf
+        if not max_added_gpus:
+            if self.increase_by_times > 0:
+                max_added_gpus = self.increase_by.get("GPU", 0) * self.increase_by_times
+            else:
+                max_added_gpus = np.inf
+        added_resources = self._get_resources_from_bundles(
+            self._get_added_bundles(bundles, base_bundles) if base_bundles else bundles
+        )
+        ret = (
+            added_resources.get("CPU", -np.inf) < max_added_cpus
+            or added_resources.get("GPU", -np.inf) < max_added_gpus
+        )
+        return ret
+
+    def _get_new_added_bundles(
+        self,
+        trial: Trial,
+        all_trials: List[Trial],
+        base_bundles: List[Dict[str, float]],
+        increase_by: Dict[str, float],
+        total_available_cpus: float,
+        total_available_gpus: float,
+        used_cpus: float,
+        used_gpus: float,
+    ) -> List[Dict[str, float]]:
+        """Returns updated added bundles."""
+        upper_limit_all_trials_bundles = [list() for _ in range(len(all_trials))]
+
+        free_cpus = total_available_cpus - used_cpus
+        free_gpus = total_available_gpus - used_gpus
+
+        base_resources = self._get_resources_from_bundles(base_bundles)
+        upper_limit_cpus_to_distribute = total_available_cpus - (
+            base_resources.get("CPU", 0) * len(all_trials)
+        )
+        upper_limit_gpus_to_distribute = total_available_gpus - (
+            base_resources.get("GPU", 0) * len(all_trials)
+        )
+        max_increase_by_times = 0
+
+        # First, calculate upper limits for uniform allocation
+        # This is done by simulating a clean slate scenario
+        # The loop runs until all resources are allocated or
+        # all trials are at their resource limits
+        i = 0
+        trials_at_limit = set()
+        while (
+            len(trials_at_limit) < len(all_trials)
+            # we have previously asserted that at least one resource has to be
+            # bigger than 0
+            and upper_limit_cpus_to_distribute >= increase_by.get("CPU", 0)
+            and upper_limit_gpus_to_distribute >= increase_by.get("GPU", 0)
+        ):
+            idx = i % len(upper_limit_all_trials_bundles)
+            old_bundles = deepcopy(upper_limit_all_trials_bundles[idx])
+            upper_limit_all_trials_bundles[
+                idx
+            ] = self._modify_bundles_with_free_resources(
+                upper_limit_all_trials_bundles[idx],
+                increase_by,
+                upper_limit_cpus_to_distribute,
+                upper_limit_gpus_to_distribute,
+                max_multiplier=1,
+            )
+            added_resources = self._get_resources_from_bundles(
+                self._get_added_bundles(
+                    upper_limit_all_trials_bundles[idx], old_bundles
+                )
+            )
+            if not added_resources.get("CPU", 0) and not added_resources.get("GPU", 0):
+                trials_at_limit.add(idx)
+            elif idx == 0:
+                max_increase_by_times += 1
+            upper_limit_cpus_to_distribute -= added_resources.get("CPU", 0)
+            upper_limit_gpus_to_distribute -= added_resources.get("GPU", 0)
+            i += 1
+
+        # Add new resourcs, but only up to calculated upper limits
+        # (max_increase_by_times)
+        return self._modify_bundles_with_free_resources(
+            self._get_added_bundles(
+                trial.placement_group_factory.bundles, base_bundles
+            ),
+            increase_by,
+            free_cpus,
+            free_gpus,
+            max_increase_by_times=max_increase_by_times,
+        )
+
+    def __call__(
+        self,
+        tune_controller: "TuneController",
+        trial: Trial,
+        result: Dict[str, Any],
+        scheduler: "ResourceChangingScheduler",
+    ) -> Optional[PlacementGroupFactory]:
+        """Run resource allocation logic.
+
+        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).
+
+        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 updated resource requirements,
+            or None if the trial's resources should be left unchanged.
+        """
+        # Get base trial resources as defined in
+        # ``tune.run(resources_per_trial)``
+        base_trial_resource = scheduler.base_trial_resources
+
+        if not self._validate(base_trial_resource=base_trial_resource, result=result):
+            return None
+
+        # default values if resources_per_trial is unspecified
+        if base_trial_resource is None:
+            base_trial_resource = PlacementGroupFactory([{"CPU": 1, "GPU": 0}])
+
+        if self.increase_by:
+            increase_by = self.increase_by
+            assert not self._is_bundle_empty(increase_by)
+            assert increase_by.get("CPU", 0) >= 0 and increase_by.get("GPU", 0) >= 0
+        elif self.add_bundles:
+            increase_by = base_trial_resource.bundles[-1]
+        elif base_trial_resource.bundles[0].get("GPU", 0):
+            increase_by = {"GPU": 1}
+        else:
+            increase_by = {"CPU": 1}
+
+        base_bundles = deepcopy(base_trial_resource.bundles)
+
+        (
+            total_available_cpus,
+            total_available_gpus,
+        ) = self._get_total_available_resources(tune_controller=tune_controller)
+
+        all_trials = tune_controller.get_live_trials()
+
+        used_cpus_and_gpus = [self._get_used_cpus_and_gpus(t) for t in all_trials]
+        used_cpus, used_gpus = zip(*used_cpus_and_gpus)
+        used_cpus = sum(used_cpus)
+        used_gpus = sum(used_gpus)
+
+        added_bundles = self._get_new_added_bundles(
+            trial,
+            all_trials,
+            base_bundles,
+            increase_by,
+            total_available_cpus,
+            total_available_gpus,
+            used_cpus,
+            used_gpus,
+        )
+
+        new_bundles = self._add_two_bundles(
+            base_bundles, added_bundles, increase_by, False
+        )
+
+        pgf = PlacementGroupFactory(
+            new_bundles,
+            strategy=base_trial_resource.strategy,
+            *base_trial_resource._args,
+            **base_trial_resource._kwargs,
+        )
+        pgf._head_bundle_is_empty = base_trial_resource._head_bundle_is_empty
+        return pgf
+
+
+@PublicAPI(stability="beta")
+class DistributeResourcesToTopJob(DistributeResources):
+    """This class creates a "TopJob" resource allocation function.
+
+    The function will assign all of the free resources to the best
+    performing trial (as defined by ``metric`` and ``mode``). The
+    previous best trials will not have their resources deallocated,
+    unless in the case outlined below.
+
+    If for some reason a trial ends up with
+    more resources than there are free ones, it will adjust downwards.
+    It will also ensure that trial as at least as many resources as
+    it started with (``base_trial_resource``).
+
+    The 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).
+
+    Args:
+        add_bundles: If True, create new bundles from free resources.
+            Otherwise, spread them among base_trial_resource bundles.
+        increase_by: A dict with key-value
+            pairs representing an atomic unit of resources (name-amount)
+            the trial will be increased by. If not set, the trial will
+            increase by 1 CPU/GPU.
+        increase_by_times: If set to >=1 and ``increase_by`` is set,
+            the trial will increase by maximum of
+            ``increase_by_times * increase_by`` resources. If set to <1,
+            no upper limit is set. Ignored if ``increase_by`` is not set.
+        reserve_resources: A dict of
+            resource_name-amount pairs representing the resources
+            that will not be allocated to resized trials.
+            is that the attribute should increase monotonically.
+        metric: The training result objective value attribute. Stopping
+            procedures will use this attribute. If None, will use the metric
+            of the scheduler.
+        mode: One of {min, max}. Determines whether objective is
+            minimizing or maximizing the metric attribute. If None, will use the metric
+            of the scheduler.
+
+    """
+
+    def __init__(
+        self,
+        add_bundles: bool = False,
+        increase_by: Optional[Dict[str, float]] = None,
+        increase_by_times: int = -1,
+        reserve_resources: Optional[Dict[str, float]] = None,
+        metric: Optional[str] = None,
+        mode: Optional[str] = None,
+    ):
+        super().__init__(add_bundles, increase_by, increase_by_times, reserve_resources)
+        self.metric = metric
+        self.mode = mode
+
+    @property
+    def _metric_op(self) -> float:
+        if self.mode not in ("min", "max"):
+            raise ValueError("The mode parameter can only be either min or max.")
+        if self.mode == "max":
+            return 1.0
+        return -1.0
+
+    def _get_new_added_bundles(
+        self,
+        trial: Trial,
+        all_trials: List[Trial],
+        base_bundles: List[Dict[str, float]],
+        increase_by: Dict[str, float],
+        total_available_cpus: float,
+        total_available_gpus: float,
+        used_cpus: float,
+        used_gpus: float,
+    ) -> List[Dict[str, float]]:
+        if self.metric is None:
+            raise ValueError(
+                "The metric parameter cannot be None. The parameter can be set in "
+                "either `DistributeResourcesToTopJob`, the base scheduler or in "
+                "`tune.TuneConfig()` (highest to lowest priority)."
+            )
+
+        free_cpus = total_available_cpus - used_cpus
+        free_gpus = total_available_gpus - used_gpus
+
+        sorted_trials = sorted(
+            all_trials,
+            key=lambda t: -self._metric_op * t.last_result.get(self.metric, np.inf),
+        )
+
+        added_bundles = self._get_added_bundles(
+            trial.placement_group_factory.bundles, base_bundles
+        )
+
+        best_trial = next(
+            (
+                t
+                for t in sorted_trials
+                if self._are_bundles_below_limit(
+                    t.placement_group_factory.bundles, base_bundles
+                )
+            ),
+            sorted_trials[0],
+        )
+
+        if (
+            trial.trial_id != best_trial.trial_id
+            # Only reduce resources here
+            and self._get_multiplier(increase_by, free_cpus, free_gpus) >= 0
+        ):
+            return added_bundles
+
+        return self._modify_bundles_with_free_resources(
+            added_bundles,
+            increase_by,
+            free_cpus,
+            free_gpus,
+        )
+
+
+_DistributeResourcesDefault = DistributeResources(add_bundles=False)
+_DistributeResourcesDistributedDefault = DistributeResources(add_bundles=True)
+
+
+@PublicAPI(stability="beta")
+class ResourceChangingScheduler(TrialScheduler):
+    """A utility scheduler to dynamically change resources of live trials.
+
+    .. versionadded:: 1.5.0
+
+    .. note::
+        Experimental. API may change in future releases.
+
+    The ResourceChangingScheduler works by wrapping around any other
+    scheduler and adjusting the resource requirements of live trials
+    in response to the decisions of the wrapped scheduler
+    through a user-specified ``resources_allocation_function``.
+
+    An example of such a function can be found in
+    :doc:`/tune/examples/includes/xgboost_dynamic_resources_example`.
+
+    If the functional API is used, the current trial resources can be obtained
+    by calling `tune.get_trial_resources()` inside the training function.
+    The function should be able to
+    :ref:`load and save checkpoints <tune-function-trainable-checkpointing>`
+    (the latter preferably every iteration).
+
+    If the Trainable (class) API is used, you can obtain the current trial
+    resources through the ``Trainable.trial_resources`` property.
+
+    Cannot be used if ``reuse_actors`` is True in ``tune.TuneConfig()``. A ValueError
+    will be raised in that case.
+
+    Args:
+        base_scheduler: The scheduler to provide decisions
+            about trials. If None, a default FIFOScheduler will be used.
+        resources_allocation_function: The callable used to change
+            live trial resource requiements during tuning. This callable
+            will be called on each trial as it finishes one step of training.
+            The callable must take four arguments: ``TrialRunner``, current
+            ``Trial``, current result :class:`dict` and the
+            ``ResourceChangingScheduler`` calling it. The callable must
+            return a ``PlacementGroupFactory``
+            or None (signifying no need for an update). If
+            ``resources_allocation_function`` is None, no resource
+            requirements will be changed at any time.
+            By default, :class:`DistributeResources` will be used,
+            distributing available CPUs and GPUs over all running trials
+            in a robust way, without any prioritization.
+
+    Warning:
+        If the ``resources_allocation_function`` sets trial resource
+        requirements to values bigger than possible, the trial will
+        not run. Ensure that your callable accounts for that possibility
+        by setting upper limits. Consult :class:`DistributeResources`
+        to see how that may be done.
+
+    Example:
+        .. code-block:: python
+
+            base_scheduler = ASHAScheduler(max_t=16)
+            def my_resources_allocation_function(
+                tune_controller: "TuneController",
+                trial: Trial,
+                result: Dict[str, Any],
+                scheduler: "ResourceChangingScheduler"
+            ) -> Optional[Union[PlacementGroupFactory, Resource]]:
+                # logic here
+                # usage of PlacementGroupFactory is strongly preferred
+                return PlacementGroupFactory(...)
+            scheduler = ResourceChangingScheduler(
+                            base_scheduler,
+                            my_resources_allocation_function
+                        )
+
+        See :doc:`/tune/examples/includes/xgboost_dynamic_resources_example` for a
+        more detailed example.
+    """
+
+    def __init__(
+        self,
+        base_scheduler: Optional[TrialScheduler] = None,
+        resources_allocation_function: Optional[
+            Callable[
+                [
+                    "TuneController",
+                    Trial,
+                    Dict[str, Any],
+                    "ResourceChangingScheduler",
+                ],
+                Optional[PlacementGroupFactory],
+            ]
+        ] = _DistributeResourcesDefault,
+    ) -> None:
+        super().__init__()
+        if resources_allocation_function is None:
+            warnings.warn(
+                "`resources_allocation_function` is None. No resource "
+                "requirements will be changed at any time. Pass a "
+                "correctly defined function to enable functionality."
+            )
+        self._resources_allocation_function = resources_allocation_function
+        self._base_scheduler = base_scheduler or FIFOScheduler()
+        self._base_trial_resources: Optional[PlacementGroupFactory] = None
+        self._trials_to_reallocate: Dict[
+            Trial, Optional[Union[dict, PlacementGroupFactory]]
+        ] = {}
+        self._reallocated_trial_ids: Set[str] = set()
+        self._metric = None
+        self._mode = None
+
+    @property
+    def metric(self):
+        return self._base_scheduler._metric
+
+    @property
+    def base_trial_resources(self) -> Optional[PlacementGroupFactory]:
+        return self._base_trial_resources
+
+    def set_search_properties(
+        self, metric: Optional[str], mode: Optional[str], **spec
+    ) -> bool:
+        self._metric = metric
+        self._mode = mode
+        return self._base_scheduler.set_search_properties(metric, mode, **spec)
+
+    def on_trial_add(self, tune_controller: "TuneController", trial: Trial, **kwargs):
+        # use the first trial resources as the base
+        if self._base_trial_resources is None:
+            self._base_trial_resources = trial.placement_group_factory
+        # Raise error if the resources of a newly added trial don't match
+        # base resources, but allow trials that have already had their
+        # resources changed by ResourceChangingScheduler
+        # (those can be added again during loading from a checkpoint)
+        elif trial.trial_id not in self._reallocated_trial_ids:
+            trial_resources = trial.placement_group_factory
+            if trial_resources != self._base_trial_resources:
+                raise RuntimeError(
+                    "ResourceChangingScheduler doesn't support trials with "
+                    "varying base resources. First trial had "
+                    f"{self._base_trial_resources}, trial {trial} has "
+                    f"{trial_resources}."
+                )
+
+        return self._base_scheduler.on_trial_add(tune_controller, trial, **kwargs)
+
+    def on_trial_error(self, tune_controller: "TuneController", trial: Trial, **kwargs):
+        return self._base_scheduler.on_trial_error(tune_controller, trial, **kwargs)
+
+    def on_trial_result(
+        self, tune_controller: "TuneController", trial: Trial, result: Dict
+    ) -> str:
+        base_scheduler_decision = self._base_scheduler.on_trial_result(
+            tune_controller, trial, result
+        )
+        if base_scheduler_decision == TrialScheduler.CONTINUE:
+            new_resources = self.reallocate_trial_resources_if_needed(
+                tune_controller, trial, result
+            )
+            if new_resources:
+                self._trials_to_reallocate[trial] = new_resources
+                return TrialScheduler.PAUSE
+        return base_scheduler_decision
+
+    def on_trial_complete(
+        self,
+        tune_controller: "TuneController",
+        trial: Trial,
+        result: Dict,
+        **kwargs,
+    ):
+        return self._base_scheduler.on_trial_complete(
+            tune_controller, trial, result, **kwargs
+        )
+
+    def on_trial_remove(
+        self, tune_controller: "TuneController", trial: Trial, **kwargs
+    ):
+        return self._base_scheduler.on_trial_remove(tune_controller, trial, **kwargs)
+
+    def choose_trial_to_run(
+        self, tune_controller: "TuneController", **kwargs
+    ) -> Optional[Trial]:
+        if getattr(tune_controller, "_reuse_actors", False):
+            raise ValueError(
+                "ResourceChangingScheduler cannot be used with "
+                "`reuse_actors=True`. FIX THIS by setting "
+                "`reuse_actors=False` in `tune.TuneConfig()`."
+            )
+
+        any_resources_changed = False
+
+        new_trials_to_reallocate = {}
+        for trial, new_resources in self._trials_to_reallocate.items():
+            if trial.status == Trial.RUNNING:
+                new_trials_to_reallocate[trial] = new_resources
+                logger.debug(f"{trial} is still running, skipping for now")
+                continue
+            any_resources_changed = any_resources_changed or self.set_trial_resources(
+                trial, new_resources
+            )
+        self._trials_to_reallocate = new_trials_to_reallocate
+
+        trial = self._base_scheduler.choose_trial_to_run(tune_controller, **kwargs)
+        return trial
+
+    def debug_string(self) -> str:
+        return "(ResourceChangingScheduler) " f"{self._base_scheduler.debug_string()}"
+
+    def save(self, checkpoint_path: str):
+        save_object = self.__dict__
+        with open(checkpoint_path, "wb") as outputFile:
+            pickle.dump(save_object, outputFile)
+
+    def restore(self, checkpoint_path: str):
+        with open(checkpoint_path, "rb") as inputFile:
+            save_object = pickle.load(inputFile)
+        self.__dict__.update(save_object)
+
+    def set_trial_resources(
+        self, trial: Trial, new_resources: Union[Dict, PlacementGroupFactory]
+    ) -> bool:
+        """Returns True if new_resources were set."""
+        if new_resources:
+            logger.info(
+                f"Setting trial {trial} resource to {new_resources} "
+                f"with {new_resources._bundles}"
+            )
+            trial.placement_group_factory = None
+            trial.update_resources(new_resources)
+            # keep track of all trials which had their resources changed
+            self._reallocated_trial_ids.add(trial.trial_id)
+            return True
+        return False
+
+    def _are_resources_the_same(
+        self,
+        trial: Trial,
+        new_resources,
+    ) -> bool:
+        """Returns True if trial's resources are value equal to new_resources.
+
+        Only checks for PlacementGroupFactories at this moment.
+        """
+        if (
+            isinstance(new_resources, PlacementGroupFactory)
+            and trial.placement_group_factory == new_resources
+        ):
+            logger.debug(
+                f"{trial} PGF "
+                f"{trial.placement_group_factory.required_resources}"
+                f" and {new_resources.required_resources}"
+                f" are the same, skipping"
+            )
+            return True
+        else:
+            return False
+
+    def reallocate_trial_resources_if_needed(
+        self, tune_controller: "TuneController", trial: Trial, result: Dict
+    ) -> Optional[Union[dict, PlacementGroupFactory]]:
+        """Calls user defined resources_allocation_function. If the returned
+        resources are not none and not the same as currently present, returns
+        them. Otherwise, returns None."""
+        if self._resources_allocation_function is None:
+            return None
+
+        if not getattr(self._resources_allocation_function, "metric", None):
+            self._resources_allocation_function.metric = getattr(
+                self._base_scheduler, "_metric", self._metric
+            )
+        if not getattr(self._resources_allocation_function, "mode", None):
+            self._resources_allocation_function.mode = getattr(
+                self._base_scheduler, "_mode", self._mode
+            )
+
+        new_resources = self._resources_allocation_function(
+            tune_controller, trial, result, self
+        )
+
+        # if we can check if the new resources are the same,
+        # we do that here and skip resource allocation
+        if new_resources and not self._are_resources_the_same(trial, new_resources):
+            return new_resources
+        return None

python/ray/tune/schedulers/trial_scheduler.py

@@ -0,0 +1,173 @@
+from typing import TYPE_CHECKING, Dict, Optional
+
+from ray.air._internal.usage import tag_scheduler
+from ray.tune.experiment import Trial
+from ray.tune.result import DEFAULT_METRIC
+from ray.util.annotations import DeveloperAPI, PublicAPI
+
+if TYPE_CHECKING:
+    from ray.tune.execution.tune_controller import TuneController
+
+
+@DeveloperAPI
+class TrialScheduler:
+    """Interface for implementing a Trial Scheduler class.
+
+    Note to Tune developers: If a new scheduler is added, please update
+    `air/_internal/usage.py`.
+    """
+
+    CONTINUE = "CONTINUE"  #: Status for continuing trial execution
+    PAUSE = "PAUSE"  #: Status for pausing trial execution
+    STOP = "STOP"  #: Status for stopping trial execution
+    # Caution: Temporary and anti-pattern! This means Scheduler calls
+    # into Executor directly without going through TrialRunner.
+    # TODO(xwjiang): Deprecate this after we control the interaction
+    #  between schedulers and executor.
+    NOOP = "NOOP"
+
+    _metric = None
+
+    _supports_buffered_results = True
+
+    def __init__(self):
+        tag_scheduler(self)
+
+    @property
+    def metric(self):
+        return self._metric
+
+    @property
+    def supports_buffered_results(self):
+        return self._supports_buffered_results
+
+    def set_search_properties(
+        self, metric: Optional[str], mode: Optional[str], **spec
+    ) -> bool:
+        """Pass search properties to scheduler.
+
+        This method acts as an alternative to instantiating schedulers
+        that react to metrics with their own `metric` and `mode` parameters.
+
+        Args:
+            metric: Metric to optimize
+            mode: One of ["min", "max"]. Direction to optimize.
+            **spec: Any kwargs for forward compatibility.
+                Info like Experiment.PUBLIC_KEYS is provided through here.
+
+        Returns:
+            True if the search properties were set successfully, False otherwise.
+        """
+        if self._metric and metric:
+            return False
+        if metric:
+            self._metric = metric
+
+        if self._metric is None:
+            # Per default, use anonymous metric
+            self._metric = DEFAULT_METRIC
+
+        return True
+
+    def on_trial_add(self, tune_controller: "TuneController", trial: Trial):
+        """Called when a new trial is added to the trial runner."""
+
+        raise NotImplementedError
+
+    def on_trial_error(self, tune_controller: "TuneController", trial: Trial):
+        """Notification for the error of trial.
+
+        This will only be called when the trial is in the RUNNING state."""
+
+        raise NotImplementedError
+
+    def on_trial_result(
+        self, tune_controller: "TuneController", trial: Trial, result: Dict
+    ) -> str:
+        """Called on each intermediate result returned by a trial.
+
+        At this point, the trial scheduler can make a decision by returning
+        one of CONTINUE, PAUSE, and STOP. This will only be called when the
+        trial is in the RUNNING state."""
+
+        raise NotImplementedError
+
+    def on_trial_complete(
+        self, tune_controller: "TuneController", trial: Trial, result: Dict
+    ):
+        """Notification for the completion of trial.
+
+        This will only be called when the trial is in the RUNNING state and
+        either completes naturally or by manual termination."""
+
+        raise NotImplementedError
+
+    def on_trial_remove(self, tune_controller: "TuneController", trial: Trial):
+        """Called to remove trial.
+
+        This is called when the trial is in PAUSED or PENDING state. Otherwise,
+        call `on_trial_complete`."""
+
+        raise NotImplementedError
+
+    def choose_trial_to_run(self, tune_controller: "TuneController") -> Optional[Trial]:
+        """Called to choose a new trial to run.
+
+        This should return one of the trials in tune_controller that is in
+        the PENDING or PAUSED state. This function must be idempotent.
+
+        If no trial is ready, return None."""
+
+        raise NotImplementedError
+
+    def debug_string(self) -> str:
+        """Returns a human readable message for printing to the console."""
+
+        raise NotImplementedError
+
+    def save(self, checkpoint_path: str):
+        """Save trial scheduler to a checkpoint"""
+        raise NotImplementedError
+
+    def restore(self, checkpoint_path: str):
+        """Restore trial scheduler from checkpoint."""
+        raise NotImplementedError
+
+
+@PublicAPI
+class FIFOScheduler(TrialScheduler):
+    """Simple scheduler that just runs trials in submission order."""
+
+    def __init__(self):
+        super().__init__()
+
+    def on_trial_add(self, tune_controller: "TuneController", trial: Trial):
+        pass
+
+    def on_trial_error(self, tune_controller: "TuneController", trial: Trial):
+        pass
+
+    def on_trial_result(
+        self, tune_controller: "TuneController", trial: Trial, result: Dict
+    ) -> str:
+        return TrialScheduler.CONTINUE
+
+    def on_trial_complete(
+        self, tune_controller: "TuneController", trial: Trial, result: Dict
+    ):
+        pass
+
+    def on_trial_remove(self, tune_controller: "TuneController", trial: Trial):
+        pass
+
+    def choose_trial_to_run(self, tune_controller: "TuneController") -> Optional[Trial]:
+        for trial in tune_controller.get_trials():
+            if trial.status == Trial.PENDING:
+                return trial
+        for trial in tune_controller.get_trials():
+            if trial.status == Trial.PAUSED:
+                return trial
+        return None
+
+    def debug_string(self) -> str:
+        return "Using FIFO scheduling algorithm."

python/ray/tune/schedulers/util.py

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