chore: import upstream snapshot with attribution

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

@@ -0,0 +1,348 @@
+import asyncio
+import unittest.mock
+from unittest.mock import create_autospec
+
+import pytest
+
+import ray
+from ray.train._checkpoint import Checkpoint
+from ray.train._internal.session import _TrainingResult
+from ray.train.v2._internal.execution.checkpoint import validation_manager
+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.training_report import (
+    _TrainingReport,
+)
+from ray.train.v2._internal.execution.worker_group.worker import Worker
+from ray.train.v2.api.reported_checkpoint import ReportedCheckpointStatus
+from ray.train.v2.api.validation_config import ValidationConfig, ValidationTaskConfig
+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()
+
+
+@unittest.mock.patch.object(ray, "wait", autospec=True)
+def test_before_controller_shutdown(mock_wait, monkeypatch):
+    monkeypatch.setattr(validation_manager, "VALIDATION_TASK_POLL_INTERVAL_S", 0)
+
+    # Create ValidationManager with mocked objects
+    checkpoint_manager = create_autospec(CheckpointManager, instance=True)
+    checkpoint1 = create_autospec(Checkpoint, instance=True)
+    checkpoint2 = create_autospec(Checkpoint, instance=True)
+    checkpoint3 = create_autospec(Checkpoint, instance=True)
+    task1 = create_autospec(ray.ObjectRef, instance=True)
+    task2 = create_autospec(ray.ObjectRef, instance=True)
+    task3 = create_autospec(ray.ObjectRef, instance=True)
+    vm = validation_manager.ValidationManager(
+        checkpoint_manager=checkpoint_manager,
+        validation_config=ValidationConfig(fn=lambda x: None),
+    )
+    vm._pending_validations = {
+        task1: validation_manager._PendingValidation(
+            checkpoint=checkpoint1, start_time=0.0, timeout_s=None
+        ),
+        task2: validation_manager._PendingValidation(
+            checkpoint=checkpoint2, start_time=0.0, timeout_s=None
+        ),
+        task3: validation_manager._PendingValidation(
+            checkpoint=checkpoint3, start_time=0.0, timeout_s=None
+        ),
+    }
+    mock_wait.side_effect = [([], [task1, task2, task3]), ([task1, task2, task3], [])]
+    monkeypatch.setattr(ray, "get", lambda x: {"score": 1})
+
+    # Call before_controller_shutdown
+    asyncio.run(vm.before_controller_shutdown())
+    assert mock_wait.call_count == 2
+    assert checkpoint_manager.update_checkpoints_with_validation_result.mock_calls == [
+        unittest.mock.call(
+            {checkpoint1: ({"score": 1}, ReportedCheckpointStatus.VALIDATED)}
+        ),
+        unittest.mock.call(
+            {
+                checkpoint2: ({"score": 1}, ReportedCheckpointStatus.VALIDATED),
+                checkpoint3: ({"score": 1}, ReportedCheckpointStatus.VALIDATED),
+            }
+        ),
+    ]
+
+
+def test_before_init_train_context():
+    checkpoint_manager = create_autospec(CheckpointManager, instance=True)
+    vm = validation_manager.ValidationManager(
+        checkpoint_manager=checkpoint_manager,
+        validation_config=ValidationConfig(fn=lambda x: None),
+    )
+    workers = [create_autospec(Worker, instance=True) for _ in range(4)]
+    assert vm.before_init_train_context(workers) == {
+        "has_validation_fn": [True] * 4,
+    }
+
+
+def test_checkpoint_validation_management_reordering(tmp_path):
+    checkpoint_manager = create_autospec(CheckpointManager, instance=True)
+
+    def validation_fn(checkpoint, score):
+        return {"score": score}
+
+    vm = validation_manager.ValidationManager(
+        checkpoint_manager=checkpoint_manager,
+        validation_config=ValidationConfig(
+            fn=validation_fn,
+            task_config=ValidationTaskConfig(fn_kwargs={"score": 100}),
+        ),
+    )
+    (
+        low_initial_high_final_training_result,
+        high_initial_low_final_training_result,
+    ) = create_dummy_training_reports(
+        num_results=2,
+        storage_context=StorageContext(
+            storage_path=tmp_path,
+            experiment_dir_name="checkpoint_validation_management_reordering_experiment",
+        ),
+    )
+
+    # Enqueue validation tasks
+    vm.after_report(
+        training_report=_TrainingReport(
+            metrics=low_initial_high_final_training_result.metrics,
+            checkpoint=low_initial_high_final_training_result.checkpoint,
+            validation=ValidationTaskConfig(fn_kwargs={"score": 200}),
+        ),
+        metrics={},
+    )
+    vm.after_report(
+        training_report=_TrainingReport(
+            metrics=high_initial_low_final_training_result.metrics,
+            checkpoint=high_initial_low_final_training_result.checkpoint,
+            validation=True,
+        ),
+        metrics={},
+    )
+
+    # Assert ValidationManager state after each poll
+    assert vm._poll_validations() == 0
+    assert vm._kick_off_validations() == 1
+    ray.wait(
+        list(vm._pending_validations.keys()),
+        num_returns=1,
+    )
+    assert vm._poll_validations() == 0
+    assert vm._kick_off_validations() == 1
+    checkpoint_manager.update_checkpoints_with_validation_result.assert_called_once_with(
+        {
+            low_initial_high_final_training_result.checkpoint: (
+                {"score": 200},
+                ReportedCheckpointStatus.VALIDATED,
+            )
+        }
+    )
+    ray.wait(
+        list(vm._pending_validations.keys()),
+        num_returns=1,
+    )
+    assert vm._poll_validations() == 0
+    assert vm._kick_off_validations() == 0
+    checkpoint_manager.update_checkpoints_with_validation_result.assert_called_with(
+        {
+            high_initial_low_final_training_result.checkpoint: (
+                {"score": 100},
+                ReportedCheckpointStatus.VALIDATED,
+            )
+        }
+    )
+
+
+def test_checkpoint_validation_management_failure(tmp_path):
+    checkpoint_manager = create_autospec(CheckpointManager, instance=True)
+
+    def failing_validation_fn(checkpoint):
+        return "invalid_return_type"
+
+    vm = validation_manager.ValidationManager(
+        checkpoint_manager=checkpoint_manager,
+        validation_config=ValidationConfig(fn=failing_validation_fn),
+    )
+    failing_training_result = create_dummy_training_reports(
+        num_results=1,
+        storage_context=StorageContext(
+            storage_path=tmp_path,
+            experiment_dir_name="checkpoint_validation_management_failure_experiment",
+        ),
+    )[0]
+
+    vm.after_report(
+        training_report=_TrainingReport(
+            metrics=failing_training_result.metrics,
+            checkpoint=failing_training_result.checkpoint,
+            validation=ValidationTaskConfig(),
+        ),
+        metrics={},
+    )
+    assert vm._poll_validations() == 0
+    assert vm._kick_off_validations() == 1
+    ray.wait(
+        list(vm._pending_validations.keys()),
+        num_returns=1,
+    )
+    assert vm._poll_validations() == 0
+    assert vm._kick_off_validations() == 0
+    checkpoint_manager.update_checkpoints_with_validation_result.assert_called_once_with(
+        {
+            failing_training_result.checkpoint: (
+                {},
+                ReportedCheckpointStatus.VALIDATION_FAILED,
+            )
+        }
+    )
+
+
+def test_checkpoint_validation_management_success_after_retry(tmp_path):
+    @ray.remote
+    class Counter:
+        def __init__(self):
+            self.value = 0
+
+        def increment(self):
+            self.value += 1
+            return self.value
+
+    counter = Counter.remote()
+
+    def one_time_failing_validation_fn(checkpoint):
+        print("one_time_failing_validation_fn called")
+        if ray.get(counter.increment.remote()) < 2:
+            raise ValueError("Fail on first attempt")
+        return {"score": 100}
+
+    checkpoint_manager = create_autospec(CheckpointManager, instance=True)
+    vm = validation_manager.ValidationManager(
+        checkpoint_manager=checkpoint_manager,
+        validation_config=ValidationConfig(
+            fn=one_time_failing_validation_fn,
+            ray_remote_kwargs={"max_retries": 1, "retry_exceptions": [ValueError]},
+        ),
+    )
+    training_result = create_dummy_training_reports(
+        num_results=1,
+        storage_context=StorageContext(
+            storage_path=tmp_path,
+            experiment_dir_name="checkpoint_validation_management_success_after_retry_experiment",
+        ),
+    )[0]
+
+    vm.after_report(
+        training_report=_TrainingReport(
+            metrics=training_result.metrics,
+            checkpoint=training_result.checkpoint,
+            validation=True,
+        ),
+        metrics={},
+    )
+    assert vm._poll_validations() == 0
+    assert vm._kick_off_validations() == 1
+    ray.wait(
+        list(vm._pending_validations.keys()),
+        num_returns=1,
+        timeout=100,
+    )
+    assert vm._poll_validations() == 0
+    assert vm._kick_off_validations() == 0
+    checkpoint_manager.update_checkpoints_with_validation_result.assert_called_once_with(
+        {
+            training_result.checkpoint: (
+                {"score": 100},
+                ReportedCheckpointStatus.VALIDATED,
+            )
+        }
+    )
+
+
+def test_checkpoint_validation_management_resume(tmp_path):
+    training_reports = create_dummy_training_reports(
+        num_results=3,
+        storage_context=StorageContext(
+            storage_path=tmp_path,
+            experiment_dir_name="checkpoint_validation_management_resume_experiment",
+        ),
+    )
+    checkpoint_manager = create_autospec(CheckpointManager, instance=True)
+    checkpoint_manager.get_pending_training_results.return_value = {
+        training_reports[0].checkpoint: (
+            _TrainingResult(
+                checkpoint=training_reports[0].checkpoint,
+                metrics=training_reports[0].metrics,
+            ),
+            True,
+        ),
+        training_reports[1].checkpoint: (
+            _TrainingResult(
+                checkpoint=training_reports[1].checkpoint,
+                metrics=training_reports[1].metrics,
+            ),
+            False,
+        ),
+        training_reports[2].checkpoint: (
+            _TrainingResult(
+                checkpoint=training_reports[2].checkpoint,
+                metrics=training_reports[2].metrics,
+            ),
+            ValidationTaskConfig(fn_kwargs={"score": 2}),
+        ),
+    }
+
+    def validation_fn(checkpoint, score):
+        return {"score": score}
+
+    vm = validation_manager.ValidationManager(
+        checkpoint_manager=checkpoint_manager,
+        validation_config=ValidationConfig(
+            fn=validation_fn,
+            task_config=ValidationTaskConfig(fn_kwargs={"score": 1}),
+        ),
+    )
+
+    assert vm._poll_validations() == 0
+    assert vm._kick_off_validations() == 1
+    ray.wait(
+        list(vm._pending_validations.keys()),
+        num_returns=1,
+    )
+    assert vm._poll_validations() == 0
+    assert vm._kick_off_validations() == 1
+    checkpoint_manager.update_checkpoints_with_validation_result.assert_called_once_with(
+        {
+            training_reports[0].checkpoint: (
+                {"score": 1},
+                ReportedCheckpointStatus.VALIDATED,
+            )
+        }
+    )
+    ray.wait(
+        list(vm._pending_validations.keys()),
+        num_returns=1,
+    )
+    assert vm._poll_validations() == 0
+    assert vm._kick_off_validations() == 0
+    checkpoint_manager.update_checkpoints_with_validation_result.assert_called_with(
+        {
+            training_reports[2].checkpoint: (
+                {"score": 2},
+                ReportedCheckpointStatus.VALIDATED,
+            )
+        }
+    )
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))