chore: import upstream snapshot with attribution

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__]))