[train+tune] Local directory refactor (2/n): Separate driver artifacts and trial working directories

python/ray/air/config.py

@@ -16,6 +16,7 @@
 
 import pyarrow.fs
 
+from ray._private.storage import _get_storage_uri
 from ray._private.thirdparty.tabulate.tabulate import tabulate
 from ray.util.annotations import PublicAPI, Deprecated
 from ray.widgets import Template, make_table_html_repr
@@ -581,10 +582,14 @@ class RunConfig:
     Args:
         name: Name of the trial or experiment. If not provided, will be deduced
             from the Trainable.
-        storage_path: [Beta] Path to store results at. Can be a local directory or
-            a destination on cloud storage. If Ray storage is set up,
-            defaults to the storage location. Otherwise, this defaults to
-            the local ``~/ray_results`` directory.
+        storage_path: [Beta] Path where all results and checkpoints are persisted.
+            Can be a local directory or a destination on cloud storage.
+            For multi-node training/tuning runs, this must be set to a
+            shared storage location (e.g., S3, NFS).
+            This defaults to the local ``~/ray_results`` directory.
+        storage_filesystem: [Beta] A custom filesystem to use for storage.
+            If this is provided, `storage_path` should be a path with its
+            prefix stripped (e.g., `s3://bucket/path` -> `bucket/path`).
         failure_config: Failure mode configuration.
         checkpoint_config: Checkpointing configuration.
         sync_config: Configuration object for syncing. See train.SyncConfig.
@@ -638,9 +643,17 @@ def __post_init__(self):
         from ray.tune.experimental.output import AirVerbosity, get_air_verbosity
 
         if self.storage_path is None:
-
             self.storage_path = DEFAULT_STORAGE_PATH
 
+            # If no remote path is set, try to get Ray Storage URI
+            ray_storage_uri: Optional[str] = _get_storage_uri()
+            if ray_storage_uri is not None:
+                logger.info(
+                    "Using configured Ray Storage URI as the `storage_path`: "
+                    f"{ray_storage_uri}"
+                )
+                self.storage_path = ray_storage_uri
+
         if not self.failure_config:
             self.failure_config = FailureConfig()
 

python/ray/train/_internal/storage.py

@@ -28,10 +28,10 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Tuple, Type, Union
 
-from ray._private.storage import _get_storage_uri
 from ray.air._internal.filelock import TempFileLock
 from ray.train._internal.syncer import SyncConfig, Syncer, _BackgroundSyncer
 from ray.train.constants import _get_ray_train_session_dir
+from ray.util.annotations import DeveloperAPI
 
 if TYPE_CHECKING:
     from ray.train._checkpoint import Checkpoint
@@ -345,16 +345,22 @@ def _delete_command(self, uri: str) -> Tuple[Callable, Dict]:
         return _delete_fs_path, dict(fs=self.storage_filesystem, fs_path=fs_path)
 
 
+@DeveloperAPI
 class StorageContext:
-    """Shared context that holds all paths and storage utilities, passed along from
-    the driver to workers.
+    """Shared context that holds the source of truth for all paths and
+    storage utilities, passed along from the driver to workers.
 
-    There are 2 types of paths:
+    This object defines a few types of paths:
     1. *_fs_path: A path on the `storage_filesystem`. This is a regular path
         which has been prefix-stripped by pyarrow.fs.FileSystem.from_uri and
         can be joined with `Path(...).as_posix()`.
-    2. *_local_staging_path: The temporary path on the local filesystem where results
-        are saved to before persisting them to storage.
+    2. *_driver_staging_path: The temporary staging directory on the local filesystem
+        where driver artifacts are saved to before persisting them to storage.
+    3. trial_working_directory: The local filesystem path that the remote
+        actors' working directories are moved to by default.
+        This is separated from the driver staging path so that driver syncing
+        does not implicitly upload the trial working directory, for trials on the
+        driver node.
 
     Example with storage_path="mock:///bucket/path?param=1":
 
@@ -370,12 +376,12 @@ class StorageContext:
         <pyarrow._fs._MockFileSystem object...
         >>> storage.experiment_fs_path
         'bucket/path/exp_name'
-        >>> storage.experiment_local_staging_path  # doctest: +ELLIPSIS
+        >>> storage.experiment_driver_staging_path  # doctest: +ELLIPSIS
         '/tmp/ray/session_.../artifacts/.../exp_name/driver_artifacts'
         >>> storage.trial_dir_name = "trial_dir"
         >>> storage.trial_fs_path
         'bucket/path/exp_name/trial_dir'
-        >>> storage.trial_local_staging_path  # doctest: +ELLIPSIS
+        >>> storage.trial_driver_staging_path  # doctest: +ELLIPSIS
         '/tmp/ray/session_.../artifacts/.../exp_name/driver_artifacts/trial_dir'
         >>> storage.trial_working_directory   # doctest: +ELLIPSIS
         '/tmp/ray/session_.../artifacts/.../exp_name/working_dirs/trial_dir'
@@ -406,6 +412,10 @@ class StorageContext:
             Path(storage.trial_fs_path, "subdir").as_posix(),
             destination_filesystem=storage.filesystem
         )
+
+    .. warning::
+        This is an experimental developer API and is subject to change
+        without notice between versions.
     """
 
     def __init__(
@@ -421,15 +431,6 @@ def __init__(
 
         self.custom_fs_provided = storage_filesystem is not None
 
-        # If no remote path is set, try to get Ray Storage URI
-        ray_storage_uri: Optional[str] = _get_storage_uri()
-        if ray_storage_uri is not None:
-            logger.info(
-                "Using configured Ray Storage URI as the `storage_path`: "
-                f"{ray_storage_uri}"
-            )
-            storage_path = ray_storage_uri
-
         # Invariant: (`storage_filesystem`, `storage_path`) is the location where
         # *all* results can be accessed.
         self.experiment_dir_name = experiment_dir_name
@@ -451,6 +452,10 @@ def __init__(
         self._create_validation_file()
         self._check_validation_file()
 
+        # Timestamp is used to create a unique session directory for the current
+        # training job. This is used to avoid conflicts when multiple training jobs
+        # run with the same name in the same cluster.
+        # This is set ONCE at the creation of the storage context, on the driver.
         self._timestamp = date_str()
 
     def __str__(self):
@@ -591,8 +596,15 @@ def experiment_fs_path(self) -> str:
         """
         return Path(self.storage_fs_path, self.experiment_dir_name).as_posix()
 
+    def _get_session_path(self) -> str:
+        """The Ray Train/Tune session local directory used to stage files
+        before persisting to the storage filesystem."""
+        return Path(
+            _get_ray_train_session_dir(), self._timestamp, self.experiment_dir_name
+        ).as_posix()
+
     @property
-    def experiment_local_staging_path(self) -> str:
+    def experiment_driver_staging_path(self) -> str:
         """The local filesystem path of the experiment directory on the driver node.
 
         The driver is the node where `Trainer.fit`/`Tuner.fit` is being called.
@@ -602,20 +614,15 @@ def experiment_local_staging_path(self) -> str:
         <experiment_dir_name>/driver_artifacts`
 
         This should be used as the temporary staging location for files *on the driver*
-        before syncing them to `(storage_filesystem, storage_path)`.
+        before syncing them to `experiment_fs_path`.
         For example, the search algorithm should dump its state to this directory.
-        See `trial_staging_path` for files that should be written to trial folders.
+        See `trial_driver_staging_path` for writing trial-specific artifacts.
 
         The directory is synced to
         `{storage_path}/{experiment_dir_name}` periodically.
         See `_ExperimentCheckpointManager.checkpoint` for where that happens.
         """
-        return Path(
-            _get_ray_train_session_dir(),
-            self._timestamp,
-            self.experiment_dir_name,
-            "driver_artifacts",
-        ).as_posix()
+        return Path(self._get_session_path(), "driver_artifacts").as_posix()
 
     @property
     def trial_fs_path(self) -> str:
@@ -630,7 +637,7 @@ def trial_fs_path(self) -> str:
         return Path(self.experiment_fs_path, self.trial_dir_name).as_posix()
 
     @property
-    def trial_local_staging_path(self) -> str:
+    def trial_driver_staging_path(self) -> str:
         """The local filesystem path of the trial directory on the driver.
 
         The driver is the node where `Trainer.fit`/`Tuner.fit` is being called.
@@ -640,17 +647,17 @@ def trial_local_staging_path(self) -> str:
         <experiment_dir_name>/driver_artifacts/<trial_dir_name>`
 
         This should be used as the temporary location for files on the driver
-        before syncing them to `(storage_filesystem, storage_path)`.
+        before persisting them to `trial_fs_path`.
 
         For example, callbacks (e.g., JsonLoggerCallback) should write trial-specific
         logfiles within this directory.
         """
         if self.trial_dir_name is None:
             raise RuntimeError(
-                "Should not access `trial_local_staging_path` "
+                "Should not access `trial_driver_staging_path` "
                 "without setting `trial_dir_name`"
             )
-        return Path(self.experiment_local_staging_path, self.trial_dir_name).as_posix()
+        return Path(self.experiment_driver_staging_path, self.trial_dir_name).as_posix()
 
     @property
     def trial_working_directory(self) -> str:
@@ -672,11 +679,7 @@ def trial_working_directory(self) -> str:
                 "setting `trial_dir_name`"
             )
         return Path(
-            _get_ray_train_session_dir(),
-            self._timestamp,
-            self.experiment_dir_name,
-            "working_dirs",
-            self.trial_dir_name,
+            self._get_session_path(), "working_dirs", self.trial_dir_name
         ).as_posix()
 
     @property

python/ray/train/tests/test_result.py

@@ -65,11 +65,7 @@ def build_dummy_tuner(configs):
 
 @pytest.mark.parametrize("storage", ["local", "remote"])
 @pytest.mark.parametrize("mode", ["trainer", "tuner"])
-def test_result_restore(
-    ray_start_4_cpus, monkeypatch, tmpdir, mock_s3_bucket_uri, storage, mode
-):
-    monkeypatch.setenv("RAY_AIR_LOCAL_CACHE_DIR", str(tmpdir / "ray_results"))
-
+def test_result_restore(ray_start_4_cpus, tmpdir, mock_s3_bucket_uri, storage, mode):
     NUM_ITERATIONS = 5
     NUM_CHECKPOINTS = 3
     if storage == "local":

python/ray/train/tests/test_storage.py

@@ -34,13 +34,6 @@ def storage(request, tmp_path) -> StorageContext:
         )
 
 
-@pytest.fixture(autouse=True)
-def local_path(tmp_path, monkeypatch):
-    local_dir = str(tmp_path / "ray_results")
-    monkeypatch.setenv("RAY_AIR_LOCAL_CACHE_DIR", local_dir)
-    yield local_dir
-
-
 @pytest.fixture(autouse=True, scope="module")
 def ray_init():
     # NOTE: This is needed to set the `/tmp/ray/session_*` directory.
@@ -171,14 +164,6 @@ def test_persist_artifacts(storage: StorageContext):
     trial_working_dir.joinpath("1.txt").touch()
 
     storage.persist_artifacts()
-
-    if not storage.syncer:
-        # No syncing is needed -- pass early if storage_path == storage_local_path
-        assert _list_at_fs_path(storage.storage_filesystem, storage.trial_fs_path) == [
-            "1.txt"
-        ]
-        return
-
     storage.syncer.wait()
 
     assert sorted(
@@ -206,11 +191,6 @@ def test_persist_artifacts(storage: StorageContext):
 
 def test_persist_artifacts_failures(storage: StorageContext):
     """Tests `StorageContext.persist_artifacts` edge cases (empty directory)."""
-    if not storage.syncer:
-        # Should be a no-op if storage_path == storage_local_path (no syncing needed)
-        storage.persist_artifacts()
-        return
-
     # Uploading before the trial directory has been created should fail
     with pytest.raises(FileNotFoundError):
         storage.persist_artifacts()