wehub-resource-sync
81 lines
3.1 KiB
Python
81 lines
3.1 KiB
Python
from typing import Any, Dict, Hashable, List
|
|
|
|
import ray
|
|
|
|
CACHED_FUNCTIONS = {}
|
|
|
|
|
|
def cached_remote_fn(fn: Any, **ray_remote_args) -> Any:
|
|
"""Lazily defines a ray.remote function.
|
|
|
|
This is used in Datasets to avoid circular import issues with ray.remote.
|
|
(ray imports ray.data in order to allow ``ray.data.read_foo()`` to work,
|
|
which means ray.remote cannot be used top-level in ray.data).
|
|
|
|
NOTE: Dynamic arguments should not be passed in directly,
|
|
and should be set with ``options`` instead:
|
|
``cached_remote_fn(fn, **static_args).options(**dynamic_args)``.
|
|
"""
|
|
|
|
# NOTE: Hash of the passed in arguments guarantees that we're caching
|
|
# complete instantiation of the Ray's remote method
|
|
#
|
|
# To compute the hash of passed in arguments and make sure it's deterministic
|
|
# - Sort all KV-pairs by the keys
|
|
# - Convert sorted list into tuple
|
|
# - Compute hash of the resulting tuple
|
|
hashable_args = _make_hashable(ray_remote_args)
|
|
args_hash = hash(hashable_args)
|
|
|
|
if (fn, args_hash) not in CACHED_FUNCTIONS:
|
|
default_ray_remote_args = {
|
|
# Use the default scheduling strategy for all tasks so that we will
|
|
# not inherit a placement group from the caller, if there is one.
|
|
# The caller of this function may override the scheduling strategy
|
|
# as needed.
|
|
"scheduling_strategy": "DEFAULT",
|
|
"max_retries": -1,
|
|
}
|
|
ray_remote_args = {**default_ray_remote_args, **ray_remote_args}
|
|
_add_system_error_to_retry_exceptions(ray_remote_args)
|
|
|
|
CACHED_FUNCTIONS[(fn, args_hash)] = ray.remote(**ray_remote_args)(fn)
|
|
|
|
return CACHED_FUNCTIONS[(fn, args_hash)]
|
|
|
|
|
|
def _make_hashable(obj):
|
|
if isinstance(obj, (List, tuple)):
|
|
return tuple([_make_hashable(o) for o in obj])
|
|
elif isinstance(obj, Dict):
|
|
converted = [(_make_hashable(k), _make_hashable(v)) for k, v in obj.items()]
|
|
return tuple(sorted(converted, key=lambda t: t[0]))
|
|
elif isinstance(obj, Hashable):
|
|
return obj
|
|
else:
|
|
raise ValueError(f"Type {type(obj)} is not hashable")
|
|
|
|
|
|
def _add_system_error_to_retry_exceptions(ray_remote_args) -> None:
|
|
"""Modify the remote args so that Ray retries `RaySystemError`s.
|
|
|
|
Ray typically automatically retries system errors. However, in some cases, Ray won't
|
|
retry system errors if they're raised from task code. To ensure that Ray Data is
|
|
fault tolerant to those errors, we need to add `RaySystemError` to the
|
|
`retry_exceptions` list.
|
|
|
|
TODO: Fix this in Ray Core. See https://github.com/ray-project/ray/pull/45079.
|
|
"""
|
|
retry_exceptions = ray_remote_args.get("retry_exceptions", False)
|
|
assert isinstance(retry_exceptions, (list, bool))
|
|
|
|
if (
|
|
isinstance(retry_exceptions, list)
|
|
and ray.exceptions.RaySystemError not in retry_exceptions
|
|
):
|
|
retry_exceptions.append(ray.exceptions.RaySystemError)
|
|
elif not retry_exceptions:
|
|
retry_exceptions = [ray.exceptions.RaySystemError]
|
|
|
|
ray_remote_args["retry_exceptions"] = retry_exceptions
|