wehub-resource-sync
435 lines
17 KiB
Python
435 lines
17 KiB
Python
import contextlib
|
|
import glob
|
|
import json
|
|
import logging
|
|
import os
|
|
import platform
|
|
import shutil
|
|
import tempfile
|
|
import traceback
|
|
import uuid
|
|
from pathlib import Path
|
|
from typing import Any, Dict, Iterator, List, Optional, Union
|
|
|
|
import pyarrow.fs
|
|
|
|
from ray.air._internal.filelock import TempFileLock
|
|
from ray.train._internal.storage import _download_from_fs_path, _exists_at_fs_path
|
|
from ray.util.annotations import PublicAPI
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# The filename of the file that stores user metadata set on the checkpoint.
|
|
_METADATA_FILE_NAME = ".metadata.json"
|
|
|
|
# The prefix of the temp checkpoint directory that `to_directory` downloads to
|
|
# on the local filesystem.
|
|
_CHECKPOINT_TEMP_DIR_PREFIX = "checkpoint_tmp_"
|
|
|
|
|
|
class _CheckpointMetaClass(type):
|
|
def __getattr__(self, item):
|
|
try:
|
|
return super().__getattribute__(item)
|
|
except AttributeError as exc:
|
|
if item in {
|
|
"from_dict",
|
|
"to_dict",
|
|
"from_bytes",
|
|
"to_bytes",
|
|
"get_internal_representation",
|
|
}:
|
|
raise _get_migration_error(item) from exc
|
|
elif item in {
|
|
"from_uri",
|
|
"to_uri",
|
|
"uri",
|
|
}:
|
|
raise _get_uri_error(item) from exc
|
|
elif item in {"get_preprocessor", "set_preprocessor"}:
|
|
raise _get_preprocessor_error(item) from exc
|
|
|
|
raise exc
|
|
|
|
|
|
@PublicAPI(stability="beta")
|
|
class Checkpoint(metaclass=_CheckpointMetaClass):
|
|
"""A reference to data persisted as a directory in local or remote storage.
|
|
|
|
Access the checkpoint contents locally using ``checkpoint.to_directory()``
|
|
or ``checkpoint.as_directory``.
|
|
|
|
Attributes
|
|
----------
|
|
path: A path on the filesystem containing the checkpoint contents.
|
|
filesystem: PyArrow FileSystem that can be used to access data at the `path`.
|
|
|
|
See Also
|
|
--------
|
|
ray.train.report : Report a checkpoint during training (with Ray Train/Tune).
|
|
ray.train.get_checkpoint : Get the latest checkpoint during training
|
|
(for restoration).
|
|
|
|
:ref:`train-checkpointing`
|
|
:ref:`persistent-storage-guide`
|
|
|
|
Examples
|
|
--------
|
|
|
|
Creating a checkpoint using ``Checkpoint.from_directory``:
|
|
|
|
>>> from ray.train import Checkpoint
|
|
>>> checkpoint = Checkpoint.from_directory("/tmp/example_checkpoint_dir")
|
|
>>> checkpoint.filesystem # doctest: +ELLIPSIS
|
|
<pyarrow._fs.LocalFileSystem object...
|
|
>>> checkpoint.path
|
|
'/tmp/example_checkpoint_dir'
|
|
|
|
Creating a checkpoint from a remote URI:
|
|
|
|
>>> checkpoint = Checkpoint("s3://bucket/path/to/checkpoint")
|
|
>>> checkpoint.filesystem # doctest: +ELLIPSIS
|
|
<pyarrow._s3fs.S3FileSystem object...
|
|
>>> checkpoint.path
|
|
'bucket/path/to/checkpoint'
|
|
|
|
Creating a checkpoint with a custom filesystem:
|
|
|
|
>>> checkpoint = Checkpoint(
|
|
... path="bucket/path/to/checkpoint",
|
|
... filesystem=pyarrow.fs.S3FileSystem(),
|
|
... )
|
|
>>> checkpoint.filesystem # doctest: +ELLIPSIS
|
|
<pyarrow._s3fs.S3FileSystem object...
|
|
>>> checkpoint.path
|
|
'bucket/path/to/checkpoint'
|
|
|
|
Accessing a checkpoint's contents:
|
|
|
|
>>> import os # doctest: +SKIP
|
|
>>> with checkpoint.as_directory() as local_checkpoint_dir: # doctest: +SKIP
|
|
... print(os.listdir(local_checkpoint_dir)) # doctest: +SKIP
|
|
['model.pt', 'optimizer.pt', 'misc.pt']
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
path: Union[str, os.PathLike],
|
|
filesystem: Optional["pyarrow.fs.FileSystem"] = None,
|
|
):
|
|
"""Construct a Checkpoint.
|
|
|
|
Args:
|
|
path: A local path or remote URI containing the checkpoint data.
|
|
If a filesystem is provided, then this path must NOT be a URI.
|
|
It should be a path on the filesystem with the prefix already stripped.
|
|
filesystem: PyArrow FileSystem to use to access data at the path.
|
|
If not specified, this is inferred from the URI scheme.
|
|
"""
|
|
self.path = str(path)
|
|
self.filesystem = filesystem
|
|
|
|
if path and not filesystem:
|
|
self.filesystem, self.path = pyarrow.fs.FileSystem.from_uri(path)
|
|
|
|
# This random UUID is used to create a temporary directory name on the
|
|
# local filesystem, which will be used for downloading checkpoint data.
|
|
# This ensures that if multiple processes download the same checkpoint object
|
|
# only one process performs the actual download while the others wait.
|
|
# This prevents duplicated download efforts and data.
|
|
# NOTE: Calling `to_directory` from multiple `Checkpoint` objects
|
|
# that point to the same (fs, path) will still download the data multiple times.
|
|
# This only ensures a canonical temp directory name for a single `Checkpoint`.
|
|
self._uuid = uuid.uuid4()
|
|
|
|
def __repr__(self):
|
|
return f"Checkpoint(filesystem={self.filesystem.type_name}, path={self.path})"
|
|
|
|
def get_metadata(self) -> Dict[str, Any]:
|
|
"""Return the metadata dict stored with the checkpoint.
|
|
|
|
If no metadata is stored, an empty dict is returned.
|
|
"""
|
|
metadata_path = Path(self.path, _METADATA_FILE_NAME).as_posix()
|
|
if not _exists_at_fs_path(self.filesystem, metadata_path):
|
|
return {}
|
|
|
|
with self.filesystem.open_input_file(metadata_path) as f:
|
|
return json.loads(f.readall().decode("utf-8"))
|
|
|
|
def set_metadata(self, metadata: Dict[str, Any]) -> None:
|
|
"""Set the metadata stored with this checkpoint.
|
|
|
|
This will overwrite any existing metadata stored with this checkpoint.
|
|
"""
|
|
metadata_path = Path(self.path, _METADATA_FILE_NAME).as_posix()
|
|
with self.filesystem.open_output_stream(metadata_path) as f:
|
|
f.write(json.dumps(metadata).encode("utf-8"))
|
|
|
|
def update_metadata(self, metadata: Dict[str, Any]) -> None:
|
|
"""Update the metadata stored with this checkpoint.
|
|
|
|
This will update any existing metadata stored with this checkpoint.
|
|
"""
|
|
existing_metadata = self.get_metadata()
|
|
existing_metadata.update(metadata)
|
|
self.set_metadata(existing_metadata)
|
|
|
|
@classmethod
|
|
def from_directory(cls, path: Union[str, os.PathLike]) -> "Checkpoint":
|
|
"""Create checkpoint object from a local directory.
|
|
|
|
Args:
|
|
path: Local directory containing checkpoint data.
|
|
|
|
Returns:
|
|
A ray.train.Checkpoint object.
|
|
"""
|
|
return cls(path, filesystem=pyarrow.fs.LocalFileSystem())
|
|
|
|
def to_directory(self, path: Optional[Union[str, os.PathLike]] = None) -> str:
|
|
"""Write checkpoint data to a local directory.
|
|
|
|
*If multiple processes on the same node call this method simultaneously,*
|
|
only a single process will perform the download, while the others
|
|
wait for the download to finish. Once the download finishes, all processes
|
|
receive the same local directory to read from.
|
|
|
|
Args:
|
|
path: Target directory to download data to. If not specified,
|
|
this method will use a temporary directory.
|
|
|
|
Returns:
|
|
str: Directory containing checkpoint data.
|
|
"""
|
|
user_provided_path = path is not None
|
|
local_path = (
|
|
path if user_provided_path else self._get_temporary_checkpoint_dir()
|
|
)
|
|
local_path = os.path.normpath(os.path.expanduser(str(local_path)))
|
|
os.makedirs(local_path, exist_ok=True)
|
|
|
|
try:
|
|
# Timeout 0 means there will be only one attempt to acquire
|
|
# the file lock. If it cannot be acquired, throw a TimeoutError
|
|
with TempFileLock(local_path, timeout=0):
|
|
_download_from_fs_path(
|
|
fs=self.filesystem, fs_path=self.path, local_path=local_path
|
|
)
|
|
except TimeoutError:
|
|
# if the directory is already locked, then wait but do not do anything.
|
|
with TempFileLock(local_path, timeout=-1):
|
|
pass
|
|
if not os.path.exists(local_path):
|
|
raise RuntimeError(
|
|
f"Checkpoint directory {local_path} does not exist, "
|
|
"even though it should have been created by "
|
|
"another process. Please raise an issue on GitHub: "
|
|
"https://github.com/ray-project/ray/issues"
|
|
)
|
|
|
|
return local_path
|
|
|
|
@contextlib.contextmanager
|
|
def as_directory(self) -> Iterator[str]:
|
|
"""Returns checkpoint contents in a local directory as a context.
|
|
|
|
This function makes checkpoint data available as a directory while avoiding
|
|
unnecessary copies and left-over temporary data.
|
|
|
|
*If the checkpoint points to a local directory*, this method just returns the
|
|
local directory path without making a copy, and nothing will be cleaned up
|
|
after exiting the context.
|
|
|
|
*If the checkpoint points to a remote directory*, this method will download the
|
|
checkpoint to a local temporary directory and return the path
|
|
to the temporary directory.
|
|
|
|
*If multiple processes on the same node call this method simultaneously,*
|
|
only a single process will perform the download, while the others
|
|
wait for the download to finish. Once the download finishes, all processes
|
|
receive the same local (temporary) directory to read from.
|
|
|
|
Once all processes have finished working with the checkpoint,
|
|
the temporary directory is cleaned up.
|
|
|
|
Users should treat the returned checkpoint directory as read-only and avoid
|
|
changing any data within it, as it may be deleted when exiting the context.
|
|
|
|
Example:
|
|
|
|
.. testcode::
|
|
:hide:
|
|
|
|
from pathlib import Path
|
|
import tempfile
|
|
|
|
from ray.train import Checkpoint
|
|
|
|
temp_dir = tempfile.mkdtemp()
|
|
(Path(temp_dir) / "example.txt").write_text("example checkpoint data")
|
|
checkpoint = Checkpoint.from_directory(temp_dir)
|
|
|
|
.. testcode::
|
|
|
|
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.
|
|
|
|
Yields:
|
|
str: The local path to the checkpoint directory.
|
|
"""
|
|
if isinstance(self.filesystem, pyarrow.fs.LocalFileSystem):
|
|
yield self.path
|
|
else:
|
|
del_lock_path = _get_del_lock_path(self._get_temporary_checkpoint_dir())
|
|
open(del_lock_path, "a").close()
|
|
|
|
temp_dir = self.to_directory()
|
|
try:
|
|
yield temp_dir
|
|
finally:
|
|
# Always cleanup the del lock after we're done with the directory.
|
|
# This avoids leaving a lock file behind in the case of an exception
|
|
# in the user code.
|
|
try:
|
|
os.remove(del_lock_path)
|
|
except Exception:
|
|
logger.warning(
|
|
f"Could not remove {del_lock_path} deletion file lock. "
|
|
f"Traceback:\n{traceback.format_exc()}"
|
|
)
|
|
|
|
# If there are no more lock files, that means there are no more
|
|
# readers of this directory, and we can safely delete it.
|
|
# In the edge case (process crash before del lock file is removed),
|
|
# we do not remove the directory at all.
|
|
# Since it's in /tmp, this is not that big of a deal.
|
|
# check if any lock files are remaining
|
|
remaining_locks = _list_existing_del_locks(temp_dir)
|
|
if not remaining_locks:
|
|
try:
|
|
# Timeout 0 means there will be only one attempt to acquire
|
|
# the file lock. If it cannot be acquired, a TimeoutError
|
|
# will be thrown.
|
|
with TempFileLock(temp_dir, timeout=0):
|
|
shutil.rmtree(temp_dir, ignore_errors=True)
|
|
except TimeoutError:
|
|
pass
|
|
|
|
def _get_temporary_checkpoint_dir(self) -> str:
|
|
"""Return the name for the temporary checkpoint dir that this checkpoint
|
|
will get downloaded to, if accessing via `to_directory` or `as_directory`.
|
|
"""
|
|
tmp_dir_path = tempfile.gettempdir()
|
|
checkpoint_dir_name = _CHECKPOINT_TEMP_DIR_PREFIX + self._uuid.hex
|
|
if platform.system() == "Windows":
|
|
# Max path on Windows is 260 chars, -1 for joining \
|
|
# Also leave a little for the del lock
|
|
del_lock_name = _get_del_lock_path("")
|
|
checkpoint_dir_name = (
|
|
_CHECKPOINT_TEMP_DIR_PREFIX
|
|
+ self._uuid.hex[
|
|
-259
|
|
+ len(_CHECKPOINT_TEMP_DIR_PREFIX)
|
|
+ len(tmp_dir_path)
|
|
+ len(del_lock_name) :
|
|
]
|
|
)
|
|
if not checkpoint_dir_name.startswith(_CHECKPOINT_TEMP_DIR_PREFIX):
|
|
raise RuntimeError(
|
|
"Couldn't create checkpoint directory due to length "
|
|
"constraints. Try specifying a shorter checkpoint path."
|
|
)
|
|
return Path(tmp_dir_path, checkpoint_dir_name).as_posix()
|
|
|
|
def __fspath__(self):
|
|
raise TypeError(
|
|
"You cannot use `Checkpoint` objects directly as paths. "
|
|
"Use `Checkpoint.to_directory()` or `Checkpoint.as_directory()` instead."
|
|
)
|
|
|
|
|
|
def _get_del_lock_path(path: str, suffix: str = None) -> str:
|
|
"""Get the path to the deletion lock file for a file/directory at `path`.
|
|
|
|
Args:
|
|
path: The path of the file or directory to generate a lock path for.
|
|
suffix: Suffix appended after ``.del_lock_``. Defaults to the current
|
|
process ID.
|
|
|
|
Returns:
|
|
The deletion lock file path.
|
|
|
|
Example:
|
|
|
|
>>> _get_del_lock_path("/tmp/checkpoint_tmp") # doctest: +ELLIPSIS
|
|
'/tmp/checkpoint_tmp.del_lock_...
|
|
>>> _get_del_lock_path("/tmp/checkpoint_tmp/") # doctest: +ELLIPSIS
|
|
'/tmp/checkpoint_tmp.del_lock_...
|
|
>>> _get_del_lock_path("/tmp/checkpoint_tmp.txt") # doctest: +ELLIPSIS
|
|
'/tmp/checkpoint_tmp.txt.del_lock_...
|
|
|
|
"""
|
|
suffix = suffix if suffix is not None else str(os.getpid())
|
|
return f"{path.rstrip('/')}.del_lock_{suffix}"
|
|
|
|
|
|
def _list_existing_del_locks(path: str) -> List[str]:
|
|
"""List all the deletion lock files for a file/directory at `path`.
|
|
|
|
For example, if 2 checkpoints are being read via `as_directory`,
|
|
then this should return a list of 2 deletion lock files.
|
|
"""
|
|
return list(glob.glob(f"{_get_del_lock_path(path, suffix='*')}"))
|
|
|
|
|
|
def _get_migration_error(name: str):
|
|
return AttributeError(
|
|
f"The new `ray.train.Checkpoint` class does not support `{name}()`. "
|
|
f"Instead, only directories are supported.\n\n"
|
|
f"Example to store a dictionary in a checkpoint:\n\n"
|
|
f"import os, tempfile\n"
|
|
f"import ray.cloudpickle as pickle\n"
|
|
f"from ray import train\n"
|
|
f"from ray.train import Checkpoint\n\n"
|
|
f"with tempfile.TemporaryDirectory() as checkpoint_dir:\n"
|
|
f" with open(os.path.join(checkpoint_dir, 'data.pkl'), 'wb') as fp:\n"
|
|
f" pickle.dump({{'data': 'value'}}, fp)\n\n"
|
|
f" checkpoint = Checkpoint.from_directory(checkpoint_dir)\n"
|
|
f" train.report(..., checkpoint=checkpoint)\n\n"
|
|
f"Example to load a dictionary from a checkpoint:\n\n"
|
|
f"if train.get_checkpoint():\n"
|
|
f" with train.get_checkpoint().as_directory() as checkpoint_dir:\n"
|
|
f" with open(os.path.join(checkpoint_dir, 'data.pkl'), 'rb') as fp:\n"
|
|
f" data = pickle.load(fp)"
|
|
)
|
|
|
|
|
|
def _get_uri_error(name: str):
|
|
return AttributeError(
|
|
f"The new `ray.train.Checkpoint` class does not support `{name}()`. "
|
|
f"To create a checkpoint from remote storage, create a `Checkpoint` using its "
|
|
f"constructor instead of `from_directory`.\n"
|
|
f'Example: `Checkpoint(path="s3://a/b/c")`.\n'
|
|
f"Then, access the contents of the checkpoint with "
|
|
f"`checkpoint.as_directory()` / `checkpoint.to_directory()`.\n"
|
|
f"To upload data to remote storage, use e.g. `pyarrow.fs.FileSystem` "
|
|
f"or your client of choice."
|
|
)
|
|
|
|
|
|
def _get_preprocessor_error(name: str):
|
|
return AttributeError(
|
|
f"The new `ray.train.Checkpoint` class does not support `{name}()`. "
|
|
f"To include preprocessor information in checkpoints, "
|
|
f"pass it as metadata in the <Framework>Trainer constructor.\n"
|
|
f"Example: `TorchTrainer(..., metadata={{...}})`.\n"
|
|
f"After training, access it in the checkpoint via `checkpoint.get_metadata()`. "
|
|
f"See here: https://docs.ray.io/en/master/train/user-guides/"
|
|
f"data-loading-preprocessing.html#preprocessing-structured-data"
|
|
)
|