chore: import upstream snapshot with attribution

python/ray/tune/tests/_test_multi_tenancy_run.py

@@ -0,0 +1,76 @@
+import os
+import time
+from pathlib import Path
+
+from ray import tune
+from ray.tune.search import BasicVariantGenerator
+
+# Hang full script until this marker is deleted
+HANG_RUN_MARKER = os.environ.get("HANG_RUN_MARKER", "")
+
+# Delete this marker when a trial is started
+DELETE_TRIAL_MARKER = os.environ.get("DELETE_TRIAL_MARKER", "")
+
+# Hang in trial until this marker is deleted
+HANG_TRIAL_MARKER = os.environ.get("HANG_TRIAL_MARKER", "")
+
+# Delete this marker after tuning finished
+DELETE_RUN_MARKER = os.environ.get("DELETE_RUN_MARKER", "")
+
+# Hang at end of run until this marker is deleted
+HANG_END_MARKER = os.environ.get("HANG_END_MARKER", "")
+
+# Report this val as the "fixed" metric in the trial.
+# This value is captured in the trainer and will conflict when a trainable
+# is overwritten!
+FIXED_VAL = int(os.environ["FIXED_VAL"])
+
+# Grid search over these vals and report as "param" metric in the trial.
+# Even with conflicting trainables, these will be reported correctly as they
+# are tracked by the driver, not the trainable.
+VALS = [int(os.environ["VAL_1"]), int(os.environ["VAL_2"])]
+
+# Wait for HANG_RUN_MARKER
+while HANG_RUN_MARKER and Path(HANG_RUN_MARKER).exists():
+    time.sleep(0.1)
+
+
+def train_func(config):
+    # Delete DELETE_TRIAL_MARKER
+    delete_marker = config["delete_marker"]
+    if delete_marker and Path(delete_marker).exists():
+        Path(delete_marker).unlink()
+
+    # Wait for HANG_TRIAL_MARKER
+    hang_marker = config["hang_marker"]
+    while hang_marker and Path(hang_marker).exists():
+        time.sleep(0.1)
+
+    # Finish trial
+    tune.report({"param": config["param"], "fixed": config["fixed"]})
+
+
+if __name__ == "__main__":
+    tuner = tune.Tuner(
+        tune.with_resources(train_func, {"CPU": 2}),
+        param_space={
+            "fixed": FIXED_VAL,
+            "param": tune.grid_search(VALS),
+            "delete_marker": DELETE_TRIAL_MARKER,
+            "hang_marker": HANG_TRIAL_MARKER,
+        },
+        tune_config=tune.TuneConfig(search_alg=BasicVariantGenerator(max_concurrent=1)),
+    )
+    results = tuner.fit()
+
+    # Delete DELETE_RUN_MARKER
+    if DELETE_RUN_MARKER and Path(DELETE_RUN_MARKER).exists():
+        Path(DELETE_RUN_MARKER).unlink()
+
+    # Wait for HANG_END_MARKER
+    while HANG_END_MARKER and Path(HANG_END_MARKER).exists():
+        time.sleep(0.1)
+
+    # Put assertions last, so we don't finish early because of failures
+    assert sorted([result.metrics["param"] for result in results]) == VALS
+    assert [result.metrics["fixed"] for result in results] == [FIXED_VAL, FIXED_VAL]

python/ray/tune/tests/_test_trial_runner_callbacks.py

@@ -0,0 +1,280 @@
+import os
+import shutil
+import sys
+import tempfile
+import time
+import unittest
+from collections import OrderedDict
+from unittest.mock import patch
+
+import ray
+from ray import tune
+from ray.air._internal.checkpoint_manager import CheckpointStorage, _TrackedCheckpoint
+from ray.air.constants import TRAINING_ITERATION
+from ray.rllib import _register_all
+from ray.tune import Callback
+from ray.tune.callback import warnings
+from ray.tune.execution.ray_trial_executor import (
+    RayTrialExecutor,
+    _ExecutorEvent,
+    _ExecutorEventType,
+)
+from ray.tune.execution.trial_runner import TrialRunner
+from ray.tune.experiment import Experiment, Trial
+
+
+class TestCallback(Callback):
+    def __init__(self):
+        self.state = OrderedDict()
+
+    def setup(self, **info):
+        self.state["setup"] = info
+
+    def on_step_begin(self, **info):
+        self.state["step_begin"] = info
+
+    def on_step_end(self, **info):
+        self.state["step_end"] = info
+
+    def on_trial_start(self, **info):
+        self.state["trial_start"] = info
+
+    def on_trial_restore(self, **info):
+        self.state["trial_restore"] = info
+
+    def on_trial_save(self, **info):
+        self.state["trial_save"] = info
+
+    def on_trial_result(self, **info):
+        self.state["trial_result"] = info
+        result = info["result"]
+        trial = info["trial"]
+        assert result.get(TRAINING_ITERATION, None) != trial.last_result.get(
+            TRAINING_ITERATION, None
+        )
+
+    def on_trial_complete(self, **info):
+        self.state["trial_complete"] = info
+
+    def on_trial_error(self, **info):
+        self.state["trial_fail"] = info
+
+    def on_experiment_end(self, **info):
+        self.state["experiment_end"] = info
+
+
+# TODO(xwjiang): Move this to a testing util.
+class _MockTrialExecutor(RayTrialExecutor):
+    def __init__(self):
+        super().__init__()
+        self.next_future_result = None
+
+    def start_trial(self, trial: Trial):
+        trial.status = Trial.RUNNING
+        return True
+
+    def continue_training(self, trial: Trial):
+        pass
+
+    def get_next_executor_event(self, live_trials, next_trial_exists):
+        return self.next_future_result
+
+
+class TrialRunnerCallbacks(unittest.TestCase):
+    def setUp(self):
+        ray.init()
+        self.tmpdir = tempfile.mkdtemp()
+        self.callback = TestCallback()
+        self.executor = _MockTrialExecutor()
+        self.trial_runner = TrialRunner(
+            trial_executor=self.executor, callbacks=[self.callback]
+        )
+        # experiment would never be None normally, but it's fine for testing
+        self.trial_runner.setup_experiments(experiments=[None], total_num_samples=1)
+
+    def tearDown(self):
+        ray.shutdown()
+        _register_all()  # re-register the evicted objects
+        if "CUDA_VISIBLE_DEVICES" in os.environ:
+            del os.environ["CUDA_VISIBLE_DEVICES"]
+        shutil.rmtree(self.tmpdir)
+
+    def testCallbackSteps(self):
+        trials = [Trial("__fake", trial_id="one"), Trial("__fake", trial_id="two")]
+        for t in trials:
+            self.trial_runner.add_trial(t)
+
+        self.executor.next_future_result = _ExecutorEvent(
+            event_type=_ExecutorEventType.PG_READY
+        )
+        self.trial_runner.step()
+
+        # Trial 1 has been started
+        self.assertEqual(self.callback.state["trial_start"]["iteration"], 0)
+        self.assertEqual(self.callback.state["trial_start"]["trial"].trial_id, "one")
+
+        # All these events haven't happened, yet
+        self.assertTrue(
+            all(
+                k not in self.callback.state
+                for k in [
+                    "trial_restore",
+                    "trial_save",
+                    "trial_result",
+                    "trial_complete",
+                    "trial_fail",
+                    "experiment_end",
+                ]
+            )
+        )
+
+        self.executor.next_future_result = _ExecutorEvent(
+            event_type=_ExecutorEventType.PG_READY
+        )
+        self.trial_runner.step()
+
+        # Iteration not increased yet
+        self.assertEqual(self.callback.state["step_begin"]["iteration"], 1)
+
+        # Iteration increased
+        self.assertEqual(self.callback.state["step_end"]["iteration"], 2)
+
+        # Second trial has been just started
+        self.assertEqual(self.callback.state["trial_start"]["iteration"], 1)
+        self.assertEqual(self.callback.state["trial_start"]["trial"].trial_id, "two")
+
+        # Just a placeholder object ref for cp.value.
+        cp = _TrackedCheckpoint(
+            dir_or_data=ray.put(1),
+            storage_mode=CheckpointStorage.PERSISTENT,
+            metrics={TRAINING_ITERATION: 0},
+        )
+        trials[0].temporary_state.saving_to = cp
+
+        # Let the first trial save a checkpoint
+        self.executor.next_future_result = _ExecutorEvent(
+            event_type=_ExecutorEventType.SAVING_RESULT,
+            trial=trials[0],
+            result={_ExecutorEvent.KEY_FUTURE_RESULT: "__checkpoint"},
+        )
+        self.trial_runner.step()
+        self.assertEqual(self.callback.state["trial_save"]["iteration"], 2)
+        self.assertEqual(self.callback.state["trial_save"]["trial"].trial_id, "one")
+
+        # Let the second trial send a result
+        result = {TRAINING_ITERATION: 1, "metric": 800, "done": False}
+        self.executor.next_future_result = _ExecutorEvent(
+            event_type=_ExecutorEventType.TRAINING_RESULT,
+            trial=trials[1],
+            result={_ExecutorEvent.KEY_FUTURE_RESULT: result},
+        )
+        self.assertTrue(not trials[1].has_reported_at_least_once)
+        self.trial_runner.step()
+        self.assertEqual(self.callback.state["trial_result"]["iteration"], 3)
+        self.assertEqual(self.callback.state["trial_result"]["trial"].trial_id, "two")
+        self.assertEqual(self.callback.state["trial_result"]["result"]["metric"], 800)
+        self.assertEqual(trials[1].last_result["metric"], 800)
+
+        # Let the second trial restore from a checkpoint
+        trials[1].temporary_state.restoring_from = cp
+        self.executor.next_future_result = _ExecutorEvent(
+            event_type=_ExecutorEventType.RESTORING_RESULT,
+            trial=trials[1],
+            result={_ExecutorEvent.KEY_FUTURE_RESULT: None},
+        )
+        self.trial_runner.step()
+        self.assertEqual(self.callback.state["trial_restore"]["iteration"], 4)
+        self.assertEqual(self.callback.state["trial_restore"]["trial"].trial_id, "two")
+
+        # Let the second trial finish
+        trials[1].temporary_state.restoring_from = None
+        self.executor.next_future_result = _ExecutorEvent(
+            event_type=_ExecutorEventType.TRAINING_RESULT,
+            trial=trials[1],
+            result={
+                _ExecutorEvent.KEY_FUTURE_RESULT: {
+                    TRAINING_ITERATION: 2,
+                    "metric": 900,
+                    "done": True,
+                }
+            },
+        )
+        self.trial_runner.step()
+        self.assertEqual(self.callback.state["trial_complete"]["iteration"], 5)
+        self.assertEqual(self.callback.state["trial_complete"]["trial"].trial_id, "two")
+
+        # Let the first trial error
+        self.executor.next_future_result = _ExecutorEvent(
+            event_type=_ExecutorEventType.TRAINING_RESULT,
+            trial=trials[0],
+            result={_ExecutorEvent.KEY_EXCEPTION: Exception()},
+        )
+        self.trial_runner.step()
+        self.assertEqual(self.callback.state["trial_fail"]["iteration"], 6)
+        self.assertEqual(self.callback.state["trial_fail"]["trial"].trial_id, "one")
+
+    def testCallbacksEndToEnd(self):
+        def train_fn(config):
+            if config["do"] == "save":
+                with tune.checkpoint_dir(0):
+                    pass
+                tune.report(metric=1)
+            elif config["do"] == "fail":
+                raise RuntimeError("I am failing on purpose.")
+            elif config["do"] == "delay":
+                time.sleep(2)
+                tune.report(metric=20)
+
+        config = {"do": tune.grid_search(["save", "fail", "delay"])}
+
+        tune.run(
+            train_fn,
+            config=config,
+            raise_on_failed_trial=False,
+            callbacks=[self.callback],
+        )
+
+        self.assertIn("setup", self.callback.state)
+        self.assertTrue(self.callback.state["setup"] is not None)
+        keys = Experiment.PUBLIC_KEYS.copy()
+        keys.add("total_num_samples")
+        for key in keys:
+            self.assertIn(key, self.callback.state["setup"])
+        # check if it was added first
+        self.assertTrue(list(self.callback.state)[0] == "setup")
+        self.assertEqual(
+            self.callback.state["trial_fail"]["trial"].config["do"], "fail"
+        )
+        self.assertEqual(
+            self.callback.state["trial_save"]["trial"].config["do"], "save"
+        )
+        self.assertEqual(
+            self.callback.state["trial_result"]["trial"].config["do"], "delay"
+        )
+        self.assertEqual(
+            self.callback.state["trial_complete"]["trial"].config["do"], "delay"
+        )
+        self.assertIn("experiment_end", self.callback.state)
+        # check if it was added last
+        self.assertTrue(list(self.callback.state)[-1] == "experiment_end")
+
+    @patch.object(warnings, "warn")
+    def testCallbackSetupBackwardsCompatible(self, mocked_warning_method):
+        class NoExperimentInSetupCallback(Callback):
+            # Old method definition didn't take in **experiment.public_spec
+            def setup(self):
+                return
+
+        callback = NoExperimentInSetupCallback()
+        trial_runner = TrialRunner(callbacks=[callback])
+        trial_runner.setup_experiments(
+            experiments=[Experiment("", lambda x: x)], total_num_samples=1
+        )
+        mocked_warning_method.assert_called_once()
+        self.assertIn("Please update", mocked_warning_method.call_args_list[0][0][0])
+
+
+if __name__ == "__main__":
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/tune/tests/_test_trial_runner_pg.py

@@ -0,0 +1,318 @@
+import os
+import sys
+import time
+import unittest
+
+import numpy as np
+
+import ray
+from ray import tune
+from ray.cluster_utils import Cluster
+from ray.rllib import _register_all
+from ray.tune import Callback
+from ray.tune.execution.placement_groups import PlacementGroupFactory
+from ray.tune.execution.ray_trial_executor import RayTrialExecutor
+from ray.tune.execution.trial_runner import TrialRunner
+from ray.tune.experiment import Trial
+from ray.util import placement_group_table
+
+
+class TrialRunnerPlacementGroupTest(unittest.TestCase):
+    def setUp(self):
+        os.environ["TUNE_GLOBAL_CHECKPOINT_S"] = "10000"
+        os.environ["TUNE_MAX_PENDING_TRIALS_PG"] = "auto"  # Reset default
+        self.head_cpus = 8
+        self.head_gpus = 4
+        self.head_custom = 16
+
+        self.cluster = Cluster(
+            initialize_head=True,
+            connect=True,
+            head_node_args={
+                "include_dashboard": False,
+                "num_cpus": self.head_cpus,
+                "num_gpus": self.head_gpus,
+                "resources": {"custom": self.head_custom},
+                "_system_config": {
+                    "health_check_initial_delay_ms": 0,
+                    "health_check_period_ms": 1000,
+                    "health_check_failure_threshold": 10,
+                },
+            },
+        )
+        # Pytest doesn't play nicely with imports
+        _register_all()
+
+    def tearDown(self):
+        ray.shutdown()
+        self.cluster.shutdown()
+        _register_all()  # re-register the evicted objects
+
+    def _assertCleanup(self, trial_executor):
+        # Assert proper cleanup
+        resource_manager = trial_executor._resource_manager
+        self.assertFalse(resource_manager._pg_to_request)
+        self.assertFalse(resource_manager._acquired_pgs)
+        self.assertFalse(resource_manager._staging_future_to_pg)
+        self.assertFalse(resource_manager._pg_to_staging_future)
+        for rr in resource_manager._request_to_staged_pgs:
+            self.assertFalse(resource_manager._request_to_staged_pgs[rr])
+        for rr in resource_manager._request_to_ready_pgs:
+            self.assertFalse(resource_manager._request_to_ready_pgs[rr])
+
+        num_non_removed_pgs = len(
+            [p for pid, p in placement_group_table().items() if p["state"] != "REMOVED"]
+        )
+        self.assertEqual(num_non_removed_pgs, 0)
+
+    def testPlacementGroupRequests(self, reuse_actors=False, scheduled=10):
+        """In this test we try to start 10 trials but only have resources
+        for 2. Placement groups should still be created and PENDING.
+
+        Eventually they should be scheduled sequentially (i.e. in pairs
+        of two)."""
+        # Since we check per-step placement groups, set the reconcilation
+        # interval to 0
+        os.environ["TUNE_PLACEMENT_GROUP_RECON_INTERVAL"] = "0"
+
+        def train_fn(config):
+            time.sleep(1)
+            now = time.time()
+            tune.report(end=now - config["start_time"])
+
+        head_bundle = {"CPU": 4, "GPU": 0, "custom": 0}
+        child_bundle = {"custom": 1}
+        # Manually calculated number of parallel trials
+        max_num_parallel = 2
+
+        placement_group_factory = PlacementGroupFactory(
+            [head_bundle, child_bundle, child_bundle]
+        )
+
+        trial_executor = RayTrialExecutor(reuse_actors=reuse_actors)
+        trial_executor.setup(max_pending_trials=max_num_parallel)
+
+        this = self
+
+        class _TestCallback(Callback):
+            def on_step_end(self, iteration, trials, **info):
+                num_finished = len(
+                    [
+                        t
+                        for t in trials
+                        if t.status == Trial.TERMINATED or t.status == Trial.ERROR
+                    ]
+                )
+
+                resource_manager = trial_executor._resource_manager
+
+                num_staging = sum(
+                    len(s) for s in resource_manager._request_to_staged_pgs.values()
+                )
+                num_ready = sum(
+                    len(s) for s in resource_manager._request_to_ready_pgs.values()
+                )
+                num_in_use = len(resource_manager._acquired_pgs)
+                num_cached = trial_executor._actor_cache.num_cached_objects
+
+                total_num_tracked = num_staging + num_ready + num_in_use + num_cached
+
+                # All trials should be scheduled
+                this.assertEqual(
+                    scheduled,
+                    min(scheduled, len(trials)),
+                    msg=f"Num trials iter {iteration}",
+                )
+
+                # The following two tests were relaxed for reuse_actors=True
+                # so that up to `max_num_parallel` more placement groups can
+                # exist than we would expect. This is because caching
+                # relies on reconciliation for cleanup to avoid overscheduling
+                # of new placement groups.
+                num_parallel_reuse = int(reuse_actors) * max_num_parallel
+
+                # The number of PGs should decrease when trials finish
+                # We allow a constant excess of 1 here because the trial will
+                # be TERMINATED and the resources only returned after the trainable
+                # cleanup future succeeded. Because num_finished will increase,
+                # this still asserts that the number of PGs goes down over time.
+                this.assertGreaterEqual(
+                    max(scheduled, len(trials)) - num_finished + 1 + num_parallel_reuse,
+                    total_num_tracked,
+                    msg=f"Num tracked iter {iteration}, {len(trials)}, "
+                    f"{scheduled}, {num_finished}, {num_parallel_reuse}",
+                )
+
+        start = time.time()
+        out = tune.run(
+            train_fn,
+            config={"start_time": start},
+            resources_per_trial=placement_group_factory,
+            num_samples=10,
+            trial_executor=trial_executor,
+            callbacks=[_TestCallback()],
+            reuse_actors=reuse_actors,
+            verbose=2,
+        )
+
+        trial_end_times = sorted(t.last_result["end"] for t in out.trials)
+        print("Trial end times:", trial_end_times)
+        max_diff = trial_end_times[-1] - trial_end_times[0]
+
+        # Not all trials have been run in parallel
+        self.assertGreater(max_diff, 3)
+
+        # Some trials should have run in parallel
+        # Todo: Re-enable when using buildkite
+        # self.assertLess(max_diff, 10)
+
+        self._assertCleanup(trial_executor)
+
+    def testPlacementGroupRequestsWithActorReuse(self):
+        """Assert that reuse actors doesn't leak placement groups"""
+        self.testPlacementGroupRequests(reuse_actors=True)
+
+    def testPlacementGroupLimitedRequests(self):
+        """Assert that maximum number of placement groups is enforced."""
+        os.environ["TUNE_MAX_PENDING_TRIALS_PG"] = "6"
+        self.testPlacementGroupRequests(scheduled=6)
+
+    def testPlacementGroupLimitedRequestsWithActorReuse(self):
+        os.environ["TUNE_MAX_PENDING_TRIALS_PG"] = "6"
+        self.testPlacementGroupRequests(reuse_actors=True, scheduled=6)
+
+    def testPlacementGroupDistributedTraining(self, reuse_actors=False):
+        """Run distributed training using placement groups.
+
+        Each trial requests 4 CPUs and starts 4 remote training workers.
+        """
+
+        head_bundle = {"CPU": 1, "GPU": 0, "custom": 0}
+        child_bundle = {"CPU": 1}
+
+        placement_group_factory = PlacementGroupFactory(
+            [head_bundle, child_bundle, child_bundle, child_bundle]
+        )
+
+        @ray.remote
+        class TrainingActor:
+            def train(self, val):
+                time.sleep(1)
+                return val
+
+        def train_fn(config):
+            base = config["base"]
+            actors = [TrainingActor.remote() for _ in range(4)]
+            futures = [
+                actor.train.remote(base + 2 * i) for i, actor in enumerate(actors)
+            ]
+            results = ray.get(futures)
+
+            end = time.time() - config["start_time"]
+            tune.report(avg=np.mean(results), end=end)
+
+        trial_executor = RayTrialExecutor(reuse_actors=reuse_actors)
+
+        start = time.time()
+        out = tune.run(
+            train_fn,
+            config={
+                "start_time": start,
+                "base": tune.grid_search(list(range(0, 100, 10))),
+            },
+            resources_per_trial=placement_group_factory,
+            num_samples=1,
+            trial_executor=trial_executor,
+            reuse_actors=reuse_actors,
+            verbose=2,
+        )
+
+        avgs = sorted(t.last_result["avg"] for t in out.trials)
+        self.assertSequenceEqual(avgs, list(range(3, 103, 10)))
+
+        trial_end_times = sorted(t.last_result["end"] for t in out.trials)
+        print("Trial end times:", trial_end_times)
+        max_diff = trial_end_times[-1] - trial_end_times[0]
+
+        # Not all trials have been run in parallel
+        self.assertGreater(max_diff, 3)
+
+        # Some trials should have run in parallel
+        # Todo: Re-enable when using buildkite
+        # self.assertLess(max_diff, 10)
+
+        self._assertCleanup(trial_executor)
+
+    def testPlacementGroupDistributedTrainingWithActorReuse(self):
+        self.testPlacementGroupDistributedTraining(reuse_actors=True)
+
+
+class TrialRunnerPlacementGroupHeterogeneousTest(unittest.TestCase):
+    def tearDown(self) -> None:
+        if ray.is_initialized:
+            ray.shutdown()
+
+    def testResourceDeadlock(self):
+        """Tests that resource deadlock is avoided for heterogeneous PGFs.
+
+        We start 4 trials in a cluster with 2 CPUs. The first two trials
+        require 1 CPU each, the third trial 2 CPUs, the fourth trial 1 CPU.
+
+        The second trial needs a bit more time to finish. This means that the
+        resources from the first trial will be freed, and the PG of the
+        _fourth_ trial becomes ready (not that of the third trial, because that
+        requires 2 CPUs - however, one is still occupied by trial 2).
+
+        After the first two trials finished, the FIFOScheduler tries to start
+        the third trial. However, it can't be started because its placement
+        group is not ready. Instead, the placement group of the fourth
+        trial is ready. Thus, we opt to run the fourth trial instead.
+        """
+
+        def train_fn(config):
+            time.sleep(config["sleep"])
+            return 4
+
+        ray.init(num_cpus=2)
+
+        tune.register_trainable("het", train_fn)
+        pgf1 = PlacementGroupFactory([{"CPU": 1}])
+        pgf2 = PlacementGroupFactory([{"CPU": 2}])
+
+        trial1 = Trial("het", config={"sleep": 0}, placement_group_factory=pgf1)
+        trial2 = Trial("het", config={"sleep": 2}, placement_group_factory=pgf1)
+        trial3 = Trial("het", config={"sleep": 0}, placement_group_factory=pgf2)
+        trial4 = Trial("het", config={"sleep": 0}, placement_group_factory=pgf1)
+
+        runner = TrialRunner(fail_fast=True)
+        runner.add_trial(trial1)
+        runner.add_trial(trial2)
+        runner.add_trial(trial3)
+        runner.add_trial(trial4)
+
+        timeout = time.monotonic() + 30
+        while not runner.is_finished():
+            # We enforce a timeout here
+            self.assertLess(
+                time.monotonic(), timeout, msg="Ran into a resource deadlock"
+            )
+
+            runner.step()
+
+
+def test_placement_group_no_cpu_trainer():
+    """Bundles with only GPU:1 but no CPU should work"""
+    ray.init(num_gpus=1, num_cpus=1)
+    pgf = PlacementGroupFactory([{"GPU": 1, "CPU": 0}, {"CPU": 1}])
+
+    def train_fn(config):
+        time.sleep(1)
+        return 5
+
+    tune.run(train_fn, resources_per_trial=pgf)
+
+
+if __name__ == "__main__":
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/tune/tests/conftest.py

@@ -0,0 +1,33 @@
+import logging
+
+import boto3
+import pytest
+
+from ray._common.test_utils import simulate_s3_bucket
+from ray.air._internal.uri_utils import URI
+
+# Trigger pytest hook to automatically zip test cluster logs to archive dir on failure
+from ray.tests.conftest import (
+    propagate_logs,  # noqa
+    pytest_runtest_makereport,  # noqa
+)
+
+
+@pytest.fixture
+def mock_s3_bucket_uri():
+    port = 5002
+    region = "us-west-2"
+    with simulate_s3_bucket(port=port, region=region) as s3_uri:
+        s3 = boto3.client(
+            "s3", region_name=region, endpoint_url=f"http://localhost:{port}"
+        )
+        # Bucket name will be autogenerated/unique per test
+        bucket_name = URI(s3_uri).name
+        s3.create_bucket(
+            Bucket=bucket_name,
+            CreateBucketConfiguration={"LocationConstraint": region},
+        )
+        # Disable server HTTP request logging
+        logging.getLogger("werkzeug").setLevel(logging.WARNING)
+        yield s3_uri
+        logging.getLogger("werkzeug").setLevel(logging.INFO)

python/ray/tune/tests/example.py

@@ -0,0 +1,61 @@
+# ruff: noqa
+
+# This is an example quickstart for Tune.
+# To connect to a cluster, uncomment below:
+
+# import ray
+# import argparse
+# parser = argparse.ArgumentParser()
+# parser.add_argument("--address")
+# args = parser.parse_args()
+# ray.init(address=args.address)
+
+# __quick_start_begin__
+from ray import tune
+
+
+def objective(config):  # <1>
+    score = config["a"] ** 2 + config["b"]
+    return {"score": score}
+
+
+search_space = {  # <2>
+    "a": tune.grid_search([0.001, 0.01, 0.1, 1.0]),
+    "b": tune.choice([1, 2, 3]),
+}
+
+tuner = tune.Tuner(objective, param_space=search_space)  # <3>
+
+results = tuner.fit()
+print(results.get_best_result(metric="score", mode="min").config)
+# __quick_start_end__
+
+# __ml_quick_start_begin__
+def objective(step, alpha, beta):
+    return (0.1 + alpha * step / 100) ** (-1) + beta * 0.1
+
+
+def training_function(config):
+    # Hyperparameters
+    alpha, beta = config["alpha"], config["beta"]
+    for step in range(10):
+        # Iterative training function - can be any arbitrary training procedure.
+        intermediate_score = objective(step, alpha, beta)
+        # Feed the score back back to Tune.
+        tune.report({"mean_loss": intermediate_score})
+
+
+tuner = tune.Tuner(
+    training_function,
+    param_space={
+        "alpha": tune.grid_search([0.001, 0.01, 0.1]),
+        "beta": tune.choice([1, 2, 3]),
+    },
+)
+results = tuner.fit()
+
+print("Best config: ", results.get_best_result(metric="mean_loss", mode="min").config)
+
+# Get a dataframe for analyzing trial results.
+df = results.get_dataframe()
+# __ml_quick_start_end__

python/ray/tune/tests/execution/conftest.py

@@ -0,0 +1,3 @@
+# Trigger pytest hook to automatically zip test cluster logs to archive dir on failure
+from ray.tests.conftest import propagate_logs  # noqa
+from ray.tests.conftest import pytest_runtest_makereport  # noqa

python/ray/tune/tests/execution/test_actor_caching.py

@@ -0,0 +1,131 @@
+import sys
+
+import pytest
+
+import ray
+from ray.tune import PlacementGroupFactory
+from ray.tune.tests.execution.utils import TestingTrial, create_execution_test_objects
+
+
+@pytest.fixture
+def ray_start_2_cpus():
+    address_info = ray.init(num_cpus=2)
+    yield address_info
+    ray.shutdown()
+
+
+def test_actor_cached(tmpdir, ray_start_2_cpus):
+    tune_controller, actor_manger, resource_manager = create_execution_test_objects(
+        max_pending_trials=8
+    )
+
+    assert not actor_manger.added_actors
+
+    tune_controller.add_trial(TestingTrial("trainable1", stub=True, trial_id="trial1"))
+    tune_controller.step()
+
+    tracked_actor, cls_name, kwargs = actor_manger.added_actors[0]
+    assert cls_name == "trainable1"
+
+
+def test_actor_reuse_unstaged(tmpdir, ray_start_2_cpus):
+    """A trial that hasn't been staged can re-use an actor.
+
+    In specific circumstances, this can lead to errors. Notably, when an
+    external source (e.g. a scheduler) directly calls TuneController APIs,
+    we can be in a situation where a trial has not been staged, but there is
+    still an actor available for it to use (because it hasn't been evicted from
+    the cache, yet).
+
+    This test constructs such a situation an asserts that actor re-use does not
+    lead to errors in those cases.
+    """
+    tune_controller, actor_manger, resource_manager = create_execution_test_objects(
+        max_pending_trials=1
+    )
+    tune_controller._reuse_actors = True
+
+    assert not actor_manger.added_actors
+
+    trialA1 = TestingTrial(
+        "trainable1",
+        stub=True,
+        trial_id="trialA1",
+        placement_group_factory=PlacementGroupFactory([{"CPU": 1}]),
+    )
+    tune_controller.add_trial(trialA1)
+    trialB1 = TestingTrial(
+        "trainable1",
+        stub=True,
+        trial_id="trialB1",
+        placement_group_factory=PlacementGroupFactory([{"CPU": 5}]),
+    )
+    tune_controller.add_trial(trialB1)
+    trialA2 = TestingTrial(
+        "trainable1",
+        stub=True,
+        trial_id="trialA2",
+        placement_group_factory=PlacementGroupFactory([{"CPU": 1}]),
+    )
+    tune_controller.add_trial(trialA2)
+    tune_controller.step()
+
+    # Prevent trial A3 from being staged by setting the number
+    # of pending actors to the maximum allowed
+    actor_manger.set_num_pending(2)
+
+    trialA3 = TestingTrial(
+        "trainable1",
+        stub=True,
+        trial_id="trialA3",
+        placement_group_factory=PlacementGroupFactory([{"CPU": 1}]),
+    )
+    tune_controller.add_trial(trialA3)
+    tune_controller.step()
+
+    tracked_actorA1, _, _ = actor_manger.added_actors[0]
+    tracked_actorB1, _, _ = actor_manger.added_actors[1]
+    tracked_actorA2, _, _ = actor_manger.added_actors[2]
+
+    # Start trial A1, report that it's done training.
+    # This will cache the actor for A1 as A2 is already scheduled.
+    tune_controller._actor_started(tracked_actorA1)
+    tune_controller._on_training_result(trialA1, {"done": True})
+
+    # Trial A2 should be in the staged trials. A3 should still not be staged.
+    assert trialA2 in tune_controller._staged_trials
+    assert trialA3 not in tune_controller._staged_trials
+
+    # The actor of A1 should be cached for re-use now.
+    assert tune_controller._actor_cache.num_cached_objects == 1
+
+    # In the meantime, actor A2 started. This will unstage it.
+    tune_controller._actor_started(tracked_actorA2)
+
+    # Now, an external source (e.g. the BOHB scheduler) wants to prematurely
+    # stop trial A2. This will leave the cached actor intact, but trial A3
+    # is still not scheduled.
+    tune_controller._schedule_trial_stop(trialA2)
+    assert tune_controller._actor_cache.num_cached_objects == 1
+
+    # Process events. This will invoke "path 3" in TuneController._maybe_add_actors
+    # and re-use the cached actor
+    tune_controller.step()
+
+    # Reset future scheduled
+    assert actor_manger.scheduled_futures[-1][2] == "reset"
+
+    # Prior to https://github.com/ray-project/ray/pull/36951, there was a bug here:
+    # Because trial A3 was never staged, the unstage ran into an error.
+    # This fails without the line: self._staged_trials.add(start_trial)
+    tune_controller._on_trial_reset(trialA3, True)
+
+    # When the actor finally stops, the cache size is adjusted and the actor is
+    # evicted. This test failed without the line:
+    # self._actor_cache.increase_max(start_trial.placement_group_factory)
+    tune_controller._actor_stopped(tracked_actorA1)
+    tune_controller.step()
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/tune/tests/execution/test_controller_callback_integration.py

@@ -0,0 +1,69 @@
+import sys
+from typing import Dict, Optional
+
+import pytest
+
+import ray
+from ray.air.execution import FixedResourceManager, PlacementGroupResourceManager
+from ray.train.tests.util import mock_storage_context
+from ray.tune import Callback, ResumeConfig
+from ray.tune.execution.tune_controller import TuneController
+from ray.tune.experiment import Trial
+from ray.tune.utils.mock_trainable import MOCK_TRAINABLE_NAME, register_mock_trainable
+
+
+@pytest.fixture(scope="function")
+def ray_start_4_cpus_2_gpus_extra():
+    address_info = ray.init(num_cpus=4, num_gpus=2, resources={"a": 2})
+    yield address_info
+    ray.shutdown()
+
+
+@pytest.fixture(autouse=True)
+def register_test_trainable():
+    register_mock_trainable()
+
+
+class StatefulCallback(Callback):
+    CKPT_FILE_TMPL = "test-callback-state-{}.json"
+
+    def __init__(self):
+        self.counter = 0
+
+    def on_trial_result(self, iteration, trials, trial, result, **info):
+        self.counter += 1
+
+    def get_state(self) -> Optional[Dict]:
+        return {"counter": self.counter}
+
+    def set_state(self, state: Dict):
+        self.counter = state["counter"]
+
+
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+def test_callback_save_restore(
+    ray_start_4_cpus_2_gpus_extra, resource_manager_cls, tmpdir
+):
+    """Check that callback state is restored correctly.
+
+    Legacy test: test_trial_runner_3.py::TrialRunnerTest::testCallbackSaveRestore
+    """
+    storage = mock_storage_context()
+    runner = TuneController(callbacks=[StatefulCallback()], storage=storage)
+    runner.add_trial(Trial(MOCK_TRAINABLE_NAME, stub=True, storage=storage))
+    for i in range(3):
+        runner._callbacks.on_trial_result(
+            iteration=i, trials=None, trial=None, result=None
+        )
+    runner.checkpoint(force=True, wait=True)
+    callback = StatefulCallback()
+    runner2 = TuneController(callbacks=[callback], storage=storage)
+    assert callback.counter == 0
+    runner2.resume(resume_config=ResumeConfig())
+    assert callback.counter == 3
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/tune/tests/execution/test_controller_checkpointing_integration.py

@@ -0,0 +1,619 @@
+import json
+import logging
+import os
+import sys
+import tempfile
+import time
+from unittest import mock
+
+import pytest
+
+import ray
+from ray.air.constants import TRAINING_ITERATION
+from ray.air.execution import FixedResourceManager, PlacementGroupResourceManager
+from ray.train._internal.session import _TrainingResult
+from ray.train._internal.storage import StorageContext
+from ray.train.tests.util import mock_storage_context
+from ray.tune import (
+    Callback,
+    Checkpoint,
+    CheckpointConfig,
+    PlacementGroupFactory,
+    ResumeConfig,
+)
+from ray.tune.execution.tune_controller import TuneController
+from ray.tune.experiment import Trial
+from ray.tune.result import DONE
+from ray.tune.schedulers import FIFOScheduler
+from ray.tune.search import BasicVariantGenerator
+from ray.tune.tests.tune_test_util import TrialResultObserver
+from ray.tune.utils.mock_trainable import MOCK_TRAINABLE_NAME, register_mock_trainable
+
+STORAGE = mock_storage_context()
+
+
+@pytest.fixture(autouse=True)
+def register_test_trainable():
+    register_mock_trainable()
+
+
+@pytest.fixture(scope="function")
+def ray_start_4_cpus_2_gpus_extra():
+    address_info = ray.init(num_cpus=4, num_gpus=2, resources={"a": 2})
+    yield address_info
+    ray.shutdown()
+
+
+def create_mock_components():
+    class _MockScheduler(FIFOScheduler):
+        errored_trials = []
+
+        def on_trial_error(self, tune_controller, trial):
+            self.errored_trials += [trial]
+
+    class _MockSearchAlg(BasicVariantGenerator):
+        errored_trials = []
+
+        def on_trial_complete(self, trial_id, error=False, **kwargs):
+            if error:
+                self.errored_trials += [trial_id]
+
+    searchalg = _MockSearchAlg()
+    scheduler = _MockScheduler()
+    return searchalg, scheduler
+
+
+def num_checkpoints(trial):
+    return sum(
+        item.startswith("checkpoint_")
+        for item in os.listdir(trial.storage.trial_fs_path)
+    )
+
+
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+def test_checkpoint_save_restore(
+    ray_start_4_cpus_2_gpus_extra, resource_manager_cls, tmpdir
+):
+    """Test that a checkpoint is saved and can be used to restore a trainable.
+
+    The trainable saves a checkpoint and terminates. We then start another trial
+    that should restore from the saved checkpoint and assert that it picks up
+    the state and continues to run to termination.
+
+    Legacy test: test_trial_runner_2.py::TrialRunnerTest::testCheckpointing
+    Legacy test: test_trial_runner_2.py::TrialRunnerTest::testRestoreMetricsAfterCheckpointing  # noqa
+    """
+    runner = TuneController(
+        resource_manager_factory=lambda: resource_manager_cls(), storage=STORAGE
+    )
+    kwargs = {
+        "stopping_criterion": {"training_iteration": 1},
+        "placement_group_factory": PlacementGroupFactory([{"CPU": 1, "GPU": 1}]),
+        "checkpoint_config": CheckpointConfig(checkpoint_frequency=1),
+        "storage": STORAGE,
+    }
+    runner.add_trial(Trial(MOCK_TRAINABLE_NAME, **kwargs))
+    trials = runner.get_trials()
+
+    runner.step()  # Start trial
+
+    while trials[0].status != Trial.RUNNING:
+        runner.step()
+
+    while trials[0].status != Trial.TERMINATED:
+        runner.step()
+
+    assert trials[0].latest_checkpoint_result.metrics[TRAINING_ITERATION] == 1
+    assert trials[0].last_result[TRAINING_ITERATION] == 1
+    assert trials[0].last_result["iterations_since_restore"] == 1
+
+    # Prepare new trial
+    kwargs["restore_path"] = trials[0].checkpoint.path
+    new_trial = Trial(MOCK_TRAINABLE_NAME, **kwargs)
+    runner.add_trial(new_trial)
+    trials = runner.get_trials()
+
+    assert trials[1].status == Trial.PENDING
+
+    # Start trial, restore, run to termination
+    while trials[1].status != Trial.RUNNING:
+        runner.step()
+
+    # Restore
+    runner.step()
+
+    # Run to termination
+    while trials[1].status != Trial.TERMINATED:
+        runner.step()
+
+    assert trials[0].latest_checkpoint_result.metrics[TRAINING_ITERATION] == 1
+    assert trials[1].last_result[TRAINING_ITERATION] == 1
+    assert trials[1].last_result["iterations_since_restore"] == 1
+
+
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+def test_checkpoint_at_end(ray_start_4_cpus_2_gpus_extra, resource_manager_cls, tmpdir):
+    """Test that a checkpoint is saved at end for class trainables with that config.
+
+    Legacy test: test_trial_runner_2.py::TrialRunnerTest::testCheckpointingAtEnd
+    Legacy test: test_trial_runner_2.py::TrialRunnerTest::testResultDone
+    """
+    runner = TuneController(
+        resource_manager_factory=lambda: resource_manager_cls(),
+        storage=STORAGE,
+    )
+    kwargs = {
+        "stopping_criterion": {"training_iteration": 2},
+        "checkpoint_config": CheckpointConfig(checkpoint_at_end=True),
+        "placement_group_factory": PlacementGroupFactory([{"CPU": 1, "GPU": 1}]),
+        "storage": STORAGE,
+    }
+    runner.add_trial(Trial(MOCK_TRAINABLE_NAME, **kwargs))
+    trials = runner.get_trials()
+
+    while not runner.is_finished():
+        runner.step()
+
+    assert trials[0].has_checkpoint()
+    assert trials[0].last_result[DONE]
+
+
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+def test_pause_resume_trial(
+    ray_start_4_cpus_2_gpus_extra, resource_manager_cls, tmpdir
+):
+    """Test that trial that is paused and resumed picks up its last checkpoint.
+
+    Legacy test: test_trial_runner_2.py::TrialRunnerTest::testPauseThenResume
+    """
+    runner = TuneController(
+        resource_manager_factory=lambda: resource_manager_cls(),
+        storage=STORAGE,
+    )
+    kwargs = {
+        "stopping_criterion": {"training_iteration": 2},
+        "placement_group_factory": PlacementGroupFactory([{"CPU": 1, "GPU": 1}]),
+        "checkpoint_config": CheckpointConfig(checkpoint_frequency=1),
+        "storage": STORAGE,
+    }
+    runner.add_trial(Trial(MOCK_TRAINABLE_NAME, **kwargs))
+    trials = runner.get_trials()
+
+    while trials[0].status != Trial.RUNNING:
+        runner.step()
+
+    runner._schedule_trial_pause(trials[0], should_checkpoint=True)
+
+    while trials[0].status != Trial.PAUSED:
+        runner.step()
+
+    assert trials[0].has_checkpoint()
+    assert not trials[0].last_result.get(DONE), trials[0].last_result
+
+    # Start again
+    runner._set_trial_status(trials[0], Trial.PENDING)
+
+    while trials[0].status != Trial.RUNNING:
+        runner.step()
+
+    while trials[0].status != Trial.TERMINATED:
+        runner.step()
+
+    assert trials[0].checkpoint
+    assert trials[0].last_result[TRAINING_ITERATION] == 2
+    assert trials[0].last_result["iterations_since_restore"] == 1
+    assert trials[0].last_result["time_since_restore"] > 0
+
+
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+def test_checkpoint_num_to_keep(
+    ray_start_4_cpus_2_gpus_extra, resource_manager_cls, tmp_path
+):
+    """Test that only num_to_keep checkpoints are kept.
+
+    This should also hold true when the experiment is resumed.
+
+    Legacy test: test_trial_runner_2.py::TrialRunnerTest::testPauseResumeCheckpointCount
+    """
+    trial = Trial(
+        MOCK_TRAINABLE_NAME,
+        checkpoint_config=CheckpointConfig(num_to_keep=2),
+        storage=STORAGE,
+    )
+    trial.init_local_path()
+
+    def write_checkpoint(trial: Trial, index: int):
+        checkpoint_dir = tmp_path / StorageContext._make_checkpoint_dir_name(index)
+        checkpoint_dir.mkdir(parents=True, exist_ok=True)
+        result = {"training_iteration": index}
+        with open(os.path.join(checkpoint_dir, "cp.json"), "w") as f:
+            json.dump(result, f)
+
+        checkpoint = Checkpoint.from_directory(checkpoint_dir)
+        return _TrainingResult(checkpoint=checkpoint, metrics=result)
+
+    def get_checkpoint_dirs(trial: Trial):
+        return [d for d in os.listdir(tmp_path) if d.startswith("checkpoint_")]
+
+    runner = TuneController(
+        resource_manager_factory=lambda: resource_manager_cls(), storage=STORAGE
+    )
+
+    runner.add_trial(trial)
+
+    # Write 1 checkpoint
+    result = write_checkpoint(trial, 1)
+    runner._on_saving_result(trial, result)
+
+    # Expect 1 checkpoint
+    cp_dirs = get_checkpoint_dirs(trial)
+    assert len(cp_dirs) == 1, f"Checkpoint dirs: {cp_dirs}"
+
+    # Write second checkpoint
+    result = write_checkpoint(trial, 2)
+    runner._on_saving_result(trial, result)
+
+    # Expect 2 checkpoints
+    cp_dirs = get_checkpoint_dirs(trial)
+    assert len(cp_dirs) == 2, f"Checkpoint dirs: {cp_dirs}"
+
+    # Write third checkpoint
+    result = write_checkpoint(trial, 3)
+    runner._on_saving_result(trial, result)
+
+    # Expect 2 checkpoints because num_to_keep = 2
+    cp_dirs = get_checkpoint_dirs(trial)
+    assert len(cp_dirs) == 2, f"Checkpoint dirs: {cp_dirs}"
+
+    # Re-instantiate trial runner and resume
+    runner.checkpoint(force=True, wait=True)
+    runner = TuneController(
+        resource_manager_factory=lambda: resource_manager_cls(),
+        storage=STORAGE,
+        resume_config=ResumeConfig(),
+    )
+
+    trial = runner.get_trials()[0]
+
+    # Write fourth checkpoint
+    result = write_checkpoint(trial, 4)
+    runner._on_saving_result(trial, result)
+
+    # Expect 2 checkpoints because num_to_keep = 2
+    cp_dirs = get_checkpoint_dirs(trial)
+    assert len(cp_dirs) == 2, f"Checkpoint dirs: {cp_dirs}"
+
+    # Write fifth checkpoint
+    result = write_checkpoint(trial, 5)
+    runner._on_saving_result(trial, result)
+
+    # Expect 2 checkpoints because num_to_keep = 2
+    cp_dirs = get_checkpoint_dirs(trial)
+    assert len(cp_dirs) == 2, f"Checkpoint dirs: {cp_dirs}"
+
+    # Checkpoints before restore should be deleted
+    assert "checkpoint_000004" in cp_dirs
+    assert "checkpoint_000005" in cp_dirs
+
+    assert "checkpoint_000002" not in cp_dirs
+    assert "checkpoint_000003" not in cp_dirs
+
+
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+def test_checkpoint_freq_buffered(
+    ray_start_4_cpus_2_gpus_extra, resource_manager_cls, tmp_path
+):
+    """Test that trial checkpoints are a lower bound for buffered training iterations.
+
+    Legacy test: test_trial_runner_3.py::TrialRunnerTest::testCheckpointFreqBuffered
+    """
+    with mock.patch.dict(
+        os.environ,
+        {"TUNE_RESULT_BUFFER_LENGTH": "7", "TUNE_RESULT_BUFFER_MIN_TIME_S": "1"},
+    ):
+        trial = Trial(
+            MOCK_TRAINABLE_NAME,
+            checkpoint_config=CheckpointConfig(checkpoint_frequency=3),
+            storage=STORAGE,
+        )
+        runner = TuneController(
+            resource_manager_factory=lambda: resource_manager_cls(),
+            storage=STORAGE,
+            checkpoint_period=0,
+        )
+        runner.add_trial(trial)
+
+        while not trial.is_saving:
+            runner.step()
+        runner.step()
+        assert trial.last_result[TRAINING_ITERATION] == 3
+        assert num_checkpoints(trial) == 1
+
+        while not trial.is_saving:
+            runner.step()
+        runner.step()
+        assert trial.last_result[TRAINING_ITERATION] == 6
+        assert num_checkpoints(trial) == 2
+
+        while not trial.is_saving:
+            runner.step()
+        runner.step()
+        assert trial.last_result[TRAINING_ITERATION] == 9
+        assert num_checkpoints(trial) == 3
+
+
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+def test_checkpoint_at_end_not_buffered(
+    ray_start_4_cpus_2_gpus_extra, resource_manager_cls, tmp_path
+):
+    """Test that trials with `checkpoint_at_end=True` are never buffered.
+
+    Legacy test: test_trial_runner_3.py::TrialRunnerTest::testCheckpointAtEndNotBuffered
+    """
+    with mock.patch.dict(
+        os.environ,
+        {"TUNE_RESULT_BUFFER_LENGTH": "7", "TUNE_RESULT_BUFFER_MIN_TIME_S": "0.5"},
+    ):
+        trial = Trial(
+            MOCK_TRAINABLE_NAME,
+            checkpoint_config=CheckpointConfig(
+                checkpoint_at_end=True,
+            ),
+            stopping_criterion={"training_iteration": 4},
+            storage=STORAGE,
+        )
+        observer = TrialResultObserver()
+        runner = TuneController(
+            resource_manager_factory=lambda: resource_manager_cls(),
+            storage=STORAGE,
+            callbacks=[observer],
+        )
+        runner.add_trial(trial)
+
+        while not observer.just_received_a_result():
+            runner.step()
+        assert trial.last_result[TRAINING_ITERATION] == 1
+        assert num_checkpoints(trial) == 0
+
+        while True:
+            runner.step()
+            if observer.just_received_a_result():
+                break
+        assert trial.last_result[TRAINING_ITERATION] == 2
+        assert num_checkpoints(trial) == 0
+
+        while True:
+            runner.step()
+            if observer.just_received_a_result():
+                break
+        assert trial.last_result[TRAINING_ITERATION] == 3
+        assert num_checkpoints(trial) == 0
+
+        while True:
+            runner.step()
+            if observer.just_received_a_result():
+                break
+        assert trial.last_result[TRAINING_ITERATION] == 4
+
+        while not runner.is_finished():
+            runner.step()
+        assert num_checkpoints(trial) == 1
+
+
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+def test_checkpoint_auto_period(
+    ray_start_4_cpus_2_gpus_extra, resource_manager_cls, tmp_path
+):
+    """Test that the checkpoint auto period is adjusted when syncing takes a long time.
+
+    Legacy test: test_trial_runner_3.py::TrialRunnerTest::testCheckpointAutoPeriod
+    """
+    storage = mock_storage_context()
+
+    with tempfile.TemporaryDirectory() as local_dir:
+        storage.storage_local_path = local_dir
+
+        runner = TuneController(
+            resource_manager_factory=lambda: resource_manager_cls(),
+            storage=storage,
+            checkpoint_period="auto",
+        )
+
+        with mock.patch.object(runner, "save_to_dir") as save_to_dir:
+            save_to_dir.side_effect = lambda *a, **kw: time.sleep(2)
+            runner.add_trial(Trial(MOCK_TRAINABLE_NAME, storage=storage))
+            runner.step()  # Run one step, this will trigger checkpointing
+
+        assert runner._checkpoint_manager._checkpoint_period > 38.0
+
+
+def test_checkpoint_force_with_num_to_keep(ray_start_4_cpus_2_gpus_extra, tmp_path):
+    """Test that cloud syncing is forced if one of the trials has made more
+    than num_to_keep checkpoints since last sync.
+    Legacy test: test_trial_runner_3.py::TrialRunnerTest::
+        testCloudCheckpointForceWithNumToKeep
+    """
+    storage = mock_storage_context()
+    # Needed to avoid infinite recursion error on CI runners
+    storage.syncer.__getstate__ = lambda *a, **kw: {}
+
+    with mock.patch.object(storage.syncer, "sync_up") as sync_up:
+        num_to_keep = 2
+        checkpoint_config = CheckpointConfig(
+            num_to_keep=num_to_keep, checkpoint_frequency=1
+        )
+
+        runner = TuneController(
+            resource_manager_factory=lambda: PlacementGroupResourceManager(),
+            storage=storage,
+            checkpoint_period=100,  # only rely on force syncing
+            trial_checkpoint_config=checkpoint_config,
+        )
+
+        class CheckpointingTrial(Trial):
+            def should_checkpoint(self):
+                return True
+
+            def get_json_state(self):
+                return "", ""
+
+        trial = CheckpointingTrial(
+            MOCK_TRAINABLE_NAME,
+            checkpoint_config=checkpoint_config,
+            stopping_criterion={"training_iteration": 10},
+            storage=storage,
+        )
+        runner.add_trial(trial)
+
+        # also check if the warning is printed
+        buffer = []
+        from ray.tune.execution.experiment_state import logger
+
+        with mock.patch.object(logger, "warning", lambda x: buffer.append(x)):
+            while not runner.is_finished():
+                runner.step()
+
+        assert any(
+            "Experiment state snapshotting has been triggered multiple times" in x
+            for x in buffer
+        )
+        # We should sync 6 times:
+        # The first checkpoint happens when the experiment starts,
+        # since no checkpoints have happened yet
+        # (This corresponds to the new_trial event in the runner loop)
+        # Then, every num_to_keep=2 checkpoints, we should perform a forced checkpoint
+        # which results in 5 more checkpoints (running for 10 iterations),
+        # giving a total of 6
+        assert sync_up.call_count == 6
+
+
+def test_checkpoint_force_by_trial_callback(ray_start_4_cpus_2_gpus_extra, tmp_path):
+    """Test that cloud syncing is forced if one of the trials has made more
+    than num_to_keep checkpoints since last sync.
+    Legacy test: test_trial_runner_3.py::TrialRunnerTest::
+        testCloudCheckpointForceWithNumToKeep
+    """
+
+    class CheckpointCallback(Callback):
+        def __init__(self):
+            self.num_checkpoints = 0
+
+        def on_trial_result(self, iteration, trials, trial: Trial, result, **info):
+            # Checkpoint every two iterations
+            if result[TRAINING_ITERATION] % 2 == 0:
+                self.num_checkpoints += 1
+                result["should_checkpoint"] = True
+
+    storage = mock_storage_context()
+
+    # disable automatic checkpointing
+    checkpoint_config = CheckpointConfig(checkpoint_frequency=0)
+    callback = CheckpointCallback()
+    runner = TuneController(
+        resource_manager_factory=PlacementGroupResourceManager,
+        storage=storage,
+        callbacks=[callback],
+        trial_checkpoint_config=checkpoint_config,
+    )
+
+    trial = Trial(
+        MOCK_TRAINABLE_NAME,
+        checkpoint_config=checkpoint_config,
+        stopping_criterion={"training_iteration": 6},
+        storage=storage,
+    )
+    runner.add_trial(trial)
+
+    while not runner.is_finished():
+        runner.step()
+
+    assert callback.num_checkpoints == 3
+    assert num_checkpoints(trial) == 3
+
+
+def test_checkpoint_sync_up_timeout(
+    ray_start_4_cpus_2_gpus_extra, tmp_path, monkeypatch
+):
+    """Test that trial runner experiment checkpointing times out correctly.
+
+    Legacy test: test_trial_runner_3.py::TrialRunnerTest::
+        testForcedCloudCheckpointSyncTimeout
+    """
+    storage = mock_storage_context(sync_config=ray.tune.SyncConfig(sync_timeout=0.5))
+    monkeypatch.setenv("TUNE_WARN_SLOW_EXPERIMENT_CHECKPOINT_SYNC_THRESHOLD_S", "0.25")
+
+    def _hanging_upload_to_fs_path(*args, **kwargs):
+        time.sleep(200)
+
+    monkeypatch.setattr(
+        ray.train._internal.storage,
+        "_upload_to_fs_path",
+        _hanging_upload_to_fs_path,
+    )
+
+    runner = TuneController(
+        resource_manager_factory=lambda: PlacementGroupResourceManager(),
+        storage=storage,
+    )
+
+    # Start a hanging sync that should not block the controller
+    runner.checkpoint()
+
+    buffer = []
+    logger = logging.getLogger("ray.tune.execution.experiment_state")
+    with mock.patch.object(logger, "error", lambda x, **kwargs: buffer.append(x)):
+        with mock.patch.object(logger, "warning", lambda x: buffer.append(x)):
+            runner.checkpoint(force=True, wait=True)
+
+    # We should see a log about the timeout
+    assert any("Saving experiment state to storage" in x for x in buffer)
+    # We should also have a warning about the slow upload
+    assert any("may be a performance bottleneck" in x for x in buffer)
+
+
+def test_checkpoint_sync_up_error(ray_start_4_cpus_2_gpus_extra, tmp_path, monkeypatch):
+    """Test that trial runner experiment checkpointing handles errors correctly."""
+    storage = mock_storage_context()
+
+    def _failing_upload_to_fs_path(*args, **kwargs):
+        raise RuntimeError("Upload failing...")
+
+    monkeypatch.setattr(
+        ray.train._internal.storage,
+        "_upload_to_fs_path",
+        _failing_upload_to_fs_path,
+    )
+
+    runner = TuneController(
+        resource_manager_factory=lambda: PlacementGroupResourceManager(),
+        storage=storage,
+    )
+
+    # Launching a failing upload task should not crash the controller / main thread
+    runner.checkpoint()
+
+    buffer = []
+    logger = logging.getLogger("ray.tune.execution.experiment_state")
+    with mock.patch.object(logger, "error", lambda x, **kwargs: buffer.append(x)):
+        runner.checkpoint(force=True)
+
+    # We should see a log about the failure
+    assert any("Saving experiment state to storage" in x for x in buffer)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/tune/tests/execution/test_controller_control_integration.py

@@ -0,0 +1,147 @@
+import sys
+from collections import Counter
+
+import pytest
+
+import ray
+from ray.air.execution import FixedResourceManager, PlacementGroupResourceManager
+from ray.train.tests.util import mock_storage_context
+from ray.tune import PlacementGroupFactory, register_trainable
+from ray.tune.execution.tune_controller import TuneController
+from ray.tune.experiment import Trial
+from ray.tune.utils.mock_trainable import MOCK_TRAINABLE_NAME, register_mock_trainable
+
+STORAGE = mock_storage_context()
+
+
+@pytest.fixture(scope="function")
+def ray_start_4_cpus_2_gpus_extra():
+    address_info = ray.init(num_cpus=4, num_gpus=2, resources={"a": 2})
+    yield address_info
+    ray.shutdown()
+
+
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+def test_stop_trial(ray_start_4_cpus_2_gpus_extra, resource_manager_cls):
+    """Stopping a trial while RUNNING or PENDING should work.
+
+    Legacy test: test_trial_runner_3.py::TrialRunnerTest::testStopTrial
+    """
+
+    register_mock_trainable()
+    runner = TuneController(
+        resource_manager_factory=lambda: resource_manager_cls(), storage=STORAGE
+    )
+    kwargs = {
+        "stopping_criterion": {"training_iteration": 10},
+        "placement_group_factory": PlacementGroupFactory([{"CPU": 2, "GPU": 1}]),
+        "config": {"sleep": 1},
+        "storage": STORAGE,
+    }
+    trials = [
+        Trial(MOCK_TRAINABLE_NAME, **kwargs),
+        Trial(MOCK_TRAINABLE_NAME, **kwargs),
+        Trial(MOCK_TRAINABLE_NAME, **kwargs),
+        Trial(MOCK_TRAINABLE_NAME, **kwargs),
+    ]
+    for t in trials:
+        runner.add_trial(t)
+
+    counter = Counter(t.status for t in trials)
+
+    # Wait until 2 trials started
+    while counter.get("RUNNING", 0) != 2:
+        runner.step()
+        counter = Counter(t.status for t in trials)
+
+    assert counter.get("RUNNING", 0) == 2
+    assert counter.get("PENDING", 0) == 2
+
+    # Stop trial that is running
+    for trial in trials:
+        if trial.status == Trial.RUNNING:
+            runner._schedule_trial_stop(trial)
+            break
+
+    counter = Counter(t.status for t in trials)
+
+    # Wait until the next trial started
+    while counter.get("RUNNING", 0) < 2:
+        runner.step()
+        counter = Counter(t.status for t in trials)
+
+    assert counter.get("RUNNING", 0) == 2
+    assert counter.get("TERMINATED", 0) == 1
+    assert counter.get("PENDING", 0) == 1
+
+    # Stop trial that is pending
+    for trial in trials:
+        if trial.status == Trial.PENDING:
+            runner._schedule_trial_stop(trial)
+            break
+
+    counter = Counter(t.status for t in trials)
+
+    # Wait until 2 trials are running again
+    while counter.get("RUNNING", 0) < 2:
+        runner.step()
+        counter = Counter(t.status for t in trials)
+
+    assert counter.get("RUNNING", 0) == 2
+    assert counter.get("TERMINATED", 0) == 2
+    assert counter.get("PENDING", 0) == 0
+
+
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+def test_remove_actor_tracking(ray_start_4_cpus_2_gpus_extra, resource_manager_cls):
+    """When we reuse actors, actors that have been requested but not started
+    should not be tracked in ``_stopping_actors``.
+
+    When actors are re-used, we cancel original actor requests for the trial.
+    If these actors haven't been alive, there won't be a stop future to be resolved,
+    and thus they would remain in ``TuneController._stopping_actors`` until they
+    get cleaned up after 600 seconds.
+
+    This test asserts that these actors are not tracked in
+    ``TuneController._stopping_actors`` at all.
+
+    We start 4 actors, and one can run at a time. Actors are re-used across trials.
+    When the experiment ends, we expect that only one actor is left to track
+    in ``self._stopping_trials``.
+    """
+    runner = TuneController(
+        resource_manager_factory=lambda: resource_manager_cls(),
+        reuse_actors=True,
+        storage=STORAGE,
+    )
+
+    def train_fn(config):
+        return 1
+
+    register_trainable("test_remove_actor_tracking", train_fn)
+
+    kwargs = {
+        "placement_group_factory": PlacementGroupFactory([{"CPU": 4, "GPU": 2}]),
+        "storage": STORAGE,
+    }
+    trials = [Trial("test_remove_actor_tracking", **kwargs) for i in range(4)]
+    for t in trials:
+        runner.add_trial(t)
+
+    while not runner.is_finished():
+        runner.step()
+
+    # Only one actor should be left to stop
+    assert len(runner._stopping_actors) == 1
+
+    runner.cleanup()
+
+    assert len(runner._stopping_actors) == 0
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", "--reruns", "3", __file__]))

python/ray/tune/tests/execution/test_controller_errors_integration.py

@@ -0,0 +1,212 @@
+import os
+import sys
+from collections import Counter
+
+import pytest
+
+import ray
+from ray.air.execution import FixedResourceManager, PlacementGroupResourceManager
+from ray.train.tests.util import mock_storage_context
+from ray.tune import CheckpointConfig, PlacementGroupFactory, TuneError
+from ray.tune.execution.tune_controller import TuneController
+from ray.tune.experiment import Trial
+from ray.tune.registry import TRAINABLE_CLASS, _global_registry
+from ray.tune.schedulers import FIFOScheduler
+from ray.tune.search import BasicVariantGenerator
+from ray.tune.tests.execution.utils import BudgetResourceManager
+from ray.tune.utils.mock_trainable import MOCK_TRAINABLE_NAME, register_mock_trainable
+
+STORAGE = mock_storage_context()
+
+
+@pytest.fixture(scope="function")
+def ray_start_4_cpus_2_gpus_extra():
+    address_info = ray.init(num_cpus=4, num_gpus=2, resources={"a": 2})
+    yield address_info
+    ray.shutdown()
+
+
+@pytest.fixture(autouse=True)
+def register_test_trainable():
+    register_mock_trainable()
+
+
+def create_mock_components():
+    class _MockScheduler(FIFOScheduler):
+        errored_trials = []
+
+        def on_trial_error(self, tune_controller, trial):
+            self.errored_trials += [trial]
+
+    class _MockSearchAlg(BasicVariantGenerator):
+        errored_trials = []
+
+        def on_trial_complete(self, trial_id, error=False, **kwargs):
+            if error:
+                self.errored_trials += [trial_id]
+
+    searchalg = _MockSearchAlg()
+    scheduler = _MockScheduler()
+    return searchalg, scheduler
+
+
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+def test_invalid_trainable(ray_start_4_cpus_2_gpus_extra, resource_manager_cls):
+    """An invalid trainable should make the trial fail on startup.
+
+    The controller itself should continue. Other trials should run.
+
+    Legacy test: test_trial_runner_2.py::TrialRunnerTest::testErrorHandling
+    """
+    runner = TuneController(
+        resource_manager_factory=lambda: resource_manager_cls(), storage=STORAGE
+    )
+    kwargs = {
+        "stopping_criterion": {"training_iteration": 1},
+        "placement_group_factory": PlacementGroupFactory([{"CPU": 1, "GPU": 1}]),
+        "storage": STORAGE,
+        "config": {"sleep": 0.5},
+    }
+    _global_registry.register(TRAINABLE_CLASS, "asdf", None)
+    trials = [Trial("asdf", **kwargs), Trial(MOCK_TRAINABLE_NAME, **kwargs)]
+    for t in trials:
+        runner.add_trial(t)
+
+    while not trials[1].status == Trial.RUNNING:
+        runner.step()
+    assert trials[0].status == Trial.ERROR
+    assert trials[1].status == Trial.RUNNING
+
+
+def test_overstep(ray_start_4_cpus_2_gpus_extra):
+    """Stepping when trials are finished should raise a TuneError.
+
+    Legacy test: test_trial_runner_2.py::TrialRunnerTest::testThrowOnOverstep
+    """
+    os.environ["TUNE_MAX_PENDING_TRIALS_PG"] = "1"
+    runner = TuneController(
+        resource_manager_factory=lambda: BudgetResourceManager({"CPU": 4}),
+        storage=STORAGE,
+    )
+    runner.step()
+    with pytest.raises(TuneError):
+        runner.step()
+
+
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+@pytest.mark.parametrize("max_failures_persistent", [(0, False), (1, False), (2, True)])
+def test_failure_recovery(
+    ray_start_4_cpus_2_gpus_extra, resource_manager_cls, max_failures_persistent
+):
+    """Test failure recover with `max_failures`.
+
+    Trials should be retried up to `max_failures` times.
+
+    Legacy test: test_trial_runner_2.py::TrialRunnerTest::testFailureRecoveryDisabled
+    Legacy test: test_trial_runner_2.py::TrialRunnerTest::testFailureRecoveryEnabled
+    Legacy test: test_trial_runner_2.py::TrialRunnerTest::testFailureRecoveryMaxFailures
+    """
+    max_failures, persistent_error = max_failures_persistent
+    searchalg, scheduler = create_mock_components()
+
+    runner = TuneController(
+        search_alg=searchalg,
+        scheduler=scheduler,
+        resource_manager_factory=lambda: resource_manager_cls(),
+        storage=STORAGE,
+    )
+    kwargs = {
+        "placement_group_factory": PlacementGroupFactory([{"CPU": 1, "GPU": 1}]),
+        "stopping_criterion": {"training_iteration": 2},
+        "checkpoint_config": CheckpointConfig(checkpoint_frequency=1),
+        "max_failures": max_failures,
+        "config": {"mock_error": True, "persistent_error": persistent_error},
+        "storage": STORAGE,
+    }
+    runner.add_trial(Trial(MOCK_TRAINABLE_NAME, **kwargs))
+    trials = runner.get_trials()
+
+    while not runner.is_finished():
+        runner.step()
+
+    if persistent_error or not max_failures:
+        assert trials[0].status == Trial.ERROR
+
+        num_failures = max_failures + 1
+        assert trials[0].num_failures == num_failures
+        # search alg receives on_complete, so only after the max failures
+        # have been exhausted. Thus, it only has errored_trials if the
+        # trial fails even in the last try.
+        assert len(searchalg.errored_trials) == 1
+        # search alg receives on_error, so every failure is registered.
+        assert len(scheduler.errored_trials) == num_failures
+    else:
+        assert trials[0].status == Trial.TERMINATED
+        assert trials[0].num_failures == 1
+        assert len(searchalg.errored_trials) == 0
+        assert len(scheduler.errored_trials) == 1
+
+
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+@pytest.mark.parametrize("fail_fast", [True, TuneController.RAISE])
+def test_fail_fast(ray_start_4_cpus_2_gpus_extra, resource_manager_cls, fail_fast):
+    """Test fail_fast feature.
+
+    If fail_fast=True, after the first failure, all other trials should be terminated
+    (because we end the experiment).
+
+    If fail_fast=RAISE, after the first failure, we should raise an error.
+
+    Legacy test: test_trial_runner_2.py::TrialRunnerTest::testFailFast
+    Legacy test: test_trial_runner_2.py::TrialRunnerTest::testFailFastRaise
+    """
+
+    runner = TuneController(
+        resource_manager_factory=lambda: resource_manager_cls(),
+        fail_fast=fail_fast,
+        storage=STORAGE,
+    )
+    kwargs = {
+        "placement_group_factory": PlacementGroupFactory([{"CPU": 1, "GPU": 1}]),
+        "checkpoint_config": CheckpointConfig(checkpoint_frequency=1),
+        "max_failures": 0,
+        "config": {
+            "mock_error": True,
+            "persistent_error": True,
+        },
+        "storage": STORAGE,
+    }
+    runner.add_trial(Trial(MOCK_TRAINABLE_NAME, **kwargs))
+    runner.add_trial(Trial(MOCK_TRAINABLE_NAME, **kwargs))
+    trials = runner.get_trials()
+
+    if fail_fast == TuneController.RAISE:
+        with pytest.raises(Exception):
+            while not runner.is_finished():
+                runner.step()
+        runner.cleanup()
+        return
+    else:
+        while not runner.is_finished():
+            runner.step()
+
+    status_count = Counter(t.status for t in trials)
+
+    # One trial failed
+    assert status_count.get(Trial.ERROR) == 1
+    # The other one was pre-empted
+    assert status_count.get(Trial.TERMINATED) == 1
+
+    # Controller finished
+    with pytest.raises(TuneError):
+        runner.step()
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/tune/tests/execution/test_controller_resources_integration.py

@@ -0,0 +1,278 @@
+import os
+import sys
+import time
+from collections import Counter
+
+import pytest
+
+import ray
+from ray import tune
+from ray.air.execution import FixedResourceManager, PlacementGroupResourceManager
+from ray.train.tests.util import mock_storage_context
+from ray.tune import PlacementGroupFactory, TuneError
+from ray.tune.execution.tune_controller import TuneController
+from ray.tune.experiment import Trial
+from ray.tune.schedulers import FIFOScheduler, TrialScheduler
+from ray.tune.search import BasicVariantGenerator
+from ray.tune.utils.mock import TrialStatusSnapshot, TrialStatusSnapshotTaker
+from ray.tune.utils.mock_trainable import MOCK_TRAINABLE_NAME, register_mock_trainable
+
+STORAGE = mock_storage_context()
+
+
+@pytest.fixture(autouse=True)
+def register_test_trainable():
+    register_mock_trainable()
+
+
+@pytest.fixture(scope="function")
+def ray_start_4_cpus_2_gpus_extra():
+    address_info = ray.init(num_cpus=4, num_gpus=2, resources={"a": 2})
+    yield address_info
+    ray.shutdown()
+
+
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+@pytest.mark.parametrize(
+    "bundles",
+    [
+        [{"CPU": 1}, {"CPU": 3, "GPU": 1}],
+        [{"CPU": 1, "a": 2}],
+        [{"CPU": 1}, {"a": 2}],
+        [{"CPU": 1, "GPU": 1}, {"GPU": 1}],
+    ],
+)
+def test_resource_parallelism_single(
+    ray_start_4_cpus_2_gpus_extra, resource_manager_cls, bundles
+):
+    """Test that extra and custom resources are respected for parallelism.
+
+    We schedule two trials with resources according to the bundle. If only
+    the head bundle or only CPU/GPU resources were considered, both trials
+    could run in parallel.
+
+    However, we assert that the resources in child bundles and extra resources
+    are respected and only one trial runs in parallel.
+
+    Legacy test: test_trial_runner.py::TrialRunnerTest::testExtraResources
+    Legacy test: test_trial_runner.py::TrialRunnerTest::testCustomResources
+    Legacy test: test_trial_runner.py::TrialRunnerTest::testExtraCustomResources
+    Legacy test: test_trial_runner.py::TrialRunnerTest::testResourceScheduler
+    """
+    snapshot = TrialStatusSnapshot()
+    runner = TuneController(
+        resource_manager_factory=lambda: resource_manager_cls(),
+        callbacks=[TrialStatusSnapshotTaker(snapshot)],
+        storage=STORAGE,
+    )
+    kwargs = {
+        "stopping_criterion": {"training_iteration": 1},
+        "placement_group_factory": PlacementGroupFactory(bundles),
+        "storage": STORAGE,
+    }
+    trials = [
+        Trial(MOCK_TRAINABLE_NAME, **kwargs),
+        Trial(MOCK_TRAINABLE_NAME, **kwargs),
+    ]
+    for t in trials:
+        runner.add_trial(t)
+
+    while not runner.is_finished():
+        runner.step()
+
+    assert snapshot.max_running_trials() == 1
+    assert snapshot.all_trials_are_terminated()
+
+
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+def test_fractional_gpus(ray_start_4_cpus_2_gpus_extra, resource_manager_cls):
+    """Test that fractional GPUs lead to more parallelism.
+
+    We schedule four trials with 0.75 GPUs each. Since our cluster has 2 GPUs,
+    we should be able to run 2 trials in parallel.
+
+    Legacy test: test_trial_runner.py::TrialRunnerTest::testFractionalGpus
+    """
+    snapshot = TrialStatusSnapshot()
+    runner = TuneController(
+        resource_manager_factory=lambda: resource_manager_cls(),
+        callbacks=[TrialStatusSnapshotTaker(snapshot)],
+        storage=STORAGE,
+    )
+    kwargs = {
+        "stopping_criterion": {"training_iteration": 1},
+        "placement_group_factory": PlacementGroupFactory([{"GPU": 0.75}]),
+        "config": {
+            "sleep": 1,
+        },
+        "storage": STORAGE,
+    }
+    trials = [Trial(MOCK_TRAINABLE_NAME, **kwargs) for i in range(4)]
+    for t in trials:
+        runner.add_trial(t)
+
+    while not runner.is_finished():
+        runner.step()
+
+    assert snapshot.max_running_trials() == 2
+    assert snapshot.all_trials_are_terminated()
+
+
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+def test_multi_step(ray_start_4_cpus_2_gpus_extra, resource_manager_cls):
+    """Test that trials can run for more than one iteration.
+
+    Todo (krfricke): This is not a resource test, so it should be moved.
+
+    Legacy test: test_trial_runner.py::TrialRunnerTest::testMultiStepRun
+    Legacy test: test_trial_runner.py::TrialRunnerTest::testMultiStepRun2
+    """
+    snapshot = TrialStatusSnapshot()
+    runner = TuneController(
+        resource_manager_factory=lambda: resource_manager_cls(),
+        callbacks=[TrialStatusSnapshotTaker(snapshot)],
+        storage=STORAGE,
+    )
+    kwargs = {
+        "stopping_criterion": {"training_iteration": 5},
+        "placement_group_factory": PlacementGroupFactory([{"CPU": 1, "GPU": 1}]),
+        "storage": STORAGE,
+    }
+    trials = [Trial(MOCK_TRAINABLE_NAME, **kwargs) for i in range(2)]
+    for t in trials:
+        runner.add_trial(t)
+
+    while not runner.is_finished():
+        runner.step()
+
+    # Overstepping should throw error
+    # test_trial_runner.py::TrialRunnerTest::testMultiStepRun2
+    with pytest.raises(TuneError):
+        runner.step()
+
+    assert snapshot.all_trials_are_terminated()
+    assert all(t.last_result["training_iteration"] == 5 for t in runner.get_trials())
+
+
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+def test_resources_changing(ray_start_4_cpus_2_gpus_extra, resource_manager_cls):
+    """Checks that resource requirements can be changed on fly.
+
+    Legacy test: test_trial_runner.py::TrialRunnerTest::testChangeResources
+    """
+
+    class ChangingScheduler(FIFOScheduler):
+        def on_trial_result(self, tune_controller, trial, result):
+            if result["training_iteration"] == 1:
+                # NOTE: This is a hack to get around the new pausing logic,
+                # which doesn't set the trial status to PAUSED immediately.
+                orig_status = trial.status
+                trial.set_status(Trial.PAUSED)
+                trial.update_resources(dict(cpu=4, gpu=0))
+                trial.set_status(orig_status)
+                return TrialScheduler.PAUSE
+            return TrialScheduler.NOOP
+
+    scheduler = ChangingScheduler()
+    runner = TuneController(
+        resource_manager_factory=lambda: resource_manager_cls(),
+        scheduler=scheduler,
+        storage=STORAGE,
+    )
+    kwargs = {
+        "stopping_criterion": {"training_iteration": 2},
+        "placement_group_factory": PlacementGroupFactory([{"CPU": 2, "GPU": 0}]),
+        "storage": STORAGE,
+    }
+    trials = [Trial(MOCK_TRAINABLE_NAME, **kwargs)]
+    for t in trials:
+        runner.add_trial(t)
+
+    while not trials[0].status == Trial.RUNNING:
+        runner.step()
+
+    assert trials[0].status == Trial.RUNNING
+    assert runner._actor_manager.get_live_actors_resources().get("CPU") == 2
+
+    with pytest.raises(ValueError):
+        trials[0].update_resources(dict(cpu=4, gpu=0))
+
+    while trials[0].status == Trial.RUNNING:
+        runner.step()
+
+    assert trials[0].status == Trial.PAUSED
+
+    while not trials[0].status == Trial.RUNNING:
+        runner.step()
+
+    assert runner._actor_manager.get_live_actors_resources().get("CPU") == 4
+
+    runner.step()
+
+
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+def test_queue_filling(ray_start_4_cpus_2_gpus_extra, resource_manager_cls):
+    """Checks that the trial queue is filled even if only 1 pending trial is allowed.
+
+    Legacy test: test_trial_runner.py::TrialRunnerTest::testQueueFilling
+    """
+    os.environ["TUNE_MAX_PENDING_TRIALS_PG"] = "1"
+
+    def f1(config):
+        for i in range(10):
+            yield i
+            time.sleep(1)
+
+    tune.register_trainable("f1", f1)
+
+    search_alg = BasicVariantGenerator()
+    search_alg.add_configurations(
+        {
+            "foo": {
+                "run": "f1",
+                "num_samples": 100,
+                "config": {
+                    "a": tune.sample_from(lambda spec: 5.0 / 7),
+                    "b": tune.sample_from(lambda spec: "long" * 40),
+                },
+                "resources_per_trial": {"cpu": 2},
+            }
+        }
+    )
+
+    runner = TuneController(
+        resource_manager_factory=lambda: resource_manager_cls(),
+        search_alg=search_alg,
+        storage=STORAGE,
+    )
+
+    while len(runner.get_trials()) < 3:
+        runner.step()
+
+    # All trials are enqueued
+    assert len(runner.get_trials()) == 3
+
+    status_count = Counter(t.status for t in runner.get_trials())
+    while status_count.get(Trial.RUNNING, 0) < 2 and not runner.is_finished():
+        runner.step()
+        status_count = Counter(t.status for t in runner.get_trials())
+
+    assert len(runner.get_trials()) == 3
+
+    status_count = Counter(t.status for t in runner.get_trials())
+    assert status_count.get(Trial.RUNNING) == 2
+    assert status_count.get(Trial.PENDING) == 1
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", "--reruns", "3", __file__]))

python/ray/tune/tests/execution/test_controller_resume_integration.py

@@ -0,0 +1,525 @@
+import os
+import sys
+from unittest.mock import patch
+
+import pandas as pd
+import pytest
+
+import ray
+from ray import tune
+from ray.air.execution import FixedResourceManager, PlacementGroupResourceManager
+from ray.train.tests.util import mock_storage_context
+from ray.tune import CheckpointConfig, Experiment, PlacementGroupFactory, ResumeConfig
+from ray.tune.execution.tune_controller import TuneController
+from ray.tune.experiment import Trial
+from ray.tune.impl.placeholder import create_resolvers_map, inject_placeholders
+from ray.tune.search import BasicVariantGenerator
+from ray.tune.utils.mock_trainable import (
+    MOCK_ERROR_KEY,
+    MOCK_TRAINABLE_NAME,
+    register_mock_trainable,
+)
+
+STORAGE = mock_storage_context()
+
+
+@pytest.fixture(autouse=True)
+def register_test_trainable():
+    register_mock_trainable()
+
+
+@pytest.fixture(scope="function")
+def ray_start_4_cpus_2_gpus_extra():
+    address_info = ray.init(num_cpus=4, num_gpus=2, resources={"a": 2})
+    yield address_info
+    ray.shutdown()
+
+
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+def test_controller_restore_dataset_references(
+    ray_start_4_cpus_2_gpus_extra, resource_manager_cls
+):
+    """Check that references to Ray Datasets are replaced on resume.
+
+    Legacy test: test_trial_runner_3.py::TrialRunnerTest::
+        testSearcherCorrectReferencesAfterRestore
+    """
+
+    class FakeDataset:
+        def __init__(self, name):
+            self.name = name
+
+    config = {
+        "param1": {
+            "param2": tune.grid_search(
+                [FakeDataset("1"), FakeDataset("2"), FakeDataset("3")]
+            ),
+        },
+        "param4": tune.sample_from(lambda: 1),
+        "param5": tune.sample_from(lambda spec: spec.config["param1"]["param2"]),
+    }
+    resolvers = create_resolvers_map()
+    config = inject_placeholders(config, resolvers)
+
+    def create_searcher():
+        search_alg = BasicVariantGenerator()
+        experiment_spec = {
+            "run": MOCK_TRAINABLE_NAME,
+            "stop": {"training_iteration": 2},
+            "config": config,
+        }
+        experiments = [Experiment.from_json("test", experiment_spec)]
+        search_alg.add_configurations(experiments)
+        return search_alg
+
+    searcher = create_searcher()
+
+    restored_config = {
+        "param1": {
+            "param2": tune.grid_search(
+                [FakeDataset("4"), FakeDataset("5"), FakeDataset("6")]
+            ),
+        },
+        "param4": tune.sample_from(lambda: 8),
+        "param5": tune.sample_from(lambda spec: spec["config"]["param1"]["param2"]),
+    }
+    replaced_resolvers = create_resolvers_map()
+    inject_placeholders(restored_config, replaced_resolvers)
+
+    runner = TuneController(
+        resource_manager_factory=lambda: resource_manager_cls(),
+        reuse_actors=False,
+        search_alg=searcher,
+        placeholder_resolvers=replaced_resolvers,
+        checkpoint_period=-1,
+        storage=STORAGE,
+    )
+
+    while len(runner.get_trials()) < 3 or any(
+        trial.status not in {Trial.RUNNING, Trial.TERMINATED}
+        for trial in runner.get_trials()
+    ):
+        runner.step()
+
+    assert len(runner.get_trials()) == 3, [t.config for t in runner.get_trials()]
+    for t in runner.get_trials():
+        # Make sure that all the trials carry updated config values.
+        assert t.config["param1"]["param2"].name in ["4", "5", "6"]
+        assert t.config["param4"] == 8
+        assert t.config["param5"].name in ["4", "5", "6"]
+
+
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+def test_controller_restore_no_error_resume(
+    ray_start_4_cpus_2_gpus_extra, resource_manager_cls
+):
+    """Check that `resume=True` does not resume errored trials.
+
+    Legacy test: test_trial_runner_3.py::TrialRunnerTest::testTrialErrorResumeFalse
+    """
+    runner = TuneController(
+        resource_manager_factory=lambda: resource_manager_cls(),
+        storage=STORAGE,
+    )
+
+    kwargs = {
+        "stopping_criterion": {"training_iteration": 4},
+        "placement_group_factory": PlacementGroupFactory([{"CPU": 1, "GPU": 0}]),
+        "storage": STORAGE,
+    }
+    trials = [
+        Trial(MOCK_TRAINABLE_NAME, config={MOCK_ERROR_KEY: True}, **kwargs),
+        Trial(MOCK_TRAINABLE_NAME, **kwargs),
+        Trial(MOCK_TRAINABLE_NAME, **kwargs),
+    ]
+    for t in trials:
+        runner.add_trial(t)
+
+    while not runner.is_finished():
+        runner.step()
+
+    runner.checkpoint(force=True, wait=True)
+
+    assert trials[0].status == Trial.ERROR
+    del runner
+
+    new_runner = TuneController(
+        resource_manager_factory=lambda: resource_manager_cls(),
+        storage=STORAGE,
+        resume_config=ResumeConfig(
+            unfinished=ResumeConfig.ResumeType.RESUME,
+            errored=ResumeConfig.ResumeType.SKIP,
+            finished=ResumeConfig.ResumeType.SKIP,
+        ),
+    )
+
+    assert len(new_runner.get_trials()) == 3
+    assert Trial.ERROR in (t.status for t in new_runner.get_trials())
+
+
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+def test_controller_restore_error_only_resume(
+    ray_start_4_cpus_2_gpus_extra, resource_manager_cls
+):
+    """Check that `resume=ERRORED_ONLY` only resumes errored trials.
+
+    Legacy test: test_trial_runner_3.py::TrialRunnerTest::testTrialErrorResumeTrue
+    """
+    runner = TuneController(
+        resource_manager_factory=lambda: resource_manager_cls(),
+        storage=STORAGE,
+    )
+    kwargs = {
+        "stopping_criterion": {"training_iteration": 4},
+        "placement_group_factory": PlacementGroupFactory([{"CPU": 1, "GPU": 0}]),
+        "storage": STORAGE,
+    }
+    trials = [
+        Trial(MOCK_TRAINABLE_NAME, config={MOCK_ERROR_KEY: True}, **kwargs),
+        Trial(MOCK_TRAINABLE_NAME, **kwargs),
+        Trial(MOCK_TRAINABLE_NAME, **kwargs),
+    ]
+    for t in trials:
+        runner.add_trial(t)
+
+    while not runner.is_finished():
+        runner.step()
+
+    runner.checkpoint(force=True, wait=True)
+
+    assert trials[0].status == Trial.ERROR
+    del runner
+
+    new_runner = TuneController(
+        resource_manager_factory=lambda: resource_manager_cls(),
+        storage=STORAGE,
+        resume_config=ResumeConfig(
+            unfinished=ResumeConfig.ResumeType.SKIP,
+            errored=ResumeConfig.ResumeType.RESUME,
+            finished=ResumeConfig.ResumeType.SKIP,
+        ),
+    )
+
+    assert len(new_runner.get_trials()) == 3
+    assert Trial.ERROR not in (t.status for t in new_runner.get_trials())
+    # The below is just a check for standard behavior.
+    disable_error = False
+    for t in new_runner.get_trials():
+        if t.config.get(MOCK_ERROR_KEY):
+            t.config[MOCK_ERROR_KEY] = False
+            disable_error = True
+    assert disable_error
+
+    while not new_runner.is_finished():
+        new_runner.step()
+    assert Trial.ERROR not in (t.status for t in new_runner.get_trials())
+
+
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+def test_controller_restore_trial_save_restore(
+    ray_start_4_cpus_2_gpus_extra, resource_manager_cls
+):
+    """Creates different trials to test runner.checkpoint/restore.
+
+    Legacy test: test_trial_runner_3.py::TrialRunnerTest::testTrialSaveRestore
+    """
+    runner = TuneController(
+        resource_manager_factory=lambda: resource_manager_cls(),
+        checkpoint_period=0,
+        storage=STORAGE,
+    )
+    trials = [
+        Trial(
+            MOCK_TRAINABLE_NAME,
+            trial_id="trial_terminate",
+            stopping_criterion={"training_iteration": 1},
+            checkpoint_config=CheckpointConfig(checkpoint_frequency=1),
+            storage=STORAGE,
+        )
+    ]
+    runner.add_trial(trials[0])
+    while not runner.is_finished():
+        # Start trial, process result, dispatch save and process save.
+        runner.step()
+    assert trials[0].status == Trial.TERMINATED
+
+    trials += [
+        Trial(
+            MOCK_TRAINABLE_NAME,
+            trial_id="trial_fail",
+            stopping_criterion={"training_iteration": 3},
+            checkpoint_config=CheckpointConfig(checkpoint_frequency=1),
+            config={MOCK_ERROR_KEY: True},
+            storage=STORAGE,
+        )
+    ]
+    runner.add_trial(trials[1])
+    while not runner.is_finished():
+        runner.step()
+    assert trials[1].status == Trial.ERROR
+
+    trials += [
+        Trial(
+            MOCK_TRAINABLE_NAME,
+            trial_id="trial_succ",
+            stopping_criterion={"training_iteration": 2},
+            checkpoint_config=CheckpointConfig(checkpoint_frequency=1),
+            storage=STORAGE,
+        )
+    ]
+    runner.add_trial(trials[2])
+
+    while not trials[2].status == Trial.RUNNING:
+        runner.step()  # Start trial
+    assert len(runner._get_trial_checkpoints()) == 3
+
+    runner.checkpoint(force=True, wait=True)
+
+    runner2 = TuneController(
+        resource_manager_factory=lambda: resource_manager_cls(),
+        storage=STORAGE,
+        resume_config=ResumeConfig(
+            unfinished=ResumeConfig.ResumeType.RESUME,
+            errored=ResumeConfig.ResumeType.SKIP,
+            finished=ResumeConfig.ResumeType.SKIP,
+        ),
+    )
+    for tid in ["trial_terminate", "trial_fail"]:
+        original_trial = runner.get_trial(tid)
+        restored_trial = runner2.get_trial(tid)
+        assert original_trial.status == restored_trial.status
+
+    restored_trial = runner2.get_trial("trial_succ")
+    assert Trial.PENDING == restored_trial.status
+
+    while not runner2.is_finished():
+        runner2.step()
+    assert restored_trial.status == Trial.TERMINATED
+
+
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+def test_controller_restore_trial_no_checkpoint_save(
+    ray_start_4_cpus_2_gpus_extra, resource_manager_cls
+):
+    """Check that non-checkpointing trials *are* saved.
+
+    Legacy test: test_trial_runner_3.py::TrialRunnerTest::testTrialNoCheckpointSave
+    """
+    with patch.dict(os.environ, {"TUNE_MAX_PENDING_TRIALS_PG": "1"}):
+        runner = TuneController(
+            resource_manager_factory=lambda: resource_manager_cls(),
+            checkpoint_period=0,
+            storage=STORAGE,
+        )
+
+        runner.add_trial(
+            Trial(
+                MOCK_TRAINABLE_NAME,
+                trial_id="non_checkpoint",
+                stopping_criterion={"training_iteration": 2},
+                storage=STORAGE,
+            )
+        )
+
+        while not all(t.status == Trial.TERMINATED for t in runner.get_trials()):
+            runner.step()
+
+        runner.add_trial(
+            Trial(
+                MOCK_TRAINABLE_NAME,
+                trial_id="checkpoint",
+                checkpoint_config=CheckpointConfig(
+                    checkpoint_at_end=True,
+                ),
+                stopping_criterion={"training_iteration": 2},
+                storage=STORAGE,
+            )
+        )
+
+        while not all(t.status == Trial.TERMINATED for t in runner.get_trials()):
+            runner.step()
+
+        runner.add_trial(
+            Trial(
+                MOCK_TRAINABLE_NAME,
+                trial_id="pending",
+                stopping_criterion={"training_iteration": 2},
+                storage=STORAGE,
+            )
+        )
+
+        old_trials = runner.get_trials()
+        while not old_trials[2].has_reported_at_least_once:
+            runner.step()
+
+        runner.checkpoint(force=True, wait=True)
+
+        runner2 = TuneController(
+            resource_manager_factory=lambda: resource_manager_cls(),
+            storage=STORAGE,
+            resume_config=ResumeConfig(
+                unfinished=ResumeConfig.ResumeType.RESUME,
+                errored=ResumeConfig.ResumeType.SKIP,
+                finished=ResumeConfig.ResumeType.SKIP,
+            ),
+        )
+        new_trials = runner2.get_trials()
+        assert len(new_trials) == 3
+        assert runner2.get_trial("non_checkpoint").status == Trial.TERMINATED
+        assert runner2.get_trial("checkpoint").status == Trial.TERMINATED
+        assert runner2.get_trial("pending").status == Trial.PENDING
+        assert runner2.get_trial("pending").has_reported_at_least_once
+        runner2.step()
+
+
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+def test_controller_restore_checkpoint_overwrite(
+    ray_start_4_cpus_2_gpus_extra, resource_manager_cls
+):
+    """Check that experiment state checkpoint are not overwritten on continue.
+
+    Legacy test: test_trial_runner_3.py::TrialRunnerTest::testCheckpointOverwrite
+    """
+    storage = mock_storage_context()
+
+    def count_checkpoints(cdir):
+        return sum(
+            (fname.startswith("experiment_state") and fname.endswith(".json"))
+            for fname in os.listdir(cdir)
+        )
+
+    tmpdir = storage.experiment_driver_staging_path
+    # The Trial `local_dir` must match the TrialRunner `local_checkpoint_dir`
+    # to match the directory structure assumed by `TrialRunner.resume`.
+    # See `test_trial_runner2.TrialRunnerTest2.testPauseResumeCheckpointCount`
+    # for more details.
+    trial = Trial(
+        MOCK_TRAINABLE_NAME,
+        checkpoint_config=CheckpointConfig(checkpoint_frequency=1),
+        storage=storage,
+    )
+    runner = TuneController(
+        resource_manager_factory=lambda: resource_manager_cls(),
+        storage=storage,
+        checkpoint_period=0,
+    )
+    runner.add_trial(trial)
+    while not trial.status == Trial.RUNNING:
+        runner.step()
+    # force checkpoint
+    runner.checkpoint(force=True, wait=True)
+    # Only one experiment state file
+    assert count_checkpoints(tmpdir) == 1
+
+    runner2 = TuneController(
+        resource_manager_factory=lambda: resource_manager_cls(),
+        storage=storage,
+        resume_config=ResumeConfig(
+            unfinished=ResumeConfig.ResumeType.RESUME,
+            errored=ResumeConfig.ResumeType.SKIP,
+            finished=ResumeConfig.ResumeType.SKIP,
+        ),
+    )
+    trial = runner2.get_trials()[0]
+    while not trial.status == Trial.RUNNING:
+        runner2.step()
+    # After resume, we have a new experiment state file in the directory
+    assert count_checkpoints(tmpdir) == 2
+
+    runner2.checkpoint()
+    assert count_checkpoints(tmpdir) == 2
+
+
+@pytest.mark.skip("TODO(justinvyu): Data lineage serialization context is broken.")
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+def test_controller_restore_with_dataset(
+    ray_start_4_cpus_2_gpus_extra, resource_manager_cls, tmp_path
+):
+    """Test trial runner checkpointing where trials contain Datasets.
+    When possible, a dataset plan should be saved (for read_* APIs).
+    See `Dataset.serialize_lineage` for more information.
+
+    If a dataset cannot be serialized, an experiment checkpoint
+    should still be created. Users can pass in the dataset again by
+    re-specifying the `param_space`.
+
+    Legacy test: test_trial_runner_3.py::TrialRunnerTest::
+        testExperimentCheckpointWithDatasets
+    """
+    # Save some test data to load
+    data_filepath = os.path.join(tmp_path, "test.csv")
+    pd.DataFrame({"x": list(range(10))}).to_csv(data_filepath)
+
+    def create_trial_config():
+        return {
+            "datasets": {
+                "with_lineage": ray.data.read_csv(data_filepath),
+                "no_lineage": ray.data.from_items([{"x": i} for i in range(10)]),
+            }
+        }
+
+    resolvers = create_resolvers_map()
+    config_with_placeholders = inject_placeholders(create_trial_config(), resolvers)
+    trial = Trial(
+        MOCK_TRAINABLE_NAME,
+        config=config_with_placeholders,
+        storage=STORAGE,
+    )
+    trial.init_local_path()
+    runner = TuneController(
+        resource_manager_factory=lambda: resource_manager_cls(),
+        storage=STORAGE,
+        placeholder_resolvers=resolvers,
+    )
+    runner.add_trial(trial)
+    # Req: TrialRunner checkpointing shouldn't error
+    runner.checkpoint(force=True, wait=True)
+
+    # Manually clear all block refs that may have been created
+    ray.shutdown()
+    ray.init(num_cpus=2)
+
+    register_mock_trainable()
+    new_runner = TuneController(
+        resource_manager_factory=lambda: resource_manager_cls(),
+        storage=STORAGE,
+    )
+    new_runner.resume(resume_config=ResumeConfig())
+    [loaded_trial] = new_runner.get_trials()
+    loaded_datasets = loaded_trial.config["datasets"]
+
+    # Req: The deserialized dataset (w/ lineage) should be usable.
+    assert [el["x"] for el in loaded_datasets["with_lineage"].take()] == list(range(10))
+
+    replaced_resolvers = create_resolvers_map()
+    inject_placeholders(create_trial_config(), replaced_resolvers)
+
+    respecified_config_runner = TuneController(
+        resource_manager_factory=lambda: resource_manager_cls(),
+        storage=STORAGE,
+        placeholder_resolvers=replaced_resolvers,
+    )
+    respecified_config_runner.resume(resume_config=ResumeConfig())
+    [loaded_trial] = respecified_config_runner.get_trials()
+    ray_ds_no_lineage = loaded_trial.config["datasets"]["no_lineage"]
+
+    # Req: The dataset (w/o lineage) can be re-specified and is usable after.
+    assert [el["x"] for el in ray_ds_no_lineage.take()] == list(range(10))
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/tune/tests/execution/test_controller_search_alg_integration.py

@@ -0,0 +1,431 @@
+import os
+import pickle
+import sys
+from collections import Counter
+
+import pytest
+
+import ray
+from ray.air.constants import TRAINING_ITERATION
+from ray.air.execution import FixedResourceManager, PlacementGroupResourceManager
+from ray.train.tests.util import mock_storage_context
+from ray.tune import Experiment, PlacementGroupFactory
+from ray.tune.execution.tune_controller import TuneController, _get_max_pending_trials
+from ray.tune.experiment import Trial
+from ray.tune.schedulers import FIFOScheduler, TrialScheduler
+from ray.tune.search import ConcurrencyLimiter, Repeater, Searcher, SearchGenerator
+from ray.tune.search._mock import _MockSearcher, _MockSuggestionAlgorithm
+from ray.tune.utils.mock_trainable import MOCK_TRAINABLE_NAME, register_mock_trainable
+
+
+class TestTuneController(TuneController):
+    def __init__(self, *args, **kwargs):
+        kwargs.update(dict(storage=mock_storage_context()))
+        super().__init__(*args, **kwargs)
+
+
+@pytest.fixture(autouse=True)
+def register_test_trainable():
+    register_mock_trainable()
+    yield
+
+
+@pytest.fixture(scope="function")
+def ray_start_8_cpus():
+    address_info = ray.init(num_cpus=8, num_gpus=0)
+    yield address_info
+    ray.shutdown()
+
+
+@pytest.fixture(scope="function")
+def ray_start_4_cpus_2_gpus_extra():
+    address_info = ray.init(num_cpus=4, num_gpus=2, resources={"a": 2})
+    yield address_info
+    ray.shutdown()
+
+
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+def test_search_alg_notification(ray_start_4_cpus_2_gpus_extra, resource_manager_cls):
+    """Check that the searchers gets notified of trial results + completions.
+
+    Also check that the searcher is "finished" before the runner, i.e. the runner
+    continues processing trials when the searcher finished.
+
+    Legacy test: test_trial_runner_3.py::TrialRunnerTest::testSearchAlgNotification
+    Legacy test: test_trial_runner_3.py::TrialRunnerTest::testSearchAlgFinished
+    """
+
+    experiment_spec = {"run": MOCK_TRAINABLE_NAME, "stop": {"training_iteration": 2}}
+    experiments = [Experiment.from_json("test", experiment_spec)]
+    search_alg = _MockSuggestionAlgorithm()
+    searcher = search_alg.searcher
+    search_alg.add_configurations(experiments)
+
+    runner = TestTuneController(
+        resource_manager_factory=lambda: resource_manager_cls(), search_alg=search_alg
+    )
+
+    # Run until trial is running
+    while not search_alg.is_finished():
+        runner.step()
+
+    trials = runner.get_trials()
+
+    # Make sure trial started
+    while trials[0].status != Trial.RUNNING:
+        runner.step()
+
+    assert trials[0].status == Trial.RUNNING
+    assert search_alg.is_finished()
+    assert not runner.is_finished()
+
+    # Run until everything finished
+    while not runner.is_finished():
+        runner.step()
+
+    assert trials[0].status == Trial.TERMINATED
+    assert search_alg.is_finished()
+    assert runner.is_finished()
+
+    assert searcher.counter["result"] == 1
+    assert searcher.counter["complete"] == 1
+
+
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+def test_search_alg_scheduler_stop(ray_start_4_cpus_2_gpus_extra, resource_manager_cls):
+    """Check that a scheduler-issued stop also notifies the search algorithm.
+
+    Legacy test: test_trial_runner_3.py::TrialRunnerTest::testSearchAlgSchedulerInteraction  # noqa
+    """
+
+    class _MockScheduler(FIFOScheduler):
+        def on_trial_result(self, *args, **kwargs):
+            return TrialScheduler.STOP
+
+    experiment_spec = {"run": MOCK_TRAINABLE_NAME, "stop": {"training_iteration": 5}}
+    experiments = [Experiment.from_json("test", experiment_spec)]
+    search_alg = _MockSuggestionAlgorithm()
+    searcher = search_alg.searcher
+    search_alg.add_configurations(experiments)
+
+    runner = TestTuneController(
+        resource_manager_factory=lambda: resource_manager_cls(),
+        search_alg=search_alg,
+        scheduler=_MockScheduler(),
+    )
+
+    trials = runner.get_trials()
+
+    while not runner.is_finished():
+        runner.step()
+
+    # Result is not processed because trial stop takes precedence
+    assert searcher.counter["result"] == 0
+    # But on_trial_complete is triggered...
+    assert searcher.counter["complete"] == 1
+    # ... and still updates the last result.
+    assert trials[0].last_result[TRAINING_ITERATION] == 1
+
+
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+def test_search_alg_stalled(ray_start_4_cpus_2_gpus_extra, resource_manager_cls):
+    """Checks that runner and searcher state is maintained when stalled.
+
+    We use a concurrency limit of 1, meaning each trial is added one-by-one
+    from the searchers.
+
+    We then run three samples. During the second trial, we stall the searcher,
+    which means we don't suggest new trials after it finished.
+
+    In this case, the runner should still be considered "running". Once we unstall,
+    the experiment finishes regularly.
+
+    Legacy test: test_trial_runner_3.py::TrialRunnerTest::testSearchAlgStalled
+    """
+    experiment_spec = {
+        "run": MOCK_TRAINABLE_NAME,
+        "num_samples": 3,
+        "stop": {"training_iteration": 1},
+    }
+    experiments = [Experiment.from_json("test", experiment_spec)]
+    search_alg = _MockSuggestionAlgorithm(max_concurrent=1)
+    search_alg.add_configurations(experiments)
+    searcher = search_alg.searcher
+    runner = TestTuneController(
+        resource_manager_factory=lambda: resource_manager_cls(),
+        search_alg=search_alg,
+    )
+    runner.step()
+    trials = runner.get_trials()
+    while trials[0].status != Trial.TERMINATED:
+        runner.step()
+
+    # On next step, trials[1] is created
+    runner.step()
+
+    trials = runner.get_trials()
+
+    while trials[1].status != Trial.RUNNING:
+        runner.step()
+
+    assert trials[1].status == Trial.RUNNING
+    assert len(searcher.live_trials) == 1
+
+    # Stall: We don't suggest new algorithms
+    searcher.stall = True
+
+    while trials[1].status != Trial.TERMINATED:
+        runner.step()
+
+    assert trials[1].status == Trial.TERMINATED
+    assert len(searcher.live_trials) == 0
+
+    assert all(trial.is_finished() for trial in trials)
+    assert not search_alg.is_finished()
+    assert not runner.is_finished()
+
+    # Unstall
+    searcher.stall = False
+
+    # Create trials[2]
+    runner.step()
+
+    trials = runner.get_trials()
+
+    while trials[2].status != Trial.RUNNING:
+        runner.step()
+
+    assert trials[2].status == Trial.RUNNING
+    assert len(searcher.live_trials) == 1
+
+    while trials[2].status != Trial.TERMINATED:
+        runner.step()
+
+    assert len(searcher.live_trials) == 0
+    assert search_alg.is_finished()
+    assert runner.is_finished()
+
+
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+def test_search_alg_finishes(ray_start_4_cpus_2_gpus_extra, resource_manager_cls):
+    """Empty SearchAlg changing state in `next_trials` does not crash.
+
+    The search algorithm changes to ``finished`` mid-run. This should not
+    affect processing of the experiment.
+
+    Legacy test: test_trial_runner_3.py::TrialRunnerTest::testSearchAlgFinishes
+    """
+    os.environ["TUNE_MAX_PENDING_TRIALS_PG"] = "1"
+
+    class FinishFastAlg(_MockSuggestionAlgorithm):
+        _index = 0
+
+        def next_trial(self):
+            spec = self._experiment.spec
+            trial = None
+            if self._index < spec["num_samples"]:
+                trial = Trial(
+                    spec.get("run"),
+                    stopping_criterion=spec.get("stop"),
+                    storage=spec.get("storage"),
+                )
+            self._index += 1
+
+            if self._index > 4:
+                self.set_finished()
+
+            return trial
+
+        def suggest(self, trial_id):
+            return {}
+
+    experiment_spec = {
+        "run": MOCK_TRAINABLE_NAME,
+        "num_samples": 2,
+        "stop": {"training_iteration": 1},
+    }
+    searcher = FinishFastAlg()
+    experiments = [Experiment.from_json("test", experiment_spec)]
+    searcher.add_configurations(experiments)
+
+    runner = TestTuneController(
+        resource_manager_factory=lambda: resource_manager_cls(),
+        search_alg=searcher,
+    )
+
+    assert not runner.is_finished()
+
+    while len(runner.get_trials()) < 2:
+        runner.step()  # Launch 2 runs
+
+    assert not searcher.is_finished()
+    assert not runner.is_finished()
+
+    searcher_finished_before = False
+    while not runner.is_finished():
+        runner.step()
+        searcher_finished_before = searcher.is_finished()
+
+    # searcher_finished_before will be True if the searcher was finished before
+    # the controller.
+    assert searcher_finished_before
+
+
+# Todo (krfricke): Fix in next batch
+@pytest.mark.skip("This test is currently flaky as it can fail due to timing issues.")
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+def test_searcher_save_restore(ray_start_8_cpus, resource_manager_cls, tmpdir):
+    """Searchers state should be saved and restored in the experiment checkpoint.
+
+    Legacy test: test_trial_runner_3.py::TrialRunnerTest::testSearcherSaveRestore
+    """
+
+    def create_searcher():
+        class TestSuggestion(Searcher):
+            def __init__(self, index):
+                self.index = index
+                self.returned_result = []
+                super().__init__(metric="episode_reward_mean", mode="max")
+
+            def suggest(self, trial_id):
+                self.index += 1
+                return {"test_variable": self.index}
+
+            def on_trial_complete(self, trial_id, result=None, **kwargs):
+                self.returned_result.append(result)
+
+            def save(self, checkpoint_path):
+                with open(checkpoint_path, "wb") as f:
+                    pickle.dump(self.__dict__, f)
+
+            def restore(self, checkpoint_path):
+                with open(checkpoint_path, "rb") as f:
+                    self.__dict__.update(pickle.load(f))
+
+        searcher = TestSuggestion(0)
+        searcher = ConcurrencyLimiter(searcher, max_concurrent=2)
+        searcher = Repeater(searcher, repeat=3, set_index=False)
+        search_alg = SearchGenerator(searcher)
+        experiment_spec = {
+            "run": MOCK_TRAINABLE_NAME,
+            "num_samples": 20,
+            "config": {"sleep": 10},
+            "stop": {"training_iteration": 2},
+            "resources_per_trial": PlacementGroupFactory([{"CPU": 1}]),
+        }
+        experiments = [Experiment.from_json("test", experiment_spec)]
+        search_alg.add_configurations(experiments)
+        return search_alg
+
+    searcher = create_searcher()
+
+    runner = TestTuneController(
+        resource_manager_factory=lambda: resource_manager_cls(),
+        search_alg=searcher,
+        checkpoint_period=-1,
+        experiment_path=str(tmpdir),
+    )
+
+    while len(runner.get_trials()) < 6:
+        runner.step()
+
+    assert len(runner.get_trials()) == 6, [t.config for t in runner.get_trials()]
+    runner.checkpoint()
+    trials = runner.get_trials()
+    [runner._schedule_trial_stop(t) for t in trials if t.status is not Trial.ERROR]
+
+    runner.cleanup()
+
+    del runner
+
+    searcher = create_searcher()
+
+    runner2 = TestTuneController(
+        resource_manager_factory=lambda: resource_manager_cls(),
+        search_alg=searcher,
+        experiment_path=str(tmpdir),
+        resume="LOCAL",
+    )
+
+    assert len(runner2.get_trials()) == 6, [t.config for t in runner2.get_trials()]
+
+    def trial_statuses():
+        return [t.status for t in runner2.get_trials()]
+
+    def num_running_trials():
+        return sum(t.status == Trial.RUNNING for t in runner2.get_trials())
+
+    while num_running_trials() < 6:
+        runner2.step()
+
+    assert len(set(trial_statuses())) == 1
+    assert Trial.RUNNING in trial_statuses()
+
+    for i in range(20):
+        runner2.step()
+        assert 1 <= num_running_trials() <= 6
+
+    evaluated = [t.evaluated_params["test_variable"] for t in runner2.get_trials()]
+    count = Counter(evaluated)
+    assert all(v <= 3 for v in count.values())
+
+
+class TestGetMaxPendingTrials:
+    """Tests for _get_max_pending_trials with custom searchers."""
+
+    def setup_method(self):
+        self._orig = os.environ.pop("TUNE_MAX_PENDING_TRIALS_PG", None)
+
+    def teardown_method(self):
+        if self._orig is not None:
+            os.environ["TUNE_MAX_PENDING_TRIALS_PG"] = self._orig
+        else:
+            os.environ.pop("TUNE_MAX_PENDING_TRIALS_PG", None)
+
+    def test_env_var_override(self):
+        os.environ["TUNE_MAX_PENDING_TRIALS_PG"] = "42"
+        sg = SearchGenerator(_MockSearcher())
+        assert _get_max_pending_trials(sg) == 42
+
+    def test_search_generator_without_concurrency_limiter(self):
+        sg = SearchGenerator(_MockSearcher())
+        assert _get_max_pending_trials(sg) == 1
+
+    def test_search_generator_with_concurrency_limiter(self):
+        limited = ConcurrencyLimiter(_MockSearcher(), max_concurrent=8)
+        sg = SearchGenerator(limited)
+        assert _get_max_pending_trials(sg) == 8
+
+    def test_search_generator_with_nested_concurrency_limiter(self):
+        limited = ConcurrencyLimiter(_MockSearcher(), max_concurrent=8)
+        repeater = Repeater(limited, repeat=3, set_index=False)
+        sg = SearchGenerator(repeater)
+        assert _get_max_pending_trials(sg) == 8
+
+    @pytest.mark.parametrize("max_concurrent", [1, 4, 16])
+    def test_various_concurrency_values(self, max_concurrent):
+        limited = ConcurrencyLimiter(_MockSearcher(), max_concurrent=max_concurrent)
+        sg = SearchGenerator(limited)
+        assert _get_max_pending_trials(sg) == max_concurrent
+
+    def test_mock_suggestion_algorithm_with_concurrency(self):
+        mock_alg = _MockSuggestionAlgorithm(max_concurrent=5)
+        assert _get_max_pending_trials(mock_alg) == 5
+
+    def test_mock_suggestion_algorithm_without_concurrency(self):
+        mock_alg = _MockSuggestionAlgorithm()
+        assert _get_max_pending_trials(mock_alg) == 1
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/tune/tests/execution/utils.py

@@ -0,0 +1,147 @@
+import os
+import uuid
+from typing import Any, Callable, Dict, Optional, Tuple, Type, Union
+
+import ray
+from ray.air.execution import FixedResourceManager
+from ray.air.execution._internal import RayActorManager
+from ray.air.execution._internal.tracked_actor import TrackedActor
+from ray.air.execution.resources import ResourceManager, ResourceRequest
+from ray.train.tests.util import mock_storage_context
+from ray.tune.execution.tune_controller import TuneController
+from ray.tune.experiment import Trial
+from ray.tune.utils.resource_updater import _ResourceUpdater
+
+
+class NoopClassCache:
+    def get(self, trainable_name: str):
+        return trainable_name
+
+
+class BudgetResourceManager(FixedResourceManager):
+    def __init__(self, total_resources: Dict[str, float]):
+        self._allow_strict_pack = True
+        self._total_resources = total_resources
+        self._requested_resources = []
+        self._used_resources = []
+
+
+class NoopActorManager(RayActorManager):
+    def __init__(self, resource_manager: ResourceManager):
+        super().__init__(resource_manager=resource_manager)
+
+        self.added_actors = []
+        self.removed_actors = []
+        self.scheduled_futures = []
+
+    def add_actor(
+        self,
+        cls: Union[Type, ray.actor.ActorClass],
+        kwargs: Dict[str, Any],
+        resource_request: ResourceRequest,
+        *,
+        on_start: Optional[Callable[[TrackedActor], None]] = None,
+        on_stop: Optional[Callable[[TrackedActor], None]] = None,
+        on_error: Optional[Callable[[TrackedActor, Exception], None]] = None,
+    ) -> TrackedActor:
+        fake_actor_ref = uuid.uuid4().int
+        tracked_actor = TrackedActor(
+            fake_actor_ref, on_start=on_start, on_stop=on_stop, on_error=on_error
+        )
+        self._live_actors_to_ray_actors_resources[tracked_actor] = (fake_actor_ref,)
+        self.added_actors.append((tracked_actor, cls, kwargs))
+        return tracked_actor
+
+    def remove_actor(
+        self,
+        tracked_actor: TrackedActor,
+        kill: bool = False,
+        stop_future: Optional[ray.ObjectRef] = None,
+    ) -> None:
+        self.removed_actors.append(tracked_actor)
+
+    def schedule_actor_task(
+        self,
+        tracked_actor: TrackedActor,
+        method_name: str,
+        args: Optional[Tuple] = None,
+        kwargs: Optional[Dict] = None,
+        on_result: Optional[Callable[[TrackedActor, Any], None]] = None,
+        on_error: Optional[Callable[[TrackedActor, Exception], None]] = None,
+        _return_future: bool = False,
+    ) -> Optional[int]:
+        fake_ref = uuid.uuid4().int
+        self.scheduled_futures.append(
+            (fake_ref, tracked_actor, method_name, args, kwargs, on_result, on_error)
+        )
+        return fake_ref
+
+    @property
+    def num_actor_tasks(self):
+        return len(self.scheduled_futures)
+
+    def get_live_actors_resources(self):
+        return {}
+
+    def next(self, timeout: Optional[Union[int, float]] = None) -> None:
+        pass
+
+    def set_num_pending(self, num_pending: int):
+        self._pending_actors_to_attrs = {i: None for i in range(num_pending)}
+
+
+class _FakeResourceUpdater(_ResourceUpdater):
+    def __init__(self, resource_manager: BudgetResourceManager):
+        self._resource_manager = resource_manager
+
+    def get_num_cpus(self):
+        return self._resource_manager._total_resources.get("CPU", 0)
+
+    def get_num_gpus(self) -> int:
+        return self._resource_manager._total_resources.get("GPU", 0)
+
+    def update_avail_resources(self, *args, **kwargs):
+        pass
+
+
+class TestingTrial(Trial):
+    def __init__(self, *args, **kwargs):
+        kwargs.setdefault("storage", mock_storage_context())
+        super().__init__(*args, **kwargs)
+
+    def get_trainable_cls(self):
+        return self.trainable_name
+
+    def create_placement_group_factory(self):
+        self.placement_group_factory = self._default_placement_group_factory
+
+    def set_ray_actor(self, ray_actor):
+        pass
+
+
+def create_execution_test_objects(
+    max_pending_trials: int = 8,
+    resources: Optional[Dict[str, float]] = None,
+    reuse_actors: bool = True,
+    tune_controller_cls: Type[TuneController] = TuneController,
+    **kwargs,
+):
+    os.environ["TUNE_MAX_PENDING_TRIALS_PG"] = str(max_pending_trials)
+
+    resources = resources or {"CPU": 4}
+
+    storage = kwargs.pop("storage", mock_storage_context())
+
+    tune_controller = tune_controller_cls(
+        reuse_actors=reuse_actors,
+        storage=storage,
+        **kwargs,
+    )
+    resource_manager = BudgetResourceManager(total_resources=resources)
+    resource_updater = _FakeResourceUpdater(resource_manager)
+    actor_manger = NoopActorManager(resource_manager)
+    tune_controller._actor_manager = actor_manger
+    tune_controller._class_cache = NoopClassCache()
+    tune_controller._resource_updater = resource_updater
+
+    return tune_controller, actor_manger, resource_manager

python/ray/tune/tests/output/conftest.py

@@ -0,0 +1,3 @@
+# Trigger pytest hook to automatically zip test cluster logs to archive dir on failure
+from ray.tests.conftest import propagate_logs  # noqa
+from ray.tests.conftest import pytest_runtest_makereport  # noqa

python/ray/tune/tests/output/test_output.py

@@ -0,0 +1,351 @@
+import argparse
+import sys
+from unittest import mock
+
+import pytest
+from freezegun import freeze_time
+
+from ray import tune
+from ray.air.constants import TRAINING_ITERATION
+from ray.tune.experiment.trial import Trial
+from ray.tune.experimental.output import (
+    AirVerbosity,
+    TrainReporter,
+    TuneTerminalReporter,
+    _best_trial_str,
+    _current_best_trial,
+    _get_dict_as_table_data,
+    _get_time_str,
+    _get_trial_info,
+    _get_trial_table_data,
+    _get_trials_by_state,
+    _infer_params,
+    _infer_user_metrics,
+    _max_len,
+)
+from ray.tune.utils.mock_trainable import MOCK_TRAINABLE_NAME
+
+LAST_RESULT = {
+    "custom_metrics": {},
+    "episode_media": {},
+    "info": {
+        "learner": {
+            "default_policy": {
+                "allreduce_latency": 0.0,
+                "grad_gnorm": 40.0,
+                "cur_lr": 0.001,
+                "total_loss": 93.35336303710938,
+                "policy_loss": -18.39633560180664,
+                "entropy": 0.5613694190979004,
+                "entropy_coeff": 0.01,
+                "var_gnorm": 23.452943801879883,
+                "vf_loss": 223.5106201171875,
+                "vf_explained_var": -0.0017577409744262695,
+                "mean_IS": 0.9987365007400513,
+                "var_IS": 0.0007558994111604989,
+            },
+        }
+    },
+    "sampler_results": {
+        "episode_reward_max": 500.0,
+        "episode_reward_min": 54.0,
+        "episode_reward_mean": 214.45,
+    },
+    "episode_reward_max": 500.0,
+    "episode_reward_min": 54.0,
+    "episode_reward_mean": 214.45,
+    "episode_len_mean": 214.45,
+    "episodes_this_iter": 66,
+    "timesteps_total": 33000,
+}
+
+
+@freeze_time("Mar 27th, 2023", auto_tick_seconds=15)
+def test_get_time_str():
+    base = 1679875200  # 2023-03-27 00:00:00
+
+    assert _get_time_str(base, base) == ("2023-03-27 00:00:00", "0s")
+    assert _get_time_str(base, base + 15) == ("2023-03-27 00:00:15", "15s")
+    assert _get_time_str(base, base + 60) == ("2023-03-27 00:01:00", "1min 0s")
+    assert _get_time_str(base, base + 65) == ("2023-03-27 00:01:05", "1min 5s")
+    assert _get_time_str(base, base + 3600) == (
+        "2023-03-27 01:00:00",
+        "1hr 0min 0s",
+    )
+    assert _get_time_str(base, base + 3605) == (
+        "2023-03-27 01:00:05",
+        "1hr 0min 5s",
+    )
+    assert _get_time_str(base, base + 3660) == (
+        "2023-03-27 01:01:00",
+        "1hr 1min 0s",
+    )
+    assert _get_time_str(base, base + 86400) == (
+        "2023-03-28 00:00:00",
+        "1d 0hr 0min 0s",
+    )
+
+
+def test_get_trials_by_state():
+    t1 = Trial(MOCK_TRAINABLE_NAME, stub=True)
+    t1.set_status(Trial.RUNNING)
+    t2 = Trial(MOCK_TRAINABLE_NAME, stub=True)
+    t2.set_status(Trial.PENDING)
+    trials = [t1, t2]
+    assert _get_trials_by_state(trials) == {"RUNNING": [t1], "PENDING": [t2]}
+
+
+def test_infer_user_metrics():
+    t = Trial(MOCK_TRAINABLE_NAME, stub=True)
+    t.run_metadata.last_result = LAST_RESULT
+    result = [
+        "episode_reward_max",
+        "episode_reward_min",
+        "episode_len_mean",
+        "episodes_this_iter",
+    ]
+    assert _infer_user_metrics([t]) == result
+
+
+def test_max_len():
+    assert _max_len("long_metrics_name", max_len=5) == "...me"
+    assert _max_len("long_metrics_name", max_len=10) == "...cs_name"
+    assert _max_len("long_metrics_name", max_len=9, wrap=True) == "long_metr\nics_name"
+    assert _max_len("long_metrics_name", max_len=8, wrap=True) == "..._metr\nics_name"
+
+
+def test_current_best_trial():
+    t1 = Trial(MOCK_TRAINABLE_NAME, stub=True)
+    t2 = Trial(MOCK_TRAINABLE_NAME, stub=True)
+    t1.run_metadata.last_result = {"metric": 2}
+    t2.run_metadata.last_result = {"metric": 1}
+    assert _current_best_trial([t1, t2], metric="metric", mode="min") == (t2, "metric")
+
+
+def test_best_trial_str():
+    t = Trial(MOCK_TRAINABLE_NAME, stub=True)
+    t.trial_id = "18ae7_00005"
+    t.run_metadata.last_result = {
+        "loss": 0.5918508041056858,
+        "config": {"train_loop_config": {"lr": 0.059253447253394785}},
+    }
+    assert (
+        _best_trial_str(t, "loss")
+        == "Current best trial: 18ae7_00005 with loss=0.5918508041056858"
+        " and params={'train_loop_config': {'lr': 0.059253447253394785}}"
+    )
+
+
+def test_get_trial_info():
+    t = Trial(MOCK_TRAINABLE_NAME, stub=True)
+    t.trial_id = "af42b609"
+    t.set_status(Trial.RUNNING)
+    t.run_metadata.last_result = LAST_RESULT
+    assert _get_trial_info(
+        t,
+        param_keys=[],
+        metric_keys=[
+            "episode_reward_mean",
+            "episode_reward_max",
+            "episode_reward_min",
+            "episode_len_mean",
+            "episodes_this_iter",
+        ],
+    ) == ["mock_trainable_af42b609", "RUNNING", 214.45, 500.0, 54.0, 214.45, 66]
+
+
+def test_get_trial_table_data_less_than_20():
+    trials = []
+    for i in range(20):
+        t = Trial(MOCK_TRAINABLE_NAME, stub=True)
+        t.trial_id = str(i)
+        t.set_status(Trial.RUNNING)
+        t.run_metadata.last_result = {"episode_reward_mean": 100 + i}
+        t.config = {"param": i}
+        trials.append(t)
+    table_data = _get_trial_table_data(trials, ["param"], ["episode_reward_mean"])
+    header = table_data.header
+    assert header == ["Trial name", "status", "param", "reward"]
+    table_data = table_data.data
+    assert len(table_data) == 1  # only the running category
+    assert len(table_data[0].trial_infos) == 20
+    assert not table_data[0].more_info
+
+
+def test_get_trial_table_data_more_than_20():
+    trials = []
+    # total of 30 trials.
+    for status in [Trial.RUNNING, Trial.TERMINATED, Trial.PENDING]:
+        for i in range(10):
+            t = Trial(MOCK_TRAINABLE_NAME, stub=True)
+            t.trial_id = str(i)
+            t.set_status(status)
+            t.run_metadata.last_result = {"episode_reward_mean": 100 + i}
+            t.config = {"param": i}
+            trials.append(t)
+    table_data = _get_trial_table_data(trials, ["param"], ["episode_reward_mean"])
+    header = table_data.header
+    assert header == ["Trial name", "status", "param", "reward"]
+    table_data = table_data.data
+    assert len(table_data) == 3  # only the running category
+    for i in range(3):
+        assert len(table_data[i].trial_infos) == 5
+    assert table_data[0].more_info == "5 more RUNNING"
+    assert table_data[1].more_info == "5 more TERMINATED"
+    assert table_data[2].more_info == "5 more PENDING"
+
+
+def test_infer_params():
+    assert _infer_params({}) == []
+    assert _infer_params({"some": "val"}) == []
+    assert _infer_params({"some": "val", "param": tune.uniform(0, 1)}) == ["param"]
+    assert _infer_params({"some": "val", "param": tune.grid_search([0, 1])}) == [
+        "param"
+    ]
+    assert sorted(
+        _infer_params(
+            {
+                "some": "val",
+                "param": tune.grid_search([0, 1]),
+                "other": tune.choice([0, 1]),
+            }
+        )
+    ) == ["other", "param"]
+
+
+def test_result_table_no_divison():
+    data = _get_dict_as_table_data(
+        {
+            "b": 6,
+            "a": 8,
+            "x": 19.123123123,
+            "c": 5,
+            "ignore": 9,
+            "nested_ignore": {"value": 5},
+            "y": 20,
+            "z": {"m": 4, "n": {"o": "p"}},
+        },
+        exclude={"ignore", "nested_ignore"},
+    )
+
+    assert data == [
+        ["a", 8],
+        ["b", 6],
+        ["c", 5],
+        ["x", "19.12312"],
+        ["y", 20],
+        ["z/m", 4],
+        ["z/n/o", "p"],
+    ]
+
+
+def test_result_table_divison():
+    data = _get_dict_as_table_data(
+        {
+            "b": 6,
+            "a": 8,
+            "x": 19.123123123,
+            "c": 5,
+            "ignore": 9,
+            "nested_ignore": {"value": 5},
+            "y": 20,
+            "z": {"m": 4, "n": {"o": "p"}},
+        },
+        exclude={"ignore", "nested_ignore"},
+        upper_keys={"x", "y", "z", "z/m", "z/n/o"},
+    )
+
+    assert data == [
+        ["x", "19.12312"],
+        ["y", 20],
+        ["z/m", 4],
+        ["z/n/o", "p"],
+        ["a", 8],
+        ["b", 6],
+        ["c", 5],
+    ]
+
+
+def test_result_include():
+    data = _get_dict_as_table_data(
+        {
+            "b": 6,
+            "a": 8,
+            "x": 19.123123123,
+            "c": 5,
+            "ignore": 9,
+            "nested_ignore": {"value": 5},
+            "y": 20,
+            "z": {"m": 4, "n": {"o": "p"}},
+        },
+        include={"y", "z"},
+        exclude={"z/n/o"},
+    )
+
+    assert data == [
+        ["y", 20],
+        ["z/m", 4],
+    ]
+
+
+def test_config_argparse():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--bool-val", action="store_true", default=True)
+    parser.add_argument("--foo", default="bar")
+    args = parser.parse_args([])
+
+    data = _get_dict_as_table_data({"parsed_args": args})
+    assert data == [
+        ["parsed_args/bool_val", True],
+        ["parsed_args/foo", "bar"],
+    ]
+
+
+@pytest.mark.parametrize("progress_reporter_cls", [TrainReporter, TuneTerminalReporter])
+def test_heartbeat_reset(progress_reporter_cls):
+    """Test heartbeat functionality in train and tune.
+
+    Tune prints a table every `heartbeat_freq` seconds.
+    Train prints a heartbeat every `heartbeat_freq` seconds, but a result
+    also resets the counter.
+    """
+    # Train heartbeats are only reporter in VERBOSE
+    reporter = progress_reporter_cls(verbosity=AirVerbosity.VERBOSE)
+    reporter._print_heartbeat = mock.MagicMock()
+
+    with freeze_time() as frozen:
+        reporter.print_heartbeat([])
+        assert reporter._print_heartbeat.call_count == 1
+
+        # Tick until heartbeat freq. Next call to print_heartbeat should trigger
+        frozen.tick(reporter._heartbeat_freq)
+        reporter.print_heartbeat([])
+        assert reporter._print_heartbeat.call_count == 2
+
+        # Not quite there, yet. This should not trigger a heartbeat.
+        frozen.tick(reporter._heartbeat_freq // 2)
+        reporter.print_heartbeat([])
+        assert reporter._print_heartbeat.call_count == 2
+
+        # Let's report a result. This will reset the heartbeat timer
+        reporter.on_trial_result(
+            0, [], Trial(MOCK_TRAINABLE_NAME, stub=True), {TRAINING_ITERATION: 1}
+        )
+
+        # Progress another half heartbeat. In Tune this triggers a heartbeat,
+        # but in train the heartbeat is reset on trial result.
+        frozen.tick(reporter._heartbeat_freq // 2 + 1)
+        reporter.print_heartbeat([])
+
+        if progress_reporter_cls == TrainReporter:
+            # Thus, train shouldn't have reported
+            assert reporter._print_heartbeat.call_count == 2
+        elif progress_reporter_cls == TuneTerminalReporter:
+            # But Tune should have.
+            assert reporter._print_heartbeat.call_count == 3
+        else:
+            raise RuntimeError("Test faulty.")
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/tune/tests/test_actor_reuse.py

@@ -0,0 +1,526 @@
+import inspect
+import os
+import sys
+import tempfile
+import time
+from pathlib import Path
+from typing import Callable
+
+import pytest
+
+import ray
+from ray import logger, tune
+from ray.train.tests.util import create_dict_checkpoint, load_dict_checkpoint
+from ray.tune import CheckpointConfig, Trainable, register_trainable, run_experiments
+from ray.tune.error import TuneError
+from ray.tune.result_grid import ResultGrid
+from ray.tune.schedulers.trial_scheduler import FIFOScheduler, TrialScheduler
+from ray.tune.tune import _check_mixin
+
+
+@pytest.fixture
+def ray_start_1_cpu():
+    address_info = ray.init(num_cpus=1)
+    os.environ["TUNE_STATE_REFRESH_PERIOD"] = "0.1"
+    yield address_info
+    ray.shutdown()
+    os.environ.pop("TUNE_STATE_REFRESH_PERIOD", None)
+
+
+@pytest.fixture
+def ray_start_2_cpus():
+    address_info = ray.init(num_cpus=2)
+    yield address_info
+    ray.shutdown()
+
+
+@pytest.fixture
+def ray_start_4_cpus_extra():
+    address_info = ray.init(num_cpus=4, resources={"extra": 4})
+    yield address_info
+    ray.shutdown()
+
+
+class FrequentPausesScheduler(FIFOScheduler):
+    def on_trial_result(self, tune_controller, trial, result):
+        return TrialScheduler.PAUSE
+
+
+class MyResettableClass(Trainable):
+    def setup(self, config):
+        self.config = config
+        self.num_resets = 0
+        self.iter = 0
+        self.msg = config.get("message", None)
+        self.sleep = int(config.get("sleep", 0))
+        self.fail = config.get("fail", False)
+
+    def step(self):
+        self.iter += 1
+
+        if self.msg:
+            print("PRINT_STDOUT: {}".format(self.msg))
+            print("PRINT_STDERR: {}".format(self.msg), file=sys.stderr)
+            logger.info("LOG_STDERR: {}".format(self.msg))
+
+        if self.fail:
+            raise RuntimeError("Failing")
+
+        if self.sleep:
+            time.sleep(self.sleep)
+
+        return {
+            "id": self.config.get("id", -1),
+            "num_resets": self.num_resets,
+            "done": self.iter > 1,
+            "iter": self.iter,
+        }
+
+    def save_checkpoint(self, chkpt_dir):
+        return {"iter": self.iter}
+
+    def load_checkpoint(self, item):
+        self.iter = item["iter"]
+
+    def reset_config(self, new_config):
+        if "fake_reset_not_supported" in self.config:
+            return False
+        self.num_resets += 1
+        self.iter = 0
+        self.msg = new_config.get("message", None)
+        self.fail = new_config.get("fail", False)
+        return True
+
+    @classmethod
+    def default_resource_request(cls, config):
+        required_resources = config.get("required_resources", None)
+        if required_resources:
+            return required_resources
+        return None
+
+
+def train_fn(config):
+    # Determine whether or not we reset to a new trial
+    marker_dir = config.get("marker_dir")
+    num_resets = 0
+    marker = Path(marker_dir) / f"{os.getpid()}.txt"
+    if marker.exists():
+        num_resets = int(marker.read_text()) + 1
+
+    checkpoint = tune.get_checkpoint()
+    it = load_dict_checkpoint(checkpoint)["iter"] if checkpoint else 0
+
+    msg = config.get("message", None)
+    sleep = int(config.get("sleep", 0))
+    fail = config.get("fail", False)
+
+    while it < 2:
+        it += 1
+        if msg:
+            print("PRINT_STDOUT: {}".format(msg))
+            print("PRINT_STDERR: {}".format(msg), file=sys.stderr)
+            logger.info("LOG_STDERR: {}".format(msg))
+
+        if fail:
+            raise RuntimeError("Failing")
+
+        # Dump the current config
+        marker.write_text(str(num_resets))
+
+        if sleep:
+            time.sleep(sleep)
+
+        metrics = {
+            "id": config.get("id", 0),
+            "num_resets": num_resets,
+            "iter": it,
+            "done": it > 1,
+        }
+        if config.get("save_checkpoint", True):
+            with create_dict_checkpoint({"iter": it}) as checkpoint:
+                tune.report(metrics, checkpoint=checkpoint)
+        else:
+            tune.report(metrics, checkpoint=checkpoint)
+
+
+@pytest.fixture(params=["function", "class"])
+def trainable(request):
+    """Fixture that sets up a function/class trainable for testing.
+    Make sure this fixture comes BEFORE the ray.init fixture in the arguments
+    so that the env var is propagated to workers correctly."""
+    trainable_type = request.param
+    if trainable_type == "function":
+        yield train_fn
+    elif trainable_type == "class":
+        yield MyResettableClass
+    else:
+        raise NotImplementedError
+
+
+def _run_trials_with_frequent_pauses(trainable, reuse=False, **kwargs):
+    tempdir = tempfile.mkdtemp()
+    marker_dir = Path(tempdir)
+    analysis = tune.run(
+        trainable,
+        num_samples=1,
+        config={
+            "id": tune.grid_search([0, 1, 2, 3]),
+            "marker_dir": marker_dir,
+        },
+        reuse_actors=reuse,
+        scheduler=FrequentPausesScheduler(),
+        verbose=0,
+        **kwargs,
+    )
+    return analysis
+
+
+def test_trial_reuse_disabled(trainable, ray_start_1_cpu):
+    """Test that reuse=False disables actor re-use.
+
+    Setup: Pass `reuse_actors=False` to tune.run()
+
+    We assert the `num_resets` of each trainable class to be 0 (no reuse).
+    """
+    analysis = _run_trials_with_frequent_pauses(trainable, reuse=False)
+    trials = analysis.trials
+    assert [t.last_result["id"] for t in trials] == [0, 1, 2, 3]
+    assert [t.last_result["iter"] for t in trials] == [2, 2, 2, 2]
+    assert [t.last_result["num_resets"] for t in trials] == [0, 0, 0, 0]
+
+
+def test_trial_reuse_enabled(trainable, ray_start_1_cpu):
+    """Test that reuse=True enables actor re-use.
+
+    Setup: Pass `reuse_actors=True` to tune.run()
+
+    We assert the `num_resets` of each trainable class to be 4, 5, 6, and 7,
+    respectively:
+
+    - Each trial runs for 2 iterations
+    - Only one trial can run at a time
+    - After each iteration, trials are paused and actors cached for reuse
+    - Thus, the first trial finishes after 4 resets, the second after 5, etc.
+    """
+    analysis = _run_trials_with_frequent_pauses(trainable, reuse=True)
+    trials = analysis.trials
+    assert [t.last_result["id"] for t in trials] == [0, 1, 2, 3]
+    assert [t.last_result["iter"] for t in trials] == [2, 2, 2, 2]
+    assert [t.last_result["num_resets"] for t in trials] == [4, 5, 6, 7]
+
+
+def test_trial_reuse_with_failing(trainable, ray_start_1_cpu, tmp_path):
+    """Test that failing actors won't be reused.
+
+    - 1 trial can run at a time
+    - Some trials are failing
+    - We assert that trials after failing trials are scheduled on fresh actors
+        (num_resets = 0)
+    - We assert that trials after successful trials are schedule on reused actors
+        (num_reset = last_num_resets + 1)
+    """
+    fail = [False, True, False, False, True, True, False, False, False]
+    trials = tune.run(
+        trainable,
+        reuse_actors=True,
+        config={
+            "id": tune.grid_search(list(range(9))),
+            "fail": tune.sample_from(lambda config: fail[config["id"]]),
+            "marker_dir": tmp_path,
+        },
+        raise_on_failed_trial=False,
+    ).trials
+
+    assert [t.last_result.get("iter") for t in trials] == [
+        2,
+        None,
+        2,
+        2,
+        None,
+        None,
+        2,
+        2,
+        2,
+    ]
+    assert [t.last_result.get("num_resets") for t in trials] == [
+        0,
+        None,
+        0,
+        1,
+        None,
+        None,
+        0,
+        1,
+        2,
+    ]
+
+
+def test_reuse_enabled_error(ray_start_1_cpu):
+    """Test that a class without reset() enabled throws an error on actor reuse."""
+    with pytest.raises(TuneError):
+        run_experiments(
+            {
+                "foo": {
+                    "run": MyResettableClass,
+                    "max_failures": 1,
+                    "num_samples": 1,
+                    "config": {
+                        "id": tune.grid_search([0, 1, 2, 3]),
+                        "fake_reset_not_supported": True,
+                    },
+                }
+            },
+            reuse_actors=True,
+            scheduler=FrequentPausesScheduler(),
+        )
+
+
+def test_trial_reuse_log_to_file(trainable, ray_start_1_cpu, tmp_path):
+    """Check that log outputs from trainables are correctly stored with actor reuse.
+
+    We run two trials with actor reuse. When the actor is reused, we expect
+    the log output to be written to the log file of the new trial - i.e. we expect
+    that the old trial logfile is closed and a new one is open.
+    """
+    register_trainable("foo2", trainable)
+
+    messages = ["First", "Second"]
+    # Log to default files
+    [trial1, trial2] = tune.run(
+        "foo2",
+        config={
+            "id": tune.grid_search(list(range(2))),
+            "message": tune.sample_from(lambda config: messages[config["id"]]),
+            "marker_dir": tmp_path,
+        },
+        log_to_file=True,
+        scheduler=FrequentPausesScheduler(),
+        reuse_actors=True,
+    ).trials
+
+    def get_trial_logfiles(trial):
+        return (
+            os.path.join(trial.storage.trial_working_directory, "stdout"),
+            os.path.join(trial.storage.trial_working_directory, "stderr"),
+        )
+
+    # Check trial 1
+    assert trial1.last_result["num_resets"] == 2
+    [stdout, stderr] = get_trial_logfiles(trial1)
+    assert os.path.exists(stdout)
+    assert os.path.exists(stderr)
+
+    # We expect that only "First" output is found in the first trial output
+    with open(stdout, "rt") as fp:
+        content = fp.read()
+        assert "PRINT_STDOUT: First" in content
+        assert "PRINT_STDOUT: Second" not in content
+    with open(stderr, "rt") as fp:
+        content = fp.read()
+        assert "PRINT_STDERR: First" in content
+        assert "LOG_STDERR: First" in content
+        assert "PRINT_STDERR: Second" not in content
+        assert "LOG_STDERR: Second" not in content
+
+    # Check trial 2
+    assert trial2.last_result["num_resets"] == 3
+    [stdout, stderr] = get_trial_logfiles(trial2)
+    assert os.path.exists(stdout)
+    assert os.path.exists(stderr)
+
+    # We expect that only "Second" output is found in the first trial output
+    with open(stdout, "rt") as fp:
+        content = fp.read()
+        assert "PRINT_STDOUT: Second" in content
+        assert "PRINT_STDOUT: First" not in content
+    with open(stderr, "rt") as fp:
+        content = fp.read()
+        assert "PRINT_STDERR: Second" in content
+        assert "LOG_STDERR: Second" in content
+        assert "PRINT_STDERR: First" not in content
+        assert "LOG_STDERR: First" not in content
+
+
+def test_multi_trial_reuse(trainable, ray_start_4_cpus_extra, monkeypatch, tmp_path):
+    """Test that actors from multiple trials running in parallel will be reused.
+
+    - 2 trials can run at the same time
+    - Trial 3 will be scheduled after trial 1 succeeded, so will reuse actor
+    - Trial 4 will be scheduled after trial 2 succeeded, so will reuse actor
+    """
+    monkeypatch.setenv("TUNE_MAX_PENDING_TRIALS_PG", "2")
+
+    register_trainable("foo2", trainable)
+
+    messages = ["First", "Second", "Third", "Fourth"]
+    # We sleep here for one second so that the third actor
+    # does not finish training before the fourth can be scheduled.
+    # This helps ensure that both remote runners are re-used and
+    # not just one.
+    [trial1, trial2, trial3, trial4] = tune.run(
+        "foo2",
+        config={
+            "id": tune.grid_search(list(range(4))),
+            "message": tune.sample_from(lambda config: messages[config["id"]]),
+            "marker_dir": tmp_path,
+            "sleep": 2,
+        },
+        reuse_actors=True,
+        resources_per_trial={"cpu": 2},
+    ).trials
+
+    assert trial3.last_result["num_resets"] == 1
+    assert trial4.last_result["num_resets"] == 1
+
+
+def test_multi_trial_reuse_with_failing(
+    trainable, ray_start_4_cpus_extra, monkeypatch, tmp_path
+):
+    """Test that failing trial's actors are not reused.
+
+    - 2 trials can run at the same time
+    - Trial 1 succeeds, trial 2 fails
+    - Trial 3 will be scheduled after trial 2 failed, so won't reuse actor
+    - Trial 4 will be scheduled after trial 1 succeeded, so will reuse actor
+    """
+    monkeypatch.setenv("TUNE_MAX_PENDING_TRIALS_PG", "2")
+
+    register_trainable("foo2", trainable)
+
+    [trial1, trial2, trial3, trial4] = tune.run(
+        "foo2",
+        config={
+            "fail": tune.grid_search([False, True, False, False]),
+            "marker_dir": tmp_path,
+            "sleep": 2,
+        },
+        reuse_actors=True,
+        resources_per_trial={"cpu": 2},
+        raise_on_failed_trial=False,
+    ).trials
+
+    assert trial1.last_result["num_resets"] == 0
+    assert trial3.last_result["num_resets"] == 0
+    assert trial4.last_result["num_resets"] == 1
+
+
+def test_multi_trial_reuse_one_by_one(trainable, ray_start_4_cpus_extra, tmp_path):
+    """Test that we still reuse actors even if we run with concurrency = 1.
+
+    - Run 6 trials, but only 1 concurrent at the time
+    - This means there won't be any PENDING trials until the trial completed
+    - We still want to reuse actors
+
+    """
+    register_trainable("foo2", trainable)
+
+    trials = tune.run(
+        "foo2",
+        config={"id": -1, "marker_dir": tmp_path},
+        reuse_actors=True,
+        num_samples=6,
+        max_concurrent_trials=1,
+    ).trials
+
+    assert sorted([t.last_result["num_resets"] for t in trials]) == [0, 1, 2, 3, 4, 5]
+
+
+def test_multi_trial_reuse_heterogeneous(ray_start_4_cpus_extra):
+    """Test that actors with heterogeneous resource requirements are reused efficiently.
+
+    - Run 6 trials in total
+    - Only 1 trial can run at the same time
+    - Trials 1 and 6, 2 and 4, and 3 and 5, have the same resource request, respectively
+    - Assert that trials 4, 5, and 6 re-use their respective previous actors
+
+    """
+    # We need to set this to 6 so that all trials will be scheduled and actors will
+    # be correctly cached.
+    os.environ["TUNE_MAX_PENDING_TRIALS_PG"] = "6"
+
+    register_trainable("foo2", MyResettableClass)
+
+    trials = tune.run(
+        "foo2",
+        config={
+            # The extra resources are selected so that only any 1 placement group
+            # can be scheduled at the same time at all times (to force sequential
+            # execution of trials)
+            "required_resources": tune.grid_search(
+                [
+                    {"cpu": 4, "custom_resources": {"extra": 4}},
+                    {"cpu": 2, "custom_resources": {"extra": 4}},
+                    {"cpu": 1, "custom_resources": {"extra": 4}},
+                    {"cpu": 2, "custom_resources": {"extra": 4}},
+                    {"cpu": 1, "custom_resources": {"extra": 4}},
+                    {"cpu": 4, "custom_resources": {"extra": 4}},
+                ]
+            ),
+            "id": -1,
+        },
+        reuse_actors=True,
+    ).trials
+
+    # Actors may be re-used in a different order as the staged_trials set is unsorted
+    assert sorted([t.last_result["num_resets"] for t in trials]) == [0, 0, 0, 1, 1, 1]
+
+
+def test_detect_reuse_mixins():
+    class DummyMixin:
+        pass
+
+    def dummy_mixin(func: Callable):
+        func.__mixins__ = (DummyMixin,)
+        return func
+
+    def train_fn(config):
+        pass
+
+    assert not _check_mixin(train_fn)
+    assert _check_mixin(dummy_mixin(train_fn))
+
+    class MyTrainable(Trainable):
+        pass
+
+    assert not _check_mixin(MyTrainable)
+    assert _check_mixin(dummy_mixin(MyTrainable))
+
+
+def test_remote_trial_dir_with_reuse_actors(trainable, ray_start_2_cpus, tmp_path):
+    """Check that the trainable has its remote directory set to the right
+    location, when new trials get swapped in on actor reuse.
+    Each trial runs for 2 iterations, with checkpoint_frequency=1, so each
+    remote trial dir should have 2 checkpoints.
+    """
+    tmp_target = str(tmp_path / "upload_dir")
+    exp_name = "remote_trial_dir_update_on_actor_reuse"
+
+    def get_remote_trial_dir(trial_id: int):
+        return os.path.join(tmp_target, exp_name, str(trial_id))
+
+    analysis = _run_trials_with_frequent_pauses(
+        trainable,
+        reuse=True,
+        max_concurrent_trials=2,
+        name=exp_name,
+        storage_path=f"file://{tmp_target}",
+        trial_dirname_creator=lambda t: str(t.config.get("id")),
+        checkpoint_config=CheckpointConfig(
+            checkpoint_frequency=1 if inspect.isclass(trainable) else 0
+        ),
+    )
+    result_grid = ResultGrid(analysis)
+    assert not result_grid.errors
+
+    # Check that each remote trial dir has 2 checkpoints.
+    for result in result_grid:
+        trial_id = result.config["id"]
+        remote_dir = get_remote_trial_dir(trial_id)
+        num_checkpoints = len(
+            [file for file in os.listdir(remote_dir) if file.startswith("checkpoint_")]
+        )
+        assert num_checkpoints == 2
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/tune/tests/test_api_checkpoint_integration.py

@@ -0,0 +1,165 @@
+import os
+import sys
+import tempfile
+from pathlib import Path
+
+import pytest
+
+import ray
+from ray.air.constants import TRAINING_ITERATION
+from ray.air.execution import FixedResourceManager
+from ray.train import ScalingConfig
+from ray.train._internal.storage import StorageContext
+from ray.train.tests.util import mock_storage_context
+from ray.tune import CheckpointConfig, Trainable, register_trainable
+from ray.tune.execution.tune_controller import TuneController
+from ray.tune.experiment import Trial
+
+STORAGE = mock_storage_context()
+
+
+@pytest.fixture(scope="function")
+def ray_start_4_cpus_2_gpus_extra():
+    address_info = ray.init(num_cpus=4, num_gpus=2, resources={"a": 2})
+    yield address_info
+    ray.shutdown()
+
+
+# TODO: [V2] Delete the `data_parallel` variant once V1 is fully removed.
+@pytest.mark.parametrize("trainable_type", ["class", "function", "data_parallel"])
+@pytest.mark.parametrize("patch_iter", [False, True])
+def test_checkpoint_freq_dir_name(
+    ray_start_4_cpus_2_gpus_extra, trainable_type, patch_iter, tmp_path
+):
+    """Test that trial checkpoint IDs are correctly set across trainable types.
+
+    This includes a current workaround to set checkpoint IDs according to reported
+    metrics.
+    """
+
+    def num_checkpoints(trial):
+        return sum(
+            item.startswith("checkpoint_")
+            for item in os.listdir(trial.storage.trial_fs_path)
+        )
+
+    def last_checkpoint_dir(trial):
+        return max(
+            item
+            for item in os.listdir(trial.storage.trial_fs_path)
+            if item.startswith("checkpoint_")
+        )
+
+    checkpoint_config = None
+
+    if trainable_type == "class":
+
+        class MyTrainable(Trainable):
+            def step(self):
+                # `training_iteration` is increased after the report, so we
+                # +1 here.
+                return {"step": self.iteration + 1}
+
+            def save_checkpoint(self, checkpoint_dir):
+                return {"test": self.iteration}
+
+            def load_checkpoint(self, checkpoint_dir):
+                pass
+
+        register_trainable("test_checkpoint_freq", MyTrainable)
+        checkpoint_config = CheckpointConfig(checkpoint_frequency=3)
+
+    elif trainable_type in {"function", "data_parallel"}:
+
+        def train_fn(config):
+            for step in range(1, 10):
+                if step > 0 and step % 3 == 0:
+                    with tempfile.TemporaryDirectory() as checkpoint_dir:
+                        (Path(checkpoint_dir) / "data.ckpt").write_text(str(step))
+                        ray.tune.report(
+                            {"step": step},
+                            checkpoint=ray.tune.Checkpoint.from_directory(
+                                checkpoint_dir
+                            ),
+                        )
+                else:
+                    ray.tune.report({"step": step})
+
+        if trainable_type == "function":
+            register_trainable("test_checkpoint_freq", train_fn)
+        elif trainable_type == "data_parallel":
+            from ray.train.data_parallel_trainer import DataParallelTrainer
+
+            trainer = DataParallelTrainer(
+                train_loop_per_worker=train_fn,
+                scaling_config=ScalingConfig(num_workers=1),
+            )
+            register_trainable("test_checkpoint_freq", trainer.as_trainable())
+
+    else:
+        raise RuntimeError("Invalid trainable type")
+
+    if patch_iter:
+
+        class CustomStorageContext(StorageContext):
+            def _update_checkpoint_index(self, metrics):
+                # Todo: Support auto-fille metrics for function trainables
+                self.current_checkpoint_index = metrics.get(
+                    "step", self.current_checkpoint_index + 1
+                )
+
+        storage = mock_storage_context(
+            storage_context_cls=CustomStorageContext,
+            storage_path=tmp_path,
+        )
+    else:
+        storage = mock_storage_context(storage_path=tmp_path)
+
+    trial = Trial(
+        "test_checkpoint_freq",
+        checkpoint_config=checkpoint_config,
+        storage=storage,
+    )
+    runner = TuneController(
+        resource_manager_factory=lambda: FixedResourceManager(),
+        storage=STORAGE,
+        checkpoint_period=0,
+    )
+    runner.add_trial(trial)
+
+    while not trial.is_saving:
+        runner.step()
+    runner.step()
+    assert trial.last_result[TRAINING_ITERATION] == 3
+    assert num_checkpoints(trial) == 1
+
+    if patch_iter:
+        assert last_checkpoint_dir(trial) == "checkpoint_000003"
+    else:
+        assert last_checkpoint_dir(trial) == "checkpoint_000000"
+
+    while not trial.is_saving:
+        runner.step()
+    runner.step()
+    assert trial.last_result[TRAINING_ITERATION] == 6
+    assert num_checkpoints(trial) == 2
+
+    if patch_iter:
+        assert last_checkpoint_dir(trial) == "checkpoint_000006"
+    else:
+        assert last_checkpoint_dir(trial) == "checkpoint_000001"
+
+    while not trial.is_saving:
+        runner.step()
+    runner.step()
+    assert trial.last_result[TRAINING_ITERATION] == 9
+    assert num_checkpoints(trial) == 3
+
+    if patch_iter:
+        assert last_checkpoint_dir(trial) == "checkpoint_000009"
+    else:
+        assert last_checkpoint_dir(trial) == "checkpoint_000002"
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/tune/tests/test_api_migrations.py

@@ -0,0 +1,98 @@
+import functools
+import importlib
+import sys
+import warnings
+
+import pytest
+
+import ray.train
+import ray.tune
+from ray.train.constants import ENABLE_V2_MIGRATION_WARNINGS_ENV_VAR
+from ray.util.annotations import RayDeprecationWarning
+
+
+@pytest.fixture(autouse=True)
+def enable_v2(monkeypatch):
+    monkeypatch.setenv("RAY_TRAIN_V2_ENABLED", "1")
+    importlib.reload(ray.train)
+    yield
+
+
+@pytest.fixture(autouse=True)
+def enable_v2_migration_deprecation_messages(monkeypatch):
+    monkeypatch.setenv(ENABLE_V2_MIGRATION_WARNINGS_ENV_VAR, "1")
+    yield
+    monkeypatch.delenv(ENABLE_V2_MIGRATION_WARNINGS_ENV_VAR)
+
+
+@pytest.mark.parametrize("v2_enabled", [False, True])
+def test_trainable_fn_utils(tmp_path, monkeypatch, v2_enabled):
+    monkeypatch.setenv("RAY_TRAIN_V2_ENABLED", str(int(v2_enabled)))
+    importlib.reload(ray.train)
+
+    dummy_checkpoint_dir = tmp_path.joinpath("dummy")
+    dummy_checkpoint_dir.mkdir()
+
+    asserting_context = (
+        functools.partial(pytest.raises, DeprecationWarning)
+        if v2_enabled
+        else functools.partial(pytest.warns, RayDeprecationWarning)
+    )
+
+    def tune_fn(config):
+        with asserting_context(match="get_checkpoint"):
+            ray.train.get_checkpoint()
+
+        with warnings.catch_warnings():
+            ray.tune.get_checkpoint()
+
+        with asserting_context(match="get_context"):
+            ray.train.get_context()
+
+        with warnings.catch_warnings():
+            ray.tune.get_context()
+
+        with asserting_context(match="report"):
+            ray.train.report({"a": 1})
+
+        with warnings.catch_warnings():
+            ray.tune.report({"a": 1})
+
+        with pytest.warns(RayDeprecationWarning, match="update your imports"):
+            ray.tune.report(
+                {"a": 1},
+                checkpoint=ray.train.Checkpoint.from_directory(dummy_checkpoint_dir),
+            )
+
+        with warnings.catch_warnings():
+            ray.tune.report(
+                {"a": 1},
+                checkpoint=ray.tune.Checkpoint.from_directory(dummy_checkpoint_dir),
+            )
+
+    tuner = ray.tune.Tuner(
+        tune_fn, run_config=ray.tune.RunConfig(storage_path=tmp_path)
+    )
+    results = tuner.fit()
+    assert not results.errors
+
+
+def test_configs():
+    with pytest.warns(RayDeprecationWarning, match="update your imports"):
+        ray.tune.Tuner(lambda c: None, run_config=ray.train.RunConfig())
+
+    with pytest.warns(RayDeprecationWarning, match="update your imports"):
+        ray.tune.Tuner(
+            lambda c: None,
+            run_config=ray.tune.RunConfig(failure_config=ray.train.FailureConfig()),
+        )
+
+    with warnings.catch_warnings():
+        ray.tune.Tuner(
+            lambda c: None,
+            run_config=ray.tune.RunConfig(failure_config=ray.tune.FailureConfig()),
+        )
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/tune/tests/test_callbacks.py

@@ -0,0 +1,60 @@
+from typing import Dict, Optional
+
+import pytest
+
+from ray.tune.callback import Callback, CallbackList
+
+
+class StatefulCallback(Callback):
+    CKPT_FILE_TMPL = "test-callback-state-{}.json"
+
+    def __init__(self):
+        self.counter = 0
+
+    def on_trial_result(self, iteration, trials, trial, result, **info):
+        self.counter += 1
+
+    def get_state(self) -> Optional[Dict]:
+        return {"counter": self.counter}
+
+    def set_state(self, state: Dict):
+        self.counter = state["counter"]
+
+
+def test_callback_list_with_stateful_callback(tmp_path):
+    """Checks that a callback list saves and restores all callbacks contained
+    inside it."""
+
+    callbacks = CallbackList([Callback(), StatefulCallback()])
+    for i in range(3):
+        callbacks.on_trial_result(iteration=i, trials=None, trial=None, result=None)
+
+    callbacks.save_to_dir(str(tmp_path))
+
+    assert list(tmp_path.glob(CallbackList.CKPT_FILE_TMPL.format("*")))
+    assert callbacks.can_restore(str(tmp_path))
+
+    restored_callbacks = CallbackList([Callback(), StatefulCallback()])
+    restored_callbacks.restore_from_dir(str(tmp_path))
+
+    assert restored_callbacks._callbacks[1].counter == 3
+
+
+def test_callback_list_without_stateful_callback(tmp_path):
+    """If no callbacks within a CallbackList are stateful, then nothing
+    should be saved."""
+
+    callbacks = CallbackList([Callback(), Callback()])
+    callbacks.save_to_dir(str(tmp_path))
+
+    assert not list(tmp_path.glob(CallbackList.CKPT_FILE_TMPL.format("*")))
+    assert not callbacks.can_restore(str(tmp_path))
+
+    with pytest.raises(RuntimeError):
+        callbacks.restore_from_dir(str(tmp_path))
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/tune/tests/test_cluster.py

@@ -0,0 +1,347 @@
+import os
+import sys
+import time
+from pathlib import Path
+from unittest.mock import MagicMock
+
+import pytest
+
+import ray
+from ray import tune
+from ray.cluster_utils import Cluster
+from ray.train._internal.storage import StorageContext
+from ray.tune import CheckpointConfig, register_trainable
+from ray.tune.error import TuneError
+from ray.tune.execution.tune_controller import TuneController
+from ray.tune.experiment import Trial
+from ray.tune.search import BasicVariantGenerator
+from ray.tune.utils.mock_trainable import MOCK_TRAINABLE_NAME, MyTrainableClass
+
+
+def _check_trial_running(trial):
+    return Path(trial.storage.trial_working_directory, "marker").exists()
+
+
+def _get_running_trials(runner):
+    return [t for t in runner.get_live_trials() if t.status == Trial.RUNNING]
+
+
+class SlowTrainable(MyTrainableClass):
+    def setup(self, config):
+        super().setup(config)
+        running_marker = Path(self._storage.trial_working_directory, "marker")
+        running_marker.touch()
+
+        self._sleep_time = config.get("sleep", 0)
+
+    def step(self):
+        time.sleep(self._sleep_time)
+        return super().step()
+
+
+def _start_new_cluster():
+    cluster = Cluster(
+        initialize_head=True,
+        connect=True,
+        head_node_args={
+            "num_cpus": 1,
+            "_system_config": {
+                "health_check_initial_delay_ms": 0,
+                "health_check_period_ms": 1000,
+                "health_check_failure_threshold": 10,
+            },
+        },
+    )
+    register_trainable(MOCK_TRAINABLE_NAME, SlowTrainable)
+    return cluster
+
+
+@pytest.fixture
+def start_connected_cluster():
+    # Start the Ray processes.
+    cluster = _start_new_cluster()
+    os.environ["TUNE_STATE_REFRESH_PERIOD"] = "0.1"
+    yield cluster
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+    cluster.shutdown()
+
+
+@pytest.fixture
+def start_connected_emptyhead_cluster():
+    """Starts head with no resources."""
+    cluster = Cluster(
+        initialize_head=True,
+        connect=True,
+        head_node_args={
+            "num_cpus": 0,
+            "_system_config": {
+                "health_check_initial_delay_ms": 0,
+                "health_check_period_ms": 1000,
+                "health_check_failure_threshold": 10,
+            },
+        },
+    )
+    os.environ["TUNE_STATE_REFRESH_PERIOD"] = "0.1"
+    yield cluster
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+    cluster.shutdown()
+
+
+@pytest.fixture
+def storage(tmp_path):
+    os.makedirs(tmp_path / "exp_name" / "trial_name", exist_ok=True)
+    yield StorageContext(
+        storage_path=str(tmp_path),
+        experiment_dir_name="exp_name",
+        trial_dir_name="trial_name",
+    )
+
+
+@pytest.fixture(autouse=True)
+def register_mock_trainable():
+    register_trainable(MOCK_TRAINABLE_NAME, SlowTrainable)
+    yield
+
+
+def test_counting_resources(start_connected_cluster, storage):
+    """Tests that Tune accounting is consistent with actual cluster."""
+
+    cluster = start_connected_cluster
+    nodes = []
+    assert ray.cluster_resources()["CPU"] == 1
+    runner = TuneController(search_alg=BasicVariantGenerator(), storage=storage)
+    kwargs = {
+        "stopping_criterion": {"training_iteration": 10},
+        "storage": storage,
+        "config": {"sleep": 1},
+    }
+
+    trials = [
+        Trial(MOCK_TRAINABLE_NAME, **kwargs),
+        Trial(MOCK_TRAINABLE_NAME, **kwargs),
+    ]
+    for t in trials:
+        runner.add_trial(t)
+
+    while not any(t.status == Trial.RUNNING for t in trials):
+        runner.step()
+
+    running_trials = _get_running_trials(runner)
+    assert len(running_trials) == 1
+    assert _check_trial_running(running_trials[0])
+    assert ray.available_resources().get("CPU", 0) == 0
+    nodes += [cluster.add_node(num_cpus=1)]
+    cluster.wait_for_nodes()
+    assert ray.cluster_resources()["CPU"] == 2
+    cluster.remove_node(nodes.pop())
+    cluster.wait_for_nodes()
+    assert ray.cluster_resources()["CPU"] == 1
+
+    while not any(t.status == Trial.RUNNING for t in trials):
+        runner.step()
+
+    # Only 1 trial can be running due to resource limitation.
+    assert sum(t.status == Trial.RUNNING for t in runner.get_trials()) == 1
+
+    for i in range(5):
+        nodes += [cluster.add_node(num_cpus=1)]
+    cluster.wait_for_nodes()
+    assert ray.cluster_resources()["CPU"] == 6
+
+    while any(t.status == Trial.PENDING for t in trials):
+        runner.step()
+    assert sum(t.status == Trial.RUNNING for t in runner.get_trials()) == 2, [
+        t.status for t in trials
+    ]
+
+
+def test_trial_processed_after_node_failure(start_connected_emptyhead_cluster, storage):
+    """Tests that Tune processes a trial as failed if its node died."""
+    cluster = start_connected_emptyhead_cluster
+    node = cluster.add_node(num_cpus=1)
+    cluster.wait_for_nodes()
+
+    runner = TuneController(search_alg=BasicVariantGenerator(), storage=storage)
+    mock_process_failure = MagicMock(side_effect=runner._process_trial_failure)
+    runner._process_trial_failure = mock_process_failure
+    # Disable recursion in magic mock when saving experiment state
+    runner.save_to_dir = lambda *args, **kwargs: None
+
+    runner.add_trial(Trial(MOCK_TRAINABLE_NAME, storage=storage))
+    trial = runner.get_trials()[0]
+
+    while trial.status != Trial.RUNNING:
+        runner.step()
+
+    assert not mock_process_failure.called
+
+    cluster.remove_node(node)
+    while not mock_process_failure.called:
+        runner.step()
+    assert mock_process_failure.called
+
+
+def test_remove_node_before_result(start_connected_emptyhead_cluster, storage):
+    """Tune continues when node is removed before trial returns."""
+    cluster = start_connected_emptyhead_cluster
+    node = cluster.add_node(num_cpus=1)
+    cluster.wait_for_nodes()
+
+    runner = TuneController(search_alg=BasicVariantGenerator(), storage=storage)
+    kwargs = {
+        "stopping_criterion": {"training_iteration": 3},
+        "checkpoint_config": CheckpointConfig(checkpoint_frequency=2),
+        "max_failures": 2,
+        "storage": storage,
+    }
+    trial = Trial(MOCK_TRAINABLE_NAME, **kwargs)
+    runner.add_trial(trial)
+
+    while trial.status != Trial.RUNNING:
+        runner.step()
+    running_trials = _get_running_trials(runner)
+    assert len(running_trials) == 1
+    assert _check_trial_running(running_trials[0])
+    assert not trial.has_reported_at_least_once
+    assert trial.status == Trial.RUNNING
+    cluster.remove_node(node)
+    cluster.add_node(num_cpus=1)
+    cluster.wait_for_nodes()
+    assert ray.cluster_resources()["CPU"] == 1
+
+    while not trial.last_result.get("training_iteration") == 1:
+        runner.step()
+    assert trial.last_result.get("training_iteration") == 1
+
+    # Process result: discover failure, recover, _train (from scratch)
+    while trial.status != Trial.TERMINATED:
+        runner.step()
+
+    assert trial.last_result.get("training_iteration") > 1
+
+    with pytest.raises(TuneError):
+        runner.step()
+
+
+def test_trial_requeue(start_connected_emptyhead_cluster, tmpdir, storage):
+    """Removing a node in full cluster causes Trial to be requeued."""
+    os.environ["TUNE_MAX_PENDING_TRIALS_PG"] = "1"
+
+    cluster = start_connected_emptyhead_cluster
+    node = cluster.add_node(num_cpus=1)
+    cluster.wait_for_nodes()
+
+    runner = TuneController(search_alg=BasicVariantGenerator(), storage=storage)
+    kwargs = {
+        "stopping_criterion": {"training_iteration": 5},
+        "checkpoint_config": CheckpointConfig(checkpoint_frequency=1),
+        "max_failures": 1,
+        "storage": storage,
+    }
+
+    trials = [
+        Trial(MOCK_TRAINABLE_NAME, **kwargs),
+        Trial(MOCK_TRAINABLE_NAME, **kwargs),
+    ]
+    for t in trials:
+        runner.add_trial(t)
+
+    while not any(t.status == Trial.RUNNING for t in trials):
+        runner.step()
+
+    runner.step()
+    runner.step()
+
+    running_trials = _get_running_trials(runner)
+    assert len(running_trials) == 1
+    assert _check_trial_running(running_trials[0])
+    cluster.remove_node(node)
+    cluster.wait_for_nodes()
+    time.sleep(0.1)  # Sleep so that next step() refreshes cluster resources
+    runner.step()  # Process result, dispatch save
+    runner.step()  # Process save (detect error), requeue trial
+    assert all(t.status == Trial.PENDING for t in trials)
+
+
+def test_migration_checkpoint_removal(
+    start_connected_emptyhead_cluster, tmpdir, storage
+):
+    """Test checks that trial restarts if checkpoint is lost w/ node fail."""
+    cluster = start_connected_emptyhead_cluster
+    node = cluster.add_node(num_cpus=1)
+    cluster.wait_for_nodes()
+
+    runner = TuneController(search_alg=BasicVariantGenerator(), storage=storage)
+    kwargs = {
+        "stopping_criterion": {"training_iteration": 4},
+        "checkpoint_config": CheckpointConfig(checkpoint_frequency=2),
+        "max_failures": 2,
+        "storage": storage,
+    }
+
+    # Test recovery of trial that has been checkpointed
+    t1 = Trial(MOCK_TRAINABLE_NAME, **kwargs)
+    runner.add_trial(t1)
+
+    # Start trial, process result (x2), process save
+    while not t1.has_checkpoint():
+        runner.step()
+
+    cluster.add_node(num_cpus=1)
+    cluster.remove_node(node)
+    cluster.wait_for_nodes()
+
+    while not runner.is_finished():
+        runner.step()
+    assert t1.status == Trial.TERMINATED
+
+
+def test_cluster_down_full(start_connected_cluster, tmpdir):
+    """Tests that run_experiment restoring works on cluster shutdown."""
+    cluster = start_connected_cluster
+
+    base_dict = dict(run=MOCK_TRAINABLE_NAME, stop=dict(training_iteration=3))
+
+    exp1_args = base_dict
+    exp2_args = dict(
+        base_dict.items(),
+        checkpoint_config=CheckpointConfig(checkpoint_frequency=1),
+    )
+    exp3_args = dict(base_dict.items(), config=dict(mock_error=True))
+    exp4_args = dict(
+        base_dict.items(),
+        config=dict(mock_error=True),
+        checkpoint_config=CheckpointConfig(checkpoint_frequency=1),
+    )
+
+    all_experiments = {
+        "exp1": exp1_args,
+        "exp2": exp2_args,
+        "exp3": exp3_args,
+        "exp4": exp4_args,
+    }
+
+    tune.run_experiments(all_experiments, raise_on_failed_trial=False)
+
+    ray.shutdown()
+    cluster.shutdown()
+    cluster = _start_new_cluster()
+
+    trials = tune.run_experiments(
+        all_experiments,
+        resume=True,
+        raise_on_failed_trial=False,
+    )
+
+    assert len(trials) == 4
+    assert all(t.status in [Trial.TERMINATED, Trial.ERROR] for t in trials)
+    ray.shutdown()
+    cluster.shutdown()
+
+
+if __name__ == "__main__":
+    import pytest
+
+    sys.exit(pytest.main(["-v", "--reruns", "3", __file__]))

python/ray/tune/tests/test_commands.py

@@ -0,0 +1,196 @@
+import os
+import random
+import subprocess
+import sys
+import time
+from unittest import mock
+
+import click
+import pytest
+
+import ray
+from ray import tune
+from ray.train.tests.util import create_dict_checkpoint
+from ray.tune.cli import commands
+from ray.tune.result import CONFIG_PREFIX
+from ray.tune.utils.mock_trainable import MyTrainableClass
+
+try:
+    from cStringIO import StringIO
+except ImportError:
+    from io import StringIO
+
+
+class Capturing:
+    def __enter__(self):
+        self._stdout = sys.stdout
+        sys.stdout = self._stringio = StringIO()
+        self.captured = []
+        return self
+
+    def __exit__(self, *args):
+        self.captured.extend(self._stringio.getvalue().splitlines())
+        del self._stringio  # free up some memory
+        sys.stdout = self._stdout
+
+
+@pytest.fixture
+def start_ray():
+    ray.init(log_to_driver=False)
+    yield
+    ray.shutdown()
+
+
+def test_time(start_ray, tmpdir, monkeypatch):
+    experiment_name = "test_time"
+    num_samples = 2
+
+    def train_fn(config):
+        for i in range(3):
+            with create_dict_checkpoint({"dummy": "data"}) as checkpoint:
+                ray.tune.report(
+                    {
+                        "epoch": i,
+                        "a": random.random(),
+                        "b/c": random.random(),
+                        "d": random.random(),
+                    },
+                    checkpoint=checkpoint,
+                )
+
+    tuner = tune.Tuner(
+        train_fn,
+        param_space={f"hp{i}": tune.uniform(0, 1) for i in range(100)},
+        tune_config=tune.TuneConfig(num_samples=num_samples),
+        run_config=ray.tune.RunConfig(name=experiment_name),
+    )
+    results = tuner.fit()
+    times = []
+    for _ in range(5):
+        start = time.time()
+        subprocess.check_call(["tune", "ls", results.experiment_path])
+        times += [time.time() - start]
+
+    print("Average CLI time: ", sum(times) / len(times))
+    assert sum(times) / len(times) < 5, "CLI is taking too long!"
+
+
+@mock.patch(
+    "ray.tune.cli.commands.print_format_output",
+    wraps=ray.tune.cli.commands.print_format_output,
+)
+def test_ls(mock_print_format_output, start_ray, tmpdir):
+    """This test captures output of list_trials."""
+    experiment_name = "test_ls"
+    experiment_path = os.path.join(str(tmpdir), experiment_name)
+    num_samples = 3
+    tune.run(
+        MyTrainableClass,
+        name=experiment_name,
+        stop={"training_iteration": 1},
+        num_samples=num_samples,
+        storage_path=str(tmpdir),
+    )
+
+    columns = ["episode_reward_mean", "training_iteration", "trial_id"]
+    limit = 2
+    commands.list_trials(experiment_path, info_keys=columns, limit=limit)
+
+    # The dataframe that is printed as a table is the first arg of the last
+    # call made to `ray.tune.cli.commands.print_format_output`.
+    mock_print_format_output.assert_called()
+    args, _ = mock_print_format_output.call_args_list[-1]
+    df = args[0]
+    assert sorted(df.columns.to_list()) == sorted(columns), df
+    assert len(df.index) == limit, df
+
+    commands.list_trials(
+        experiment_path,
+        sort=["trial_id"],
+        info_keys=("trial_id", "training_iteration"),
+        filter_op="training_iteration == 1",
+    )
+    args, _ = mock_print_format_output.call_args_list[-1]
+    df = args[0]
+    assert sorted(df.columns.to_list()) == sorted(["trial_id", "training_iteration"])
+    assert len(df.index) == num_samples
+
+    with pytest.raises(click.ClickException):
+        commands.list_trials(
+            experiment_path, sort=["trial_id"], info_keys=("training_iteration",)
+        )
+
+    with pytest.raises(click.ClickException):
+        commands.list_trials(experiment_path, info_keys=("asdf",))
+
+
+@mock.patch(
+    "ray.tune.cli.commands.print_format_output",
+    wraps=ray.tune.cli.commands.print_format_output,
+)
+def test_ls_with_cfg(mock_print_format_output, start_ray, tmpdir):
+    experiment_name = "test_ls_with_cfg"
+    experiment_path = os.path.join(str(tmpdir), experiment_name)
+    tune.run(
+        MyTrainableClass,
+        name=experiment_name,
+        stop={"training_iteration": 1},
+        config={"test_variable": tune.grid_search(list(range(5)))},
+        storage_path=str(tmpdir),
+    )
+
+    columns = [CONFIG_PREFIX + "/test_variable", "trial_id"]
+    limit = 4
+
+    commands.list_trials(experiment_path, info_keys=columns, limit=limit)
+
+    # The dataframe that is printed as a table is the first arg of the last
+    # call made to `ray.tune.cli.commands.print_format_output`.
+    mock_print_format_output.assert_called()
+    args, _ = mock_print_format_output.call_args_list[-1]
+    df = args[0]
+    assert sorted(df.columns.to_list()) == sorted(columns), df
+    assert len(df.index) == limit, df
+
+
+def test_lsx(start_ray, tmpdir):
+    """This test captures output of list_experiments."""
+    project_path = str(tmpdir)
+    num_experiments = 3
+    for i in range(num_experiments):
+        experiment_name = "test_lsx{}".format(i)
+        tune.run(
+            MyTrainableClass,
+            name=experiment_name,
+            stop={"training_iteration": 1},
+            num_samples=1,
+            storage_path=project_path,
+        )
+
+    limit = 2
+    with Capturing() as output:
+        commands.list_experiments(
+            project_path, info_keys=("total_trials",), limit=limit
+        )
+    lines = output.captured
+    assert "total_trials" in lines[1]
+    assert lines[1].count("|") == 2
+    assert len(lines) == 3 + limit + 1
+
+    with Capturing() as output:
+        commands.list_experiments(
+            project_path,
+            sort=["total_trials"],
+            info_keys=("total_trials",),
+            filter_op="total_trials == 1",
+        )
+    lines = output.captured
+    assert sum("1" in line for line in lines) >= num_experiments
+    assert len(lines) == 3 + num_experiments + 1
+
+
+if __name__ == "__main__":
+    # Make click happy in bazel.
+    os.environ["LC_ALL"] = "en_US.UTF-8"
+    os.environ["LANG"] = "en_US.UTF-8"
+    sys.exit(pytest.main([__file__]))

python/ray/tune/tests/test_convergence.py

@@ -0,0 +1,144 @@
+import math
+import sys
+import unittest
+
+import numpy as np
+import pytest
+
+import ray
+from ray import tune
+from ray.tune.search import ConcurrencyLimiter
+from ray.tune.stopper import ExperimentPlateauStopper
+
+
+def loss(config):
+    x = config.get("x")
+    tune.report({"loss": x**2})  # A simple function to optimize
+
+
+class ConvergenceTest(unittest.TestCase):
+    """Test convergence in gaussian process."""
+
+    @classmethod
+    def setUpClass(cls) -> None:
+        ray.init(num_cpus=1, num_gpus=0)
+
+    @classmethod
+    def tearDownClass(cls) -> None:
+        ray.shutdown()
+
+    def _testConvergence(self, searcher, top=3, patience=20):
+        # This is the space of parameters to explore
+        space = {"x": tune.uniform(0, 20)}
+
+        resources_per_trial = {"cpu": 1, "gpu": 0}
+
+        analysis = tune.run(
+            loss,
+            metric="loss",
+            mode="min",
+            stop=ExperimentPlateauStopper(metric="loss", top=top, patience=patience),
+            search_alg=searcher,
+            config=space,
+            num_samples=max(100, patience),  # Number of iterations
+            resources_per_trial=resources_per_trial,
+            raise_on_failed_trial=False,
+            fail_fast=True,
+            reuse_actors=True,
+            verbose=1,
+        )
+        print(
+            f"Num trials: {len(analysis.trials)}. "
+            f"Best result: {analysis.best_config['x']}"
+        )
+
+        return analysis
+
+    @unittest.skip("ax warm start tests currently failing (need to upgrade ax)")
+    def testConvergenceAx(self):
+        from ray.tune.search.ax import AxSearch
+
+        np.random.seed(0)
+
+        searcher = AxSearch()
+        analysis = self._testConvergence(searcher, patience=10)
+
+        assert math.isclose(analysis.best_config["x"], 0, abs_tol=1e-5)
+
+    def testConvergenceBayesOpt(self):
+        from ray.tune.search.bayesopt import BayesOptSearch
+
+        np.random.seed(0)
+
+        # Following bayesian optimization
+        searcher = BayesOptSearch(random_search_steps=10)
+        searcher.repeat_float_precision = 5
+        searcher = ConcurrencyLimiter(searcher, 1)
+
+        analysis = self._testConvergence(searcher, patience=100)
+
+        assert len(analysis.trials) < 50
+        assert math.isclose(analysis.best_config["x"], 0, abs_tol=1e-5)
+
+    @pytest.mark.skipif(
+        sys.version_info >= (3, 12), reason="HEBO doesn't support py312"
+    )
+    def testConvergenceHEBO(self):
+        from ray.tune.search.hebo import HEBOSearch
+
+        np.random.seed(0)
+        searcher = HEBOSearch()
+        analysis = self._testConvergence(searcher)
+
+        assert len(analysis.trials) < 100
+        assert math.isclose(analysis.best_config["x"], 0, abs_tol=1e-2)
+
+    def testConvergenceHyperopt(self):
+        from ray.tune.search.hyperopt import HyperOptSearch
+
+        np.random.seed(0)
+        searcher = HyperOptSearch(random_state_seed=1234)
+        analysis = self._testConvergence(searcher, patience=50, top=5)
+
+        assert math.isclose(analysis.best_config["x"], 0, abs_tol=1e-2)
+
+    def testConvergenceNevergrad(self):
+        import nevergrad as ng
+
+        from ray.tune.search.nevergrad import NevergradSearch
+
+        np.random.seed(0)
+        searcher = NevergradSearch(optimizer=ng.optimizers.PSO)
+        analysis = self._testConvergence(searcher, patience=50, top=5)
+
+        assert math.isclose(analysis.best_config["x"], 0, abs_tol=1e-3)
+
+    def testConvergenceOptuna(self):
+        from ray.tune.search.optuna import OptunaSearch
+
+        np.random.seed(1)
+        searcher = OptunaSearch(seed=1)
+        analysis = self._testConvergence(
+            searcher,
+            top=5,
+        )
+
+        # This assertion is much weaker than in the BO case, but TPE
+        # don't converge too close. It is still unlikely to get to this
+        # tolerance with random search (5 * 0.1 = 0.5% chance)
+        assert len(analysis.trials) < 100
+        assert math.isclose(analysis.best_config["x"], 0, abs_tol=1e-1)
+
+    def testConvergenceZoopt(self):
+        from ray.tune.search.zoopt import ZOOptSearch
+
+        np.random.seed(0)
+        searcher = ZOOptSearch(budget=100)
+        analysis = self._testConvergence(searcher)
+
+        assert len(analysis.trials) < 100
+        assert math.isclose(analysis.best_config["x"], 0, abs_tol=1e-3)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/tune/tests/test_dependency.py

@@ -0,0 +1,27 @@
+#!/usr/bin/env python
+
+import sys
+
+import pytest
+
+import ray
+import ray.tune
+from ray.tune import register_trainable, run_experiments
+
+
+def f(config):
+    ray.tune.report(dict(timesteps_total=1))
+
+
+def test_dependency():
+    ray.init(num_cpus=2)
+
+    register_trainable("my_class", f)
+    run_experiments({"test": {"run": "my_class", "stop": {"training_iteration": 1}}})
+    assert "ray.rllib" not in sys.modules, "RLlib should not be imported"
+    assert "mlflow" not in sys.modules, "MLflow should not be imported"
+    ray.shutdown()
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/tune/tests/test_env_callbacks.py

@@ -0,0 +1,108 @@
+import os
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from ray.tune.constants import RAY_TUNE_CALLBACKS_ENV_VAR
+from ray.tune.utils.callback import Callback, _initialize_env_callbacks
+
+
+class MockCallback(Callback):
+    pass
+
+
+@pytest.mark.parametrize(
+    "env_value,expected_callback_count",
+    [
+        ("my.module.Callback1", 1),
+        ("module1.Callback1,module2.Callback2", 2),
+        ("", 0),
+        ("   ", 0),
+        ("module.Callback1, ,module.Callback2", 2),
+    ],
+)
+@patch("importlib.import_module")
+def test_env_callbacks_loading(mock_import, env_value, expected_callback_count):
+    """Test loading execution callbacks from environment variable with various inputs."""
+    if env_value:
+        with patch.dict(os.environ, {RAY_TUNE_CALLBACKS_ENV_VAR: env_value}):
+            mock_module = MagicMock()
+            mock_module.Callback1 = MockCallback
+            mock_module.Callback2 = MockCallback
+            mock_import.return_value = mock_module
+
+            callbacks = _initialize_env_callbacks()
+
+            assert len(callbacks) == expected_callback_count
+            for callback in callbacks:
+                assert isinstance(callback, MockCallback)
+    else:
+        with patch.dict(
+            os.environ, {RAY_TUNE_CALLBACKS_ENV_VAR: env_value}, clear=True
+        ):
+            callbacks = _initialize_env_callbacks()
+            assert len(callbacks) == 0
+
+
+@pytest.mark.parametrize(
+    "env_value,original_error_type",
+    [
+        ("invalid_module", ValueError),
+        ("module.Class", TypeError),
+        ("module.NonExistentClass", AttributeError),
+    ],
+)
+@patch("importlib.import_module")
+def test_callback_loading_errors(mock_import, env_value, original_error_type):
+    """Test handling of various error conditions when loading callbacks."""
+    with patch.dict(os.environ, {RAY_TUNE_CALLBACKS_ENV_VAR: env_value}):
+        if "invalid_module" in env_value:
+            pass
+        elif "NonExistentClass" in env_value:
+            mock_module = MagicMock()
+            del mock_module.NonExistentClass
+            mock_import.return_value = mock_module
+        else:
+            mock_module = MagicMock()
+
+            class RegularClass:
+                pass
+
+            mock_module.Class = RegularClass
+            mock_import.return_value = mock_module
+
+        with pytest.raises(
+            ValueError, match=f"Failed to import callback from '{env_value}'"
+        ) as exc_info:
+            _initialize_env_callbacks()
+
+        assert isinstance(exc_info.value.__cause__, original_error_type)
+
+
+def test_import_error_handling():
+    """Test handling of import errors when loading callbacks."""
+    with patch.dict(
+        os.environ, {RAY_TUNE_CALLBACKS_ENV_VAR: "nonexistent.module.TestCallback"}
+    ):
+        with pytest.raises(
+            ValueError,
+            match="Failed to import callback from 'nonexistent.module.TestCallback'",
+        ) as exc_info:
+            _initialize_env_callbacks()
+
+        assert isinstance(exc_info.value.__cause__, ImportError)
+
+
+def test_no_env_variable():
+    """Test that no callbacks are loaded when environment variable is not set."""
+    if RAY_TUNE_CALLBACKS_ENV_VAR in os.environ:
+        del os.environ[RAY_TUNE_CALLBACKS_ENV_VAR]
+
+    callbacks = _initialize_env_callbacks()
+    assert len(callbacks) == 0
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/tune/tests/test_experiment.py

@@ -0,0 +1,117 @@
+import threading
+import unittest
+
+import ray
+from ray.tune import CheckpointConfig, register_trainable
+from ray.tune.error import TuneError
+from ray.tune.experiment import Experiment, _convert_to_experiment_list
+from ray.tune.utils import diagnose_serialization
+
+
+class ExperimentTest(unittest.TestCase):
+    def tearDown(self):
+        ray.shutdown()
+
+    def setUp(self):
+        def train_fn(config):
+            for i in range(100):
+                ray.tune.report(dict(timesteps_total=i))
+
+        register_trainable("f1", train_fn)
+
+    def testConvertExperimentFromExperiment(self):
+        exp1 = Experiment(
+            **{"name": "foo", "run": "f1", "config": {"script_min_iter_time_s": 0}}
+        )
+        result = _convert_to_experiment_list(exp1)
+        self.assertEqual(len(result), 1)
+        self.assertEqual(type(result), list)
+
+    def testConvertExperimentNone(self):
+        result = _convert_to_experiment_list(None)
+        self.assertEqual(len(result), 0)
+        self.assertEqual(type(result), list)
+
+    def testConvertExperimentList(self):
+        exp1 = Experiment(
+            **{"name": "foo", "run": "f1", "config": {"script_min_iter_time_s": 0}}
+        )
+        result = _convert_to_experiment_list([exp1, exp1])
+        self.assertEqual(len(result), 2)
+        self.assertEqual(type(result), list)
+
+    def testConvertExperimentJSON(self):
+        experiment = {
+            "name": {"run": "f1", "config": {"script_min_iter_time_s": 0}},
+            "named": {"run": "f1", "config": {"script_min_iter_time_s": 0}},
+        }
+        result = _convert_to_experiment_list(experiment)
+        self.assertEqual(len(result), 2)
+        self.assertEqual(type(result), list)
+
+    def testConvertExperimentIncorrect(self):
+        self.assertRaises(TuneError, lambda: _convert_to_experiment_list("hi"))
+
+    def testFuncTrainableCheckpointConfigValidation(self):
+        """Raise an error when trying to specify checkpoint_at_end/checkpoint_frequency
+        with a function trainable."""
+        with self.assertRaises(ValueError):
+            Experiment(
+                name="foo",
+                run="f1",  # Will point to a wrapped function trainable
+                checkpoint_config=CheckpointConfig(checkpoint_at_end=True),
+            )
+        with self.assertRaises(ValueError):
+            Experiment(
+                name="foo",
+                run="f1",
+                checkpoint_config=CheckpointConfig(checkpoint_frequency=1),
+            )
+        with self.assertRaises(ValueError):
+            Experiment(
+                name="foo",
+                run=lambda config: 1,
+                checkpoint_config=CheckpointConfig(checkpoint_at_end=True),
+            )
+
+    def testInvalidExperimentConfig(self):
+        with self.assertRaises(ValueError):
+            Experiment(name="foo", run="f1", config="invalid")
+
+        class InvalidClass:
+            def to_dict(self):
+                return {"valid": 1}
+
+        with self.assertRaises(ValueError):
+            Experiment(name="foo", run="f1", config=InvalidClass())
+
+        Experiment(name="foo", run="f1", config=InvalidClass().to_dict())
+
+
+class ValidateUtilTest(unittest.TestCase):
+    def testDiagnoseSerialization(self):
+        # this is not serializable
+        e = threading.Event()
+
+        def test(config):
+            print(e)
+
+        assert diagnose_serialization(test) is not True
+
+        # should help identify that 'e' should be moved into
+        # the `test` scope.
+
+        # correct implementation
+        def test(config):
+            e = threading.Event()
+            print(e)
+
+        assert diagnose_serialization(test) is True
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/tune/tests/test_experiment_analysis.py

@@ -0,0 +1,297 @@
+import os
+import pickle
+import tempfile
+from contextlib import contextmanager
+from pathlib import Path
+from typing import List
+
+import numpy as np
+import pandas as pd
+import pytest
+
+from ray import tune
+from ray.air._internal.uri_utils import URI
+from ray.air.constants import EXPR_PROGRESS_FILE, EXPR_RESULT_FILE
+from ray.train._internal.storage import _delete_fs_path
+from ray.train.tests.test_new_persistence import mock_s3_bucket_uri
+from ray.train.tests.util import create_dict_checkpoint, load_dict_checkpoint
+from ray.tune.analysis.experiment_analysis import ExperimentAnalysis
+from ray.tune.experiment import Trial
+from ray.tune.utils import flatten_dict
+
+NUM_TRIALS = 3
+NON_NAN_VALUE = 42
+PEAK_VALUE = 100
+
+
+def train_fn(config):
+    def report(metrics, should_checkpoint=True):
+        if should_checkpoint:
+            with create_dict_checkpoint(metrics) as checkpoint:
+                tune.report(metrics, checkpoint=checkpoint)
+        else:
+            tune.report(metrics)
+
+    id = config["id"]
+
+    report({"ascending": 1 * id, "peak": 0, "maybe_nan": np.nan, "iter": 1})
+    report({"ascending": 2 * id, "peak": 0, "maybe_nan": np.nan, "iter": 2})
+    report({"ascending": 3 * id, "peak": 0, "maybe_nan": np.nan, "iter": 3})
+    report(
+        {
+            "ascending": 4 * id,
+            "peak": 0,
+            "maybe_nan": NON_NAN_VALUE,
+            "iter": 4,
+        }
+    )
+    report({"ascending": 5 * id, "peak": PEAK_VALUE, "maybe_nan": np.nan, "iter": 5})
+
+    report(
+        {"ascending": 6 * id, "peak": -PEAK_VALUE, "maybe_nan": np.nan, "iter": 6},
+        should_checkpoint=False,
+    )
+    report(
+        {"ascending": 7 * id, "peak": 0, "maybe_nan": np.nan, "iter": 7},
+        should_checkpoint=False,
+    )
+
+
+def _get_trial_with_id(trials: List[Trial], id: int) -> Trial:
+    return [trial for trial in trials if trial.config["id"] == id][0]
+
+
+@contextmanager
+def dummy_context_manager():
+    yield "dummy value"
+
+
+@pytest.fixture(scope="module", params=["dir", "memory", "cloud"])
+def experiment_analysis(request):
+    load_from = request.param
+    tmp_path = Path(tempfile.mkdtemp())
+
+    context_manager = (
+        mock_s3_bucket_uri if load_from == "cloud" else dummy_context_manager
+    )
+
+    with context_manager() as cloud_storage_path:
+        storage_path = (
+            str(cloud_storage_path)
+            if load_from == "cloud"
+            else str(tmp_path / "fake_nfs")
+        )
+        ea = tune.run(
+            train_fn,
+            config={"id": tune.grid_search(list(range(1, NUM_TRIALS + 1)))},
+            metric="ascending",
+            mode="max",
+            storage_path=storage_path,
+            name="test_experiment_analysis",
+        )
+
+        if load_from in ["dir", "cloud"]:
+            # Test init without passing in in-memory trials.
+            # Load them from an experiment directory instead.
+            yield ExperimentAnalysis(
+                str(URI(storage_path) / "test_experiment_analysis"),
+                default_metric="ascending",
+                default_mode="max",
+            )
+        elif load_from == "memory":
+            yield ea
+        else:
+            raise NotImplementedError(f"Invalid param: {load_from}")
+
+
+@pytest.mark.parametrize("filetype", ["json", "csv"])
+def test_fetch_trial_dataframes(experiment_analysis, filetype):
+    if filetype == "csv":
+        # Delete all json files so that we can test fallback to csv loading
+        for trial in experiment_analysis.trials:
+            _delete_fs_path(
+                fs=trial.storage.storage_filesystem,
+                fs_path=os.path.join(trial.storage.trial_fs_path, EXPR_RESULT_FILE),
+            )
+    else:
+        assert filetype == "json"
+
+    dfs = experiment_analysis._fetch_trial_dataframes()
+    assert len(dfs) == NUM_TRIALS
+    assert all(isinstance(df, pd.DataFrame) for df in dfs.values())
+    assert {trial.trial_id for trial in experiment_analysis.trials} == set(dfs)
+
+    for trial_id, df in dfs.items():
+        trial_config = experiment_analysis.get_all_configs()[trial_id]
+        assert np.all(
+            df["ascending"].to_numpy() == np.arange(1, 8) * trial_config["id"]
+        )
+
+
+def test_fetch_trial_dataframes_with_errors(
+    experiment_analysis, tmp_path, propagate_logs, caplog
+):
+    # Add "corrupted" json files)
+    for trial in experiment_analysis.trials:
+        fs = trial.storage.storage_filesystem
+        with fs.open_output_stream(
+            os.path.join(trial.storage.trial_fs_path, EXPR_RESULT_FILE)
+        ) as f:
+            f.write(b"malformed")
+
+    experiment_analysis._fetch_trial_dataframes()
+    assert "Failed to fetch metrics" in caplog.text
+    caplog.clear()
+
+    # Delete ALL metrics files to check that a warning gets logged.
+    for trial in experiment_analysis.trials:
+        fs = trial.storage.storage_filesystem
+
+        # Delete ALL metrics files to check that a warning gets logged.
+        _delete_fs_path(
+            fs=trial.storage.storage_filesystem,
+            fs_path=os.path.join(trial.storage.trial_fs_path, EXPR_RESULT_FILE),
+        )
+        _delete_fs_path(
+            fs=trial.storage.storage_filesystem,
+            fs_path=os.path.join(trial.storage.trial_fs_path, EXPR_PROGRESS_FILE),
+        )
+
+    experiment_analysis._fetch_trial_dataframes()
+    assert "Could not fetch metrics for" in caplog.text
+    assert "FileNotFoundError" in caplog.text
+    caplog.clear()
+
+
+def test_get_all_configs(experiment_analysis):
+    configs = experiment_analysis.get_all_configs()
+    assert len(configs) == NUM_TRIALS
+    assert all(isinstance(config, dict) for config in configs.values())
+    assert {trial.trial_id for trial in experiment_analysis.trials} == set(configs)
+    for trial_id, config in configs.items():
+        trial = [
+            trial for trial in experiment_analysis.trials if trial.trial_id == trial_id
+        ][0]
+        assert trial.config == config
+
+
+def test_dataframe(experiment_analysis):
+    with pytest.raises(ValueError):
+        # Invalid mode
+        df = experiment_analysis.dataframe(mode="bad")
+
+    with pytest.raises(ValueError):
+        # Should raise because we didn't pass a metric
+        df = experiment_analysis.dataframe(mode="max")
+
+    # If we specify `max`, we expect the largets ever observed result
+    df = experiment_analysis.dataframe(metric="peak", mode="max")
+    assert df.iloc[0]["peak"] == PEAK_VALUE
+
+    # If we specify `min`, we expect the lowest ever observed result
+    df = experiment_analysis.dataframe(metric="peak", mode="min")
+    assert df.iloc[0]["peak"] == -PEAK_VALUE
+
+    # If we don't pass a mode, we just fetch the last result
+    df = experiment_analysis.dataframe(metric="peak")
+    assert df.iloc[0]["peak"] == 0
+    assert df.iloc[0]["iter"] == 7
+
+
+def test_default_properties(experiment_analysis):
+    # The last trial has the highest score (according to the default metric/mode).
+    best_trial = _get_trial_with_id(experiment_analysis.trials, NUM_TRIALS)
+    assert experiment_analysis.best_trial == best_trial
+    assert experiment_analysis.best_config == best_trial.config
+    # The last (most recent) checkpoint has the highest score.
+    assert experiment_analysis.best_checkpoint == best_trial.checkpoint
+    # NaN != NaN, so fill them in for this equality check.
+    assert experiment_analysis.best_dataframe.fillna(-1).equals(
+        experiment_analysis.trial_dataframes[best_trial.trial_id].fillna(-1)
+    )
+    assert experiment_analysis.best_result == best_trial.last_result
+
+    result_df_dict = experiment_analysis.best_result_df.iloc[0].to_dict()
+    # Converting -> pandas -> dict flattens the dict.
+    best_trial_dict = flatten_dict(best_trial.last_result, delimiter="/")
+    assert result_df_dict["ascending"] == best_trial_dict["ascending"]
+
+    assert len(experiment_analysis.results) == NUM_TRIALS
+    assert len(experiment_analysis.results_df) == NUM_TRIALS
+
+
+def test_get_best_config(experiment_analysis):
+    assert experiment_analysis.get_best_config()["id"] == NUM_TRIALS
+    assert (
+        experiment_analysis.get_best_config(metric="ascending", mode="min")["id"] == 1
+    )
+
+    assert not experiment_analysis.get_best_config(metric="maybe_nan", scope="last")
+
+
+def test_get_best_trial(experiment_analysis):
+    assert (
+        experiment_analysis.get_best_trial().config
+        == experiment_analysis.get_best_config()
+    )
+
+    assert not experiment_analysis.get_best_trial(metric="maybe_nan")
+    assert experiment_analysis.get_best_trial(
+        metric="maybe_nan", filter_nan_and_inf=False
+    )
+
+
+def test_get_best_checkpoint(experiment_analysis):
+    best_trial = experiment_analysis.get_best_trial()
+    best_checkpoint = load_dict_checkpoint(
+        experiment_analysis.get_best_checkpoint(best_trial)
+    )
+    # NOTE: There are 7 reports, but only the first 5 include a checkpoint.
+    assert best_checkpoint["ascending"] == 5 * NUM_TRIALS
+
+    best_checkpoint = load_dict_checkpoint(
+        experiment_analysis.get_best_checkpoint(
+            best_trial, metric="ascending", mode="min"
+        )
+    )
+    assert best_checkpoint["ascending"] == 1 * NUM_TRIALS
+
+    # Filter checkpoints w/ NaN metrics
+    best_checkpoint = load_dict_checkpoint(
+        experiment_analysis.get_best_checkpoint(best_trial, metric="maybe_nan")
+    )
+    assert best_checkpoint["maybe_nan"] == NON_NAN_VALUE
+
+
+def test_get_last_checkpoint(experiment_analysis):
+    # Defaults to getting the last checkpoint of the best trial.
+    last_checkpoint = load_dict_checkpoint(experiment_analysis.get_last_checkpoint())
+    assert last_checkpoint["iter"] == 5  # See note
+
+    last_checkpoint = load_dict_checkpoint(
+        experiment_analysis.get_last_checkpoint(
+            trial=_get_trial_with_id(experiment_analysis.trials, 1)
+        )
+    )
+    assert last_checkpoint["ascending"] == 5 * 1  # See note
+
+
+def test_pickle(experiment_analysis, tmp_path):
+    pickle_path = os.path.join(tmp_path, "analysis.pkl")
+    with open(pickle_path, "wb") as f:
+        pickle.dump(experiment_analysis, f)
+
+    assert experiment_analysis.get_best_trial(metric="ascending", mode="max")
+
+    with open(pickle_path, "rb") as f:
+        loaded_analysis = pickle.load(f)
+
+    assert loaded_analysis.get_best_trial(metric="ascending", mode="max")
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/tune/tests/test_function_api.py

@@ -0,0 +1,277 @@
+import json
+import os
+import sys
+import tempfile
+import unittest
+
+import ray
+from ray import tune
+from ray.air.constants import TRAINING_ITERATION
+from ray.rllib import _register_all
+from ray.train.tests.util import mock_storage_context
+from ray.tune import Checkpoint, CheckpointConfig
+from ray.tune.execution.placement_groups import PlacementGroupFactory
+from ray.tune.result import DEFAULT_METRIC
+from ray.tune.schedulers import ResourceChangingScheduler
+from ray.tune.trainable import with_parameters, wrap_function
+
+
+class FunctionCheckpointingTest(unittest.TestCase):
+    def setUp(self):
+        self.tmpdir = tempfile.TemporaryDirectory()
+
+    def create_trainable(self, train_fn):
+        return wrap_function(train_fn)(storage=mock_storage_context())
+
+    def tearDown(self):
+        self.tmpdir.cleanup()
+
+    def testCheckpointReuse(self):
+        """Test that repeated save/restore never reuses same checkpoint dir."""
+
+        def train_fn(config):
+            checkpoint = ray.tune.get_checkpoint()
+            if checkpoint:
+                with checkpoint.as_directory() as checkpoint_dir:
+                    count = sum(
+                        "checkpoint-" in path for path in os.listdir(checkpoint_dir)
+                    )
+                    assert count == 1, os.listdir(checkpoint_dir)
+
+            for step in range(20):
+                with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
+                    path = os.path.join(
+                        temp_checkpoint_dir, "checkpoint-{}".format(step)
+                    )
+                    open(path, "a").close()
+                    ray.tune.report(
+                        dict(test=step),
+                        checkpoint=Checkpoint.from_directory(temp_checkpoint_dir),
+                    )
+
+        checkpoint = None
+        for i in range(5):
+            new_trainable = self.create_trainable(train_fn)
+            if checkpoint:
+                new_trainable.restore(checkpoint)
+            for i in range(2):
+                result = new_trainable.train()
+            checkpoint = new_trainable.save()
+            new_trainable.stop()
+        assert result[TRAINING_ITERATION] == 10
+
+    def testFunctionRecurringSave(self):
+        """This tests that save and restore are commutative."""
+
+        def train_fn(config):
+            for step in range(10):
+                with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
+                    if step % 3 == 0:
+                        path = os.path.join(temp_checkpoint_dir, "checkpoint.json")
+                        with open(path, "w") as f:
+                            json.dump({"step": step}, f)
+                    ray.tune.report(
+                        dict(test=step),
+                        checkpoint=Checkpoint.from_directory(temp_checkpoint_dir),
+                    )
+
+        new_trainable = self.create_trainable(train_fn)
+        new_trainable.train()
+        checkpoint_obj = new_trainable.save()
+        new_trainable.restore(checkpoint_obj)
+        checkpoint = new_trainable.save()
+
+        new_trainable.stop()
+
+        new_trainable2 = self.create_trainable(train_fn)
+        new_trainable2.restore(checkpoint)
+        new_trainable2.train()
+        new_trainable2.stop()
+
+
+class FunctionApiTest(unittest.TestCase):
+    def setUp(self):
+        ray.init(num_cpus=4, num_gpus=0, object_store_memory=150 * 1024 * 1024)
+
+    def tearDown(self):
+        ray.shutdown()
+        _register_all()  # re-register the evicted objects
+
+    def testCheckpointError(self):
+        def train_fn(config):
+            pass
+
+        with self.assertRaises(ValueError):
+            tune.run(
+                train_fn, checkpoint_config=CheckpointConfig(checkpoint_frequency=1)
+            )
+        with self.assertRaises(ValueError):
+            tune.run(
+                train_fn, checkpoint_config=CheckpointConfig(checkpoint_at_end=True)
+            )
+
+    def testWithParameters(self):
+        class Data:
+            def __init__(self):
+                self.data = [0] * 500_000
+
+        data = Data()
+        data.data[100] = 1
+
+        def train_fn(config, data=None):
+            data.data[101] = 2  # Changes are local
+            ray.tune.report(dict(metric=len(data.data), hundred=data.data[100]))
+
+        trial_1, trial_2 = tune.run(
+            with_parameters(train_fn, data=data), num_samples=2
+        ).trials
+
+        self.assertEqual(data.data[101], 0)
+        self.assertEqual(trial_1.last_result["metric"], 500_000)
+        self.assertEqual(trial_1.last_result["hundred"], 1)
+        self.assertEqual(trial_2.last_result["metric"], 500_000)
+        self.assertEqual(trial_2.last_result["hundred"], 1)
+        self.assertTrue(str(trial_1).startswith("train_"))
+
+        # With checkpoint dir parameter
+        def train_fn(config, data=None):
+            data.data[101] = 2  # Changes are local
+            ray.tune.report(dict(metric=len(data.data)))
+
+        trial_1, trial_2 = tune.run(
+            with_parameters(train_fn, data=data), num_samples=2
+        ).trials
+
+        self.assertEqual(data.data[101], 0)
+        self.assertEqual(trial_1.last_result["metric"], 500_000)
+        self.assertEqual(trial_2.last_result["metric"], 500_000)
+        self.assertTrue(str(trial_1).startswith("train_"))
+
+    def testWithParameters2(self):
+        class Data:
+            def __init__(self):
+                import numpy as np
+
+                self.data = np.random.rand((2 * 1024 * 1024))
+
+        def train_fn(config, data=None):
+            pass
+
+        trainable = tune.with_parameters(train_fn, data=Data())
+        # ray.cloudpickle will crash for some reason
+        import cloudpickle as cp
+
+        dumped = cp.dumps(trainable)
+        assert sys.getsizeof(dumped) < 100 * 1024
+
+    def testNewResources(self):
+        sched = ResourceChangingScheduler(
+            resources_allocation_function=(
+                lambda a, b, c, d: PlacementGroupFactory([{"CPU": 2}])
+            )
+        )
+
+        def train_fn(config):
+            ray.tune.report(
+                dict(metric=1, resources=ray.tune.get_context().get_trial_resources())
+            )
+
+        analysis = tune.run(
+            train_fn,
+            scheduler=sched,
+            stop={"training_iteration": 2},
+            resources_per_trial=PlacementGroupFactory([{"CPU": 1}]),
+            num_samples=1,
+        )
+
+        results_list = list(analysis.results.values())
+        assert results_list[0]["resources"].head_cpus == 2.0
+
+    def testWithParametersTwoRuns1(self):
+        # Makes sure two runs in the same script but different ray sessions
+        # pass (https://github.com/ray-project/ray/issues/16609)
+        def train_fn(config, extra=4):
+            ray.tune.report(dict(metric=extra))
+
+        trainable = tune.with_parameters(train_fn, extra=8)
+        out = tune.run(trainable, metric="metric", mode="max")
+        self.assertEqual(out.best_result["metric"], 8)
+
+        self.tearDown()
+        self.setUp()
+
+        def train_fn_2(config, extra=5):
+            ray.tune.report(dict(metric=extra))
+
+        trainable = tune.with_parameters(train_fn_2, extra=9)
+        out = tune.run(trainable, metric="metric", mode="max")
+        self.assertEqual(out.best_result["metric"], 9)
+
+    def testWithParametersTwoRuns2(self):
+        # Makes sure two runs in the same script
+        # pass (https://github.com/ray-project/ray/issues/16609)
+        def train_fn(config, extra=4):
+            ray.tune.report(dict(metric=extra))
+
+        def train_fn_2(config, extra=5):
+            ray.tune.report(dict(metric=extra))
+
+        trainable1 = tune.with_parameters(train_fn, extra=8)
+        trainable2 = tune.with_parameters(train_fn_2, extra=9)
+
+        out1 = tune.run(trainable1, metric="metric", mode="max")
+        out2 = tune.run(trainable2, metric="metric", mode="max")
+        self.assertEqual(out1.best_result["metric"], 8)
+        self.assertEqual(out2.best_result["metric"], 9)
+
+    def testReturnAnonymous(self):
+        def train_fn(config):
+            return config["a"]
+
+        trial_1, trial_2 = tune.run(
+            train_fn, config={"a": tune.grid_search([4, 8])}
+        ).trials
+
+        self.assertEqual(trial_1.last_result[DEFAULT_METRIC], 4)
+        self.assertEqual(trial_2.last_result[DEFAULT_METRIC], 8)
+
+    def testReturnSpecific(self):
+        def train_fn(config):
+            return {"m": config["a"]}
+
+        trial_1, trial_2 = tune.run(
+            train_fn, config={"a": tune.grid_search([4, 8])}
+        ).trials
+
+        self.assertEqual(trial_1.last_result["m"], 4)
+        self.assertEqual(trial_2.last_result["m"], 8)
+
+    def testYieldAnonymous(self):
+        def train_fn(config):
+            for i in range(10):
+                yield config["a"] + i
+
+        trial_1, trial_2 = tune.run(
+            train_fn, config={"a": tune.grid_search([4, 8])}
+        ).trials
+
+        self.assertEqual(trial_1.last_result[DEFAULT_METRIC], 4 + 9)
+        self.assertEqual(trial_2.last_result[DEFAULT_METRIC], 8 + 9)
+
+    def testYieldSpecific(self):
+        def train_fn(config):
+            for i in range(10):
+                yield {"m": config["a"] + i}
+
+        trial_1, trial_2 = tune.run(
+            train_fn, config={"a": tune.grid_search([4, 8])}
+        ).trials
+
+        self.assertEqual(trial_1.last_result["m"], 4 + 9)
+        self.assertEqual(trial_2.last_result["m"], 8 + 9)
+
+
+if __name__ == "__main__":
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/tune/tests/test_integration_pytorch_lightning.py

@@ -0,0 +1,136 @@
+import shutil
+import tempfile
+import unittest
+
+import lightning.pytorch as pl
+import torch
+from torch.utils.data import DataLoader, Dataset
+
+from ray import tune
+from ray.air.constants import TRAINING_ITERATION
+from ray.tune import CheckpointConfig
+from ray.tune.integration.pytorch_lightning import TuneReportCheckpointCallback
+
+
+class _MockDataset(Dataset):
+    def __init__(self, values):
+        self.values = values
+
+    def __getitem__(self, index):
+        return self.values[index]
+
+    def __len__(self):
+        return len(self.values)
+
+
+class _MockModule(pl.LightningModule):
+    def __init__(self, loss, acc):
+        super().__init__()
+
+        self._dummy = torch.nn.Parameter(torch.zeros(1))
+        self.loss_val = loss
+        self.acc_val = acc
+
+    def forward(self, *args, **kwargs):
+        return self._dummy
+
+    def training_step(self, train_batch, batch_idx):
+        loss = self._dummy.sum() * 0 + self.loss_val
+        self.log("loss", loss)
+        self.log("acc", torch.tensor(self.acc_val))
+        return loss
+
+    def validation_step(self, val_batch, batch_idx):
+        self.log("avg_val_loss", torch.tensor(self.loss_val * 1.1))
+        self.log("avg_val_acc", torch.tensor(self.acc_val * 0.9))
+
+    def configure_optimizers(self):
+        return torch.optim.SGD([self._dummy], lr=0.001)
+
+    def train_dataloader(self):
+        return DataLoader(_MockDataset(list(range(10))), batch_size=1)
+
+    def val_dataloader(self):
+        return DataLoader(_MockDataset(list(range(10))), batch_size=1)
+
+
+class PyTorchLightningIntegrationTest(unittest.TestCase):
+    def setUp(self):
+        pass
+
+    def tearDown(self):
+        pass
+
+    def testReportCallbackUnnamed(self):
+        def train_fn(config):
+            module = _MockModule(10.0, 20.0)
+            trainer = pl.Trainer(
+                max_epochs=1,
+                callbacks=[
+                    TuneReportCheckpointCallback(
+                        on="validation_end", save_checkpoints=False
+                    )
+                ],
+            )
+            trainer.fit(module)
+
+        analysis = tune.run(train_fn, stop={TRAINING_ITERATION: 1})
+
+        self.assertEqual(analysis.trials[0].last_result["avg_val_loss"], 10.0 * 1.1)
+
+    def testReportCallbackNamed(self):
+        def train_fn(config):
+            module = _MockModule(10.0, 20.0)
+            trainer = pl.Trainer(
+                max_epochs=1,
+                callbacks=[
+                    TuneReportCheckpointCallback(
+                        metrics={"tune_loss": "avg_val_loss"},
+                        on="validation_end",
+                        save_checkpoints=False,
+                    )
+                ],
+            )
+            trainer.fit(module)
+
+        analysis = tune.run(train_fn, stop={TRAINING_ITERATION: 1})
+
+        self.assertEqual(analysis.trials[0].last_result["tune_loss"], 10.0 * 1.1)
+
+    def testCheckpointCallback(self):
+        tmpdir = tempfile.mkdtemp()
+        self.addCleanup(lambda: shutil.rmtree(tmpdir))
+
+        def train_fn(config):
+            module = _MockModule(10.0, 20.0)
+            trainer = pl.Trainer(
+                max_epochs=10,
+                callbacks=[
+                    TuneReportCheckpointCallback(
+                        filename="trainer.ckpt", on=["train_epoch_end"]
+                    )
+                ],
+            )
+            trainer.fit(module)
+
+        checkpoint_config = CheckpointConfig(num_to_keep=100)
+        tuner = tune.Tuner(
+            train_fn,
+            run_config=tune.RunConfig(
+                stop={TRAINING_ITERATION: 10},
+                storage_path=tmpdir,
+                checkpoint_config=checkpoint_config,
+            ),
+        )
+        results = tuner.fit()
+
+        # 1 checkpoint per epoch
+        self.assertEqual(len(results[0].best_checkpoints), 10)
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(sys.argv[1:] + ["-v", __file__]))

python/ray/tune/tests/test_logger.py

@@ -0,0 +1,385 @@
+import csv
+import glob
+import json
+import os
+import shutil
+import sys
+import tempfile
+import unittest
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Optional
+
+import numpy as np
+import pytest
+
+import ray
+from ray.air.constants import (
+    EXPR_PARAM_FILE,
+    EXPR_PARAM_PICKLE_FILE,
+    EXPR_PROGRESS_FILE,
+    EXPR_RESULT_FILE,
+)
+from ray.cloudpickle import cloudpickle
+from ray.tune import Checkpoint
+from ray.tune.logger import (
+    CSVLoggerCallback,
+    JsonLoggerCallback,
+    TBXLoggerCallback,
+)
+from ray.tune.logger.aim import AimLoggerCallback
+from ray.tune.utils import flatten_dict
+
+
+@dataclass
+class Trial:
+    evaluated_params: dict
+    trial_id: str
+    logdir: str
+    experiment_path: Optional[str] = None
+    experiment_dir_name: Optional[str] = None
+    path: Optional[str] = None
+    checkpoint: Optional[Checkpoint] = None
+
+    @property
+    def config(self):
+        return self.evaluated_params
+
+    def init_local_path(self):
+        return
+
+    @property
+    def local_path(self):
+        if self.logdir:
+            return self.logdir
+        if not self.experiment_dir_name:
+            return None
+        return str(Path(self.experiment_path) / self.experiment_dir_name)
+
+    @property
+    def local_experiment_path(self):
+        return self.experiment_path
+
+    def __hash__(self):
+        return hash(self.trial_id)
+
+    def get_ray_actor_ip(self) -> str:
+        return ray.util.get_node_ip_address()
+
+
+def result(t, rew, **kwargs):
+    results = dict(
+        time_total_s=t,
+        episode_reward_mean=rew,
+        mean_accuracy=rew * 2,
+        training_iteration=int(t),
+    )
+    results.update(kwargs)
+    return results
+
+
+class LoggerSuite(unittest.TestCase):
+    """Test built-in loggers."""
+
+    def setUp(self):
+        self.test_dir = tempfile.mkdtemp()
+
+    def tearDown(self):
+        shutil.rmtree(self.test_dir, ignore_errors=True)
+
+    def testCSV(self):
+        config = {"a": 2, "b": 5, "c": {"c": {"D": 123}, "e": None}}
+        t = Trial(evaluated_params=config, trial_id="csv", logdir=self.test_dir)
+        logger = CSVLoggerCallback()
+        logger.on_trial_result(0, [], t, result(0, 4))
+        logger.on_trial_result(1, [], t, result(1, 5))
+        logger.on_trial_result(
+            2, [], t, result(2, 6, score=[1, 2, 3], hello={"world": 1})
+        )
+
+        logger.on_trial_complete(3, [], t)
+        self._validate_csv_result()
+
+    def testCSVEmptyHeader(self):
+        """Test that starting a trial twice does not lead to empty CSV headers.
+
+        In a previous bug, the CSV header was sometimes missing when a trial
+        crashed before reporting results. See
+        https://github.com/ray-project/ray/issues/15106
+        """
+        config = {"a": 2, "b": 5, "c": {"c": {"D": 123}, "e": None}}
+        t = Trial(evaluated_params=config, trial_id="csv", logdir=self.test_dir)
+        logger = CSVLoggerCallback()
+        logger.on_trial_start(0, [], t)
+        logger.on_trial_start(0, [], t)
+        logger.on_trial_result(1, [], t, result(1, 5))
+
+        with open(os.path.join(self.test_dir, "progress.csv"), "rt") as f:
+            csv_contents = f.read()
+
+        csv_lines = csv_contents.split("\n")
+
+        # Assert header has been written to progress.csv
+        assert "training_iteration" in csv_lines[0]
+
+    def _validate_csv_result(self):
+        results = []
+        result_file = os.path.join(self.test_dir, EXPR_PROGRESS_FILE)
+        with open(result_file, "rt") as fp:
+            reader = csv.DictReader(fp)
+            for row in reader:
+                results.append(row)
+
+        self.assertEqual(len(results), 3)
+        self.assertSequenceEqual(
+            [int(row["episode_reward_mean"]) for row in results], [4, 5, 6]
+        )
+
+    def testJSON(self):
+        config = {"a": 2, "b": 5, "c": {"c": {"D": 123}, "e": None}}
+        t = Trial(evaluated_params=config, trial_id="json", logdir=self.test_dir)
+        logger = JsonLoggerCallback()
+        logger.on_trial_result(0, [], t, result(0, 4))
+        logger.on_trial_result(1, [], t, result(1, 5))
+        logger.on_trial_result(
+            2, [], t, result(2, 6, score=[1, 2, 3], hello={"world": 1})
+        )
+
+        logger.on_trial_complete(3, [], t)
+        self._validate_json_result(config)
+
+    def _validate_json_result(self, config):
+        # Check result logs
+        results = []
+        result_file = os.path.join(self.test_dir, EXPR_RESULT_FILE)
+        with open(result_file, "rt") as fp:
+            for row in fp.readlines():
+                results.append(json.loads(row))
+
+        self.assertEqual(len(results), 3)
+        self.assertSequenceEqual(
+            [int(row["episode_reward_mean"]) for row in results], [4, 5, 6]
+        )
+
+        # Check json saved config file
+        config_file = os.path.join(self.test_dir, EXPR_PARAM_FILE)
+        with open(config_file, "rt") as fp:
+            loaded_config = json.load(fp)
+
+        self.assertEqual(loaded_config, config)
+
+        # Check pickled config file
+        config_file = os.path.join(self.test_dir, EXPR_PARAM_PICKLE_FILE)
+        with open(config_file, "rb") as fp:
+            loaded_config = cloudpickle.load(fp)
+
+        self.assertEqual(loaded_config, config)
+
+    def testTBX(self):
+        config = {
+            "a": 2,
+            "b": [1, 2],
+            "c": {"c": {"D": 123}},
+            "int32": np.int32(1),
+            "int64": np.int64(2),
+            "bool8": np.bool_(True),
+            "float32": np.float32(3),
+            "float64": np.float64(4),
+            "bad": np.float128(4),
+        }
+        t = Trial(evaluated_params=config, trial_id="tbx", logdir=self.test_dir)
+        logger = TBXLoggerCallback()
+        logger.on_trial_result(0, [], t, result(0, 4))
+        logger.on_trial_result(1, [], t, result(1, 5))
+        logger.on_trial_result(
+            2, [], t, result(2, 6, score=[1, 2, 3], hello={"world": 1})
+        )
+
+        logger.on_trial_complete(3, [], t)
+
+        self._validate_tbx_result(
+            params=(b"float32", b"float64", b"int32", b"int64", b"bool8"),
+            excluded_params=(b"bad",),
+        )
+
+    def _validate_tbx_result(self, params=None, excluded_params=None):
+        try:
+            from tensorflow.python.summary.summary_iterator import summary_iterator
+        except ImportError:
+            print("Skipping rest of test as tensorflow is not installed.")
+            return
+
+        events_file = list(glob.glob(f"{self.test_dir}/events*"))[0]
+        results = []
+        excluded_params = excluded_params or []
+        for event in summary_iterator(events_file):
+            for v in event.summary.value:
+                if v.tag == "ray/tune/episode_reward_mean":
+                    results.append(v.simple_value)
+                elif v.tag == "_hparams_/experiment" and params:
+                    for key in params:
+                        self.assertIn(key, v.metadata.plugin_data.content)
+                    for key in excluded_params:
+                        self.assertNotIn(key, v.metadata.plugin_data.content)
+                elif v.tag == "_hparams_/session_start_info" and params:
+                    for key in params:
+                        self.assertIn(key, v.metadata.plugin_data.content)
+                    for key in excluded_params:
+                        self.assertNotIn(key, v.metadata.plugin_data.content)
+
+        self.assertEqual(len(results), 3)
+        self.assertSequenceEqual([int(res) for res in results], [4, 5, 6])
+
+    def testBadTBX(self):
+        config = {"b": (1, 2, 3)}
+        t = Trial(evaluated_params=config, trial_id="tbx", logdir=self.test_dir)
+        logger = TBXLoggerCallback()
+        logger.on_trial_result(0, [], t, result(0, 4))
+        logger.on_trial_result(1, [], t, result(1, 5))
+        logger.on_trial_result(
+            2, [], t, result(2, 6, score=[1, 2, 3], hello={"world": 1})
+        )
+        with self.assertLogs("ray.tune.logger", level="INFO") as cm:
+            logger.on_trial_complete(3, [], t)
+        assert "INFO" in cm.output[0]
+
+
+@pytest.mark.skipif(sys.version_info >= (3, 12), reason="Aim doesn't support py312")
+class AimLoggerSuite(unittest.TestCase):
+    """Test Aim integration."""
+
+    def setUp(self):
+        self.test_dir = tempfile.mkdtemp()
+
+    def tearDown(self):
+        shutil.rmtree(self.test_dir, ignore_errors=True)
+
+    def initialize_logger(self, repo=None, experiment_name=None, metrics=None):
+        try:
+            from aim import Repo
+        except ImportError:
+            print("Skipping rest of test as aim is not installed.")
+            return
+
+        class Dummy:
+            pass
+
+        self.config = {
+            "a": 2,
+            "b": [1, 2],
+            "c": {"d": {"e": 123}},
+            "int32": np.int32(1),
+            "int64": np.int64(2),
+            "bool8": np.bool_(True),
+            "float32": np.float32(3),
+            "float64": np.float64(4),
+            "bad": Dummy(),
+        }
+        trial_logdir = os.path.join(self.test_dir, "trial_logdir")
+        trials = [
+            Trial(
+                evaluated_params=self.config,
+                trial_id="aim_1",
+                experiment_path=self.test_dir,
+                logdir=trial_logdir,
+                experiment_dir_name="aim_test",
+                path="bucket/aim_test/trial_0_logdir",
+            ),
+            Trial(
+                evaluated_params=self.config,
+                trial_id="aim_2",
+                experiment_path=self.test_dir,
+                logdir=trial_logdir,
+                experiment_dir_name="aim_test",
+                path="bucket/aim_test/trial_1_logdir",
+            ),
+        ]
+
+        # Test that aim repo is saved to the experiment directory
+        # (one up from the trial directory) as the default.
+        # In this example, this is `self.test_dir`.
+        repo = repo or self.test_dir
+        logger = AimLoggerCallback(
+            repo=repo, experiment_name=experiment_name, metrics=metrics
+        )
+
+        for i, t in enumerate(trials):
+            with self.assertLogs("ray.tune.logger", level="INFO") as cm:
+                logger.log_trial_start(t)
+                # Check that we log that the "bad" hparam gets thrown away
+                assert "INFO" in cm.output[0]
+
+            logger.on_trial_result(0, [], t, result(0, 3 * i + 1))
+            logger.on_trial_result(1, [], t, result(1, 3 * i + 2))
+            logger.on_trial_result(
+                2, [], t, result(2, 3 * i + 3, score=[1, 2, 3], hello={"world": 1})
+            )
+            logger.on_trial_complete(3, [], t)
+
+        aim_repo = Repo(repo)
+        runs = list(aim_repo.iter_runs())
+        assert len(runs) == 2
+        runs.sort(key=lambda r: r["trial_id"])
+        return runs
+
+    def validateLogs(self, runs: list, metrics: list = None):
+        expected_logged_hparams = set(flatten_dict(self.config)) - {"bad"}
+
+        for i, run in enumerate(runs):
+            assert set(run["hparams"]) == expected_logged_hparams
+            assert run.get("trial_log_dir")
+            assert run.get("trial_ip")
+
+            results = None
+            all_tune_metrics = set()
+            for metric in run.metrics():
+                if metric.name.startswith("ray/tune/"):
+                    all_tune_metrics.add(metric.name.replace("ray/tune/", ""))
+                if metric.name == "ray/tune/episode_reward_mean":
+                    results = metric.values.values_list()
+
+            assert results
+            # Make sure that the set of reported metrics matches with the
+            # set of metric names passed in
+            # If None is passed in, then all Tune metrics get reported
+            assert metrics is None or set(metrics) == all_tune_metrics
+
+            results = [int(res) for res in results]
+            if i == 0:
+                self.assertSequenceEqual(results, [1, 2, 3])
+            elif i == 1:
+                self.assertSequenceEqual(results, [4, 5, 6])
+
+    def testDefault(self):
+        """Test AimLoggerCallback with default settings.
+        - Req: a repo gets created at the experiment-level directory.
+        - Req: the experiment param passed into each aim Run is the Tune experiment name
+        """
+        runs = self.initialize_logger()
+        self.validateLogs(runs)
+        for run in runs:
+            assert run.repo.path == os.path.join(self.test_dir, ".aim")
+            assert run.experiment == "aim_test"
+
+    def testFilteredMetrics(self):
+        """Test AimLoggerCallback, logging only a subset of metrics."""
+        metrics_to_log = ("episode_reward_mean",)
+        runs = self.initialize_logger(metrics=metrics_to_log)
+        self.validateLogs(runs=runs, metrics=metrics_to_log)
+
+    def testCustomConfigurations(self):
+        """Test AimLoggerCallback, setting a custom repo and experiment name."""
+        custom_repo = os.path.join(self.test_dir, "custom_repo")
+        runs = self.initialize_logger(repo=custom_repo, experiment_name="custom")
+        self.validateLogs(runs)
+        for run in runs:
+            assert run.repo.path == os.path.join(custom_repo, ".aim")
+            assert run.experiment == "custom"
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__] + sys.argv[1:]))

python/ray/tune/tests/test_multi_tenancy.py

@@ -0,0 +1,128 @@
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+import pytest
+
+import ray
+
+
+@pytest.fixture
+def ray_start_4_cpus():
+    address_info = ray.init(num_cpus=4)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+@pytest.mark.parametrize("exit_same", [False, True])
+def test_registry_conflict(ray_start_4_cpus, tmpdir, exit_same):
+    """Two concurrent Tune runs can conflict with each other when they
+    use a trainable with the same name.
+
+    This test starts two runs in parallel and asserts that our fix in
+    https://github.com/ray-project/ray/pull/33095 resolves the issue.
+
+    This is how we schedule the runs:
+
+    - We have two runs. Every run starts two trials.
+    - Run 1 starts 1 trial immediately. This trial starts with
+      the correct parameters for the script. The trial hangs until the file
+      ``run_2_finished`` is deleted.
+    - Run 2 starts as soon as the first trial of Run 1 runs (by waiting
+      until the ``run_1_running`` file is deleted by that trial). It will overwrite
+      the global registry trainable with the same name.
+    - Run 2 finishes both trials. The script finishes with the expected
+      parameters.
+    - Run 2 then deletes the ``run_2_finished`` marker, allowing Run 1 trial 1
+      to continue training. When training finishes, the second trial launches.
+      This second trial then uses the overwritten trainable, that is, the wrong
+      parameters unless you use the workaround.
+    - Run 1 finally finishes, and we compare the expected results with the actual
+      results.
+
+    NOTE: Two errors can occur with registry conflicts. First,
+    the trainable can be overwritten and captured, for example, when a fixed value
+    is included in the trainable. The second trial of run 1 then has a wrong
+    parameter and reports a wrong metric (from run 2).
+
+    The second error occurs when the second run finishes fully and its objects
+    are garbage collected. In this case, the first run tries to find the trainable
+    registered by run 2, but fails lookup because the objects have been
+    removed already. Note that these objects are registered with
+    ``tune.with_parameters()`` (not the global registry store).
+    We test both scenarios using the ``exit_same`` parameter.
+    """
+    # Create file markers
+    run_1_running = tmpdir / "run_1_running"
+    run_1_finished = tmpdir / "run_1_finished"
+    run_2_finished = tmpdir / "run_2_finished"
+
+    run_1_running.write_text("", encoding="utf-8")
+    run_1_finished.write_text("", encoding="utf-8")
+    run_2_finished.write_text("", encoding="utf-8")
+
+    ray_address = ray_start_4_cpus.address_info["address"]
+
+    run_1_env = os.environ.copy()
+    run_1_env.update(
+        {
+            "RAY_ADDRESS": ray_address,
+            "FIXED_VAL": str(1),
+            "VAL_1": str(2),
+            "VAL_2": str(3),
+            # Run 1 can start immediately
+            "HANG_RUN_MARKER": "",
+            # Allow second run to start once first trial of first run is started
+            "DELETE_TRIAL_MARKER": str(run_1_running),
+            # Hang in first trial until the second run finished
+            "HANG_TRIAL_MARKER": str(run_2_finished),
+            # Mark run 1 as completed
+            "DELETE_RUN_MARKER": str(run_1_finished),
+            # Do not wait at end
+            "HANG_END_MARKER": "",
+        }
+    )
+
+    run_2_env = os.environ.copy()
+    run_2_env.update(
+        {
+            "RAY_ADDRESS": ray_address,
+            "FIXED_VAL": str(4),
+            "VAL_1": str(5),
+            "VAL_2": str(6),
+            # Wait until first trial of first run is running
+            "HANG_RUN_MARKER": str(run_1_running),
+            # Don't delete during run
+            "DELETE_TRIAL_MARKER": "",
+            # No need to hang in trial
+            "HANG_TRIAL_MARKER": "",
+            # After full run finished, allow first run to continue
+            "DELETE_RUN_MARKER": str(run_2_finished),
+            # Wait until first run finished
+            # If we don't do this, we actually don't die because of parameter conflict
+            # but because of "The object's owner has exited" - so we test this
+            # separately
+            "HANG_END_MARKER": str(run_1_finished) if exit_same else "",
+        }
+    )
+
+    script_path = Path(__file__).parent / "_test_multi_tenancy_run.py"
+
+    run_1 = subprocess.Popen(
+        [sys.executable, script_path], env=run_1_env, stderr=subprocess.PIPE
+    )
+    print("Started run 1:", run_1.pid)
+
+    run_2 = subprocess.Popen([sys.executable, script_path], env=run_2_env)
+    print("Started run 2:", run_2.pid)
+
+    assert run_2.wait() == 0
+    assert run_1.wait() == 0
+
+
+if __name__ == "__main__":
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/tune/tests/test_placeholder.py

@@ -0,0 +1,316 @@
+import unittest
+
+import numpy as np
+
+from ray import tune
+from ray.tune.impl.placeholder import (
+    _FunctionResolver,
+    _RefResolver,
+    create_resolvers_map,
+    inject_placeholders,
+    resolve_placeholders,
+)
+from ray.tune.search.sample import Float, Integer
+
+
+class Dummy:
+    def __init__(self, value):
+        self.value = value
+
+
+class PlaceholderTest(unittest.TestCase):
+    def testNotReplaced(self):
+        config = {
+            "param1": "ok",
+            "param2": ["not ok", tune.grid_search(["ok", "not ok"])],
+            "param3": {
+                "param4": tune.choice(["ok", "not ok"]),
+            },
+        }
+
+        replaced = create_resolvers_map()
+        config = inject_placeholders(config, replaced)
+
+        # Primitive typed choices are not replaced.
+        self.assertEqual(config["param2"][1]["grid_search"], ["ok", "not ok"])
+        self.assertEqual(config["param3"]["param4"].categories, ["ok", "not ok"])
+
+    def testGridSearch(self):
+        config = {
+            "param1": "ok",
+            "param2": ["not ok", tune.grid_search(["ok", Dummy("not ok")])],
+            "param3": {
+                "param4": tune.grid_search([Dummy("ok"), "not ok"]),
+            },
+        }
+
+        replaced = create_resolvers_map()
+        config = inject_placeholders(config, replaced)
+
+        self.assertEqual(
+            config["param2"][1]["grid_search"],
+            ["ok", (_RefResolver.TOKEN, "1870fa9b")],
+        )
+        self.assertEqual(
+            config["param3"]["param4"]["grid_search"],
+            [(_RefResolver.TOKEN, "8515e998"), "not ok"],
+        )
+
+        # Pretend we picked a choice from the grid searches.
+        config["param2"][1] = (_RefResolver.TOKEN, "1870fa9b")
+        config["param3"]["param4"] = "not ok"
+
+        resolve_placeholders(config, replaced)
+
+        self.assertEqual(config["param2"][1].value, "not ok")
+        self.assertEqual(config["param3"]["param4"], "not ok")
+
+    def testCategorical(self):
+        config = {
+            "param1": "ok",
+            "param2": ["not ok", tune.choice([Dummy("ok"), "not ok"])],
+            "param3": {
+                "param4": tune.choice([Dummy("ok"), "not ok"]),
+            },
+        }
+
+        replaced = create_resolvers_map()
+        config = inject_placeholders(config, replaced)
+
+        self.assertEqual(
+            config["param2"][1].categories,
+            [(_RefResolver.TOKEN, "ec0e030c"), "not ok"],
+        )
+        self.assertEqual(
+            config["param3"]["param4"].categories,
+            [(_RefResolver.TOKEN, "8515e998"), "not ok"],
+        )
+
+        # Pretend we picked a choice from the categoricals.
+        config["param2"][1] = (_RefResolver.TOKEN, "ec0e030c")
+        config["param3"]["param4"] = "not ok"
+
+        resolve_placeholders(config, replaced)
+
+        self.assertEqual(config["param2"][1].value, "ok")
+        self.assertEqual(config["param3"]["param4"], "not ok")
+
+    def _testNonSearchSpaceRef(self, value):
+        """Tests that non-primitives (numpy, lambda fn) get replaced by a reference."""
+        config = {"param": tune.choice([value, "other", value])}
+
+        replaced = create_resolvers_map()
+        config = inject_placeholders(config, replaced)
+
+        self.assertEqual(
+            config["param"].categories,
+            [
+                (_RefResolver.TOKEN, "ec15c422"),
+                "other",
+                (_RefResolver.TOKEN, "3c7edff5"),
+            ],
+        )
+
+    def testNumpyToRef(self):
+        self._testNonSearchSpaceRef(np.arange(10))
+
+    def testLambdaToRef(self):
+        self._testNonSearchSpaceRef(lambda x: x)
+
+    def testFunction(self):
+        config = {
+            "param1": "ok",
+            "param2": ["not ok", tune.sample_from(lambda: "not ok")],
+            # Both lambdas, either taking spec or config, should work.
+            "param3": {
+                "param4": tune.sample_from(lambda spec: spec["config"]["param1"]),
+            },
+            "param4": {
+                "param4": tune.sample_from(lambda config: config["param1"]),
+            },
+            # Make sure dot notation also works with spec passed in.
+            "param5": {
+                "param4": tune.sample_from(lambda spec: spec.config["param1"]),
+            },
+        }
+
+        replaced = create_resolvers_map()
+        config = inject_placeholders(config, replaced)
+
+        self.assertEqual(config["param2"][1][0], _FunctionResolver.TOKEN)
+        self.assertEqual(config["param3"]["param4"][0], _FunctionResolver.TOKEN)
+        self.assertEqual(config["param4"]["param4"][0], _FunctionResolver.TOKEN)
+        self.assertEqual(config["param5"]["param4"][0], _FunctionResolver.TOKEN)
+
+        resolve_placeholders(config, replaced)
+
+        self.assertEqual(config["param2"][1], "not ok")
+        self.assertEqual(config["param3"]["param4"], "ok")
+        self.assertEqual(config["param4"]["param4"], "ok")
+        self.assertEqual(config["param5"]["param4"], "ok")
+
+    def testRefValue(self):
+        config = {
+            "param1": "ok",
+            "param2": ["not ok", Dummy("ok")],
+            "param3": {
+                "param4": Dummy("not ok"),
+            },
+        }
+
+        replaced = create_resolvers_map()
+        config = inject_placeholders(config, replaced)
+
+        self.assertEqual(config["param2"][1][0], _RefResolver.TOKEN)
+        self.assertEqual(config["param3"]["param4"][0], _RefResolver.TOKEN)
+
+        resolve_placeholders(config, replaced)
+
+        self.assertEqual(config["param2"][1].value, "ok")
+        self.assertEqual(config["param3"]["param4"].value, "not ok")
+
+    def testTuple(self):
+        class Dummy:
+            def __init__(self, value):
+                self.value = value
+
+        config = {
+            "param1": ("ok", "not ok"),
+            "param2": ["not ok", (1, Dummy("ok"))],
+            "param3": {
+                "param4": (1, [2, Dummy("not ok")], 3),
+            },
+        }
+
+        replaced = create_resolvers_map()
+        config = inject_placeholders(config, replaced)
+
+        self.assertTrue(isinstance(config["param1"], tuple))
+        self.assertEqual(config["param1"], ("ok", "not ok"))
+        self.assertTrue(isinstance(config["param2"][1], tuple))
+        self.assertTrue(isinstance(config["param3"]["param4"], tuple))
+
+        resolve_placeholders(config, replaced)
+
+        self.assertTrue(isinstance(config["param2"][1], tuple))
+        self.assertEqual(config["param2"][1][1].value, "ok")
+        self.assertTrue(isinstance(config["param3"]["param4"], tuple))
+        self.assertEqual(config["param3"]["param4"][1][1].value, "not ok")
+
+    def testOtherDomains(self):
+        config = {
+            "param1": tune.uniform(0, 1),
+            "param2": tune.randint(2, 3),
+            "param3": tune.qrandn(0, 1, 0.1),
+        }
+
+        replaced = create_resolvers_map()
+        config = inject_placeholders(config, replaced)
+
+        # Normal params are not replaced.
+        self.assertTrue(isinstance(config["param1"], Float))
+        self.assertTrue(isinstance(config["param2"], Integer))
+        self.assertTrue(isinstance(config["param3"], Float))
+
+    def testPointToEval(self):
+        config = {
+            "param1": "ok",
+            "param2": ["not ok", tune.choice([Dummy("ok"), "not ok"])],
+            "param3": {
+                "param4": tune.sample_from(lambda spec: spec["config"]["param1"]),
+            },
+        }
+
+        replaced = create_resolvers_map()
+        config = inject_placeholders(config, replaced)
+
+        # Normal params are not replaced.
+        self.assertEqual(
+            config["param2"][1].categories,
+            [(_RefResolver.TOKEN, "ec0e030c"), "not ok"],
+        )
+        self.assertEqual(
+            config["param3"]["param4"], (_FunctionResolver.TOKEN, "134aff3a")
+        )
+
+        # Now, say we manually resolved the placeholders based on
+        # points_to_evaluate.
+        config["param2"][1] = "not_ok"
+        config["param3"]["param4"] = "ok"
+
+        resolve_placeholders(config, replaced)
+
+        # Params stays the same.
+        self.assertEqual(config["param2"][1], "not_ok")
+        self.assertEqual(config["param3"]["param4"], "ok")
+
+    def testSimpleNestedSearchSpaces(self):
+        config = {
+            "param1": "ok",
+            "param2": tune.choice(
+                [
+                    tune.choice([Dummy(1), 2, 3]),
+                    tune.uniform(5, 6),
+                ]
+            ),
+        }
+
+        replaced = create_resolvers_map()
+        config = inject_placeholders(config, replaced)
+
+        # Manually resolve. Select the Dummy value.
+        config["param2"] = (_RefResolver.TOKEN, "6f33af83")
+
+        resolve_placeholders(config, replaced)
+
+        self.assertEqual(config["param2"].value, 1)
+
+    def testSimpleNestedSearchSpaces2(self):
+        config = {
+            "param1": "ok",
+            "param2": tune.choice(
+                [
+                    (None, Dummy(1), None),
+                    (Dummy(2), None, None),
+                    (None, None, Dummy(3)),
+                ]
+            ),
+        }
+
+        replaced = create_resolvers_map()
+        config = inject_placeholders(config, replaced)
+
+        # Manually resolve. Select the Dummy value.
+        config["param2"] = (None, None, (_RefResolver.TOKEN, "a1433750"))
+
+        resolve_placeholders(config, replaced)
+
+        self.assertEqual(config["param2"][2].value, 3)
+
+    def testResolveFunctionAfterRef(self):
+        config = {
+            "param1": "ok",
+            "param2": tune.choice([Dummy("ok"), "not ok"]),
+            "param3": {
+                "param4": tune.sample_from(lambda config: config["param2"]),
+            },
+        }
+
+        replaced = create_resolvers_map()
+        config = inject_placeholders(config, replaced)
+
+        # Manually resolve param2.
+        config["param2"] = (_RefResolver.TOKEN, "07cb6238")
+
+        resolve_placeholders(config, replaced)
+
+        # param3.param4 should get the same value as resolved param2.
+        self.assertEqual(config["param3"]["param4"].value, "ok")
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/tune/tests/test_progress_reporter.py

@@ -0,0 +1,862 @@
+import collections
+import os
+import unittest
+from unittest.mock import MagicMock, Mock, patch
+
+import numpy as np
+import pytest
+import regex as re
+
+from ray import tune
+from ray._common.test_utils import run_string_as_driver
+from ray.tune.experiment.trial import Trial
+from ray.tune.progress_reporter import (
+    CLIReporter,
+    JupyterNotebookReporter,
+    ProgressReporter,
+    TuneReporterBase,
+    _best_trial_str,
+    _detect_reporter,
+    _fair_filter_trials,
+    _max_len,
+    _time_passed_str,
+    _trial_progress_str,
+)
+from ray.tune.result import AUTO_RESULT_KEYS
+
+EXPECTED_RESULT_1 = """Result logdir: /foo
+Number of trials: 5 (1 PENDING, 3 RUNNING, 1 TERMINATED)
++--------------+------------+-------+-----+-----+------------+
+|   Trial name | status     | loc   |   a |   b |   metric_1 |
+|--------------+------------+-------+-----+-----+------------|
+|        00002 | RUNNING    | here  |   2 |   4 |        1   |
+|        00001 | PENDING    | here  |   1 |   2 |        0.5 |
+|        00000 | TERMINATED | here  |   0 |   0 |        0   |
++--------------+------------+-------+-----+-----+------------+
+... 2 more trials not shown (2 RUNNING)"""
+
+EXPECTED_RESULT_2 = """Result logdir: /foo
+Number of trials: 5 (1 PENDING, 3 RUNNING, 1 TERMINATED)
++--------------+------------+-------+-----+-----+---------+---------+
+|   Trial name | status     | loc   |   a |   b |   n/k/0 |   n/k/1 |
+|--------------+------------+-------+-----+-----+---------+---------|
+|        00002 | RUNNING    | here  |   2 |   4 |       2 |       4 |
+|        00003 | RUNNING    | here  |   3 |   6 |       3 |       6 |
+|        00004 | RUNNING    | here  |   4 |   8 |       4 |       8 |
+|        00001 | PENDING    | here  |   1 |   2 |       1 |       2 |
+|        00000 | TERMINATED | here  |   0 |   0 |       0 |       0 |
++--------------+------------+-------+-----+-----+---------+---------+"""
+
+EXPECTED_RESULT_3 = """Result logdir: /foo
+Number of trials: 5 (1 PENDING, 3 RUNNING, 1 TERMINATED)
++--------------+------------+-------+-----+-----------+------------+
+|   Trial name | status     | loc   |   A |   NestSub |   Metric 2 |
+|--------------+------------+-------+-----+-----------+------------|
+|        00002 | RUNNING    | here  |   2 |       1   |       0.5  |
+|        00001 | PENDING    | here  |   1 |       0.5 |       0.25 |
+|        00000 | TERMINATED | here  |   0 |       0   |       0    |
++--------------+------------+-------+-----+-----------+------------+
+... 2 more trials not shown (2 RUNNING)"""
+
+EXPECTED_RESULT_4 = """Result logdir: /foo
+Number of trials: 5 (1 PENDING, 3 RUNNING, 1 TERMINATED)
++--------------+------------+-------+-----+-----+------------+
+|   Trial name | status     | loc   |   a |   b |   metric_1 |
+|--------------+------------+-------+-----+-----+------------|
+|        00002 | RUNNING    | here  |   2 |   4 |        1   |
+|        00003 | RUNNING    | here  |   3 |   6 |        1.5 |
+|        00004 | RUNNING    | here  |   4 |   8 |        2   |
+|        00001 | PENDING    | here  |   1 |   2 |        0.5 |
+|        00000 | TERMINATED | here  |   0 |   0 |        0   |
++--------------+------------+-------+-----+-----+------------+"""
+
+END_TO_END_COMMAND = """
+import ray
+from ray import tune
+from ray.tune.experiment.trial import _Location
+from ray.tune.progress_reporter import _get_trial_location
+from unittest.mock import patch
+
+
+def mock_get_trial_location(trial, result):
+    location = _get_trial_location(trial, result)
+    if location.pid:
+        return _Location("123.123.123.123", "1")
+    return location
+
+
+with patch("ray.tune.progress_reporter._get_trial_location",
+           mock_get_trial_location):
+    reporter = tune.progress_reporter.CLIReporter(metric_columns=["done"])
+
+    def f(config):
+        return {"done": True}
+
+    ray.init(num_cpus=1)
+    tune.run_experiments(
+        {
+            "one": {
+                "run": f,
+                "config": {
+                    "a": tune.grid_search(list(range(10))),
+                },
+            },
+            "two": {
+                "run": f,
+                "config": {
+                    "b": tune.grid_search(list(range(10))),
+                },
+            },
+            "three": {
+                "run": f,
+                "config": {
+                    "c": tune.grid_search(list(range(10))),
+                },
+            },
+        },
+        verbose=3,
+        progress_reporter=reporter)"""
+
+EXPECTED_END_TO_END_START = """Number of trials: 30/30 (29 PENDING, 1 RUNNING)
++---------------+----------+-------------------+-----+-----+
+| Trial name    | status   | loc               |   a |   b |
+|---------------+----------+-------------------+-----+-----|
+| f_xxxxx_00000 | RUNNING  | 123.123.123.123:1 |   0 |     |
+| f_xxxxx_00001 | PENDING  |                   |   1 |     |"""
+
+EXPECTED_END_TO_END_END = """Number of trials: 30/30 (30 TERMINATED)
++---------------+------------+-------------------+-----+-----+-----+--------+
+| Trial name    | status     | loc               |   a |   b |   c | done   |
+|---------------+------------+-------------------+-----+-----+-----+--------|
+| f_xxxxx_00000 | TERMINATED | 123.123.123.123:1 |   0 |     |     | True   |
+| f_xxxxx_00001 | TERMINATED | 123.123.123.123:1 |   1 |     |     | True   |
+| f_xxxxx_00002 | TERMINATED | 123.123.123.123:1 |   2 |     |     | True   |
+| f_xxxxx_00003 | TERMINATED | 123.123.123.123:1 |   3 |     |     | True   |
+| f_xxxxx_00004 | TERMINATED | 123.123.123.123:1 |   4 |     |     | True   |
+| f_xxxxx_00005 | TERMINATED | 123.123.123.123:1 |   5 |     |     | True   |
+| f_xxxxx_00006 | TERMINATED | 123.123.123.123:1 |   6 |     |     | True   |
+| f_xxxxx_00007 | TERMINATED | 123.123.123.123:1 |   7 |     |     | True   |
+| f_xxxxx_00008 | TERMINATED | 123.123.123.123:1 |   8 |     |     | True   |
+| f_xxxxx_00009 | TERMINATED | 123.123.123.123:1 |   9 |     |     | True   |
+| f_xxxxx_00010 | TERMINATED | 123.123.123.123:1 |     |   0 |     | True   |
+| f_xxxxx_00011 | TERMINATED | 123.123.123.123:1 |     |   1 |     | True   |
+| f_xxxxx_00012 | TERMINATED | 123.123.123.123:1 |     |   2 |     | True   |
+| f_xxxxx_00013 | TERMINATED | 123.123.123.123:1 |     |   3 |     | True   |
+| f_xxxxx_00014 | TERMINATED | 123.123.123.123:1 |     |   4 |     | True   |
+| f_xxxxx_00015 | TERMINATED | 123.123.123.123:1 |     |   5 |     | True   |
+| f_xxxxx_00016 | TERMINATED | 123.123.123.123:1 |     |   6 |     | True   |
+| f_xxxxx_00017 | TERMINATED | 123.123.123.123:1 |     |   7 |     | True   |
+| f_xxxxx_00018 | TERMINATED | 123.123.123.123:1 |     |   8 |     | True   |
+| f_xxxxx_00019 | TERMINATED | 123.123.123.123:1 |     |   9 |     | True   |
+| f_xxxxx_00020 | TERMINATED | 123.123.123.123:1 |     |     |   0 | True   |
+| f_xxxxx_00021 | TERMINATED | 123.123.123.123:1 |     |     |   1 | True   |
+| f_xxxxx_00022 | TERMINATED | 123.123.123.123:1 |     |     |   2 | True   |
+| f_xxxxx_00023 | TERMINATED | 123.123.123.123:1 |     |     |   3 | True   |
+| f_xxxxx_00024 | TERMINATED | 123.123.123.123:1 |     |     |   4 | True   |
+| f_xxxxx_00025 | TERMINATED | 123.123.123.123:1 |     |     |   5 | True   |
+| f_xxxxx_00026 | TERMINATED | 123.123.123.123:1 |     |     |   6 | True   |
+| f_xxxxx_00027 | TERMINATED | 123.123.123.123:1 |     |     |   7 | True   |
+| f_xxxxx_00028 | TERMINATED | 123.123.123.123:1 |     |     |   8 | True   |
+| f_xxxxx_00029 | TERMINATED | 123.123.123.123:1 |     |     |   9 | True   |
++---------------+------------+-------------------+-----+-----+-----+--------+"""  # noqa
+
+EXPECTED_END_TO_END_AC = """Number of trials: 30/30 (30 TERMINATED)
++---------------+------------+-------+-----+-----+-----+
+| Trial name    | status     | loc   |   a |   b |   c |
+|---------------+------------+-------+-----+-----+-----|
+| f_xxxxx_00000 | TERMINATED |       |   0 |     |     |
+| f_xxxxx_00001 | TERMINATED |       |   1 |     |     |
+| f_xxxxx_00002 | TERMINATED |       |   2 |     |     |
+| f_xxxxx_00003 | TERMINATED |       |   3 |     |     |
+| f_xxxxx_00004 | TERMINATED |       |   4 |     |     |
+| f_xxxxx_00005 | TERMINATED |       |   5 |     |     |
+| f_xxxxx_00006 | TERMINATED |       |   6 |     |     |
+| f_xxxxx_00007 | TERMINATED |       |   7 |     |     |
+| f_xxxxx_00008 | TERMINATED |       |   8 |     |     |
+| f_xxxxx_00009 | TERMINATED |       |   9 |     |     |
+| f_xxxxx_00010 | TERMINATED |       |     |   0 |     |
+| f_xxxxx_00011 | TERMINATED |       |     |   1 |     |
+| f_xxxxx_00012 | TERMINATED |       |     |   2 |     |
+| f_xxxxx_00013 | TERMINATED |       |     |   3 |     |
+| f_xxxxx_00014 | TERMINATED |       |     |   4 |     |
+| f_xxxxx_00015 | TERMINATED |       |     |   5 |     |
+| f_xxxxx_00016 | TERMINATED |       |     |   6 |     |
+| f_xxxxx_00017 | TERMINATED |       |     |   7 |     |
+| f_xxxxx_00018 | TERMINATED |       |     |   8 |     |
+| f_xxxxx_00019 | TERMINATED |       |     |   9 |     |
+| f_xxxxx_00020 | TERMINATED |       |     |     |   0 |
+| f_xxxxx_00021 | TERMINATED |       |     |     |   1 |
+| f_xxxxx_00022 | TERMINATED |       |     |     |   2 |
+| f_xxxxx_00023 | TERMINATED |       |     |     |   3 |
+| f_xxxxx_00024 | TERMINATED |       |     |     |   4 |
+| f_xxxxx_00025 | TERMINATED |       |     |     |   5 |
+| f_xxxxx_00026 | TERMINATED |       |     |     |   6 |
+| f_xxxxx_00027 | TERMINATED |       |     |     |   7 |
+| f_xxxxx_00028 | TERMINATED |       |     |     |   8 |
+| f_xxxxx_00029 | TERMINATED |       |     |     |   9 |
++---------------+------------+-------+-----+-----+-----+"""
+
+EXPECTED_BEST_1 = (
+    "Current best trial: 00001 with metric_1=0.5 and "
+    "parameters={'a': 1, 'b': 2, 'n': {'k': [1, 2]}}"
+)
+
+EXPECTED_BEST_2 = "Current best trial: 00004 with metric_1=2.0 and parameters={'a': 4}"
+
+EXPECTED_SORT_RESULT_UNSORTED = """Number of trials: 5 (1 PENDING, 1 RUNNING, 3 TERMINATED)
++--------------+------------+-------+-----+------------+
+|   Trial name | status     | loc   |   a |   metric_1 |
+|--------------+------------+-------+-----+------------|
+|        00004 | RUNNING    | here  |   4 |            |
+|        00003 | PENDING    | here  |   3 |            |
+|        00000 | TERMINATED | here  |   0 |        0.3 |
+|        00001 | TERMINATED | here  |   1 |        0.2 |
++--------------+------------+-------+-----+------------+
+... 1 more trials not shown (1 TERMINATED)"""
+
+EXPECTED_SORT_RESULT_ASC = """Number of trials: 5 (1 PENDING, 1 RUNNING, 3 TERMINATED)
++--------------+------------+-------+-----+------------+
+|   Trial name | status     | loc   |   a |   metric_1 |
+|--------------+------------+-------+-----+------------|
+|        00004 | RUNNING    | here  |   4 |            |
+|        00003 | PENDING    | here  |   3 |            |
+|        00001 | TERMINATED | here  |   1 |        0.2 |
+|        00000 | TERMINATED | here  |   0 |        0.3 |
++--------------+------------+-------+-----+------------+
+... 1 more trials not shown (1 TERMINATED)"""
+
+EXPECTED_NESTED_SORT_RESULT = """Number of trials: 5 (1 PENDING, 1 RUNNING, 3 TERMINATED)
++--------------+------------+-------+-----+-------------------+
+|   Trial name | status     | loc   |   a |   nested/metric_2 |
+|--------------+------------+-------+-----+-------------------|
+|        00004 | RUNNING    | here  |   4 |                   |
+|        00003 | PENDING    | here  |   3 |                   |
+|        00001 | TERMINATED | here  |   1 |               0.2 |
+|        00000 | TERMINATED | here  |   0 |               0.3 |
++--------------+------------+-------+-----+-------------------+
+... 1 more trials not shown (1 TERMINATED)"""
+
+EXPECTED_SORT_RESULT_DESC = """Number of trials: 5 (1 PENDING, 1 RUNNING, 3 TERMINATED)
++--------------+------------+-------+-----+------------+
+|   Trial name | status     | loc   |   a |   metric_1 |
+|--------------+------------+-------+-----+------------|
+|        00004 | RUNNING    | here  |   4 |            |
+|        00003 | PENDING    | here  |   3 |            |
+|        00002 | TERMINATED | here  |   2 |        0.4 |
+|        00000 | TERMINATED | here  |   0 |        0.3 |
++--------------+------------+-------+-----+------------+
+... 1 more trials not shown (1 TERMINATED)"""
+
+VERBOSE_EXP_OUT_1 = "Number of trials: 3/3 (2 PENDING, 1 RUNNING)"
+VERBOSE_EXP_OUT_2 = "Number of trials: 3/3 (3 TERMINATED)"
+
+VERBOSE_TRIAL_NORM_1 = (
+    "Trial train_fn_xxxxx_00000 reported acc=5 "
+    "with parameters={'do': 'complete'}. This trial completed.\n"
+)
+
+# NOTE: We use Regex for `VERBOSE_TRIAL_NORM_2` to make the test deterministic.
+# `"Trial train_fn_xxxxx_00001 reported..."` and
+# `"Trial train_fn_xxxxx_00001 completed..."`
+# are printed in separate calls. Sometimes, a status update is printed between the
+# calls. For more information, see #29693.
+VERBOSE_TRIAL_NORM_2_PATTERN = (
+    r"Trial train_fn_xxxxx_00001 reported _metric=6 "
+    r"with parameters=\{'do': 'once'\}\.\n"
+    r"(?s).*"
+    r"Trial train_fn_xxxxx_00001 completed\. Last result: _metric=6\n"
+)
+
+VERBOSE_TRIAL_NORM_3 = (
+    "Trial train_fn_xxxxx_00002 reported acc=7 with parameters={'do': 'twice'}.\n"
+)
+
+VERBOSE_TRIAL_NORM_4 = (
+    "Trial train_fn_xxxxx_00002 reported acc=8 "
+    "with parameters={'do': 'twice'}. This trial completed.\n"
+)
+
+VERBOSE_TRIAL_WITH_ONCE_RESULT = "Result for train_fn_xxxxx_00001"
+VERBOSE_TRIAL_WITH_ONCE_COMPLETED = "Trial train_fn_xxxxx_00001 completed."
+
+VERBOSE_TRIAL_DETAIL = """+-------------------+----------+-------------------+----------+
+| Trial name        | status   | loc               | do       |
+|-------------------+----------+-------------------+----------|
+| train_fn_xxxxx_00000 | RUNNING  | 123.123.123.123:1 | complete |"""
+
+VERBOSE_CMD = """import ray.tune
+import random
+import numpy as np
+import time
+from ray.tune.experiment.trial import _Location
+from ray.tune.progress_reporter import _get_trial_location
+from unittest.mock import patch
+
+
+def mock_get_trial_location(trial, result):
+    location = _get_trial_location(trial, result)
+    if location.pid:
+        return _Location("123.123.123.123", "1")
+    return location
+
+
+def train_fn(config):
+    if config["do"] == "complete":
+        time.sleep(0.1)
+        ray.tune.report(dict(acc=5, done=True))
+    elif config["do"] == "once":
+        time.sleep(0.5)
+        return 6
+    else:
+        time.sleep(1.0)
+        ray.tune.report(dict(acc=7))
+        ray.tune.report(dict(acc=8))
+
+random.seed(1234)
+np.random.seed(1234)
+
+
+with patch("ray.tune.progress_reporter._get_trial_location",
+           mock_get_trial_location):
+    ray.tune.run(
+        train_fn,
+        config={
+            "do": ray.tune.grid_search(["complete", "once", "twice"])
+        },"""
+
+# Add "verbose=3)" etc
+
+
+class ProgressReporterTest(unittest.TestCase):
+    def setUp(self) -> None:
+        os.environ["TUNE_MAX_PENDING_TRIALS_PG"] = "auto"
+        os.environ["RAY_AIR_NEW_OUTPUT"] = "0"
+
+    def mock_trial(self, status, i):
+        mock = MagicMock()
+        mock.status = status
+        mock.trial_id = "%05d" % i
+        return mock
+
+    def testFairFilterTrials(self):
+        """Tests that trials are represented fairly."""
+        trials_by_state = collections.defaultdict(list)
+        # States for which trials are under and overrepresented
+        states_under = (Trial.PAUSED, Trial.ERROR)
+        states_over = (Trial.PENDING, Trial.RUNNING, Trial.TERMINATED)
+
+        max_trials = 13
+        num_trials_under = 2  # num of trials for each underrepresented state
+        num_trials_over = 10  # num of trials for each overrepresented state
+
+        i = 0
+        for state in states_under:
+            for _ in range(num_trials_under):
+                trials_by_state[state].append(self.mock_trial(state, i))
+                i += 1
+        for state in states_over:
+            for _ in range(num_trials_over):
+                trials_by_state[state].append(self.mock_trial(state, i))
+                i += 1
+
+        filtered_trials_by_state = _fair_filter_trials(
+            trials_by_state, max_trials=max_trials
+        )
+        for state in trials_by_state:
+            if state in states_under:
+                expected_num_trials = num_trials_under
+            else:
+                expected_num_trials = (
+                    max_trials - num_trials_under * len(states_under)
+                ) / len(states_over)
+            state_trials = filtered_trials_by_state[state]
+            self.assertEqual(len(state_trials), expected_num_trials)
+            # Make sure trials are sorted newest-first within state.
+            for i in range(len(state_trials) - 1):
+                assert state_trials[i].trial_id < state_trials[i + 1].trial_id
+
+    def testAddMetricColumn(self):
+        """Tests edge cases of add_metric_column."""
+
+        # Test list-initialized metric columns.
+        reporter = CLIReporter(metric_columns=["foo", "bar"])
+        with self.assertRaises(ValueError):
+            reporter.add_metric_column("bar")
+
+        with self.assertRaises(ValueError):
+            reporter.add_metric_column("baz", "qux")
+
+        reporter.add_metric_column("baz")
+        self.assertIn("baz", reporter._metric_columns)
+
+        # Test default-initialized (dict) metric columns.
+        reporter = CLIReporter()
+        reporter.add_metric_column("foo", "bar")
+        self.assertIn("foo", reporter._metric_columns)
+
+    def testInfer(self):
+        reporter = CLIReporter()
+        test_result = dict(foo_result=1, baz_result=4123, bar_result="testme")
+
+        def test(config):
+            for i in range(3):
+                tune.report(test_result)
+
+        analysis = tune.run(test, num_samples=3, verbose=3)
+        all_trials = analysis.trials
+        inferred_results = reporter._infer_user_metrics(all_trials)
+
+        for metric in inferred_results:
+            self.assertNotIn(metric, AUTO_RESULT_KEYS)
+            self.assertTrue(metric in test_result)
+
+        class TestReporter(CLIReporter):
+            _output = []
+
+            def __init__(self, *args, **kwargs):
+                super().__init__(*args, **kwargs)
+                self._max_report_freqency = 0
+
+            def report(self, *args, **kwargs):
+                progress_str = self._progress_str(*args, **kwargs)
+                self._output.append(progress_str)
+
+        reporter = TestReporter()
+        analysis = tune.run(test, num_samples=3, progress_reporter=reporter, verbose=3)
+        found = {k: False for k in test_result}
+        for output in reporter._output:
+            for key in test_result:
+                if key in output:
+                    found[key] = True
+        assert found["foo_result"]
+        assert found["baz_result"]
+        assert not found["bar_result"]
+
+    def testProgressStr(self):
+        trials = []
+        for i in range(5):
+            t = Mock()
+            if i == 0:
+                t.status = "TERMINATED"
+            elif i == 1:
+                t.status = "PENDING"
+            else:
+                t.status = "RUNNING"
+            t.trial_id = "%05d" % i
+            t.local_experiment_path = "/foo"
+            t.temporary_state = Mock()
+            t.temporary_state.location = "here"
+            t.config = {"a": i, "b": i * 2, "n": {"k": [i, 2 * i]}}
+            t.evaluated_params = {"a": i, "b": i * 2, "n/k/0": i, "n/k/1": 2 * i}
+            t.last_result = {
+                "config": {"a": i, "b": i * 2, "n": {"k": [i, 2 * i]}},
+                "metric_1": i / 2,
+                "metric_2": i / 4,
+                "nested": {"sub": i / 2},
+            }
+            t.__str__ = lambda self: self.trial_id
+            trials.append(t)
+        # One metric, two parameters
+        prog1 = _trial_progress_str(
+            trials, ["metric_1"], ["a", "b"], fmt="psql", max_rows=3, force_table=True
+        )
+        print(prog1)
+        assert prog1 == EXPECTED_RESULT_1
+
+        # No metric, all parameters
+        prog2 = _trial_progress_str(
+            trials, [], None, fmt="psql", max_rows=None, force_table=True
+        )
+        print(prog2)
+        assert prog2 == EXPECTED_RESULT_2
+
+        # Two metrics, one parameter, all with custom representation
+        prog3 = _trial_progress_str(
+            trials,
+            {"nested/sub": "NestSub", "metric_2": "Metric 2"},
+            {"a": "A"},
+            fmt="psql",
+            max_rows=3,
+            force_table=True,
+        )
+        print(prog3)
+        assert prog3 == EXPECTED_RESULT_3
+
+        # Current best trial
+        best1 = _best_trial_str(trials[1], "metric_1")
+        assert best1 == EXPECTED_BEST_1
+
+    def testBestTrialStr(self):
+        """Assert that custom nested parameter columns are printed correctly"""
+        config = {"nested": {"conf": "nested_value"}, "toplevel": "toplevel_value"}
+
+        trial = Trial("", config=config, stub=True)
+        trial.run_metadata.last_result = {
+            "metric": 1,
+            "config": config,
+            "nested": {"metric": 2},
+        }
+
+        result = _best_trial_str(trial, "metric")
+        self.assertIn("nested_value", result)
+
+        result = _best_trial_str(trial, "metric", parameter_columns=["nested/conf"])
+        self.assertIn("nested_value", result)
+
+        # Test that this works with a nested metric
+        result = _best_trial_str(
+            trial, "nested/metric", parameter_columns=["nested/conf"]
+        )
+        self.assertIn("nested_value", result)
+
+    def testBestTrialZero(self):
+        trial1 = Trial("", config={}, stub=True)
+        trial1.run_metadata.last_result = {"metric": 7, "config": {}}
+
+        trial2 = Trial("", config={}, stub=True)
+        trial2.run_metadata.last_result = {"metric": 0, "config": {}}
+
+        trial3 = Trial("", config={}, stub=True)
+        trial3.run_metadata.last_result = {"metric": 2, "config": {}}
+
+        reporter = TuneReporterBase(metric="metric", mode="min")
+        best_trial, metric = reporter._current_best_trial([trial1, trial2, trial3])
+        assert best_trial == trial2
+
+    def testBestTrialNan(self):
+        trial1 = Trial("", config={}, stub=True)
+        trial1.run_metadata.last_result = {"metric": np.nan, "config": {}}
+
+        trial2 = Trial("", config={}, stub=True)
+        trial2.run_metadata.last_result = {"metric": 0, "config": {}}
+
+        trial3 = Trial("", config={}, stub=True)
+        trial3.run_metadata.last_result = {"metric": 2, "config": {}}
+
+        reporter = TuneReporterBase(metric="metric", mode="min")
+        best_trial, metric = reporter._current_best_trial([trial1, trial2, trial3])
+        assert best_trial == trial2
+
+        trial1 = Trial("", config={}, stub=True)
+        trial1.run_metadata.last_result = {"metric": np.nan, "config": {}}
+
+        trial2 = Trial("", config={}, stub=True)
+        trial2.run_metadata.last_result = {"metric": 0, "config": {}}
+
+        trial3 = Trial("", config={}, stub=True)
+        trial3.run_metadata.last_result = {"metric": 2, "config": {}}
+
+        reporter = TuneReporterBase(metric="metric", mode="max")
+        best_trial, metric = reporter._current_best_trial([trial1, trial2, trial3])
+        assert best_trial == trial3
+
+    def testTimeElapsed(self):
+        # Sun Feb 7 14:18:40 2016 -0800
+        # (time of the first Ray commit)
+        time_start = 1454825920
+        time_now = (
+            time_start
+            + 1 * 60 * 60  # 1 hour
+            + 31 * 60  # 31 minutes
+            + 22  # 22 seconds
+        )  # time to second commit
+
+        # Local timezone output can be tricky, so we don't check the
+        # day and the hour in this test.
+        output = _time_passed_str(time_start, time_now)
+        self.assertIn("Current time: 2016-02-", output)
+        self.assertIn(":50:02 (running for 01:31:22.00)", output)
+
+        time_now += 2 * 60 * 60 * 24  # plus two days
+        output = _time_passed_str(time_start, time_now)
+        self.assertIn("Current time: 2016-02-", output)
+        self.assertIn(":50:02 (running for 2 days, 01:31:22.00)", output)
+
+    def testCurrentBestTrial(self):
+        trials = []
+        for i in range(5):
+            t = Mock()
+            t.status = "RUNNING"
+            t.trial_id = "%05d" % i
+            t.local_experiment_path = "/foo"
+            t.temporary_state = Mock()
+            t.temporary_state.location = "here"
+            t.config = {"a": i, "b": i * 2, "n": {"k": [i, 2 * i]}}
+            t.evaluated_params = {"a": i}
+            t.last_result = {"config": {"a": i}, "metric_1": i / 2}
+            t.__str__ = lambda self: self.trial_id
+            trials.append(t)
+
+        class TestReporter(CLIReporter):
+            _output = []
+
+            def __init__(self, *args, **kwargs):
+                super().__init__(*args, **kwargs)
+                self._max_report_freqency = 0
+
+            def report(self, *args, **kwargs):
+                progress_str = self._progress_str(*args, **kwargs)
+                self._output.append(progress_str)
+
+        reporter = TestReporter(mode="max")
+        reporter.report(trials, done=False)
+        assert EXPECTED_BEST_2 in reporter._output[0]
+
+    def testSortByMetric(self):
+        trials = []
+        for i in range(5):
+            t = Mock()
+            if i < 3:
+                t.status = "TERMINATED"
+            elif i == 3:
+                t.status = "PENDING"
+            else:
+                t.status = "RUNNING"
+            t.trial_id = "%05d" % i
+            t.local_experiment_path = "/foo"
+            t.temporary_state = Mock()
+            t.temporary_state.location = "here"
+            t.run_metadata = Mock()
+            t.config = {"a": i}
+            t.evaluated_params = {"a": i}
+            t.last_result = {"config": {"a": i}}
+            t.__str__ = lambda self: self.trial_id
+            trials.append(t)
+        # Set `metric_1` for terminated trails
+        trials[0].last_result["metric_1"] = 0.3
+        trials[0].last_result["nested"] = {"metric_2": 0.3}
+        trials[1].last_result["metric_1"] = 0.2
+        trials[1].last_result["nested"] = {"metric_2": 0.2}
+        trials[2].last_result["metric_1"] = 0.4
+        trials[2].last_result["nested"] = {"metric_2": 0.4}
+
+        class TestReporter(CLIReporter):
+            def __init__(self, *args, **kwargs):
+                super().__init__(*args, **kwargs)
+                self._max_report_freqency = 0
+                self._output = ""
+
+            def report(self, *args, **kwargs):
+                progress_str = self._progress_str(*args, **kwargs)
+                self._output = progress_str
+
+        # Default reporter
+        reporter1 = TestReporter(max_progress_rows=4, mode="max", metric="metric_1")
+        reporter1.report(trials, done=False)
+        assert EXPECTED_SORT_RESULT_UNSORTED in reporter1._output
+
+        # Sort by metric (asc)
+        reporter2 = TestReporter(
+            max_progress_rows=4, mode="min", metric="metric_1", sort_by_metric=True
+        )
+        reporter2.report(trials, done=False)
+        assert EXPECTED_SORT_RESULT_ASC in reporter2._output
+
+        # Sort by metric (desc)
+        reporter3 = TestReporter(
+            max_progress_rows=4, mode="max", metric="metric_1", sort_by_metric=True
+        )
+        reporter3.report(trials, done=False)
+        assert EXPECTED_SORT_RESULT_DESC in reporter3._output
+
+        # Sort by metric when mode is None
+        reporter4 = TestReporter(
+            max_progress_rows=4, metric="metric_1", sort_by_metric=True
+        )
+        reporter4.report(trials, done=False)
+        assert EXPECTED_SORT_RESULT_UNSORTED in reporter4._output
+
+        # Sort by metric when metric is None
+        reporter5 = TestReporter(max_progress_rows=4, mode="max", sort_by_metric=True)
+        reporter5.report(trials, done=False)
+        assert EXPECTED_SORT_RESULT_UNSORTED in reporter5._output
+
+        # Sort by metric when metric is passed using
+        # reporter.setup (called from tune.run)
+        # calling repoter.set_search_properties
+        reporter6 = TestReporter(max_progress_rows=4, sort_by_metric=True)
+        reporter6.set_search_properties(metric="metric_1", mode="max")
+        reporter6.report(trials, done=False)
+        assert EXPECTED_SORT_RESULT_DESC in reporter6._output
+
+        # Sort by nested metric (asc)
+        reporter7 = TestReporter(
+            max_progress_rows=4,
+            mode="min",
+            metric="nested/metric_2",
+            sort_by_metric=True,
+            metric_columns=["nested/metric_2"],
+        )
+        reporter7.report(trials, done=False)
+        assert EXPECTED_NESTED_SORT_RESULT in reporter7._output
+
+    def testEndToEndReporting(self):
+        try:
+            os.environ["_TEST_TUNE_TRIAL_UUID"] = "xxxxx"
+            os.environ["TUNE_MAX_PENDING_TRIALS_PG"] = "100"
+            output = run_string_as_driver(END_TO_END_COMMAND)
+            try:
+                # New execution path is too fast, trials are already terminated
+                if os.environ.get("TUNE_NEW_EXECUTION") == "0":
+                    assert EXPECTED_END_TO_END_START in output
+                assert EXPECTED_END_TO_END_END in output
+                for line in output.splitlines():
+                    if "(raylet)" in line:
+                        assert "Setting" in line, "Unexpected raylet log messages"
+            except Exception:
+                print("*** BEGIN OUTPUT ***")
+                print(output)
+                print("*** END OUTPUT ***")
+                raise
+        finally:
+            del os.environ["_TEST_TUNE_TRIAL_UUID"]
+
+    def testVerboseReporting(self):
+        try:
+            os.environ["_TEST_TUNE_TRIAL_UUID"] = "xxxxx"
+
+            verbose_0_cmd = VERBOSE_CMD + "verbose=0)"
+            output = run_string_as_driver(verbose_0_cmd)
+            try:
+                self.assertNotIn(VERBOSE_EXP_OUT_1, output)
+                self.assertNotIn(VERBOSE_EXP_OUT_2, output)
+                self.assertNotIn(VERBOSE_TRIAL_NORM_1, output)
+                self.assertIsNone(re.search(VERBOSE_TRIAL_NORM_2_PATTERN, output))
+                self.assertNotIn(VERBOSE_TRIAL_NORM_3, output)
+                self.assertNotIn(VERBOSE_TRIAL_NORM_4, output)
+                if os.environ.get("TUNE_NEW_EXECUTION") == "0":
+                    self.assertNotIn(VERBOSE_TRIAL_DETAIL, output)
+            except Exception:
+                print("*** BEGIN OUTPUT ***")
+                print(output)
+                print("*** END OUTPUT ***")
+                raise
+
+            verbose_1_cmd = VERBOSE_CMD + "verbose=1)"
+            output = run_string_as_driver(verbose_1_cmd)
+            try:
+                # New execution path is too fast, trials are already terminated
+                if os.environ.get("TUNE_NEW_EXECUTION") == "0":
+                    self.assertIn(VERBOSE_EXP_OUT_1, output)
+                self.assertIn(VERBOSE_EXP_OUT_2, output)
+                self.assertNotIn(VERBOSE_TRIAL_NORM_1, output)
+                self.assertIsNone(re.search(VERBOSE_TRIAL_NORM_2_PATTERN, output))
+                self.assertNotIn(VERBOSE_TRIAL_NORM_3, output)
+                self.assertNotIn(VERBOSE_TRIAL_NORM_4, output)
+                if os.environ.get("TUNE_NEW_EXECUTION") == "0":
+                    self.assertNotIn(VERBOSE_TRIAL_DETAIL, output)
+            except Exception:
+                print("*** BEGIN OUTPUT ***")
+                print(output)
+                print("*** END OUTPUT ***")
+                raise
+
+            verbose_2_cmd = VERBOSE_CMD + "verbose=2)"
+            output = run_string_as_driver(verbose_2_cmd)
+            try:
+                if os.environ.get("TUNE_NEW_EXECUTION") == "0":
+                    self.assertIn(VERBOSE_EXP_OUT_1, output)
+                self.assertIn(VERBOSE_EXP_OUT_2, output)
+                self.assertIn(VERBOSE_TRIAL_NORM_1, output)
+                self.assertIsNotNone(re.search(VERBOSE_TRIAL_NORM_2_PATTERN, output))
+                self.assertIn(VERBOSE_TRIAL_NORM_3, output)
+                self.assertIn(VERBOSE_TRIAL_NORM_4, output)
+                self.assertNotIn(VERBOSE_TRIAL_DETAIL, output)
+            except Exception:
+                print("*** BEGIN OUTPUT ***")
+                print(output)
+                print("*** END OUTPUT ***")
+                raise
+
+            verbose_3_cmd = VERBOSE_CMD + "verbose=3)"
+            output = run_string_as_driver(verbose_3_cmd)
+            try:
+                if os.environ.get("TUNE_NEW_EXECUTION") == "0":
+                    self.assertIn(VERBOSE_EXP_OUT_1, output)
+                self.assertIn(VERBOSE_EXP_OUT_2, output)
+                self.assertNotIn(VERBOSE_TRIAL_NORM_1, output)
+                self.assertIsNone(re.search(VERBOSE_TRIAL_NORM_2_PATTERN, output))
+                self.assertNotIn(VERBOSE_TRIAL_NORM_3, output)
+                self.assertNotIn(VERBOSE_TRIAL_NORM_4, output)
+                if os.environ.get("TUNE_NEW_EXECUTION") == "0":
+                    self.assertIn(VERBOSE_TRIAL_DETAIL, output)
+                # Check that we don't print duplicate results at the end
+                self.assertTrue(output.count(VERBOSE_TRIAL_WITH_ONCE_RESULT) == 1)
+                self.assertIn(VERBOSE_TRIAL_WITH_ONCE_COMPLETED, output)
+            except Exception:
+                print("*** BEGIN OUTPUT ***")
+                print(output)
+                print("*** END OUTPUT ***")
+                raise
+        finally:
+            del os.environ["_TEST_TUNE_TRIAL_UUID"]
+
+    def testReporterDetection(self):
+        """Test if correct reporter is returned from ``detect_reporter()``"""
+        reporter = _detect_reporter()
+        self.assertTrue(isinstance(reporter, CLIReporter))
+        self.assertFalse(isinstance(reporter, JupyterNotebookReporter))
+
+        with patch("ray.tune.progress_reporter.IS_NOTEBOOK", True):
+            reporter = _detect_reporter()
+            self.assertFalse(isinstance(reporter, CLIReporter))
+            self.assertTrue(isinstance(reporter, JupyterNotebookReporter))
+            trainer_reporter = _detect_reporter(_trainer_api=True)
+            self.assertFalse(isinstance(trainer_reporter, JupyterNotebookReporter))
+            self.assertTrue(isinstance(trainer_reporter, CLIReporter))
+
+    def testProgressReporterAPI(self):
+        class CustomReporter(ProgressReporter):
+            def should_report(self, trials, done=False):
+                return True
+
+            def report(self, trials, done, *sys_info):
+                pass
+
+        tune.run(
+            lambda config: 2,
+            num_samples=1,
+            progress_reporter=CustomReporter(),
+            verbose=3,
+        )
+
+    def testMaxLen(self):
+        trials = []
+        for i in range(5):
+            t = Mock()
+            t.status = "TERMINATED"
+            t.trial_id = "%05d" % i
+            t.local_experiment_path = "/foo"
+            t.temporary_state = Mock()
+            t.temporary_state.location = "here"
+            t.config = {"verylong" * 20: i}
+            t.evaluated_params = {"verylong" * 20: i}
+            t.last_result = {"some_metric": "evenlonger" * 100}
+            t.__str__ = lambda self: self.trial_id
+            trials.append(t)
+
+        progress_str = _trial_progress_str(
+            trials, metric_columns=["some_metric"], force_table=True
+        )
+        assert any(len(row) <= 90 for row in progress_str.split("\n"))
+
+
+def test_max_len():
+    assert (
+        _max_len("some_long_string/even_longer", max_len=28)
+        == "some_long_string/even_longer"
+    )
+    assert _max_len("some_long_string/even_longer", max_len=15) == ".../even_longer"
+
+    assert (
+        _max_len(
+            "19_character_string/19_character_string/too_long", max_len=20, wrap=True
+        )
+        == "...r_string/19_chara\ncter_string/too_long"
+    )
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/tune/tests/test_remote.py

@@ -0,0 +1,109 @@
+import inspect
+import unittest
+from unittest.mock import patch
+
+import ray
+from ray.tune import choice, register_trainable, run, run_experiments
+from ray.tune.experiment import Experiment, Trial
+from ray.tune.result import TIMESTEPS_TOTAL
+from ray.tune.search.hyperopt import HyperOptSearch
+from ray.util.client.ray_client_helpers import ray_start_client_server
+
+
+def train_fn(config):
+    for i in range(100):
+        ray.tune.report(dict(timesteps_total=i))
+
+
+class RemoteTest(unittest.TestCase):
+    def tearDown(self):
+        ray.shutdown()
+
+    def testRemoteRunExperiments(self):
+        register_trainable("f1", train_fn)
+        exp1 = Experiment(
+            **{
+                "name": "foo",
+                "run": "f1",
+            }
+        )
+        [trial] = run_experiments(exp1, _remote=True)
+        self.assertEqual(trial.status, Trial.TERMINATED)
+        self.assertEqual(trial.last_result[TIMESTEPS_TOTAL], 99)
+
+    def testRemoteRun(self):
+        analysis = run(train_fn, _remote=True)
+        [trial] = analysis.trials
+        self.assertEqual(trial.status, Trial.TERMINATED)
+        self.assertEqual(trial.last_result[TIMESTEPS_TOTAL], 99)
+
+    def testRemoteRunArguments(self):
+        def mocked_run(*args, **kwargs):
+            capture_args_kwargs = (args, kwargs)
+            return run(*args, **kwargs), capture_args_kwargs
+
+        with patch("ray.tune.tune.run", mocked_run):
+            analysis, capture_args_kwargs = run(train_fn, _remote=True)
+        args, kwargs = capture_args_kwargs
+        self.assertFalse(args)
+        kwargs.pop("run_or_experiment")
+        kwargs.pop("_remote")
+        kwargs.pop("progress_reporter")  # gets autodetected and set
+
+        default_kwargs = {
+            k: v.default for k, v in inspect.signature(run).parameters.items()
+        }
+        default_kwargs.pop("run_or_experiment")
+        default_kwargs.pop("_remote")
+        default_kwargs.pop("progress_reporter")
+
+        self.assertDictEqual(kwargs, default_kwargs)
+
+    def testRemoteRunWithSearcher(self):
+        analysis = run(
+            train_fn,
+            search_alg=HyperOptSearch(),
+            config={"a": choice(["a", "b"])},
+            metric="timesteps_total",
+            mode="max",
+            _remote=True,
+        )
+        [trial] = analysis.trials
+        self.assertEqual(trial.status, Trial.TERMINATED)
+        self.assertEqual(trial.last_result[TIMESTEPS_TOTAL], 99)
+
+    def testRemoteRunExperimentsInClient(self):
+        ray.init()
+        assert not ray.util.client.ray.is_connected()
+        with ray_start_client_server():
+            assert ray.util.client.ray.is_connected()
+
+            register_trainable("f1", train_fn)
+            exp1 = Experiment(
+                **{
+                    "name": "foo",
+                    "run": "f1",
+                }
+            )
+            [trial] = run_experiments(exp1)
+            self.assertEqual(trial.status, Trial.TERMINATED)
+            self.assertEqual(trial.last_result[TIMESTEPS_TOTAL], 99)
+
+    def testRemoteRunInClient(self):
+        ray.init()
+        assert not ray.util.client.ray.is_connected()
+        with ray_start_client_server():
+            assert ray.util.client.ray.is_connected()
+
+            analysis = run(train_fn)
+            [trial] = analysis.trials
+            self.assertEqual(trial.status, Trial.TERMINATED)
+            self.assertEqual(trial.last_result[TIMESTEPS_TOTAL], 99)
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/tune/tests/test_resource_updater.py

@@ -0,0 +1,151 @@
+from unittest import mock
+
+import ray
+from ray.tests.conftest import *  # noqa
+from ray.tune.utils.resource_updater import _Resources, _ResourceUpdater
+
+
+def test_resources_numerical_error():
+    resource = _Resources(cpu=0.99, gpu=0.99, custom_resources={"a": 0.99})
+    small_resource = _Resources(cpu=0.33, gpu=0.33, custom_resources={"a": 0.33})
+    for i in range(3):
+        resource = _Resources.subtract(resource, small_resource)
+    assert resource.is_nonnegative()
+
+
+def test_resources_subtraction():
+    resource_1 = _Resources(
+        1,
+        0,
+        0,
+        1,
+        custom_resources={"a": 1, "b": 2},
+        extra_custom_resources={"a": 1, "b": 1},
+    )
+    resource_2 = _Resources(
+        1,
+        0,
+        0,
+        1,
+        custom_resources={"a": 1, "b": 2},
+        extra_custom_resources={"a": 1, "b": 1},
+    )
+    new_res = _Resources.subtract(resource_1, resource_2)
+    assert new_res.cpu == 0
+    assert new_res.gpu == 0
+    assert new_res.extra_cpu == 0
+    assert new_res.extra_gpu == 0
+
+    assert all(k == 0 for k in new_res.custom_resources.values())
+    assert all(k == 0 for k in new_res.extra_custom_resources.values())
+
+
+def test_resources_different():
+    resource_1 = _Resources(1, 0, 0, 1, custom_resources={"a": 1, "b": 2})
+    resource_2 = _Resources(1, 0, 0, 1, custom_resources={"a": 1, "c": 2})
+    new_res = _Resources.subtract(resource_1, resource_2)
+    assert "c" in new_res.custom_resources
+    assert "b" in new_res.custom_resources
+
+    assert new_res.cpu == 0
+    assert new_res.gpu == 0
+    assert new_res.extra_cpu == 0
+    assert new_res.extra_gpu == 0
+    assert new_res.get("a") == 0
+
+
+def test_resource_updater(ray_start_cluster):
+    cluster = ray_start_cluster
+
+    resource_updater = _ResourceUpdater(refresh_period=100)
+    # Before initialization, all resources are 0.
+    assert resource_updater.get_num_cpus() == 0
+    assert resource_updater.get_num_gpus() == 0
+
+    cluster.add_node(num_cpus=1, num_gpus=2)
+    cluster.wait_for_nodes()
+    ray.init(address=cluster.address)
+
+    # Resource updater will update resource immediately
+    # after ray is initialized for the first time.
+    assert resource_updater.get_num_cpus() == 1
+    assert resource_updater.get_num_gpus() == 2
+
+    # It will not update the resource before "refresh_period".
+    cluster.add_node(num_cpus=1, num_gpus=1)
+    cluster.wait_for_nodes()
+    assert resource_updater.get_num_cpus() == 1
+    assert resource_updater.get_num_gpus() == 2
+
+    resource_updater = _ResourceUpdater(refresh_period=0)
+    assert resource_updater.get_num_cpus() == 2
+    assert resource_updater.get_num_gpus() == 3
+
+    cluster.add_node(num_cpus=1, num_gpus=1)
+    cluster.wait_for_nodes()
+    assert resource_updater.get_num_cpus() == 3
+    assert resource_updater.get_num_gpus() == 4
+
+
+def test_resource_updater_automatic():
+    """Test that resources are automatically updated when they get out of sync.
+
+    We instantiate a resource updater. When the reported resources are less than
+    what is available, we don't force an update.
+    However, if any of the resources (cpu, gpu, or custom) are higher than what
+    the updater currently think is available, we force an update from the
+    Ray cluster.
+    """
+    resource_updater = _ResourceUpdater()
+    resource_updater._avail_resources = _Resources(
+        cpu=2,
+        gpu=1,
+        memory=1,
+        object_store_memory=1,
+        custom_resources={"a": 4},
+    )
+    resource_updater._last_resource_refresh = 2
+
+    # Should not trigger
+    with mock.patch.object(
+        _ResourceUpdater,
+        "update_avail_resources",
+        wraps=resource_updater.update_avail_resources,
+    ) as upd:
+        # No update
+        assert "2/2 CPUs" in resource_updater.debug_string(
+            total_allocated_resources={"CPU": 2, "GPU": 1, "a": 4}
+        )
+        assert upd.call_count == 0
+
+        # Too many CPUs
+        assert "4/2 CPUs" in resource_updater.debug_string(
+            total_allocated_resources={"CPU": 4, "GPU": 1, "a": 0}
+        )
+        assert upd.call_count == 1
+
+        # Too many GPUs
+        assert "8/1 GPUs" in resource_updater.debug_string(
+            total_allocated_resources={"CPU": 2, "GPU": 8, "a": 0}
+        )
+        assert upd.call_count == 2
+
+        # Too many `a`
+        assert "6/4 a" in resource_updater.debug_string(
+            total_allocated_resources={"CPU": 2, "GPU": 1, "a": 6}
+        )
+        assert upd.call_count == 3
+
+        # No update again
+        assert "2/2 CPUs" in resource_updater.debug_string(
+            total_allocated_resources={"CPU": 2, "GPU": 1, "a": 4}
+        )
+        assert upd.call_count == 3
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/tune/tests/test_result_grid.py

@@ -0,0 +1,214 @@
+import pytest
+
+import ray
+from ray import tune
+from ray.train.tests.util import create_dict_checkpoint, load_dict_checkpoint
+from ray.tune import Checkpoint, Result
+from ray.tune.result_grid import ResultGrid
+
+
+@pytest.fixture
+def ray_start_2_cpus():
+    address_info = ray.init(num_cpus=2)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+def test_result_grid_api(ray_start_2_cpus, tmp_path):
+    def train_fn(config):
+        peak_fn = [0, config["id"], -config["id"], 0]
+        for i in range(len(peak_fn)):
+            with create_dict_checkpoint({"iter": i}) as checkpoint:
+                tune.report(
+                    {"iter": i, "score": config["id"], "peak": peak_fn[i]},
+                    checkpoint=checkpoint,
+                )
+
+    tuner = tune.Tuner(
+        train_fn,
+        param_space={"id": tune.grid_search([1, 2])},
+        run_config=tune.RunConfig(
+            storage_path=str(tmp_path),
+            name="test_result_grid_api",
+            checkpoint_config=tune.CheckpointConfig(num_to_keep=2),
+        ),
+    )
+    result_grid = tuner.fit()
+
+    assert len(result_grid) == 2
+    assert result_grid.experiment_path == str(tmp_path / "test_result_grid_api")
+
+    with pytest.raises(ValueError):
+        result_grid.get_best_result()
+    with pytest.raises(ValueError):
+        result_grid.get_best_result(metric="score")
+    assert result_grid.get_best_result(metric="score", mode="max").config["id"] == 2
+
+    df = result_grid.get_dataframe()
+    assert len(df) == 2
+    assert df["iter"].to_list() == [3, 3]
+
+    df = result_grid.get_dataframe(filter_metric="peak", filter_mode="max")
+    assert df["iter"].to_list() == [1, 1]
+    df = result_grid.get_dataframe(filter_metric="peak", filter_mode="min")
+    assert df["iter"].to_list() == [2, 2]
+
+    assert not result_grid.errors
+    assert result_grid.num_errors == 0
+    assert result_grid.num_terminated == 2
+
+    for result in result_grid:
+        assert result.checkpoint is not None
+        assert result.error is None
+        assert load_dict_checkpoint(result.checkpoint)["iter"] == 3
+        assert {metrics["iter"] for _, metrics in result.best_checkpoints} == {2, 3}
+        assert {
+            load_dict_checkpoint(checkpoint)["iter"]
+            for checkpoint, _ in result.best_checkpoints
+        } == {2, 3}
+
+
+def test_result_grid_no_checkpoint(ray_start_2_cpus):
+    def f(config):
+        pass
+
+    analysis = tune.run(f)
+    result_grid = ResultGrid(analysis)
+    result = result_grid[0]
+    assert result.checkpoint is None
+
+
+def test_best_result_no_report(ray_start_2_cpus):
+    def f(config):
+        pass
+
+    analysis = tune.run(f, config={"x": tune.grid_search([1, 2])})
+    result_grid = ResultGrid(analysis)
+    with pytest.raises(RuntimeError, match="No best trial found*"):
+        result_grid.get_best_result(metric="x", mode="max")
+
+
+def test_result_repr(ray_start_2_cpus):
+    def f(config):
+        tune.report({"loss": 1})
+
+    tuner = tune.Tuner(f, param_space={"x": tune.grid_search([1, 2])})
+    result_grid = tuner.fit()
+    result = result_grid[0]
+
+    from ray.tune.experimental.output import BLACKLISTED_KEYS
+    from ray.tune.result import AUTO_RESULT_KEYS
+
+    representation = result.__repr__()
+    assert not any(key in representation for key in AUTO_RESULT_KEYS)
+    assert not any(key in representation for key in BLACKLISTED_KEYS)
+
+
+def test_result_grid_repr(tmp_path):
+    class MockExperimentAnalysis:
+        trials = []
+
+    result_grid = ResultGrid(experiment_analysis=MockExperimentAnalysis())
+
+    result_grid._results = [
+        Result(
+            metrics={"loss": 1.0},
+            checkpoint=Checkpoint("/tmp/ckpt1"),
+            path="log_1",
+            error=None,
+            metrics_dataframe=None,
+        ),
+        Result(
+            metrics={"loss": 2.0},
+            checkpoint=Checkpoint("/tmp/ckpt2"),
+            path="log_2",
+            error=RuntimeError(),
+            metrics_dataframe=None,
+            best_checkpoints=None,
+        ),
+    ]
+
+    from ray.tune.result import AUTO_RESULT_KEYS
+
+    assert len(result_grid) == 2
+    assert not any(key in repr(result_grid) for key in AUTO_RESULT_KEYS)
+
+    expected_repr = """ResultGrid<[
+  Result(
+    metrics={'loss': 1.0},
+    path='log_1',
+    filesystem='local',
+    checkpoint=Checkpoint(filesystem=local, path=/tmp/ckpt1)
+  ),
+  Result(
+    error='RuntimeError',
+    metrics={'loss': 2.0},
+    path='log_2',
+    filesystem='local',
+    checkpoint=Checkpoint(filesystem=local, path=/tmp/ckpt2)
+  )
+]>"""
+
+    assert repr(result_grid) == expected_repr
+
+
+def test_no_metric_mode_one_trial(ray_start_2_cpus):
+    def f(config):
+        tune.report(dict(x=1))
+
+    results = tune.Tuner(f, tune_config=tune.TuneConfig(num_samples=1)).fit()
+    # This should not throw any exception
+    best_result = results.get_best_result()
+    assert best_result
+
+
+def test_result_grid_df(ray_start_2_cpus):
+    def f(config):
+        tune.report(dict(metric=config["nested"]["param"] * 1))
+        tune.report(dict(metric=config["nested"]["param"] * 4))
+        tune.report(dict(metric=config["nested"]["param"] * 3))
+
+    analysis = tune.run(f, config={"nested": {"param": tune.grid_search([1, 2])}})
+    result_grid = ResultGrid(analysis)
+
+    assert len(result_grid) == 2
+
+    # Last result
+    df = result_grid.get_dataframe()
+    assert sorted(df["metric"]) == [3, 6]
+
+    # Best result (max)
+    df = result_grid.get_dataframe(filter_metric="metric", filter_mode="max")
+    assert sorted(df["metric"]) == [4, 8]
+
+    # Best result (min)
+    df = result_grid.get_dataframe(filter_metric="metric", filter_mode="min")
+    assert sorted(df["metric"]) == [1, 2]
+
+    assert sorted(df["config/nested/param"]) == [1, 2]
+
+
+def test_num_errors_terminated(ray_start_2_cpus, tmp_path):
+    def train_fn(config):
+        if config["id"] == 1:
+            raise RuntimeError()
+        else:
+            tune.report({"score": config["id"]})
+
+    tuner = tune.Tuner(
+        train_fn,
+        param_space={"id": tune.grid_search([1, 2])},
+        run_config=tune.RunConfig(storage_path=str(tmp_path)),
+    )
+
+    result_grid = tuner.fit()
+    assert result_grid.num_errors == 1
+    assert result_grid.num_terminated == 1
+    assert isinstance(result_grid.errors[0], RuntimeError)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/tune/tests/test_run_experiment.py

@@ -0,0 +1,238 @@
+import os
+import unittest
+
+import ray
+from ray.tune import (
+    CheckpointConfig,
+    Trainable,
+    TuneError,
+    register_trainable,
+    run_experiments,
+)
+from ray.tune.experiment import Experiment
+from ray.tune.experiment.trial import ExportFormat, Trial
+from ray.tune.logger import LoggerCallback
+from ray.tune.result import TIMESTEPS_TOTAL
+
+
+def train_fn(config):
+    for i in range(100):
+        ray.tune.report(dict(timesteps_total=i))
+
+
+class RunExperimentTest(unittest.TestCase):
+    def setUp(self):
+        os.environ["TUNE_STATE_REFRESH_PERIOD"] = "0.1"
+        register_trainable("f1", train_fn)
+
+    def tearDown(self):
+        ray.shutdown()
+
+    def testDict(self):
+        trials = run_experiments(
+            {
+                "foo": {
+                    "run": "f1",
+                },
+                "bar": {
+                    "run": "f1",
+                },
+            }
+        )
+        for trial in trials:
+            self.assertEqual(trial.status, Trial.TERMINATED)
+            self.assertEqual(trial.last_result[TIMESTEPS_TOTAL], 99)
+
+    def testExperiment(self):
+        exp1 = Experiment(
+            **{
+                "name": "foo",
+                "run": "f1",
+            }
+        )
+        [trial] = run_experiments(exp1)
+        self.assertEqual(trial.status, Trial.TERMINATED)
+        self.assertEqual(trial.last_result[TIMESTEPS_TOTAL], 99)
+
+    def testExperimentList(self):
+        exp1 = Experiment(
+            **{
+                "name": "foo",
+                "run": "f1",
+            }
+        )
+        exp2 = Experiment(
+            **{
+                "name": "bar",
+                "run": "f1",
+            }
+        )
+        trials = run_experiments([exp1, exp2])
+        for trial in trials:
+            self.assertEqual(trial.status, Trial.TERMINATED)
+            self.assertEqual(trial.last_result[TIMESTEPS_TOTAL], 99)
+
+    def testAutoregisterTrainable(self):
+        class B(Trainable):
+            def step(self):
+                return {"timesteps_this_iter": 1, "done": True}
+
+        trials = run_experiments(
+            {
+                "foo": {
+                    "run": train_fn,
+                },
+                "bar": {"run": B},
+            }
+        )
+        for trial in trials:
+            self.assertEqual(trial.status, Trial.TERMINATED)
+
+    def testCheckpointAtEnd(self):
+        class MyTrainable(Trainable):
+            def step(self):
+                return {"timesteps_this_iter": 1, "done": True}
+
+            def save_checkpoint(self, path):
+                checkpoint = os.path.join(path, "checkpoint")
+                with open(checkpoint, "w") as f:
+                    f.write("OK")
+
+        trials = run_experiments(
+            {
+                "foo": {
+                    "run": MyTrainable,
+                    "checkpoint_config": CheckpointConfig(checkpoint_at_end=True),
+                }
+            }
+        )
+        for trial in trials:
+            self.assertEqual(trial.status, Trial.TERMINATED)
+            self.assertTrue(trial.checkpoint)
+
+    def testExportFormats(self):
+        class train_fn(Trainable):
+            def step(self):
+                return {"timesteps_this_iter": 1, "done": True}
+
+            def _export_model(self, export_formats, export_dir):
+                path = os.path.join(export_dir, "exported")
+                with open(path, "w") as f:
+                    f.write("OK")
+                return {export_formats[0]: path}
+
+        trials = run_experiments(
+            {"foo": {"run": train_fn, "export_formats": ["format"]}}
+        )
+        for trial in trials:
+            self.assertEqual(trial.status, Trial.TERMINATED)
+            self.assertTrue(
+                os.path.exists(
+                    os.path.join(trial.storage.trial_working_directory, "exported")
+                )
+            )
+
+    def testInvalidExportFormats(self):
+        class MyTrainable(Trainable):
+            def step(self):
+                return {"timesteps_this_iter": 1, "done": True}
+
+            def _export_model(self, export_formats, export_dir):
+                ExportFormat.validate(export_formats)
+                return {}
+
+        def fail_trial():
+            run_experiments({"foo": {"run": MyTrainable, "export_formats": ["format"]}})
+
+        self.assertRaises(TuneError, fail_trial)
+
+    def testCustomResources(self):
+        ray.shutdown()
+        ray.init(resources={"hi": 3})
+
+        class MyTrainable(Trainable):
+            def step(self):
+                return {"timesteps_this_iter": 1, "done": True}
+
+        trials = run_experiments(
+            {
+                "foo": {
+                    "run": MyTrainable,
+                    "resources_per_trial": {"cpu": 1, "custom_resources": {"hi": 2}},
+                }
+            }
+        )
+        for trial in trials:
+            self.assertEqual(trial.status, Trial.TERMINATED)
+
+    def testCustomLoggerNoAutoLogging(self):
+        """Does not create CSV/JSON logger callbacks automatically"""
+        os.environ["TUNE_DISABLE_AUTO_CALLBACK_LOGGERS"] = "1"
+
+        class CustomLoggerCallback(LoggerCallback):
+            def log_trial_result(self, iteration, trial, result):
+                with open(os.path.join(trial.local_path, "test.log"), "w") as f:
+                    f.write("hi")
+
+        [trial] = run_experiments(
+            {"foo": {"run": "f1", "stop": {"training_iteration": 1}}},
+            callbacks=[CustomLoggerCallback()],
+        )
+        self.assertTrue(os.path.exists(os.path.join(trial.local_path, "test.log")))
+        self.assertFalse(os.path.exists(os.path.join(trial.local_path, "params.json")))
+
+        [trial] = run_experiments(
+            {"foo": {"run": "f1", "stop": {"training_iteration": 1}}}
+        )
+        self.assertFalse(os.path.exists(os.path.join(trial.local_path, "params.json")))
+
+        [trial] = run_experiments(
+            {"foo": {"run": "f1", "stop": {"training_iteration": 1}}},
+        )
+        self.assertFalse(os.path.exists(os.path.join(trial.local_path, "params.json")))
+
+    def testCustomLoggerWithAutoLogging(self):
+        """Creates CSV/JSON logger callbacks automatically"""
+        if "TUNE_DISABLE_AUTO_CALLBACK_LOGGERS" in os.environ:
+            del os.environ["TUNE_DISABLE_AUTO_CALLBACK_LOGGERS"]
+
+        class CustomLoggerCallback(LoggerCallback):
+            def log_trial_result(self, iteration, trial, result):
+                with open(os.path.join(trial.local_path, "test.log"), "w") as f:
+                    f.write("hi")
+
+        [trial] = run_experiments(
+            {"foo": {"run": "f1", "stop": {"training_iteration": 1}}},
+            callbacks=[CustomLoggerCallback()],
+        )
+        self.assertTrue(os.path.exists(os.path.join(trial.local_path, "test.log")))
+        self.assertTrue(os.path.exists(os.path.join(trial.local_path, "params.json")))
+
+        [trial] = run_experiments(
+            {"foo": {"run": "f1", "stop": {"training_iteration": 1}}}
+        )
+        self.assertTrue(os.path.exists(os.path.join(trial.local_path, "params.json")))
+
+    def testCustomTrialString(self):
+        [trial] = run_experiments(
+            {
+                "foo": {
+                    "run": "f1",
+                    "stop": {"training_iteration": 1},
+                    "trial_name_creator": lambda t: "{}_{}_321".format(
+                        t.trainable_name, t.trial_id
+                    ),
+                }
+            }
+        )
+        self.assertEqual(
+            str(trial), "{}_{}_321".format(trial.trainable_name, trial.trial_id)
+        )
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/tune/tests/test_searcher_utils.py

@@ -0,0 +1,315 @@
+import pytest
+
+from ray.tune.search import BasicVariantGenerator, ConcurrencyLimiter, Searcher
+from ray.tune.search.repeater import Repeater
+from ray.tune.search.search_generator import SearchGenerator
+from ray.tune.utils.mock_trainable import MOCK_TRAINABLE_NAME, register_mock_trainable
+
+
+@pytest.fixture(autouse=True)
+def register_test_trainable():
+    register_mock_trainable()
+
+
+def test_nested_suggestion():
+    class TestSuggestion(Searcher):
+        def suggest(self, trial_id):
+            return {"a": {"b": {"c": {"d": 4, "e": 5}}}}
+
+    searcher = TestSuggestion()
+    alg = SearchGenerator(searcher)
+    alg.add_configurations({"test": {"run": MOCK_TRAINABLE_NAME}})
+    trial = alg.next_trial()
+    assert "e=5" in trial.experiment_tag
+    assert "d=4" in trial.experiment_tag
+
+
+def _repeat_trials(num_samples: int, repeat: int):
+    class TestSuggestion(Searcher):
+        index = 0
+
+        def suggest(self, trial_id):
+            self.index += 1
+            return {"test_variable": 5 + self.index}
+
+        def on_trial_complete(self, *args, **kwargs):
+            return
+
+    searcher = TestSuggestion(metric="episode_reward_mean")
+    repeat_searcher = Repeater(searcher, repeat=repeat, set_index=False)
+    alg = SearchGenerator(repeat_searcher)
+
+    alg.add_configurations(
+        {
+            "test": {
+                "run": MOCK_TRAINABLE_NAME,
+                "num_samples": num_samples,
+                "stop": {"training_iteration": 1},
+            }
+        }
+    )
+
+    trials = []
+    while not alg.is_finished():
+        trials.append(alg.next_trial())
+
+    return trials
+
+
+def test_repeat_1():
+    trials = _repeat_trials(num_samples=2, repeat=1)
+    assert len(trials) == 2
+    parameter_set = {t.evaluated_params["test_variable"] for t in trials}
+    assert len(parameter_set) == 2
+
+
+def test_repeat_4():
+    trials = _repeat_trials(num_samples=12, repeat=4)
+    assert len(trials) == 12
+    parameter_set = {t.evaluated_params["test_variable"] for t in trials}
+    assert len(parameter_set) == 3
+
+
+def test_odd_repeat():
+    trials = _repeat_trials(num_samples=11, repeat=5)
+    assert len(trials) == 11
+    parameter_set = {t.evaluated_params["test_variable"] for t in trials}
+    assert len(parameter_set) == 3
+
+
+def test_set_get_repeater():
+    class TestSuggestion(Searcher):
+        def __init__(self, index):
+            self.index = index
+            self.returned_result = []
+            super().__init__(metric="result", mode="max")
+
+        def suggest(self, trial_id):
+            self.index += 1
+            return {"score": self.index}
+
+        def on_trial_complete(self, trial_id, result=None, **kwargs):
+            self.returned_result.append(result)
+
+    searcher = TestSuggestion(0)
+    repeater1 = Repeater(searcher, repeat=3, set_index=False)
+    for i in range(3):
+        assert repeater1.suggest(f"test_{i}")["score"] == 1
+    for i in range(2):  # An incomplete set of results
+        assert repeater1.suggest(f"test_{i}_2")["score"] == 2
+
+    # Restore a new one
+    state = repeater1.get_state()
+    del repeater1
+    new_repeater = Repeater(searcher, repeat=1, set_index=True)
+    new_repeater.set_state(state)
+    assert new_repeater.repeat == 3
+    assert new_repeater.suggest("test_2_2")["score"] == 2
+    assert new_repeater.suggest("test_x")["score"] == 3
+
+    # Report results
+    for i in range(3):
+        new_repeater.on_trial_complete(f"test_{i}", {"result": 2})
+
+    for i in range(3):
+        new_repeater.on_trial_complete(f"test_{i}_2", {"result": -i * 10})
+
+    assert len(new_repeater.searcher.returned_result) == 2
+    assert new_repeater.searcher.returned_result[-1] == {"result": -10}
+
+    # Finish the rest of the last trial group
+    new_repeater.on_trial_complete("test_x", {"result": 3})
+    assert new_repeater.suggest("test_y")["score"] == 3
+    new_repeater.on_trial_complete("test_y", {"result": 3})
+    assert len(new_repeater.searcher.returned_result) == 2
+    assert new_repeater.suggest("test_z")["score"] == 3
+    new_repeater.on_trial_complete("test_z", {"result": 3})
+    assert len(new_repeater.searcher.returned_result) == 3
+    assert new_repeater.searcher.returned_result[-1] == {"result": 3}
+
+
+def test_set_get_limiter():
+    class TestSuggestion(Searcher):
+        def __init__(self, index):
+            self.index = index
+            self.returned_result = []
+            super().__init__(metric="result", mode="max")
+
+        def suggest(self, trial_id):
+            self.index += 1
+            return {"score": self.index}
+
+        def on_trial_complete(self, trial_id, result=None, **kwargs):
+            self.returned_result.append(result)
+
+    searcher = TestSuggestion(0)
+    limiter = ConcurrencyLimiter(searcher, max_concurrent=2)
+    assert limiter.suggest("test_1")["score"] == 1
+    assert limiter.suggest("test_2")["score"] == 2
+    assert limiter.suggest("test_3") is None
+
+    state = limiter.get_state()
+    del limiter
+    limiter2 = ConcurrencyLimiter(searcher, max_concurrent=3)
+    limiter2.set_state(state)
+    assert limiter2.suggest("test_4") is None
+    assert limiter2.suggest("test_5") is None
+    limiter2.on_trial_complete("test_1", {"result": 3})
+    limiter2.on_trial_complete("test_2", {"result": 3})
+    assert limiter2.suggest("test_3")["score"] == 3
+
+
+def test_basic_variant_limiter():
+    search_alg = BasicVariantGenerator(max_concurrent=2)
+
+    experiment_spec = {
+        "run": MOCK_TRAINABLE_NAME,
+        "num_samples": 5,
+        "stop": {"training_iteration": 1},
+    }
+    search_alg.add_configurations({"test": experiment_spec})
+
+    trial1 = search_alg.next_trial()
+    assert trial1
+
+    trial2 = search_alg.next_trial()
+    assert trial2
+
+    # Returns None because of limiting
+    trial3 = search_alg.next_trial()
+    assert not trial3
+
+    # Finish trial, now trial 3 should be created
+    search_alg.on_trial_complete(trial1.trial_id, None, False)
+    trial3 = search_alg.next_trial()
+    assert trial3
+
+    trial4 = search_alg.next_trial()
+    assert not trial4
+
+    search_alg.on_trial_complete(trial2.trial_id, None, False)
+    search_alg.on_trial_complete(trial3.trial_id, None, False)
+
+    trial4 = search_alg.next_trial()
+    assert trial4
+
+    trial5 = search_alg.next_trial()
+    assert trial5
+
+    search_alg.on_trial_complete(trial4.trial_id, None, False)
+
+    # Should also be None because search is finished
+    trial6 = search_alg.next_trial()
+    assert not trial6
+
+
+def test_batch_limiter():
+    class TestSuggestion(Searcher):
+        def __init__(self, index):
+            self.index = index
+            self.returned_result = []
+            super().__init__(metric="result", mode="max")
+
+        def suggest(self, trial_id):
+            self.index += 1
+            return {"score": self.index}
+
+        def on_trial_complete(self, trial_id, result=None, **kwargs):
+            self.returned_result.append(result)
+
+    searcher = TestSuggestion(0)
+    limiter = ConcurrencyLimiter(searcher, max_concurrent=2, batch=True)
+    assert limiter.suggest("test_1")["score"] == 1
+    assert limiter.suggest("test_2")["score"] == 2
+    assert limiter.suggest("test_3") is None
+
+    limiter.on_trial_complete("test_1", {"result": 3})
+    assert limiter.suggest("test_3") is None
+    limiter.on_trial_complete("test_2", {"result": 3})
+    assert limiter.suggest("test_3") is not None
+
+
+def test_batch_limiter_infinite_loop():
+    """Check whether an infinite loop when less than max_concurrent trials
+    are suggested with batch mode is avoided.
+    """
+
+    class TestSuggestion(Searcher):
+        def __init__(self, index, max_suggestions=10):
+            self.index = index
+            self.max_suggestions = max_suggestions
+            self.returned_result = []
+            super().__init__(metric="result", mode="max")
+
+        def suggest(self, trial_id):
+            self.index += 1
+            if self.index > self.max_suggestions:
+                return None
+            return {"score": self.index}
+
+        def on_trial_complete(self, trial_id, result=None, **kwargs):
+            self.returned_result.append(result)
+            self.index = 0
+
+    searcher = TestSuggestion(0, 2)
+    limiter = ConcurrencyLimiter(searcher, max_concurrent=5, batch=True)
+    limiter.suggest("test_1")
+    limiter.suggest("test_2")
+    limiter.suggest("test_3")  # TestSuggestion return None
+
+    limiter.on_trial_complete("test_1", {"result": 3})
+    limiter.on_trial_complete("test_2", {"result": 3})
+    assert limiter.searcher.returned_result
+
+    searcher = TestSuggestion(0, 10)
+    limiter = ConcurrencyLimiter(searcher, max_concurrent=5, batch=True)
+    limiter.suggest("test_1")
+    limiter.suggest("test_2")
+    limiter.suggest("test_3")
+
+    limiter.on_trial_complete("test_1", {"result": 3})
+    limiter.on_trial_complete("test_2", {"result": 3})
+    assert not limiter.searcher.returned_result
+
+
+def test_set_max_concurrency():
+    """Test whether ``set_max_concurrency`` is called by the
+    ``ConcurrencyLimiter`` and works correctly.
+    """
+
+    class TestSuggestion(Searcher):
+        def __init__(self, index):
+            self.index = index
+            self.returned_result = []
+            self._max_concurrent = 1
+            super().__init__(metric="result", mode="max")
+
+        def suggest(self, trial_id):
+            self.index += 1
+            return {"score": self.index}
+
+        def on_trial_complete(self, trial_id, result=None, **kwargs):
+            self.returned_result.append(result)
+
+        def set_max_concurrency(self, max_concurrent: int) -> bool:
+            self._max_concurrent = max_concurrent
+            return True
+
+    searcher = TestSuggestion(0)
+    limiter_max_concurrent = 2
+    limiter = ConcurrencyLimiter(
+        searcher, max_concurrent=limiter_max_concurrent, batch=True
+    )
+    assert limiter.searcher._max_concurrent == limiter_max_concurrent
+    # Since set_max_concurrency returns True, ConcurrencyLimiter should not
+    # be limiting concurrency itself
+    assert not limiter._limit_concurrency
+    assert limiter.suggest("test_1")["score"] == 1
+    assert limiter.suggest("test_2")["score"] == 2
+    assert limiter.suggest("test_3")["score"] == 3
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/tune/tests/test_searchers.py

@@ -0,0 +1,877 @@
+import contextlib
+import os
+import shutil
+import sys
+import tempfile
+import unittest
+from copy import deepcopy
+from unittest.mock import patch
+
+import numpy as np
+import pandas
+import pytest
+from packaging.version import Version
+
+import ray
+from ray import tune
+from ray.air.constants import TRAINING_ITERATION
+from ray.tune.search import ConcurrencyLimiter
+
+
+def _invalid_objective(config):
+    metric = "report"
+
+    if config[metric] > 4:
+        tune.report({"_metric": float("inf")})
+    elif config[metric] > 3:
+        tune.report({"_metric": float("-inf")})
+    elif config[metric] > 2:
+        tune.report({"_metric": np.nan})
+    else:
+        tune.report({"_metric": float(config[metric]) or 0.1})
+
+
+def _multi_objective(config):
+    tune.report(dict(a=config["a"] * 100, b=config["b"] * -100, c=config["c"]))
+
+
+def _dummy_objective(config):
+    tune.report(dict(metric=config["report"]))
+
+
+class InvalidValuesTest(unittest.TestCase):
+    """
+    Test searcher handling of invalid values (NaN, -inf, inf).
+    Implicitly tests automatic config conversion and default (anonymous)
+    mode handling.
+    Also tests that searcher save doesn't throw any errors during
+    experiment checkpointing.
+    """
+
+    def setUp(self):
+        self.config = {"report": tune.uniform(0.0, 5.0), "list": [1, 2, 3], "num": 4}
+
+    def tearDown(self):
+        pass
+
+    @classmethod
+    def setUpClass(cls):
+        ray.init(num_cpus=4, num_gpus=0, include_dashboard=False)
+
+    @classmethod
+    def tearDownClass(cls):
+        ray.shutdown()
+
+    def assertCorrectExperimentOutput(self, analysis):
+        best_trial = analysis.best_trial
+        self.assertLessEqual(best_trial.config["report"], 2.0)
+        # Make sure that constant parameters aren't lost
+        # Hyperopt converts lists to tuples, so check for either
+        self.assertIn(best_trial.config["list"], ([1, 2, 3], (1, 2, 3)))
+        self.assertEqual(best_trial.config["num"], 4)
+
+    @contextlib.contextmanager
+    def check_searcher_checkpoint_errors_scope(self):
+        buffer = []
+        from ray.tune.execution.tune_controller import logger
+
+        with patch.object(logger, "warning", lambda x: buffer.append(x)):
+            yield
+
+        assert not any(
+            "Experiment state snapshotting failed: Can't pickle local object" in x
+            for x in buffer
+        ), "Searcher checkpointing failed (unable to serialize)."
+
+    def testAxManualSetup(self):
+        from ax.service.ax_client import AxClient, ObjectiveProperties
+
+        from ray.tune.search.ax import AxSearch
+
+        config = self.config.copy()
+        config["mixed_list"] = [1, tune.uniform(2, 3), 4]
+        converted_config = AxSearch.convert_search_space(config)
+        # At least one nan, inf, -inf and float
+        client = AxClient(random_seed=4321)
+
+        client.create_experiment(
+            parameters=converted_config,
+            objectives={"_metric": ObjectiveProperties(minimize=False)},
+        )
+        searcher = AxSearch(ax_client=client)
+
+        out = tune.run(
+            _invalid_objective,
+            search_alg=searcher,
+            metric="_metric",
+            mode="max",
+            num_samples=4,
+            reuse_actors=False,
+        )
+        self.assertCorrectExperimentOutput(out)
+        self.assertEqual(out.best_trial.config["mixed_list"][0], 1)
+        self.assertGreaterEqual(out.best_trial.config["mixed_list"][1], 2)
+        self.assertLess(out.best_trial.config["mixed_list"][1], 3)
+        self.assertEqual(out.best_trial.config["mixed_list"][2], 4)
+
+    def testAx(self):
+        from ray.tune.search.ax import AxSearch
+
+        searcher = ConcurrencyLimiter(AxSearch(random_seed=4321), max_concurrent=2)
+
+        with self.check_searcher_checkpoint_errors_scope():
+            # Make sure enough samples are used so that Ax actually fits a model
+            # for config suggestion
+            out = tune.run(
+                _invalid_objective,
+                search_alg=searcher,
+                metric="_metric",
+                mode="max",
+                num_samples=16,
+                reuse_actors=False,
+                config=self.config,
+            )
+
+        self.assertCorrectExperimentOutput(out)
+
+    def testBayesOpt(self):
+        from ray.tune.search.bayesopt import BayesOptSearch
+
+        with self.check_searcher_checkpoint_errors_scope():
+            out = tune.run(
+                _invalid_objective,
+                # At least one nan, inf, -inf and float
+                search_alg=BayesOptSearch(random_state=1234),
+                config=self.config,
+                metric="_metric",
+                mode="max",
+                num_samples=8,
+                reuse_actors=False,
+            )
+        self.assertCorrectExperimentOutput(out)
+
+    @pytest.mark.skipif(
+        sys.version_info >= (3, 12),
+        reason="BOHB not yet supported for python 3.12+",
+    )
+    def testBOHB(self):
+        from ray.tune.search.bohb import TuneBOHB
+
+        with self.check_searcher_checkpoint_errors_scope():
+            out = tune.run(
+                _invalid_objective,
+                search_alg=TuneBOHB(seed=1000),
+                config=self.config,
+                metric="_metric",
+                mode="max",
+                num_samples=8,
+                reuse_actors=False,
+            )
+        self.assertCorrectExperimentOutput(out)
+
+    @pytest.mark.skipif(
+        sys.version_info >= (3, 12), reason="HEBO doesn't support py312"
+    )
+    def testHEBO(self):
+        if Version(pandas.__version__) >= Version("2.0.0"):
+            pytest.skip("HEBO does not support pandas>=2.0.0")
+
+        from ray.tune.search.hebo import HEBOSearch
+
+        with self.check_searcher_checkpoint_errors_scope():
+            out = tune.run(
+                _invalid_objective,
+                # At least one nan, inf, -inf and float
+                search_alg=HEBOSearch(random_state_seed=123),
+                config=self.config,
+                metric="_metric",
+                mode="max",
+                num_samples=8,
+                reuse_actors=False,
+            )
+        self.assertCorrectExperimentOutput(out)
+
+    def testHyperopt(self):
+        from ray.tune.search.hyperopt import HyperOptSearch
+
+        with self.check_searcher_checkpoint_errors_scope():
+            out = tune.run(
+                _invalid_objective,
+                # At least one nan, inf, -inf and float
+                search_alg=HyperOptSearch(random_state_seed=1234),
+                config=self.config,
+                metric="_metric",
+                mode="max",
+                num_samples=8,
+                reuse_actors=False,
+            )
+        self.assertCorrectExperimentOutput(out)
+
+    def testNevergrad(self):
+        import nevergrad as ng
+
+        from ray.tune.search.nevergrad import NevergradSearch
+
+        np.random.seed(2020)  # At least one nan, inf, -inf and float
+
+        with self.check_searcher_checkpoint_errors_scope():
+            out = tune.run(
+                _invalid_objective,
+                search_alg=NevergradSearch(optimizer=ng.optimizers.RandomSearch),
+                config=self.config,
+                mode="max",
+                num_samples=16,
+                reuse_actors=False,
+            )
+        self.assertCorrectExperimentOutput(out)
+
+    def testNevergradWithRequiredOptimizerKwargs(self):
+        import nevergrad as ng
+
+        from ray.tune.search.nevergrad import NevergradSearch
+
+        NevergradSearch(optimizer=ng.optimizers.CM, optimizer_kwargs=dict(budget=16))
+
+    def testOptuna(self):
+        from optuna.samplers import RandomSampler
+
+        from ray.tune.search.optuna import OptunaSearch
+
+        np.random.seed(1000)  # At least one nan, inf, -inf and float
+
+        with self.check_searcher_checkpoint_errors_scope():
+            out = tune.run(
+                _invalid_objective,
+                search_alg=OptunaSearch(sampler=RandomSampler(seed=1234), storage=None),
+                config=self.config,
+                metric="_metric",
+                mode="max",
+                num_samples=8,
+                reuse_actors=False,
+            )
+        self.assertCorrectExperimentOutput(out)
+
+    def testOptunaWithStorage(self):
+        from optuna.samplers import RandomSampler
+        from optuna.storages import JournalStorage
+        from optuna.storages.journal import JournalFileBackend
+
+        from ray.tune.search.optuna import OptunaSearch
+
+        np.random.seed(1000)  # At least one nan, inf, -inf and float
+        storage_file_path = "/tmp/my_test_study.log"
+
+        with self.check_searcher_checkpoint_errors_scope():
+            out = tune.run(
+                _invalid_objective,
+                search_alg=OptunaSearch(
+                    sampler=RandomSampler(seed=1234),
+                    study_name="my_test_study",
+                    storage=JournalStorage(
+                        JournalFileBackend(file_path=storage_file_path)
+                    ),
+                ),
+                config=self.config,
+                metric="_metric",
+                mode="max",
+                num_samples=8,
+                reuse_actors=False,
+            )
+        self.assertCorrectExperimentOutput(out)
+        self.assertTrue(os.path.exists(storage_file_path))
+
+    def testOptunaReportTooOften(self):
+        from optuna.samplers import RandomSampler
+
+        from ray.tune.search.optuna import OptunaSearch
+
+        searcher = OptunaSearch(
+            sampler=RandomSampler(seed=1234),
+            space=OptunaSearch.convert_search_space(self.config),
+            metric="metric",
+            mode="max",
+        )
+        searcher.suggest("trial_1")
+        searcher.on_trial_result("trial_1", {"training_iteration": 1, "metric": 1})
+        searcher.on_trial_complete("trial_1", {"training_iteration": 2, "metric": 1})
+
+        # Report after complete should not fail
+        searcher.on_trial_result("trial_1", {"training_iteration": 3, "metric": 1})
+
+        searcher.on_trial_complete("trial_1", {"training_iteration": 4, "metric": 1})
+
+    def testZOOpt(self):
+        self.skipTest(
+            "Recent ZOOpt versions fail handling invalid values gracefully. "
+            "Skipping until a fix is added in a future ZOOpt release."
+        )
+        from ray.tune.search.zoopt import ZOOptSearch
+
+        # This seed tests that a nan result doesn't cause an error if it shows
+        # up after the initial data collection phase.
+        np.random.seed(1002)  # At least one nan, inf, -inf and float
+
+        with self.check_searcher_checkpoint_errors_scope():
+            out = tune.run(
+                _invalid_objective,
+                search_alg=ZOOptSearch(budget=25, parallel_num=4),
+                config=self.config,
+                metric="_metric",
+                mode="max",
+                num_samples=16,
+                reuse_actors=False,
+            )
+        self.assertCorrectExperimentOutput(out)
+
+
+class AddEvaluatedPointTest(unittest.TestCase):
+    """
+    Test add_evaluated_point method in searchers that support it.
+    """
+
+    def setUp(self):
+        self.param_name = "report"
+        self.valid_value = 1.0
+        self.space = {self.param_name: tune.uniform(0.0, 5.0)}
+
+        self.analysis = tune.run(
+            _dummy_objective,
+            config=self.space,
+            metric="metric",
+            num_samples=4,
+            verbose=0,
+        )
+
+    def tearDown(self):
+        pass
+
+    @classmethod
+    def setUpClass(cls):
+        ray.init(num_cpus=4, num_gpus=0, include_dashboard=False)
+
+    @classmethod
+    def tearDownClass(cls):
+        ray.shutdown()
+
+    def run_add_evaluated_point(self, point, searcher, get_len_X, get_len_y):
+        searcher = deepcopy(searcher)
+        len_X = get_len_X(searcher)
+        len_y = get_len_y(searcher)
+        self.assertEqual(len_X, 0)
+        self.assertEqual(len_y, 0)
+
+        searcher.add_evaluated_point(point, 1.0)
+
+        len_X = get_len_X(searcher)
+        len_y = get_len_y(searcher)
+        self.assertEqual(len_X, 1)
+        self.assertEqual(len_y, 1)
+
+        searcher.suggest("1")
+
+    def run_add_evaluated_trials(self, searcher, get_len_X, get_len_y):
+        searcher_copy = deepcopy(searcher)
+        searcher_copy.add_evaluated_trials(self.analysis, "metric")
+        self.assertEqual(get_len_X(searcher_copy), 4)
+        self.assertEqual(get_len_y(searcher_copy), 4)
+        searcher_copy.suggest("1")
+
+        searcher_copy = deepcopy(searcher)
+        searcher_copy.add_evaluated_trials(self.analysis.trials, "metric")
+        self.assertEqual(get_len_X(searcher_copy), 4)
+        self.assertEqual(get_len_y(searcher_copy), 4)
+        searcher_copy.suggest("1")
+
+        searcher_copy = deepcopy(searcher)
+        searcher_copy.add_evaluated_trials(self.analysis.trials[0], "metric")
+        self.assertEqual(get_len_X(searcher_copy), 1)
+        self.assertEqual(get_len_y(searcher_copy), 1)
+        searcher_copy.suggest("1")
+
+    def testOptuna(self):
+        from optuna.storages import JournalStorage
+        from optuna.storages.journal import JournalFileBackend
+        from optuna.trial import TrialState
+
+        from ray.tune.search.optuna import OptunaSearch
+
+        # OptunaSearch with in-memory storage
+        searcher = OptunaSearch(
+            space=self.space,
+            storage=None,
+            metric="metric",
+            mode="max",
+            points_to_evaluate=[{self.param_name: self.valid_value}],
+            evaluated_rewards=[1.0],
+        )
+
+        get_len = lambda s: len(s._ot_study.trials)  # noqa E731
+
+        self.assertGreater(get_len(searcher), 0)
+
+        # OptunaSearch with external storage
+        storage_file_path = "/tmp/my_test_study.log"
+        searcher = OptunaSearch(
+            space=self.space,
+            study_name="my_test_study",
+            storage=JournalStorage(JournalFileBackend(file_path=storage_file_path)),
+            metric="metric",
+            mode="max",
+            points_to_evaluate=[{self.param_name: self.valid_value}],
+            evaluated_rewards=[1.0],
+        )
+
+        get_len = lambda s: len(s._ot_study.trials)  # noqa E731
+
+        self.assertGreater(get_len(searcher), 0)
+        self.assertTrue(os.path.exists(storage_file_path))
+
+        searcher = OptunaSearch(
+            space=self.space,
+            metric="metric",
+            mode="max",
+        )
+
+        point = {
+            self.param_name: self.valid_value,
+        }
+
+        self.assertEqual(get_len(searcher), 0)
+
+        searcher.add_evaluated_point(point, 1.0, intermediate_values=[0.8, 0.9])
+        self.assertEqual(get_len(searcher), 1)
+        self.assertTrue(searcher._ot_study.trials[-1].state == TrialState.COMPLETE)
+
+        searcher.add_evaluated_point(
+            point, 1.0, intermediate_values=[0.8, 0.9], error=True
+        )
+        self.assertEqual(get_len(searcher), 2)
+        self.assertTrue(searcher._ot_study.trials[-1].state == TrialState.FAIL)
+
+        searcher.add_evaluated_point(
+            point, 1.0, intermediate_values=[0.8, 0.9], pruned=True
+        )
+        self.assertEqual(get_len(searcher), 3)
+        self.assertTrue(searcher._ot_study.trials[-1].state == TrialState.PRUNED)
+
+        searcher.suggest("1")
+
+        searcher = OptunaSearch(
+            space=self.space,
+            metric="metric",
+            mode="max",
+        )
+
+        self.run_add_evaluated_trials(searcher, get_len, get_len)
+
+        def dbr_space(trial):
+            return {self.param_name: trial.suggest_float(self.param_name, 0.0, 5.0)}
+
+        dbr_searcher = OptunaSearch(
+            space=dbr_space,
+            metric="metric",
+            mode="max",
+        )
+        with self.assertRaises(TypeError):
+            dbr_searcher.add_evaluated_point(point, 1.0)
+
+    @pytest.mark.skipif(
+        sys.version_info >= (3, 12), reason="HEBO doesn't support py312"
+    )
+    def testHEBO(self):
+        if Version(pandas.__version__) >= Version("2.0.0"):
+            pytest.skip("HEBO does not support pandas>=2.0.0")
+
+        from ray.tune.search.hebo import HEBOSearch
+
+        searcher = HEBOSearch(
+            space=self.space,
+            metric="metric",
+            mode="max",
+        )
+
+        point = {
+            self.param_name: self.valid_value,
+        }
+
+        get_len_X = lambda s: len(s._opt.X)  # noqa E731
+        get_len_y = lambda s: len(s._opt.y)  # noqa E731
+
+        self.run_add_evaluated_point(point, searcher, get_len_X, get_len_y)
+        self.run_add_evaluated_trials(searcher, get_len_X, get_len_y)
+
+
+class SaveRestoreCheckpointTest(unittest.TestCase):
+    """
+    Test searcher save and restore functionality.
+    """
+
+    def setUp(self):
+        self.tempdir = tempfile.mkdtemp()
+        self.checkpoint_path = os.path.join(self.tempdir, "checkpoint.pkl")
+        self.metric_name = "metric"
+        self.config = {"a": tune.uniform(0.0, 5.0)}
+
+    def tearDown(self):
+        shutil.rmtree(self.tempdir)
+
+    @classmethod
+    def setUpClass(cls):
+        ray.init(num_cpus=4, num_gpus=0, include_dashboard=False)
+
+    @classmethod
+    def tearDownClass(cls):
+        ray.shutdown()
+
+    def _on_trial_callbacks(self, searcher, trial_id):
+        result = {
+            TRAINING_ITERATION: 1,
+            self.metric_name: 1,
+            "config/a": 1.0,
+            "time_total_s": 1,
+        }
+        searcher.on_trial_result(trial_id, result)
+        searcher.on_trial_complete(trial_id, result)
+
+    def _save(self, searcher):
+        searcher.set_search_properties(
+            metric=self.metric_name, mode="max", config=self.config
+        )
+
+        searcher.suggest("1")
+        searcher.suggest("2")
+        searcher.suggest("not_completed")
+        self._on_trial_callbacks(searcher, "1")
+
+        searcher.save(self.checkpoint_path)
+
+    def _restore(self, searcher):
+        # Restoration shouldn't require another call to `searcher.set_search_properties`
+        searcher.restore(self.checkpoint_path)
+
+        self._on_trial_callbacks(searcher, "2")
+        searcher.suggest("3")
+        self._on_trial_callbacks(searcher, "3")
+
+        # NOTE: Trial "not_completed" that was suggested before saving never completes
+        # We expect that it should still be tracked in the searcher state,
+        # which is usually done in the searcher's `_live_trial_mapping`.
+        # See individual searcher tests below for the special cases (e.g. Optuna, BOHB).
+        if hasattr(searcher, "_live_trial_mapping"):
+            assert "not_completed" in searcher._live_trial_mapping
+
+    def testAx(self):
+        from ax.service.ax_client import AxClient, ObjectiveProperties
+
+        from ray.tune.search.ax import AxSearch
+
+        converted_config = AxSearch.convert_search_space(self.config)
+        client = AxClient()
+        client.create_experiment(
+            parameters=converted_config,
+            objectives={self.metric_name: ObjectiveProperties(minimize=False)},
+        )
+        searcher = AxSearch(ax_client=client)
+
+        self._save(searcher)
+
+        client = AxClient()
+        client.create_experiment(
+            parameters=converted_config,
+            objectives={self.metric_name: ObjectiveProperties(minimize=False)},
+        )
+        searcher = AxSearch(ax_client=client)
+        self._restore(searcher)
+
+    def testBayesOpt(self):
+        from ray.tune.search.bayesopt import BayesOptSearch
+
+        searcher = BayesOptSearch(
+            space=self.config, metric=self.metric_name, mode="max"
+        )
+        self._save(searcher)
+
+        searcher = BayesOptSearch()
+        self._restore(searcher)
+
+    @pytest.mark.skipif(
+        sys.version_info >= (3, 12),
+        reason="BOHB not yet supported for python 3.12+",
+    )
+    def testBOHB(self):
+        from ray.tune.search.bohb import TuneBOHB
+
+        searcher = TuneBOHB(space=self.config, metric=self.metric_name, mode="max")
+
+        self._save(searcher)
+
+        searcher = TuneBOHB()
+        self._restore(searcher)
+
+        assert "not_completed" in searcher.trial_to_params
+
+    @pytest.mark.skipif(
+        sys.version_info >= (3, 12), reason="HEBO doesn't support py312"
+    )
+    def testHEBO(self):
+        if Version(pandas.__version__) >= Version("2.0.0"):
+            pytest.skip("HEBO does not support pandas>=2.0.0")
+
+        from ray.tune.search.hebo import HEBOSearch
+
+        searcher = HEBOSearch(
+            space=self.config,
+            metric=self.metric_name,
+            mode="max",
+            random_state_seed=1234,
+        )
+
+        self._save(searcher)
+
+        searcher = HEBOSearch()
+        self._restore(searcher)
+
+    def testHyperopt(self):
+        from ray.tune.search.hyperopt import HyperOptSearch
+
+        searcher = HyperOptSearch(
+            space=self.config,
+            metric=self.metric_name,
+            mode="max",
+        )
+        self._save(searcher)
+
+        searcher = HyperOptSearch()
+        self._restore(searcher)
+
+    def testNevergrad(self):
+        import nevergrad as ng
+
+        from ray.tune.search.nevergrad import NevergradSearch
+
+        searcher = NevergradSearch(
+            space=self.config,
+            metric=self.metric_name,
+            mode="max",
+            optimizer=ng.optimizers.RandomSearch,
+        )
+        self._save(searcher)
+
+        # `optimizer` is the only required argument
+        searcher = NevergradSearch(optimizer=ng.optimizers.RandomSearch)
+        self._restore(searcher)
+
+    def testOptuna(self):
+        from ray.tune.search.optuna import OptunaSearch
+
+        searcher = OptunaSearch(
+            space=self.config,
+            storage=None,
+            metric=self.metric_name,
+            mode="max",
+        )
+        self._save(searcher)
+
+        searcher = OptunaSearch()
+        self._restore(searcher)
+
+        assert "not_completed" in searcher._ot_trials
+
+    def testOptunaWithStorage(self):
+        from optuna.storages import JournalStorage
+        from optuna.storages.journal import JournalFileBackend
+
+        from ray.tune.search.optuna import OptunaSearch
+
+        storage_file_path = "/tmp/my_test_study.log"
+        searcher = OptunaSearch(
+            space=self.config,
+            study_name="my_test_study",
+            storage=JournalStorage(JournalFileBackend(file_path=storage_file_path)),
+            metric=self.metric_name,
+            mode="max",
+        )
+        self._save(searcher)
+
+        searcher = OptunaSearch()
+        self._restore(searcher)
+
+        assert "not_completed" in searcher._ot_trials
+        self.assertTrue(os.path.exists(storage_file_path))
+
+    def testZOOpt(self):
+        from ray.tune.search.zoopt import ZOOptSearch
+
+        searcher = ZOOptSearch(
+            space=self.config,
+            metric=self.metric_name,
+            mode="max",
+            budget=100,
+            parallel_num=4,
+        )
+
+        self._save(searcher)
+
+        # `budget` is the only required argument - will get replaced on restore
+        searcher = ZOOptSearch(budget=0)
+        self._restore(searcher)
+        assert searcher._budget == 100
+
+
+class MultiObjectiveTest(unittest.TestCase):
+    """
+    Test multi-objective optimization in searchers that support it.
+    """
+
+    def setUp(self):
+        self.config = {
+            "a": tune.uniform(0, 1),
+            "b": tune.uniform(0, 1),
+            "c": tune.uniform(0, 1),
+        }
+
+    def tearDown(self):
+        pass
+
+    @classmethod
+    def setUpClass(cls):
+        ray.init(num_cpus=4, num_gpus=0, include_dashboard=False)
+
+    @classmethod
+    def tearDownClass(cls):
+        ray.shutdown()
+
+    def testOptuna(self):
+        from optuna.samplers import RandomSampler
+
+        from ray.tune.search.optuna import OptunaSearch
+
+        np.random.seed(1000)
+
+        out = tune.run(
+            _multi_objective,
+            search_alg=OptunaSearch(
+                sampler=RandomSampler(seed=1234),
+                storage=None,
+                metric=["a", "b", "c"],
+                mode=["max", "min", "max"],
+            ),
+            config=self.config,
+            num_samples=16,
+            reuse_actors=False,
+        )
+
+        best_trial_a = out.get_best_trial("a", "max")
+        self.assertGreaterEqual(best_trial_a.config["a"], 0.8)
+        best_trial_b = out.get_best_trial("b", "min")
+        self.assertGreaterEqual(best_trial_b.config["b"], 0.8)
+        best_trial_c = out.get_best_trial("c", "max")
+        self.assertGreaterEqual(best_trial_c.config["c"], 0.8)
+
+    def testOptunaWithStorage(self):
+        from optuna.samplers import RandomSampler
+        from optuna.storages import JournalStorage
+        from optuna.storages.journal import JournalFileBackend
+
+        from ray.tune.search.optuna import OptunaSearch
+
+        np.random.seed(1000)
+        storage_file_path = "/tmp/my_test_study.log"
+
+        out = tune.run(
+            _multi_objective,
+            search_alg=OptunaSearch(
+                sampler=RandomSampler(seed=1234),
+                study_name="my_test_study",
+                storage=JournalStorage(JournalFileBackend(file_path=storage_file_path)),
+                metric=["a", "b", "c"],
+                mode=["max", "min", "max"],
+            ),
+            config=self.config,
+            num_samples=16,
+            reuse_actors=False,
+        )
+
+        best_trial_a = out.get_best_trial("a", "max")
+        self.assertGreaterEqual(best_trial_a.config["a"], 0.8)
+        best_trial_b = out.get_best_trial("b", "min")
+        self.assertGreaterEqual(best_trial_b.config["b"], 0.8)
+        best_trial_c = out.get_best_trial("c", "max")
+        self.assertGreaterEqual(best_trial_c.config["c"], 0.8)
+        self.assertTrue(os.path.exists(storage_file_path))
+
+
+class BayesOptHashPrecisionTest(unittest.TestCase):
+    def testDictHashPrecisionDistinguishesNearFloats(self):
+        from ray.tune.search.bayesopt.bayesopt_search import _dict_hash
+
+        a = {"lr": 1.00001e-05}
+        b = {"lr": 1.46532e-05}
+        # The default precision of 5 rounds both to the same string, so the
+        # two distinct configs collide and one suggestion would be skipped.
+        self.assertEqual(_dict_hash(a, 5), _dict_hash(b, 5))
+        # A higher precision keeps them apart.
+        self.assertNotEqual(_dict_hash(a, 16), _dict_hash(b, 16))
+
+    def testRepeatFloatPrecisionIsConfigurable(self):
+        pytest.importorskip("bayes_opt")
+        from ray.tune.search.bayesopt import BayesOptSearch
+
+        # Default stays at 5 for backward compatibility.
+        self.assertEqual(BayesOptSearch().repeat_float_precision, 5)
+        searcher = BayesOptSearch(repeat_float_precision=16)
+        self.assertEqual(searcher.repeat_float_precision, 16)
+
+    def testInvalidRepeatFloatPrecisionRaises(self):
+        pytest.importorskip("bayes_opt")
+        from ray.tune.search.bayesopt import BayesOptSearch
+
+        with self.assertRaises(ValueError):
+            BayesOptSearch(repeat_float_precision=-1)
+
+        with self.assertRaises(TypeError):
+            BayesOptSearch(repeat_float_precision="5")
+
+        with self.assertRaises(TypeError):
+            BayesOptSearch(repeat_float_precision=5.5)
+
+        with self.assertRaises(TypeError):
+            BayesOptSearch(repeat_float_precision=True)
+
+
+class BayesOptConvergenceWarningTest(unittest.TestCase):
+    def testWarnsAndStopsOnConvergence(self):
+        """BayesOptSearch should warn (not silently stop) when it converges."""
+        from ray.tune.search import Searcher
+        from ray.tune.search.bayesopt import BayesOptSearch
+
+        space = {"width": tune.uniform(0, 20), "height": tune.uniform(-100, 100)}
+        # patience=1 -> the search stops as soon as a config first repeats,
+        # making convergence deterministic and quick to reach.
+        searcher = BayesOptSearch(
+            space=space, metric="loss", mode="min", random_state=42, patience=1
+        )
+        logger_name = "ray.tune.search.bayesopt.bayesopt_search"
+        finished = False
+        with self.assertLogs(logger_name, level="WARNING") as cm:
+            for i in range(50):
+                config = searcher.suggest(f"trial_{i}")
+                if config == Searcher.FINISHED:
+                    finished = True
+                    break
+                if config is None:
+                    continue
+                searcher.on_trial_complete(
+                    f"trial_{i}", {"loss": config["width"] + config["height"]}
+                )
+        self.assertTrue(finished, "BayesOptSearch should finish once converged")
+        self.assertTrue(
+            any("stopping early" in msg for msg in cm.output),
+            f"Expected a convergence warning, got: {cm.output}",
+        )
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/tune/tests/test_soft_imports.py

@@ -0,0 +1,25 @@
+import sys
+import unittest
+
+
+class TestSoftImports(unittest.TestCase):
+    """Tests whether it's possible to use Ray Tune without soft dependencies"""
+
+    def testSoftImports(self):
+        import ray.tune.schedulers  # noqa: F401
+        from ray.tune.search import SEARCH_ALG_IMPORT
+
+        for name, import_func in SEARCH_ALG_IMPORT.items():
+            print(f"testing searcher {name}")
+            searcher = import_func()
+
+            # ensure that the dependencies aren't actually installed
+            if searcher and name not in ("variant_generator", "random"):
+                with self.assertRaises((AssertionError, ImportError)):
+                    searcher()
+
+
+if __name__ == "__main__":
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/tune/tests/test_stopper.py

@@ -0,0 +1,64 @@
+import pickle
+
+from freezegun import freeze_time
+
+from ray.tune.stopper import TimeoutStopper
+
+
+def test_timeout_stopper_timeout():
+    with freeze_time() as frozen:
+        stopper = TimeoutStopper(timeout=60)
+        assert not stopper.stop_all()
+        frozen.tick(40)
+        assert not stopper.stop_all()
+        frozen.tick(22)
+        assert stopper.stop_all()
+
+
+def test_timeout_stopper_recover_before_timeout():
+    """If checkpointed before timeout, should continue where we left."""
+    with freeze_time() as frozen:
+        stopper = TimeoutStopper(timeout=60)
+        assert not stopper.stop_all()
+        frozen.tick(40)
+        assert not stopper.stop_all()
+        checkpoint = pickle.dumps(stopper)
+
+        # Continue sometime in the future. This is after start_time + timeout
+        # but we should still continue training.
+        frozen.tick(200)
+
+        # Continue, so we shouldn't time out
+        stopper = pickle.loads(checkpoint)
+        assert not stopper.stop_all()
+        frozen.tick(10)
+        assert not stopper.stop_all()
+        frozen.tick(12)
+        assert stopper.stop_all()
+
+
+def test_timeout_stopper_recover_after_timeout():
+    """If checkpointed after timeout, should still stop after recover."""
+    with freeze_time() as frozen:
+        stopper = TimeoutStopper(timeout=60)
+        assert not stopper.stop_all()
+        frozen.tick(62)
+        assert stopper.stop_all()
+        checkpoint = pickle.dumps(stopper)
+
+        # Continue sometime in the future
+        frozen.tick(200)
+
+        # Continue, so we should still time out.
+        stopper = pickle.loads(checkpoint)
+        assert stopper.stop_all()
+        frozen.tick(10)
+        assert stopper.stop_all()
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__] + sys.argv[1:]))

python/ray/tune/tests/test_syncer.py

@@ -0,0 +1,404 @@
+import logging
+import os
+import time
+from typing import List, Optional
+
+import pytest
+from freezegun import freeze_time
+
+import ray
+import ray.cloudpickle as pickle
+from ray.train._internal.storage import (
+    _download_from_fs_path,
+    _FilesystemSyncer,
+    _upload_to_fs_path,
+    get_fs_and_path,
+)
+from ray.train._internal.syncer import _BackgroundProcess
+from ray.train.tests.test_new_persistence import _create_mock_custom_fs
+
+
+@pytest.fixture
+def propagate_logs():
+    # Ensure that logs are propagated to ancestor handles. This is required if using the
+    # caplog fixture with Ray's logging.
+    # NOTE: This only enables log propagation in the driver process, not the workers!
+    logger = logging.getLogger("ray")
+    logger.propagate = True
+    yield
+    logger.propagate = False
+
+
+@pytest.fixture
+def ray_start_4_cpus():
+    address_info = ray.init(num_cpus=4, configure_logging=False)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+@pytest.fixture
+def ray_start_2_cpus():
+    address_info = ray.init(num_cpus=2, configure_logging=False)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+@pytest.fixture
+def shutdown_only():
+    yield None
+    ray.shutdown()
+
+
+@pytest.fixture
+def temp_data_dirs(tmp_path):
+    tmp_source = tmp_path / "source"
+    tmp_target = tmp_path / "target"
+    tmp_target.mkdir()
+
+    os.makedirs(os.path.join(tmp_source, "subdir", "nested"))
+    os.makedirs(os.path.join(tmp_source, "subdir_exclude", "something"))
+
+    files = [
+        "level0.txt",
+        "level0_exclude.txt",
+        "subdir/level1.txt",
+        "subdir/level1_exclude.txt",
+        "subdir/nested/level2.txt",
+        "subdir_nested_level2_exclude.txt",
+        "subdir_exclude/something/somewhere.txt",
+    ]
+
+    for file in files:
+        with open(os.path.join(tmp_source, file), "w") as f:
+            f.write("Data")
+
+    yield str(tmp_source), str(tmp_target)
+
+
+@pytest.fixture
+def syncer(tmp_path):
+    yield _FilesystemSyncer(storage_filesystem=_create_mock_custom_fs(tmp_path))
+
+
+def assert_file(exists: bool, root: str, path: str):
+    full_path = os.path.join(root, path)
+
+    if exists:
+        assert os.path.exists(full_path)
+    else:
+        assert not os.path.exists(full_path)
+
+
+def test_syncer_sync_up(temp_data_dirs, syncer):
+    """Check that syncing up works"""
+    tmp_source, tmp_target = temp_data_dirs
+
+    syncer.sync_up(local_dir=tmp_source, remote_dir="/test/test_syncer_sync_up_down")
+    syncer.wait()
+
+    _download_from_fs_path(
+        syncer.storage_filesystem, "/test/test_syncer_sync_up_down", tmp_target
+    )
+
+    # Target dir should have all files
+    assert_file(True, tmp_target, "level0.txt")
+    assert_file(True, tmp_target, "level0_exclude.txt")
+    assert_file(True, tmp_target, "subdir/level1.txt")
+    assert_file(True, tmp_target, "subdir/level1_exclude.txt")
+    assert_file(True, tmp_target, "subdir/nested/level2.txt")
+    assert_file(True, tmp_target, "subdir_nested_level2_exclude.txt")
+    assert_file(True, tmp_target, "subdir_exclude/something/somewhere.txt")
+
+
+def test_syncer_sync_exclude(temp_data_dirs, syncer):
+    """Check that the exclude parameter works"""
+    tmp_source, tmp_target = temp_data_dirs
+
+    syncer.sync_up(
+        local_dir=tmp_source,
+        remote_dir="/test/test_syncer_sync_exclude",
+        exclude=["*_exclude*"],
+    )
+    syncer.wait()
+
+    _download_from_fs_path(
+        syncer.storage_filesystem, "/test/test_syncer_sync_exclude", tmp_target
+    )
+
+    # Excluded files should not be found in target
+    assert_file(True, tmp_target, "level0.txt")
+    assert_file(False, tmp_target, "level0_exclude.txt")
+    assert_file(True, tmp_target, "subdir/level1.txt")
+    assert_file(False, tmp_target, "subdir/level1_exclude.txt")
+    assert_file(True, tmp_target, "subdir/nested/level2.txt")
+    assert_file(False, tmp_target, "subdir_nested_level2_exclude.txt")
+    assert_file(False, tmp_target, "subdir_exclude/something/somewhere.txt")
+
+
+def test_sync_up_if_needed(temp_data_dirs, tmp_path):
+    """Check that we only sync up again after sync period"""
+    tmp_source, tmp_target = temp_data_dirs
+
+    with freeze_time() as frozen:
+        syncer = _FilesystemSyncer(
+            storage_filesystem=_create_mock_custom_fs(tmp_path), sync_period=60
+        )
+
+        assert syncer.sync_up_if_needed(
+            local_dir=tmp_source, remote_dir="/test/test_sync_up_not_needed"
+        )
+        syncer.wait()
+
+        frozen.tick(30)
+
+        # Sync period not over, yet
+        assert not syncer.sync_up_if_needed(
+            local_dir=tmp_source, remote_dir="/test/test_sync_up_not_needed"
+        )
+
+        frozen.tick(30)
+
+        # Sync period over, sync again
+        assert syncer.sync_up_if_needed(
+            local_dir=tmp_source, remote_dir="/test/test_sync_up_not_needed"
+        )
+
+
+def test_syncer_still_running_no_sync(temp_data_dirs, tmp_path):
+    """Check that no new sync is issued if old sync is still running"""
+    tmp_source, tmp_target = temp_data_dirs
+
+    class FakeSyncProcess:
+        @property
+        def is_running(self):
+            return True
+
+        @property
+        def start_time(self):
+            # Don't consider the sync process timeout
+            return float("inf")
+
+    syncer = _FilesystemSyncer(
+        storage_filesystem=_create_mock_custom_fs(tmp_path), sync_period=60
+    )
+    syncer._sync_process = FakeSyncProcess()
+    assert not syncer.sync_up_if_needed(
+        local_dir=tmp_source,
+        remote_dir="/test/test_syncer_still_running_no_sync",
+    )
+
+
+def test_syncer_not_running_sync(temp_data_dirs, tmp_path):
+    """Check that new sync is issued if old sync completed"""
+    tmp_source, tmp_target = temp_data_dirs
+
+    class FakeSyncProcess:
+        @property
+        def is_running(self):
+            return False
+
+        def wait(self):
+            return True
+
+    syncer = _FilesystemSyncer(
+        storage_filesystem=_create_mock_custom_fs(tmp_path), sync_period=60
+    )
+    syncer._sync_process = FakeSyncProcess()
+    assert syncer.sync_up_if_needed(
+        local_dir=tmp_source,
+        remote_dir="/test/test_syncer_not_running_sync",
+    )
+
+
+def test_syncer_hanging_sync_with_timeout(temp_data_dirs, tmp_path):
+    """Check that syncing times out when the sync process is hanging."""
+    tmp_source, tmp_target = temp_data_dirs
+
+    def _hanging_sync_up_command(*args, **kwargs):
+        time.sleep(200)
+
+    class _HangingSyncer(_FilesystemSyncer):
+        def _sync_up_command(
+            self, local_path: str, uri: str, exclude: Optional[List] = None
+        ):
+            return _hanging_sync_up_command, {}
+
+    syncer = _HangingSyncer(
+        storage_filesystem=_create_mock_custom_fs(tmp_path),
+        sync_period=60,
+        sync_timeout=10,
+    )
+
+    def sync_up():
+        return syncer.sync_up(
+            local_dir=tmp_source, remote_dir="/test/test_syncer_timeout"
+        )
+
+    with freeze_time() as frozen:
+        assert sync_up()
+        frozen.tick(5)
+        # 5 seconds - initial sync hasn't reached the timeout yet
+        # It should continue running without launching a new sync
+        assert not sync_up()
+        frozen.tick(5)
+        # Reached the timeout - start running a new sync command
+        assert sync_up()
+        frozen.tick(20)
+        # We're 10 seconds past the timeout, waiting should result in a timeout error
+        with pytest.raises(TimeoutError):
+            syncer.wait()
+
+
+def test_syncer_not_running_sync_last_failed(
+    propagate_logs, caplog, temp_data_dirs, tmp_path
+):
+    """Check that new sync is issued if old sync completed"""
+    caplog.set_level(logging.WARNING)
+
+    tmp_source, tmp_target = temp_data_dirs
+
+    class FakeSyncProcess(_BackgroundProcess):
+        @property
+        def is_running(self):
+            return False
+
+        def wait(self, *args, **kwargs):
+            raise RuntimeError("Sync failed")
+
+    syncer = _FilesystemSyncer(
+        storage_filesystem=_create_mock_custom_fs(tmp_path), sync_period=60
+    )
+    syncer._sync_process = FakeSyncProcess(lambda: None)
+    assert syncer.sync_up_if_needed(
+        local_dir=tmp_source,
+        remote_dir="/test/test_syncer_not_running_sync",
+    )
+    assert "Last sync command failed" in caplog.text
+
+
+def test_syncer_wait_or_retry_failure(temp_data_dirs, tmp_path):
+    """Check that the wait or retry API fails after max_retries."""
+    tmp_source, tmp_target = temp_data_dirs
+
+    syncer = _FilesystemSyncer(storage_filesystem=lambda: "error", sync_period=60)
+
+    # Will fail since the storage filesystem is invalid
+    syncer.sync_up(local_dir=tmp_source, remote_dir="/test/test_syncer_wait_or_retry")
+    with pytest.raises(RuntimeError) as e:
+        syncer.wait_or_retry(max_retries=3, backoff_s=0)
+
+    assert "Failed sync even after 3 retries." in str(e.value)
+
+
+def test_syncer_wait_or_retry_timeout(temp_data_dirs, tmp_path):
+    """Check that the wait or retry API raises a timeout error after `sync_timeout`."""
+    tmp_source, tmp_target = temp_data_dirs
+
+    def slow_upload(*args, **kwargs):
+        time.sleep(5)
+
+    class HangingSyncer(_FilesystemSyncer):
+        def _sync_up_command(
+            self, local_path: str, uri: str, exclude: Optional[List] = None
+        ):
+            return (
+                slow_upload,
+                dict(local_path=local_path, uri=uri, exclude=exclude),
+            )
+
+    syncer = HangingSyncer(
+        storage_filesystem=_create_mock_custom_fs(tmp_path),
+        sync_period=60,
+        sync_timeout=0.1,
+    )
+
+    syncer.sync_up(local_dir=tmp_source, remote_dir="/test/timeout")
+    with pytest.raises(RuntimeError) as e:
+        syncer.wait_or_retry(max_retries=3, backoff_s=0)
+        assert "Failed sync even after 3 retries." in str(e.value)
+        assert isinstance(e.value.__cause__, TimeoutError)
+
+
+def test_syncer_wait_or_retry_eventual_success(temp_data_dirs, tmp_path):
+    """Check that the wait or retry API succeeds for a sync_down that
+    fails, times out, then succeeds."""
+    tmp_source, tmp_target = temp_data_dirs
+
+    success = tmp_path / "success"
+    fail_marker = tmp_path / "fail_marker"
+    hang_marker = tmp_path / "hang_marker"
+
+    def eventual_upload(*args, **kwargs):
+        if not fail_marker.exists():
+            fail_marker.write_text(".", encoding="utf-8")
+            raise RuntimeError("Failing")
+        elif not hang_marker.exists():
+            hang_marker.write_text(".", encoding="utf-8")
+            time.sleep(5)
+        else:
+            success.write_text(".", encoding="utf-8")
+
+    class EventualSuccessSyncer(_FilesystemSyncer):
+        def _sync_up_command(
+            self, local_path: str, uri: str, exclude: Optional[List] = None
+        ):
+            return (
+                eventual_upload,
+                dict(local_path=local_path, uri=uri, exclude=exclude),
+            )
+
+    syncer = EventualSuccessSyncer(
+        storage_filesystem=_create_mock_custom_fs(tmp_path),
+        sync_period=60,
+        sync_timeout=0.5,
+    )
+
+    syncer.sync_up(local_dir=tmp_source, remote_dir="/test/eventual_success")
+    # The syncer will retry 2 times, running 3 times in total and eventually succeeding.
+    syncer.wait_or_retry(max_retries=2, backoff_s=0)
+    assert success.exists()
+
+
+def test_syncer_serialize(temp_data_dirs, syncer):
+    tmp_source, tmp_target = temp_data_dirs
+
+    syncer.sync_up(local_dir=tmp_source, remote_dir="/test/serialize")
+
+    serialized = pickle.dumps(syncer)
+    loaded_syncer = pickle.loads(serialized)
+    assert not loaded_syncer._sync_process
+
+
+def test_sync_many_files_local_to_cloud(mock_s3_bucket_uri, tmp_path):
+    source_dir = tmp_path / "source"
+    check_dir = tmp_path / "check"
+    source_dir.mkdir()
+    check_dir.mkdir()
+
+    # Create 256 files to upload
+    for i in range(256):
+        (source_dir / str(i)).write_text("", encoding="utf-8")
+
+    fs, fs_path = get_fs_and_path(mock_s3_bucket_uri)
+    _upload_to_fs_path(source_dir, fs, fs_path)
+    _download_from_fs_path(fs, fs_path, check_dir)
+    assert (check_dir / "255").exists()
+
+
+def test_sync_many_files_local_to_local(tmp_path):
+    (tmp_path / "source").mkdir()
+    # Create 256 files to upload
+    for i in range(256):
+        (tmp_path / "source" / str(i)).write_text("", encoding="utf-8")
+
+    fs, fs_path = get_fs_and_path(str(tmp_path / "destination"))
+    _upload_to_fs_path(str(tmp_path / "source"), fs, fs_path)
+    assert (tmp_path / "destination" / "255").exists()
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/tune/tests/test_telemetry.py

@@ -0,0 +1,54 @@
+import sys
+
+import pytest
+
+import ray
+import ray._common.usage.usage_lib as ray_usage_lib
+from ray import tune
+from ray._common.test_utils import TelemetryCallsite, check_library_usage_telemetry
+
+
+@pytest.fixture
+def reset_usage_lib():
+    yield
+    ray.shutdown()
+    ray_usage_lib.reset_global_state()
+
+
+@pytest.mark.parametrize("callsite", list(TelemetryCallsite))
+def test_not_used_on_import(reset_usage_lib, callsite: TelemetryCallsite):
+    if callsite in {TelemetryCallsite.ACTOR, TelemetryCallsite.TASK}:
+        pytest.skip("TODO: train usage is exported when importing in an actor or task.")
+
+    def _import_ray_tune():
+        from ray import tune  # noqa: F401
+
+    check_library_usage_telemetry(
+        _import_ray_tune, callsite=callsite, expected_library_usages=[set(), {"core"}]
+    )
+
+
+@pytest.mark.parametrize("callsite", list(TelemetryCallsite))
+def test_used_on_tuner_fit(reset_usage_lib, callsite: TelemetryCallsite):
+    def _call_tuner_fit():
+        def objective(*args):
+            pass
+
+        tuner = tune.Tuner(objective)
+        tuner.fit()
+
+    check_library_usage_telemetry(
+        _call_tuner_fit,
+        callsite=callsite,
+        expected_library_usages=[{"tune"}, {"core", "tune"}],
+        expected_extra_usage_tags={
+            "tune_scheduler": "FIFOScheduler",
+            "tune_searcher": "BasicVariantGenerator",
+            "air_entrypoint": "Tuner.fit",
+            "air_storage_configuration": "local",
+        },
+    )
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", "-s", __file__]))

python/ray/tune/tests/test_train_v2_integration.py

@@ -0,0 +1,139 @@
+import sys
+
+import pytest
+
+import ray.train
+import ray.tune
+from ray.cluster_utils import Cluster
+from ray.train.tests.util import create_dict_checkpoint
+from ray.train.v2._internal.constants import HEALTH_CHECK_INTERVAL_S_ENV_VAR
+from ray.train.v2.api.data_parallel_trainer import DataParallelTrainer
+from ray.tune.integration.ray_train import CHECKPOINT_PATH_KEY, TuneReportCallback
+
+TRAIN_DRIVER_RESOURCE_NAME = "train_driver_resource"
+NUM_GPUS_IN_CLUSTER = 4
+
+
+@pytest.fixture()
+def ray_start_4_cpus():
+    ray.init(num_cpus=4)
+    yield
+    ray.shutdown()
+
+
+@pytest.fixture()
+def ray_cpu_head_gpu_worker():
+    cluster = Cluster()
+    cluster.add_node(resources={TRAIN_DRIVER_RESOURCE_NAME: 1})
+    cluster.add_node(num_cpus=0, num_gpus=NUM_GPUS_IN_CLUSTER)
+
+    ray.init(address=cluster.address)
+
+    yield
+
+    ray.shutdown()
+    cluster.shutdown()
+
+
+@pytest.fixture(autouse=True)
+def speed_up_tests(monkeypatch):
+    monkeypatch.setenv(HEALTH_CHECK_INTERVAL_S_ENV_VAR, "0.1")
+
+
+@pytest.mark.parametrize("num_workers_grid_search", [[1], [1, 2, 4]])
+@pytest.mark.parametrize("limit_concurrency", [True, False])
+def test_e2e(
+    ray_cpu_head_gpu_worker,
+    tmp_path,
+    num_workers_grid_search,
+    limit_concurrency,
+):
+    num_non_checkpoint_reports = 2
+    num_checkpoint_reports = 1
+
+    def train_fn_per_worker(train_fn_config):
+        assert "lr" in train_fn_config
+
+        world_size = ray.train.get_context().get_world_size()
+        for i in range(num_non_checkpoint_reports):
+            ray.train.report({"idx": i})
+
+        for i in range(num_checkpoint_reports):
+            with create_dict_checkpoint({"model": "dummy"}) as checkpoint:
+                ray.train.report(
+                    {"loss": 0.1, "world_size": world_size}, checkpoint=checkpoint
+                )
+
+    def launch_training(tune_config):
+        trainer = DataParallelTrainer(
+            train_loop_per_worker=train_fn_per_worker,
+            train_loop_config=tune_config["train_loop_config"],
+            scaling_config=ray.train.ScalingConfig(
+                num_workers=tune_config["num_workers"], use_gpu=True
+            ),
+            run_config=ray.train.RunConfig(
+                storage_path=tmp_path,
+                name=f"train-{ray.tune.get_context().get_trial_id()}",
+                callbacks=[TuneReportCallback()],
+            ),
+        )
+        trainer.fit()
+
+    tuner = ray.tune.Tuner(
+        ray.tune.with_resources(launch_training, {TRAIN_DRIVER_RESOURCE_NAME: 0.01}),
+        param_space={
+            # Search over parameters passed into each train worker.
+            "train_loop_config": {"lr": ray.tune.choice([0.01, 0.001])},
+            # Search over Train "run level" parameters.
+            "num_workers": ray.tune.grid_search(num_workers_grid_search),
+        },
+        tune_config=ray.tune.TuneConfig(
+            max_concurrent_trials=(
+                NUM_GPUS_IN_CLUSTER // max(num_workers_grid_search)
+                if limit_concurrency
+                else None
+            )
+        ),
+        run_config=ray.tune.RunConfig(storage_path=tmp_path, name="tune"),
+    )
+    result_grid = tuner.fit()
+    assert len(result_grid) == len(num_workers_grid_search)
+
+    world_sizes = set()
+    for result in result_grid:
+        assert (
+            len(result.metrics_dataframe)
+            == num_non_checkpoint_reports + num_checkpoint_reports
+        )
+        assert "loss" in result.metrics
+        assert CHECKPOINT_PATH_KEY in result.metrics
+        world_sizes.add(result.metrics["world_size"])
+    assert world_sizes == set(num_workers_grid_search)
+
+
+def test_errors(ray_start_4_cpus):
+    """Test that errors in training are properly captured and reported."""
+
+    def train_worker_fn():
+        raise RuntimeError("Simulated training error")
+
+    def train_fn(config):
+        trainer = DataParallelTrainer(train_worker_fn)
+        trainer.fit()
+
+    tuner = ray.tune.Tuner(train_fn)
+
+    results = tuner.fit()
+
+    assert results.errors, "Expected errors to be captured"
+    assert len(results.errors) == 1, "Expected exactly one error"
+
+    error = results.errors[0]
+    assert "RuntimeError" in str(error), f"Expected RuntimeError, got: {error}"
+    assert "Simulated training error" in str(
+        error
+    ), f"Expected specific error message, got: {error}"
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/tune/tests/test_trainable.py

@@ -0,0 +1,134 @@
+import json
+import os
+from typing import Dict, Union
+
+import pytest
+
+import ray
+from ray import tune
+from ray.train._internal.storage import StorageContext
+from ray.train.tests.util import create_dict_checkpoint
+from ray.tune.trainable import wrap_function
+
+
+@pytest.fixture
+def ray_start_2_cpus():
+    address_info = ray.init(num_cpus=2)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+class SavingTrainable(tune.Trainable):
+    def __init__(self, return_type: str, *args, **kwargs):
+        self.return_type = return_type
+        super(SavingTrainable, self).__init__(*args, **kwargs)
+
+    def step(self):
+        return {"iter": self.training_iteration}
+
+    def save_checkpoint(self, tmp_checkpoint_dir: str):
+        checkpoint_data = {"data": 1}
+
+        if self.return_type == "object":
+            return checkpoint_data
+
+        subdir = os.path.join(tmp_checkpoint_dir, "subdir")
+        os.makedirs(subdir, exist_ok=True)
+        checkpoint_file = os.path.join(subdir, "checkpoint.pkl")
+        with open(checkpoint_file, "w") as f:
+            f.write(json.dumps(checkpoint_data))
+
+        if self.return_type == "root":
+            return tmp_checkpoint_dir
+        elif self.return_type == "subdir":
+            return subdir
+        elif self.return_type == "checkpoint":
+            return checkpoint_file
+
+    def load_checkpoint(self, checkpoint: Union[Dict, str]):
+        if self.return_type == "object":
+            assert isinstance(checkpoint, dict)
+            checkpoint_data = checkpoint
+            checkpoint_file = None
+        elif self.return_type == "root":
+            assert "subdir" not in checkpoint
+            checkpoint_file = os.path.join(checkpoint, "subdir", "checkpoint.pkl")
+        elif self.return_type == "subdir":
+            assert "subdir" in checkpoint
+            assert "checkpoint.pkl" not in checkpoint
+            checkpoint_file = os.path.join(checkpoint, "checkpoint.pkl")
+        else:  # self.return_type == "checkpoint"
+            assert checkpoint.endswith("subdir/checkpoint.pkl")
+            checkpoint_file = checkpoint
+
+        if checkpoint_file:
+            with open(checkpoint_file, "rb") as f:
+                checkpoint_data = json.load(f)
+
+        checkpoint_data = {
+            key: value
+            for key, value in checkpoint_data.items()
+            if not key.startswith("_")
+        }
+        assert checkpoint_data == {"data": 1}, checkpoint_data
+
+
+def function_trainable(config):
+    with create_dict_checkpoint({"checkpoint_data": 5}) as checkpoint:
+        tune.report({"metric": 4}, checkpoint=checkpoint)
+
+
+@pytest.mark.parametrize("return_type", ["object", "root"])
+def test_save_load_checkpoint_path_class(ray_start_2_cpus, return_type, tmpdir):
+    """Assert that restoring from a Trainable.save() future works with
+    class trainables.
+
+    Needs Ray cluster so we get actual futures.
+    """
+    trainable = ray.remote(SavingTrainable).remote(return_type=return_type)
+
+    # Train one step
+    ray.get(trainable.train.remote())
+
+    # Save checkpoint
+    saving_future = trainable.save.remote()
+
+    # Check for errors
+    ray.get(saving_future)
+
+    restoring_future = trainable.restore.remote(saving_future)
+
+    ray.get(restoring_future)
+
+
+def test_save_load_checkpoint_path_fn(ray_start_2_cpus, tmp_path):
+    """Assert that restoring from a Trainable.save() future works with
+    function trainables.
+
+    Needs Ray cluster so we get actual futures.
+    """
+    trainable_cls = wrap_function(function_trainable)
+    trainable = ray.remote(trainable_cls).remote(
+        storage=StorageContext(
+            storage_path=str(tmp_path),
+            experiment_dir_name="exp",
+            trial_dir_name="trial",
+        )
+    )
+    ray.get(trainable.train.remote())
+
+    saving_future = trainable.save.remote()
+
+    # Check for errors
+    ray.get(saving_future)
+
+    restoring_future = trainable.restore.remote(saving_future)
+
+    ray.get(restoring_future)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/tune/tests/test_trainable_util.py

@@ -0,0 +1,184 @@
+import copy
+import sys
+import unittest
+from collections import OrderedDict
+from unittest.mock import patch
+
+import pytest
+
+from ray.tune.utils.util import (
+    flatten_dict,
+    unflatten_dict,
+    unflatten_list_dict,
+    wait_for_gpu,
+)
+
+
+class FlattenDictTest(unittest.TestCase):
+    def test_output_type(self):
+        in_ = OrderedDict({"a": {"b": 1}, "c": {"d": 2}, "e": 3})
+        out = flatten_dict(in_)
+        assert type(in_) is type(out)
+
+    def test_one_level_nested(self):
+        ori_in = OrderedDict({"a": {"b": 1}, "c": {"d": 2}, "e": 3})
+        in_ = copy.deepcopy(ori_in)
+        result = flatten_dict(in_)
+        assert in_ == ori_in
+        assert result == {"a/b": 1, "c/d": 2, "e": 3}
+
+    def test_multi_level_nested(self):
+        ori_in = OrderedDict(
+            {
+                "a": {
+                    "b": {
+                        "c": {
+                            "d": 1,
+                        },
+                    },
+                },
+                "b": {
+                    "c": {
+                        "d": 2,
+                    },
+                },
+                "c": {
+                    "d": 3,
+                },
+                "e": 4,
+            }
+        )
+        in_ = copy.deepcopy(ori_in)
+        result = flatten_dict(in_)
+        assert in_ == ori_in
+        assert result == {"a/b/c/d": 1, "b/c/d": 2, "c/d": 3, "e": 4}
+
+
+class UnflattenDictTest(unittest.TestCase):
+    def test_output_type(self):
+        in_ = OrderedDict({"a/b": 1, "c/d": 2, "e": 3})
+        out = unflatten_dict(in_)
+        assert type(in_) is type(out)
+
+    def test_one_level_nested(self):
+        result = unflatten_dict({"a/b": 1, "c/d": 2, "e": 3})
+        assert result == {"a": {"b": 1}, "c": {"d": 2}, "e": 3}
+
+    def test_multi_level_nested(self):
+        result = unflatten_dict({"a/b/c/d": 1, "b/c/d": 2, "c/d": 3, "e": 4})
+        assert result == {
+            "a": {
+                "b": {
+                    "c": {
+                        "d": 1,
+                    },
+                },
+            },
+            "b": {
+                "c": {
+                    "d": 2,
+                },
+            },
+            "c": {
+                "d": 3,
+            },
+            "e": 4,
+        }
+
+    def test_unflatten_list_dict_output_type(self):
+        in_ = OrderedDict({"a/0": 0, "a/1": 1, "c/d": 2, "e": 3})
+        out = unflatten_list_dict(in_)
+        assert type(out) is OrderedDict
+
+        in_ = OrderedDict({"0/a": 0, "1/b": 1, "2/c": 2, "3/d": 3})
+        out = unflatten_list_dict(in_)
+        assert type(out) is list
+
+    def test_unflatten_list_dict_one_level_nested(self):
+        result = unflatten_list_dict({"a/0": 0, "a/1": 1, "c/d": 2, "e": 3})
+        assert result == {"a": [0, 1], "c": {"d": 2}, "e": 3}
+
+        result = unflatten_list_dict({"0/a": 0, "1/b": 1, "2/c": 2, "3": 3})
+        assert result == [{"a": 0}, {"b": 1}, {"c": 2}, 3]
+
+    def test_unflatten_list_dict_multi_level_nested(self):
+        result = unflatten_list_dict({"a/0/c/d": 1, "a/1/c": 2, "a/2": 3, "e": 4})
+        assert result == {"a": [{"c": {"d": 1}}, {"c": 2}, 3], "e": 4}
+
+        result = unflatten_list_dict(
+            {"0/a/0/b": 1, "0/a/1": 2, "1/0": 3, "1/1": 4, "1/2/c": 5, "2": 6}
+        )
+        assert result == [{"a": [{"b": 1}, 2]}, [3, 4, {"c": 5}], 6]
+
+    def test_unflatten_noop(self):
+        """Unflattening an already unflattened dict should be a noop."""
+        unflattened = {"a": 1, "b": {"c": {"d": [1, 2]}, "e": 3}, "f": {"g": 3}}
+        assert unflattened == unflatten_dict(unflattened)
+        assert unflattened == unflatten_list_dict(unflattened)
+
+    def test_raises_error_on_key_conflict(self):
+        """Ensure that an informative exception is raised on key conflict."""
+        with self.assertRaisesRegex(TypeError, r"Cannot unflatten dict"):
+            unflatten_dict({"a": 1, "a/b": 2, "a/c": 3})
+
+        with self.assertRaisesRegex(TypeError, r"Cannot unflatten dict"):
+            unflatten_dict({"a/b": 2, "a/b/c": 3})
+
+
+class GPUUtilMock:
+    class GPU:
+        def __init__(self, id, uuid, util=None):
+            self.id = id
+            self.uuid = uuid
+            self.util = [0.5, 0.0]
+
+        @property
+        def memoryUtil(self):
+            if self.util:
+                return self.util.pop(0)
+            return 0
+
+    def __init__(self, gpus, gpu_uuids):
+        self.gpus = gpus
+        self.uuids = gpu_uuids
+        self.gpu_list = [
+            self.GPU(gpu, uuid) for gpu, uuid in zip(self.gpus, self.uuids)
+        ]
+
+    def getGPUs(self):
+        return self.gpu_list
+
+
+class GPUTest(unittest.TestCase):
+    def setUp(self):
+        sys.modules["GPUtil"] = GPUUtilMock([0, 1], ["GPU-aaa", "GPU-bbb"])
+
+    def testGPUWait1(self):
+        wait_for_gpu(0, delay_s=0)
+
+    def testGPUWait2(self):
+        wait_for_gpu("1", delay_s=0)
+
+    def testGPUWait3(self):
+        wait_for_gpu("GPU-aaa", delay_s=0)
+
+    def testGPUWaitFail(self):
+        with self.assertRaises(ValueError):
+            wait_for_gpu(2, delay_s=0)
+
+        with self.assertRaises(ValueError):
+            wait_for_gpu("4", delay_s=0)
+
+        with self.assertRaises(ValueError):
+            wait_for_gpu(1.23, delay_s=0)
+
+    @patch("ray.get_gpu_ids", lambda: ["0"])
+    def testDefaultGPU(self):
+        import sys
+
+        sys.modules["GPUtil"] = GPUUtilMock([0], ["GPU-aaa"])
+        wait_for_gpu(delay_s=0)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/tune/tests/test_trial.py

@@ -0,0 +1,156 @@
+import logging
+import sys
+
+import pytest
+
+from ray.exceptions import RayActorError, RayTaskError
+from ray.tests.conftest import propagate_logs  # noqa
+from ray.train._internal.session import _TrainingResult
+from ray.train._internal.storage import StorageContext
+from ray.train.constants import RAY_TRAIN_COUNT_PREEMPTION_AS_FAILURE
+from ray.train.tests.util import mock_storage_context
+from ray.tune import Checkpoint
+from ray.tune.experiment import Trial
+
+
+@pytest.fixture
+def trial(tmp_path):
+    yield Trial(
+        "mock",
+        stub=True,
+        storage=mock_storage_context(storage_path=str(tmp_path)),
+    )
+
+
+@pytest.mark.parametrize("count_preemption_errors", [False, True])
+def test_handle_preemption_error(
+    trial: Trial, count_preemption_errors: bool, monkeypatch
+):
+    """Check that the Trial counts preemption errors correctly."""
+    if count_preemption_errors:
+        monkeypatch.setenv(RAY_TRAIN_COUNT_PREEMPTION_AS_FAILURE, "1")
+
+    # Case 1: Directly raised (preemption) RayActorError
+    class PreemptionRayActorError(RayActorError):
+        def preempted(self) -> bool:
+            return True
+
+    err = PreemptionRayActorError()
+    trial.handle_error(err)
+    assert trial.num_failures == (1 if count_preemption_errors else 0)
+
+    # Case 2: RayTaskError, where the cause is a (preemption) RayActorError
+    wrapped_err = RayTaskError(
+        function_name="test", traceback_str="traceback_str", cause=err
+    )
+    trial.handle_error(wrapped_err)
+    assert trial.num_failures == (2 if count_preemption_errors else 0)
+
+    # Case 3: Non-preemption error
+    non_preempted_err = RayActorError()
+    trial.handle_error(non_preempted_err)
+    assert trial.num_failures == (3 if count_preemption_errors else 1)
+
+
+def test_load_trial_from_json_state():
+    """Check that serializing a trial to a JSON string with `Trial.get_json_state`
+    and then creating a new trial using the `Trial.from_json_state` alternate
+    constructor loads the trial with equivalent state."""
+    trial = Trial(
+        "MockTrainable",
+        stub=True,
+        trial_id="abcd1234",
+        storage=mock_storage_context(),
+    )
+    trial.create_placement_group_factory()
+    trial.init_local_path()
+    trial.status = Trial.TERMINATED
+
+    # After loading, the trial state should be the same
+    json_state, _ = trial.get_json_state()
+    new_trial = Trial.from_json_state(json_state, stub=True)
+    assert new_trial.get_json_state()[0] == json_state
+
+
+def test_set_storage(tmp_path):
+    """Test that setting the trial's storage context will update the tracked
+    checkpoint paths."""
+    original_storage = mock_storage_context()
+    trial = Trial(
+        "MockTrainable",
+        stub=True,
+        trial_id="abcd1234",
+        storage=original_storage,
+    )
+
+    result_1 = _TrainingResult(
+        checkpoint=Checkpoint.from_directory(original_storage.checkpoint_fs_path),
+        metrics={},
+    )
+    trial.on_checkpoint(result_1)
+
+    result_2 = _TrainingResult(
+        checkpoint=Checkpoint.from_directory(original_storage.checkpoint_fs_path),
+        metrics={},
+    )
+    trial.on_checkpoint(result_2)
+
+    new_storage = StorageContext(
+        storage_path=tmp_path / "new_storage_path",
+        experiment_dir_name="new_name",
+        trial_dir_name="new_trial",
+    )
+    trial.set_storage(new_storage)
+
+    assert result_1.checkpoint.path.startswith(new_storage.trial_fs_path)
+    assert result_2.checkpoint.path.startswith(new_storage.trial_fs_path)
+
+
+def test_trial_logdir_length():
+    """Test that trial local paths with a long logdir are truncated"""
+    trial = Trial(
+        trainable_name="none",
+        stub=True,
+        config={"a" * 50: 5.0 / 7, "b" * 50: "long" * 40},
+        storage=mock_storage_context(),
+    )
+    trial.init_local_path()
+    assert len(trial.storage.trial_dir_name) < 200
+
+
+def test_should_stop(caplog, propagate_logs):  # noqa
+    """Test whether `Trial.should_stop()` works as expected given a result dict."""
+    trial = Trial(
+        "MockTrainable",
+        stub=True,
+        trial_id="abcd1234",
+        stopping_criterion={"a": 10.0, "b/c": 20.0},
+    )
+
+    # Criterion is not reached yet -> don't stop.
+    result = {"a": 9.999, "b/c": 0.0, "some_other_key": True}
+    assert not trial.should_stop(result)
+
+    # Criterion is exactly reached -> stop.
+    result = {"a": 10.0, "b/c": 0.0, "some_other_key": False}
+    assert trial.should_stop(result)
+
+    # Criterion is exceeded -> stop.
+    result = {"a": 10000.0, "b/c": 0.0, "some_other_key": False}
+    assert trial.should_stop(result)
+
+    # Test nested criterion.
+    result = {"a": 5.0, "b/c": 1000.0, "some_other_key": False}
+    assert trial.should_stop(result)
+
+    # Test criterion NOT found in result metrics.
+    result = {"b/c": 1000.0}
+    with caplog.at_level(logging.WARNING):
+        trial.should_stop(result)
+    assert (
+        "Stopping criterion 'a' not found in result dict! Available keys are ['b/c']."
+    ) in caplog.text
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/tune/tests/test_trial_scheduler_resource_changing.py

@@ -0,0 +1,619 @@
+import shutil
+import tempfile
+import unittest
+
+from ray.train.tests.util import mock_storage_context
+from ray.tune import PlacementGroupFactory
+from ray.tune.execution.tune_controller import TuneController
+from ray.tune.experiment import Trial
+from ray.tune.schedulers.resource_changing_scheduler import (
+    DistributeResources,
+    DistributeResourcesToTopJob,
+    ResourceChangingScheduler,
+)
+from ray.tune.schedulers.trial_scheduler import TrialScheduler
+from ray.tune.tests.execution.utils import create_execution_test_objects
+
+
+class MockTuneController(TuneController):
+    def get_live_trials(self):
+        return [t for t in self._trials if t.status != "TERMINATED"]
+
+
+class TestUniformResourceAllocation(unittest.TestCase):
+    def setUp(self):
+        self.tmpdir = tempfile.mkdtemp()
+        self.tune_controller, *_ = create_execution_test_objects(
+            resources={"CPU": 8, "GPU": 8},
+            reuse_actors=False,
+            tune_controller_cls=MockTuneController,
+            storage=mock_storage_context(),
+        )
+
+    def tearDown(self) -> None:
+        shutil.rmtree(self.tmpdir)
+
+    def _prepareTrials(self, scheduler, base_pgf):
+        trial1 = Trial("mock", config=dict(num=1), stub=True)
+        trial1.placement_group_factory = base_pgf
+        trial2 = Trial("mock", config=dict(num=2), stub=True)
+        trial2.placement_group_factory = base_pgf
+        trial3 = Trial("mock", config=dict(num=3), stub=True)
+        trial3.placement_group_factory = base_pgf
+        trial4 = Trial("mock", config=dict(num=4), stub=True)
+        trial4.placement_group_factory = base_pgf
+
+        self.tune_controller._trials = [trial1, trial2, trial3, trial4]
+
+        scheduler.on_trial_add(self.tune_controller, trial1)
+        scheduler.on_trial_add(self.tune_controller, trial2)
+        scheduler.on_trial_add(self.tune_controller, trial3)
+        scheduler.on_trial_add(self.tune_controller, trial4)
+
+        trial1.status = Trial.RUNNING
+        trial2.status = Trial.RUNNING
+        trial3.status = Trial.RUNNING
+        trial4.status = Trial.RUNNING
+        return trial1, trial2, trial3, trial4
+
+    def _allocateAndAssertNewResources(self, trial, scheduler, target_pgf, metric=1):
+        result = {"metric": metric, "training_iteration": 4}
+        trial.run_metadata.last_result = result
+        decision = scheduler.on_trial_result(self.tune_controller, trial, result)
+        assert decision == TrialScheduler.PAUSE
+        trial.status = Trial.PENDING
+        scheduler.choose_trial_to_run(self.tune_controller)
+        assert trial.placement_group_factory == target_pgf
+        trial.status = Trial.RUNNING
+
+    def testAllocateFreeResources(self):
+        scheduler = ResourceChangingScheduler(
+            resources_allocation_function=DistributeResources(add_bundles=False)
+        )
+
+        base_pgf = PlacementGroupFactory([{"CPU": 1, "GPU": 0}])
+        trial1, trial2, trial3, trial4 = self._prepareTrials(scheduler, base_pgf)
+
+        self._allocateAndAssertNewResources(
+            trial1, scheduler, PlacementGroupFactory([{"CPU": 2}])
+        )
+        self._allocateAndAssertNewResources(
+            trial2, scheduler, PlacementGroupFactory([{"CPU": 2}])
+        )
+
+        trial4.status = Trial.TERMINATED
+
+        self._allocateAndAssertNewResources(
+            trial1, scheduler, PlacementGroupFactory([{"CPU": 3}])
+        )
+
+        trial3.status = Trial.TERMINATED
+
+        self._allocateAndAssertNewResources(
+            trial1, scheduler, PlacementGroupFactory([{"CPU": 4}])
+        )
+
+        trial2.status = Trial.TERMINATED
+
+        self._allocateAndAssertNewResources(
+            trial1, scheduler, PlacementGroupFactory([{"CPU": 8}])
+        )
+
+    def testAllocateFreeResourcesWithIncreaseBy(self):
+        scheduler = ResourceChangingScheduler(
+            resources_allocation_function=DistributeResources(
+                add_bundles=False, increase_by={"CPU": 2, "GPU": 2}
+            )
+        )
+
+        base_pgf = PlacementGroupFactory([{"CPU": 2, "GPU": 2}])
+        trial1, trial2, trial3, trial4 = self._prepareTrials(scheduler, base_pgf)
+
+        decision = scheduler.on_trial_result(
+            self.tune_controller, trial1, {"metric": 1, "training_iteration": 4}
+        )
+        assert decision == TrialScheduler.CONTINUE
+
+        trial4.status = Trial.TERMINATED
+
+        self._allocateAndAssertNewResources(
+            trial1, scheduler, PlacementGroupFactory([{"CPU": 4, "GPU": 4}])
+        )
+
+        trial3.status = Trial.TERMINATED
+
+        self._allocateAndAssertNewResources(
+            trial2, scheduler, PlacementGroupFactory([{"CPU": 4, "GPU": 4}])
+        )
+
+        trial2.status = Trial.TERMINATED
+
+        self._allocateAndAssertNewResources(
+            trial1, scheduler, PlacementGroupFactory([{"CPU": 8, "GPU": 8}])
+        )
+
+    def testAllocateFreeResourcesWithIncreaseByTimes(self):
+        scheduler = ResourceChangingScheduler(
+            resources_allocation_function=DistributeResources(
+                add_bundles=False, increase_by={"GPU": 2}, increase_by_times=2
+            )
+        )
+
+        base_pgf = PlacementGroupFactory([{"CPU": 1, "GPU": 2}])
+        trial1, trial2, trial3, trial4 = self._prepareTrials(scheduler, base_pgf)
+
+        decision = scheduler.on_trial_result(
+            self.tune_controller, trial1, {"metric": 1, "training_iteration": 4}
+        )
+        assert decision == TrialScheduler.CONTINUE
+
+        trial4.status = Trial.TERMINATED
+
+        self._allocateAndAssertNewResources(
+            trial1, scheduler, PlacementGroupFactory([{"CPU": 1, "GPU": 4}])
+        )
+
+        trial3.status = Trial.TERMINATED
+
+        self._allocateAndAssertNewResources(
+            trial2, scheduler, PlacementGroupFactory([{"CPU": 1, "GPU": 4}])
+        )
+
+        trial2.status = Trial.TERMINATED
+
+        self._allocateAndAssertNewResources(
+            trial1, scheduler, PlacementGroupFactory([{"CPU": 1, "GPU": 6}])
+        )
+
+    def testDeallocateResources(self):
+        scheduler = ResourceChangingScheduler(
+            resources_allocation_function=DistributeResources(
+                add_bundles=False, increase_by={"GPU": 2}
+            )
+        )
+
+        base_pgf = PlacementGroupFactory([{"CPU": 1, "GPU": 2}])
+        trial1, trial2, trial3, trial4 = self._prepareTrials(scheduler, base_pgf)
+        trial1.placement_group_factory = PlacementGroupFactory([{"CPU": 1, "GPU": 4}])
+        trial4.status = Trial.PENDING
+
+        self._allocateAndAssertNewResources(
+            trial1, scheduler, PlacementGroupFactory([{"CPU": 1, "GPU": 2}])
+        )
+
+
+class TestUniformResourceAllocationAddBundles(TestUniformResourceAllocation):
+    def testAllocateFreeResources(self):
+        scheduler = ResourceChangingScheduler(
+            resources_allocation_function=DistributeResources(add_bundles=True)
+        )
+
+        base_pgf = PlacementGroupFactory([{"CPU": 1, "GPU": 0}])
+        trial1, trial2, trial3, trial4 = self._prepareTrials(scheduler, base_pgf)
+
+        self._allocateAndAssertNewResources(
+            trial1, scheduler, PlacementGroupFactory([{"CPU": 1}] * 2)
+        )
+        self._allocateAndAssertNewResources(
+            trial2, scheduler, PlacementGroupFactory([{"CPU": 1}] * 2)
+        )
+
+        trial4.status = Trial.TERMINATED
+
+        self._allocateAndAssertNewResources(
+            trial1, scheduler, PlacementGroupFactory([{"CPU": 1}] * 3)
+        )
+
+        trial3.status = Trial.TERMINATED
+
+        self._allocateAndAssertNewResources(
+            trial1, scheduler, PlacementGroupFactory([{"CPU": 1}] * 4)
+        )
+
+        trial2.status = Trial.TERMINATED
+
+        self._allocateAndAssertNewResources(
+            trial1, scheduler, PlacementGroupFactory([{"CPU": 1}] * 8)
+        )
+
+    def testAllocateFreeResourcesWithIncreaseBy(self):
+        scheduler = ResourceChangingScheduler(
+            resources_allocation_function=DistributeResources(
+                add_bundles=True, increase_by={"CPU": 2, "GPU": 2}
+            )
+        )
+
+        base_pgf = PlacementGroupFactory([{}, {"CPU": 2, "GPU": 2}])
+        trial1, trial2, trial3, trial4 = self._prepareTrials(scheduler, base_pgf)
+
+        decision = scheduler.on_trial_result(
+            self.tune_controller, trial1, {"metric": 1, "training_iteration": 4}
+        )
+        assert decision == TrialScheduler.CONTINUE
+
+        trial4.status = Trial.TERMINATED
+
+        self._allocateAndAssertNewResources(
+            trial1, scheduler, PlacementGroupFactory([{}] + [{"CPU": 2, "GPU": 2}] * 2)
+        )
+
+        trial3.status = Trial.TERMINATED
+
+        self._allocateAndAssertNewResources(
+            trial2, scheduler, PlacementGroupFactory([{}] + [{"CPU": 2, "GPU": 2}] * 2)
+        )
+
+        trial2.status = Trial.TERMINATED
+
+        self._allocateAndAssertNewResources(
+            trial1, scheduler, PlacementGroupFactory([{}] + [{"CPU": 2, "GPU": 2}] * 4)
+        )
+
+    def testAllocateFreeResourcesWithIncreaseByTimes(self):
+        scheduler = ResourceChangingScheduler(
+            resources_allocation_function=DistributeResources(
+                add_bundles=True, increase_by={"GPU": 2}, increase_by_times=2
+            )
+        )
+
+        base_pgf = PlacementGroupFactory([{"CPU": 1}, {"GPU": 2}])
+        trial1, trial2, trial3, trial4 = self._prepareTrials(scheduler, base_pgf)
+
+        decision = scheduler.on_trial_result(
+            self.tune_controller, trial1, {"metric": 1, "training_iteration": 4}
+        )
+        assert decision == TrialScheduler.CONTINUE
+
+        trial4.status = Trial.TERMINATED
+
+        self._allocateAndAssertNewResources(
+            trial1, scheduler, PlacementGroupFactory([{"CPU": 1}] + [{"GPU": 2}] * 2)
+        )
+
+        trial3.status = Trial.TERMINATED
+
+        self._allocateAndAssertNewResources(
+            trial2, scheduler, PlacementGroupFactory([{"CPU": 1}] + [{"GPU": 2}] * 2)
+        )
+
+        trial2.status = Trial.TERMINATED
+
+        self._allocateAndAssertNewResources(
+            trial1, scheduler, PlacementGroupFactory([{"CPU": 1}] + [{"GPU": 2}] * 3)
+        )
+
+    def testDeallocateResources(self):
+        scheduler = ResourceChangingScheduler(
+            resources_allocation_function=DistributeResources(
+                add_bundles=True, increase_by={"GPU": 2}
+            )
+        )
+
+        base_pgf = PlacementGroupFactory([{"CPU": 1}, {"GPU": 2}])
+        trial1, trial2, trial3, trial4 = self._prepareTrials(scheduler, base_pgf)
+        trial1.placement_group_factory = PlacementGroupFactory(
+            [{"CPU": 1}] + [{"GPU": 2}] * 2
+        )
+        trial4.status = Trial.PENDING
+
+        self._allocateAndAssertNewResources(
+            trial1, scheduler, PlacementGroupFactory([{"CPU": 1}, {"GPU": 2}])
+        )
+
+
+class TestTopJobResourceAllocation(TestUniformResourceAllocation):
+    def _prepareTrials(self, scheduler, base_pgf):
+        t1, t2, t3, t4 = super()._prepareTrials(scheduler, base_pgf)
+        t1.run_metadata.last_result = {"metric": 1, "training_iteration": 3}
+        t2.run_metadata.last_result = {"metric": 0.9, "training_iteration": 3}
+        t3.run_metadata.last_result = {"metric": 0.8, "training_iteration": 3}
+        t4.run_metadata.last_result = {"metric": 0.7, "training_iteration": 3}
+        return t1, t2, t3, t4
+
+    def testAllocateFreeResources(self):
+        scheduler = ResourceChangingScheduler(
+            resources_allocation_function=DistributeResourcesToTopJob(
+                add_bundles=False, metric="metric", mode="max"
+            )
+        )
+
+        base_pgf = PlacementGroupFactory([{"CPU": 1, "GPU": 0}])
+        trial1, trial2, trial3, trial4 = self._prepareTrials(scheduler, base_pgf)
+
+        decision = scheduler.on_trial_result(
+            self.tune_controller, trial2, {"metric": 0.9, "training_iteration": 4}
+        )
+        assert decision == TrialScheduler.CONTINUE
+
+        self._allocateAndAssertNewResources(
+            trial1, scheduler, PlacementGroupFactory([{"CPU": 5}])
+        )
+        decision = scheduler.on_trial_result(
+            self.tune_controller, trial2, {"metric": 1.1, "training_iteration": 4}
+        )
+        assert decision == TrialScheduler.CONTINUE
+        trial4.status = Trial.TERMINATED
+
+        self._allocateAndAssertNewResources(
+            trial2, scheduler, PlacementGroupFactory([{"CPU": 2}]), metric=1.1
+        )
+        trial3.status = Trial.TERMINATED
+
+        self._allocateAndAssertNewResources(
+            trial1, scheduler, PlacementGroupFactory([{"CPU": 6}]), metric=1.2
+        )
+
+        trial2.status = Trial.TERMINATED
+
+        self._allocateAndAssertNewResources(
+            trial1, scheduler, PlacementGroupFactory([{"CPU": 8}])
+        )
+
+    def testAllocateFreeResourcesWithIncreaseBy(self):
+        scheduler = ResourceChangingScheduler(
+            resources_allocation_function=DistributeResourcesToTopJob(
+                add_bundles=False,
+                increase_by={"CPU": 2, "GPU": 2},
+                metric="metric",
+                mode="max",
+            )
+        )
+
+        base_pgf = PlacementGroupFactory([{"CPU": 2, "GPU": 2}])
+        trial1, trial2, trial3, trial4 = self._prepareTrials(scheduler, base_pgf)
+
+        decision = scheduler.on_trial_result(
+            self.tune_controller, trial2, {"metric": 0.9, "training_iteration": 4}
+        )
+        assert decision == TrialScheduler.CONTINUE
+
+        decision = scheduler.on_trial_result(
+            self.tune_controller, trial1, {"metric": 1.0, "training_iteration": 4}
+        )
+        assert decision == TrialScheduler.CONTINUE
+
+        trial4.status = Trial.TERMINATED
+
+        self._allocateAndAssertNewResources(
+            trial1, scheduler, PlacementGroupFactory([{"CPU": 4, "GPU": 4}])
+        )
+        decision = scheduler.on_trial_result(
+            self.tune_controller, trial2, {"metric": 1.1, "training_iteration": 4}
+        )
+        assert decision == TrialScheduler.CONTINUE
+        trial3.status = Trial.TERMINATED
+
+        self._allocateAndAssertNewResources(
+            trial2, scheduler, PlacementGroupFactory([{"CPU": 4, "GPU": 4}]), metric=1.1
+        )
+        trial2.status = Trial.TERMINATED
+
+        self._allocateAndAssertNewResources(
+            trial1, scheduler, PlacementGroupFactory([{"CPU": 8, "GPU": 8}]), metric=1.2
+        )
+
+    def testAllocateFreeResourcesWithIncreaseByTimes(self):
+        scheduler = ResourceChangingScheduler(
+            resources_allocation_function=DistributeResourcesToTopJob(
+                add_bundles=False,
+                increase_by={"GPU": 2},
+                increase_by_times=2,
+                metric="metric",
+                mode="max",
+            )
+        )
+
+        base_pgf = PlacementGroupFactory([{"CPU": 1, "GPU": 2}])
+        trial1, trial2, trial3, trial4 = self._prepareTrials(scheduler, base_pgf)
+
+        decision = scheduler.on_trial_result(
+            self.tune_controller, trial2, {"metric": 0.9, "training_iteration": 4}
+        )
+        assert decision == TrialScheduler.CONTINUE
+
+        decision = scheduler.on_trial_result(
+            self.tune_controller, trial1, {"metric": 1.0, "training_iteration": 4}
+        )
+        assert decision == TrialScheduler.CONTINUE
+
+        trial4.status = Trial.TERMINATED
+
+        self._allocateAndAssertNewResources(
+            trial1, scheduler, PlacementGroupFactory([{"CPU": 1, "GPU": 4}])
+        )
+        decision = scheduler.on_trial_result(
+            self.tune_controller, trial2, {"metric": 1.1, "training_iteration": 4}
+        )
+        assert decision == TrialScheduler.CONTINUE
+        trial3.status = Trial.TERMINATED
+
+        self._allocateAndAssertNewResources(
+            trial2, scheduler, PlacementGroupFactory([{"CPU": 1, "GPU": 4}]), metric=1.1
+        )
+        trial2.status = Trial.TERMINATED
+
+        self._allocateAndAssertNewResources(
+            trial1, scheduler, PlacementGroupFactory([{"CPU": 1, "GPU": 6}]), metric=1.2
+        )
+
+    def testDeallocateResources(self):
+        scheduler = ResourceChangingScheduler(
+            resources_allocation_function=DistributeResourcesToTopJob(
+                add_bundles=False, increase_by={"GPU": 2}, metric="metric", mode="max"
+            )
+        )
+
+        base_pgf = PlacementGroupFactory([{"CPU": 1, "GPU": 2}])
+        trial1, trial2, trial3, trial4 = self._prepareTrials(scheduler, base_pgf)
+        trial1.placement_group_factory = PlacementGroupFactory([{"CPU": 1, "GPU": 4}])
+        trial4.status = Trial.PENDING
+
+        self._allocateAndAssertNewResources(
+            trial1, scheduler, PlacementGroupFactory([{"CPU": 1, "GPU": 2}])
+        )
+
+
+class TestTopJobResourceAllocationAddBundles(TestTopJobResourceAllocation):
+    def testAllocateFreeResources(self):
+        scheduler = ResourceChangingScheduler(
+            resources_allocation_function=DistributeResourcesToTopJob(
+                add_bundles=True, metric="metric", mode="max"
+            )
+        )
+
+        base_pgf = PlacementGroupFactory([{"CPU": 1, "GPU": 0}])
+        trial1, trial2, trial3, trial4 = self._prepareTrials(scheduler, base_pgf)
+
+        decision = scheduler.on_trial_result(
+            self.tune_controller, trial2, {"metric": 0.9, "training_iteration": 4}
+        )
+        assert decision == TrialScheduler.CONTINUE
+
+        self._allocateAndAssertNewResources(
+            trial1, scheduler, PlacementGroupFactory([{"CPU": 1}] * 5)
+        )
+        decision = scheduler.on_trial_result(
+            self.tune_controller, trial2, {"metric": 1.1, "training_iteration": 4}
+        )
+        assert decision == TrialScheduler.CONTINUE
+        trial4.status = Trial.TERMINATED
+
+        self._allocateAndAssertNewResources(
+            trial2, scheduler, PlacementGroupFactory([{"CPU": 1}] * 2), metric=1.1
+        )
+        trial3.status = Trial.TERMINATED
+
+        self._allocateAndAssertNewResources(
+            trial1, scheduler, PlacementGroupFactory([{"CPU": 1}] * 6), metric=1.2
+        )
+
+        trial2.status = Trial.TERMINATED
+
+        self._allocateAndAssertNewResources(
+            trial1, scheduler, PlacementGroupFactory([{"CPU": 1}] * 8)
+        )
+
+    def testAllocateFreeResourcesWithIncreaseBy(self):
+        scheduler = ResourceChangingScheduler(
+            resources_allocation_function=DistributeResourcesToTopJob(
+                add_bundles=True,
+                increase_by={"CPU": 2, "GPU": 2},
+                metric="metric",
+                mode="max",
+            )
+        )
+
+        base_pgf = PlacementGroupFactory([{}, {"CPU": 2, "GPU": 2}])
+        trial1, trial2, trial3, trial4 = self._prepareTrials(scheduler, base_pgf)
+
+        decision = scheduler.on_trial_result(
+            self.tune_controller, trial2, {"metric": 0.9, "training_iteration": 4}
+        )
+        assert decision == TrialScheduler.CONTINUE
+
+        decision = scheduler.on_trial_result(
+            self.tune_controller, trial1, {"metric": 1.0, "training_iteration": 4}
+        )
+        assert decision == TrialScheduler.CONTINUE
+
+        trial4.status = Trial.TERMINATED
+
+        self._allocateAndAssertNewResources(
+            trial1, scheduler, PlacementGroupFactory([{}] + [{"CPU": 2, "GPU": 2}] * 2)
+        )
+        decision = scheduler.on_trial_result(
+            self.tune_controller, trial2, {"metric": 1.1, "training_iteration": 4}
+        )
+        assert decision == TrialScheduler.CONTINUE
+        trial3.status = Trial.TERMINATED
+
+        self._allocateAndAssertNewResources(
+            trial2,
+            scheduler,
+            PlacementGroupFactory([{}] + [{"CPU": 2, "GPU": 2}] * 2),
+            metric=1.1,
+        )
+        trial2.status = Trial.TERMINATED
+
+        self._allocateAndAssertNewResources(
+            trial1,
+            scheduler,
+            PlacementGroupFactory([{}] + [{"CPU": 2, "GPU": 2}] * 4),
+            metric=1.2,
+        )
+
+    def testAllocateFreeResourcesWithIncreaseByTimes(self):
+        scheduler = ResourceChangingScheduler(
+            resources_allocation_function=DistributeResourcesToTopJob(
+                add_bundles=True,
+                increase_by={"GPU": 2},
+                increase_by_times=2,
+                metric="metric",
+                mode="max",
+            )
+        )
+
+        base_pgf = PlacementGroupFactory([{"CPU": 1}, {"GPU": 2}])
+        trial1, trial2, trial3, trial4 = self._prepareTrials(scheduler, base_pgf)
+
+        decision = scheduler.on_trial_result(
+            self.tune_controller, trial2, {"metric": 0.9, "training_iteration": 4}
+        )
+        assert decision == TrialScheduler.CONTINUE
+
+        decision = scheduler.on_trial_result(
+            self.tune_controller, trial1, {"metric": 1.0, "training_iteration": 4}
+        )
+        assert decision == TrialScheduler.CONTINUE
+
+        trial4.status = Trial.TERMINATED
+
+        self._allocateAndAssertNewResources(
+            trial1, scheduler, PlacementGroupFactory([{"CPU": 1}] + [{"GPU": 2}] * 2)
+        )
+        decision = scheduler.on_trial_result(
+            self.tune_controller, trial2, {"metric": 1.1, "training_iteration": 4}
+        )
+        assert decision == TrialScheduler.CONTINUE
+        trial3.status = Trial.TERMINATED
+
+        self._allocateAndAssertNewResources(
+            trial2,
+            scheduler,
+            PlacementGroupFactory([{"CPU": 1}] + [{"GPU": 2}] * 2),
+            metric=1.1,
+        )
+        trial2.status = Trial.TERMINATED
+
+        self._allocateAndAssertNewResources(
+            trial1,
+            scheduler,
+            PlacementGroupFactory([{"CPU": 1}] + [{"GPU": 2}] * 3),
+            metric=1.2,
+        )
+
+    def testDeallocateResources(self):
+        scheduler = ResourceChangingScheduler(
+            resources_allocation_function=DistributeResourcesToTopJob(
+                add_bundles=True, increase_by={"GPU": 2}, metric="metric", mode="max"
+            )
+        )
+
+        base_pgf = PlacementGroupFactory([{"CPU": 1}, {"GPU": 2}])
+        trial1, trial2, trial3, trial4 = self._prepareTrials(scheduler, base_pgf)
+        trial1.placement_group_factory = PlacementGroupFactory(
+            [{"CPU": 1}] + [{"GPU": 2}] * 2
+        )
+        trial4.status = Trial.PENDING
+
+        self._allocateAndAssertNewResources(
+            trial1, scheduler, PlacementGroupFactory([{"CPU": 1}, {"GPU": 2}])
+        )
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/tune/tests/test_tune_restore.py

@@ -0,0 +1,678 @@
+# coding: utf-8
+import multiprocessing
+import os
+import shutil
+import signal
+import subprocess
+import tempfile
+import threading
+import time
+import unittest
+from collections import Counter
+from typing import List
+from unittest import mock
+
+import numpy as np
+import pytest
+
+import ray
+from ray import tune
+from ray._common.test_utils import run_string_as_driver
+from ray.exceptions import RayTaskError
+from ray.train._internal.session import _TrainingResult
+from ray.tune import Checkpoint, TuneError
+from ray.tune.callback import Callback
+from ray.tune.execution.tune_controller import TuneController
+from ray.tune.experiment import Trial
+from ray.tune.search import Searcher
+from ray.tune.search.basic_variant import BasicVariantGenerator
+from ray.tune.utils import validate_save_restore
+from ray.tune.utils.mock_trainable import MyTrainableClass
+
+
+# Defining the callbacks at the file level, so they can be pickled and spawned
+# in a separate process.
+class SteppingCallback(Callback):
+    def __init__(self, driver_semaphore, trainer_semaphore):
+        self.driver_semaphore = driver_semaphore
+        self.trainer_semaphore = trainer_semaphore
+
+    def on_step_end(self, iteration, trials, **info):
+        self.driver_semaphore.release()  # Driver should continue
+        self.trainer_semaphore.acquire()  # Wait until released
+
+
+def _run(local_dir, driver_semaphore, trainer_semaphore):
+    def _train(config):
+        for i in range(7):
+            ray.tune.report(dict(val=i))
+
+    tune.run(
+        _train,
+        storage_path=local_dir,
+        name="interrupt",
+        callbacks=[SteppingCallback(driver_semaphore, trainer_semaphore)],
+    )
+
+
+class TuneInterruptionTest(unittest.TestCase):
+    def testExperimentInterrupted(self):
+        local_dir = tempfile.mkdtemp()
+        # Unix platforms may default to "fork", which is problematic with
+        # multithreading and GRPC. The child process should always be spawned.
+        mp_ctx = multiprocessing.get_context("spawn")
+        driver_semaphore = mp_ctx.Semaphore()
+        trainer_semaphore = mp_ctx.Semaphore()
+        process = mp_ctx.Process(
+            target=_run,
+            args=(local_dir, driver_semaphore, trainer_semaphore),
+            name="tune_interrupt",
+        )
+        process.daemon = False
+        process.start()
+
+        exp_dir = os.path.join(local_dir, "interrupt")
+
+        # Skip first five steps
+        for i in range(5):
+            driver_semaphore.acquire()  # Wait for callback
+            trainer_semaphore.release()  # Continue training
+
+        driver_semaphore.acquire()
+
+        experiment_state_file = None
+        for file in os.listdir(exp_dir):
+            if file.startswith("experiment_state"):
+                experiment_state_file = os.path.join(exp_dir, file)
+                break
+
+        self.assertTrue(experiment_state_file)
+        last_mtime = os.path.getmtime(experiment_state_file)
+
+        # Now send kill signal
+        os.kill(process.pid, signal.SIGINT)
+        # Release trainer. It should handle the signal and try to
+        # checkpoint the experiment
+        trainer_semaphore.release()
+
+        time.sleep(2)  # Wait for checkpoint
+        new_mtime = os.path.getmtime(experiment_state_file)
+
+        self.assertNotEqual(last_mtime, new_mtime)
+
+        shutil.rmtree(local_dir)
+
+    def testInterruptDisabledInWorkerThread(self):
+        # https://github.com/ray-project/ray/issues/22295
+        # This test will hang without the proper patch because tune.run will fail.
+
+        event = threading.Event()
+
+        def run_in_thread():
+            def _train(config):
+                for i in range(7):
+                    ray.tune.report(dict(val=i))
+
+            tune.run(_train)
+            event.set()
+
+        thread = threading.Thread(target=run_in_thread)
+        thread.start()
+        event.wait()
+        thread.join()
+
+        ray.shutdown()
+        os.environ.pop("TUNE_DISABLE_SIGINT_HANDLER", None)
+
+
+class TuneFailResumeGridTest(unittest.TestCase):
+    class FailureInjectorCallback(Callback):
+        """Adds random failure injection to the TrialExecutor."""
+
+        def __init__(self, num_trials=20, delay_s=0.3):
+            self.num_trials = num_trials
+            self.delay_s = delay_s
+            self.fail_at = None
+
+        def on_step_end(self, trials, **kwargs):
+            if self.fail_at:
+                if time.monotonic() >= self.fail_at:
+                    raise RuntimeError(f"Failing after {self.delay_s}")
+                return
+
+            if len(trials) >= self.num_trials:
+                print(
+                    f"Reached {self.num_trials} trials. "
+                    f"Scheduling failure in {self.delay_s} seconds."
+                )
+                self.fail_at = time.monotonic() + self.delay_s
+
+    class CheckStateCallback(Callback):
+        """Checks state for the experiment initialization."""
+
+        def __init__(self, expected_trials=20):
+            self.expected_trials = expected_trials
+            self._checked = False
+
+        def on_step_begin(self, iteration, trials, **kwargs):
+            if not self._checked:
+                assert len(trials) == self.expected_trials
+                self._checked = True
+
+    class CheckTrialResourcesCallback(Callback):
+        """Checks if pending trials are requesting the right amount of
+        resources.
+
+        The check happens exactly once after `check_after` number of calls
+        to on_step_begin(). Note, we deliberately delay the check to after
+        `check_after` number of steps. This is because when we start a
+        tuning job from fresh (rather than restored), trial list is still
+        empty - any check now would be trivial and thus wasted.
+        """
+
+        def __init__(self, expected_cpu: int, check_after: int = 1):
+            self._expected_cpu = expected_cpu
+            self._checked = False
+            self._check_after = check_after
+
+        def on_step_begin(self, iteration: int, trials: List["Trial"], **info):
+            if not self._checked and iteration >= self._check_after:
+                for trial in trials:
+                    if trial.status == Trial.PENDING:
+                        assert (
+                            trial.placement_group_factory.required_resources.get(
+                                "CPU", 0
+                            )
+                            == self._expected_cpu
+                        )
+                self._checked = True
+
+    def setUp(self):
+        self.logdir = tempfile.mkdtemp()
+
+        # These tests need driver syncing to happen before the crash happens
+        # so that they can pick up from the *exact* state it left off at.
+        # We do this by failing after a delay of 0.3s > TUNE_GLOBAL_CHECKPOINT_S
+        os.environ["TUNE_GLOBAL_CHECKPOINT_S"] = "0.1"
+
+        ray.init(num_cpus=2)
+
+        from ray.tune import register_trainable
+
+        register_trainable("trainable", MyTrainableClass)
+
+    def tearDown(self):
+        os.environ.pop("TUNE_GLOBAL_CHECKPOINT_S")
+        os.environ.pop("TUNE_MAX_PENDING_TRIALS_PG", None)
+        shutil.rmtree(self.logdir)
+        ray.shutdown()
+
+    def testFailResumeGridSearch(self):
+        os.environ["TUNE_MAX_PENDING_TRIALS_PG"] = "1"
+
+        config = dict(
+            num_samples=3,
+            fail_fast=True,
+            config={
+                "test": tune.grid_search([1, 2, 3]),
+                "test2": tune.grid_search([1, 2, 3]),
+            },
+            stop={"training_iteration": 2},
+            name="testFailResumeGridSearch",
+            verbose=1,
+        )
+
+        with self.assertRaises(RuntimeError):
+            tune.run("trainable", callbacks=[self.FailureInjectorCallback()], **config)
+
+        analysis = tune.run(
+            "trainable", resume=True, callbacks=[self.CheckStateCallback()], **config
+        )
+        assert len(analysis.trials) == 27
+        test_counter = Counter([t.config["test"] for t in analysis.trials])
+        assert all(v == 9 for v in test_counter.values())
+        test2_counter = Counter([t.config["test2"] for t in analysis.trials])
+        assert all(v == 9 for v in test2_counter.values())
+
+    # Unfinished trials' resources should be updated.
+    def testResourceUpdateInResume(self):
+        os.environ["TUNE_MAX_PENDING_TRIALS_PG"] = "1"
+
+        config = dict(
+            num_samples=3,
+            fail_fast=True,
+            config={
+                "test": tune.grid_search([1, 2, 3]),
+                "test2": tune.grid_search([1, 2, 3]),
+            },
+            stop={"training_iteration": 2},
+            name="testResourceUpdateInResume",
+            verbose=1,
+        )
+
+        with self.assertRaises(RuntimeError):
+            tune.run(
+                "trainable",
+                callbacks=[
+                    self.FailureInjectorCallback(),
+                    self.CheckTrialResourcesCallback(1),
+                ],
+                **config,
+            )
+
+        analysis = tune.run(
+            "trainable",
+            resume=True,
+            resources_per_trial={"cpu": 2},
+            callbacks=[self.CheckTrialResourcesCallback(2)],
+            **config,
+        )
+        assert len(analysis.trials) == 27
+
+    @mock.patch.dict(os.environ, {"TUNE_MAX_PENDING_TRIALS_PG": "1"})
+    def testConfigUpdateInResume(self):
+        class FakeDataset:
+            def __init__(self, name):
+                self.name = name
+
+        config = dict(
+            num_samples=1,
+            fail_fast=True,
+            config={
+                "test": tune.grid_search(
+                    [FakeDataset("1"), FakeDataset("2"), FakeDataset("3")]
+                ),
+                "test2": tune.grid_search(
+                    [
+                        FakeDataset("4"),
+                        FakeDataset("5"),
+                        FakeDataset("6"),
+                        FakeDataset("7"),
+                    ]
+                ),
+            },
+            stop={"training_iteration": 2},
+            name="testConfigUpdateInResume",
+            verbose=1,
+        )
+
+        with self.assertRaises(RuntimeError):
+            tune.run(
+                "trainable",
+                callbacks=[
+                    self.FailureInjectorCallback(num_trials=1),
+                    self.CheckTrialResourcesCallback(1),
+                ],
+                **config,
+            )
+
+        config["config"] = {
+            "test": tune.grid_search(
+                [FakeDataset("8"), FakeDataset("9"), FakeDataset("10")]
+            ),
+            "test2": tune.grid_search(
+                [
+                    FakeDataset("11"),
+                    FakeDataset("12"),
+                    FakeDataset("13"),
+                    FakeDataset("14"),
+                ]
+            ),
+        }
+
+        analysis = tune.run(
+            "trainable",
+            resume=True,
+            **config,
+        )
+        assert len(analysis.trials) == 12
+        for t in analysis.trials:
+            # Make sure that test and test2 are updated.
+            assert t.config["test"].name in ["8", "9", "10"]
+            assert t.config["test2"].name in ["11", "12", "13", "14"]
+
+    def testFailResumeWithPreset(self):
+        os.environ["TUNE_MAX_PENDING_TRIALS_PG"] = "1"
+
+        search_alg = BasicVariantGenerator(
+            points_to_evaluate=[{"test": -1, "test2": -1}, {"test": -1}, {"test2": -1}]
+        )
+
+        config = dict(
+            num_samples=3 + 3,  # 3 preset, 3 samples
+            fail_fast=True,
+            config={
+                "test": tune.grid_search([1, 2, 3]),
+                "test2": tune.grid_search([1, 2, 3]),
+            },
+            stop={"training_iteration": 2},
+            name="testFailResumeWithPreset",
+            verbose=1,
+        )
+        with self.assertRaises(RuntimeError):
+            tune.run(
+                "trainable",
+                callbacks=[self.FailureInjectorCallback(5)],
+                search_alg=search_alg,
+                **config,
+            )
+
+        print("---- RESTARTING RUN ----")
+
+        analysis = tune.run(
+            "trainable",
+            resume=True,
+            callbacks=[self.CheckStateCallback(expected_trials=5)],
+            search_alg=search_alg,
+            **config,
+        )
+        assert len(analysis.trials) == 34
+        test_counter = Counter([t.config["test"] for t in analysis.trials])
+        assert test_counter.pop(-1) == 4
+        assert all(v == 10 for v in test_counter.values())
+        test2_counter = Counter([t.config["test2"] for t in analysis.trials])
+        assert test2_counter.pop(-1) == 4
+        assert all(v == 10 for v in test2_counter.values())
+
+    def testFailResumeAfterPreset(self):
+        os.environ["TUNE_MAX_PENDING_TRIALS_PG"] = "1"
+
+        search_alg = BasicVariantGenerator(
+            points_to_evaluate=[{"test": -1, "test2": -1}, {"test": -1}, {"test2": -1}]
+        )
+
+        config = dict(
+            num_samples=3 + 3,  # 3 preset, 3 samples
+            fail_fast=True,
+            config={
+                "test": tune.grid_search([1, 2, 3]),
+                "test2": tune.grid_search([1, 2, 3]),
+            },
+            stop={"training_iteration": 2},
+            name="testFailResumeAfterPreset",
+            verbose=1,
+        )
+
+        with self.assertRaises(RuntimeError):
+            tune.run(
+                "trainable",
+                callbacks=[self.FailureInjectorCallback(15)],
+                search_alg=search_alg,
+                **config,
+            )
+
+        print("---- RESTARTING RUN ----")
+
+        analysis = tune.run(
+            "trainable",
+            resume=True,
+            callbacks=[self.CheckStateCallback(expected_trials=15)],
+            search_alg=search_alg,
+            **config,
+        )
+        assert len(analysis.trials) == 34
+        test_counter = Counter([t.config["test"] for t in analysis.trials])
+        assert test_counter.pop(-1) == 4
+        assert all(v == 10 for v in test_counter.values())
+        test2_counter = Counter([t.config["test2"] for t in analysis.trials])
+        assert test2_counter.pop(-1) == 4
+        assert all(v == 10 for v in test2_counter.values())
+
+    def testMultiExperimentFail(self):
+        os.environ["TUNE_MAX_PENDING_TRIALS_PG"] = "1"
+
+        experiments = []
+        for i in range(3):
+            experiments.append(
+                tune.Experiment(
+                    run=MyTrainableClass,
+                    name="testMultiExperimentFail",
+                    num_samples=2,
+                    config={
+                        "test": tune.grid_search([1, 2, 3]),
+                    },
+                    stop={"training_iteration": 1},
+                )
+            )
+
+        with self.assertRaises(RuntimeError):
+            tune.run(
+                experiments,
+                callbacks=[self.FailureInjectorCallback(10)],
+                fail_fast=True,
+            )
+
+        analysis = tune.run(
+            experiments,
+            resume=True,
+            callbacks=[self.CheckStateCallback(expected_trials=10)],
+            fail_fast=True,
+        )
+        assert len(analysis.trials) == 18
+
+    def testWarningLargeGrid(self):
+        config = dict(
+            num_samples=3,
+            fail_fast=True,
+            config={
+                "test": tune.grid_search(list(range(20))),
+                "test2": tune.grid_search(list(range(20))),
+                "test3": tune.grid_search(list(range(20))),
+                "test4": tune.grid_search(list(range(20))),
+                "test5": tune.grid_search(list(range(20))),
+            },
+            stop={"training_iteration": 2},
+            name="testWarningLargeGrid",
+            verbose=1,
+        )
+        with self.assertWarnsRegex(UserWarning, "exceeds the serialization threshold"):
+            with self.assertRaises(RuntimeError):
+                tune.run(
+                    "trainable", callbacks=[self.FailureInjectorCallback(10)], **config
+                )
+
+
+class TuneExampleTest(unittest.TestCase):
+    def setUp(self):
+        ray.init(num_cpus=2)
+
+    def tearDown(self):
+        ray.shutdown()
+
+    def testPBTKeras(self):
+        from tensorflow.keras.datasets import cifar10
+
+        from ray.tune.examples.pbt_tune_cifar10_with_keras import Cifar10Model
+
+        cifar10.load_data()
+        validate_save_restore(Cifar10Model)
+
+    def testPyTorchMNIST(self):
+        from torchvision import datasets
+
+        from ray.tune.examples.mnist_pytorch_trainable import TrainMNIST
+
+        datasets.MNIST("~/data", train=True, download=True)
+        validate_save_restore(TrainMNIST)
+
+    def testHyperbandExample(self):
+        validate_save_restore(MyTrainableClass)
+
+    def testAsyncHyperbandExample(self):
+        validate_save_restore(MyTrainableClass)
+
+
+class AutoInitTest(unittest.TestCase):
+    def testTuneRestore(self):
+        self.assertFalse(ray.is_initialized())
+        tune.run(MyTrainableClass, name="TestAutoInit", stop={"training_iteration": 1})
+        self.assertTrue(ray.is_initialized())
+
+    def tearDown(self):
+        ray.shutdown()
+
+
+class SearcherTest(unittest.TestCase):
+    class MockSearcher(Searcher):
+        def __init__(self, data):
+            self.data = data
+
+        def save(self, path):
+            with open(path, "w") as f:
+                f.write(self.data)
+
+        def restore(self, path):
+            with open(path, "r") as f:
+                self.data = f.read()
+
+    def testSaveRestoreDir(self):
+        tmpdir = tempfile.mkdtemp()
+        original_data = "hello-its-me"
+        searcher = self.MockSearcher(original_data)
+        searcher.save_to_dir(tmpdir)
+        searcher_2 = self.MockSearcher("no-its-not-me")
+        searcher_2.restore_from_dir(tmpdir)
+        assert searcher_2.data == original_data
+
+
+class WorkingDirectoryTest(unittest.TestCase):
+    def testWorkingDir(self):
+        """Trainables should know the original working dir through env variable."""
+
+        os.environ.pop("TUNE_ORIG_WORKING_DIR", None)
+        working_dir = os.getcwd()
+
+        def f(config):
+            assert os.environ.get("TUNE_ORIG_WORKING_DIR") == working_dir
+
+        ray.init(num_cpus=1)
+        tune.run(f)
+        ray.shutdown()
+
+
+class TrainableCrashWithFailFast(unittest.TestCase):
+    def test(self):
+        """Trainable crashes with fail_fast flag and the original crash message
+        should bubble up."""
+
+        def f(config):
+            ray.tune.report({"a": 1})
+            time.sleep(0.1)
+            raise RuntimeError("Error happens in trainable!!")
+
+        with self.assertRaisesRegex(RayTaskError, "Error happens in trainable!!"):
+            tune.run(f, fail_fast=TuneController.RAISE)
+
+
+@pytest.mark.parametrize(
+    "trial_config", [{}, {"attr": 4}, {"nested": {"key": "value"}}]
+)
+def test_trial_last_result_restore(trial_config):
+    metrics = {"metric1": 4, "nested2": {"metric3": 6}}
+    metrics["config"] = trial_config
+
+    trial = Trial(trainable_name="stub", config=trial_config, stub=True)
+    trial.update_last_result(metrics)
+
+    result = _TrainingResult(
+        checkpoint=Checkpoint(path="file:///tmp/no_data"), metrics=metrics
+    )
+
+    trial.temporary_state.restoring_from = result
+    trial.on_restore()
+    assert trial.run_metadata.last_result == metrics
+
+
+def test_stacktrace():
+    """Test proper stacktrace is printed for RayTaskError."""
+    CMD = """
+from ray import tune
+
+def train_fn(config):
+    raise Exception("Inducing exception for testing purposes.")
+
+tune.run(train_fn, num_samples=1)
+    """
+    with pytest.raises(subprocess.CalledProcessError) as exc_info:
+        run_string_as_driver(CMD)
+    assert "Inducing exception for testing purposes." in exc_info.value.output.decode()
+
+
+@pytest.mark.parametrize(
+    "resume",
+    [
+        True,
+        "AUTO",
+        "AUTO+ERRORED",
+        "AUTO+ERRORED_ONLY",
+        "AUTO+RESTART_ERRORED",
+        "AUTO+RESTART_ERRORED_ONLY",
+    ],
+)
+def test_resume_options(tmp_path, resume):
+    tmp_path.joinpath("dummy_ckpt").mkdir()
+
+    def train_fn(config):
+        checkpoint = ray.tune.get_checkpoint()
+        if not checkpoint:
+            ray.tune.report(
+                {"finish_marker": False},
+                checkpoint=Checkpoint.from_directory(tmp_path / "dummy_ckpt"),
+            )
+            raise RuntimeError("failing on the first run!!")
+        ray.tune.report({"finish_marker": True})
+
+    analysis = tune.run(
+        train_fn,
+        storage_path=str(tmp_path),
+        name="test_resume_options",
+        raise_on_failed_trial=False,
+    )
+    results = ray.tune.ResultGrid(analysis)
+    assert not results[0].metrics.get("finish_marker", False)
+    analysis = tune.run(
+        train_fn,
+        storage_path=str(tmp_path),
+        name="test_resume_options",
+        resume=resume,
+        raise_on_failed_trial=False,
+    )
+    results = ray.tune.ResultGrid(analysis)
+    if resume in [True, "AUTO", "AUTO+RESTART_ERRORED", "AUTO+RESTART_ERRORED_ONLY"]:
+        # These options either don't resume the errored trial,
+        # or restart it without a checkpoint --> leading to the RuntimeError again
+        assert not results[0].metrics.get("finish_marker")
+    else:
+        assert results[0].metrics.get("finish_marker")
+
+
+# For some reason, different tests are coupled through tune.registry.
+# After running `ResourceExhaustedTest`, there is always a super huge `training_func` to
+# be put through GCS, which will fail subsequent tests.
+# tldr, make sure that this test is the last test in the file.
+class ResourceExhaustedTest(unittest.TestCase):
+    def test_resource_exhausted_info(self):
+        """This is to test if helpful information is displayed when
+        the objects captured in trainable/training function are too
+        large and RESOURCES_EXHAUSTED error of gRPC is triggered."""
+
+        a_large_array = []
+        for _ in range(50):
+            a_large_array.append(np.random.rand(400, 4096))
+
+        def training_func(config):
+            del config  # unused var
+            for item in a_large_array:
+                assert item
+
+        with self.assertRaisesRegex(
+            TuneError,
+            "The Trainable/training function is too large for grpc resource limit.",
+        ):
+            tune.run(training_func)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__] + sys.argv[1:]))

python/ray/tune/tests/test_tune_restore_warm_start.py

@@ -0,0 +1,373 @@
+# coding: utf-8
+import os
+import shutil
+import sys
+import tempfile
+import unittest
+
+import numpy as np
+import pandas
+import pytest
+from hyperopt import hp
+from nevergrad.optimization import optimizerlib
+from packaging.version import Version
+from zoopt import ValueType
+
+import ray
+from ray import tune
+from ray.rllib import _register_all
+from ray.tune.schedulers.hb_bohb import HyperBandForBOHB
+from ray.tune.search import ConcurrencyLimiter
+from ray.tune.search.ax import AxSearch
+from ray.tune.search.bayesopt import BayesOptSearch
+from ray.tune.search.bohb import TuneBOHB
+from ray.tune.search.hebo import HEBOSearch
+from ray.tune.search.hyperopt import HyperOptSearch
+from ray.tune.search.nevergrad import NevergradSearch
+from ray.tune.search.optuna import OptunaSearch
+from ray.tune.search.zoopt import ZOOptSearch
+
+
+class AbstractWarmStartTest:
+    def setUp(self):
+        ray.init(num_cpus=1)
+        self.tmpdir = tempfile.mkdtemp()
+        self.experiment_name = "results"
+
+    def tearDown(self):
+        shutil.rmtree(self.tmpdir)
+        ray.shutdown()
+        _register_all()
+
+    def set_basic_conf(self):
+        raise NotImplementedError()
+
+    def get_scheduler(self):
+        return None
+
+    def treat_trial_config(self, trial_config):
+        return trial_config
+
+    def run_part_from_scratch(self):
+        np.random.seed(162)
+        search_alg, cost = self.set_basic_conf()
+        if not isinstance(search_alg, ConcurrencyLimiter):
+            search_alg = ConcurrencyLimiter(search_alg, 1)
+        results_exp_1 = tune.run(
+            cost,
+            num_samples=5,
+            search_alg=search_alg,
+            scheduler=self.get_scheduler(),
+            verbose=0,
+            name=self.experiment_name,
+            storage_path=self.tmpdir,
+            reuse_actors=True,
+        )
+        checkpoint_path = os.path.join(self.tmpdir, "warmStartTest.pkl")
+        search_alg.save(checkpoint_path)
+        return results_exp_1, np.random.get_state(), checkpoint_path
+
+    def run_from_experiment_restore(self, random_state):
+        search_alg, cost = self.set_basic_conf()
+        if not isinstance(search_alg, ConcurrencyLimiter):
+            search_alg = ConcurrencyLimiter(search_alg, 1)
+        search_alg.restore_from_dir(os.path.join(self.tmpdir, self.experiment_name))
+        results = tune.run(
+            cost,
+            num_samples=5,
+            search_alg=search_alg,
+            scheduler=self.get_scheduler(),
+            verbose=0,
+            name=self.experiment_name,
+            storage_path=self.tmpdir,
+            reuse_actors=True,
+        )
+        return results
+
+    def run_explicit_restore(self, random_state, checkpoint_path):
+        np.random.set_state(random_state)
+        search_alg2, cost = self.set_basic_conf()
+        if not isinstance(search_alg2, ConcurrencyLimiter):
+            search_alg2 = ConcurrencyLimiter(search_alg2, 1)
+        search_alg2.restore(checkpoint_path)
+        return tune.run(
+            cost,
+            num_samples=5,
+            search_alg=search_alg2,
+            scheduler=self.get_scheduler(),
+            verbose=0,
+            reuse_actors=True,
+        )
+
+    def run_full(self):
+        np.random.seed(162)
+        search_alg3, cost = self.set_basic_conf()
+        if not isinstance(search_alg3, ConcurrencyLimiter):
+            search_alg3 = ConcurrencyLimiter(search_alg3, 1)
+        return tune.run(
+            cost,
+            num_samples=10,
+            search_alg=search_alg3,
+            scheduler=self.get_scheduler(),
+            verbose=0,
+            reuse_actors=True,
+        )
+
+    def testWarmStart(self):
+        results_exp_1, r_state, checkpoint_path = self.run_part_from_scratch()
+        results_exp_2 = self.run_explicit_restore(r_state, checkpoint_path)
+        results_exp_3 = self.run_full()
+        trials_1_config = self.treat_trial_config(
+            [trial.config for trial in results_exp_1.trials]
+        )
+        trials_2_config = self.treat_trial_config(
+            [trial.config for trial in results_exp_2.trials]
+        )
+        trials_3_config = self.treat_trial_config(
+            [trial.config for trial in results_exp_3.trials]
+        )
+        self.assertEqual(trials_1_config + trials_2_config, trials_3_config)
+
+    def testRestore(self):
+        results_exp_1, r_state, checkpoint_path = self.run_part_from_scratch()
+        results_exp_2 = self.run_from_experiment_restore(r_state)
+        results_exp_3 = self.run_full()
+
+        trials_1_config = self.treat_trial_config(
+            [trial.config for trial in results_exp_1.trials]
+        )
+        trials_2_config = self.treat_trial_config(
+            [trial.config for trial in results_exp_2.trials]
+        )
+        trials_3_config = self.treat_trial_config(
+            [trial.config for trial in results_exp_3.trials]
+        )
+        self.assertEqual(trials_1_config + trials_2_config, trials_3_config)
+
+
+class HyperoptWarmStartTest(AbstractWarmStartTest, unittest.TestCase):
+    def set_basic_conf(self):
+        space = {
+            "x": hp.uniform("x", 0, 10),
+            "y": hp.uniform("y", -10, 10),
+            "z": hp.uniform("z", -10, 0),
+        }
+
+        def cost(space):
+            loss = space["x"] ** 2 + space["y"] ** 2 + space["z"] ** 2
+            tune.report(dict(loss=loss))
+
+        search_alg = HyperOptSearch(
+            space,
+            metric="loss",
+            mode="min",
+            random_state_seed=5,
+            n_initial_points=1,
+        )
+        return search_alg, cost
+
+
+class BayesoptWarmStartTest(AbstractWarmStartTest, unittest.TestCase):
+    def set_basic_conf(self, analysis=None):
+        space = {"width": (0, 20), "height": (-100, 100)}
+
+        def cost(space):
+            tune.report(
+                dict(loss=(space["height"] - 14) ** 2 - abs(space["width"] - 3))
+            )
+
+        search_alg = BayesOptSearch(space, metric="loss", mode="min", analysis=analysis)
+        return search_alg, cost
+
+    def testBootStrapAnalysis(self):
+        analysis = self.run_full()
+        search_alg3, cost = self.set_basic_conf(analysis)
+        if not isinstance(search_alg3, ConcurrencyLimiter):
+            search_alg3 = ConcurrencyLimiter(search_alg3, 1)
+        tune.run(
+            cost, num_samples=10, search_alg=search_alg3, verbose=0, reuse_actors=True
+        )
+
+
+class NevergradWarmStartTest(AbstractWarmStartTest, unittest.TestCase):
+    def set_basic_conf(self):
+        instrumentation = 2
+        parameter_names = ["height", "width"]
+        optimizer = optimizerlib.OnePlusOne(instrumentation)
+
+        def cost(space):
+            tune.report(
+                dict(loss=(space["height"] - 14) ** 2 - abs(space["width"] - 3))
+            )
+
+        search_alg = NevergradSearch(
+            optimizer,
+            space=parameter_names,
+            metric="loss",
+            mode="min",
+        )
+        return search_alg, cost
+
+
+class OptunaWarmStartTest(AbstractWarmStartTest, unittest.TestCase):
+    def set_basic_conf(self):
+        from optuna.samplers import TPESampler
+
+        space = OptunaSearch.convert_search_space(
+            {"width": tune.uniform(0, 20), "height": tune.uniform(-100, 100)}
+        )
+
+        def cost(space):
+            tune.report(
+                dict(loss=(space["height"] - 14) ** 2 - abs(space["width"] - 3))
+            )
+
+        search_alg = OptunaSearch(
+            space, sampler=TPESampler(seed=10), metric="loss", mode="min"
+        )
+        return search_alg, cost
+
+
+class ZOOptWarmStartTest(AbstractWarmStartTest, unittest.TestCase):
+    def set_basic_conf(self):
+        dim_dict = {
+            "height": (ValueType.CONTINUOUS, [-100, 100], 1e-2),
+            "width": (ValueType.DISCRETE, [0, 20], False),
+        }
+
+        def cost(param):
+            tune.report(
+                dict(loss=(param["height"] - 14) ** 2 - abs(param["width"] - 3))
+            )
+
+        search_alg = ZOOptSearch(
+            algo="Asracos",  # only support ASRacos currently
+            budget=200,
+            dim_dict=dim_dict,
+            metric="loss",
+            mode="min",
+        )
+
+        return search_alg, cost
+
+
+@pytest.mark.skipif(sys.version_info >= (3, 12), reason="HEBO doesn't support py312")
+class HEBOWarmStartTest(AbstractWarmStartTest, unittest.TestCase):
+    def set_basic_conf(self):
+        if Version(pandas.__version__) >= Version("2.0.0"):
+            pytest.skip("HEBO does not support pandas>=2.0.0")
+
+        from hebo.design_space.design_space import DesignSpace as HEBODesignSpace
+
+        space_config = [
+            {"name": "width", "type": "num", "lb": 0, "ub": 20},
+            {"name": "height", "type": "num", "lb": -100, "ub": 100},
+        ]
+        space = HEBODesignSpace().parse(space_config)
+
+        def cost(param):
+            tune.report(
+                dict(loss=(param["height"] - 14) ** 2 - abs(param["width"] - 3))
+            )
+
+        search_alg = HEBOSearch(
+            space=space, metric="loss", mode="min", random_state_seed=5
+        )
+        # This is done on purpose to speed up the test, as HEBO will
+        # cache suggestions
+        search_alg = ConcurrencyLimiter(search_alg, max_concurrent=10)
+        return search_alg, cost
+
+
+class AxWarmStartTest(AbstractWarmStartTest, unittest.TestCase):
+    def set_basic_conf(self):
+        from ax.service.ax_client import AxClient, ObjectiveProperties
+
+        space = AxSearch.convert_search_space(
+            {"width": tune.uniform(0, 20), "height": tune.uniform(-100, 100)}
+        )
+
+        try:
+            # ax 1.0+: ax.modelbridge was removed
+            from ax.adapter.registry import Generators as Models
+            from ax.generation_strategy.generation_node import GenerationStep
+            from ax.generation_strategy.generation_strategy import GenerationStrategy
+        except ImportError:
+            # ax 0.x
+            from ax.modelbridge.generation_strategy import (
+                GenerationStep,
+                GenerationStrategy,
+            )
+            from ax.modelbridge.registry import Models
+
+        # set generation strategy to sobol to ensure reproductibility
+        # ax 1.0+ renamed 'model' to 'generator'; ax <0.2.0 used 'num_arms'
+        try:
+            gs = GenerationStrategy(
+                steps=[
+                    GenerationStep(
+                        generator=Models.SOBOL,
+                        num_trials=-1,
+                        model_kwargs={"seed": 42},
+                    ),
+                ]
+            )
+        except TypeError:
+            try:
+                gs = GenerationStrategy(
+                    steps=[
+                        GenerationStep(
+                            model=Models.SOBOL,
+                            num_trials=-1,
+                            model_kwargs={"seed": 42},
+                        ),
+                    ]
+                )
+            except TypeError:
+                # ax-platform<0.2.0
+                gs = GenerationStrategy(
+                    steps=[
+                        GenerationStep(
+                            model=Models.SOBOL,
+                            num_arms=-1,
+                            model_kwargs={"seed": 42},
+                        ),
+                    ]
+                )
+
+        client = AxClient(random_seed=42, generation_strategy=gs)
+        client.create_experiment(
+            parameters=space,
+            objectives={"loss": ObjectiveProperties(minimize=True)},
+        )
+
+        def cost(space):
+            tune.report(
+                dict(loss=(space["height"] - 14) ** 2 - abs(space["width"] - 3))
+            )
+
+        search_alg = AxSearch(ax_client=client)
+        return search_alg, cost
+
+
+@pytest.mark.skipif(sys.version_info >= (3, 12), reason="BOHB doesn't support py312")
+class BOHBWarmStartTest(AbstractWarmStartTest, unittest.TestCase):
+    def set_basic_conf(self):
+        space = {"width": tune.uniform(0, 20), "height": tune.uniform(-100, 100)}
+
+        def cost(space):
+            for i in range(10):
+                tune.report(
+                    dict(loss=(space["height"] - 14) ** 2 - abs(space["width"] - 3 - i))
+                )
+
+        search_alg = TuneBOHB(space=space, metric="loss", mode="min", seed=1)
+
+        return search_alg, cost
+
+    def get_scheduler(self):
+        return HyperBandForBOHB(max_t=100, metric="loss", mode="min")
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__] + sys.argv[1:]))

python/ray/tune/tests/test_tune_save_restore.py

@@ -0,0 +1,144 @@
+# coding: utf-8
+import os
+import pickle
+import shutil
+import tempfile
+import unittest
+
+import ray
+from ray import tune
+from ray.tune import CheckpointConfig, Trainable
+from ray.tune.utils import validate_save_restore
+
+
+class SerialTuneRelativeLocalDirTest(unittest.TestCase):
+    prefix = "Serial"
+
+    class MockTrainable(Trainable):
+        _name = "MockTrainable"
+
+        def setup(self, config):
+            self.state = {"hi": 1}
+
+        def step(self):
+            return {"timesteps_this_iter": 1, "done": True}
+
+        def save_checkpoint(self, checkpoint_dir):
+            checkpoint_path = os.path.join(checkpoint_dir, "checkpoint.pkl")
+            with open(checkpoint_path, "wb") as f:
+                pickle.dump(self.state, f)
+
+        def load_checkpoint(self, checkpoint_dir):
+            checkpoint_path = os.path.join(checkpoint_dir, "checkpoint.pkl")
+            with open(checkpoint_path, "rb") as f:
+                extra_data = pickle.load(f)
+            self.state.update(extra_data)
+
+    def setUp(self):
+        self.absolute_local_dir = None
+        ray.init(num_cpus=2, num_gpus=0)
+
+    def tearDown(self):
+        if self.absolute_local_dir is not None:
+            shutil.rmtree(self.absolute_local_dir, ignore_errors=True)
+            self.absolute_local_dir = None
+        ray.shutdown()
+
+    def _get_trial_dir(self, absoulte_exp_dir):
+        print("looking for", self.MockTrainable._name)
+        print("in", os.listdir(absoulte_exp_dir))
+        trial_dirname = next(
+            (
+                child_dir
+                for child_dir in os.listdir(absoulte_exp_dir)
+                if (
+                    os.path.isdir(os.path.join(absoulte_exp_dir, child_dir))
+                    and child_dir.startswith(self.MockTrainable._name)
+                )
+            )
+        )
+
+        trial_absolute_dir = os.path.join(absoulte_exp_dir, trial_dirname)
+
+        return trial_dirname, trial_absolute_dir
+
+    def _train(self, exp_name, local_dir, absolute_local_dir):
+        (trial,) = tune.run(
+            self.MockTrainable,
+            name=exp_name,
+            stop={"training_iteration": 1},
+            checkpoint_config=CheckpointConfig(checkpoint_frequency=1),
+            storage_path=local_dir,
+            config={"env": "CartPole-v0", "log_level": "DEBUG"},
+        ).trials
+
+        exp_dir = os.path.join(absolute_local_dir, exp_name)
+        _, abs_trial_dir = self._get_trial_dir(exp_dir)
+
+        self.assertIsNone(trial.error_file)
+        self.assertEqual(trial.path, abs_trial_dir)
+
+        self.assertTrue(os.path.isdir(absolute_local_dir), absolute_local_dir)
+        self.assertTrue(os.path.isdir(exp_dir))
+        self.assertTrue(os.path.isdir(abs_trial_dir))
+        self.assertTrue(
+            os.path.isfile(
+                os.path.join(abs_trial_dir, "checkpoint_000000/checkpoint.pkl")
+            )
+        )
+
+    def _restore(self, exp_name, local_dir, absolute_local_dir):
+        trial_name, abs_trial_dir = self._get_trial_dir(
+            os.path.join(absolute_local_dir, exp_name)
+        )
+
+        checkpoint_path = os.path.join(
+            local_dir, exp_name, trial_name, "checkpoint_000000"
+        )
+        assert os.path.exists(os.path.expanduser(checkpoint_path))
+
+        (trial,) = tune.run(
+            self.MockTrainable,
+            name=exp_name,
+            stop={"training_iteration": 2},  # train one more iteration.
+            restore=checkpoint_path,  # Restore the checkpoint
+            config={"env": "CartPole-v0", "log_level": "DEBUG"},
+        ).trials
+        self.assertIsNone(trial.error_file)
+
+    def testTempfile(self):
+        local_dir = tempfile.mkdtemp()
+        exp_name = self.prefix + "Tempfile"
+        self.absolute_local_dir = local_dir
+        self._train(exp_name, local_dir, local_dir)
+        self._restore(exp_name, local_dir, local_dir)
+
+    def testCheckpointWithNoop(self):
+        """Tests that passing the checkpoint_dir right back works."""
+
+        class MockTrainable(Trainable):
+            def setup(self, config):
+                pass
+
+            def step(self):
+                return {"score": 1}
+
+            def save_checkpoint(self, checkpoint_dir):
+                with open(os.path.join(checkpoint_dir, "test.txt"), "wb") as f:
+                    pickle.dump("test", f)
+
+            def load_checkpoint(self, checkpoint_dir):
+                with open(os.path.join(checkpoint_dir, "test.txt"), "rb") as f:
+                    x = pickle.load(f)
+
+                assert x == "test"
+
+        validate_save_restore(MockTrainable)
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/tune/tests/test_tuner.py

@@ -0,0 +1,543 @@
+import os
+import unittest
+from pathlib import Path
+from typing import Optional
+from unittest.mock import patch
+
+import pytest
+from sklearn.datasets import load_breast_cancer
+from sklearn.utils import shuffle
+
+import ray
+from ray import tune
+from ray.data import Dataset, Datasource, ReadTask, from_pandas, read_datasource
+from ray.data.block import BlockMetadata
+from ray.train.data_parallel_trainer import DataParallelTrainer
+from ray.train.examples.pytorch.torch_linear_example import (
+    train_func as linear_train_func,
+)
+from ray.train.torch import TorchTrainer
+from ray.train.trainer import BaseTrainer
+from ray.train.xgboost import XGBoostTrainer
+from ray.tune import Callback, CheckpointConfig, CLIReporter, RunConfig
+from ray.tune.tune_config import TuneConfig
+from ray.tune.tuner import Tuner
+
+
+@pytest.fixture
+def shutdown_only():
+    yield None
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+@pytest.fixture
+def chdir_tmpdir(tmpdir):
+    old_cwd = os.getcwd()
+    os.chdir(tmpdir)
+    yield tmpdir
+    os.chdir(old_cwd)
+
+
+class DummyTrainer(BaseTrainer):
+    _scaling_config_allowed_keys = BaseTrainer._scaling_config_allowed_keys + [
+        "num_workers",
+        "use_gpu",
+        "resources_per_worker",
+        "placement_strategy",
+    ]
+
+    def training_loop(self) -> None:
+        for i in range(5):
+            tune.report({"step": i})
+
+
+class FailingTrainer(DummyTrainer):
+    def training_loop(self) -> None:
+        raise RuntimeError("There is an error in trainer!")
+
+
+class TestDatasource(Datasource):
+    def __init__(self, do_shuffle: bool):
+        self._shuffle = do_shuffle
+
+    def prepare_read(self, parallelism: int, **read_args):
+        import pyarrow as pa
+
+        def load_data():
+            data_raw = load_breast_cancer(as_frame=True)
+            dataset_df = data_raw["data"]
+            dataset_df["target"] = data_raw["target"]
+            if self._shuffle:
+                dataset_df = shuffle(dataset_df)
+            return [pa.Table.from_pandas(dataset_df)]
+
+        meta = BlockMetadata(
+            num_rows=None,
+            size_bytes=None,
+            input_files=None,
+            exec_stats=None,
+        )
+        return [ReadTask(load_data, meta)]
+
+
+def gen_dataset_func(do_shuffle: Optional[bool] = False) -> Dataset:
+    test_datasource = TestDatasource(do_shuffle)
+    return read_datasource(test_datasource, override_num_blocks=1)
+
+
+def gen_dataset_func_eager():
+    data_raw = load_breast_cancer(as_frame=True)
+    dataset_df = data_raw["data"]
+    dataset_df["target"] = data_raw["target"]
+    dataset = from_pandas(dataset_df)
+    return dataset
+
+
+class TunerTest(unittest.TestCase):
+    """The e2e test for hparam tuning using Tuner API."""
+
+    @pytest.fixture(autouse=True)
+    def tmp_path(self, tmp_path):
+        self.tmp_path = tmp_path
+
+    def setUp(self):
+        ray.init()
+
+    def tearDown(self):
+        ray.shutdown()
+
+    def test_tuner_with_xgboost_trainer(self):
+        """Test a successful run."""
+        trainer = XGBoostTrainer(
+            label_column="target",
+            params={},
+            datasets={"train": gen_dataset_func_eager()},
+        )
+        param_space = {
+            "scaling_config": ray.train.ScalingConfig(
+                num_workers=tune.grid_search([1, 2])
+            ),
+            "datasets": {
+                "train": tune.grid_search(
+                    [gen_dataset_func(), gen_dataset_func(do_shuffle=True)]
+                ),
+            },
+            "params": {
+                "objective": "binary:logistic",
+                "tree_method": "approx",
+                "eval_metric": ["logloss", "error"],
+                "eta": tune.loguniform(1e-4, 1e-1),
+                "subsample": tune.uniform(0.5, 1.0),
+                "max_depth": tune.randint(1, 9),
+            },
+        }
+        tuner = Tuner(
+            trainable=trainer,
+            run_config=RunConfig(name="test_tuner"),
+            param_space=param_space,
+            tune_config=TuneConfig(mode="min", metric="train-error"),
+            # limiting the number of trials running at one time.
+            # As the unit test only has access to 4 CPUs on Buildkite.
+            _tuner_kwargs={"max_concurrent_trials": 1},
+        )
+        results = tuner.fit()
+        assert len(results) == 4
+
+    def test_tuner_with_xgboost_trainer_driver_fail_and_resume(self):
+        # So that we have some global checkpointing happening.
+        os.environ["TUNE_GLOBAL_CHECKPOINT_S"] = "1"
+        trainer = XGBoostTrainer(
+            label_column="target",
+            params={},
+            datasets={"train": gen_dataset_func_eager()},
+        )
+        # prep_v1 = StandardScaler(["worst radius", "worst area"])
+        # prep_v2 = StandardScaler(["worst concavity", "worst smoothness"])
+        param_space = {
+            "scaling_config": ray.train.ScalingConfig(
+                num_workers=tune.grid_search([1, 2])
+            ),
+            "datasets": {
+                "train": tune.grid_search(
+                    [gen_dataset_func(), gen_dataset_func(do_shuffle=True)]
+                ),
+            },
+            "params": {
+                "objective": "binary:logistic",
+                "tree_method": "approx",
+                "eval_metric": ["logloss", "error"],
+                "eta": tune.loguniform(1e-4, 1e-1),
+                "subsample": tune.uniform(0.5, 1.0),
+                "max_depth": tune.randint(1, 9),
+            },
+        }
+
+        class FailureInjectionCallback(Callback):
+            """Inject failure at the configured iteration number."""
+
+            def __init__(self, num_iters=10):
+                self.num_iters = num_iters
+
+            def on_step_end(self, iteration, trials, **kwargs):
+                if iteration == self.num_iters:
+                    print(f"Failing after {self.num_iters} iters.")
+                    raise RuntimeError
+
+        tuner = Tuner(
+            trainable=trainer,
+            run_config=RunConfig(
+                name="test_tuner_driver_fail",
+                storage_path=str(self.tmp_path),
+                callbacks=[FailureInjectionCallback()],
+            ),
+            param_space=param_space,
+            tune_config=TuneConfig(mode="min", metric="train-error"),
+            # limiting the number of trials running at one time.
+            # As the unit test only has access to 4 CPUs on Buildkite.
+            _tuner_kwargs={"max_concurrent_trials": 1},
+        )
+        with self.assertRaises(RuntimeError):
+            tuner.fit()
+
+        # Test resume
+        restore_path = os.path.join(self.tmp_path, "test_tuner_driver_fail")
+        tuner = Tuner.restore(restore_path, trainable=trainer, param_space=param_space)
+        # A hack before we figure out RunConfig semantics across resumes.
+        tuner._local_tuner._run_config.callbacks = None
+        results = tuner.fit()
+        assert len(results) == 4
+        assert not results.errors
+
+    def test_tuner_with_torch_trainer(self):
+        """Test a successful run using torch trainer."""
+        # The following two should be tunable.
+        config = {"lr": 1e-2, "hidden_size": 1, "batch_size": 4, "epochs": 10}
+        scaling_config = ray.train.ScalingConfig(num_workers=1, use_gpu=False)
+        trainer = TorchTrainer(
+            train_loop_per_worker=linear_train_func,
+            train_loop_config=config,
+            scaling_config=scaling_config,
+        )
+        param_space = {
+            "scaling_config": ray.train.ScalingConfig(
+                num_workers=tune.grid_search([1, 2])
+            ),
+            "train_loop_config": {
+                "batch_size": tune.grid_search([4, 8]),
+                "epochs": tune.grid_search([5, 10]),
+            },
+        }
+        tuner = Tuner(
+            trainable=trainer,
+            run_config=RunConfig(name="test_tuner"),
+            param_space=param_space,
+            tune_config=TuneConfig(mode="min", metric="loss"),
+        )
+        results = tuner.fit()
+        assert len(results) == 8
+
+    def test_tuner_run_config_override(self):
+        trainer = DummyTrainer(run_config=RunConfig(stop={"metric": 4}))
+        tuner = Tuner(trainer)
+
+        assert tuner._local_tuner._run_config.stop == {"metric": 4}
+
+
+@pytest.mark.parametrize(
+    "params_expected",
+    [
+        (
+            {"run_config": RunConfig(progress_reporter=CLIReporter())},
+            lambda kw: isinstance(kw["progress_reporter"], CLIReporter),
+        ),
+        (
+            {"tune_config": TuneConfig(reuse_actors=True)},
+            lambda kw: kw["reuse_actors"] is True,
+        ),
+        (
+            {"run_config": RunConfig(log_to_file="some_file")},
+            lambda kw: kw["log_to_file"] == "some_file",
+        ),
+        (
+            {"tune_config": TuneConfig(max_concurrent_trials=3)},
+            lambda kw: kw["max_concurrent_trials"] == 3,
+        ),
+        (
+            {"tune_config": TuneConfig(time_budget_s=60)},
+            lambda kw: kw["time_budget_s"] == 60,
+        ),
+    ],
+)
+def test_tuner_api_kwargs(shutdown_only, params_expected):
+    tuner_params, assertion = params_expected
+
+    tuner = Tuner(lambda config: 1, **tuner_params)
+
+    caught_kwargs = {}
+
+    class MockExperimentAnalysis:
+        trials = []
+
+    def catch_kwargs(**kwargs):
+        caught_kwargs.update(kwargs)
+        return MockExperimentAnalysis()
+
+    with patch("ray.tune.impl.tuner_internal.run", catch_kwargs):
+        tuner.fit()
+
+    assert assertion(caught_kwargs)
+
+
+def test_tuner_fn_trainable_invalid_checkpoint_config(shutdown_only):
+    tuner = Tuner(
+        lambda config: 1,
+        run_config=RunConfig(
+            checkpoint_config=CheckpointConfig(checkpoint_at_end=True)
+        ),
+    )
+    with pytest.raises(ValueError):
+        tuner.fit()
+
+    tuner = Tuner(
+        lambda config: 1,
+        run_config=RunConfig(
+            checkpoint_config=CheckpointConfig(checkpoint_frequency=1)
+        ),
+    )
+    with pytest.raises(ValueError):
+        tuner.fit()
+
+
+def test_tuner_trainer_checkpoint_config(shutdown_only):
+    custom_training_loop_trainer = DataParallelTrainer(
+        train_loop_per_worker=lambda config: 1
+    )
+    tuner = Tuner(
+        custom_training_loop_trainer,
+        run_config=RunConfig(
+            checkpoint_config=CheckpointConfig(checkpoint_at_end=True)
+        ),
+    )
+    with pytest.raises(ValueError):
+        tuner.fit()
+
+    tuner = Tuner(
+        custom_training_loop_trainer,
+        run_config=RunConfig(
+            checkpoint_config=CheckpointConfig(checkpoint_frequency=1)
+        ),
+    )
+    with pytest.raises(ValueError):
+        tuner.fit()
+
+    handles_checkpoints_trainer = XGBoostTrainer(
+        label_column="target",
+        params={},
+        datasets={"train": ray.data.from_items(list(range(5)))},
+    )
+    tuner = Tuner(
+        handles_checkpoints_trainer,
+        run_config=RunConfig(
+            checkpoint_config=CheckpointConfig(
+                checkpoint_at_end=True, checkpoint_frequency=1
+            )
+        ),
+    )._local_tuner
+    # Check that validation passes for a Trainer that does handle checkpointing
+    tuner._get_tune_run_arguments(tuner.converted_trainable)
+
+
+def test_tuner_fn_trainable_checkpoint_at_end_false(shutdown_only):
+    tuner = Tuner(
+        lambda config: 1,
+        run_config=RunConfig(
+            checkpoint_config=CheckpointConfig(checkpoint_at_end=False)
+        ),
+    )
+    tuner.fit()
+
+
+def test_tuner_fn_trainable_checkpoint_at_end_none(shutdown_only):
+    tuner = Tuner(
+        lambda config: 1,
+        run_config=RunConfig(
+            checkpoint_config=CheckpointConfig(checkpoint_at_end=None)
+        ),
+    )
+    tuner.fit()
+
+
+def test_nonserializable_trainable():
+    import threading
+
+    lock = threading.Lock()
+    # Check that the `inspect_serializability` trace was printed
+    with pytest.raises(TypeError, match=r".*was found to be non-serializable.*"):
+        Tuner(lambda config: print(lock))
+
+
+# TODO: [V2] Delete the `trainer` variant once V1 is fully removed.
+def _test_no_chdir(runner_type, runtime_env, use_deprecated_config=False):
+    # Write a data file that we want to read in our training loop
+    with open("./read.txt", "w") as f:
+        f.write("data")
+
+    ray.init(num_cpus=4, runtime_env=runtime_env)
+
+    def train_func(config):
+        # Make sure we can access the data from the original working dir
+        assert os.path.exists("./read.txt") and open("./read.txt", "r").read() == "data"
+
+        # Write operations should happen in each trial's independent logdir to
+        # prevent write conflicts
+        trial_dir = Path(tune.get_context().get_trial_dir())
+        trial_dir.joinpath("write.txt").touch()
+
+    if runner_type == "trainer":
+        trainer = DataParallelTrainer(
+            train_func, scaling_config=ray.train.ScalingConfig(num_workers=2)
+        )
+        result = trainer.fit()
+        results = [result]
+    elif runner_type == "tuner":
+        tuner = Tuner(
+            train_func,
+            param_space={"id": tune.grid_search(list(range(4)))},
+            tune_config=(
+                TuneConfig(chdir_to_trial_dir=False) if use_deprecated_config else None
+            ),
+        )
+        results = tuner.fit()
+        assert not results.errors
+    else:
+        raise NotImplementedError(f"Invalid: {runner_type}")
+
+    for result in results:
+        assert os.path.exists(os.path.join(result.path, "write.txt"))
+
+
+def test_tuner_no_chdir_to_trial_dir_deprecated(shutdown_only, chdir_tmpdir):
+    """Test the deprecated `chdir_to_trial_dir` config."""
+    with pytest.raises(DeprecationWarning):
+        _test_no_chdir("tuner", {}, use_deprecated_config=True)
+
+
+@pytest.mark.parametrize("runtime_env", [{}, {"working_dir": "."}])
+def test_tuner_no_chdir_to_trial_dir(
+    shutdown_only, chdir_tmpdir, monkeypatch, runtime_env
+):
+    """Tests that disabling the env var to keep the working directory the same
+    works for a Tuner run."""
+    from ray.train.constants import RAY_CHDIR_TO_TRIAL_DIR
+
+    monkeypatch.setenv(RAY_CHDIR_TO_TRIAL_DIR, "0")
+    _test_no_chdir("tuner", runtime_env)
+
+
+@pytest.mark.parametrize("runtime_env", [{}, {"working_dir": "."}])
+def test_trainer_no_chdir_to_trial_dir(
+    shutdown_only, chdir_tmpdir, monkeypatch, runtime_env
+):
+    """Tests that disabling the env var to keep the working directory the same
+    works for a Trainer run."""
+    from ray.train.constants import RAY_CHDIR_TO_TRIAL_DIR
+
+    monkeypatch.setenv(RAY_CHDIR_TO_TRIAL_DIR, "0")
+    _test_no_chdir("trainer", runtime_env)
+
+
+@pytest.mark.parametrize("runtime_env", [{}, {"working_dir": "."}])
+def test_tuner_relative_pathing_with_env_vars(
+    shutdown_only, chdir_tmpdir, tmp_path, runtime_env
+):
+    """Tests that `TUNE_ORIG_WORKING_DIR` environment variable can be used to access
+    relative paths to the original working directory.
+    """
+    # Write a data file that we want to read in our training loop
+    with open("./read.txt", "w") as f:
+        f.write("data")
+
+    # Even if we set our runtime_env `{"working_dir": "."}` to the current directory,
+    # Tune should still chdir to the trial directory.
+    ray.init(num_cpus=1, runtime_env=runtime_env)
+
+    def train_func(config):
+        orig_working_dir = Path(os.environ["TUNE_ORIG_WORKING_DIR"])
+        assert (
+            str(orig_working_dir) != os.getcwd()
+        ), f"Working directory should have changed from {orig_working_dir}"
+
+        # Make sure we can access the data from the original working dir
+        # Different from above: create an absolute path using the env variable
+        data_path = orig_working_dir / "read.txt"
+        assert os.path.exists(data_path) and open(data_path, "r").read() == "data"
+
+        # Tune chdirs to the trial working directory
+        storage = tune.get_context().get_storage()
+        assert Path(storage.trial_working_directory).resolve() == Path.cwd().resolve()
+
+        with open("write.txt", "w") as f:
+            f.write(f"{config['id']}")
+
+    tuner = Tuner(
+        train_func,
+        param_space={"id": tune.grid_search(list(range(4)))},
+        run_config=RunConfig(
+            storage_path=str(tmp_path),
+            sync_config=tune.SyncConfig(sync_artifacts=True),
+        ),
+    )
+    results = tuner.fit()
+    assert not results.errors
+    for result in results:
+        artifact_data = open(os.path.join(result.path, "write.txt"), "r").read()
+        assert artifact_data == f"{result.config['id']}"
+
+
+def test_invalid_param_space(shutdown_only):
+    """Check that Tune raises an error on invalid param_space types."""
+
+    def trainable(config):
+        return {"metric": 1}
+
+    with pytest.raises(ValueError):
+        Tuner(trainable, param_space="not allowed")
+
+    from ray.tune.tune import _Config
+
+    class CustomConfig(_Config):
+        def to_dict(self) -> dict:
+            return {"hparam": 1}
+
+    with pytest.raises(ValueError):
+        Tuner(trainable, param_space="not allowed").fit()
+
+    with pytest.raises(ValueError):
+        tune.run(trainable, config="not allowed")
+
+    # Dict and custom _Config subclasses are fine
+    Tuner(trainable, param_space={}).fit()
+    Tuner(trainable, param_space=CustomConfig()).fit()
+    tune.run(trainable, config=CustomConfig())
+
+
+def test_tuner_restore_classmethod():
+    tuner = Tuner(lambda x: None)
+
+    # Calling `tuner.restore()` on an instance should raise an AttributeError
+    with pytest.raises(AttributeError):
+        tuner.restore("/", lambda x: None)
+
+    # Calling `Tuner.restore()` on the class should work. This will throw a
+    # FileNotFoundError because no checkpoint exists at that location. Since
+    # this happens in the downstream restoration code, this means that the
+    # classmethod check successfully passed.
+    with pytest.raises(FileNotFoundError):
+        tuner = Tuner.restore("/invalid", lambda x: None)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__] + sys.argv[1:]))

python/ray/tune/tests/test_util_file_transfer.py

@@ -0,0 +1,247 @@
+import io
+import os
+import shutil
+import tarfile
+import tempfile
+
+import pytest
+
+import ray.util
+from ray.exceptions import RayTaskError
+from ray.tune.utils.file_transfer import (
+    _sync_dir_between_different_nodes,
+    _sync_dir_on_same_node,
+)
+
+
+@pytest.fixture
+def ray_start_2_cpus():
+    address_info = ray.init(num_cpus=2, configure_logging=False)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+@pytest.fixture
+def temp_data_dirs():
+    tmp_source = os.path.realpath(tempfile.mkdtemp())
+    tmp_target = os.path.realpath(tempfile.mkdtemp())
+
+    os.makedirs(os.path.join(tmp_source, "subdir", "nested"))
+    os.makedirs(os.path.join(tmp_source, "subdir_exclude", "something"))
+
+    files = [
+        "level0.txt",
+        "level0_exclude.txt",
+        "subdir/level1.txt",
+        "subdir/level1_exclude.txt",
+        "subdir/nested/level2.txt",
+        "subdir_nested_level2_exclude.txt",
+        "subdir_exclude/something/somewhere.txt",
+    ]
+
+    for file in files:
+        with open(os.path.join(tmp_source, file), "w") as f:
+            f.write("Data")
+
+    yield tmp_source, tmp_target
+
+    shutil.rmtree(tmp_source)
+    shutil.rmtree(tmp_target)
+
+
+def assert_file(exists: bool, root: str, path: str):
+    full_path = os.path.join(root, path)
+
+    if exists:
+        assert os.path.exists(full_path)
+    else:
+        assert not os.path.exists(full_path)
+
+
+def test_sync_nodes(ray_start_2_cpus, temp_data_dirs):
+    """Check that syncing between nodes works (data is found in target directory)"""
+    tmp_source, tmp_target = temp_data_dirs
+
+    assert_file(True, tmp_source, "level0.txt")
+    assert_file(True, tmp_source, "subdir/level1.txt")
+    assert_file(False, tmp_target, "level0.txt")
+    assert_file(False, tmp_target, "subdir/level1.txt")
+
+    node_ip = ray.util.get_node_ip_address()
+    _sync_dir_between_different_nodes(
+        source_ip=node_ip,
+        source_path=tmp_source,
+        target_ip=node_ip,
+        target_path=tmp_target,
+    )
+
+    assert_file(True, tmp_target, "level0.txt")
+    assert_file(True, tmp_target, "subdir/level1.txt")
+
+
+def test_sync_nodes_only_diff(ray_start_2_cpus, temp_data_dirs):
+    """Check that only differing files are synced between nodes"""
+    tmp_source, tmp_target = temp_data_dirs
+
+    # Sanity check
+    assert_file(True, tmp_source, "level0.txt")
+    assert_file(True, tmp_source, "subdir/level1.txt")
+    assert_file(False, tmp_target, "level0.txt")
+    assert_file(False, tmp_target, "level0_new.txt")
+
+    node_ip = ray.util.get_node_ip_address()
+    _sync_dir_between_different_nodes(
+        source_ip=node_ip,
+        source_path=tmp_source,
+        target_ip=node_ip,
+        target_path=tmp_target,
+    )
+
+    assert_file(True, tmp_source, "level0.txt")
+    assert_file(True, tmp_target, "level0.txt")
+    assert_file(True, tmp_target, "subdir/level1.txt")
+    assert_file(False, tmp_target, "level0_new.txt")
+
+    # Add new file
+    with open(os.path.join(tmp_source, "level0_new.txt"), "w") as f:
+        f.write("Data\n")
+
+    # Modify existing file
+    with open(os.path.join(tmp_source, "subdir", "level1.txt"), "w") as f:
+        f.write("New data\n")
+
+    unpack, pack_actor, files_stats = _sync_dir_between_different_nodes(
+        source_ip=node_ip,
+        source_path=tmp_source,
+        target_ip=node_ip,
+        target_path=tmp_target,
+        return_futures=True,
+    )
+
+    files_stats = ray.get(files_stats)
+    tarball = ray.get(pack_actor.get_full_data.remote())
+
+    assert "./level0.txt" in files_stats
+    assert "./level0_new.txt" not in files_stats  # Was not in target dir
+    assert "subdir/level1.txt" in files_stats
+
+    with tarfile.open(fileobj=io.BytesIO(tarball)) as tar:
+        files_in_tar = tar.getnames()
+        assert "./level0.txt" not in files_in_tar
+        assert "./level0_new.txt" in files_in_tar
+        assert "subdir/level1.txt" in files_in_tar
+        assert len(files_in_tar) == 7  # 3 files, 4 dirs (including root)
+
+    ray.get(unpack)  # Wait until finished for teardown
+
+
+@pytest.mark.parametrize("exclude", [["subdir/*"], ["*/level1.txt"]])
+def test_sync_nodes_exclude_different_node(ray_start_2_cpus, temp_data_dirs, exclude):
+    """Check that excluding files works"""
+    tmp_source, tmp_target = temp_data_dirs
+
+    assert_file(True, tmp_source, "level0.txt")
+    assert_file(True, tmp_source, "subdir/level1.txt")
+    assert_file(False, tmp_target, "level0.txt")
+    assert_file(False, tmp_target, "subdir/level1.txt")
+
+    node_ip = ray.util.get_node_ip_address()
+    _sync_dir_between_different_nodes(
+        source_ip=node_ip,
+        source_path=tmp_source,
+        target_ip=node_ip,
+        target_path=tmp_target,
+        exclude=exclude,
+    )
+
+    assert_file(True, tmp_target, "level0.txt")
+    assert_file(False, tmp_target, "subdir/level1.txt")
+
+
+@pytest.mark.parametrize("exclude", [["subdir/*"], ["*/level1.txt"]])
+def test_sync_nodes_exclude_same_node(ray_start_2_cpus, temp_data_dirs, exclude):
+    """Check that excluding files works"""
+    tmp_source, tmp_target = temp_data_dirs
+
+    assert_file(True, tmp_source, "level0.txt")
+    assert_file(True, tmp_source, "subdir/level1.txt")
+    assert_file(False, tmp_target, "level0.txt")
+    assert_file(False, tmp_target, "subdir/level1.txt")
+
+    node_ip = ray.util.get_node_ip_address()
+    _sync_dir_on_same_node(
+        ip=node_ip, source_path=tmp_source, target_path=tmp_target, exclude=exclude
+    )
+
+    assert_file(True, tmp_target, "level0.txt")
+    assert_file(False, tmp_target, "subdir/level1.txt")
+
+
+@pytest.mark.parametrize("num_workers", [1, 8])
+def test_multi_sync_same_node(ray_start_2_cpus, temp_data_dirs, num_workers):
+    """Check that multiple competing syncs to the same node+dir don't interfere"""
+    tmp_source, tmp_target = temp_data_dirs
+
+    assert_file(True, tmp_source, "level0.txt")
+    assert_file(True, tmp_source, "subdir/level1.txt")
+
+    node_ip = ray.util.get_node_ip_address()
+    futures = [
+        _sync_dir_on_same_node(
+            ip=node_ip,
+            source_path=tmp_source,
+            target_path=tmp_target,
+            return_futures=True,
+        )
+        for _ in range(num_workers)
+    ]
+    ray.get(futures)
+
+    assert_file(True, tmp_target, "level0.txt")
+    assert_file(True, tmp_target, "subdir/level1.txt")
+
+
+@pytest.mark.parametrize("num_workers", [1, 8])
+def test_multi_sync_different_node(ray_start_2_cpus, temp_data_dirs, num_workers):
+    """Check that multiple competing syncs to the same dir don't interfere"""
+    tmp_source, tmp_target = temp_data_dirs
+
+    assert_file(True, tmp_source, "level0.txt")
+    assert_file(True, tmp_source, "subdir/level1.txt")
+
+    node_ip = ray.util.get_node_ip_address()
+    futures = [
+        _sync_dir_between_different_nodes(
+            source_ip=node_ip,
+            source_path=tmp_source,
+            target_ip=node_ip,
+            target_path=tmp_target,
+            return_futures=True,
+        )[0]
+        for _ in range(num_workers)
+    ]
+    ray.get(futures)
+
+    assert_file(True, tmp_target, "level0.txt")
+    assert_file(True, tmp_target, "subdir/level1.txt")
+
+
+def test_max_size_exceeded(ray_start_2_cpus, temp_data_dirs):
+    tmp_source, tmp_target = temp_data_dirs
+
+    node_ip = ray.util.get_node_ip_address()
+    with pytest.raises(RayTaskError):
+        _sync_dir_between_different_nodes(
+            source_ip=node_ip,
+            source_path=tmp_source,
+            target_ip=node_ip,
+            target_path=tmp_target,
+            max_size_bytes=2,
+        )
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/tune/tests/test_util_object_cache.py

@@ -0,0 +1,124 @@
+import pytest
+
+from ray.tune.utils.object_cache import _ObjectCache
+
+
+@pytest.mark.parametrize("eager", [False, True])
+def test_no_may_keep_one(eager):
+    """Test object caching.
+
+    - After init, no objects are cached (as max cached is 0), except when eager caching
+    - After increasing max to 2, up to 2 objects are cached
+    - Decreasing max objects will evict them on flush
+    """
+    cache = _ObjectCache(may_keep_one=eager)
+
+    # max(A) = 0, so we we only cache when eager caching
+    assert cache.cache_object("A", 1) == eager
+    assert cache.num_cached_objects == int(eager)
+
+    # Set max(A) = 2
+    cache.increase_max("A", 2)
+
+    # max(A) = 2, so we cache up to two objects
+    if not eager:
+        assert cache.cache_object("A", 1)
+
+    assert cache.cache_object("A", 2)
+    assert not cache.cache_object("A", 3)
+
+    assert cache.num_cached_objects == 2
+
+    # Nothing has to be evicted
+    assert not list(cache.flush_cached_objects())
+
+    # Set max(A) = 1, so we have one object too much
+    cache.decrease_max("A", 1)
+
+    # First cached object is evicted
+    assert list(cache.flush_cached_objects()) == [1]
+    assert cache.num_cached_objects == 1
+
+    # Set max(A) = 0
+    cache.decrease_max("A", 1)
+
+    # Second cached object is evicted if not eager caching
+    assert list(cache.flush_cached_objects()) == ([2] if not eager else [])
+    assert cache.num_cached_objects == (0 if not eager else 1)
+
+
+@pytest.mark.parametrize("eager", [False, True])
+def test_multi(eager):
+    """Test caching with multiple objects"""
+    cache = _ObjectCache(may_keep_one=eager)
+
+    # max(A) = 0, so we we only cache when eager caching
+    assert cache.cache_object("A", 1) == eager
+    assert cache.num_cached_objects == int(eager)
+
+    # max(B) = 0, so no caching
+    assert not cache.cache_object("B", 5)
+    assert cache.num_cached_objects == int(eager)
+
+    # Increase maximums levels
+    cache.increase_max("A", 1)
+    cache.increase_max("B", 1)
+
+    # Cache objects (A is already cached if eager)
+    assert cache.cache_object("A", 1) != eager
+    assert cache.cache_object("B", 5)
+
+    # No further objects can be cached
+    assert not cache.cache_object("A", 2)
+    assert not cache.cache_object("B", 6)
+
+    assert cache.num_cached_objects == 2
+
+    # Decrease
+    cache.decrease_max("A", 1)
+
+    # Evict A object
+    assert list(cache.flush_cached_objects()) == [1]
+
+    cache.decrease_max("B", 1)
+
+    # If eager, keep B object, otherwise, evict B
+    assert list(cache.flush_cached_objects()) == ([5] if not eager else [])
+    assert cache.num_cached_objects == (0 if not eager else 1)
+
+
+def test_multi_eager_other():
+    """On eager caching, only cache an object if no other object is expected.
+
+    - Expect up to one cached A object
+    - Try to cache object B --> doesn't get cached
+    - Remove expectation for A object
+    - Try to cache object B --> get's cached
+    """
+    cache = _ObjectCache(may_keep_one=True)
+
+    cache.increase_max("A", 1)
+    assert not cache.cache_object("B", 2)
+
+    cache.decrease_max("A", 1)
+    assert cache.cache_object("B", 3)
+
+
+@pytest.mark.parametrize("eager", [False, True])
+def test_force_all(eager):
+    """Assert that force_all=True will always evict all object."""
+    cache = _ObjectCache(may_keep_one=eager)
+
+    cache.increase_max("A", 2)
+
+    assert cache.cache_object("A", 1)
+    assert cache.cache_object("A", 2)
+
+    assert list(cache.flush_cached_objects(force_all=True)) == [1, 2]
+    assert cache.num_cached_objects == 0
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/tune/tests/test_utils.py

@@ -0,0 +1,111 @@
+import io
+import logging
+import sys
+import time
+
+import pytest
+
+from ray.tune.search.variant_generator import format_vars
+from ray.tune.utils.util import Tee, logger as util_logger, retry_fn
+
+
+def test_format_vars():
+
+    # Format brackets correctly
+    assert (
+        format_vars(
+            {
+                ("a", "b", "c"): 8.1234567,
+                ("a", "b", "d"): [7, 8],
+                ("a", "b", "e"): [[[3, 4]]],
+            }
+        )
+        == "c=8.1235,d=7_8,e=3_4"
+    )
+    # Sorted by full keys, but only last key is reported
+    assert (
+        format_vars(
+            {
+                ("a", "c", "x"): [7, 8],
+                ("a", "b", "x"): 8.1234567,
+            }
+        )
+        == "x=8.1235,x=7_8"
+    )
+    # Filter out invalid chars. It's ok to have empty keys or values.
+    assert (
+        format_vars(
+            {
+                ("a  c?x",): " <;%$ok ",
+                ("some",): " ",
+            }
+        )
+        == "a_c_x=ok,some="
+    )
+
+
+def test_retry_fn_repeat(tmpdir):
+    success = tmpdir / "success"
+    marker = tmpdir / "marker"
+
+    def _fail_once():
+        if marker.exists():
+            success.write_text(".", encoding="utf-8")
+            return
+        marker.write_text(".", encoding="utf-8")
+        raise RuntimeError("Failing")
+
+    assert not success.exists()
+    assert not marker.exists()
+
+    assert retry_fn(
+        fn=_fail_once,
+        exception_type=RuntimeError,
+        sleep_time=0,
+    )
+
+    assert success.exists()
+    assert marker.exists()
+
+
+def test_retry_fn_timeout(tmpdir):
+    success = tmpdir / "success"
+    marker = tmpdir / "marker"
+
+    def _fail_once():
+        if not marker.exists():
+            marker.write_text(".", encoding="utf-8")
+            raise RuntimeError("Failing")
+        time.sleep(5)
+        success.write_text(".", encoding="utf-8")
+        return
+
+    assert not success.exists()
+    assert not marker.exists()
+
+    assert not retry_fn(
+        fn=_fail_once, exception_type=RuntimeError, sleep_time=0, timeout=0.1
+    )
+
+    assert not success.exists()
+    assert marker.exists()
+
+
+def test_tee_recursion():
+    f = io.StringIO()
+    g = io.StringIO()
+
+    tee = Tee(f, g)
+
+    hdlr = logging.StreamHandler(tee)
+    util_logger.addHandler(hdlr)
+
+    util_logger.info("BEFORE")
+    f.close()
+    util_logger.info("AFTER")
+
+    util_logger.removeHandler(hdlr)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/tune/tests/test_var.py

@@ -0,0 +1,370 @@
+import os
+import random
+import unittest
+
+import numpy as np
+
+import ray
+from ray import tune
+from ray.train.constants import DEFAULT_STORAGE_PATH
+from ray.tune.search import BasicVariantGenerator, grid_search
+from ray.tune.search.variant_generator import (
+    RecursiveDependencyError,
+    _resolve_nested_dict,
+)
+from ray.tune.utils.mock_trainable import MOCK_TRAINABLE_NAME, register_mock_trainable
+
+
+class VariantGeneratorTest(unittest.TestCase):
+    def setUp(self):
+        ray.init(num_cpus=2)
+
+        register_mock_trainable()
+
+    def tearDown(self):
+        ray.shutdown()
+
+    def generate_trials(self, spec, name):
+        suggester = BasicVariantGenerator()
+        suggester.add_configurations({name: spec})
+        trials = []
+        while not suggester.is_finished():
+            trial = suggester.next_trial()
+            if trial:
+                trials.append(trial)
+            else:
+                break
+        return trials
+
+    def testParseToTrials(self):
+        trials = self.generate_trials(
+            {
+                "run": MOCK_TRAINABLE_NAME,
+                "num_samples": 2,
+                "max_failures": 5,
+                "config": {"env": "Pong-v0", "foo": "bar"},
+            },
+            "tune-pong",
+        )
+        trials = list(trials)
+        self.assertEqual(len(trials), 2)
+        self.assertTrue(MOCK_TRAINABLE_NAME in str(trials[0]))
+        self.assertEqual(trials[0].config, {"foo": "bar", "env": "Pong-v0"})
+        self.assertEqual(trials[0].trainable_name, MOCK_TRAINABLE_NAME)
+        self.assertEqual(trials[0].experiment_tag, "0")
+        self.assertEqual(trials[0].max_failures, 5)
+        self.assertEqual(trials[0].evaluated_params, {})
+        self.assertEqual(
+            trials[0].storage.experiment_fs_path,
+            os.path.join(DEFAULT_STORAGE_PATH, "tune-pong"),
+        )
+        self.assertEqual(trials[1].experiment_tag, "1")
+
+    def testEval(self):
+        trials = self.generate_trials(
+            {
+                "run": MOCK_TRAINABLE_NAME,
+                "config": {
+                    "foo": {"eval": "2 + 2"},
+                },
+            },
+            "eval",
+        )
+        trials = list(trials)
+        self.assertEqual(len(trials), 1)
+        self.assertEqual(trials[0].config, {"foo": 4})
+        self.assertEqual(trials[0].evaluated_params, {"foo": 4})
+        self.assertEqual(trials[0].experiment_tag, "0_foo=4")
+
+    def testGridSearch(self):
+        trials = self.generate_trials(
+            {
+                "run": MOCK_TRAINABLE_NAME,
+                "config": {
+                    "bar": {"grid_search": [True, False]},
+                    "foo": {"grid_search": [1, 2, 3]},
+                    "baz": "asd",
+                },
+            },
+            "grid_search",
+        )
+        trials = list(trials)
+        self.assertEqual(len(trials), 6)
+        self.assertEqual(
+            trials[0].config,
+            {
+                "bar": True,
+                "foo": 1,
+                "baz": "asd",
+            },
+        )
+        self.assertEqual(
+            trials[0].evaluated_params,
+            {
+                "bar": True,
+                "foo": 1,
+            },
+        )
+        self.assertEqual(trials[0].experiment_tag, "0_bar=True,foo=1")
+
+        self.assertEqual(
+            trials[1].config,
+            {
+                "bar": False,
+                "foo": 1,
+                "baz": "asd",
+            },
+        )
+        self.assertEqual(
+            trials[1].evaluated_params,
+            {
+                "bar": False,
+                "foo": 1,
+            },
+        )
+        self.assertEqual(trials[1].experiment_tag, "1_bar=False,foo=1")
+
+        self.assertEqual(
+            trials[2].config,
+            {
+                "bar": True,
+                "foo": 2,
+                "baz": "asd",
+            },
+        )
+        self.assertEqual(
+            trials[2].evaluated_params,
+            {
+                "bar": True,
+                "foo": 2,
+            },
+        )
+
+        self.assertEqual(
+            trials[3].config,
+            {
+                "bar": False,
+                "foo": 2,
+                "baz": "asd",
+            },
+        )
+        self.assertEqual(
+            trials[3].evaluated_params,
+            {
+                "bar": False,
+                "foo": 2,
+            },
+        )
+
+        self.assertEqual(
+            trials[4].config,
+            {
+                "bar": True,
+                "foo": 3,
+                "baz": "asd",
+            },
+        )
+        self.assertEqual(
+            trials[4].evaluated_params,
+            {
+                "bar": True,
+                "foo": 3,
+            },
+        )
+
+        self.assertEqual(
+            trials[5].config,
+            {
+                "bar": False,
+                "foo": 3,
+                "baz": "asd",
+            },
+        )
+        self.assertEqual(
+            trials[5].evaluated_params,
+            {
+                "bar": False,
+                "foo": 3,
+            },
+        )
+
+    def testGridSearchAndEval(self):
+        trials = self.generate_trials(
+            {
+                "run": MOCK_TRAINABLE_NAME,
+                "config": {
+                    "qux": tune.sample_from(lambda spec: 2 + 2),
+                    "bar": grid_search([True, False]),
+                    "foo": grid_search([1, 2, 3]),
+                    "baz": "asd",
+                },
+            },
+            "grid_eval",
+        )
+        trials = list(trials)
+        self.assertEqual(len(trials), 6)
+        self.assertEqual(
+            trials[0].config,
+            {
+                "bar": True,
+                "foo": 1,
+                "qux": 4,
+                "baz": "asd",
+            },
+        )
+        self.assertEqual(
+            trials[0].evaluated_params,
+            {
+                "bar": True,
+                "foo": 1,
+                "qux": 4,
+            },
+        )
+        self.assertEqual(trials[0].experiment_tag, "0_bar=True,foo=1,qux=4")
+
+    def testConditionResolution(self):
+        trials = self.generate_trials(
+            {
+                "run": MOCK_TRAINABLE_NAME,
+                "config": {
+                    "x": 1,
+                    "y": tune.sample_from(lambda spec: spec.config.x + 1),
+                    "z": tune.sample_from(lambda spec: spec.config.y + 1),
+                },
+            },
+            "condition_resolution",
+        )
+        trials = list(trials)
+        self.assertEqual(len(trials), 1)
+        self.assertEqual(trials[0].config, {"x": 1, "y": 2, "z": 3})
+        self.assertEqual(trials[0].evaluated_params, {"y": 2, "z": 3})
+        self.assertEqual(trials[0].experiment_tag, "0_y=2,z=3")
+
+    def testDependentLambda(self):
+        trials = self.generate_trials(
+            {
+                "run": MOCK_TRAINABLE_NAME,
+                "config": {
+                    "x": grid_search([1, 2]),
+                    "y": tune.sample_from(lambda spec: spec.config.x * 100),
+                },
+            },
+            "dependent_lambda",
+        )
+        trials = list(trials)
+        self.assertEqual(len(trials), 2)
+        self.assertEqual(trials[0].config, {"x": 1, "y": 100})
+        self.assertEqual(trials[1].config, {"x": 2, "y": 200})
+
+    def testDependentGridSearch(self):
+        trials = self.generate_trials(
+            {
+                "run": MOCK_TRAINABLE_NAME,
+                "config": {
+                    "x": grid_search(
+                        [
+                            tune.sample_from(lambda spec: spec.config.y * 100),
+                            tune.sample_from(lambda spec: spec.config.y * 200),
+                        ]
+                    ),
+                    "y": tune.sample_from(lambda spec: 1),
+                },
+            },
+            "dependent_grid_search",
+        )
+        trials = list(trials)
+        self.assertEqual(len(trials), 2)
+        self.assertEqual(trials[0].config, {"x": 100, "y": 1})
+        self.assertEqual(trials[1].config, {"x": 200, "y": 1})
+
+    def testDependentGridSearchCallable(self):
+        class Normal:
+            def __call__(self, _config):
+                return random.normalvariate(mu=0, sigma=1)
+
+        class Single:
+            def __call__(self, _config):
+                return 20
+
+        trials = self.generate_trials(
+            {
+                "run": MOCK_TRAINABLE_NAME,
+                "config": {
+                    "x": grid_search(
+                        [tune.sample_from(Normal()), tune.sample_from(Normal())]
+                    ),
+                    "y": tune.sample_from(Single()),
+                },
+            },
+            "dependent_grid_search",
+        )
+        trials = list(trials)
+        self.assertEqual(len(trials), 2)
+        self.assertEqual(trials[0].config["y"], 20)
+        self.assertEqual(trials[1].config["y"], 20)
+
+    def testNestedValues(self):
+        trials = self.generate_trials(
+            {
+                "run": MOCK_TRAINABLE_NAME,
+                "config": {
+                    "x": {"y": {"z": tune.sample_from(lambda spec: 1)}},
+                    "y": tune.sample_from(lambda spec: 12),
+                    "z": tune.sample_from(lambda spec: spec.config.x.y.z * 100),
+                },
+            },
+            "nested_values",
+        )
+        trials = list(trials)
+        self.assertEqual(len(trials), 1)
+        self.assertEqual(trials[0].config, {"x": {"y": {"z": 1}}, "y": 12, "z": 100})
+        self.assertEqual(trials[0].evaluated_params, {"x/y/z": 1, "y": 12, "z": 100})
+
+    def testLogUniform(self):
+        sampler = tune.loguniform(1e-10, 1e-1)
+        results = sampler.sample(None, 1000)
+        assert abs(np.log(min(results)) / np.log(10) - -10) < 0.1
+        assert abs(np.log(max(results)) / np.log(10) - -1) < 0.1
+
+        sampler_e = tune.loguniform(np.e**-4, np.e)
+        results_e = sampler_e.sample(None, 1000)
+        assert abs(np.log(min(results_e)) - -4) < 0.1
+        assert abs(np.log(max(results_e)) - 1) < 0.1
+
+    def test_resolve_dict(self):
+        config = {
+            "a": {
+                "b": 1,
+                "c": 2,
+            },
+            "b": {"a": 3},
+        }
+        resolved = _resolve_nested_dict(config)
+        for k, v in [(("a", "b"), 1), (("a", "c"), 2), (("b", "a"), 3)]:
+            self.assertEqual(resolved.get(k), v)
+
+    def testRecursiveDep(self):
+        try:
+            list(
+                self.generate_trials(
+                    {
+                        "run": MOCK_TRAINABLE_NAME,
+                        "config": {
+                            "foo": tune.sample_from(lambda spec: spec.config.foo),
+                        },
+                    },
+                    "recursive_dep",
+                )
+            )
+        except RecursiveDependencyError as e:
+            assert "`foo` recursively depends on" in str(e), e
+        else:
+            raise
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/tune/tests/test_warnings.py

@@ -0,0 +1,143 @@
+import pytest
+
+import ray
+from ray import tune
+from ray.data.context import DataContext
+from ray.tune.error import TuneError
+from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
+
+
+def test_nowarn_zero_cpu():
+    def f(*a):
+        @ray.remote(num_cpus=0)
+        def f():
+            pass
+
+        @ray.remote(num_cpus=0)
+        class Actor:
+            def f(self):
+                pass
+
+        ray.get(f.remote())
+        a = Actor.remote()
+        ray.get(a.f.remote())
+
+    tune.run(f, verbose=0)
+
+
+def test_warn_cpu():
+    def f(*a):
+        @ray.remote(num_cpus=1)
+        def f():
+            pass
+
+        ray.get(f.remote())
+
+    with pytest.raises(TuneError):
+        tune.run(f, verbose=0)
+
+    with pytest.raises(TuneError):
+        tune.run(
+            f, resources_per_trial=tune.PlacementGroupFactory([{"CPU": 1}]), verbose=0
+        )
+
+    def g(*a):
+        @ray.remote(num_cpus=1)
+        class Actor:
+            def f(self):
+                pass
+
+        a = Actor.remote()
+        ray.get(a.f.remote())
+
+    with pytest.raises(TuneError):
+        tune.run(g, verbose=0)
+
+    with pytest.raises(TuneError):
+        tune.run(
+            g, resources_per_trial=tune.PlacementGroupFactory([{"CPU": 1}]), verbose=0
+        )
+
+
+def test_pg_slots_ok():
+    def f(*a):
+        @ray.remote(num_cpus=1)
+        def f():
+            pass
+
+        @ray.remote(num_cpus=1)
+        class Actor:
+            def f(self):
+                pass
+
+        ray.get(f.remote())
+        a = Actor.remote()
+        ray.get(a.f.remote())
+
+    tune.run(
+        f, resources_per_trial=tune.PlacementGroupFactory([{"CPU": 1}] * 2), verbose=0
+    )
+
+
+def test_bad_pg_slots():
+    def f(*a):
+        @ray.remote(num_cpus=2)
+        def f():
+            pass
+
+        ray.get(f.remote())
+
+    with pytest.raises(TuneError):
+        tune.run(
+            f,
+            resources_per_trial=tune.PlacementGroupFactory([{"CPU": 1}] * 2),
+            verbose=0,
+        )
+
+
+def test_dataset_ok():
+    def f(*a):
+        ray.data.range(10).show()
+
+    tune.run(f, verbose=0)
+
+    def g(*a):
+        ctx = DataContext.get_current()
+        ctx.scheduling_strategy = PlacementGroupSchedulingStrategy(
+            ray.util.get_current_placement_group()
+        )
+        ray.data.range(10).show()
+
+    with pytest.raises(TuneError):
+        tune.run(g, verbose=0)
+
+    tune.run(
+        g, resources_per_trial=tune.PlacementGroupFactory([{"CPU": 1}] * 2), verbose=0
+    )
+
+
+def test_scheduling_strategy_override():
+    def f(*a):
+        @ray.remote(num_cpus=1, scheduling_strategy="SPREAD")
+        def f():
+            pass
+
+        @ray.remote(num_cpus=1, scheduling_strategy="SPREAD")
+        class Actor:
+            def f(self):
+                pass
+
+        # SPREAD tasks are not captured by placement groups, so don't warn.
+        ray.get(f.remote())
+
+        # SPREAD actors are not captured by placement groups, so don't warn.
+        a = Actor.remote()
+        ray.get(a.f.remote())
+
+    tune.run(f, verbose=0)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/tune/tests/tune_test_util.py

@@ -0,0 +1,51 @@
+import os
+import tempfile
+
+from ray.tune import Callback
+from ray.tune.execution.tune_controller import TuneController
+
+
+class TrialResultObserver(Callback):
+    """Helper class to control runner.step() count."""
+
+    def __init__(self):
+        self._counter = 0
+        self._last_counter = 0
+
+    def reset(self):
+        self._last_counter = self._counter
+
+    def just_received_a_result(self):
+        if self._last_counter == self._counter:
+            return False
+        else:
+            self._last_counter = self._counter
+            return True
+
+    def on_trial_result(self, **kwargs):
+        self._counter += 1
+
+
+def create_tune_experiment_checkpoint(trials: list, **runner_kwargs) -> str:
+    experiment_dir = tempfile.mkdtemp()
+    runner_kwargs.setdefault("experiment_path", experiment_dir)
+
+    # Update environment
+    orig_env = os.environ.copy()
+
+    # Set to 1 to disable ray cluster resource lookup. That way we can
+    # create experiment checkpoints without initializing ray.
+    os.environ["TUNE_MAX_PENDING_TRIALS_PG"] = "1"
+
+    try:
+        runner = TuneController(**runner_kwargs)
+
+        for trial in trials:
+            runner.add_trial(trial)
+
+        runner.checkpoint(force=True, wait=True)
+    finally:
+        os.environ.clear()
+        os.environ.update(orig_env)
+
+    return experiment_dir

python/ray/tune/tests/tutorial.py

@@ -0,0 +1,227 @@
+# ruff: noqa
+# isort: skip_file
+# Original Code: https://github.com/pytorch/examples/blob/master/mnist/main.py
+
+# fmt: off
+# __tutorial_imports_begin__
+import numpy as np
+import torch
+import torch.optim as optim
+import torch.nn as nn
+from torchvision import datasets, transforms
+from torch.utils.data import DataLoader
+import torch.nn.functional as F
+
+from ray import tune
+from ray.tune.schedulers import ASHAScheduler
+# __tutorial_imports_end__
+# fmt: on
+
+
+# fmt: off
+# __model_def_begin__
+class ConvNet(nn.Module):
+    def __init__(self):
+        super(ConvNet, self).__init__()
+        # In this example, we don't change the model architecture
+        # due to simplicity.
+        self.conv1 = nn.Conv2d(1, 3, kernel_size=3)
+        self.fc = nn.Linear(192, 10)
+
+    def forward(self, x):
+        x = F.relu(F.max_pool2d(self.conv1(x), 3))
+        x = x.view(-1, 192)
+        x = self.fc(x)
+        return F.log_softmax(x, dim=1)
+# __model_def_end__
+# fmt: on
+
+# fmt: off
+# __train_def_begin__
+
+# Change these values if you want the training to run quicker or slower.
+EPOCH_SIZE = 512
+TEST_SIZE = 256
+
+def train_func(model, optimizer, train_loader):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model.train()
+    for batch_idx, (data, target) in enumerate(train_loader):
+        # We set this just for the example to run quickly.
+        if batch_idx * len(data) > EPOCH_SIZE:
+            return
+        data, target = data.to(device), target.to(device)
+        optimizer.zero_grad()
+        output = model(data)
+        loss = F.nll_loss(output, target)
+        loss.backward()
+        optimizer.step()
+
+
+def test_func(model, data_loader):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model.eval()
+    correct = 0
+    total = 0
+    with torch.no_grad():
+        for batch_idx, (data, target) in enumerate(data_loader):
+            # We set this just for the example to run quickly.
+            if batch_idx * len(data) > TEST_SIZE:
+                break
+            data, target = data.to(device), target.to(device)
+            outputs = model(data)
+            _, predicted = torch.max(outputs.data, 1)
+            total += target.size(0)
+            correct += (predicted == target).sum().item()
+
+    return correct / total
+# __train_def_end__
+
+
+# __train_func_begin__
+import os
+import tempfile
+
+from ray.tune import Checkpoint
+
+def train_mnist(config):
+    # Data Setup
+    mnist_transforms = transforms.Compose(
+        [transforms.ToTensor(),
+         transforms.Normalize((0.1307, ), (0.3081, ))])
+
+    train_loader = DataLoader(
+        datasets.MNIST("~/data", train=True, download=True, transform=mnist_transforms),
+        batch_size=64,
+        shuffle=True)
+    test_loader = DataLoader(
+        datasets.MNIST("~/data", train=False, transform=mnist_transforms),
+        batch_size=64,
+        shuffle=True)
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    model = ConvNet()
+    model.to(device)
+
+    optimizer = optim.SGD(
+        model.parameters(), lr=config["lr"], momentum=config["momentum"])
+    for i in range(10):
+        train_func(model, optimizer, train_loader)
+        acc = test_func(model, test_loader)
+
+
+        with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
+            checkpoint = None
+            if (i + 1) % 5 == 0:
+                # This saves the model to the trial directory
+                torch.save(
+                    model.state_dict(),
+                    os.path.join(temp_checkpoint_dir, "model.pth")
+                )
+                checkpoint = Checkpoint.from_directory(temp_checkpoint_dir)
+
+            # Send the current training result back to Tune
+            tune.report({"mean_accuracy": acc}, checkpoint=checkpoint)
+
+# __train_func_end__
+# fmt: on
+
+# __eval_func_begin__
+search_space = {
+    "lr": tune.sample_from(lambda spec: 10 ** (-10 * np.random.rand())),
+    "momentum": tune.uniform(0.1, 0.9),
+}
+
+# Uncomment this to enable distributed execution
+# `ray.init(address="auto")`
+
+# Download the dataset first
+datasets.MNIST("~/data", train=True, download=True)
+
+tuner = tune.Tuner(
+    train_mnist,
+    param_space=search_space,
+)
+results = tuner.fit()
+# __eval_func_end__
+
+# __plot_begin__
+dfs = {result.path: result.metrics_dataframe for result in results}
+[d.mean_accuracy.plot() for d in dfs.values()]
+# __plot_end__
+
+# __run_scheduler_begin__
+tuner = tune.Tuner(
+    train_mnist,
+    tune_config=tune.TuneConfig(
+        num_samples=20,
+        scheduler=ASHAScheduler(metric="mean_accuracy", mode="max"),
+    ),
+    param_space=search_space,
+)
+results = tuner.fit()
+
+# Obtain a trial dataframe from all run trials of this `tune.run` call.
+dfs = {result.path: result.metrics_dataframe for result in results}
+# __run_scheduler_end__
+
+# fmt: off
+# __plot_scheduler_begin__
+# Plot by epoch
+ax = None  # This plots everything on the same plot
+for d in dfs.values():
+    ax = d.mean_accuracy.plot(ax=ax, legend=False)
+# __plot_scheduler_end__
+# fmt: on
+
+# __run_searchalg_begin__
+from hyperopt import hp
+from ray.tune.search.hyperopt import HyperOptSearch
+
+space = {
+    "lr": hp.loguniform("lr", -10, -1),
+    "momentum": hp.uniform("momentum", 0.1, 0.9),
+}
+
+hyperopt_search = HyperOptSearch(space, metric="mean_accuracy", mode="max")
+
+tuner = tune.Tuner(
+    train_mnist,
+    tune_config=tune.TuneConfig(
+        num_samples=10,
+        search_alg=hyperopt_search,
+    ),
+)
+results = tuner.fit()
+
+# To enable GPUs, use this instead:
+# analysis = tune.run(
+#     train_mnist, config=search_space, resources_per_trial={'gpu': 1})
+
+# __run_searchalg_end__
+
+# __run_analysis_begin__
+best_result = results.get_best_result("mean_accuracy", mode="max")
+with best_result.checkpoint.as_directory() as checkpoint_dir:
+    state_dict = torch.load(os.path.join(checkpoint_dir, "model.pth"))
+
+model = ConvNet()
+model.load_state_dict(state_dict)
+# __run_analysis_end__
+
+from ray.tune.examples.mnist_pytorch_trainable import TrainMNIST
+
+# __trainable_run_begin__
+search_space = {
+    "lr": tune.sample_from(lambda spec: 10 ** (-10 * np.random.rand())),
+    "momentum": tune.uniform(0.1, 0.9),
+}
+
+tuner = tune.Tuner(
+    TrainMNIST,
+    run_config=tune.RunConfig(stop={"training_iteration": 10}),
+    param_space=search_space,
+)
+results = tuner.fit()
+# __trainable_run_end__