[train][V2] Implement Result::from_path in v2 (ray-project#58216)

## Description

In this function, `Result::from_path` is implemented in ray train v2,
which reconstructs a `Result` object from the checkpoints. This
implementation leverages `CheckpointManager` and refers to
https://github.com/ray-project/ray/blob/master/python/ray/train/v2/_internal/execution/controller/controller.py#L512-L540

---------

Signed-off-by: xgui <xgui@anyscale.com>
Signed-off-by: Justin Yu <justinvyu@anyscale.com>
Co-authored-by: Justin Yu <justinvyu@anyscale.com>

Author: 2 people

python/ray/train/v2/BUILD.bazel

@@ -343,7 +343,7 @@ py_test(
 
 py_test(
     name = "test_result",
-    size = "small",
+    size = "medium",
     srcs = ["tests/test_result.py"],
     env = {"RAY_TRAIN_V2_ENABLED": "1"},
     tags = [

python/ray/train/v2/_internal/execution/callback.py

@@ -4,7 +4,6 @@
 from ray.train.v2._internal.execution.training_report import _TrainingReport
 from ray.train.v2.api.callback import RayTrainCallback
 from ray.train.v2.api.config import ScalingConfig
-from ray.train.v2.api.result import Result
 from ray.util.annotations import DeveloperAPI
 
 if TYPE_CHECKING:
@@ -20,6 +19,7 @@
         WorkerGroupContext,
         WorkerGroupPollStatus,
     )
+    from ray.train.v2.api.result import Result
 
 
 @DeveloperAPI
@@ -128,7 +128,7 @@ def before_controller_execute_resize_decision(
         """Called before the controller executes a resize decision."""
         pass
 
-    def after_controller_finish(self, result: Result):
+    def after_controller_finish(self, result: "Result"):
         """Called after the training run completes, providing access to the final result.
 
         Args:

python/ray/train/v2/_internal/execution/checkpoint/checkpoint_manager.py

@@ -90,6 +90,8 @@ def __init__(
         self._condition = asyncio.Condition()
         super().__init__(checkpoint_config)
         # If the snapshot is found, the checkpoint manager will restore its state.
+        # TODO(xgui): CheckpointManager is used to save or restore the checkpoint manager state.
+        # We should sanity check if we should see old state in the storage folder.
         self._maybe_load_state_from_storage()
 
     def register_checkpoint(

python/ray/train/v2/api/result.py

@@ -3,10 +3,21 @@
 from dataclasses import dataclass
 from typing import Any, Dict, List, Optional, Tuple, Union
 
+import pandas as pd
 import pyarrow
 
 import ray
 from ray.air.result import Result as ResultV1
+from ray.train import Checkpoint, CheckpointConfig
+from ray.train.v2._internal.constants import CHECKPOINT_MANAGER_SNAPSHOT_FILENAME
+from ray.train.v2._internal.execution.checkpoint.checkpoint_manager import (
+    CheckpointManager,
+)
+from ray.train.v2._internal.execution.storage import (
+    StorageContext,
+    _exists_at_fs_path,
+    get_fs_and_path,
+)
 from ray.train.v2.api.exceptions import TrainingFailedError
 from ray.util.annotations import Deprecated, PublicAPI
 
@@ -15,11 +26,9 @@
 
 @dataclass
 class Result(ResultV1):
-    checkpoint: Optional["ray.train.Checkpoint"]
+    checkpoint: Optional[Checkpoint]
     error: Optional[TrainingFailedError]
-    best_checkpoints: Optional[
-        List[Tuple["ray.train.Checkpoint", Dict[str, Any]]]
-    ] = None
+    best_checkpoints: Optional[List[Tuple[Checkpoint, Dict[str, Any]]]] = None
 
     @PublicAPI(stability="alpha")
     def get_best_checkpoint(
@@ -33,7 +42,90 @@ def from_path(
         path: Union[str, os.PathLike],
         storage_filesystem: Optional[pyarrow.fs.FileSystem] = None,
     ) -> "Result":
-        raise NotImplementedError("`Result.from_path` is not implemented yet.")
+        """Restore a training result from a previously saved training run path.
+
+        Args:
+            path: Path to the run output directory
+            storage_filesystem: Optional filesystem to use for accessing the path
+
+        Returns:
+            Result object with restored checkpoints and metrics
+        """
+        fs, fs_path = get_fs_and_path(str(path), storage_filesystem)
+
+        # Validate that the experiment directory exists
+        if not _exists_at_fs_path(fs, fs_path):
+            raise RuntimeError(f"Experiment folder {fs_path} doesn't exist.")
+
+        # Remove trailing slashes to handle paths correctly
+        # os.path.basename() returns empty string for paths with trailing slashes
+        fs_path = fs_path.rstrip("/")
+        storage_path, experiment_dir_name = os.path.dirname(fs_path), os.path.basename(
+            fs_path
+        )
+
+        storage_context = StorageContext(
+            storage_path=storage_path,
+            experiment_dir_name=experiment_dir_name,
+            storage_filesystem=fs,
+        )
+
+        # Validate that the checkpoint manager snapshot file exists
+        if not _exists_at_fs_path(
+            storage_context.storage_filesystem,
+            storage_context.checkpoint_manager_snapshot_path,
+        ):
+            raise RuntimeError(
+                f"Failed to restore the Result object: "
+                f"{CHECKPOINT_MANAGER_SNAPSHOT_FILENAME} doesn't exist in the "
+                f"experiment folder. Make sure that this is an output directory created by a Ray Train run."
+            )
+
+        checkpoint_manager = CheckpointManager(
+            storage_context=storage_context,
+            checkpoint_config=CheckpointConfig(),
+        )
+
+        # When we build a Result object from checkpoints, the error is not loaded.
+        return cls._from_checkpoint_manager(
+            checkpoint_manager=checkpoint_manager,
+            storage_context=storage_context,
+        )
+
+    @classmethod
+    def _from_checkpoint_manager(
+        cls,
+        checkpoint_manager: CheckpointManager,
+        storage_context: StorageContext,
+        error: Optional[TrainingFailedError] = None,
+    ) -> "Result":
+        """Create a Result object from a CheckpointManager."""
+        latest_checkpoint_result = checkpoint_manager.latest_checkpoint_result
+        if latest_checkpoint_result:
+            latest_metrics = latest_checkpoint_result.metrics
+            latest_checkpoint = latest_checkpoint_result.checkpoint
+        else:
+            latest_metrics = None
+            latest_checkpoint = None
+        best_checkpoints = [
+            (r.checkpoint, r.metrics)
+            for r in checkpoint_manager.best_checkpoint_results
+        ]
+
+        # Provide the history of metrics attached to checkpoints as a dataframe.
+        metrics_dataframe = None
+        if best_checkpoints:
+            metrics_dataframe = pd.DataFrame([m for _, m in best_checkpoints])
+
+        return Result(
+            metrics=latest_metrics,
+            checkpoint=latest_checkpoint,
+            error=error,
+            path=storage_context.experiment_fs_path,
+            best_checkpoints=best_checkpoints,
+            metrics_dataframe=metrics_dataframe,
+            _storage_filesystem=storage_context.storage_filesystem,
+        )
 
     @property
     @Deprecated

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

@@ -1,9 +1,11 @@
 import logging
 
+import boto3
 import pytest
 
 import ray
 from ray import runtime_context
+from ray._common.test_utils import simulate_s3_bucket
 from ray.cluster_utils import Cluster
 from ray.train.v2._internal.constants import (
     ENABLE_STATE_ACTOR_RECONCILIATION_ENV_VAR,
@@ -80,3 +82,25 @@ def mock_current_actor(self):
     )
 
     yield
+
+
+@pytest.fixture
+def mock_s3_bucket_uri():
+    from ray.air._internal.uri_utils import URI
+
+    port = 5002
+    region = "us-west-2"
+    with simulate_s3_bucket(port=port, region=region) as s3_uri:
+        s3 = boto3.client(
+            "s3", region_name=region, endpoint_url=f"http://localhost:{port}"
+        )
+        # Bucket name will be autogenerated/unique per test
+        bucket_name = URI(s3_uri).name
+        s3.create_bucket(
+            Bucket=bucket_name,
+            CreateBucketConfiguration={"LocationConstraint": region},
+        )
+        # Disable server HTTP request logging
+        logging.getLogger("werkzeug").setLevel(logging.WARNING)
+        yield s3_uri
+        logging.getLogger("werkzeug").setLevel(logging.INFO)