chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,369 @@
|
||||
import asyncio
|
||||
import contextlib
|
||||
import functools
|
||||
import logging
|
||||
import threading
|
||||
import time
|
||||
import traceback
|
||||
from datetime import datetime
|
||||
from enum import Enum
|
||||
from typing import (
|
||||
Any,
|
||||
Callable,
|
||||
ContextManager,
|
||||
Dict,
|
||||
Generator,
|
||||
Generic,
|
||||
Iterator,
|
||||
List,
|
||||
Optional,
|
||||
TypeVar,
|
||||
Union,
|
||||
)
|
||||
|
||||
import ray
|
||||
from ray.train._internal.utils import count_required_parameters
|
||||
from ray.train.v2._internal.exceptions import UserExceptionWithTraceback
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
def bundle_to_remote_args(bundle: dict) -> dict:
|
||||
"""Convert a bundle of resources to Ray actor/task arguments.
|
||||
|
||||
>>> bundle_to_remote_args({"GPU": 1, "memory": 1, "custom": 0.1})
|
||||
{'num_cpus': 0, 'num_gpus': 1, 'memory': 1, 'resources': {'custom': 0.1}}
|
||||
"""
|
||||
bundle = bundle.copy()
|
||||
args = {
|
||||
"num_cpus": bundle.pop("CPU", 0),
|
||||
"num_gpus": bundle.pop("GPU", 0),
|
||||
"memory": bundle.pop("memory", 0),
|
||||
}
|
||||
if bundle:
|
||||
args["resources"] = bundle
|
||||
return args
|
||||
|
||||
|
||||
def construct_train_func(
|
||||
train_func: Union[Callable[[], T], Callable[[Dict[str, Any]], T]],
|
||||
config: Optional[Dict[str, Any]],
|
||||
train_func_context: ContextManager,
|
||||
fn_arg_name: Optional[str] = "train_loop_per_worker",
|
||||
) -> Callable[[], T]:
|
||||
"""Validates and constructs the training function to execute.
|
||||
|
||||
Args:
|
||||
train_func: The training function to execute.
|
||||
This can either take in no arguments or a ``config`` dict.
|
||||
config: Configurations to pass into ``train_func``. If None then an
|
||||
empty Dict will be created.
|
||||
train_func_context: Context manager for user's `train_func`, which executes
|
||||
backend-specific logic before and after the training function.
|
||||
fn_arg_name: The name of training function to use for error messages.
|
||||
|
||||
Returns:
|
||||
A valid training function.
|
||||
|
||||
Raises:
|
||||
ValueError: if the input ``train_func`` is invalid.
|
||||
"""
|
||||
num_required_params = count_required_parameters(train_func)
|
||||
|
||||
if num_required_params > 1:
|
||||
err_msg = (
|
||||
f"{fn_arg_name} should take in 0 or 1 required arguments, but it accepts "
|
||||
f"{num_required_params} required arguments instead."
|
||||
)
|
||||
raise ValueError(err_msg)
|
||||
|
||||
if num_required_params == 1:
|
||||
config = config or {}
|
||||
|
||||
@functools.wraps(train_func)
|
||||
def train_fn():
|
||||
with train_func_context():
|
||||
return train_func(config)
|
||||
|
||||
else: # num_params == 0
|
||||
|
||||
@functools.wraps(train_func)
|
||||
def train_fn():
|
||||
with train_func_context():
|
||||
return train_func()
|
||||
|
||||
return train_fn
|
||||
|
||||
|
||||
class TrainingFramework(Enum):
|
||||
TORCH = "torch"
|
||||
JAX = "jax"
|
||||
TENSORFLOW = "tensorflow"
|
||||
XGBOOST = "xgboost"
|
||||
LIGHTGBM = "lightgbm"
|
||||
|
||||
def module_names(self) -> tuple[str, ...]:
|
||||
"""Returns the relevant module names for the training framework.
|
||||
|
||||
These module names are used by Train state version collection (see
|
||||
`_get_framework_version`) to gather versions of key framework-related packages.
|
||||
|
||||
Note: If adding a new module, make sure to use the module name rather than
|
||||
the distribution name. (e.g. sklearn instead of scikit-learn)
|
||||
"""
|
||||
if self is TrainingFramework.TORCH:
|
||||
return ("torch",)
|
||||
if self is TrainingFramework.JAX:
|
||||
return ("jax", "jaxlib")
|
||||
if self is TrainingFramework.TENSORFLOW:
|
||||
return ("tensorflow", "keras")
|
||||
if self is TrainingFramework.XGBOOST:
|
||||
return ("xgboost",)
|
||||
if self is TrainingFramework.LIGHTGBM:
|
||||
return ("lightgbm",)
|
||||
|
||||
return (self.value,)
|
||||
|
||||
|
||||
class ObjectRefWrapper(Generic[T]):
|
||||
"""Thin wrapper around ray.put to manually control dereferencing."""
|
||||
|
||||
def __init__(self, obj: T):
|
||||
self._ref = ray.put(obj)
|
||||
|
||||
def get(self) -> T:
|
||||
return ray.get(self._ref)
|
||||
|
||||
|
||||
def date_str(include_ms: bool = False):
|
||||
pattern = "%Y-%m-%d_%H-%M-%S"
|
||||
if include_ms:
|
||||
pattern += ".%f"
|
||||
return datetime.today().strftime(pattern)
|
||||
|
||||
|
||||
def time_monotonic():
|
||||
return time.monotonic()
|
||||
|
||||
|
||||
def _copy_doc(copy_func):
|
||||
def wrapped(func):
|
||||
func.__doc__ = copy_func.__doc__
|
||||
return func
|
||||
|
||||
return wrapped
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def invoke_context_managers(
|
||||
context_managers: List[ContextManager],
|
||||
) -> Generator[None, None, None]:
|
||||
"""
|
||||
Utility to invoke a list of context managers and yield sequentially.
|
||||
|
||||
Args:
|
||||
context_managers: List of context managers to invoke.
|
||||
"""
|
||||
with contextlib.ExitStack() as stack:
|
||||
for context_manager in context_managers:
|
||||
stack.enter_context(context_manager())
|
||||
yield
|
||||
|
||||
|
||||
def get_module_name(obj: object) -> str:
|
||||
"""Returns the full module name of the given object, including its qualified name.
|
||||
|
||||
Args:
|
||||
obj: The object (class, function, etc.) whose module name is required.
|
||||
|
||||
Returns:
|
||||
Full module and qualified name as a string.
|
||||
"""
|
||||
return f"{obj.__module__}.{obj.__qualname__}"
|
||||
|
||||
|
||||
def get_callable_name(fn: Callable) -> str:
|
||||
"""Returns a readable name for any callable.
|
||||
|
||||
Args:
|
||||
fn: The callable to extract a name from.
|
||||
|
||||
Returns:
|
||||
A human-readable name for the callable.
|
||||
|
||||
Examples:
|
||||
|
||||
>>> get_callable_name(lambda x: x)
|
||||
'<lambda>'
|
||||
>>> def foo(a, b): pass
|
||||
>>> get_callable_name(foo)
|
||||
'foo'
|
||||
>>> from functools import partial
|
||||
>>> bar = partial(partial(foo, a=1), b=2)
|
||||
>>> get_callable_name(bar)
|
||||
'foo'
|
||||
>>> class Dummy:
|
||||
... def __call__(self, a, b): pass
|
||||
>>> get_callable_name(Dummy())
|
||||
'Dummy'
|
||||
"""
|
||||
if isinstance(fn, functools.partial):
|
||||
return get_callable_name(fn.func)
|
||||
|
||||
# Use __name__ for regular functions and lambdas
|
||||
if hasattr(fn, "__name__"):
|
||||
return fn.__name__
|
||||
|
||||
# Fallback to the class name for objects that implement __call__
|
||||
return fn.__class__.__name__
|
||||
|
||||
|
||||
def construct_user_exception_with_traceback(
|
||||
e: BaseException, exclude_frames: int = 0
|
||||
) -> UserExceptionWithTraceback:
|
||||
"""Construct a UserExceptionWithTraceback from a base exception.
|
||||
|
||||
Args:
|
||||
e: The base exception to construct a UserExceptionWithTraceback from.
|
||||
exclude_frames: The number of frames to exclude from the beginnning of
|
||||
the traceback.
|
||||
|
||||
Returns:
|
||||
A UserExceptionWithTraceback object.
|
||||
"""
|
||||
# TODO(justinvyu): This is brittle and may break if the call stack
|
||||
# changes. Figure out a more robust way to exclude these frames.
|
||||
exc_traceback_str = traceback.format_exc(
|
||||
limit=-(len(traceback.extract_tb(e.__traceback__)) - exclude_frames)
|
||||
)
|
||||
logger.error(f"Error in training function:\n{exc_traceback_str}")
|
||||
return UserExceptionWithTraceback(e, traceback_str=exc_traceback_str)
|
||||
|
||||
|
||||
def _in_ray_train_worker() -> bool:
|
||||
"""Check if the current process is a Ray Train V2 worker."""
|
||||
from ray.train.v2._internal.execution.train_fn_utils import get_train_fn_utils
|
||||
|
||||
try:
|
||||
get_train_fn_utils()
|
||||
return True
|
||||
except RuntimeError:
|
||||
return False
|
||||
|
||||
|
||||
def requires_train_worker(raise_in_tune_session: bool = False) -> Callable:
|
||||
"""Check that the caller is a Ray Train worker spawned by Ray Train,
|
||||
with access to training function utilities.
|
||||
|
||||
Args:
|
||||
raise_in_tune_session: Whether to raise a specific error message if the caller
|
||||
is in a Tune session. If True, will raise a DeprecationWarning.
|
||||
|
||||
Returns:
|
||||
A decorator that performs this check, which raises an error if the caller
|
||||
is not a Ray Train worker.
|
||||
"""
|
||||
|
||||
def _wrap(fn: Callable) -> Callable:
|
||||
@functools.wraps(fn)
|
||||
def _wrapped_fn(*args, **kwargs):
|
||||
from ray.tune.trainable.trainable_fn_utils import _in_tune_session
|
||||
|
||||
if raise_in_tune_session and _in_tune_session():
|
||||
raise DeprecationWarning(
|
||||
f"`ray.train.{fn.__name__}` is deprecated when running in a function "
|
||||
"passed to Ray Tune. Please use the equivalent `ray.tune` API instead. "
|
||||
"See this issue for more context: "
|
||||
"https://github.com/ray-project/ray/issues/49454"
|
||||
)
|
||||
|
||||
if not _in_ray_train_worker():
|
||||
raise RuntimeError(
|
||||
f"`{fn.__name__}` cannot be used outside of a Ray Train training function. "
|
||||
"You are calling this API from the driver or another non-training process. "
|
||||
"These utilities are only available within a function launched by `trainer.fit()`."
|
||||
)
|
||||
return fn(*args, **kwargs)
|
||||
|
||||
return _wrapped_fn
|
||||
|
||||
return _wrap
|
||||
|
||||
|
||||
async def wait_with_logging(
|
||||
condition: asyncio.Condition,
|
||||
predicate: Optional[Callable[[], bool]] = None,
|
||||
generate_warning_message: Optional[Callable[[], str]] = None,
|
||||
warn_interval_s: float = 60,
|
||||
timeout_s: Optional[float] = None,
|
||||
):
|
||||
"""Waits for condition to be notified, logging warnings and eventually timing out.
|
||||
|
||||
You must acquire the condition before calling this function.
|
||||
|
||||
Args:
|
||||
condition: The condition to wait for.
|
||||
predicate: Wait until this predicate is True. If None, wait until the condition
|
||||
is notified.
|
||||
generate_warning_message: A function that generates the warning message to log.
|
||||
If None, no warning is logged.
|
||||
warn_interval_s: The interval in seconds to log a warning.
|
||||
timeout_s: The timeout in seconds. Defaults to``None`` to not time out.
|
||||
"""
|
||||
|
||||
async def _wait_loop():
|
||||
while True:
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
condition.wait()
|
||||
if predicate is None
|
||||
else condition.wait_for(predicate),
|
||||
timeout=warn_interval_s,
|
||||
)
|
||||
return
|
||||
# asyncio.wait_for() raises `asyncio.TimeoutError` for asyncio<=3.10
|
||||
# and raises `TimeoutError` for asyncio>=3.11
|
||||
# https://docs.python.org/3/library/asyncio-task.html#asyncio.wait_for
|
||||
except (asyncio.TimeoutError, TimeoutError):
|
||||
if generate_warning_message is not None:
|
||||
warning_message = generate_warning_message()
|
||||
logger.warning(warning_message)
|
||||
|
||||
await asyncio.wait_for(
|
||||
_wait_loop(),
|
||||
timeout=timeout_s,
|
||||
)
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def context_watchdog(fn: Callable, *args: Any) -> Iterator[None]:
|
||||
"""Run a function in a background thread for the duration of the context.
|
||||
|
||||
The function is started in a daemon thread on entry. On exit, a
|
||||
threading.Event is set to signal the thread to stop. The function is
|
||||
responsible for checking the event and returning promptly once it is set.
|
||||
|
||||
Args:
|
||||
fn: A function whose first argument is a threading.Event stop signal.
|
||||
The function should return when stop_event.is_set() or
|
||||
stop_event.wait(...) returns True.
|
||||
*args: Additional arguments forwarded to fn after the stop event.
|
||||
|
||||
Yields:
|
||||
None: Control is yielded to the caller while the watchdog thread runs.
|
||||
"""
|
||||
stop_event = threading.Event()
|
||||
thread = threading.Thread(
|
||||
target=fn,
|
||||
args=(stop_event, *args),
|
||||
daemon=True, # thread will end even if the finally is bypassed by an abnormal exit
|
||||
)
|
||||
thread.start()
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
stop_event.set()
|
||||
thread.join()
|
||||
Reference in New Issue
Block a user