chore: import upstream snapshot with attribution

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