wehub-resource-sync
574 lines
21 KiB
Python
574 lines
21 KiB
Python
# Try import ray[train] core requirements (defined in setup.py)
|
|
# isort: off
|
|
try:
|
|
import fsspec # noqa
|
|
from fsspec.implementations.local import LocalFileSystem
|
|
|
|
except (ImportError, ModuleNotFoundError) as e:
|
|
raise RuntimeError(
|
|
"fsspec is a required dependency of Ray Train and Ray Tune. "
|
|
"Please install with: `pip install fsspec`"
|
|
) from e
|
|
|
|
try:
|
|
import pyarrow
|
|
import pyarrow.fs
|
|
|
|
except (ImportError, ModuleNotFoundError) as e:
|
|
raise RuntimeError(
|
|
"pyarrow is a required dependency of Ray Train and Ray Tune. "
|
|
"Please install with: `pip install pyarrow`"
|
|
) from e
|
|
# isort: on
|
|
|
|
import fnmatch
|
|
import logging
|
|
import os
|
|
import shutil
|
|
from pathlib import Path
|
|
from typing import TYPE_CHECKING, Callable, List, Optional, Tuple, Type, Union
|
|
|
|
from ray.air._internal.filelock import TempFileLock
|
|
from ray.train.constants import _get_ray_train_session_dir
|
|
from ray.train.v2._internal.constants import (
|
|
CHECKPOINT_MANAGER_SNAPSHOT_FILENAME,
|
|
VALIDATE_STORAGE_MARKER_FILENAME,
|
|
)
|
|
from ray.train.v2._internal.util import date_str
|
|
from ray.util.annotations import DeveloperAPI
|
|
|
|
if TYPE_CHECKING:
|
|
from ray.train import Checkpoint
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class _ExcludingLocalFilesystem(LocalFileSystem):
|
|
"""LocalFileSystem wrapper to exclude files according to patterns.
|
|
|
|
Args:
|
|
root_path: Root path to strip when matching with the exclude pattern.
|
|
Ex: root_path="/tmp/a/b/c", exclude=["*a*"], will exclude
|
|
/tmp/a/b/c/_a_.txt but not ALL of /tmp/a/*.
|
|
exclude: List of patterns that are applied to files returned by
|
|
``self.find()``. If a file path matches this pattern, it will
|
|
be excluded.
|
|
**kwargs: Additional keyword arguments forwarded to
|
|
``pyarrow.fs.LocalFileSystem``.
|
|
"""
|
|
|
|
def __init__(self, root_path: Path, exclude: List[str], **kwargs):
|
|
super().__init__(**kwargs)
|
|
self._exclude = exclude
|
|
self._root_path = root_path
|
|
|
|
@property
|
|
def fsid(self):
|
|
return "_excluding_local"
|
|
|
|
def _should_exclude(self, path: str) -> bool:
|
|
"""Return True if `path` (relative to `root_path`) matches any of the
|
|
`self._exclude` patterns."""
|
|
path = Path(path)
|
|
relative_path = path.relative_to(self._root_path).as_posix()
|
|
match_candidates = [relative_path]
|
|
if path.is_dir():
|
|
# Everything is in posix path format ('/')
|
|
match_candidates.append(relative_path + "/")
|
|
|
|
for excl in self._exclude:
|
|
if any(fnmatch.fnmatch(candidate, excl) for candidate in match_candidates):
|
|
return True
|
|
return False
|
|
|
|
def find(self, path, maxdepth=None, withdirs=False, detail=False, **kwargs):
|
|
"""Call parent find() and exclude from result."""
|
|
paths = super().find(
|
|
path, maxdepth=maxdepth, withdirs=withdirs, detail=detail, **kwargs
|
|
)
|
|
if detail:
|
|
return {
|
|
path: out
|
|
for path, out in paths.items()
|
|
if not self._should_exclude(path)
|
|
}
|
|
else:
|
|
return [path for path in paths if not self._should_exclude(path)]
|
|
|
|
|
|
def _pyarrow_fs_copy_files(
|
|
source, destination, source_filesystem=None, destination_filesystem=None, **kwargs
|
|
):
|
|
if isinstance(destination_filesystem, pyarrow.fs.S3FileSystem):
|
|
# Workaround multi-threading issue with pyarrow. Note that use_threads=True
|
|
# is safe for download, just not for uploads, see:
|
|
# https://github.com/apache/arrow/issues/32372
|
|
kwargs.setdefault("use_threads", False)
|
|
|
|
# Use a large chunk size to speed up large checkpoint transfers.
|
|
kwargs.setdefault("chunk_size", 64 * 1024 * 1024)
|
|
|
|
return pyarrow.fs.copy_files(
|
|
source,
|
|
destination,
|
|
source_filesystem=source_filesystem,
|
|
destination_filesystem=destination_filesystem,
|
|
**kwargs,
|
|
)
|
|
|
|
|
|
# TODO(justinvyu): Add unit tests for all these utils.
|
|
|
|
|
|
def delete_fs_path(fs: pyarrow.fs.FileSystem, fs_path: str):
|
|
"""Deletes (fs, fs_path) or raises FileNotFoundError if it doesn't exist."""
|
|
is_dir = _is_directory(fs, fs_path)
|
|
|
|
try:
|
|
if is_dir:
|
|
fs.delete_dir(fs_path)
|
|
else:
|
|
fs.delete_file(fs_path)
|
|
except Exception:
|
|
logger.exception(f"Caught exception when deleting path at ({fs}, {fs_path}):")
|
|
|
|
|
|
def _download_from_fs_path(
|
|
fs: pyarrow.fs.FileSystem,
|
|
fs_path: str,
|
|
local_path: str,
|
|
filelock: bool = True,
|
|
):
|
|
"""Downloads a directory or file from (fs, fs_path) to a local path.
|
|
|
|
If fs_path points to a directory:
|
|
- The full directory contents are downloaded directly into `local_path`,
|
|
rather than to a subdirectory of `local_path`.
|
|
|
|
If fs_path points to a file:
|
|
- The file is downloaded to `local_path`, which is expected to be a file path.
|
|
|
|
If the download fails, the `local_path` contents are
|
|
cleaned up before raising, if the directory did not previously exist.
|
|
|
|
NOTE: This method creates `local_path`'s parent directories if they do not
|
|
already exist. If the download fails, this does NOT clean up all the parent
|
|
directories that were created.
|
|
|
|
Args:
|
|
fs: The filesystem to download from.
|
|
fs_path: The filesystem path (either a directory or a file) to download.
|
|
local_path: The local path to download to.
|
|
filelock: Whether to require a file lock before downloading, useful for
|
|
multiple downloads to the same directory that may be happening in parallel.
|
|
|
|
Raises:
|
|
FileNotFoundError: if (fs, fs_path) doesn't exist.
|
|
"""
|
|
|
|
_local_path = Path(local_path).resolve()
|
|
exists_before = _local_path.exists()
|
|
if _is_directory(fs=fs, fs_path=fs_path):
|
|
_local_path.mkdir(parents=True, exist_ok=True)
|
|
else:
|
|
_local_path.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
try:
|
|
if filelock:
|
|
with TempFileLock(f"{os.path.normpath(local_path)}.lock"):
|
|
_pyarrow_fs_copy_files(fs_path, local_path, source_filesystem=fs)
|
|
else:
|
|
_pyarrow_fs_copy_files(fs_path, local_path, source_filesystem=fs)
|
|
except Exception as e:
|
|
# Clean up the directory if downloading was unsuccessful
|
|
if not exists_before:
|
|
shutil.rmtree(local_path, ignore_errors=True)
|
|
raise e
|
|
|
|
|
|
def _upload_to_fs_path(
|
|
local_path: str,
|
|
fs: pyarrow.fs.FileSystem,
|
|
fs_path: str,
|
|
exclude: Optional[List[str]] = None,
|
|
) -> None:
|
|
"""Uploads a local directory or file to (fs, fs_path).
|
|
|
|
NOTE: This will create all necessary parent directories at the destination.
|
|
|
|
Args:
|
|
local_path: The local path to upload.
|
|
fs: The filesystem to upload to.
|
|
fs_path: The filesystem path where the dir/file will be uploaded to.
|
|
exclude: A list of filename matches to exclude from upload. This includes
|
|
all files under subdirectories as well.
|
|
This pattern will match with the relative paths of all files under
|
|
`local_path`.
|
|
Ex: ["*.png"] to exclude all .png images.
|
|
"""
|
|
|
|
if not exclude:
|
|
# TODO(justinvyu): uploading a single file doesn't work
|
|
# (since we always create a directory at fs_path)
|
|
_create_directory(fs=fs, fs_path=fs_path)
|
|
_pyarrow_fs_copy_files(local_path, fs_path, destination_filesystem=fs)
|
|
return
|
|
|
|
_upload_to_uri_with_exclude_fsspec(
|
|
local_path=local_path, fs=fs, fs_path=fs_path, exclude=exclude
|
|
)
|
|
|
|
|
|
def _upload_to_uri_with_exclude_fsspec(
|
|
local_path: str, fs: "pyarrow.fs", fs_path: str, exclude: Optional[List[str]]
|
|
) -> None:
|
|
local_fs = _ExcludingLocalFilesystem(root_path=local_path, exclude=exclude)
|
|
handler = pyarrow.fs.FSSpecHandler(local_fs)
|
|
source_fs = pyarrow.fs.PyFileSystem(handler)
|
|
|
|
_create_directory(fs=fs, fs_path=fs_path)
|
|
_pyarrow_fs_copy_files(
|
|
local_path, fs_path, source_filesystem=source_fs, destination_filesystem=fs
|
|
)
|
|
|
|
|
|
def _list_at_fs_path(
|
|
fs: pyarrow.fs.FileSystem,
|
|
fs_path: str,
|
|
file_filter: Callable[[pyarrow.fs.FileInfo], bool] = lambda x: True,
|
|
) -> List[str]:
|
|
"""Returns the list of filenames at (fs, fs_path), similar to os.listdir.
|
|
|
|
If the path doesn't exist, returns an empty list.
|
|
"""
|
|
selector = pyarrow.fs.FileSelector(fs_path, allow_not_found=True, recursive=False)
|
|
return [
|
|
os.path.relpath(file_info.path.lstrip("/"), start=fs_path.lstrip("/"))
|
|
for file_info in fs.get_file_info(selector)
|
|
if file_filter(file_info)
|
|
]
|
|
|
|
|
|
def _exists_at_fs_path(fs: pyarrow.fs.FileSystem, fs_path: str) -> bool:
|
|
"""Returns True if (fs, fs_path) exists."""
|
|
|
|
valid = fs.get_file_info(fs_path)
|
|
return valid.type != pyarrow.fs.FileType.NotFound
|
|
|
|
|
|
def _is_directory(fs: pyarrow.fs.FileSystem, fs_path: str) -> bool:
|
|
"""Checks if (fs, fs_path) is a directory or a file.
|
|
|
|
Args:
|
|
fs: The filesystem to query.
|
|
fs_path: The path on the filesystem.
|
|
|
|
Returns:
|
|
``True`` if the path is a directory, ``False`` if it is a file.
|
|
|
|
Raises:
|
|
FileNotFoundError: if (fs, fs_path) doesn't exist.
|
|
"""
|
|
|
|
file_info = fs.get_file_info(fs_path)
|
|
if file_info.type == pyarrow.fs.FileType.NotFound:
|
|
raise FileNotFoundError(f"Path not found: ({fs}, {fs_path})")
|
|
|
|
return not file_info.is_file
|
|
|
|
|
|
def _create_directory(fs: pyarrow.fs.FileSystem, fs_path: str) -> None:
|
|
"""Create directory at (fs, fs_path).
|
|
|
|
Some external filesystems require directories to already exist, or at least
|
|
the `netloc` to be created (e.g. PyArrows ``mock://`` filesystem).
|
|
|
|
Generally this should be done before and outside of Ray applications. This
|
|
utility is thus primarily used in testing, e.g. of ``mock://` URIs.
|
|
"""
|
|
try:
|
|
fs.create_dir(fs_path)
|
|
except Exception:
|
|
logger.exception(
|
|
f"Caught exception when creating directory at ({fs}, {fs_path}):"
|
|
)
|
|
|
|
|
|
def get_fs_and_path(
|
|
storage_path: Union[str, os.PathLike],
|
|
storage_filesystem: Optional[pyarrow.fs.FileSystem] = None,
|
|
) -> Tuple[pyarrow.fs.FileSystem, str]:
|
|
"""Returns the fs and path from a storage path and an optional custom fs.
|
|
|
|
Args:
|
|
storage_path: A storage path or URI. (ex: s3://bucket/path or /tmp/ray_results)
|
|
storage_filesystem: A custom filesystem to use. If not provided,
|
|
this will be auto-resolved by pyarrow. If provided, the storage_path
|
|
is assumed to be prefix-stripped already, and must be a valid path
|
|
on the filesystem.
|
|
|
|
Returns:
|
|
A ``(filesystem, path)`` tuple.
|
|
"""
|
|
storage_path = str(storage_path)
|
|
|
|
if storage_filesystem:
|
|
return storage_filesystem, storage_path
|
|
|
|
return pyarrow.fs.FileSystem.from_uri(storage_path)
|
|
|
|
|
|
@DeveloperAPI
|
|
class StorageContext:
|
|
"""Shared context that holds the source of truth for all paths and
|
|
storage utilities, passed along from the driver to workers.
|
|
|
|
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. *_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":
|
|
|
|
>>> import ray
|
|
>>> from ray.train._internal.storage import StorageContext
|
|
>>> import os
|
|
>>> _ = ray.init()
|
|
>>> storage = StorageContext(
|
|
... storage_path="mock://netloc/bucket/path?param=1",
|
|
... experiment_dir_name="exp_name",
|
|
... )
|
|
>>> storage.storage_filesystem # Auto-resolved # doctest: +ELLIPSIS
|
|
<pyarrow._fs._MockFileSystem object...
|
|
>>> storage.experiment_fs_path
|
|
'bucket/path/exp_name'
|
|
>>> 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_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'
|
|
>>> ray.shutdown()
|
|
|
|
Example with storage_path="/tmp/ray_results":
|
|
|
|
>>> from ray.train._internal.storage import StorageContext
|
|
>>> storage = StorageContext(
|
|
... storage_path="/tmp/ray_results",
|
|
... experiment_dir_name="exp_name",
|
|
... )
|
|
>>> storage.storage_fs_path
|
|
'/tmp/ray_results'
|
|
>>> storage.experiment_fs_path
|
|
'/tmp/ray_results/exp_name'
|
|
>>> storage.storage_filesystem # Auto-resolved # doctest: +ELLIPSIS
|
|
<pyarrow._fs.LocalFileSystem object...
|
|
|
|
Internal Usage Examples:
|
|
- To copy files to the trial directory on the storage filesystem:
|
|
|
|
pyarrow.fs.copy_files(
|
|
local_dir,
|
|
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__(
|
|
self,
|
|
storage_path: Union[str, os.PathLike],
|
|
experiment_dir_name: str,
|
|
storage_filesystem: Optional[pyarrow.fs.FileSystem] = None,
|
|
read_only: bool = False,
|
|
):
|
|
self.custom_fs_provided = storage_filesystem is not None
|
|
|
|
# Invariant: (`storage_filesystem`, `storage_path`) is the location where
|
|
# *all* results can be accessed.
|
|
self.experiment_dir_name = experiment_dir_name
|
|
|
|
self.storage_filesystem, self.storage_fs_path = get_fs_and_path(
|
|
storage_path, storage_filesystem
|
|
)
|
|
self.storage_fs_path = Path(self.storage_fs_path).as_posix()
|
|
|
|
self.read_only = read_only
|
|
if not self.read_only:
|
|
self._create_validation_file()
|
|
self._check_validation_file()
|
|
|
|
def __str__(self):
|
|
return (
|
|
"StorageContext<\n"
|
|
f" storage_filesystem='{self.storage_filesystem.type_name}',\n"
|
|
f" storage_fs_path='{self.storage_fs_path}',\n"
|
|
f" experiment_dir_name='{self.experiment_dir_name}',\n"
|
|
">"
|
|
)
|
|
|
|
def _create_validation_file(self):
|
|
"""On the creation of a storage context, create a validation file at the
|
|
storage path to verify that the storage path can be written to.
|
|
This validation file is also used to check whether the storage path is
|
|
accessible by all nodes in the cluster."""
|
|
valid_file = Path(
|
|
self.experiment_fs_path, VALIDATE_STORAGE_MARKER_FILENAME
|
|
).as_posix()
|
|
self.storage_filesystem.create_dir(self.experiment_fs_path)
|
|
with self.storage_filesystem.open_output_stream(valid_file):
|
|
pass
|
|
|
|
def _check_validation_file(self):
|
|
"""Checks that the validation file exists at the storage path."""
|
|
valid_file = Path(
|
|
self.experiment_fs_path, VALIDATE_STORAGE_MARKER_FILENAME
|
|
).as_posix()
|
|
if not _exists_at_fs_path(fs=self.storage_filesystem, fs_path=valid_file):
|
|
raise RuntimeError(
|
|
f"Unable to set up cluster storage with the following settings:\n{self}"
|
|
"\nCheck that all nodes in the cluster have read/write access "
|
|
"to the configured storage path. `RunConfig(storage_path)` should be "
|
|
"set to a cloud storage URI or a shared filesystem path accessible "
|
|
"by all nodes in your cluster ('s3://bucket' or '/mnt/nfs'). "
|
|
"A local path on the head node is not accessible by worker nodes. "
|
|
"See: https://docs.ray.io/en/latest/train/user-guides/persistent-storage.html" # noqa: E501
|
|
)
|
|
|
|
def persist_current_checkpoint(
|
|
self, checkpoint: "Checkpoint", checkpoint_dir_name: str
|
|
) -> "Checkpoint":
|
|
"""Persists a given checkpoint to the current checkpoint path on the filesystem.
|
|
|
|
This method copies the checkpoint files to the storage location.
|
|
It's up to the user to delete the original checkpoint files if desired.
|
|
|
|
For example, the original directory is typically a local temp directory.
|
|
|
|
Args:
|
|
checkpoint: The checkpoint to persist to
|
|
(fs, experiment_fs_path / checkpoint_dir_name).
|
|
checkpoint_dir_name: Name of the destination directory for the
|
|
checkpoint, relative to ``experiment_fs_path``.
|
|
|
|
Returns:
|
|
Checkpoint: A Checkpoint pointing to the persisted checkpoint location.
|
|
"""
|
|
if self.read_only:
|
|
raise RuntimeError(
|
|
"Cannot perform write/validation operations as the StorageContext is read-only."
|
|
)
|
|
|
|
# TODO(justinvyu): Fix this cyclical import.
|
|
from ray.train import Checkpoint
|
|
|
|
checkpoint_fs_path = self.build_checkpoint_path_from_name(checkpoint_dir_name)
|
|
|
|
logger.debug(
|
|
"Copying checkpoint files to storage path:\n"
|
|
"({source_fs}, {source}) -> ({dest_fs}, {destination})".format(
|
|
source=checkpoint.path,
|
|
destination=checkpoint_fs_path,
|
|
source_fs=checkpoint.filesystem,
|
|
dest_fs=self.storage_filesystem,
|
|
)
|
|
)
|
|
|
|
# Raise an error if the storage path is not accessible when
|
|
# attempting to upload a checkpoint from a remote worker.
|
|
# Ex: If storage_path is a local path, then a validation marker
|
|
# will only exist on the head node but not the worker nodes.
|
|
self._check_validation_file()
|
|
|
|
self.storage_filesystem.create_dir(checkpoint_fs_path)
|
|
_pyarrow_fs_copy_files(
|
|
source=checkpoint.path,
|
|
destination=checkpoint_fs_path,
|
|
source_filesystem=checkpoint.filesystem,
|
|
destination_filesystem=self.storage_filesystem,
|
|
)
|
|
|
|
persisted_checkpoint = Checkpoint(
|
|
filesystem=self.storage_filesystem,
|
|
path=checkpoint_fs_path,
|
|
)
|
|
logger.info(f"Checkpoint successfully created at: {persisted_checkpoint}")
|
|
return persisted_checkpoint
|
|
|
|
@property
|
|
def experiment_fs_path(self) -> str:
|
|
"""The path on the `storage_filesystem` to the experiment directory.
|
|
|
|
NOTE: This does not have a URI prefix anymore, since it has been stripped
|
|
by pyarrow.fs.FileSystem.from_uri already. The URI scheme information is
|
|
kept in `storage_filesystem` instead.
|
|
"""
|
|
return Path(self.storage_fs_path, self.experiment_dir_name).as_posix()
|
|
|
|
@property
|
|
def local_working_directory(self) -> str:
|
|
"""Every ray train worker will set this directory as its working directory."""
|
|
if self.experiment_dir_name is None:
|
|
raise RuntimeError(
|
|
"Cannot access `local_working_directory` without "
|
|
"setting `experiment_dir_name`"
|
|
)
|
|
return Path(_get_ray_train_session_dir(), self.experiment_dir_name).as_posix()
|
|
|
|
@property
|
|
def checkpoint_manager_snapshot_path(self) -> str:
|
|
"""The path to the checkpoint manager snapshot file."""
|
|
return Path(
|
|
self.experiment_fs_path, CHECKPOINT_MANAGER_SNAPSHOT_FILENAME
|
|
).as_posix()
|
|
|
|
@staticmethod
|
|
def get_experiment_dir_name(run_obj: Union[str, Callable, Type]) -> str:
|
|
from ray.tune.experiment import Experiment
|
|
|
|
run_identifier = Experiment.get_trainable_name(run_obj)
|
|
|
|
if bool(int(os.environ.get("TUNE_DISABLE_DATED_SUBDIR", 0))):
|
|
dir_name = run_identifier
|
|
else:
|
|
dir_name = "{}_{}".format(run_identifier, date_str())
|
|
return dir_name
|
|
|
|
@staticmethod
|
|
def make_default_checkpoint_dir_name():
|
|
"""Get the name of the checkpoint directory by timestamp."""
|
|
return f"checkpoint_{date_str(include_ms=True)}"
|
|
|
|
def extract_checkpoint_dir_name_from_path(self, checkpoint_path: str) -> str:
|
|
"""Get the checkpoint name from the checkpoint path.
|
|
The parent directory of the checkpoint path should be the experiment directory.
|
|
"""
|
|
# TODO: Use Pathlib to extract the name when supports at least Python 3.9
|
|
experiment_fs_path = self.experiment_fs_path + "/"
|
|
if not checkpoint_path.startswith(experiment_fs_path):
|
|
raise ValueError(
|
|
f"Checkpoint path {checkpoint_path} is not under the experiment "
|
|
f"directory {self.experiment_fs_path}."
|
|
)
|
|
return checkpoint_path[len(experiment_fs_path) :]
|
|
|
|
def build_checkpoint_path_from_name(self, checkpoint_name: str) -> str:
|
|
"""Get the checkpoint path from the checkpoint name.
|
|
The parent directory of the checkpoint path should be the experiment directory.
|
|
"""
|
|
return Path(self.experiment_fs_path, checkpoint_name).as_posix()
|