wehub/ray-project--ray
1
0
Fork 0
Code Issues Pull Requests Actions 1 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:17:40 +08:00
commit f1825c8ceb
10096 changed files with 2364182 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/ray/tune/utils/__init__.py
+27
View File
@@ -0,0 +1,27 @@
from ray.tune.utils.util import (
UtilMonitor,
_detect_config_single,
date_str,
deep_update,
diagnose_serialization,
flatten_dict,
merge_dicts,
unflattened_lookup,
validate_save_restore,
wait_for_gpu,
warn_if_slow,
)
__all__ = [
"deep_update",
"date_str",
"flatten_dict",
"merge_dicts",
"unflattened_lookup",
"UtilMonitor",
"validate_save_restore",
"warn_if_slow",
"diagnose_serialization",
"_detect_config_single",
"wait_for_gpu",
]
python/ray/tune/utils/callback.py
+172
View File
@@ -0,0 +1,172 @@
import importlib
import logging
import os
from typing import TYPE_CHECKING, Collection, List, Optional, Type, Union
from ray.tune.callback import Callback, CallbackList
from ray.tune.constants import RAY_TUNE_CALLBACKS_ENV_VAR
from ray.tune.logger import (
CSVLoggerCallback,
JsonLoggerCallback,
TBXLoggerCallback,
)
logger = logging.getLogger(__name__)
if TYPE_CHECKING:
from ray.tune.experimental.output import AirVerbosity
DEFAULT_CALLBACK_CLASSES = (
CSVLoggerCallback,
JsonLoggerCallback,
TBXLoggerCallback,
)
def _get_artifact_templates_for_callbacks(
callbacks: Union[List[Callback], List[Type[Callback]], CallbackList]
) -> List[str]:
templates = []
for callback in callbacks:
templates += list(callback._SAVED_FILE_TEMPLATES)
return templates
def _create_default_callbacks(
callbacks: Optional[List[Callback]],
*,
air_verbosity: Optional["AirVerbosity"] = None,
entrypoint: Optional[str] = None,
metric: Optional[str] = None,
mode: Optional[str] = None,
config: Optional[dict] = None,
progress_metrics: Optional[Collection[str]] = None,
) -> List[Callback]:
"""Create default callbacks for `Tuner.fit()`.
This function takes a list of existing callbacks and adds default
callbacks to it.
Specifically, three kinds of callbacks will be added:
1. Loggers. Ray Tune's experiment analysis relies on CSV and JSON logging.
2. Syncer. Ray Tune synchronizes logs and checkpoint between workers and
the head node.
2. Trial progress reporter. For reporting intermediate progress, like trial
results, Ray Tune uses a callback.
These callbacks will only be added if they don't already exist, i.e. if
they haven't been passed (and configured) by the user. A notable case
is when a Logger is passed, which is not a CSV or JSON logger - then
a CSV and JSON logger will still be created.
Lastly, this function will ensure that the Syncer callback comes after all
Logger callbacks, to ensure that the most up-to-date logs and checkpoints
are synced across nodes.
"""
callbacks = callbacks or []
# Initialize callbacks from environment variable
env_callbacks = _initialize_env_callbacks()
callbacks.extend(env_callbacks)
has_csv_logger = False
has_json_logger = False
has_tbx_logger = False
from ray.tune.progress_reporter import TrialProgressCallback
has_trial_progress_callback = any(
isinstance(c, TrialProgressCallback) for c in callbacks
)
if has_trial_progress_callback and air_verbosity is not None:
logger.warning(
"AIR_VERBOSITY is set, ignoring passed-in TrialProgressCallback."
)
new_callbacks = [
c for c in callbacks if not isinstance(c, TrialProgressCallback)
]
callbacks = new_callbacks
if air_verbosity is not None: # new flow
from ray.tune.experimental.output import (
_detect_reporter as _detect_air_reporter,
)
air_progress_reporter = _detect_air_reporter(
air_verbosity,
num_samples=1, # Update later with setup()
entrypoint=entrypoint,
metric=metric,
mode=mode,
config=config,
progress_metrics=progress_metrics,
)
callbacks.append(air_progress_reporter)
elif not has_trial_progress_callback: # old flow
trial_progress_callback = TrialProgressCallback(
metric=metric, progress_metrics=progress_metrics
)
callbacks.append(trial_progress_callback)
# Check if we have a CSV, JSON and TensorboardX logger
for callback in callbacks:
if isinstance(callback, CSVLoggerCallback):
has_csv_logger = True
elif isinstance(callback, JsonLoggerCallback):
has_json_logger = True
elif isinstance(callback, TBXLoggerCallback):
has_tbx_logger = True
# If CSV, JSON or TensorboardX loggers are missing, add
if os.environ.get("TUNE_DISABLE_AUTO_CALLBACK_LOGGERS", "0") != "1":
if not has_csv_logger:
callbacks.append(CSVLoggerCallback())
if not has_json_logger:
callbacks.append(JsonLoggerCallback())
if not has_tbx_logger:
try:
callbacks.append(TBXLoggerCallback())
except ImportError:
logger.warning(
"The TensorboardX logger cannot be instantiated because "
"either TensorboardX or one of it's dependencies is not "
"installed. Please make sure you have the latest version "
"of TensorboardX installed: `pip install -U tensorboardx`"
)
return callbacks
def _initialize_env_callbacks() -> List[Callback]:
"""Initialize callbacks from environment variable.
Returns:
List of callbacks initialized from environment variable.
"""
callbacks = []
callbacks_str = os.environ.get(RAY_TUNE_CALLBACKS_ENV_VAR, "")
if not callbacks_str:
return callbacks
for callback_path in callbacks_str.split(","):
callback_path = callback_path.strip()
if not callback_path:
continue
try:
module_path, class_name = callback_path.rsplit(".", 1)
module = importlib.import_module(module_path)
callback_cls = getattr(module, class_name)
if not issubclass(callback_cls, Callback):
raise TypeError(
f"Callback class '{callback_path}' must be a subclass of "
f"Callback, got {type(callback_cls).__name__}"
)
callback = callback_cls()
callbacks.append(callback)
except (ImportError, AttributeError, ValueError, TypeError) as e:
raise ValueError(f"Failed to import callback from '{callback_path}'") from e
return callbacks
python/ray/tune/utils/file_transfer.py
+481
View File
@@ -0,0 +1,481 @@
import fnmatch
import io
import os
import shutil
import tarfile
from typing import Dict, Generator, List, Optional, Tuple, Union
import ray
from ray.air._internal.filelock import TempFileLock
from ray.air.util.node import _force_on_node, _get_node_id_from_node_ip
from ray.util.annotations import DeveloperAPI
_DEFAULT_CHUNK_SIZE_BYTES = 500 * 1024 * 1024 # 500 MiB
_DEFAULT_MAX_SIZE_BYTES = 1 * 1024 * 1024 * 1024 # 1 GiB
@DeveloperAPI
def sync_dir_between_nodes(
source_ip: str,
source_path: str,
target_ip: str,
target_path: str,
force_all: bool = False,
exclude: Optional[List] = None,
chunk_size_bytes: int = _DEFAULT_CHUNK_SIZE_BYTES,
max_size_bytes: Optional[int] = _DEFAULT_MAX_SIZE_BYTES,
return_futures: bool = False,
) -> Union[
None,
Tuple[ray.ObjectRef, ray.ActorID, ray.ObjectRef],
Tuple[ray.ObjectRef, None, None],
]:
"""Synchronize directory on source node to directory on target node.
Per default, this function will collect information about already existing
files in the target directory. Only files that differ in either mtime or
filesize will be transferred, unless ``force_all=True``.
If ``source_ip==target_ip``, shutil will be used to copy the directory. Otherwise,
the directory will be packed and sent through the Ray Object Store to the target
node.
Args:
source_ip: IP of source node.
source_path: Path to directory on source node.
target_ip: IP of target node.
target_path: Path to directory on target node.
force_all: If True, all files will be transferred (not just differing files).
Ignored if ``source_ip==target_ip``.
exclude: Pattern of files to exclude, e.g.
``["*/checkpoint_*]`` to exclude trial checkpoints.
chunk_size_bytes: Chunk size for data transfer. Ignored if
``source_ip==target_ip``.
max_size_bytes: If packed data exceeds this value, raise an error before
transfer. If ``None``, no limit is enforced. Ignored if
``source_ip==target_ip``.
return_futures: If True, returns a tuple of the unpack future,
the pack actor, and the files_stats future. If False (default) will
block until synchronization finished and return None.
Returns:
None, or Tuple of unpack future, pack actor, and files_stats future.
If ``source_ip==target_ip``, pack actor and files_stats future will be None.
"""
if source_ip != target_ip:
return _sync_dir_between_different_nodes(
source_ip=source_ip,
source_path=source_path,
target_ip=target_ip,
target_path=target_path,
force_all=force_all,
exclude=exclude,
chunk_size_bytes=chunk_size_bytes,
max_size_bytes=max_size_bytes,
return_futures=return_futures,
)
elif source_path != target_path:
ret = _sync_dir_on_same_node(
ip=source_ip,
source_path=source_path,
target_path=target_path,
exclude=exclude,
return_futures=return_futures,
)
if return_futures:
return ret, None, None
return ret
def _sync_dir_on_same_node(
ip: str,
source_path: str,
target_path: str,
exclude: Optional[List] = None,
return_futures: bool = False,
) -> Optional[ray.ObjectRef]:
"""Synchronize directory to another directory on the same node.
Per default, this function will collect information about already existing
files in the target directory. All files will be copied over.
Args:
ip: IP of the node.
source_path: Path to source directory.
target_path: Path to target directory.
exclude: Pattern of files to exclude, e.g.
``["*/checkpoint_*]`` to exclude trial checkpoints.
return_futures: If True, returns a future of the copy task.
Returns:
None, or future of the copy task.
"""
node_id = _get_node_id_from_node_ip(ip)
copy_on_node = _remote_copy_dir.options(num_cpus=0, **_force_on_node(node_id))
copy_future = copy_on_node.remote(
source_dir=source_path, target_dir=target_path, exclude=exclude
)
if return_futures:
return copy_future
return ray.get(copy_future)
def _sync_dir_between_different_nodes(
source_ip: str,
source_path: str,
target_ip: str,
target_path: str,
force_all: bool = False,
exclude: Optional[List] = None,
chunk_size_bytes: int = _DEFAULT_CHUNK_SIZE_BYTES,
max_size_bytes: Optional[int] = _DEFAULT_MAX_SIZE_BYTES,
return_futures: bool = False,
) -> Union[None, Tuple[ray.ObjectRef, ray.ActorID, ray.ObjectRef]]:
"""Synchronize directory on source node to directory on target node.
Per default, this function will collect information about already existing
files in the target directory. Only files that differ in either mtime or
filesize will be transferred, unless ``force_all=True``.
Args:
source_ip: IP of source node.
source_path: Path to directory on source node.
target_ip: IP of target node.
target_path: Path to directory on target node.
force_all: If True, all files will be transferred (not just differing files).
exclude: Pattern of files to exclude, e.g.
``["*/checkpoint_*]`` to exclude trial checkpoints.
chunk_size_bytes: Chunk size for data transfer.
max_size_bytes: If packed data exceeds this value, raise an error before
transfer. If ``None``, no limit is enforced.
return_futures: If True, returns a tuple of the unpack future,
the pack actor, and the files_stats future. If False (default) will
block until synchronization finished and return None.
Returns:
None, or Tuple of unpack future, pack actor, and files_stats future.
"""
source_node_id = _get_node_id_from_node_ip(source_ip)
target_node_id = _get_node_id_from_node_ip(target_ip)
pack_actor_on_source_node = _PackActor.options(
num_cpus=0, **_force_on_node(source_node_id)
)
unpack_on_target_node = _unpack_from_actor.options(
num_cpus=0, **_force_on_node(target_node_id)
)
if force_all:
files_stats = None
else:
files_stats = _remote_get_recursive_files_and_stats.options(
num_cpus=0, **_force_on_node(target_node_id)
).remote(target_path)
pack_actor = pack_actor_on_source_node.remote(
source_dir=source_path,
files_stats=files_stats,
chunk_size_bytes=chunk_size_bytes,
max_size_bytes=max_size_bytes,
exclude=exclude,
)
unpack_future = unpack_on_target_node.remote(pack_actor, target_path)
if return_futures:
return unpack_future, pack_actor, files_stats
return ray.get(unpack_future)
def _get_recursive_files_and_stats(path: str) -> Dict[str, Tuple[float, int]]:
"""Return dict of files mapping to stats in ``path``.
This function scans a directory ``path`` recursively and returns a dict
mapping each contained file to a tuple of (mtime, filesize).
mtime and filesize are returned from ``os.lstat`` and are usually a
floating point number (timestamp) and an int (filesize in bytes).
"""
files_stats = {}
for root, dirs, files in os.walk(path, topdown=False):
rel_root = os.path.relpath(root, path)
for file in files:
try:
key = os.path.join(rel_root, file)
stat = os.lstat(os.path.join(path, key))
files_stats[key] = stat.st_mtime, stat.st_size
except FileNotFoundError:
# Race condition: If a file is deleted while executing this
# method, just continue and don't include the file in the stats
pass
return files_stats
# Only export once
_remote_get_recursive_files_and_stats = ray.remote(_get_recursive_files_and_stats)
def _pack_dir(
source_dir: str,
exclude: Optional[List] = None,
files_stats: Optional[Dict[str, Tuple[float, int]]] = None,
) -> io.BytesIO:
"""Pack whole directory contents into an uncompressed tarfile.
This function accepts a ``files_stats`` argument. If given, only files
whose stats differ from these stats will be packed.
The main use case for this is that we can collect information about files
already existing in the target directory, and only pack files that have
been updated. This is similar to how cloud syncing utilities decide
which files to transfer.
Args:
source_dir: Path to local directory to pack into tarfile.
exclude: Pattern of files to exclude, e.g.
``["*/checkpoint_*]`` to exclude trial checkpoints.
files_stats: Dict of relative filenames mapping to a tuple of
(mtime, filesize). Only files that differ from these stats
will be packed.
Returns:
Tarfile as a stream object.
"""
def _should_exclude(candidate: str) -> bool:
if not exclude:
return False
for excl in exclude:
if fnmatch.fnmatch(candidate, excl):
return True
return False
stream = io.BytesIO()
with tarfile.open(fileobj=stream, mode="w", format=tarfile.PAX_FORMAT) as tar:
if not files_stats and not exclude:
# If no `files_stats` is passed, pack whole directory
tar.add(source_dir, arcname="", recursive=True)
else:
files_stats = files_stats or {}
# Otherwise, only pack differing files
tar.add(source_dir, arcname="", recursive=False)
for root, dirs, files in os.walk(source_dir, topdown=False):
rel_root = os.path.relpath(root, source_dir)
# Always add all directories
for dir in dirs:
key = os.path.join(rel_root, dir)
tar.add(os.path.join(source_dir, key), arcname=key, recursive=False)
# Add files where our information differs
for file in files:
key = os.path.join(rel_root, file)
stat = os.lstat(os.path.join(source_dir, key))
file_stat = stat.st_mtime, stat.st_size
if _should_exclude(key):
# If the file matches an exclude pattern, skip
continue
if key in files_stats and files_stats[key] == file_stat:
# If the file did not change, skip
continue
tar.add(os.path.join(source_dir, key), arcname=key)
return stream
def _gib_string(num_bytes: float) -> str:
return f"{float(num_bytes / 1024 ** 3):.2f}GiB"
@ray.remote
class _PackActor:
"""Actor wrapping around a packing job.
This actor is used for chunking the packed data into smaller chunks that
can be transferred via the object store more efficiently.
The actor will start packing the directory when initialized, and separate
chunks can be received by calling the remote ``next()`` task.
Args:
source_dir: Path to local directory to pack into tarfile.
exclude: Pattern of files to exclude, e.g.
``["*/checkpoint_*]`` to exclude trial checkpoints.
files_stats: Dict of relative filenames mapping to a tuple of
(mtime, filesize). Only files that differ from these stats
will be packed.
chunk_size_bytes: Cut bytes stream into chunks of this size in bytes.
max_size_bytes: If packed data exceeds this value, raise an error before
transfer. If ``None``, no limit is enforced.
"""
def __init__(
self,
source_dir: str,
exclude: Optional[List] = None,
files_stats: Optional[Dict[str, Tuple[float, int]]] = None,
chunk_size_bytes: int = _DEFAULT_CHUNK_SIZE_BYTES,
max_size_bytes: Optional[int] = _DEFAULT_MAX_SIZE_BYTES,
):
self.stream = _pack_dir(
source_dir=source_dir, exclude=exclude, files_stats=files_stats
)
# Get buffer size
self.stream.seek(0, 2)
file_size = self.stream.tell()
if max_size_bytes and file_size > max_size_bytes:
raise RuntimeError(
f"Packed directory {source_dir} content has a size of "
f"{_gib_string(file_size)}, which exceeds the limit "
f"of {_gib_string(max_size_bytes)}. Please check the directory "
f"contents. If you want to transfer everything, you can increase "
f"or disable the limit by passing the `max_size` argument."
)
self.chunk_size = chunk_size_bytes
self.max_size = max_size_bytes
self.iter = None
def get_full_data(self) -> bytes:
return self.stream.getvalue()
def _chunk_generator(self) -> Generator[bytes, None, None]:
self.stream.seek(0)
data = self.stream.read(self.chunk_size)
while data:
yield data
data = self.stream.read(self.chunk_size)
def next(self) -> Optional[bytes]:
if not self.iter:
self.iter = iter(self._chunk_generator())
try:
return next(self.iter)
except StopIteration:
return None
def _iter_remote(actor: ray.ActorID) -> Generator[bytes, None, None]:
"""Iterate over actor task and return as generator."""
while True:
buffer = ray.get(actor.next.remote())
if buffer is None:
return
yield buffer
def _unpack_dir(stream: io.BytesIO, target_dir: str, *, _retry: bool = True) -> None:
"""Unpack tarfile stream into target directory."""
stream.seek(0)
target_dir = os.path.normpath(target_dir)
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(f"{target_dir}.lock", timeout=0):
with tarfile.open(fileobj=stream) as tar:
tar.extractall(target_dir)
except TimeoutError:
# wait, but do not do anything
with TempFileLock(f"{target_dir}.lock"):
pass
# if the dir was locked due to being deleted,
# recreate
if not os.path.exists(target_dir):
if _retry:
_unpack_dir(stream, target_dir, _retry=False)
else:
raise RuntimeError(
f"Target directory {target_dir} does not exist "
"and couldn't be recreated. "
"Please raise an issue on GitHub: "
"https://github.com/ray-project/ray/issues"
)
@ray.remote
def _unpack_from_actor(pack_actor: ray.ActorID, target_dir: str) -> None:
"""Iterate over chunks received from pack actor and unpack."""
stream = io.BytesIO()
for buffer in _iter_remote(pack_actor):
stream.write(buffer)
_unpack_dir(stream, target_dir=target_dir)
def _copy_dir(
source_dir: str,
target_dir: str,
*,
exclude: Optional[List] = None,
_retry: bool = True,
) -> None:
"""Copy dir with shutil on the actor."""
target_dir = os.path.normpath(target_dir)
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(f"{target_dir}.lock", timeout=0):
_delete_path_unsafe(target_dir)
_ignore_func = None
if exclude:
def _ignore(path, names):
ignored_names = set()
rel_path = os.path.relpath(path, source_dir)
for name in names:
candidate = os.path.join(rel_path, name)
for excl in exclude:
if fnmatch.fnmatch(candidate, excl):
ignored_names.add(name)
break
return ignored_names
_ignore_func = _ignore
shutil.copytree(source_dir, target_dir, ignore=_ignore_func)
except TimeoutError:
# wait, but do not do anything
with TempFileLock(f"{target_dir}.lock"):
pass
# if the dir was locked due to being deleted,
# recreate
if not os.path.exists(target_dir):
if _retry:
_copy_dir(source_dir, target_dir, _retry=False)
else:
raise RuntimeError(
f"Target directory {target_dir} does not exist "
"and couldn't be recreated. "
"Please raise an issue on GitHub: "
"https://github.com/ray-project/ray/issues"
)
# Only export once
_remote_copy_dir = ray.remote(_copy_dir)
def _delete_path_unsafe(target_path: str):
"""Delete path (files and directories). No filelock."""
if os.path.exists(target_path):
if os.path.isdir(target_path):
shutil.rmtree(target_path)
else:
os.remove(target_path)
return True
return False
python/ray/tune/utils/log.py
+64
View File
@@ -0,0 +1,64 @@
import time
from enum import Enum
from typing import Dict, Tuple, Union
from ray.util import PublicAPI
from ray.util.annotations import DeveloperAPI
@PublicAPI
class Verbosity(Enum):
V0_MINIMAL = 0
V1_EXPERIMENT = 1
V2_TRIAL_NORM = 2
V3_TRIAL_DETAILS = 3
def __int__(self):
return self.value
verbosity: Union[int, Verbosity] = Verbosity.V3_TRIAL_DETAILS
@DeveloperAPI
def set_verbosity(level: Union[int, Verbosity]):
global verbosity
if isinstance(level, int):
verbosity = Verbosity(level)
else:
verbosity = level
@DeveloperAPI
def has_verbosity(level: Union[int, Verbosity]) -> bool:
"""Return True if passed level exceeds global verbosity level."""
global verbosity
log_level = int(level)
verbosity_level = int(verbosity)
return verbosity_level >= log_level
@DeveloperAPI
def disable_ipython():
"""Disable output of IPython HTML objects."""
try:
from IPython.core.interactiveshell import InteractiveShell
InteractiveShell.clear_instance()
except Exception:
pass
_log_cache_count: Dict[str, Tuple[str, float]] = {}
def _dedup_logs(domain: str, value: str, repeat_after_s: int = 5) -> bool:
cur_val, ts = _log_cache_count.get(domain, (None, None))
if value == cur_val and time.monotonic() - repeat_after_s < ts:
return False
else:
_log_cache_count[domain] = value, time.monotonic()
return True
python/ray/tune/utils/mock.py
+124
View File
@@ -0,0 +1,124 @@
import logging
import os
import random
import time
from collections import defaultdict
from pathlib import Path
from typing import Dict
from ray.tune.callback import Callback
from ray.tune.experiment import Trial
logger = logging.getLogger(__name__)
class FailureInjectorCallback(Callback):
"""Adds random failure injection to the TrialExecutor."""
def __init__(
self,
config_path="~/ray_bootstrap_config.yaml",
probability=0.1,
time_between_checks=0,
disable=False,
):
self.probability = probability
self.config_path = Path(config_path).expanduser().as_posix()
self.disable = disable
self.time_between_checks = time_between_checks
# Initialize with current time so we don't fail right away
self.last_fail_check = time.monotonic()
def on_step_begin(self, **info):
if not os.path.exists(self.config_path):
return
if time.monotonic() < self.last_fail_check + self.time_between_checks:
return
self.last_fail_check = time.monotonic()
import click
from ray.autoscaler._private.commands import kill_node
failures = 0
max_failures = 3
# With 10% probability inject failure to a worker.
if random.random() < self.probability and not self.disable:
# With 10% probability fully terminate the node.
should_terminate = random.random() < self.probability
while failures < max_failures:
try:
kill_node(
self.config_path,
yes=True,
hard=should_terminate,
override_cluster_name=None,
)
return
except click.exceptions.ClickException:
failures += 1
logger.exception(
"Killing random node failed in attempt "
"{}. "
"Retrying {} more times".format(
str(failures), str(max_failures - failures)
)
)
class TrialStatusSnapshot:
"""A sequence of statuses of trials as they progress.
If all trials keep previous status, no snapshot is taken.
"""
def __init__(self):
self._snapshot = []
def append(self, new_snapshot: Dict[str, str]):
"""May append a new snapshot to the sequence."""
if not new_snapshot:
# Don't add an empty snapshot.
return
if not self._snapshot or new_snapshot != self._snapshot[-1]:
self._snapshot.append(new_snapshot)
def max_running_trials(self) -> int:
"""Outputs the max number of running trials at a given time.
Usually used to assert certain number given resource restrictions.
"""
result = 0
for snapshot in self._snapshot:
count = 0
for trial_id in snapshot:
if snapshot[trial_id] == Trial.RUNNING:
count += 1
result = max(result, count)
return result
def all_trials_are_terminated(self) -> bool:
"""True if all trials are terminated."""
if not self._snapshot:
return False
last_snapshot = self._snapshot[-1]
return all(
last_snapshot[trial_id] == Trial.TERMINATED for trial_id in last_snapshot
)
class TrialStatusSnapshotTaker(Callback):
"""Collects a sequence of statuses of trials as they progress.
If all trials keep previous status, no snapshot is taken.
"""
def __init__(self, snapshot: TrialStatusSnapshot):
self._snapshot = snapshot
def on_step_end(self, iteration, trials, **kwargs):
new_snapshot = defaultdict(str)
for trial in trials:
new_snapshot[trial.trial_id] = trial.status
self._snapshot.append(new_snapshot)
python/ray/tune/utils/mock_trainable.py
+63
View File
@@ -0,0 +1,63 @@
import json
import os
import time
import numpy as np
from ray.tune import Trainable
MOCK_TRAINABLE_NAME = "mock_trainable"
MOCK_ERROR_KEY = "mock_error"
class MyTrainableClass(Trainable):
"""Example agent whose learning curve is a random sigmoid.
The dummy hyperparameters "width" and "height" determine the slope and
maximum reward value reached.
"""
def setup(self, config):
self._sleep_time = config.get("sleep", 0)
self._mock_error = config.get(MOCK_ERROR_KEY, False)
self._persistent_error = config.get("persistent_error", False)
self.timestep = 0
self.restored = False
def step(self):
if (
self._mock_error
and self.timestep > 0 # allow at least 1 successful checkpoint.
and (self._persistent_error or not self.restored)
):
raise RuntimeError(f"Failing on purpose! {self.timestep=}")
if self._sleep_time > 0:
time.sleep(self._sleep_time)
self.timestep += 1
v = np.tanh(float(self.timestep) / self.config.get("width", 1))
v *= self.config.get("height", 1)
# Here we use `episode_reward_mean`, but you can also report other
# objectives such as loss or accuracy.
return {"episode_reward_mean": v}
def save_checkpoint(self, checkpoint_dir):
path = os.path.join(checkpoint_dir, "checkpoint")
with open(path, "w") as f:
f.write(json.dumps({"timestep": self.timestep}))
def load_checkpoint(self, checkpoint_dir):
path = os.path.join(checkpoint_dir, "checkpoint")
with open(path, "r") as f:
self.timestep = json.loads(f.read())["timestep"]
self.restored = True
def register_mock_trainable():
from ray.tune import register_trainable
register_trainable(MOCK_TRAINABLE_NAME, MyTrainableClass)
python/ray/tune/utils/object_cache.py
+173
View File
@@ -0,0 +1,173 @@
from collections import Counter, defaultdict
from typing import Dict, Generator, List, Optional, TypeVar
# Grouping key - must be hashable
T = TypeVar("T")
# Objects to cache
U = TypeVar("U")
class _ObjectCache:
"""Cache up to some maximum count given a grouping key.
This object cache can e.g. be used to cache Ray Tune trainable actors
given their resource requirements (reuse_actors=True).
If the max number of cached objects for a grouping key is reached,
no more objects for this group will be cached.
However, if `may_keep_one=True`, one object (globally across all grouping
keys) may be cached, even if the max number of objects is 0. This is to
allow to cache an object if the max number of objects of this key
will increase shortly after (as is the case e.g. in the Ray Tune control
loop).
Args:
may_keep_one: If True, one object (globally) may be cached if no desired
maximum objects are defined.
"""
def __init__(self, may_keep_one: bool = True):
self._num_cached_objects: int = 0
self._cached_objects: Dict[T, List[U]] = defaultdict(list)
self._max_num_objects: Counter[T] = Counter()
self._may_keep_one = may_keep_one
@property
def num_cached_objects(self):
return self._num_cached_objects
@property
def total_max_objects(self):
# Counter.total() is only available for python 3.10+
return sum(self._max_num_objects.values())
def increase_max(self, key: T, by: int = 1) -> None:
"""Increase number of max objects for this key.
Args:
key: Group key.
by: Decrease by this amount.
"""
self._max_num_objects[key] += by
def decrease_max(self, key: T, by: int = 1) -> None:
"""Decrease number of max objects for this key.
Args:
key: Group key.
by: Decrease by this amount.
"""
self._max_num_objects[key] -= by
def has_cached_object(self, key: T) -> bool:
"""Return True if at least one cached object exists for this key.
Args:
key: Group key.
Returns:
True if at least one cached object exists for this key.
"""
return bool(self._cached_objects[key])
def cache_object(self, key: T, obj: U) -> bool:
"""Cache object for a given key.
This will put the object into a cache, assuming the number
of cached objects for this key is less than the number of
max objects for this key.
An exception is made if `max_keep_one=True` and no other
objects are cached globally. In that case, the object can
still be cached.
Args:
key: Group key.
obj: Object to cache.
Returns:
True if the object has been cached. False otherwise.
"""
# If we have more objects cached already than we desire
if len(self._cached_objects[key]) >= self._max_num_objects[key]:
# If may_keep_one is False, never cache
if not self._may_keep_one:
return False
# If we have more than one other cached object, don't cache
if self._num_cached_objects > 0:
return False
# If any other objects are expected to be cached, don't cache
if any(v for v in self._max_num_objects.values()):
return False
# Otherwise, cache (for now).
self._cached_objects[key].append(obj)
self._num_cached_objects += 1
return True
def pop_cached_object(self, key: T) -> Optional[U]:
"""Get one cached object for a key.
This will remove the object from the cache.
Args:
key: Group key.
Returns:
Cached object.
"""
if not self.has_cached_object(key):
return None
self._num_cached_objects -= 1
return self._cached_objects[key].pop(0)
def flush_cached_objects(self, force_all: bool = False) -> Generator[U, None, None]:
"""Return a generator over cached objects evicted from the cache.
This method yields all cached objects that should be evicted from the
cache for cleanup by the caller.
If the number of max objects is lower than the number of
cached objects for a given key, objects are evicted until
the numbers are equal.
If `max_keep_one=True` (and ``force_all=False``), one cached object
may be retained.
Objects are evicted FIFO.
If ``force_all=True``, all objects are evicted.
Args:
force_all: If True, all objects are flushed. This takes precedence
over ``keep_one``.
Yields:
U: Evicted objects to be cleaned up by caller.
"""
# If force_all=True, don't keep one.
keep_one = self._may_keep_one and not force_all
for key, objs in self._cached_objects.items():
max_cached = self._max_num_objects[key] if not force_all else 0
if (
self._num_cached_objects == 1
and keep_one
# Only keep this object if we don't expect a different one
and not any(v for v in self._max_num_objects.values())
):
break
while len(objs) > max_cached:
self._num_cached_objects -= 1
yield objs.pop(0)
python/ray/tune/utils/release_test_util.py
+190
View File
@@ -0,0 +1,190 @@
import json
import os
import pickle
import tempfile
import time
from collections import Counter
import numpy as np
from ray import tune
from ray._private.test_utils import safe_write_to_results_json
from ray.tune import Checkpoint
from ray.tune.callback import Callback
class ProgressCallback(Callback):
def __init__(self):
self.last_update = 0
self.update_interval = 60
def on_step_end(self, iteration, trials, **kwargs):
if time.time() - self.last_update > self.update_interval:
now = time.time()
result = {
"last_update": now,
"iteration": iteration,
"trial_states": dict(Counter([trial.status for trial in trials])),
}
safe_write_to_results_json(result, "/tmp/release_test_out.json")
self.last_update = now
class TestDurableTrainable(tune.Trainable):
def __init__(self, *args, **kwargs):
self.setup_env()
super(TestDurableTrainable, self).__init__(*args, **kwargs)
def setup_env(self):
pass
def setup(self, config):
self._num_iters = int(config["num_iters"])
self._sleep_time = config["sleep_time"]
self._score = config["score"]
self._checkpoint_iters = config["checkpoint_iters"]
self._checkpoint_size_b = config["checkpoint_size_b"]
self._checkpoint_num_items = self._checkpoint_size_b // 8 # np.float64
self._iter = 0
def step(self):
if self._iter > 0:
time.sleep(self._sleep_time)
res = dict(score=self._iter + self._score)
if self._iter >= self._num_iters:
res["done"] = True
self._iter += 1
return res
def save_checkpoint(self, tmp_checkpoint_dir):
checkpoint_file = os.path.join(tmp_checkpoint_dir, "bogus.ckpt")
checkpoint_data = np.random.uniform(0, 1, size=self._checkpoint_num_items)
with open(checkpoint_file, "wb") as fp:
pickle.dump(checkpoint_data, fp)
def load_checkpoint(self, checkpoint):
pass
def function_trainable(config):
num_iters = int(config["num_iters"])
sleep_time = config["sleep_time"]
score = config["score"]
checkpoint_iters = config["checkpoint_iters"]
checkpoint_size_b = config["checkpoint_size_b"]
checkpoint_num_items = checkpoint_size_b // 8 # np.float64
checkpoint_num_files = config["checkpoint_num_files"]
for i in range(num_iters):
metrics = {"score": i + score}
if (
checkpoint_iters >= 0
and checkpoint_size_b > 0
and i % checkpoint_iters == 0
):
with tempfile.TemporaryDirectory() as tmpdir:
for i in range(checkpoint_num_files):
checkpoint_file = os.path.join(tmpdir, f"bogus_{i}.ckpt")
checkpoint_data = np.random.uniform(0, 1, size=checkpoint_num_items)
with open(checkpoint_file, "wb") as fp:
pickle.dump(checkpoint_data, fp)
tune.report(metrics, checkpoint=Checkpoint.from_directory(tmpdir))
else:
tune.report(metrics)
time.sleep(sleep_time)
def timed_tune_run(
name: str,
num_samples: int,
results_per_second: int = 1,
trial_length_s: int = 1,
max_runtime: int = 300,
checkpoint_freq_s: int = -1,
checkpoint_size_b: int = 0,
checkpoint_num_files: int = 1,
**tune_kwargs,
) -> bool:
durable = (
"storage_path" in tune_kwargs
and tune_kwargs["storage_path"]
and (
tune_kwargs["storage_path"].startswith("s3://")
or tune_kwargs["storage_path"].startswith("gs://")
)
)
sleep_time = 1.0 / results_per_second
num_iters = int(trial_length_s / sleep_time)
checkpoint_iters = -1
if checkpoint_freq_s >= 0:
checkpoint_iters = int(checkpoint_freq_s / sleep_time)
config = {
"score": tune.uniform(0.0, 1.0),
"num_iters": num_iters,
"sleep_time": sleep_time,
"checkpoint_iters": checkpoint_iters,
"checkpoint_size_b": checkpoint_size_b,
"checkpoint_num_files": checkpoint_num_files,
}
print(f"Starting benchmark with config: {config}")
run_kwargs = {"reuse_actors": True, "verbose": 2}
run_kwargs.update(tune_kwargs)
_train = function_trainable
if durable:
_train = TestDurableTrainable
run_kwargs["checkpoint_freq"] = checkpoint_iters
start_time = time.monotonic()
analysis = tune.run(
_train,
config=config,
num_samples=num_samples,
raise_on_failed_trial=False,
**run_kwargs,
)
time_taken = time.monotonic() - start_time
result = {
"time_taken": time_taken,
"trial_states": dict(Counter([trial.status for trial in analysis.trials])),
"last_update": time.time(),
}
test_output_json = os.environ.get("TEST_OUTPUT_JSON", "/tmp/tune_test.json")
with open(test_output_json, "wt") as f:
json.dump(result, f)
success = time_taken <= max_runtime
if not success:
print(
f"The {name} test took {time_taken:.2f} seconds, but should not "
f"have exceeded {max_runtime:.2f} seconds. Test failed. \n\n"
f"--- FAILED: {name.upper()} ::: "
f"{time_taken:.2f} > {max_runtime:.2f} ---"
)
else:
print(
f"The {name} test took {time_taken:.2f} seconds, which "
f"is below the budget of {max_runtime:.2f} seconds. "
f"Test successful. \n\n"
f"--- PASSED: {name.upper()} ::: "
f"{time_taken:.2f} <= {max_runtime:.2f} ---"
)
return success
python/ray/tune/utils/resource_updater.py
+369
View File
@@ -0,0 +1,369 @@
import logging
import os
import time
from collections import namedtuple
from numbers import Number
from typing import Any, Dict, Optional
import ray
from ray._common.constants import NODE_ID_PREFIX
logger = logging.getLogger(__name__)
TUNE_STATE_REFRESH_PERIOD = 10 # Refresh resources every 10 s
def _to_gb(n_bytes):
return round(n_bytes / (1024**3), 2)
class _Resources(
namedtuple(
"_Resources",
[
"cpu",
"gpu",
"memory",
"object_store_memory",
"extra_cpu",
"extra_gpu",
"extra_memory",
"extra_object_store_memory",
"custom_resources",
"extra_custom_resources",
"has_placement_group",
],
)
):
"""Ray resources required to schedule a trial.
Parameters:
cpu: Number of CPUs to allocate to the trial.
gpu: Number of GPUs to allocate to the trial.
memory: Memory to reserve for the trial.
object_store_memory: Object store memory to reserve.
extra_cpu: Extra CPUs to reserve in case the trial needs to
launch additional Ray actors that use CPUs.
extra_gpu: Extra GPUs to reserve in case the trial needs to
launch additional Ray actors that use GPUs.
extra_memory: Memory to reserve for the trial launching
additional Ray actors that use memory.
extra_object_store_memory: Object store memory to reserve for
the trial launching additional Ray actors that use object store
memory.
custom_resources: Mapping of resource to quantity to allocate
to the trial.
extra_custom_resources: Extra custom resources to reserve in
case the trial needs to launch additional Ray actors that use
any of these custom resources.
has_placement_group: Bool indicating if the trial also
has an associated placement group.
"""
__slots__ = ()
def __new__(
cls,
cpu: float,
gpu: float,
memory: float = 0,
object_store_memory: float = 0.0,
extra_cpu: float = 0.0,
extra_gpu: float = 0.0,
extra_memory: float = 0.0,
extra_object_store_memory: float = 0.0,
custom_resources: Optional[dict] = None,
extra_custom_resources: Optional[dict] = None,
has_placement_group: bool = False,
):
custom_resources = custom_resources or {}
extra_custom_resources = extra_custom_resources or {}
leftovers = set(custom_resources) ^ set(extra_custom_resources)
for value in leftovers:
custom_resources.setdefault(value, 0)
extra_custom_resources.setdefault(value, 0)
cpu = round(cpu, 2)
gpu = round(gpu, 2)
memory = round(memory, 2)
object_store_memory = round(object_store_memory, 2)
extra_cpu = round(extra_cpu, 2)
extra_gpu = round(extra_gpu, 2)
extra_memory = round(extra_memory, 2)
extra_object_store_memory = round(extra_object_store_memory, 2)
custom_resources = {
resource: round(value, 2) for resource, value in custom_resources.items()
}
extra_custom_resources = {
resource: round(value, 2)
for resource, value in extra_custom_resources.items()
}
all_values = [
cpu,
gpu,
memory,
object_store_memory,
extra_cpu,
extra_gpu,
extra_memory,
extra_object_store_memory,
]
all_values += list(custom_resources.values())
all_values += list(extra_custom_resources.values())
assert len(custom_resources) == len(extra_custom_resources)
for entry in all_values:
assert isinstance(entry, Number), ("Improper resource value.", entry)
return super(_Resources, cls).__new__(
cls,
cpu,
gpu,
memory,
object_store_memory,
extra_cpu,
extra_gpu,
extra_memory,
extra_object_store_memory,
custom_resources,
extra_custom_resources,
has_placement_group,
)
def summary_string(self):
summary = "{} CPUs, {} GPUs".format(
self.cpu + self.extra_cpu, self.gpu + self.extra_gpu
)
if self.memory or self.extra_memory:
summary += ", {} GiB heap".format(
round((self.memory + self.extra_memory) / (1024**3), 2)
)
if self.object_store_memory or self.extra_object_store_memory:
summary += ", {} GiB objects".format(
round(
(self.object_store_memory + self.extra_object_store_memory)
/ (1024**3),
2,
)
)
custom_summary = ", ".join(
[
"{} {}".format(self.get_res_total(res), res)
for res in self.custom_resources
if not res.startswith(NODE_ID_PREFIX)
]
)
if custom_summary:
summary += " ({})".format(custom_summary)
return summary
def cpu_total(self):
return self.cpu + self.extra_cpu
def gpu_total(self):
return self.gpu + self.extra_gpu
def memory_total(self):
return self.memory + self.extra_memory
def object_store_memory_total(self):
return self.object_store_memory + self.extra_object_store_memory
def get_res_total(self, key):
return self.custom_resources.get(key, 0) + self.extra_custom_resources.get(
key, 0
)
def get(self, key):
return self.custom_resources.get(key, 0)
def is_nonnegative(self):
all_values = [self.cpu, self.gpu, self.extra_cpu, self.extra_gpu]
all_values += list(self.custom_resources.values())
all_values += list(self.extra_custom_resources.values())
return all(v >= 0 for v in all_values)
@classmethod
def subtract(cls, original, to_remove):
cpu = original.cpu - to_remove.cpu
gpu = original.gpu - to_remove.gpu
memory = original.memory - to_remove.memory
object_store_memory = (
original.object_store_memory - to_remove.object_store_memory
)
extra_cpu = original.extra_cpu - to_remove.extra_cpu
extra_gpu = original.extra_gpu - to_remove.extra_gpu
extra_memory = original.extra_memory - to_remove.extra_memory
extra_object_store_memory = (
original.extra_object_store_memory - to_remove.extra_object_store_memory
)
all_resources = set(original.custom_resources).union(
set(to_remove.custom_resources)
)
new_custom_res = {
k: original.custom_resources.get(k, 0)
- to_remove.custom_resources.get(k, 0)
for k in all_resources
}
extra_custom_res = {
k: original.extra_custom_resources.get(k, 0)
- to_remove.extra_custom_resources.get(k, 0)
for k in all_resources
}
return _Resources(
cpu,
gpu,
memory,
object_store_memory,
extra_cpu,
extra_gpu,
extra_memory,
extra_object_store_memory,
new_custom_res,
extra_custom_res,
)
class _ResourceUpdater:
"""Periodic Resource updater for Tune.
Initially, all resources are set to 0. The updater will try to update resources
when (1) init ResourceUpdater (2) call "update_avail_resources", "num_cpus"
or "num_gpus".
The update takes effect when (1) Ray is initialized (2) the interval between
this and last update is larger than "refresh_period"
"""
def __init__(self, refresh_period: Optional[float] = None):
self._avail_resources = _Resources(cpu=0, gpu=0)
if refresh_period is None:
refresh_period = float(
os.environ.get("TUNE_STATE_REFRESH_PERIOD", TUNE_STATE_REFRESH_PERIOD)
)
self._refresh_period = refresh_period
self._last_resource_refresh = float("-inf")
self.update_avail_resources()
def update_avail_resources(self, num_retries: int = 5, force: bool = False):
if not ray.is_initialized():
return
if (
time.time() - self._last_resource_refresh < self._refresh_period
and not force
):
return
logger.debug("Checking Ray cluster resources.")
resources = None
for i in range(num_retries):
if i > 0:
logger.warning(
f"Cluster resources not detected or are 0. Attempt #{i + 1}...",
)
time.sleep(0.5)
resources = ray.cluster_resources()
if resources:
break
if not resources:
# NOTE: This hides the possibility that Ray may be waiting for
# clients to connect.
resources.setdefault("CPU", 0)
resources.setdefault("GPU", 0)
logger.warning(
"Cluster resources cannot be detected or are 0. "
"You can resume this experiment by passing in `resume=True` to `run`."
)
resources = resources.copy()
num_cpus = resources.pop("CPU", 0)
num_gpus = resources.pop("GPU", 0)
memory = resources.pop("memory", 0)
object_store_memory = resources.pop("object_store_memory", 0)
custom_resources = resources
self._avail_resources = _Resources(
int(num_cpus),
int(num_gpus),
memory=int(memory),
object_store_memory=int(object_store_memory),
custom_resources=custom_resources,
)
self._last_resource_refresh = time.time()
def _get_used_avail_resources(self, total_allocated_resources: Dict[str, Any]):
total_allocated_resources = total_allocated_resources.copy()
used_cpu = total_allocated_resources.pop("CPU", 0)
total_cpu = self._avail_resources.cpu
used_gpu = total_allocated_resources.pop("GPU", 0)
total_gpu = self._avail_resources.gpu
custom_used_total = {
name: (
total_allocated_resources.get(name, 0.0),
self._avail_resources.get_res_total(name),
)
for name in self._avail_resources.custom_resources
if not name.startswith(NODE_ID_PREFIX)
and (total_allocated_resources.get(name, 0.0) > 0 or "_group_" not in name)
}
return used_cpu, total_cpu, used_gpu, total_gpu, custom_used_total
def debug_string(self, total_allocated_resources: Dict[str, Any]) -> str:
"""Returns a human readable message for printing to the console."""
if self._last_resource_refresh > 0:
(
used_cpu,
total_cpu,
used_gpu,
total_gpu,
custom_used_total,
) = self._get_used_avail_resources(total_allocated_resources)
if (
used_cpu > total_cpu
or used_gpu > total_gpu
or any(used > total for (used, total) in custom_used_total.values())
):
# If any of the used resources are higher than what we currently think
# is available, update our state and re-fetch
self.update_avail_resources(force=True)
(
used_cpu,
total_cpu,
used_gpu,
total_gpu,
custom_used_total,
) = self._get_used_avail_resources(total_allocated_resources)
status = (
f"Logical resource usage: {used_cpu}/{total_cpu} CPUs, "
f"{used_gpu}/{total_gpu} GPUs"
)
customs = ", ".join(
f"{used}/{total} {name}"
for name, (used, total) in custom_used_total.items()
)
if customs:
status += f" ({customs})"
return status
else:
return "Logical resource usage: ?"
def get_num_cpus(self) -> int:
self.update_avail_resources()
return self._avail_resources.cpu
def get_num_gpus(self) -> int:
self.update_avail_resources()
return self._avail_resources.gpu
def __reduce__(self):
# Do not need to serialize resources, because we can always
# update it again. This also prevents keeping outdated resources
# when deserialized.
return _ResourceUpdater, (self._refresh_period,)
python/ray/tune/utils/serialization.py
+96
View File
@@ -0,0 +1,96 @@
import json
import logging
import types
from ray import cloudpickle as cloudpickle
from ray._common.utils import binary_to_hex, hex_to_binary
from ray.util.annotations import DeveloperAPI
from ray.util.debug import log_once
logger = logging.getLogger(__name__)
# Marker used by TuneFunctionEncoder/Decoder to embed a cloudpickle blob
# (hex-encoded) inside an otherwise-JSON document for objects that JSON
# cannot represent (e.g. functions). Deserializing such a blob is equivalent
# to ``cloudpickle.loads`` -- i.e. arbitrary code execution -- and must only
# be done for inputs from a trusted source.
_CLOUDPICKLE_FALLBACK_TYPE = "CLOUDPICKLE_FALLBACK"
@DeveloperAPI
class TuneFunctionEncoder(json.JSONEncoder):
def default(self, obj):
if isinstance(obj, types.FunctionType):
return self._to_cloudpickle(obj)
try:
return super(TuneFunctionEncoder, self).default(obj)
except Exception:
if log_once(f"tune_func_encode:{str(obj)}"):
logger.debug("Unable to encode. Falling back to cloudpickle.")
return self._to_cloudpickle(obj)
def _to_cloudpickle(self, obj):
return {
"_type": _CLOUDPICKLE_FALLBACK_TYPE,
"value": binary_to_hex(cloudpickle.dumps(obj)),
}
@DeveloperAPI
class TuneFunctionDecoder(json.JSONDecoder):
"""JSON decoder that mirrors :class:`TuneFunctionEncoder`.
.. warning::
``TuneFunctionEncoder`` may embed a cloudpickle blob inside the JSON
output (under the ``CLOUDPICKLE_FALLBACK`` marker) for objects that
cannot be JSON-encoded. Deserializing such a blob runs
``cloudpickle.loads`` on attacker-controllable bytes, which is
equivalent to arbitrary code execution.
For that reason, this decoder **refuses** to expand
``CLOUDPICKLE_FALLBACK`` payloads by default and raises ``ValueError``
instead. Tune-internal callers that load state from a trusted source
(for example, a path that the same process just wrote) opt in via
``TuneFunctionDecoder(allow_cloudpickle=True)``.
"""
def __init__(self, *args, allow_cloudpickle: bool = False, **kwargs):
self._allow_cloudpickle = allow_cloudpickle
json.JSONDecoder.__init__(self, object_hook=self.object_hook, *args, **kwargs)
def object_hook(self, obj):
if obj.get("_type") == _CLOUDPICKLE_FALLBACK_TYPE:
if not self._allow_cloudpickle:
raise ValueError(
"Refusing to deserialize an embedded cloudpickle payload "
f"({_CLOUDPICKLE_FALLBACK_TYPE!r}) from JSON: this is "
"equivalent to executing arbitrary Python code from the "
"input. If the input comes from a trusted source, opt in "
"explicitly via `TuneFunctionDecoder(allow_cloudpickle=True)`."
)
return self._from_cloudpickle(obj)
return obj
def _from_cloudpickle(self, obj):
return cloudpickle.loads(hex_to_binary(obj["value"]))
def _loads_with_cloudpickle(s):
"""Decode a JSON document written by :class:`TuneFunctionEncoder`.
Accepts ``str`` or ``bytes``/``bytearray`` (decoded as UTF-8), mirroring
:func:`json.loads`.
Internal helper: opts in to expanding ``CLOUDPICKLE_FALLBACK`` payloads
embedded in the JSON document, which executes arbitrary code from those
payloads. Only call this on input the caller trusts (for example,
Tune-internal state written by the same process).
External callers should use ``json.loads(s, cls=TuneFunctionDecoder)``
instead, which raises ``ValueError`` on any embedded cloudpickle blob.
"""
if isinstance(s, (bytes, bytearray)):
s = s.decode("utf-8")
return TuneFunctionDecoder(allow_cloudpickle=True).decode(s)
python/ray/tune/utils/util.py
+660
View File
@@ -0,0 +1,660 @@
import copy
import glob
import inspect
import logging
import os
import threading
import time
import uuid
from collections import defaultdict
from datetime import datetime
from numbers import Number
from threading import Thread
from typing import Any, Callable, Dict, List, Optional, Sequence, Type, Union
import numpy as np
import ray
from ray._private.dict import ( # noqa: F401
deep_update,
flatten_dict,
merge_dicts,
unflatten_dict,
unflatten_list_dict,
unflattened_lookup,
)
from ray.air._internal.json import SafeFallbackEncoder # noqa
from ray.air._internal.util import is_nan, is_nan_or_inf # noqa: F401
from ray.util.annotations import DeveloperAPI, PublicAPI
import psutil
logger = logging.getLogger(__name__)
def _import_gputil():
try:
import GPUtil
except ImportError:
GPUtil = None
return GPUtil
START_OF_TIME = time.time()
@DeveloperAPI
class UtilMonitor(Thread):
"""Class for system usage utilization monitoring.
It keeps track of CPU, RAM, GPU, VRAM usage (each gpu separately) by
pinging for information every x seconds in a separate thread.
Requires psutil and GPUtil to be installed. Can be enabled with
Tuner(param_space={"log_sys_usage": True}).
"""
def __init__(self, start=True, delay=0.7):
self.stopped = True
GPUtil = _import_gputil()
self.GPUtil = GPUtil
if GPUtil is None and start:
logger.warning("Install gputil for GPU system monitoring.")
if psutil is None and start:
logger.warning("Install psutil to monitor system performance.")
if GPUtil is None and psutil is None:
return
super(UtilMonitor, self).__init__()
self.delay = delay # Time between calls to GPUtil
self.values = defaultdict(list)
self.lock = threading.Lock()
self.daemon = True
if start:
self.start()
def _read_utilization(self):
with self.lock:
if psutil is not None:
self.values["cpu_util_percent"].append(
float(psutil.cpu_percent(interval=None))
)
self.values["ram_util_percent"].append(
float(psutil.virtual_memory().percent)
)
if self.GPUtil is not None:
gpu_list = []
try:
gpu_list = self.GPUtil.getGPUs()
except Exception:
logger.debug("GPUtil failed to retrieve GPUs.")
for gpu in gpu_list:
self.values["gpu_util_percent" + str(gpu.id)].append(
float(gpu.load)
)
self.values["vram_util_percent" + str(gpu.id)].append(
float(gpu.memoryUtil)
)
def get_data(self):
if self.stopped:
return {}
with self.lock:
ret_values = copy.deepcopy(self.values)
for key, val in self.values.items():
del val[:]
return {"perf": {k: np.mean(v) for k, v in ret_values.items() if len(v) > 0}}
def run(self):
self.stopped = False
while not self.stopped:
self._read_utilization()
time.sleep(self.delay)
def stop(self):
self.stopped = True
@DeveloperAPI
def retry_fn(
fn: Callable[[], Any],
exception_type: Union[Type[Exception], Sequence[Type[Exception]]] = Exception,
num_retries: int = 3,
sleep_time: int = 1,
timeout: Optional[Number] = None,
) -> bool:
errored = threading.Event()
def _try_fn():
try:
fn()
except exception_type as e:
logger.warning(e)
errored.set()
for i in range(num_retries):
errored.clear()
proc = threading.Thread(target=_try_fn)
proc.daemon = True
proc.start()
proc.join(timeout=timeout)
if proc.is_alive():
logger.debug(
f"Process timed out (try {i+1}/{num_retries}): "
f"{getattr(fn, '__name__', None)}"
)
elif not errored.is_set():
return True
# Timed out, sleep and try again
time.sleep(sleep_time)
# Timed out, so return False
return False
@DeveloperAPI
class warn_if_slow:
"""Prints a warning if a given operation is slower than 500ms.
Example:
>>> from ray.tune.utils.util import warn_if_slow
>>> something = ... # doctest: +SKIP
>>> with warn_if_slow("some_operation"): # doctest: +SKIP
... ray.get(something) # doctest: +SKIP
"""
DEFAULT_THRESHOLD = float(os.environ.get("TUNE_WARN_THRESHOLD_S", 0.5))
DEFAULT_MESSAGE = (
"The `{name}` operation took {duration:.3f} s, "
"which may be a performance bottleneck."
)
def __init__(
self,
name: str,
threshold: Optional[float] = None,
message: Optional[str] = None,
disable: bool = False,
):
"""Initialize the context manager.
Args:
name: Identifier for the operation, used in the warning message.
threshold: Duration in seconds above which to warn. Defaults to
``DEFAULT_THRESHOLD``.
message: Optional override for the warning message format. Receives
``name`` and ``duration`` as format kwargs.
disable: If True, suppress warnings entirely.
"""
self.name = name
self.threshold = threshold or self.DEFAULT_THRESHOLD
self.message = message or self.DEFAULT_MESSAGE
self.too_slow = False
self.disable = disable
def __enter__(self):
self.start = time.time()
return self
def __exit__(self, type, value, traceback):
now = time.time()
if self.disable:
return
if now - self.start > self.threshold and now - START_OF_TIME > 60.0:
self.too_slow = True
duration = now - self.start
logger.warning(self.message.format(name=self.name, duration=duration))
@DeveloperAPI
class Tee(object):
def __init__(self, stream1, stream2):
self.stream1 = stream1
self.stream2 = stream2
# If True, we are currently handling a warning.
# We use this flag to avoid infinite recursion.
self._handling_warning = False
def _warn(self, op, s, args, kwargs):
# If we are already handling a warning, this is because
# `logger.warning` below triggered the same object again
# (e.g. because stderr is redirected to this object).
# In that case, exit early to avoid recursion.
if self._handling_warning:
return
msg = f"ValueError when calling '{op}' on stream ({s}). "
msg += f"args: {args} kwargs: {kwargs}"
self._handling_warning = True
logger.warning(msg)
self._handling_warning = False
def seek(self, *args, **kwargs):
for s in [self.stream1, self.stream2]:
try:
s.seek(*args, **kwargs)
except ValueError:
self._warn("seek", s, args, kwargs)
def write(self, *args, **kwargs):
for s in [self.stream1, self.stream2]:
try:
s.write(*args, **kwargs)
except ValueError:
self._warn("write", s, args, kwargs)
def flush(self, *args, **kwargs):
for s in [self.stream1, self.stream2]:
try:
s.flush(*args, **kwargs)
except ValueError:
self._warn("flush", s, args, kwargs)
@property
def encoding(self):
if hasattr(self.stream1, "encoding"):
return self.stream1.encoding
return self.stream2.encoding
@property
def error(self):
if hasattr(self.stream1, "error"):
return self.stream1.error
return self.stream2.error
@property
def newlines(self):
if hasattr(self.stream1, "newlines"):
return self.stream1.newlines
return self.stream2.newlines
def detach(self):
raise NotImplementedError
def read(self, *args, **kwargs):
raise NotImplementedError
def readline(self, *args, **kwargs):
raise NotImplementedError
def tell(self, *args, **kwargs):
raise NotImplementedError
@DeveloperAPI
def date_str():
return datetime.today().strftime("%Y-%m-%d_%H-%M-%S")
def _to_pinnable(obj):
"""Converts obj to a form that can be pinned in object store memory.
Currently only numpy arrays are pinned in memory, if you have a strong
reference to the array value.
"""
return (obj, np.zeros(1))
def _from_pinnable(obj):
"""Retrieve from _to_pinnable format."""
return obj[0]
@DeveloperAPI
def diagnose_serialization(trainable: Callable):
"""Utility for detecting why your trainable function isn't serializing.
Args:
trainable: The trainable object passed to
tune.Tuner(trainable). Currently only supports
Function API.
Returns:
bool | set of unserializable objects.
Example:
.. code-block:: python
import threading
# this is not serializable
e = threading.Event()
def test():
print(e)
diagnose_serialization(test)
# should help identify that 'e' should be moved into
# the `test` scope.
# correct implementation
def test():
e = threading.Event()
print(e)
assert diagnose_serialization(test) is True
"""
from ray.tune.registry import _check_serializability, register_trainable
def check_variables(objects, failure_set, printer):
for var_name, variable in objects.items():
msg = None
try:
_check_serializability(var_name, variable)
status = "PASSED"
except Exception as e:
status = "FAILED"
msg = f"{e.__class__.__name__}: {str(e)}"
failure_set.add(var_name)
printer(f"{str(variable)}[name='{var_name}'']... {status}")
if msg:
printer(msg)
print(f"Trying to serialize {trainable}...")
try:
register_trainable("__test:" + str(trainable), trainable, warn=False)
print("Serialization succeeded!")
return True
except Exception as e:
print(f"Serialization failed: {e}")
print(
"Inspecting the scope of the trainable by running "
f"`inspect.getclosurevars({str(trainable)})`..."
)
closure = inspect.getclosurevars(trainable)
failure_set = set()
if closure.globals:
print(
f"Detected {len(closure.globals)} global variables. "
"Checking serializability..."
)
check_variables(closure.globals, failure_set, lambda s: print(" " + s))
if closure.nonlocals:
print(
f"Detected {len(closure.nonlocals)} nonlocal variables. "
"Checking serializability..."
)
check_variables(closure.nonlocals, failure_set, lambda s: print(" " + s))
if not failure_set:
print(
"Nothing was found to have failed the diagnostic test, though "
"serialization did not succeed. Feel free to raise an "
"issue on github."
)
return failure_set
else:
print(
f"Variable(s) {failure_set} was found to be non-serializable. "
"Consider either removing the instantiation/imports "
"of these objects or moving them into the scope of "
"the trainable. "
)
return failure_set
def _atomic_save(state: Dict, checkpoint_dir: str, file_name: str, tmp_file_name: str):
"""Atomically saves the state object to the checkpoint directory.
This is automatically used by Tuner().fit during a Tune job.
Args:
state: Object state to be serialized.
checkpoint_dir: Directory location for the checkpoint.
file_name: Final name of file.
tmp_file_name: Temporary name of file. We prepend a .uuid- prefix.
"""
import ray.cloudpickle as cloudpickle
tmp_search_ckpt_path = os.path.join(
checkpoint_dir, f".{str(uuid.uuid4())}-{tmp_file_name}"
)
with open(tmp_search_ckpt_path, "wb") as f:
cloudpickle.dump(state, f)
os.replace(tmp_search_ckpt_path, os.path.join(checkpoint_dir, file_name))
def _load_newest_checkpoint(dirpath: str, ckpt_pattern: str) -> Optional[Dict]:
"""Returns the most recently modified checkpoint.
Assumes files are saved with an ordered name, most likely by
:obj:atomic_save.
Args:
dirpath: Directory in which to look for the checkpoint file.
ckpt_pattern: File name pattern to match to find checkpoint
files.
Returns:
(dict) Deserialized state dict.
"""
import ray.cloudpickle as cloudpickle
full_paths = glob.glob(os.path.join(dirpath, ckpt_pattern))
if not full_paths:
return
most_recent_checkpoint = max(full_paths)
with open(most_recent_checkpoint, "rb") as f:
checkpoint_state = cloudpickle.load(f)
return checkpoint_state
@PublicAPI(stability="beta")
def wait_for_gpu(
gpu_id: Optional[Union[int, str]] = None,
target_util: float = 0.01,
retry: int = 20,
delay_s: int = 5,
gpu_memory_limit: Optional[float] = None,
) -> bool:
"""Checks if a given GPU has freed memory.
Requires ``gputil`` to be installed: ``pip install gputil``.
Args:
gpu_id: GPU id or uuid to check.
Must be found within GPUtil.getGPUs(). If none, resorts to
the first item returned from `ray.get_gpu_ids()`.
target_util: The utilization threshold to reach to unblock.
Set this to 0 to block until the GPU is completely free.
retry: Number of times to check GPU limit. Sleeps `delay_s`
seconds between checks.
delay_s: Seconds to wait before check.
gpu_memory_limit: Deprecated. No longer used.
Returns:
bool: True if free.
Raises:
RuntimeError: If GPUtil is not found, if no GPUs are detected
or if the check fails.
Example:
.. code-block:: python
def tune_func(config):
tune.utils.wait_for_gpu()
train()
tuner = tune.Tuner(
tune.with_resources(
tune_func,
resources={"gpu": 1}
),
tune_config=tune.TuneConfig(num_samples=10)
)
tuner.fit()
"""
GPUtil = _import_gputil()
if GPUtil is None:
raise RuntimeError("GPUtil must be installed if calling `wait_for_gpu`.")
if gpu_id is None:
gpu_id_list = ray.get_gpu_ids()
if not gpu_id_list:
raise RuntimeError(
"No GPU ids found from `ray.get_gpu_ids()`. "
"Did you set Tune resources correctly?"
)
gpu_id = gpu_id_list[0]
gpu_attr = "id"
if isinstance(gpu_id, str):
if gpu_id.isdigit():
# GPU ID returned from `ray.get_gpu_ids()` is a str representation
# of the int GPU ID
gpu_id = int(gpu_id)
else:
# Could not coerce gpu_id to int, so assume UUID
# and compare against `uuid` attribute e.g.,
# 'GPU-04546190-b68d-65ac-101b-035f8faed77d'
gpu_attr = "uuid"
elif not isinstance(gpu_id, int):
raise ValueError(f"gpu_id ({type(gpu_id)}) must be type str/int.")
def gpu_id_fn(g):
# Returns either `g.id` or `g.uuid` depending on
# the format of the input `gpu_id`
return getattr(g, gpu_attr)
gpu_ids = {gpu_id_fn(g) for g in GPUtil.getGPUs()}
if gpu_id not in gpu_ids:
raise ValueError(
f"{gpu_id} not found in set of available GPUs: {gpu_ids}. "
"`wait_for_gpu` takes either GPU ordinal ID (e.g., '0') or "
"UUID (e.g., 'GPU-04546190-b68d-65ac-101b-035f8faed77d')."
)
for i in range(int(retry)):
gpu_object = next(g for g in GPUtil.getGPUs() if gpu_id_fn(g) == gpu_id)
if gpu_object.memoryUtil > target_util:
logger.info(
f"Waiting for GPU util to reach {target_util}. "
f"Util: {gpu_object.memoryUtil:0.3f}"
)
time.sleep(delay_s)
else:
return True
raise RuntimeError("GPU memory was not freed.")
@DeveloperAPI
def validate_save_restore(
trainable_cls: Type,
config: Optional[Dict] = None,
num_gpus: int = 0,
) -> bool:
"""Helper method to check if your Trainable class will resume correctly.
Args:
trainable_cls: Trainable class for evaluation.
config: Config to pass to Trainable when testing.
num_gpus: GPU resources to allocate when testing.
Returns:
True if the save/restore round-trip succeeded.
"""
assert ray.is_initialized(), "Need Ray to be initialized."
remote_cls = ray.remote(num_gpus=num_gpus)(trainable_cls)
trainable_1 = remote_cls.remote(config=config)
trainable_2 = remote_cls.remote(config=config)
from ray.air.constants import TRAINING_ITERATION
for _ in range(3):
res = ray.get(trainable_1.train.remote())
assert res.get(TRAINING_ITERATION), (
"Validation will not pass because it requires `training_iteration` "
"to be returned."
)
ray.get(trainable_2.restore.remote(trainable_1.save.remote()))
res = ray.get(trainable_2.train.remote())
assert res[TRAINING_ITERATION] == 4
res = ray.get(trainable_2.train.remote())
assert res[TRAINING_ITERATION] == 5
return True
def _detect_config_single(func):
"""Check if func({}) works."""
func_sig = inspect.signature(func)
use_config_single = True
try:
func_sig.bind({})
except Exception as e:
logger.debug(str(e))
use_config_single = False
return use_config_single
@PublicAPI()
def validate_warmstart(
parameter_names: List[str],
points_to_evaluate: List[Union[List, Dict]],
evaluated_rewards: List,
validate_point_name_lengths: bool = True,
):
"""Generic validation of a Searcher's warm start functionality.
Raises exceptions in case of type and length mismatches between
parameters.
If ``validate_point_name_lengths`` is False, the equality of lengths
between ``points_to_evaluate`` and ``parameter_names`` will not be
validated.
"""
if points_to_evaluate:
if not isinstance(points_to_evaluate, list):
raise TypeError(
"points_to_evaluate expected to be a list, got {}.".format(
type(points_to_evaluate)
)
)
for point in points_to_evaluate:
if not isinstance(point, (dict, list)):
raise TypeError(
f"points_to_evaluate expected to include list or dict, "
f"got {point}."
)
if validate_point_name_lengths and (not len(point) == len(parameter_names)):
raise ValueError(
"Dim of point {}".format(point)
+ " and parameter_names {}".format(parameter_names)
+ " do not match."
)
if points_to_evaluate and evaluated_rewards:
if not isinstance(evaluated_rewards, list):
raise TypeError(
"evaluated_rewards expected to be a list, got {}.".format(
type(evaluated_rewards)
)
)
if not len(evaluated_rewards) == len(points_to_evaluate):
raise ValueError(
"Dim of evaluated_rewards {}".format(evaluated_rewards)
+ " and points_to_evaluate {}".format(points_to_evaluate)
+ " do not match."
)