wehub-resource-sync
222 lines
7.6 KiB
Python
222 lines
7.6 KiB
Python
import os
|
|
import threading
|
|
from contextlib import contextmanager
|
|
from functools import wraps
|
|
from typing import Any, Callable, Optional, Tuple, TypeVar, cast, overload
|
|
|
|
from ray._private.auto_init_hook import auto_init_ray
|
|
|
|
F = TypeVar("F", bound=Callable[..., Any])
|
|
|
|
# Attr set on func defs to mark they have been converted to client mode.
|
|
RAY_CLIENT_MODE_ATTR = "__ray_client_mode_key__"
|
|
|
|
# Global setting of whether client mode is enabled. This default to OFF,
|
|
# but is enabled upon ray.client(...).connect() or in tests.
|
|
is_client_mode_enabled = os.environ.get("RAY_CLIENT_MODE", "0") == "1"
|
|
|
|
# When RAY_CLIENT_MODE == 1, we treat it as default enabled client mode
|
|
# This is useful for testing
|
|
is_client_mode_enabled_by_default = is_client_mode_enabled
|
|
os.environ.update({"RAY_CLIENT_MODE": "0"})
|
|
|
|
is_init_called = False
|
|
|
|
# Local setting of whether to ignore client hook conversion. This defaults
|
|
# to TRUE and is disabled when the underlying 'real' Ray function is needed.
|
|
_client_hook_status_on_thread = threading.local()
|
|
_client_hook_status_on_thread.status = True
|
|
|
|
|
|
def _get_client_hook_status_on_thread():
|
|
"""Get's the value of `_client_hook_status_on_thread`.
|
|
Since `_client_hook_status_on_thread` is a thread-local variable, we may
|
|
need to add and set the 'status' attribute.
|
|
"""
|
|
global _client_hook_status_on_thread
|
|
if not hasattr(_client_hook_status_on_thread, "status"):
|
|
_client_hook_status_on_thread.status = True
|
|
return _client_hook_status_on_thread.status
|
|
|
|
|
|
def _set_client_hook_status(val: bool):
|
|
global _client_hook_status_on_thread
|
|
_client_hook_status_on_thread.status = val
|
|
|
|
|
|
def _disable_client_hook():
|
|
global _client_hook_status_on_thread
|
|
out = _get_client_hook_status_on_thread()
|
|
_client_hook_status_on_thread.status = False
|
|
return out
|
|
|
|
|
|
def _explicitly_enable_client_mode():
|
|
"""Force client mode to be enabled.
|
|
NOTE: This should not be used in tests, use `enable_client_mode`.
|
|
"""
|
|
global is_client_mode_enabled
|
|
is_client_mode_enabled = True
|
|
|
|
|
|
def _explicitly_disable_client_mode():
|
|
global is_client_mode_enabled
|
|
is_client_mode_enabled = False
|
|
|
|
|
|
@contextmanager
|
|
def disable_client_hook():
|
|
val = _disable_client_hook()
|
|
try:
|
|
yield None
|
|
finally:
|
|
_set_client_hook_status(val)
|
|
|
|
|
|
@contextmanager
|
|
def enable_client_mode():
|
|
_explicitly_enable_client_mode()
|
|
try:
|
|
yield None
|
|
finally:
|
|
_explicitly_disable_client_mode()
|
|
|
|
|
|
@overload
|
|
def client_mode_hook(func: F) -> F:
|
|
...
|
|
|
|
|
|
@overload
|
|
def client_mode_hook(*, local_only_kwargs: Tuple[str, ...] = ()) -> Callable[[F], F]:
|
|
...
|
|
|
|
|
|
def client_mode_hook(
|
|
func: Optional[F] = None, *, local_only_kwargs: Tuple[str, ...] = ()
|
|
):
|
|
"""Decorator for whether to use the 'regular' ray version of a function,
|
|
or the Ray Client version of that function.
|
|
|
|
Args:
|
|
func: This function. This is set when this function is used
|
|
as a bare decorator.
|
|
local_only_kwargs: Names of keyword arguments that apply only to the
|
|
local (non-client) implementation. They are stripped before the
|
|
call is redirected to the Ray Client, so the client API does not
|
|
need to accept them. On the local path they pass through unchanged.
|
|
as a decorator.
|
|
|
|
Returns:
|
|
The wrapped function that dispatches to the regular or client version.
|
|
"""
|
|
|
|
def decorator(func: F) -> F:
|
|
from ray.util.client import ray
|
|
|
|
@wraps(func)
|
|
def wrapper(*args, **kwargs):
|
|
# NOTE(hchen): DO NOT use "import" inside this function.
|
|
# Because when it's called within a `__del__` method, this error
|
|
# will be raised (see #35114):
|
|
# ImportError: sys.meta_path is None, Python is likely shutting down.
|
|
if client_mode_should_convert():
|
|
# Legacy code
|
|
# we only convert init function if RAY_CLIENT_MODE=1
|
|
if func.__name__ != "init" or is_client_mode_enabled_by_default:
|
|
if local_only_kwargs:
|
|
kwargs = {
|
|
k: v
|
|
for k, v in kwargs.items()
|
|
if k not in local_only_kwargs
|
|
}
|
|
return getattr(ray, func.__name__)(*args, **kwargs)
|
|
return func(*args, **kwargs)
|
|
|
|
return cast(F, wrapper)
|
|
|
|
# Support both `@client_mode_hook` and
|
|
# `@client_mode_hook(local_only_kwargs=...)` usage.
|
|
if func is not None:
|
|
return decorator(func)
|
|
return decorator
|
|
|
|
|
|
def client_mode_should_convert():
|
|
"""Determines if functions should be converted to client mode."""
|
|
|
|
# `is_client_mode_enabled_by_default` is used for testing with
|
|
# `RAY_CLIENT_MODE=1`. This flag means all tests run with client mode.
|
|
return (
|
|
is_client_mode_enabled or is_client_mode_enabled_by_default
|
|
) and _get_client_hook_status_on_thread()
|
|
|
|
|
|
def client_mode_wrap(func):
|
|
"""Wraps a function called during client mode for execution as a remote
|
|
task.
|
|
|
|
Can be used to implement public features of ray client which do not
|
|
belong in the main ray API (`ray.*`), yet require server-side execution.
|
|
An example is the creation of placement groups:
|
|
`ray.util.placement_group.placement_group()`. When called on the client
|
|
side, this function is wrapped in a task to facilitate interaction with
|
|
the GCS.
|
|
"""
|
|
|
|
@wraps(func)
|
|
def wrapper(*args, **kwargs):
|
|
from ray.util.client import ray
|
|
|
|
auto_init_ray()
|
|
# Directly pass this through since `client_mode_wrap` is for
|
|
# Placement Group APIs
|
|
if client_mode_should_convert():
|
|
f = ray.remote(num_cpus=0)(func)
|
|
ref = f.remote(*args, **kwargs)
|
|
return ray.get(ref)
|
|
return func(*args, **kwargs)
|
|
|
|
return wrapper
|
|
|
|
|
|
def client_mode_convert_function(func_cls, in_args, in_kwargs, **kwargs):
|
|
"""Runs a preregistered ray RemoteFunction through the ray client.
|
|
|
|
The common case for this is to transparently convert that RemoteFunction
|
|
to a ClientRemoteFunction. This happens in circumstances where the
|
|
RemoteFunction is declared early, in a library and only then is Ray used in
|
|
client mode -- necessitating a conversion.
|
|
"""
|
|
from ray.util.client import ray
|
|
|
|
key = getattr(func_cls, RAY_CLIENT_MODE_ATTR, None)
|
|
|
|
# Second part of "or" is needed in case func_cls is reused between Ray
|
|
# client sessions in one Python interpreter session.
|
|
if (key is None) or (not ray._converted_key_exists(key)):
|
|
key = ray._convert_function(func_cls)
|
|
setattr(func_cls, RAY_CLIENT_MODE_ATTR, key)
|
|
client_func = ray._get_converted(key)
|
|
return client_func._remote(in_args, in_kwargs, **kwargs)
|
|
|
|
|
|
def client_mode_convert_actor(actor_cls, in_args, in_kwargs, **kwargs):
|
|
"""Runs a preregistered actor class on the ray client
|
|
|
|
The common case for this decorator is for instantiating an ActorClass
|
|
transparently as a ClientActorClass. This happens in circumstances where
|
|
the ActorClass is declared early, in a library and only then is Ray used in
|
|
client mode -- necessitating a conversion.
|
|
"""
|
|
from ray.util.client import ray
|
|
|
|
key = getattr(actor_cls, RAY_CLIENT_MODE_ATTR, None)
|
|
# Second part of "or" is needed in case actor_cls is reused between Ray
|
|
# client sessions in one Python interpreter session.
|
|
if (key is None) or (not ray._converted_key_exists(key)):
|
|
key = ray._convert_actor(actor_cls)
|
|
setattr(actor_cls, RAY_CLIENT_MODE_ATTR, key)
|
|
client_actor = ray._get_converted(key)
|
|
return client_actor._remote(in_args, in_kwargs, **kwargs)
|