chore: import upstream snapshot with attribution

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

@@ -0,0 +1,360 @@
+import stat
+import uuid
+from pathlib import Path
+from urllib.parse import urlparse, urlunparse
+
+import pytest
+
+import ray
+import ray.train.collective
+from ray import train
+from ray.train import Checkpoint, CheckpointConfig, RunConfig, ScalingConfig
+from ray.train.tests.util import create_dict_checkpoint, load_dict_checkpoint
+from ray.train.torch import TorchTrainer
+from ray.train.v2._internal.constants import CHECKPOINT_MANAGER_SNAPSHOT_FILENAME
+from ray.train.v2._internal.execution.storage import StorageContext
+from ray.train.v2.api.data_parallel_trainer import DataParallelTrainer
+from ray.train.v2.api.exceptions import WorkerGroupError
+from ray.train.v2.api.result import Result
+
+
+def uri_join(base_uri: str, *paths: str) -> str:
+    """
+    Join a base URI (local or remote) with one or more subpaths.
+    Preserves query parameters and scheme.
+    """
+    parsed = urlparse(base_uri)
+    new_path = "/".join([p.strip("/") for p in [parsed.path, *paths] if p])
+
+    # If it's a local path (no scheme), ensure we preserve the leading /
+    if not parsed.scheme and not new_path.startswith("/"):
+        new_path = "/" + new_path
+
+    return urlunparse(
+        (
+            parsed.scheme,
+            parsed.netloc,
+            new_path,
+            parsed.params,
+            parsed.query,
+            parsed.fragment,
+        )
+    )
+
+
+def build_dummy_trainer(
+    exp_name: str,
+    storage_path: str,
+    num_iterations: int,
+    num_checkpoints: int,
+    train_loop_config: dict,
+):
+    """Build a dummy TorchTrainer for testing purposes."""
+
+    def worker_loop(_config):
+        for i in range(num_iterations):
+            # Do some random reports in between checkpoints.
+            train.report({"metric_a": -100, "metric_b": -100})
+
+            if ray.train.get_context().get_world_rank() == 0:
+                with create_dict_checkpoint({"iter": i}) as checkpoint:
+                    train.report(
+                        metrics={"metric_a": i, "metric_b": -i},
+                        checkpoint=checkpoint,
+                    )
+            else:
+                train.report(metrics={"metric_a": i, "metric_b": -i})
+
+        # Ensure that all checkpoints are saved before the RuntimeError
+        ray.train.collective.barrier()
+        raise RuntimeError()
+
+    trainer = TorchTrainer(
+        train_loop_per_worker=worker_loop,
+        train_loop_config=train_loop_config,
+        scaling_config=ScalingConfig(num_workers=2, use_gpu=False),
+        run_config=RunConfig(
+            name=exp_name,
+            storage_path=storage_path,
+            checkpoint_config=CheckpointConfig(
+                num_to_keep=num_checkpoints,
+                checkpoint_score_attribute="metric_a",
+                checkpoint_score_order="max",
+            ),
+        ),
+    )
+    return trainer
+
+
+def test_result_repr():
+    """Test that the Result __repr__ function can return a string."""
+    res = Result(
+        metrics={"iter": 0, "metric": 1.0},
+        checkpoint=Checkpoint("/bucket/path/ckpt0"),
+        error=None,
+        path="/bucket/path",
+    )
+    assert isinstance(repr(res), str)
+    assert "Checkpoint(filesystem=local, path=/bucket/path/ckpt0)" in repr(res)
+    assert "metrics={'iter': 0, 'metric': 1.0}" in repr(res)
+
+
+def test_get_best_checkpoint():
+    res = Result(
+        metrics={},
+        checkpoint=None,
+        error=None,
+        path="/bucket/path",
+        best_checkpoints=[
+            (Checkpoint("/bucket/path/ckpt0"), {"iter": 0, "metric": 1.0}),
+            (Checkpoint("/bucket/path/ckpt1"), {"iter": 1, "metric": 2.0}),
+            (Checkpoint("/bucket/path/ckpt2"), {"iter": 2, "metric": 3.0}),
+            (Checkpoint("/bucket/path/ckpt3"), {"iter": 3, "metric": 4.0}),
+        ],
+    )
+    assert (
+        res.get_best_checkpoint(metric="metric", mode="max").path
+        == "/bucket/path/ckpt3"
+    )
+    assert (
+        res.get_best_checkpoint(metric="metric", mode="min").path
+        == "/bucket/path/ckpt0"
+    )
+
+
+def test_get_best_checkpoint_nested_metrics():
+    """Test that get_best_checkpoint works with nested metric dictionaries."""
+    # Test with nested metric structure
+    res = Result(
+        metrics={},
+        checkpoint=None,
+        error=None,
+        path="/bucket/path",
+        best_checkpoints=[
+            (
+                Checkpoint("/bucket/path/ckpt0"),
+                {
+                    "iter": 0,
+                    "env_runners": {"episode_return_mean": 100.0, "num_episodes": 10},
+                },
+            ),
+            (
+                Checkpoint("/bucket/path/ckpt1"),
+                {
+                    "iter": 1,
+                    "env_runners": {"episode_return_mean": 200.0, "num_episodes": 10},
+                },
+            ),
+            (
+                Checkpoint("/bucket/path/ckpt2"),
+                {
+                    "iter": 2,
+                    "env_runners": {"episode_return_mean": 300.0, "num_episodes": 10},
+                },
+            ),
+            (
+                Checkpoint("/bucket/path/ckpt3"),
+                {
+                    "iter": 3,
+                    "env_runners": {"episode_return_mean": 400.0, "num_episodes": 10},
+                },
+            ),
+        ],
+    )
+
+    # Test max mode with nested metric
+    assert (
+        res.get_best_checkpoint(
+            metric="env_runners/episode_return_mean", mode="max"
+        ).path
+        == "/bucket/path/ckpt3"
+    )
+
+    # Test min mode with nested metric
+    assert (
+        res.get_best_checkpoint(
+            metric="env_runners/episode_return_mean", mode="min"
+        ).path
+        == "/bucket/path/ckpt0"
+    )
+
+    # Test that flat keys still work (backwards compatibility)
+    res_flat = Result(
+        metrics={},
+        checkpoint=None,
+        error=None,
+        path="/bucket/path",
+        best_checkpoints=[
+            (
+                Checkpoint("/bucket/path/ckpt0"),
+                {"iter": 0, "env_runners/episode_return_mean": 100.0},
+            ),
+            (
+                Checkpoint("/bucket/path/ckpt1"),
+                {"iter": 1, "env_runners/episode_return_mean": 200.0},
+            ),
+        ],
+    )
+    assert (
+        res_flat.get_best_checkpoint(
+            metric="env_runners/episode_return_mean", mode="max"
+        ).path
+        == "/bucket/path/ckpt1"
+    )
+
+
+@pytest.mark.parametrize("path_type", ["str", "PathLike"])
+@pytest.mark.parametrize("pass_storage_filesystem", [True, False])
+@pytest.mark.parametrize("trailing_slash", [False, True])
+def test_result_restore(
+    ray_start_4_cpus,
+    tmp_path,
+    path_type,
+    pass_storage_filesystem,
+    trailing_slash,
+):
+    """Test Result.from_path functionality similar to v1 test_result_restore."""
+
+    num_iterations = 3
+    num_checkpoints = 2
+
+    storage_path = str(tmp_path)
+
+    # Add UUID to ensure test isolation when sharing module-scoped S3 mock
+    exp_name = f"test_result_restore_v2-{uuid.uuid4().hex[:8]}"
+
+    trainer = build_dummy_trainer(
+        exp_name,
+        storage_path,
+        num_iterations,
+        num_checkpoints,
+        train_loop_config={"a": 1, "b": 2},
+    )
+    with pytest.raises(WorkerGroupError):
+        trainer.fit()
+
+    if pass_storage_filesystem:
+        storage_context = StorageContext(
+            storage_path=storage_path,
+            experiment_dir_name=exp_name,
+        )
+
+        trial_dir = storage_context.experiment_fs_path
+        file_system = storage_context.storage_filesystem
+    else:
+        trial_dir = uri_join(storage_path, exp_name)
+        file_system = None
+
+    # Add trailing slash if parameterized to test that case
+    if trailing_slash:
+        trial_dir = trial_dir + "/"
+
+    # For PathLike test, only use Path() for local paths, not URIs
+    if path_type == "PathLike":
+        trial_dir_arg = Path(trial_dir)
+    else:
+        trial_dir_arg = trial_dir
+
+    result = Result.from_path(
+        trial_dir_arg,
+        storage_filesystem=file_system,
+    )
+
+    assert result.checkpoint
+    assert len(result.best_checkpoints) == num_checkpoints
+
+    """
+    Top-2 checkpoints with metrics:
+
+                        | iter   | metric_a    metric_b
+    checkpoint_000002        2            2          -2
+    checkpoint_000001        1            1          -1
+    """
+    # Check if the checkpoints bounded with correct metrics
+    best_ckpt_a = result.get_best_checkpoint(metric="metric_a", mode="max")
+    assert load_dict_checkpoint(best_ckpt_a)["iter"] == num_iterations - 1
+
+    best_ckpt_b = result.get_best_checkpoint(metric="metric_b", mode="max")
+    assert load_dict_checkpoint(best_ckpt_b)["iter"] == num_iterations - num_checkpoints
+
+    with pytest.raises(RuntimeError, match="Invalid metric name.*"):
+        result.get_best_checkpoint(metric="invalid_metric", mode="max")
+
+
+def test_result_from_path_read_only_storage(
+    ray_start_4_cpus,
+    tmp_path,
+):
+    """Reproduces https://github.com/ray-project/ray/issues/64305.
+
+    If checkpoints live on write-once / read-only storage, check that
+    ``Result.from_path`` / ``get_best_checkpoint`` are read-only operations.
+    """
+
+    def train_fn():
+        with create_dict_checkpoint({"iter": 1}) as ckpt:
+            ray.train.report({"epoch": 1}, ckpt)
+        with create_dict_checkpoint({"iter": 2}) as ckpt:
+            ray.train.report({"epoch": 2}, ckpt)
+
+    trainer = DataParallelTrainer(
+        train_fn,
+        run_config=RunConfig(
+            name="test_read_only_storage",
+            storage_path=str(tmp_path),
+        ),
+    )
+    trainer.fit()
+
+    # Strip write permissions from the experiment directory, so that any new write fails.
+    no_write_mask = ~(stat.S_IWUSR | stat.S_IWGRP | stat.S_IWOTH)
+    experiment_dir = Path(tmp_path) / "test_read_only_storage"
+    assert experiment_dir.exists()
+    original_modes = {
+        path: stat.S_IMODE(path.stat().st_mode)
+        for path in [experiment_dir, *experiment_dir.rglob("*")]
+    }
+    for path in original_modes:
+        path.chmod(original_modes[path] & no_write_mask)
+
+    try:
+        # Reading the best checkpoint must not require write access.
+        result = Result.from_path(str(experiment_dir))
+        assert result.checkpoint
+        assert len(result.best_checkpoints) == 2
+
+        best_ckpt = result.get_best_checkpoint(metric="epoch", mode="max")
+        assert load_dict_checkpoint(best_ckpt)["iter"] == 2
+    finally:
+        for path, mode in original_modes.items():
+            path.chmod(mode)
+
+
+def test_result_from_path_validation(
+    ray_start_4_cpus,
+    tmp_path,
+):
+    """Test that Result.from_path raises RuntimeError when folder or snapshot file doesn't exist."""
+
+    nonexistent_folder = str(tmp_path / "nonexistent_experiment")
+    existing_folder = str(tmp_path / "existing_experiment")
+
+    # Test 1: Folder doesn't exist
+    with pytest.raises(RuntimeError, match="Experiment folder .* doesn't exist."):
+        Result.from_path(nonexistent_folder)
+
+    # Test 2: Folder exists but snapshot file doesn't exist
+    Path(existing_folder).mkdir(parents=True, exist_ok=True)
+
+    with pytest.raises(
+        RuntimeError,
+        match=f"Failed to restore the Result object: {CHECKPOINT_MANAGER_SNAPSHOT_FILENAME} doesn't exist in the experiment folder. Make sure that this is an output directory created "
+        "by a Ray Train run.",
+    ):
+        Result.from_path(existing_folder)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))