chore: import upstream snapshot with attribution

python/ray/train/v2/tests/test_checkpoint_manager.py

@@ -0,0 +1,359 @@
+import uuid
+from typing import Optional
+from unittest.mock import create_autospec
+
+import pytest
+
+import ray
+from ray.train import CheckpointConfig
+from ray.train._internal.session import _TrainingResult
+from ray.train.v2._internal.exceptions import CheckpointManagerInitializationError
+from ray.train.v2._internal.execution.checkpoint.checkpoint_manager import (
+    CheckpointManager,
+)
+from ray.train.v2._internal.execution.storage import StorageContext
+from ray.train.v2._internal.execution.worker_group import Worker
+from ray.train.v2.api.reported_checkpoint import ReportedCheckpointStatus
+from ray.train.v2.tests.util import create_dummy_training_reports
+
+
+@pytest.fixture(autouse=True, scope="module")
+def ray_start_4_cpus():
+    ray.init(num_cpus=4)
+    yield
+    ray.shutdown()
+
+
+def _checkpoint_managers_equal(cm1: CheckpointManager, cm2: CheckpointManager) -> bool:
+    """
+    Compare two checkpoint managers for equality.
+    Ignore uuid differences of all the checkpoints recorded.
+    """
+
+    def _training_results_equal(
+        tr1: Optional[_TrainingResult], tr2: Optional[_TrainingResult]
+    ) -> bool:
+        if not tr1 and not tr2:
+            return True
+        if not tr1 or not tr2:
+            return False
+        return (
+            tr1.metrics == tr2.metrics
+            and tr1.checkpoint.path == tr2.checkpoint.path
+            and tr1.checkpoint.filesystem == tr2.checkpoint.filesystem
+        )
+
+    def _checkpoint_to_report_indices_equal(
+        cm1: CheckpointManager, cm2: CheckpointManager
+    ) -> bool:
+        # Do this because Checkpoint and Filesystem are not hashable.
+        cm1_dict = {
+            checkpoint.path: report_index
+            for checkpoint, report_index in cm1._checkpoint_to_report_index.items()
+        }
+        cm2_dict = {
+            checkpoint.path: report_index
+            for checkpoint, report_index in cm2._checkpoint_to_report_index.items()
+        }
+        return cm1_dict == cm2_dict
+
+    if cm1._checkpoint_config != cm2._checkpoint_config:
+        return False
+    if not _training_results_equal(
+        cm1.latest_checkpoint_result, cm2.latest_checkpoint_result
+    ):
+        return False
+    if not _training_results_equal(
+        cm1.best_checkpoint_result, cm2.best_checkpoint_result
+    ):
+        return False
+    if len(cm1.best_checkpoint_results) != len(cm2.best_checkpoint_results):
+        return False
+    for tr1, tr2 in zip(cm1.best_checkpoint_results, cm2.best_checkpoint_results):
+        if not _training_results_equal(tr1, tr2):
+            return False
+    if cm1._current_report_index != cm2._current_report_index:
+        return False
+    if not _checkpoint_to_report_indices_equal(cm1, cm2):
+        return False
+    return True
+
+
+@pytest.mark.parametrize(
+    "checkpoint_config",
+    [
+        CheckpointConfig(),
+        CheckpointConfig(
+            num_to_keep=1,
+            checkpoint_score_attribute="score",
+            checkpoint_score_order="max",
+        ),
+    ],
+)
+@pytest.mark.asyncio
+async def test_save_load_state_equivalence(
+    monkeypatch, tmp_path, checkpoint_config: CheckpointConfig
+):
+    # Use async here because register_checkpoint creates an async task
+
+    # Mock the delete function as we don't want report checkpoints to be deleted.
+    monkeypatch.setattr(
+        ray.train.v2._internal.execution.checkpoint.checkpoint_manager,
+        "delete_fs_path",
+        lambda *args, **kwargs: None,
+    )
+    exp_name = f"checkpoint_manager_test-{uuid.uuid4().hex}"
+
+    storage_context = StorageContext(
+        storage_path=tmp_path,
+        experiment_dir_name=exp_name,
+    )
+    checkpoint_manager = CheckpointManager(
+        storage_context=storage_context,
+        checkpoint_config=checkpoint_config,
+    )
+    training_reports = create_dummy_training_reports(
+        num_results=2, storage_context=storage_context
+    ) + create_dummy_training_reports(
+        num_results=1,
+        storage_context=storage_context,
+        include_validation=True,
+        starting_checkpoint_index=2,
+    )
+
+    # Register the training results into checkpoint manager
+    for i, tr in enumerate(training_reports):
+        checkpoint_manager.register_checkpoint(tr)
+        assert checkpoint_manager._current_report_index == i + 1
+        loaded_checkpoint_manager = CheckpointManager(
+            storage_context=storage_context,
+            checkpoint_config=checkpoint_config,
+        )
+        assert _checkpoint_managers_equal(checkpoint_manager, loaded_checkpoint_manager)
+
+
+@pytest.mark.parametrize(
+    "json_state,match",
+    [
+        (
+            '{"dummy": "1", "dummy_dict": {"key": "value"}}',
+            "You are loading a checkpoint manager snapshot saved with an unknown Ray version but",
+        ),
+        ('{"ray_version": "2.0.0", "dummy": "1", "dummy_dict": {"key": "value"', None),
+        (
+            '{"ray_version": "2.0.0", "dummy": "1", "dummy_dict": {"key": "value"}}',
+            "You are loading a checkpoint manager snapshot saved with Ray version 2.0.0 but",
+        ),
+    ],
+)
+def test_load_state_error(tmp_path, json_state, match):
+
+    storage_context = StorageContext(
+        storage_path=tmp_path,
+        experiment_dir_name="load_state_error_experiment",
+    )
+    checkpoint_manager = CheckpointManager(
+        storage_context=storage_context,
+        checkpoint_config=CheckpointConfig(),
+    )
+    with pytest.raises(
+        CheckpointManagerInitializationError,
+        match=match,
+    ):
+        checkpoint_manager._load_state(json_state)
+
+
+@pytest.mark.asyncio
+async def test_before_init_train_context(tmp_path):
+
+    storage_context = StorageContext(
+        storage_path=tmp_path,
+        experiment_dir_name="my_experiment_name",
+    )
+    checkpoint_manager = CheckpointManager(
+        storage_context=storage_context,
+        checkpoint_config=CheckpointConfig(),
+    )
+    workers = [create_autospec(Worker, instance=True) for _ in range(4)]
+
+    # Assert without a checkpoint.
+    assert checkpoint_manager.before_init_train_context(workers) == {
+        "checkpoint": [None] * 4,
+        "current_report_index": [0] * 4,
+    }
+
+    # Assert with a checkpoint
+    latest_checkpoint_report = create_dummy_training_reports(1, storage_context)[0]
+    checkpoint_manager.register_checkpoint(latest_checkpoint_report)
+    assert checkpoint_manager.before_init_train_context(workers) == {
+        "checkpoint": [latest_checkpoint_report.checkpoint] * 4,
+        "current_report_index": [1] * 4,
+    }
+
+
+@pytest.mark.asyncio
+async def test_pending_checkpoint_management(tmp_path):
+    storage_context = StorageContext(
+        storage_path=tmp_path,
+        experiment_dir_name="pending_checkpoint_management_experiment",
+    )
+    checkpoint_config = CheckpointConfig(
+        num_to_keep=1,
+        checkpoint_score_attribute="score",
+        checkpoint_score_order="max",
+    )
+    checkpoint_manager = CheckpointManager(
+        storage_context=storage_context,
+        checkpoint_config=checkpoint_config,
+    )
+    (
+        low_initial_high_final_training_report,
+        high_initial_low_final_training_report,
+        final_training_report,
+    ) = create_dummy_training_reports(
+        num_results=3, storage_context=storage_context, include_validation=True
+    )
+    final_training_report.validation = False
+    scoreless_training_report = create_dummy_training_reports(
+        num_results=1,
+        storage_context=storage_context,
+        include_metrics=False,
+        starting_checkpoint_index=3,
+    )[0]
+
+    # Register pending/final/unknown checkpoints and verify their storage
+    checkpoint_manager.register_checkpoint(low_initial_high_final_training_report)
+    checkpoint_manager.register_checkpoint(final_training_report)
+    checkpoint_manager.register_checkpoint(scoreless_training_report)
+    checkpoint_manager.register_checkpoint(high_initial_low_final_training_report)
+    assert [tr.checkpoint for tr in checkpoint_manager._checkpoint_results] == [
+        low_initial_high_final_training_report.checkpoint,  # keep pending
+        high_initial_low_final_training_report.checkpoint,  # keep pending/latest
+        final_training_report.checkpoint,  # keep highest final score so far
+    ]
+
+    # Assert checkpoint state after all tasks are done
+    checkpoint_manager.update_checkpoints_with_validation_result(
+        {
+            low_initial_high_final_training_report.checkpoint: (
+                {"score": 200},
+                ReportedCheckpointStatus.VALIDATED,
+            ),
+            high_initial_low_final_training_report.checkpoint: (
+                {"score": 100},
+                ReportedCheckpointStatus.VALIDATED,
+            ),
+        }
+    )
+    assert [tr.checkpoint for tr in checkpoint_manager._checkpoint_results] == [
+        high_initial_low_final_training_report.checkpoint,  # keep latest checkpoint
+        low_initial_high_final_training_report.checkpoint,  # keep highest score checkpoint
+    ]
+
+
+@pytest.mark.asyncio
+async def test_pending_checkpoint_management_break_ties_by_report_index(tmp_path):
+    storage_context = StorageContext(
+        storage_path=tmp_path,
+        experiment_dir_name="pending_checkpoint_management_break_ties_by_report_index_experiment",
+    )
+    checkpoint_manager = CheckpointManager(
+        storage_context=storage_context,
+        checkpoint_config=CheckpointConfig(),
+    )
+    training_reports = create_dummy_training_reports(
+        num_results=2, storage_context=storage_context, include_validation=True
+    )
+    checkpoint_manager.register_checkpoint(training_reports[0])
+    checkpoint_manager.register_checkpoint(training_reports[1])
+    assert [tr.checkpoint for tr in checkpoint_manager._checkpoint_results] == [
+        training_reports[0].checkpoint,
+        training_reports[1].checkpoint,
+    ]
+    checkpoint_manager.update_checkpoints_with_validation_result(
+        {
+            training_reports[1].checkpoint: ({}, ReportedCheckpointStatus.VALIDATED),
+        }
+    )
+    assert [tr.checkpoint for tr in checkpoint_manager._checkpoint_results] == [
+        training_reports[0].checkpoint,
+        training_reports[1].checkpoint,
+    ]
+    checkpoint_manager.update_checkpoints_with_validation_result(
+        {
+            training_reports[0].checkpoint: ({}, ReportedCheckpointStatus.VALIDATED),
+        }
+    )
+    assert [tr.checkpoint for tr in checkpoint_manager._checkpoint_results] == [
+        training_reports[0].checkpoint,
+        training_reports[1].checkpoint,
+    ]
+
+
+@pytest.mark.asyncio
+async def test_pending_checkpoint_management_finalized_checkpoint(tmp_path):
+    storage_context = StorageContext(
+        storage_path=tmp_path,
+        experiment_dir_name="pending_checkpoint_management_experiment",
+    )
+    checkpoint_manager = CheckpointManager(
+        storage_context=storage_context,
+        checkpoint_config=CheckpointConfig(
+            checkpoint_score_attribute="score",
+            checkpoint_score_order="max",
+        ),
+    )
+    training_reports = create_dummy_training_reports(
+        num_results=2, storage_context=storage_context
+    )
+    checkpoint_manager.register_checkpoint(training_reports[0])
+    checkpoint_manager.register_checkpoint(training_reports[1])
+    assert [tr.checkpoint for tr in checkpoint_manager._checkpoint_results] == [
+        training_reports[0].checkpoint,
+        training_reports[1].checkpoint,
+    ]
+    checkpoint_manager.update_checkpoints_with_validation_result(
+        {
+            training_reports[0].checkpoint: (
+                {"score": 100},
+                ReportedCheckpointStatus.VALIDATED,
+            ),
+        }
+    )
+    assert [tr.checkpoint for tr in checkpoint_manager._checkpoint_results] == [
+        training_reports[0].checkpoint,
+        training_reports[1].checkpoint,
+    ]
+
+
+def test_update_checkpoints_with_metrics_not_in_checkpoint_results(tmp_path):
+    storage_context = StorageContext(
+        storage_path=tmp_path,
+        experiment_dir_name="update_checkpoints_with_metrics_error_experiment",
+    )
+    checkpoint_manager = CheckpointManager(
+        storage_context=storage_context,
+        checkpoint_config=CheckpointConfig(),
+    )
+    training_reports = create_dummy_training_reports(
+        num_results=1, storage_context=storage_context
+    )
+    checkpoint_manager._pending_training_results[training_reports[0].checkpoint] = (
+        _TrainingResult(training_reports[0].checkpoint, training_reports[0].metrics),
+        training_reports[0].validation,
+    )
+    with pytest.raises(ValueError):
+        checkpoint_manager.update_checkpoints_with_validation_result(
+            {
+                training_reports[0].checkpoint: (
+                    {"score": 100},
+                    ReportedCheckpointStatus.VALIDATED,
+                )
+            }
+        )
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))