[air] Add Checkpoint.as_directory() for efficient checkpoint fs processing

python/ray/ml/checkpoint.py

@@ -1,17 +1,18 @@
+import contextlib
 import io
 import os
 import shutil
 import tarfile
 import tempfile
-from typing import Optional, Union, Tuple
+from typing import Iterator, Optional, Tuple, Union
 
 import ray
 from ray import cloudpickle as pickle
 from ray.ml.utils.remote_storage import (
-    upload_to_uri,
-    is_non_local_path_uri,
     download_from_uri,
     fs_hint,
+    is_non_local_path_uri,
+    upload_to_uri,
 )
 from ray.util.annotations import DeveloperAPI
 
@@ -338,6 +339,37 @@ def to_directory(self, path: Optional[str] = None) -> str:
 
         return path
 
+    @contextlib.contextmanager
+    def as_directory(self) -> Iterator[str]:
+        """Return checkpoint directory path in a context.
+
+        This function makes checkpoint data available as a directory while avoiding
+        unnecessary copies and left-over temporary data.
+
+        If the checkpoint is already a directory checkpoint, it will return
+        the existing path. If it is not, it will create a temporary directory,
+        which will be deleted after the context is exited.
+
+        Users should treat the returned checkpoint directory as read-only and avoid
+        changing any data within it, as it might get deleted when exiting the context.
+
+        Example:
+
+            with checkpoint.as_directory() as checkpoint_dir:
+                # Do some read-only processing of files within checkpoint_dir
+                pass
+
+            # At this point, if a temporary directory was created, it will have
+            # been deleted.
+
+        """
+        if self._local_path:
+            yield self._local_path
+        else:
+            temp_dir = self.to_directory()
+            yield temp_dir
+            shutil.rmtree(temp_dir)
+
     @classmethod
     def from_uri(cls, uri: str) -> "Checkpoint":
         """Create checkpoint object from location URI (e.g. cloud storage).

python/ray/ml/tests/test_checkpoints.py

@@ -270,6 +270,22 @@ def test_fs_delete_at_uri(self):
         with self.assertRaises(FileNotFoundError):
             checkpoint.to_directory()
 
+    def test_fs_cp_as_directory(self):
+        checkpoint = self._prepare_fs_checkpoint()
+
+        with checkpoint.as_directory() as checkpoint_dir:
+            assert checkpoint._local_path == checkpoint_dir
+
+        assert os.path.exists(checkpoint_dir)
+
+    def test_dict_cp_as_directory(self):
+        checkpoint = self._prepare_dict_checkpoint()
+
+        with checkpoint.as_directory() as checkpoint_dir:
+            assert os.path.exists(checkpoint_dir)
+
+        assert not os.path.exists(checkpoint_dir)
+
 
 class CheckpointsSerdeTest(unittest.TestCase):
     def setUp(self) -> None: