wehub/ray-project--ray
1
0
Fork 0
Code Issues Pull Requests Actions 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/util/client/ARCHITECTURE.md
+133
View File
@@ -0,0 +1,133 @@
# Ray Client Architecture Guide
A quick development primer on the codebase layout
## General
The Ray client is a gRPC client and server.
The server runs `ray.init()` and acts like a normal Ray driver, controlled by the gRPC connection.
It does all the bookkeeping and keeps things in scope for the clients that connect.
Generally, the client side lives in `ray/util/client` and the server lives in `ray/util/client/server`.
By convention, the `ray/util/client` avoids importing `ray` directly, but the server side, being just another Ray application, is allowed to do so.
This separation exists both for dependency cycle reasons and also to, if desired in the future, pull either portion out into its own repo or sub-installation.
(eg, `pip install ray_client`)
The `ray` global variable of type `RayAPIStub` in [`ray/util/client/__init__.py`](./__init__.py) acts as the equivalent API surface as does the `ray` package.
Functions in the `ray` namespace are methods on the RayAPIStub object.
For many of the objects in the root `ray` namespace, there is an equivalent client object. These are mostly contained in [`ray/util/client/common.py`](./common.py).
These objects are client stand-ins for their server-side objects. For example:
```
ObjectRef <-> ClientObjectRef
ActorID <-> ClientActorRef
RemoteFunc <-> ClientRemoteFunc
```
This means that, if the type of the object you're looking at (say, in a bug report) is a ClientObjectRef, it should have come from the client code (ie, constructed by returning from the server).
How the two interchange is talked about under Protocol.
## Protocol
There's one gRPC spec for the client, and that lives at [`/src/ray/protobuf/ray_client.proto`](/src/ray/protobuf/ray_client.proto).
There's another protocol at play, however, and that is _the way in which functions and data are encoded_.
This is particularly important in the context of the Ray client; since both ends are Python, this is a `pickle`.
The key separation to understand is that `pickle` is how client objects, including the client-side stubs, get serialized and transported, opaquely, to the server.
The gRPC service is the API surface to implement remote, thin, client functionality.
The `pickle` protocol is how the data is encoded for that API.
### gRPC services
The gRPC side is the most straightforward and easiest to follow from the proto file.
#### get, put, and function calls
Client started life as a set of unary RPCs, with just enough functionality to implement the most-used APIs.
As an introduction to the RPC API, they're a good place to start to understand how the protocol works.
The proto file is well-commented with every field describing what it does.
The Unary RPCs are still around, but they are ripe for deprecation.
The problem they have is that they are not tied to a persistent connection.
If you imagine a load balancer in front of a couple client-servers, then any client could hit any state on any server with a unary RPC.
As we need to keep handles to ray ObjectRefs and similar so that they don't go out of scope and dropped, these must stay in sync for connected clients.
With unary RPCs, that means one RPC could go to one server (say, a `x = f.remote()`), and the follow up (`ray.get(x)`) wouldn't have the corresponding ObjectRef on the other server.
Get, Put, and Wait are pretty standard.
The more interesting one is Schedule, which implies a Put before it (the function to execute) and then executes it.
#### Data Channel
Which brings us to the data channel.
The data channel is a bidirectional streaming connection for the client.
It wraps all the same Request/Response patterns as the Unary RPCs.
At the start, the client associates itself with the server with a UUID-generated ClientID.
As long as the channel is open, the client is connected.
Tracking the ClientID then allows us to track all the resources we're holding for a particular client, and we know if the client has disconnected (the channel drops).
It's also through this mechanism we can do reference counting.
The client can keep track of how many references it has to various Ray client objects, and, if they fall out of scope, can send a `ReleaseRequest` to the server to optimistically clean up after itself.
Otherwise, all reference counting is done on the client side and the server only needs to know when to clean up.
The server can also clean up all references held for a client whenever it thinks it's safe to do so.
In the future, having an explicit "ClientDisconnection" message may help here, to delineate between a client that's intentionally done and will never come back and one that's experiencing a connectivity issue.
#### Logs Channel
Similar to the data channel, there's an associated logs channel which will pipe logs back to the client.
It's a separate channel as it's ancillary to the continued connection of the client, and if some logs are dropped due to disconnection, that's generally okay.
It's also then a separate way for a log aggregator to connect without implementing the full API.
This is also a bidirectional stream, where the client sends messages to control the verbosity and type of content, and the server streams back all the logs as they come in.
#### On CloudPickle
As per the introduction to this section, `pickle` and `cloudpickle` are the way to encode to Python-executable data for the generic transport of gRPC.
Ray Client provides its own pickle/unpickle subclasses in `client_pickler.py` and `server/server_pickler.py`.
The reason it has its own subclasses is to solve the problem of mixing `Client*` stub objects.
This is easier to describe by example.
Suppose a `RemoteFunc`, `f()` calls another `RemoteFunc`, `g()`
`f()` has to have a reference to `g()`, and knows it's a `RemoteFunc` (calls `g.remote()`, say) and so the function `f()` gets serialized with pickle, and in that serialization data is an object of class `RemoteFunc` to be deserialized on the worker side.
In Ray client, `f()` and `g()` are `ClientRemoteFunc`s and work the same way.
In early versions of the Ray Client, a `ClientRemoteFunc` had to know whether it was on the server or client side, and either run like a normal `RemoteFunc` (on the server) or a client call on the client.
This led to some interesting bugs where passing, or especially returning, client-side stub objects around.
(Imagine if `f()` calls `g()` which builds and returns a new closure `h()`)
To simplify all of this, the pickler subclasses were written.
Now, whenever a Client-stub-object is serialized, a struct (as of writing, a tuple) is stored in its place, and when deserialized, the server "fills in" the appropriate non-stub object.
And vice versa -- if the server is encoding a return/response that is an `ObjectRef`, a tuple is passed on the wire instead and the deserializer on the client side turns it back into a `ClientObjectRef`
This means that the client side deals as much as possible in its stub objects, and the server side never sees a stub object, and there is a clean separation between the two. Now, a `ClientObjectRef` existing on the server is an error case, and not a case to be handled specially.
It also means that the server side works just like normal Ray and deals in the normal Ray objects and can encode them transparently into client objects as they are sent.
Because never the twain shall meet, it's much easier to model and debug what's going on.
## Integration points with Ray core
In order to provide a seamless client experience with Ray core, we need to wrap some of the core Ray functions (eg, `ray.get()`).
Python's dynamic nature helps us here. As mentioned, the `RayAPIStub` is a class, not a module, which is a subtle difference.
This also allows us to have a `__getattr__` on the API level to redirect whereever we'd like, which we can't do in modules ([at least until Python 3.6 is deprecated](https://www.python.org/dev/peps/pep-0562/))
If the `ray` core were an object instead of functions in a namespace, we wouldn't need to wrap them to integrate, we'd simply swap the implementation.
But we have backwards compatibility to maintain.
All the interesting integration points with Ray core live within `ray/_private/client_mode_hook.py`.
In that file are contextmanagers and decorators meant to wrap Ray core functions.
If `client_mode_should_convert()` returns `True`, based on the environment variables as they've been set, then the decorators spring into action, and forward the calls to the `ray/util/client` object.
## Testing
There are two primary ways to test client code.
The first is to approach it from the context of knowing that we're testing a client/server app.
We can run both ends of the connection, call the client side (that we know to be the client side), and see the effect on the server and vice-versa.
The set of tests of the form `test_client*.py` take this approach.
The other way to approach them is as a fixture where we're testing the API of Ray, with Ray's own tests, _as though it were normal Ray_.
It's a highly powerful pattern, in that a test passing there means that the user can feel confident that things work the way they always have, client or not.
It's also a more difficult pattern to implement, as hooking the setup and shutdown of a client/server pair and a single-node ray instance are different, especially when these tests may make assumptions about how the fixtures were set up in the first place.
It is, however, the only way to test the integration points.
So generally speaking, if it's implementing a feature that makes the client work, it probably should be a test in the `test_client` series where one controls both ends and tests the client code.
If it's fixing a user-side API bug or an integration with Ray core, it's probably adapting or including a pre-existing unit test as part of Ray Client.
python/ray/util/client/__init__.py
+424
View File
@@ -0,0 +1,424 @@
import logging
import os
import threading
from typing import Any, Dict, List, Optional, Tuple
import ray._private.ray_constants as ray_constants
from ray._common.network_utils import build_address, get_localhost_ip
from ray._private.client_mode_hook import (
_explicitly_disable_client_mode,
_explicitly_enable_client_mode,
)
from ray._private.ray_logging import setup_logger
from ray._private.utils import check_version_info
from ray.job_config import JobConfig
from ray.util.annotations import DeveloperAPI
logger = logging.getLogger(__name__)
def _apply_uv_hook_for_client(
runtime_env: Optional[Dict[str, Any]],
) -> Optional[Dict[str, Any]]:
"""Apply UV runtime env hook on client side before connection.
UV (https://docs.astral.sh/uv/) is a modern Python package manager that
manages dependencies via pyproject.toml and uv.lock files. This function
detects when the client is running under 'uv run' and automatically
propagates the UV configuration to cluster workers so they can install
the same dependencies.
How it works:
1. Detects 'uv run' in the parent process tree
2. Extracts UV command-line arguments (e.g., --python, --locked)
3. Sets py_executable to 'uv run [args]' in runtime_env
4. Workers will use this UV command to install dependencies
Precedence rules:
- If user provides py_executable, UV hook is skipped entirely to avoid
unintended side effects (e.g., auto-setting working_dir)
- User-provided working_dir is preserved when UV hook runs
- Other runtime_env settings are merged with UV config
Feature flag:
Controlled by RAY_ENABLE_UV_RUN_RUNTIME_ENV constant (default: enabled)
Args:
runtime_env: The runtime environment dict to potentially modify.
Can be None if no runtime_env was specified.
Returns:
Modified runtime_env dict with UV configuration if detected,
otherwise the original runtime_env unchanged. Returns None if
input was None.
Raises:
RuntimeError: If UV environment is detected but configuration is invalid
(e.g., pyproject.toml not in working_dir, conflicting runtime_env).
Validation errors fail fast to provide clear feedback.
Note:
ImportError and other environmental errors are caught and logged,
allowing connection to proceed without UV propagation.
Example:
Client running under: uv run --python 3.11 my_script.py
>>> runtime_env = {"working_dir": "/tmp/myapp"}
>>> result = _apply_uv_hook_for_client(runtime_env)
>>> result
{'working_dir': '/tmp/myapp', 'py_executable': 'uv run --python 3.11'}
See Also:
- Issue: https://github.com/ray-project/ray/issues/57991
- UV docs: https://docs.astral.sh/uv/
"""
if not ray_constants.RAY_ENABLE_UV_RUN_RUNTIME_ENV:
return runtime_env
# If user provided py_executable, skip UV hook entirely to avoid side effects
# (e.g., auto-setting working_dir which triggers unwanted directory upload)
if runtime_env and "py_executable" in runtime_env:
logger.debug(
"User-provided py_executable found, skipping UV hook to avoid "
"unintended runtime_env modifications"
)
return runtime_env
# Import hook here (not at module level) to:
# 1. Avoid circular import issues with ray._private modules
# 2. Only load UV hook code when feature flag is enabled
from ray._private.runtime_env.uv_runtime_env_hook import hook
try:
result = hook(runtime_env)
except Exception as e:
raise RuntimeError(
f"Failed to apply UV runtime env hook for Ray Client: {e} "
"If you want the driver to use UV without propagating to workers, "
"set RAY_ENABLE_UV_RUN_RUNTIME_ENV=0."
) from e
if "py_executable" in result:
# UV environment was detected and applied by the hook
logger.debug(
f"UV environment detected for Ray Client: "
f"py_executable={result['py_executable']}"
)
return result
return runtime_env
class _ClientContext:
def __init__(self):
from ray.util.client.api import _ClientAPI
self.api = _ClientAPI()
self.client_worker = None
self._server = None
self._connected_with_init = False
self._inside_client_test = False
def connect(
self,
conn_str: str,
job_config: JobConfig = None,
secure: bool = False,
metadata: List[Tuple[str, str]] = None,
connection_retries: int = 3,
namespace: str = None,
*,
ignore_version: bool = False,
_credentials: Optional["grpc.ChannelCredentials"] = None, # noqa: F821
ray_init_kwargs: Optional[Dict[str, Any]] = None,
) -> Dict[str, Any]:
"""Connect the Ray Client to a server.
Args:
conn_str: Connection string, in the form "[host]:port"
job_config: The job config of the server.
secure: Whether to use a TLS secured gRPC channel
metadata: gRPC metadata to send on connect
connection_retries: number of connection attempts to make
namespace: The namespace to connect to.
ignore_version: whether to ignore Python or Ray version mismatches.
This should only be used for debugging purposes.
_credentials: Optional gRPC channel credentials for secure connection.
ray_init_kwargs: Optional additional keyword arguments for ray.init().
Returns:
Dictionary of connection info, e.g., {"num_clients": 1}.
"""
# Delay imports until connect to avoid circular imports.
from ray.util.client.worker import Worker
if self.client_worker is not None:
if self._connected_with_init:
return
raise Exception("ray.init() called, but ray client is already connected")
if not self._inside_client_test:
# If we're calling a client connect specifically and we're not
# currently in client mode, ensure we are.
_explicitly_enable_client_mode()
if namespace is not None:
job_config = job_config or JobConfig()
job_config.set_ray_namespace(namespace)
logging_level = ray_constants.LOGGER_LEVEL
logging_format = ray_constants.LOGGER_FORMAT
if ray_init_kwargs is None:
ray_init_kwargs = {}
# Apply UV hook client-side before connection.
# UV detection must happen on client side where 'uv run' process exists.
# See: https://github.com/ray-project/ray/issues/57991
#
# Runtime env can come from two sources:
# 1. ray_init_kwargs["runtime_env"] - directly passed to connect()
# 2. job_config.runtime_env - passed via JobConfig object
# We need to handle both sources and update them appropriately after UV hook.
runtime_env = ray_init_kwargs.get("runtime_env")
if runtime_env is None and job_config and job_config.runtime_env is not None:
runtime_env = job_config.runtime_env
runtime_env = _apply_uv_hook_for_client(runtime_env)
if runtime_env is not None:
# Update both ray_init_kwargs and job_config with UV modifications.
# This is necessary because _server_init() reads runtime_env from
# job_config.runtime_env, not from ray_init_kwargs["runtime_env"].
ray_init_kwargs["runtime_env"] = runtime_env
if job_config:
job_config.set_runtime_env(runtime_env)
# NOTE(architkulkarni): Custom env_hook is not supported with Ray Client.
# However, UV hook is now applied client-side above.
ray_init_kwargs["_skip_env_hook"] = True
if ray_init_kwargs.get("logging_level") is not None:
logging_level = ray_init_kwargs["logging_level"]
if ray_init_kwargs.get("logging_format") is not None:
logging_format = ray_init_kwargs["logging_format"]
setup_logger(logging_level, logging_format)
try:
self.client_worker = Worker(
conn_str,
secure=secure,
_credentials=_credentials,
metadata=metadata,
connection_retries=connection_retries,
)
self.api.worker = self.client_worker
self.client_worker._server_init(job_config, ray_init_kwargs)
conn_info = self.client_worker.connection_info()
self._check_versions(conn_info, ignore_version)
self._register_serializers()
return conn_info
except Exception:
self.disconnect()
raise
def _register_serializers(self):
"""Register the custom serializer addons at the client side.
The server side should have already registered the serializers via
regular worker's serialization_context mechanism.
"""
import ray.util.serialization_addons
from ray.util.serialization import StandaloneSerializationContext
ctx = StandaloneSerializationContext()
ray.util.serialization_addons.apply(ctx)
def _check_versions(self, conn_info: Dict[str, Any], ignore_version: bool) -> None:
# conn_info has "python_version" and "ray_version" so it can be used to compare.
ignore_version = ignore_version or ("RAY_IGNORE_VERSION_MISMATCH" in os.environ)
check_version_info(
conn_info,
"Ray Client",
raise_on_mismatch=not ignore_version,
python_version_match_level="minor",
)
def disconnect(self):
"""Disconnect the Ray Client."""
from ray.util.client.api import _ClientAPI
if self.client_worker is not None:
self.client_worker.close()
self.api = _ClientAPI()
self.client_worker = None
# remote can be called outside of a connection, which is why it
# exists on the same API layer as connect() itself.
def remote(self, *args: Any, **kwargs: Any):
"""remote is the hook stub passed on to replace `ray.remote`.
This sets up remote functions or actors, as the decorator,
but does not execute them.
Args:
*args: opaque arguments forwarded to ``_ClientAPI.remote``.
**kwargs: opaque keyword arguments forwarded to ``_ClientAPI.remote``.
Returns:
A client-side stub for the remote function or actor, or a
decorator that produces one when applied.
"""
return self.api.remote(*args, **kwargs)
def __getattr__(self, key: str):
if self.is_connected():
return getattr(self.api, key)
elif key in ["is_initialized", "_internal_kv_initialized"]:
# Client is not connected, thus Ray is not considered initialized.
return lambda: False
else:
raise Exception(
"Ray Client is not connected. Please connect by calling `ray.init`."
)
def is_connected(self) -> bool:
if self.client_worker is None:
return False
return self.client_worker.is_connected()
def init(self, *args, **kwargs):
if self._server is not None:
raise Exception("Trying to start two instances of ray via client")
import ray.util.client.server.server as ray_client_server
server_handle, address_info = ray_client_server.init_and_serve(
get_localhost_ip(), 50051, *args, **kwargs
)
self._server = server_handle.grpc_server
self.connect(build_address(get_localhost_ip(), 50051))
self._connected_with_init = True
return address_info
def shutdown(self, _exiting_interpreter=False):
self.disconnect()
import ray.util.client.server.server as ray_client_server
if self._server is None:
return
ray_client_server.shutdown_with_server(self._server, _exiting_interpreter)
self._server = None
# All connected context will be put here
# This struct will be guarded by a lock for thread safety
_all_contexts = set()
_lock = threading.Lock()
# This is the default context which is used when allow_multiple is not True
_default_context = _ClientContext()
@DeveloperAPI
class RayAPIStub:
"""This class stands in as the replacement API for the `import ray` module.
Much like the ray module, this mostly delegates the work to the
_client_worker. As parts of the ray API are covered, they are piped through
here or on the client worker API.
"""
def __init__(self):
self._cxt = threading.local()
self._cxt.handler = _default_context
self._inside_client_test = False
def get_context(self):
try:
return self._cxt.__getattribute__("handler")
except AttributeError:
self._cxt.handler = _default_context
return self._cxt.handler
def set_context(self, cxt):
old_cxt = self.get_context()
if cxt is None:
self._cxt.handler = _ClientContext()
else:
self._cxt.handler = cxt
return old_cxt
def is_default(self):
return self.get_context() == _default_context
def connect(self, *args, **kw_args):
self.get_context()._inside_client_test = self._inside_client_test
conn = self.get_context().connect(*args, **kw_args)
global _lock, _all_contexts
with _lock:
_all_contexts.add(self._cxt.handler)
return conn
def disconnect(self, *args, **kw_args):
global _lock, _all_contexts, _default_context
with _lock:
if _default_context == self.get_context():
for cxt in _all_contexts:
cxt.disconnect(*args, **kw_args)
_all_contexts = set()
else:
self.get_context().disconnect(*args, **kw_args)
if self.get_context() in _all_contexts:
_all_contexts.remove(self.get_context())
if len(_all_contexts) == 0:
_explicitly_disable_client_mode()
def remote(self, *args, **kwargs):
return self.get_context().remote(*args, **kwargs)
def __getattr__(self, name):
return self.get_context().__getattr__(name)
def is_connected(self, *args, **kwargs):
return self.get_context().is_connected(*args, **kwargs)
def init(self, *args, **kwargs):
ret = self.get_context().init(*args, **kwargs)
global _lock, _all_contexts
with _lock:
_all_contexts.add(self._cxt.handler)
return ret
def shutdown(self, *args, **kwargs):
global _lock, _all_contexts
with _lock:
if _default_context == self.get_context():
for cxt in _all_contexts:
cxt.shutdown(*args, **kwargs)
_all_contexts = set()
else:
self.get_context().shutdown(*args, **kwargs)
if self.get_context() in _all_contexts:
_all_contexts.remove(self.get_context())
if len(_all_contexts) == 0:
_explicitly_disable_client_mode()
ray = RayAPIStub()
@DeveloperAPI
def num_connected_contexts():
"""Return the number of client connections active."""
global _lock, _all_contexts
with _lock:
return len(_all_contexts)
# Someday we might add methods in this module so that someone who
# tries to `import ray_client as ray` -- as a module, instead of
# `from ray_client import ray` -- as the API stub
# still gets expected functionality. This is the way the ray package
# worked in the past.
#
# This really calls for PEP 562: https://www.python.org/dev/peps/pep-0562/
# But until Python 3.6 is EOL, here we are.
python/ray/util/client/api.py
+460
View File
@@ -0,0 +1,460 @@
"""This file defines the interface between the ray client worker
and the overall ray module API.
"""
import json
import logging
from concurrent.futures import Future
from typing import TYPE_CHECKING, Any, Callable, List, Optional, Union
from ray._common import ray_option_utils
from ray.util.client.runtime_context import _ClientWorkerPropertyAPI
if TYPE_CHECKING:
from ray.actor import ActorClass
from ray.core.generated.ray_client_pb2 import DataResponse
from ray.remote_function import RemoteFunction
from ray.util.client.common import ClientActorHandle, ClientObjectRef, ClientStub
logger = logging.getLogger(__name__)
def _as_bytes(value):
if isinstance(value, str):
return value.encode("utf-8")
return value
class _ClientAPI:
"""The Client-side methods corresponding to the ray API. Delegates
to the Client Worker that contains the connection to the ClientServer.
"""
def __init__(self, worker=None):
self.worker = worker
def get(
self,
vals: Union["ClientObjectRef", List["ClientObjectRef"]],
*,
timeout: Optional[float] = None,
):
"""get is the hook stub passed on to replace `ray.get`
Args:
vals: [Client]ObjectRef or list of these refs to retrieve.
timeout: Optional timeout in seconds
Returns:
The Python object(s) corresponding to ``vals``.
"""
return self.worker.get(vals, timeout=timeout)
def put(self, *args: Any, **kwargs: Any):
"""put is the hook stub passed on to replace `ray.put`
Args:
*args: opaque arguments forwarded to the worker's ``put``.
**kwargs: opaque keyword arguments forwarded to the worker's ``put``.
Returns:
A ``ClientObjectRef`` for the stored value.
"""
return self.worker.put(*args, **kwargs)
def wait(self, *args: Any, **kwargs: Any):
"""wait is the hook stub passed on to replace `ray.wait`
Args:
*args: opaque arguments forwarded to the worker's ``wait``.
**kwargs: opaque keyword arguments forwarded to the worker's ``wait``.
Returns:
A tuple ``(ready, remaining)`` of object refs, mirroring ``ray.wait``.
"""
return self.worker.wait(*args, **kwargs)
def _wait_generators_bulk(self, *args, **kwargs):
raise RuntimeError(
"ray._private.worker._wait_generators_bulk is not supported on Ray Client. "
"Connect with ray.init(address=...) instead, or use ray.wait."
)
def remote(self, *args: Any, **kwargs: Any):
"""remote is the hook stub passed on to replace `ray.remote`.
This sets up remote functions or actors, as the decorator,
but does not execute them.
Args:
*args: opaque arguments; when used as ``@ray.remote`` with no
parentheses, contains the wrapped function or class.
**kwargs: opaque keyword arguments; the options forwarded to
``ray.remote(...)``.
Returns:
A client-side stub for the remote function or actor, or a
decorator that produces one when applied.
"""
# Delayed import to avoid a cyclic import
from ray.util.client.common import remote_decorator
if len(args) == 1 and len(kwargs) == 0 and callable(args[0]):
# This is the case where the decorator is just @ray.remote.
return remote_decorator(options=None)(args[0])
assert (
len(args) == 0 and len(kwargs) > 0
), ray_option_utils.remote_args_error_string
return remote_decorator(options=kwargs)
# TODO(mwtian): consider adding _internal_ prefix to call_remote /
# call_release / call_retain.
def call_remote(
self, instance: "ClientStub", *args: Any, **kwargs: Any
) -> List[Future]:
"""call_remote is called by stub objects to execute them remotely.
This is used by stub objects in situations where they're called
with .remote, eg, `f.remote()` or `actor_cls.remote()`.
This allows the client stub objects to delegate execution to be
implemented in the most effective way whether it's in the client,
clientserver, or raylet worker.
Args:
instance: The Client-side stub reference to a remote object
*args: opaque arguments forwarded to the remote invocation.
**kwargs: opaque keyword arguments forwarded to the remote invocation.
Returns:
A list of futures, one per return value of the remote call.
"""
return self.worker.call_remote(instance, *args, **kwargs)
def call_release(self, id: bytes) -> None:
"""Attempts to release an object reference.
When client references are destructed, they release their reference,
which can opportunistically send a notification through the datachannel
to release the reference being held for that object on the server.
Args:
id: The id of the reference to release on the server side.
"""
return self.worker.call_release(id)
def call_retain(self, id: bytes) -> None:
"""Attempts to retain a client object reference.
Increments the reference count on the client side, to prevent
the client worker from attempting to release the server reference.
Args:
id: The id of the reference to retain on the client side.
"""
return self.worker.call_retain(id)
def close(self) -> None:
"""close cleans up an API connection by closing any channels or
shutting down any servers gracefully.
"""
return self.worker.close()
def get_actor(
self, name: str, namespace: Optional[str] = None
) -> "ClientActorHandle":
"""Returns a handle to an actor by name.
Args:
name: The name passed to this actor by
Actor.options(name="name").remote()
namespace: The namespace the named actor was created in.
Defaults to the current namespace.
Returns:
A ``ClientActorHandle`` for the named actor.
"""
return self.worker.get_actor(name, namespace)
def list_named_actors(self, all_namespaces: bool = False) -> List[str]:
"""List all named actors in the system.
Actors must have been created with Actor.options(name="name").remote().
This works for both detached & non-detached actors.
By default, only actors in the current namespace will be returned
and the returned entries will simply be their name.
If `all_namespaces` is set to True, all actors in the cluster will be
returned regardless of namespace, and the retunred entries will be of
the form '<namespace>/<name>'.
"""
return self.worker.list_named_actors(all_namespaces)
def kill(self, actor: "ClientActorHandle", *, no_restart: bool = True):
"""kill forcibly stops an actor running in the cluster
Args:
actor: The client-side handle of the actor to kill.
no_restart: Whether this actor should be restarted if it's a
restartable actor.
Returns:
The result of the underlying ``terminate_actor`` call.
"""
return self.worker.terminate_actor(actor, no_restart)
def cancel(
self,
obj: "ClientObjectRef",
*,
force: bool = False,
recursive: bool = True,
):
"""Cancels a task on the cluster.
If the specified task is pending execution, it will not be executed. If
the task is currently executing, the behavior depends on the ``force``
flag, as per `ray.cancel()`
Only non-actor tasks can be canceled. Canceled tasks will not be
retried (max_retries will not be respected).
Args:
obj: ObjectRef returned by the task that should be canceled.
force: Whether to force-kill a running task by killing
the worker that is running the task.
recursive: Whether to try to cancel tasks submitted by
the task specified.
Returns:
The result of the underlying ``terminate_task`` call.
"""
return self.worker.terminate_task(obj, force, recursive)
# Various metadata methods for the client that are defined in the protocol.
def is_initialized(self) -> bool:
"""True if our client is connected, and if the server is initialized.
Returns:
A boolean determining if the client is connected and
server initialized.
"""
return self.worker.is_initialized()
def nodes(self):
"""Get a list of the nodes in the cluster (for debugging only).
Returns:
Information about the Ray clients in the cluster.
"""
# This should be imported here, otherwise, it will error doc build.
import ray.core.generated.ray_client_pb2 as ray_client_pb2
return self.worker.get_cluster_info(ray_client_pb2.ClusterInfoType.NODES)
def method(self, *args: Any, **kwargs: Any):
"""Annotate an actor method
Args:
*args: Positional arguments are not supported; ``@ray.method``
must be invoked with at least one keyword argument.
**kwargs: Supported keyword arguments are ``num_returns`` (the
number of object refs that should be returned by invocations
of this actor method) and ``concurrency_group``.
Returns:
A decorator that annotates an actor method with the supplied
options.
"""
# NOTE: So this follows the same logic as in ray/actor.py::method()
# The reason to duplicate it here is to simplify the client mode
# redirection logic. As the annotated method gets pickled and sent to
# the server from the client it carries this private variable, it
# activates the same logic on the server side; so there's no need to
# pass anything else. It's inside the class definition that becomes an
# actor. Similar annotations would follow the same way.
valid_kwargs = ["num_returns", "concurrency_group"]
error_string = (
"The @ray.method decorator must be applied using at least one of "
f"the arguments in the list {valid_kwargs}, for example "
"'@ray.method(num_returns=2)'."
)
assert len(args) == 0 and len(kwargs) > 0, error_string
for key in kwargs:
key_error_string = (
f'Unexpected keyword argument to @ray.method: "{key}". The '
f"supported keyword arguments are {valid_kwargs}"
)
assert key in valid_kwargs, key_error_string
def annotate_method(method):
if "num_returns" in kwargs:
method.__ray_num_returns__ = kwargs["num_returns"]
if "concurrency_group" in kwargs:
method.__ray_concurrency_group__ = kwargs["concurrency_group"]
return method
return annotate_method
def cluster_resources(self):
"""Get the current total cluster resources.
Note that this information can grow stale as nodes are added to or
removed from the cluster.
Returns:
A dictionary mapping resource name to the total quantity of that
resource in the cluster.
"""
# This should be imported here, otherwise, it will error doc build.
import ray.core.generated.ray_client_pb2 as ray_client_pb2
return self.worker.get_cluster_info(
ray_client_pb2.ClusterInfoType.CLUSTER_RESOURCES
)
def available_resources(self):
"""Get the current available cluster resources.
This is different from `cluster_resources` in that this will return
idle (available) resources rather than total resources.
Note that this information can grow stale as tasks start and finish.
Returns:
A dictionary mapping resource name to the total quantity of that
resource in the cluster.
"""
# This should be imported here, otherwise, it will error doc build.
import ray.core.generated.ray_client_pb2 as ray_client_pb2
return self.worker.get_cluster_info(
ray_client_pb2.ClusterInfoType.AVAILABLE_RESOURCES
)
def get_runtime_context(self):
"""Return a Ray RuntimeContext describing the state on the server
Returns:
A RuntimeContext wrapping a client making get_cluster_info calls.
"""
return _ClientWorkerPropertyAPI(self.worker).build_runtime_context()
# Client process isn't assigned any GPUs.
def get_gpu_ids(self) -> list:
return []
def timeline(self, filename: Optional[str] = None) -> Optional[List[Any]]:
logger.warning(
"Timeline will include events from other clients using this server."
)
# This should be imported here, otherwise, it will error doc build.
import ray.core.generated.ray_client_pb2 as ray_client_pb2
all_events = self.worker.get_cluster_info(
ray_client_pb2.ClusterInfoType.TIMELINE
)
if filename is not None:
with open(filename, "w") as outfile:
json.dump(all_events, outfile)
else:
return all_events
def _internal_kv_initialized(self) -> bool:
"""Hook for internal_kv._internal_kv_initialized."""
# NOTE(edoakes): the kv is always initialized because we initialize it
# manually in the proxier with a GCS client if Ray hasn't been
# initialized yet.
return True
def _internal_kv_exists(
self, key: Union[str, bytes], *, namespace: Optional[Union[str, bytes]] = None
) -> bool:
"""Hook for internal_kv._internal_kv_exists."""
return self.worker.internal_kv_exists(
_as_bytes(key), namespace=_as_bytes(namespace)
)
def _internal_kv_get(
self, key: Union[str, bytes], *, namespace: Optional[Union[str, bytes]] = None
) -> bytes:
"""Hook for internal_kv._internal_kv_get."""
return self.worker.internal_kv_get(
_as_bytes(key), namespace=_as_bytes(namespace)
)
def _internal_kv_put(
self,
key: Union[str, bytes],
value: Union[str, bytes],
overwrite: bool = True,
*,
namespace: Optional[Union[str, bytes]] = None,
) -> bool:
"""Hook for internal_kv._internal_kv_put."""
return self.worker.internal_kv_put(
_as_bytes(key), _as_bytes(value), overwrite, namespace=_as_bytes(namespace)
)
def _internal_kv_del(
self,
key: Union[str, bytes],
*,
del_by_prefix: bool = False,
namespace: Optional[Union[str, bytes]] = None,
) -> int:
"""Hook for internal_kv._internal_kv_del."""
return self.worker.internal_kv_del(
_as_bytes(key), del_by_prefix=del_by_prefix, namespace=_as_bytes(namespace)
)
def _internal_kv_list(
self,
prefix: Union[str, bytes],
*,
namespace: Optional[Union[str, bytes]] = None,
) -> List[bytes]:
"""Hook for internal_kv._internal_kv_list."""
return self.worker.internal_kv_list(
_as_bytes(prefix), namespace=_as_bytes(namespace)
)
def _pin_runtime_env_uri(self, uri: str, expiration_s: int) -> None:
"""Hook for internal_kv._pin_runtime_env_uri."""
return self.worker.pin_runtime_env_uri(uri, expiration_s)
def _convert_actor(self, actor: "ActorClass") -> str:
"""Register a ClientActorClass for the ActorClass and return a UUID"""
return self.worker._convert_actor(actor)
def _convert_function(self, func: "RemoteFunction") -> str:
"""Register a ClientRemoteFunc for the ActorClass and return a UUID"""
return self.worker._convert_function(func)
def _get_converted(self, key: str) -> "ClientStub":
"""Given a UUID, return the converted object"""
return self.worker._get_converted(key)
def _converted_key_exists(self, key: str) -> bool:
"""Check if a key UUID is present in the store of converted objects."""
return self.worker._converted_key_exists(key)
def __getattr__(self, key: str):
if not key.startswith("_"):
raise NotImplementedError(
"Not available in Ray client: `ray.{}`. This method is only "
"available within Ray remote functions and is not yet "
"implemented in the client API.".format(key)
)
return self.__getattribute__(key)
def _register_callback(
self, ref: "ClientObjectRef", callback: Callable[["DataResponse"], None]
) -> None:
self.worker.register_callback(ref, callback)
def _get_dashboard_url(self) -> str:
import ray.core.generated.ray_client_pb2 as ray_client_pb2
return self.worker.get_cluster_info(
ray_client_pb2.ClusterInfoType.DASHBOARD_URL
).get("dashboard_url", "")
python/ray/util/client/client_app.py
+91
View File
@@ -0,0 +1,91 @@
from typing import Tuple
from ray.util.client import ray
ray.connect("localhost:50051")
@ray.remote
class HelloActor:
def __init__(self):
self.count = 0
def say_hello(self, whom: str) -> Tuple[str, int]:
self.count += 1
return ("Hello " + whom, self.count)
actor = HelloActor.remote()
s, count = ray.get(actor.say_hello.remote("you"))
print(s, count)
assert s == "Hello you"
assert count == 1
s, count = ray.get(actor.say_hello.remote("world"))
print(s, count)
assert s == "Hello world"
assert count == 2
@ray.remote
def plus2(x):
return x + 2
@ray.remote
def fact(x):
print(x, type(fact))
if x <= 0:
return 1
# This hits the "nested tasks" issue
# https://github.com/ray-project/ray/issues/3644
# So we're on the right track!
return ray.get(fact.remote(x - 1)) * x
@ray.remote
def get_nodes():
return ray.nodes() # Can access the full Ray API in remote methods.
print("Cluster nodes", ray.get(get_nodes.remote()))
print(ray.nodes())
objectref = ray.put("hello world")
# `ClientObjectRef(...)`
print(objectref)
# `hello world`
print(ray.get(objectref))
ref2 = plus2.remote(234)
# `ClientObjectRef(...)`
print(ref2)
# `236`
print(ray.get(ref2))
ref3 = fact.remote(20)
# `ClientObjectRef(...)`
print(ref3)
# `2432902008176640000`
print(ray.get(ref3))
# Reuse the cached ClientRemoteFunc object
ref4 = fact.remote(5)
# `120`
print(ray.get(ref4))
ref5 = fact.remote(10)
print([ref2, ref3, ref4, ref5])
# should return ref2, ref3, ref4
res = ray.wait([ref5, ref2, ref3, ref4], num_returns=3)
print(res)
assert [ref2, ref3, ref4] == res[0]
assert [ref5] == res[1]
# should return ref2, ref3, ref4, ref5
res = ray.wait([ref2, ref3, ref4, ref5], num_returns=4)
print(res)
assert [ref2, ref3, ref4, ref5] == res[0]
assert [] == res[1]
python/ray/util/client/client_pickler.py
+175
View File
@@ -0,0 +1,175 @@
"""Implements the client side of the client/server pickling protocol.
All ray client client/server data transfer happens through this pickling
protocol. The model is as follows:
* All Client objects (eg ClientObjectRef) always live on the client and
are never represented in the server
* All Ray objects (eg, ray.ObjectRef) always live on the server and are
never returned to the client
* In order to translate between these two references, PickleStub tuples
are generated as persistent ids in the data blobs during the pickling
and unpickling of these objects.
The PickleStubs have just enough information to find or generate their
associated partner object on either side.
This also has the advantage of avoiding predefined pickle behavior for ray
objects, which may include ray internal reference counting.
ClientPickler dumps things from the client into the appropriate stubs
ServerUnpickler loads stubs from the server into their client counterparts.
"""
import io
import pickle # noqa: F401
from typing import Any, Dict, NamedTuple, Optional
import ray.cloudpickle as cloudpickle
import ray.core.generated.ray_client_pb2 as ray_client_pb2
from ray.util.client import RayAPIStub
from ray.util.client.common import (
ClientActorClass,
ClientActorHandle,
ClientActorRef,
ClientObjectRef,
ClientRemoteFunc,
ClientRemoteMethod,
InProgressSentinel,
OptionWrapper,
)
# NOTE(barakmich): These PickleStubs are really close to
# the data for an execution, with no arguments. Combine the two?
class PickleStub(
NamedTuple(
"PickleStub",
[
("type", str),
("client_id", str),
("ref_id", bytes),
("name", Optional[str]),
("baseline_options", Optional[Dict]),
],
)
):
def __reduce__(self):
# PySpark's namedtuple monkey patch breaks compatibility with
# cloudpickle. Thus we revert this patch here if it exists.
return object.__reduce__(self)
class ClientPickler(cloudpickle.CloudPickler):
def __init__(self, client_id, *args, **kwargs):
super().__init__(*args, **kwargs)
self.client_id = client_id
def persistent_id(self, obj):
if isinstance(obj, RayAPIStub):
return PickleStub(
type="Ray",
client_id=self.client_id,
ref_id=b"",
name=None,
baseline_options=None,
)
elif isinstance(obj, ClientObjectRef):
return PickleStub(
type="Object",
client_id=self.client_id,
ref_id=obj.id,
name=None,
baseline_options=None,
)
elif isinstance(obj, ClientActorHandle):
return PickleStub(
type="Actor",
client_id=self.client_id,
ref_id=obj._actor_id.id,
name=None,
baseline_options=None,
)
elif isinstance(obj, ClientRemoteFunc):
if obj._ref is None:
obj._ensure_ref()
if type(obj._ref) is InProgressSentinel:
return PickleStub(
type="RemoteFuncSelfReference",
client_id=self.client_id,
ref_id=obj._client_side_ref.id,
name=None,
baseline_options=None,
)
return PickleStub(
type="RemoteFunc",
client_id=self.client_id,
ref_id=obj._ref.id,
name=None,
baseline_options=obj._options,
)
elif isinstance(obj, ClientActorClass):
if obj._ref is None:
obj._ensure_ref()
if type(obj._ref) is InProgressSentinel:
return PickleStub(
type="RemoteActorSelfReference",
client_id=self.client_id,
ref_id=obj._client_side_ref.id,
name=None,
baseline_options=None,
)
return PickleStub(
type="RemoteActor",
client_id=self.client_id,
ref_id=obj._ref.id,
name=None,
baseline_options=obj._options,
)
elif isinstance(obj, ClientRemoteMethod):
return PickleStub(
type="RemoteMethod",
client_id=self.client_id,
ref_id=obj._actor_handle.actor_ref.id,
name=obj._method_name,
baseline_options=None,
)
elif isinstance(obj, OptionWrapper):
raise NotImplementedError("Sending a partial option is unimplemented")
return None
class ServerUnpickler(pickle.Unpickler):
def persistent_load(self, pid):
assert isinstance(pid, PickleStub)
if pid.type == "Object":
return ClientObjectRef(pid.ref_id)
elif pid.type == "Actor":
return ClientActorHandle(ClientActorRef(pid.ref_id))
else:
raise NotImplementedError("Being passed back an unknown stub")
def dumps_from_client(obj: Any, client_id: str, protocol=None) -> bytes:
with io.BytesIO() as file:
cp = ClientPickler(client_id, file, protocol=protocol)
cp.dump(obj)
return file.getvalue()
def loads_from_server(
data: bytes, *, fix_imports=True, encoding="ASCII", errors="strict"
) -> Any:
if isinstance(data, str):
raise TypeError("Can't load pickle from unicode string")
file = io.BytesIO(data)
return ServerUnpickler(
file, fix_imports=fix_imports, encoding=encoding, errors=errors
).load()
def convert_to_arg(val: Any, client_id: str) -> ray_client_pb2.Arg:
out = ray_client_pb2.Arg()
out.local = ray_client_pb2.Arg.Locality.INTERNED
out.data = dumps_from_client(val, client_id)
return out
python/ray/util/client/common.py
+957
View File
@@ -0,0 +1,957 @@
import inspect
import logging
import os
import pickle
import threading
import uuid
from collections import OrderedDict
from concurrent.futures import Future
from dataclasses import dataclass
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
import grpc
import ray._raylet as raylet
import ray.core.generated.ray_client_pb2 as ray_client_pb2
import ray.core.generated.ray_client_pb2_grpc as ray_client_pb2_grpc
from ray._common.signature import extract_signature, get_signature
from ray._private import ray_constants
from ray._private.inspect_util import (
is_class_method,
is_cython,
is_function_or_method,
is_static_method,
)
from ray._private.utils import check_oversized_function
from ray.util.client import ray
from ray.util.client.options import validate_options
from ray.util.common import INT32_MAX
logger = logging.getLogger(__name__)
# gRPC status codes that the client shouldn't attempt to recover from
# Resource exhausted: Server is low on resources, or has hit the max number
# of client connections
# Invalid argument: Reserved for application errors
# Not found: Set if the client is attempting to reconnect to a session that
# does not exist
# Failed precondition: Reserverd for application errors
# Aborted: Set when an error is serialized into the details of the context,
# signals that error should be deserialized on the client side
GRPC_UNRECOVERABLE_ERRORS = (
grpc.StatusCode.RESOURCE_EXHAUSTED,
grpc.StatusCode.INVALID_ARGUMENT,
grpc.StatusCode.NOT_FOUND,
grpc.StatusCode.FAILED_PRECONDITION,
grpc.StatusCode.ABORTED,
)
# TODO: Instead of just making the max message size large, the right thing to
# do is to split up the bytes representation of serialized data into multiple
# messages and reconstruct them on either end. That said, since clients are
# drivers and really just feed initial things in and final results out, (when
# not going to S3 or similar) then a large limit will suffice for many use
# cases.
#
# Currently, this is 2GiB, the max for a signed int.
GRPC_MAX_MESSAGE_SIZE = (2 * 1024 * 1024 * 1024) - 1
# 30 seconds because ELB timeout is 60 seconds
GRPC_KEEPALIVE_TIME_MS = 1000 * 30
# Long timeout because we do not want gRPC ending a connection.
GRPC_KEEPALIVE_TIMEOUT_MS = 1000 * 600
GRPC_OPTIONS = [
*ray_constants.GLOBAL_GRPC_OPTIONS,
("grpc.max_send_message_length", GRPC_MAX_MESSAGE_SIZE),
("grpc.max_receive_message_length", GRPC_MAX_MESSAGE_SIZE),
("grpc.keepalive_time_ms", GRPC_KEEPALIVE_TIME_MS),
("grpc.keepalive_timeout_ms", GRPC_KEEPALIVE_TIMEOUT_MS),
("grpc.keepalive_permit_without_calls", 1),
# Send an infinite number of pings
("grpc.http2.max_pings_without_data", 0),
("grpc.http2.min_ping_interval_without_data_ms", GRPC_KEEPALIVE_TIME_MS - 50),
# Allow many strikes
("grpc.http2.max_ping_strikes", 0),
]
CLIENT_SERVER_MAX_THREADS = float(os.getenv("RAY_CLIENT_SERVER_MAX_THREADS", 100))
# Large objects are chunked into 5 MiB messages, ref PR #35025
OBJECT_TRANSFER_CHUNK_SIZE = 5 * 2**20
# Warn the user if the object being transferred is larger than 2 GiB
OBJECT_TRANSFER_WARNING_SIZE = 2 * 2**30
class ClientObjectRef(raylet.ObjectRef):
def __init__(self, id: Union[bytes, Future]):
self._mutex = threading.Lock()
self._worker = ray.get_context().client_worker
self._id_future = None
if isinstance(id, bytes):
self._set_id(id)
elif isinstance(id, Future):
self._id_future = id
else:
raise TypeError("Unexpected type for id {}".format(id))
def __del__(self):
if self._worker is not None and self._worker.is_connected():
try:
if not self.is_nil():
self._worker.call_release(self.id)
except Exception:
logger.info(
"Exception in ObjectRef is ignored in destructor. "
"To receive this exception in application code, call "
"a method on the actor reference before its destructor "
"is run."
)
def binary(self):
self._wait_for_id()
return super().binary()
def hex(self):
self._wait_for_id()
return super().hex()
def is_nil(self):
self._wait_for_id()
return super().is_nil()
def __hash__(self):
self._wait_for_id()
return hash(self.id)
def task_id(self):
self._wait_for_id()
return super().task_id()
@property
def id(self):
return self.binary()
def future(self) -> Future:
fut = Future()
def set_future(data: Any) -> None:
"""Schedules a callback to set the exception or result
in the Future."""
if isinstance(data, Exception):
fut.set_exception(data)
else:
fut.set_result(data)
self._on_completed(set_future)
# Prevent this object ref from being released.
fut.object_ref = self
return fut
def _on_completed(self, py_callback: Callable[[Any], None]) -> None:
"""Register a callback that will be called after Object is ready.
If the ObjectRef is already ready, the callback will be called soon.
The callback should take the result as the only argument. The result
can be an exception object in case of task error.
"""
def deserialize_obj(
resp: Union[ray_client_pb2.DataResponse, Exception]
) -> None:
from ray.util.client.client_pickler import loads_from_server
if isinstance(resp, Exception):
data = resp
elif isinstance(resp, bytearray):
data = loads_from_server(resp)
else:
obj = resp.get
data = None
if not obj.valid:
data = loads_from_server(resp.get.error)
else:
data = loads_from_server(resp.get.data)
py_callback(data)
self._worker.register_callback(self, deserialize_obj)
def _set_id(self, id):
super()._set_id(id)
self._worker.call_retain(id)
def _wait_for_id(self, timeout=None):
if self._id_future:
with self._mutex:
if self._id_future:
self._set_id(self._id_future.result(timeout=timeout))
self._id_future = None
class ClientActorRef(raylet.ActorID):
def __init__(
self,
id: Union[bytes, Future],
weak_ref: Optional[bool] = False,
):
self._weak_ref = weak_ref
self._mutex = threading.Lock()
self._worker = ray.get_context().client_worker
if isinstance(id, bytes):
self._set_id(id)
self._id_future = None
elif isinstance(id, Future):
self._id_future = id
else:
raise TypeError("Unexpected type for id {}".format(id))
def __del__(self):
if self._weak_ref:
return
if self._worker is not None and self._worker.is_connected():
try:
if not self.is_nil():
self._worker.call_release(self.id)
except Exception:
logger.debug(
"Exception from actor creation is ignored in destructor. "
"To receive this exception in application code, call "
"a method on the actor reference before its destructor "
"is run."
)
def binary(self):
self._wait_for_id()
return super().binary()
def hex(self):
self._wait_for_id()
return super().hex()
def is_nil(self):
self._wait_for_id()
return super().is_nil()
def __hash__(self):
self._wait_for_id()
return hash(self.id)
@property
def id(self):
return self.binary()
def _set_id(self, id):
super()._set_id(id)
self._worker.call_retain(id)
def _wait_for_id(self, timeout=None):
if self._id_future:
with self._mutex:
if self._id_future:
self._set_id(self._id_future.result(timeout=timeout))
self._id_future = None
class ClientStub:
pass
class ClientRemoteFunc(ClientStub):
"""A stub created on the Ray Client to represent a remote
function that can be exectued on the cluster.
This class is allowed to be passed around between remote functions.
Args:
f: The actual function to execute remotely.
options: Optional ``ray.remote`` options applied to this function.
"""
def __init__(self, f: Callable, options: Optional[Dict[str, Any]] = None):
self._lock = threading.Lock()
self._func = f
self._name = f.__name__
self._signature = get_signature(f)
self._ref = None
self._client_side_ref = ClientSideRefID.generate_id()
self._options = validate_options(options)
def __call__(self, *args, **kwargs):
raise TypeError(
"Remote function cannot be called directly. "
f"Use {self._name}.remote method instead"
)
def remote(self, *args, **kwargs):
# Check if supplied parameters match the function signature. Same case
# at the other callsites.
self._signature.bind(*args, **kwargs)
return return_refs(ray.call_remote(self, *args, **kwargs))
def options(self, **kwargs):
return OptionWrapper(self, kwargs)
def _remote(self, args=None, kwargs=None, **option_args):
if args is None:
args = []
if kwargs is None:
kwargs = {}
return self.options(**option_args).remote(*args, **kwargs)
def __repr__(self):
return "ClientRemoteFunc(%s, %s)" % (self._name, self._ref)
def _ensure_ref(self):
with self._lock:
if self._ref is None:
# While calling ray.put() on our function, if
# our function is recursive, it will attempt to
# encode the ClientRemoteFunc -- itself -- and
# infinitely recurse on _ensure_ref.
#
# So we set the state of the reference to be an
# in-progress self reference value, which
# the encoding can detect and handle correctly.
self._ref = InProgressSentinel()
data = ray.worker._dumps_from_client(self._func)
# Check pickled size before sending it to server, which is more
# efficient and can be done synchronously inside remote() call.
check_oversized_function(data, self._name, "remote function", None)
self._ref = ray.worker._put_pickled(
data, client_ref_id=self._client_side_ref.id
)
def _prepare_client_task(self) -> ray_client_pb2.ClientTask:
self._ensure_ref()
task = ray_client_pb2.ClientTask()
task.type = ray_client_pb2.ClientTask.FUNCTION
task.name = self._name
task.payload_id = self._ref.id
set_task_options(task, self._options, "baseline_options")
return task
def _num_returns(self) -> int:
if not self._options:
return None
return self._options.get("num_returns")
class ClientActorClass(ClientStub):
"""A stub created on the Ray Client to represent an actor class.
It is wrapped by ray.remote and can be executed on the cluster.
Args:
actor_cls: The actual class to execute remotely.
options: Optional ``ray.remote`` options applied to this actor class.
"""
def __init__(self, actor_cls: type, options: Optional[Dict[str, Any]] = None):
self.actor_cls = actor_cls
self._lock = threading.Lock()
self._name = actor_cls.__name__
self._init_signature = inspect.Signature(
parameters=extract_signature(actor_cls.__init__, ignore_first=True)
)
self._ref = None
self._client_side_ref = ClientSideRefID.generate_id()
self._options = validate_options(options)
def __call__(self, *args, **kwargs):
raise TypeError(
"Remote actor cannot be instantiated directly. "
f"Use {self._name}.remote() instead"
)
def _ensure_ref(self):
with self._lock:
if self._ref is None:
# As before, set the state of the reference to be an
# in-progress self reference value, which
# the encoding can detect and handle correctly.
self._ref = InProgressSentinel()
data = ray.worker._dumps_from_client(self.actor_cls)
# Check pickled size before sending it to server, which is more
# efficient and can be done synchronously inside remote() call.
check_oversized_function(data, self._name, "actor", None)
self._ref = ray.worker._put_pickled(
data, client_ref_id=self._client_side_ref.id
)
def remote(self, *args, **kwargs) -> "ClientActorHandle":
self._init_signature.bind(*args, **kwargs)
# Actually instantiate the actor
futures = ray.call_remote(self, *args, **kwargs)
assert len(futures) == 1
return ClientActorHandle(ClientActorRef(futures[0]), actor_class=self)
def options(self, **kwargs):
return ActorOptionWrapper(self, kwargs)
def _remote(self, args=None, kwargs=None, **option_args):
if args is None:
args = []
if kwargs is None:
kwargs = {}
return self.options(**option_args).remote(*args, **kwargs)
def __repr__(self):
return "ClientActorClass(%s, %s)" % (self._name, self._ref)
def __getattr__(self, key):
if key not in self.__dict__:
raise AttributeError("Not a class attribute")
raise NotImplementedError("static methods")
def _prepare_client_task(self) -> ray_client_pb2.ClientTask:
self._ensure_ref()
task = ray_client_pb2.ClientTask()
task.type = ray_client_pb2.ClientTask.ACTOR
task.name = self._name
task.payload_id = self._ref.id
set_task_options(task, self._options, "baseline_options")
return task
@staticmethod
def _num_returns() -> int:
return 1
class ClientActorHandle(ClientStub):
"""Client-side stub for instantiated actor.
A stub created on the Ray Client to represent a remote actor that
has been started on the cluster. This class is allowed to be passed
around between remote functions.
Args:
actor_ref: A reference to the running actor given to the client. This
is a serialized version of the actual handle as an opaque token.
actor_class: Optional ``ClientActorClass`` used to populate method
signatures and ``num_returns`` metadata without a server round-trip.
"""
def __init__(
self,
actor_ref: ClientActorRef,
actor_class: Optional[ClientActorClass] = None,
):
self.actor_ref = actor_ref
self._dir: Optional[List[str]] = None
if actor_class is not None:
self._method_num_returns = {}
self._method_signatures = {}
for method_name, method_obj in inspect.getmembers(
actor_class.actor_cls, is_function_or_method
):
self._method_num_returns[method_name] = getattr(
method_obj, "__ray_num_returns__", None
)
self._method_signatures[method_name] = inspect.Signature(
parameters=extract_signature(
method_obj,
ignore_first=(
not (
is_class_method(method_obj)
or is_static_method(actor_class.actor_cls, method_name)
)
),
)
)
else:
self._method_num_returns = None
self._method_signatures = None
def __dir__(self) -> List[str]:
if self._method_num_returns is not None:
return self._method_num_returns.keys()
if ray.is_connected():
self._init_class_info()
return self._method_num_returns.keys()
return super().__dir__()
# For compatibility with core worker ActorHandle._actor_id which returns
# ActorID
@property
def _actor_id(self) -> ClientActorRef:
return self.actor_ref
def __hash__(self) -> int:
return hash(self._actor_id)
def __eq__(self, __value) -> bool:
return hash(self) == hash(__value)
def __getattr__(self, key):
if key == "_method_num_returns":
# We need to explicitly handle this value since it is used below,
# otherwise we may end up infinitely recursing when deserializing.
# This can happen after unpickling an object but before
# _method_num_returns is correctly populated.
raise AttributeError(f"ClientActorRef has no attribute '{key}'")
if self._method_num_returns is None:
self._init_class_info()
if key not in self._method_signatures:
raise AttributeError(f"ClientActorRef has no attribute '{key}'")
return ClientRemoteMethod(
self,
key,
self._method_num_returns.get(key),
self._method_signatures.get(key),
)
def __repr__(self):
return "ClientActorHandle(%s)" % (self.actor_ref.id.hex())
def _init_class_info(self):
# TODO: fetch Ray method decorators
@ray.remote(num_cpus=0)
def get_class_info(x):
return x._ray_method_num_returns, x._ray_method_signatures
self._method_num_returns, method_parameters = ray.get(
get_class_info.remote(self)
)
self._method_signatures = {}
for method, parameters in method_parameters.items():
self._method_signatures[method] = inspect.Signature(parameters=parameters)
class ClientRemoteMethod(ClientStub):
"""A stub for a method on a remote actor.
Can be annotated with execution options.
Args:
actor_handle: A reference to the ClientActorHandle that generated
this method and will have this method called upon it.
method_name: The name of this method.
num_returns: Number of object refs returned by invocations of this
method.
signature: The method's bound signature, used to validate call args.
"""
def __init__(
self,
actor_handle: ClientActorHandle,
method_name: str,
num_returns: int,
signature: inspect.Signature,
):
self._actor_handle = actor_handle
self._method_name = method_name
self._method_num_returns = num_returns
self._signature = signature
def __call__(self, *args, **kwargs):
raise TypeError(
"Actor methods cannot be called directly. Instead "
f"of running 'object.{self._method_name}()', try "
f"'object.{self._method_name}.remote()'."
)
def remote(self, *args, **kwargs):
self._signature.bind(*args, **kwargs)
return return_refs(ray.call_remote(self, *args, **kwargs))
def __repr__(self):
return "ClientRemoteMethod(%s, %s, %s)" % (
self._method_name,
self._actor_handle,
self._method_num_returns,
)
def options(self, **kwargs):
return OptionWrapper(self, kwargs)
def _remote(self, args=None, kwargs=None, **option_args):
if args is None:
args = []
if kwargs is None:
kwargs = {}
return self.options(**option_args).remote(*args, **kwargs)
def _prepare_client_task(self) -> ray_client_pb2.ClientTask:
task = ray_client_pb2.ClientTask()
task.type = ray_client_pb2.ClientTask.METHOD
task.name = self._method_name
task.payload_id = self._actor_handle.actor_ref.id
return task
def _num_returns(self) -> int:
return self._method_num_returns
class OptionWrapper:
def __init__(self, stub: ClientStub, options: Optional[Dict[str, Any]]):
self._remote_stub = stub
self._options = validate_options(options)
def remote(self, *args, **kwargs):
self._remote_stub._signature.bind(*args, **kwargs)
return return_refs(ray.call_remote(self, *args, **kwargs))
def __getattr__(self, key):
return getattr(self._remote_stub, key)
def _prepare_client_task(self):
task = self._remote_stub._prepare_client_task()
set_task_options(task, self._options)
return task
def _num_returns(self) -> int:
if self._options:
num = self._options.get("num_returns")
if num is not None:
return num
return self._remote_stub._num_returns()
class ActorOptionWrapper(OptionWrapper):
def remote(self, *args, **kwargs):
self._remote_stub._init_signature.bind(*args, **kwargs)
futures = ray.call_remote(self, *args, **kwargs)
assert len(futures) == 1
actor_class = None
if isinstance(self._remote_stub, ClientActorClass):
actor_class = self._remote_stub
return ClientActorHandle(ClientActorRef(futures[0]), actor_class=actor_class)
def set_task_options(
task: ray_client_pb2.ClientTask,
options: Optional[Dict[str, Any]],
field: str = "options",
) -> None:
if options is None:
task.ClearField(field)
return
getattr(task, field).pickled_options = pickle.dumps(options)
def return_refs(
futures: List[Future],
) -> Union[None, ClientObjectRef, List[ClientObjectRef]]:
if not futures:
return None
if len(futures) == 1:
return ClientObjectRef(futures[0])
return [ClientObjectRef(fut) for fut in futures]
class InProgressSentinel:
def __repr__(self) -> str:
return self.__class__.__name__
class ClientSideRefID:
"""An ID generated by the client for objects not yet given an ObjectRef"""
def __init__(self, id: bytes):
assert len(id) != 0
self.id = id
@staticmethod
def generate_id() -> "ClientSideRefID":
tid = uuid.uuid4()
return ClientSideRefID(b"\xcc" + tid.bytes)
def remote_decorator(options: Optional[Dict[str, Any]]):
def decorator(function_or_class) -> ClientStub:
if inspect.isfunction(function_or_class) or is_cython(function_or_class):
return ClientRemoteFunc(function_or_class, options=options)
elif inspect.isclass(function_or_class):
return ClientActorClass(function_or_class, options=options)
else:
raise TypeError(
"The @ray.remote decorator must be applied to "
"either a function or to a class."
)
return decorator
@dataclass
class ClientServerHandle:
"""Holds the handles to the registered gRPC servicers and their server."""
task_servicer: ray_client_pb2_grpc.RayletDriverServicer
data_servicer: ray_client_pb2_grpc.RayletDataStreamerServicer
logs_servicer: ray_client_pb2_grpc.RayletLogStreamerServicer
grpc_server: grpc.Server
def stop(self, grace: int) -> None:
# The data servicer might be sleeping while waiting for clients to
# reconnect. Signal that they no longer have to sleep and can exit
# immediately, since the RPC server is stopped.
self.grpc_server.stop(grace)
self.data_servicer.stopped.set()
# Add a hook for all the cases that previously
# expected simply a gRPC server
def __getattr__(self, attr):
return getattr(self.grpc_server, attr)
def _get_client_id_from_context(context: Any) -> str:
"""
Get `client_id` from gRPC metadata. If the `client_id` is not present,
this function logs an error and sets the status_code.
"""
metadata = dict(context.invocation_metadata())
client_id = metadata.get("client_id") or ""
if client_id == "":
logger.error("Client connecting with no client_id")
context.set_code(grpc.StatusCode.FAILED_PRECONDITION)
return client_id
def _propagate_error_in_context(e: Exception, context: Any) -> bool:
"""
Encode an error into the context of an RPC response. Returns True
if the error can be recovered from, false otherwise
"""
try:
if isinstance(e, grpc.RpcError):
# RPC error, propagate directly by copying details into context
context.set_code(e.code())
context.set_details(e.details())
return e.code() not in GRPC_UNRECOVERABLE_ERRORS
except Exception:
# Extra precaution -- if encoding the RPC directly fails fallback
# to treating it as a regular error
pass
context.set_code(grpc.StatusCode.FAILED_PRECONDITION)
context.set_details(str(e))
return False
def _id_is_newer(id1: int, id2: int) -> bool:
"""
We should only replace cache entries with the responses for newer IDs.
Most of the time newer IDs will be the ones with higher value, except when
the req_id counter rolls over. We check for this case by checking the
distance between the two IDs. If the distance is significant, then it's
likely that the req_id counter rolled over, and the smaller id should
still be used to replace the one in cache.
"""
diff = abs(id2 - id1)
# Int32 max is also the maximum number of simultaneous in-flight requests.
if diff > (INT32_MAX // 2):
# Rollover likely occurred. In this case the smaller ID is newer
return id1 < id2
return id1 > id2
class ResponseCache:
"""
Cache for blocking method calls. Needed to prevent retried requests from
being applied multiple times on the server, for example when the client
disconnects. This is used to cache requests/responses sent through
unary-unary RPCs to the RayletServicer.
Note that no clean up logic is used, the last response for each thread
will always be remembered, so at most the cache will hold N entries,
where N is the number of threads on the client side. This relies on the
assumption that a thread will not make a new blocking request until it has
received a response for a previous one, at which point it's safe to
overwrite the old response.
The high level logic is:
1. Before making a call, check the cache for the current thread.
2. If present in the cache, check the request id of the cached
response.
a. If it matches the current request_id, then the request has been
received before and we shouldn't re-attempt the logic. Wait for
the response to become available in the cache, and then return it
b. If it doesn't match, then this is a new request and we can
proceed with calling the real stub. While the response is still
being generated, temporarily keep (req_id, None) in the cache.
Once the call is finished, update the cache entry with the
new (req_id, response) pair. Notify other threads that may
have been waiting for the response to be prepared.
"""
def __init__(self):
self.cv = threading.Condition()
self.cache: Dict[int, Tuple[int, Any]] = {}
def check_cache(self, thread_id: int, request_id: int) -> Optional[Any]:
"""
Check the cache for a given thread, and see if the entry in the cache
matches the current request_id. Returns None if the request_id has
not been seen yet, otherwise returns the cached result.
Throws an error if the placeholder in the cache doesn't match the
request_id -- this means that a new request evicted the old value in
the cache, and that the RPC for `request_id` is redundant and the
result can be discarded, i.e.:
1. Request A is sent (A1)
2. Channel disconnects
3. Request A is resent (A2)
4. A1 is received
5. A2 is received, waits for A1 to finish
6. A1 finishes and is sent back to client
7. Request B is sent
8. Request B overwrites cache entry
9. A2 wakes up extremely late, but cache is now invalid
In practice this is VERY unlikely to happen, but the error can at
least serve as a sanity check or catch invalid request id's.
"""
with self.cv:
if thread_id in self.cache:
cached_request_id, cached_resp = self.cache[thread_id]
if cached_request_id == request_id:
while cached_resp is None:
# The call was started, but the response hasn't yet
# been added to the cache. Let go of the lock and
# wait until the response is ready.
self.cv.wait()
cached_request_id, cached_resp = self.cache[thread_id]
if cached_request_id != request_id:
raise RuntimeError(
"Cached response doesn't match the id of the "
"original request. This might happen if this "
"request was received out of order. The "
"result of the caller is no longer needed. "
f"({request_id} != {cached_request_id})"
)
return cached_resp
if not _id_is_newer(request_id, cached_request_id):
raise RuntimeError(
"Attempting to replace newer cache entry with older "
"one. This might happen if this request was received "
"out of order. The result of the caller is no "
f"longer needed. ({request_id} != {cached_request_id}"
)
self.cache[thread_id] = (request_id, None)
return None
def update_cache(self, thread_id: int, request_id: int, response: Any) -> None:
"""
Inserts `response` into the cache for `request_id`.
"""
with self.cv:
cached_request_id, cached_resp = self.cache[thread_id]
if cached_request_id != request_id or cached_resp is not None:
# The cache was overwritten by a newer requester between
# our call to check_cache and our call to update it.
# This can't happen if the assumption that the cached requests
# are all blocking on the client side, so if you encounter
# this, check if any async requests are being cached.
raise RuntimeError(
"Attempting to update the cache, but placeholder's "
"do not match the current request_id. This might happen "
"if this request was received out of order. The result "
f"of the caller is no longer needed. ({request_id} != "
f"{cached_request_id})"
)
self.cache[thread_id] = (request_id, response)
self.cv.notify_all()
class OrderedResponseCache:
"""
Cache for streaming RPCs, i.e. the DataServicer. Relies on explicit
ack's from the client to determine when it can clean up cache entries.
"""
def __init__(self):
self.last_received = 0
self.cv = threading.Condition()
self.cache: Dict[int, Any] = OrderedDict()
def check_cache(self, req_id: int) -> Optional[Any]:
"""
Check the cache for a given thread, and see if the entry in the cache
matches the current request_id. Returns None if the request_id has
not been seen yet, otherwise returns the cached result.
"""
with self.cv:
if _id_is_newer(self.last_received, req_id) or self.last_received == req_id:
# Request is for an id that has already been cleared from
# cache/acknowledged.
raise RuntimeError(
"Attempting to accesss a cache entry that has already "
"cleaned up. The client has already acknowledged "
f"receiving this response. ({req_id}, "
f"{self.last_received})"
)
if req_id in self.cache:
cached_resp = self.cache[req_id]
while cached_resp is None:
# The call was started, but the response hasn't yet been
# added to the cache. Let go of the lock and wait until
# the response is ready
self.cv.wait()
if req_id not in self.cache:
raise RuntimeError(
"Cache entry was removed. This likely means that "
"the result of this call is no longer needed."
)
cached_resp = self.cache[req_id]
return cached_resp
self.cache[req_id] = None
return None
def update_cache(self, req_id: int, resp: Any) -> None:
"""
Inserts `response` into the cache for `request_id`.
"""
with self.cv:
self.cv.notify_all()
if req_id not in self.cache:
raise RuntimeError(
"Attempting to update the cache, but placeholder is "
"missing. This might happen on a redundant call to "
f"update_cache. ({req_id})"
)
self.cache[req_id] = resp
def invalidate(self, e: Exception) -> bool:
"""
Invalidate any partially populated cache entries, replacing their
placeholders with the passed in exception. Useful to prevent a thread
from waiting indefinitely on a failed call.
Returns True if the cache contains an error, False otherwise
"""
with self.cv:
invalid = False
for req_id in self.cache:
if self.cache[req_id] is None:
self.cache[req_id] = e
if isinstance(self.cache[req_id], Exception):
invalid = True
self.cv.notify_all()
return invalid
def cleanup(self, last_received: int) -> None:
"""
Cleanup all of the cached requests up to last_received. Assumes that
the cache entries were inserted in ascending order.
"""
with self.cv:
if _id_is_newer(last_received, self.last_received):
self.last_received = last_received
to_remove = []
for req_id in self.cache:
if _id_is_newer(last_received, req_id) or last_received == req_id:
to_remove.append(req_id)
else:
break
for req_id in to_remove:
del self.cache[req_id]
self.cv.notify_all()
python/ray/util/client/dataclient.py
+598
View File
@@ -0,0 +1,598 @@
"""This file implements a threaded stream controller to abstract a data stream
back to the ray clientserver.
"""
import logging
import math
import queue
import threading
import warnings
from collections import OrderedDict
from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Union
import grpc
import ray.core.generated.ray_client_pb2 as ray_client_pb2
import ray.core.generated.ray_client_pb2_grpc as ray_client_pb2_grpc
from ray.util.client.common import (
INT32_MAX,
OBJECT_TRANSFER_CHUNK_SIZE,
OBJECT_TRANSFER_WARNING_SIZE,
)
from ray.util.debug import log_once
if TYPE_CHECKING:
from ray.util.client.worker import Worker
logger = logging.getLogger(__name__)
ResponseCallable = Callable[[Union[ray_client_pb2.DataResponse, Exception]], None]
# Send an acknowledge on every 32nd response received
ACKNOWLEDGE_BATCH_SIZE = 32
def chunk_put(req: ray_client_pb2.DataRequest):
"""
Chunks a put request. Doing this lazily is important for large objects,
since taking slices of bytes objects does a copy. This means if we
immediately materialized every chunk of a large object and inserted them
into the result_queue, we would effectively double the memory needed
on the client to handle the put.
"""
# When accessing a protobuf field, deserialization is performed, which will
# generate a copy. So we need to avoid accessing the `data` field multiple
# times in the loop
request_data = req.put.data
total_size = len(request_data)
assert total_size > 0, "Cannot chunk object with missing data"
if total_size >= OBJECT_TRANSFER_WARNING_SIZE and log_once(
"client_object_put_size_warning"
):
size_gb = total_size / 2**30
warnings.warn(
"Ray Client is attempting to send a "
f"{size_gb:.2f} GiB object over the network, which may "
"be slow. Consider serializing the object and using a remote "
"URI to transfer via S3 or Google Cloud Storage instead. "
"Documentation for doing this can be found here: "
"https://docs.ray.io/en/latest/handling-dependencies.html#remote-uris",
UserWarning,
)
total_chunks = math.ceil(total_size / OBJECT_TRANSFER_CHUNK_SIZE)
for chunk_id in range(0, total_chunks):
start = chunk_id * OBJECT_TRANSFER_CHUNK_SIZE
end = min(total_size, (chunk_id + 1) * OBJECT_TRANSFER_CHUNK_SIZE)
chunk = ray_client_pb2.PutRequest(
client_ref_id=req.put.client_ref_id,
data=request_data[start:end],
chunk_id=chunk_id,
total_chunks=total_chunks,
total_size=total_size,
)
yield ray_client_pb2.DataRequest(req_id=req.req_id, put=chunk)
def chunk_task(req: ray_client_pb2.DataRequest):
"""
Chunks a client task. Doing this lazily is important with large arguments,
since taking slices of bytes objects does a copy. This means if we
immediately materialized every chunk of a large argument and inserted them
into the result_queue, we would effectively double the memory needed
on the client to handle the task.
"""
# When accessing a protobuf field, deserialization is performed, which will
# generate a copy. So we need to avoid accessing the `data` field multiple
# times in the loop
request_data = req.task.data
total_size = len(request_data)
assert total_size > 0, "Cannot chunk object with missing data"
total_chunks = math.ceil(total_size / OBJECT_TRANSFER_CHUNK_SIZE)
for chunk_id in range(0, total_chunks):
start = chunk_id * OBJECT_TRANSFER_CHUNK_SIZE
end = min(total_size, (chunk_id + 1) * OBJECT_TRANSFER_CHUNK_SIZE)
chunk = ray_client_pb2.ClientTask(
type=req.task.type,
name=req.task.name,
payload_id=req.task.payload_id,
client_id=req.task.client_id,
options=req.task.options,
baseline_options=req.task.baseline_options,
namespace=req.task.namespace,
data=request_data[start:end],
chunk_id=chunk_id,
total_chunks=total_chunks,
)
yield ray_client_pb2.DataRequest(req_id=req.req_id, task=chunk)
class ChunkCollector:
"""
This object collects chunks from async get requests via __call__, and
calls the underlying callback when the object is fully received, or if an
exception while retrieving the object occurs.
This is not used in synchronous gets (synchronous gets interact with the
raylet servicer directly, not through the datapath).
__call__ returns true once the underlying call back has been called.
"""
def __init__(self, callback: ResponseCallable, request: ray_client_pb2.DataRequest):
# Bytearray containing data received so far
self.data = bytearray()
# The callback that will be called once all data is received
self.callback = callback
# The id of the last chunk we've received, or -1 if haven't seen any yet
self.last_seen_chunk = -1
# The GetRequest that initiated the transfer. start_chunk_id will be
# updated as chunks are received to avoid re-requesting chunks that
# we've already received.
self.request = request
def __call__(self, response: Union[ray_client_pb2.DataResponse, Exception]) -> bool:
if isinstance(response, Exception):
self.callback(response)
return True
get_resp = response.get
if not get_resp.valid:
self.callback(response)
return True
if get_resp.total_size > OBJECT_TRANSFER_WARNING_SIZE and log_once(
"client_object_transfer_size_warning"
):
size_gb = get_resp.total_size / 2**30
warnings.warn(
"Ray Client is attempting to retrieve a "
f"{size_gb:.2f} GiB object over the network, which may "
"be slow. Consider serializing the object to a file and "
"using rsync or S3 instead.",
UserWarning,
)
chunk_data = get_resp.data
chunk_id = get_resp.chunk_id
if chunk_id == self.last_seen_chunk + 1:
self.data.extend(chunk_data)
self.last_seen_chunk = chunk_id
# If we disconnect partway through, restart the get request
# at the first chunk we haven't seen
self.request.get.start_chunk_id = self.last_seen_chunk + 1
elif chunk_id > self.last_seen_chunk + 1:
# A chunk was skipped. This shouldn't happen in practice since
# grpc guarantees that chunks will arrive in order.
msg = (
f"Received chunk {chunk_id} when we expected "
f"{self.last_seen_chunk + 1} for request {response.req_id}"
)
logger.warning(msg)
self.callback(RuntimeError(msg))
return True
else:
# We received a chunk that've already seen before. Ignore, since
# it should already be appended to self.data.
logger.debug(
f"Received a repeated chunk {chunk_id} "
f"from request {response.req_id}."
)
if get_resp.chunk_id == get_resp.total_chunks - 1:
self.callback(self.data)
return True
else:
# Not done yet
return False
class DataClient:
def __init__(self, client_worker: "Worker", client_id: str, metadata: list):
"""Initializes a thread-safe datapath over a Ray Client gRPC channel.
Args:
client_worker: The Ray Client worker that manages this client
client_id: the generated ID representing this client
metadata: metadata to pass to gRPC requests
"""
self.client_worker = client_worker
self._client_id = client_id
self._metadata = metadata
self.data_thread = self._start_datathread()
# Track outstanding requests to resend in case of disconnection
self.outstanding_requests: Dict[int, Any] = OrderedDict()
# Serialize access to all mutable internal states: self.request_queue,
# self.ready_data, self.asyncio_waiting_data,
# self._in_shutdown, self._req_id, self.outstanding_requests and
# calling self._next_id()
self.lock = threading.Lock()
# Waiting for response or shutdown.
self.cv = threading.Condition(lock=self.lock)
self.request_queue = self._create_queue()
self.ready_data: Dict[int, Any] = {}
# NOTE: Dictionary insertion is guaranteed to complete before lookup
# and/or removal because of synchronization via the request_queue.
self.asyncio_waiting_data: Dict[int, ResponseCallable] = {}
self._in_shutdown = False
self._req_id = 0
self._last_exception = None
self._acknowledge_counter = 0
self.data_thread.start()
# Must hold self.lock when calling this function.
def _next_id(self) -> int:
assert self.lock.locked()
self._req_id += 1
if self._req_id > INT32_MAX:
self._req_id = 1
# Responses that aren't tracked (like opportunistic releases)
# have req_id=0, so make sure we never mint such an id.
assert self._req_id != 0
return self._req_id
def _start_datathread(self) -> threading.Thread:
return threading.Thread(
target=self._data_main,
name="ray_client_streaming_rpc",
args=(),
daemon=True,
)
# A helper that takes requests from queue. If the request wraps a PutRequest,
# lazily chunks and yields the request. Otherwise, yields the request directly.
def _requests(self):
while True:
req = self.request_queue.get()
if req is None:
# Stop when client signals shutdown.
return
req_type = req.WhichOneof("type")
if req_type == "put":
yield from chunk_put(req)
elif req_type == "task":
yield from chunk_task(req)
else:
yield req
def _data_main(self) -> None:
reconnecting = False
try:
while not self.client_worker._in_shutdown:
stub = ray_client_pb2_grpc.RayletDataStreamerStub(
self.client_worker.channel
)
metadata = self._metadata + [("reconnecting", str(reconnecting))]
resp_stream = stub.Datapath(
self._requests(),
metadata=metadata,
wait_for_ready=True,
)
try:
for response in resp_stream:
self._process_response(response)
return
except grpc.RpcError as e:
reconnecting = self._can_reconnect(e)
if not reconnecting:
self._last_exception = e
return
self._reconnect_channel()
except Exception as e:
self._last_exception = e
finally:
logger.debug("Shutting down data channel.")
self._shutdown()
def _process_response(self, response: Any) -> None:
"""
Process responses from the data servicer.
"""
if response.req_id == 0:
# This is not being waited for.
logger.debug(f"Got unawaited response {response}")
return
if response.req_id in self.asyncio_waiting_data:
can_remove = True
try:
callback = self.asyncio_waiting_data[response.req_id]
if isinstance(callback, ChunkCollector):
can_remove = callback(response)
elif callback:
callback(response)
if can_remove:
# NOTE: calling del self.asyncio_waiting_data results
# in the destructor of ClientObjectRef running, which
# calls ReleaseObject(). So self.asyncio_waiting_data
# is accessed without holding self.lock. Holding the
# lock shouldn't be necessary either.
del self.asyncio_waiting_data[response.req_id]
except Exception:
logger.exception("Callback error:")
with self.lock:
# Update outstanding requests
if response.req_id in self.outstanding_requests and can_remove:
del self.outstanding_requests[response.req_id]
# Acknowledge response
self._acknowledge(response.req_id)
else:
with self.lock:
self.ready_data[response.req_id] = response
self.cv.notify_all()
def _can_reconnect(self, e: grpc.RpcError) -> bool:
"""
Processes RPC errors that occur while reading from data stream.
Returns True if the error can be recovered from, False otherwise.
"""
if not self.client_worker._can_reconnect(e):
logger.error("Unrecoverable error in data channel.")
logger.debug(e)
return False
logger.debug("Recoverable error in data channel.")
logger.debug(e)
return True
def _shutdown(self) -> None:
"""
Shutdown the data channel
"""
with self.lock:
self._in_shutdown = True
self.cv.notify_all()
callbacks = self.asyncio_waiting_data.values()
self.asyncio_waiting_data = {}
if self._last_exception:
# Abort async requests with the error.
err = ConnectionError(
"Failed during this or a previous request. Exception that "
f"broke the connection: {self._last_exception}"
)
else:
err = ConnectionError(
"Request cannot be fulfilled because the data client has "
"disconnected."
)
for callback in callbacks:
if callback:
callback(err)
# Since self._in_shutdown is set to True, no new item
# will be added to self.asyncio_waiting_data
def _acknowledge(self, req_id: int) -> None:
"""
Puts an acknowledge request on the request queue periodically.
Lock should be held before calling this. Used when an async or
blocking response is received.
"""
if not self.client_worker._reconnect_enabled:
# Skip ACKs if reconnect isn't enabled
return
assert self.lock.locked()
self._acknowledge_counter += 1
if self._acknowledge_counter % ACKNOWLEDGE_BATCH_SIZE == 0:
self.request_queue.put(
ray_client_pb2.DataRequest(
acknowledge=ray_client_pb2.AcknowledgeRequest(req_id=req_id)
)
)
def _reconnect_channel(self) -> None:
"""
Attempts to reconnect the gRPC channel and resend outstanding
requests. First, the server is pinged to see if the current channel
still works. If the ping fails, then the current channel is closed
and replaced with a new one.
Once a working channel is available, a new request queue is made
and filled with any outstanding requests to be resent to the server.
"""
try:
# Ping the server to see if the current channel is reuseable, for
# example if gRPC reconnected the channel on its own or if the
# RPC error was transient and the channel is still open
ping_succeeded = self.client_worker.ping_server(timeout=5)
except grpc.RpcError:
ping_succeeded = False
if not ping_succeeded:
# Ping failed, try refreshing the data channel
logger.warning(
"Encountered connection issues in the data channel. "
"Attempting to reconnect."
)
try:
self.client_worker._connect_channel(reconnecting=True)
except ConnectionError:
logger.warning("Failed to reconnect the data channel")
raise
logger.debug("Reconnection succeeded!")
# Recreate the request queue, and resend outstanding requests
with self.lock:
self.request_queue = self._create_queue()
for request in self.outstanding_requests.values():
# Resend outstanding requests
self.request_queue.put(request)
# Use SimpleQueue to avoid deadlocks when appending to queue from __del__()
@staticmethod
def _create_queue():
return queue.SimpleQueue()
def close(self) -> None:
thread = None
with self.lock:
self._in_shutdown = True
# Notify blocking operations to fail.
self.cv.notify_all()
# Add sentinel to terminate streaming RPC.
if self.request_queue is not None:
# Intentional shutdown, tell server it can clean up the
# connection immediately and ignore the reconnect grace period.
cleanup_request = ray_client_pb2.DataRequest(
connection_cleanup=ray_client_pb2.ConnectionCleanupRequest()
)
self.request_queue.put(cleanup_request)
self.request_queue.put(None)
if self.data_thread is not None:
thread = self.data_thread
# Wait until streaming RPCs are done.
if thread is not None:
thread.join()
def _blocking_send(
self, req: ray_client_pb2.DataRequest
) -> ray_client_pb2.DataResponse:
with self.lock:
self._check_shutdown()
req_id = self._next_id()
req.req_id = req_id
self.request_queue.put(req)
self.outstanding_requests[req_id] = req
self.cv.wait_for(lambda: req_id in self.ready_data or self._in_shutdown)
self._check_shutdown()
data = self.ready_data[req_id]
del self.ready_data[req_id]
del self.outstanding_requests[req_id]
self._acknowledge(req_id)
return data
def _async_send(
self,
req: ray_client_pb2.DataRequest,
callback: Optional[ResponseCallable] = None,
) -> None:
with self.lock:
self._check_shutdown()
req_id = self._next_id()
req.req_id = req_id
self.asyncio_waiting_data[req_id] = callback
self.outstanding_requests[req_id] = req
self.request_queue.put(req)
# Must hold self.lock when calling this function.
def _check_shutdown(self):
assert self.lock.locked()
if not self._in_shutdown:
return
self.lock.release()
# Do not try disconnect() or throw exceptions in self.data_thread.
# Otherwise deadlock can occur.
if threading.current_thread().ident == self.data_thread.ident:
return
from ray.util import disconnect
disconnect()
self.lock.acquire()
if self._last_exception is not None:
msg = (
"Request can't be sent because the Ray client has already "
"been disconnected due to an error. Last exception: "
f"{self._last_exception}"
)
else:
msg = (
"Request can't be sent because the Ray client has already "
"been disconnected."
)
raise ConnectionError(msg)
def Init(
self, request: ray_client_pb2.InitRequest, context=None
) -> ray_client_pb2.InitResponse:
datareq = ray_client_pb2.DataRequest(
init=request,
)
resp = self._blocking_send(datareq)
return resp.init
def PrepRuntimeEnv(
self, request: ray_client_pb2.PrepRuntimeEnvRequest, context=None
) -> ray_client_pb2.PrepRuntimeEnvResponse:
datareq = ray_client_pb2.DataRequest(
prep_runtime_env=request,
)
resp = self._blocking_send(datareq)
return resp.prep_runtime_env
def ConnectionInfo(self, context=None) -> ray_client_pb2.ConnectionInfoResponse:
datareq = ray_client_pb2.DataRequest(
connection_info=ray_client_pb2.ConnectionInfoRequest()
)
resp = self._blocking_send(datareq)
return resp.connection_info
def GetObject(
self, request: ray_client_pb2.GetRequest, context=None
) -> ray_client_pb2.GetResponse:
datareq = ray_client_pb2.DataRequest(
get=request,
)
resp = self._blocking_send(datareq)
return resp.get
def RegisterGetCallback(
self, request: ray_client_pb2.GetRequest, callback: ResponseCallable
) -> None:
if len(request.ids) != 1:
raise ValueError(
"RegisterGetCallback() must have exactly 1 Object ID. "
f"Actual: {request}"
)
datareq = ray_client_pb2.DataRequest(
get=request,
)
collector = ChunkCollector(callback=callback, request=datareq)
self._async_send(datareq, collector)
# TODO: convert PutObject to async
def PutObject(
self, request: ray_client_pb2.PutRequest, context=None
) -> ray_client_pb2.PutResponse:
datareq = ray_client_pb2.DataRequest(
put=request,
)
resp = self._blocking_send(datareq)
return resp.put
def ReleaseObject(
self, request: ray_client_pb2.ReleaseRequest, context=None
) -> None:
datareq = ray_client_pb2.DataRequest(
release=request,
)
self._async_send(datareq)
def Schedule(self, request: ray_client_pb2.ClientTask, callback: ResponseCallable):
datareq = ray_client_pb2.DataRequest(task=request)
self._async_send(datareq, callback)
def Terminate(
self, request: ray_client_pb2.TerminateRequest
) -> ray_client_pb2.TerminateResponse:
req = ray_client_pb2.DataRequest(
terminate=request,
)
resp = self._blocking_send(req)
return resp.terminate
def ListNamedActors(
self, request: ray_client_pb2.ClientListNamedActorsRequest
) -> ray_client_pb2.ClientListNamedActorsResponse:
req = ray_client_pb2.DataRequest(
list_named_actors=request,
)
resp = self._blocking_send(req)
return resp.list_named_actors
python/ray/util/client/examples/run_tune.py
+6
View File
@@ -0,0 +1,6 @@
from ray.tune import tune
from ray.util.client import ray
ray.connect("localhost:50051")
tune.run("PG", config={"env": "CartPole-v0"})
python/ray/util/client/logsclient.py
+135
View File
@@ -0,0 +1,135 @@
"""This file implements a threaded stream controller to return logs back from
the ray clientserver.
"""
import logging
import queue
import sys
import threading
import time
from typing import TYPE_CHECKING
import grpc
import ray.core.generated.ray_client_pb2 as ray_client_pb2
import ray.core.generated.ray_client_pb2_grpc as ray_client_pb2_grpc
from ray.util.debug import log_once
if TYPE_CHECKING:
from ray.util.client.worker import Worker
logger = logging.getLogger(__name__)
# TODO(barakmich): Running a logger in a logger causes loopback.
# The client logger need its own root -- possibly this one.
# For the moment, let's just not propagate beyond this point.
logger.propagate = False
class LogstreamClient:
def __init__(self, client_worker: "Worker", metadata: list):
"""Initializes a thread-safe log stream over a Ray Client gRPC channel.
Args:
client_worker: The Ray Client worker that manages this client
metadata: metadata to pass to gRPC requests
"""
self.client_worker = client_worker
self._metadata = metadata
self.request_queue = queue.Queue()
self.log_thread = self._start_logthread()
self.log_thread.start()
self.last_req = None
def _start_logthread(self) -> threading.Thread:
return threading.Thread(target=self._log_main, args=(), daemon=True)
def _log_main(self) -> None:
reconnecting = False
while not self.client_worker._in_shutdown:
if reconnecting:
# Refresh queue and retry last request
self.request_queue = queue.Queue()
if self.last_req:
self.request_queue.put(self.last_req)
stub = ray_client_pb2_grpc.RayletLogStreamerStub(self.client_worker.channel)
try:
log_stream = stub.Logstream(
iter(self.request_queue.get, None), metadata=self._metadata
)
except ValueError:
# Trying to use the stub on a cancelled channel will raise
# ValueError. This should only happen when the data client
# is attempting to reset the connection -- sleep and try
# again.
time.sleep(0.5)
continue
try:
for record in log_stream:
if record.level < 0:
self.stdstream(level=record.level, msg=record.msg)
self.log(level=record.level, msg=record.msg)
return
except grpc.RpcError as e:
reconnecting = self._process_rpc_error(e)
if not reconnecting:
return
def _process_rpc_error(self, e: grpc.RpcError) -> bool:
"""
Processes RPC errors that occur while reading from data stream.
Returns True if the error can be recovered from, False otherwise.
"""
if self.client_worker._can_reconnect(e):
if log_once("lost_reconnect_logs"):
logger.warning(
"Log channel is reconnecting. Logs produced while "
"the connection was down can be found on the head "
"node of the cluster in "
"`ray_client_server_[port].out`"
)
logger.debug("Log channel dropped, retrying.")
time.sleep(0.5)
return True
logger.debug("Shutting down log channel.")
if not self.client_worker._in_shutdown:
logger.exception("Unexpected exception:")
return False
def log(self, level: int, msg: str):
"""Log the message from the log stream.
By default, calls logger.log but this can be overridden.
Args:
level: The loglevel of the received log message
msg: The content of the message
"""
logger.log(level=level, msg=msg)
def stdstream(self, level: int, msg: str):
"""Log the stdout/stderr entry from the log stream.
By default, calls print but this can be overridden.
Args:
level: The loglevel of the received log message
msg: The content of the message
"""
print_file = sys.stderr if level == -2 else sys.stdout
print(msg, file=print_file, end="")
def set_logstream_level(self, level: int):
logger.setLevel(level)
req = ray_client_pb2.LogSettingsRequest()
req.enabled = True
req.loglevel = level
self.request_queue.put(req)
self.last_req = req
def close(self) -> None:
self.request_queue.put(None)
if self.log_thread is not None:
self.log_thread.join()
def disable_logs(self) -> None:
req = ray_client_pb2.LogSettingsRequest()
req.enabled = False
self.request_queue.put(req)
self.last_req = req
python/ray/util/client/options.py
+45
View File
@@ -0,0 +1,45 @@
from typing import Any, Dict, Optional
from ray._common import ray_option_utils
from ray.util.placement_group import PlacementGroup, check_placement_group_index
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
def validate_options(kwargs_dict: Optional[Dict[str, Any]]) -> Optional[Dict[str, Any]]:
if kwargs_dict is None:
return None
if len(kwargs_dict) == 0:
return None
out = {}
for k, v in kwargs_dict.items():
if k not in ray_option_utils.valid_options:
raise ValueError(
f"Invalid option keyword: '{k}'. "
f"{ray_option_utils.remote_args_error_string}"
)
ray_option_utils.valid_options[k].validate(k, v)
out[k] = v
# Validate placement setting similar to the logic in ray/actor.py and
# ray/remote_function.py. The difference is that when
# placement_group = default and placement_group_capture_child_tasks
# specified, placement group cannot be resolved at client. So this check
# skips this case and relies on server to enforce any condition.
bundle_index = out.get("placement_group_bundle_index", None)
pg = out.get("placement_group", None)
scheduling_strategy = out.get("scheduling_strategy", None)
if isinstance(scheduling_strategy, PlacementGroupSchedulingStrategy):
pg = scheduling_strategy.placement_group
bundle_index = scheduling_strategy.placement_group_bundle_index
if bundle_index is not None:
if pg is None:
pg = PlacementGroup.empty()
if pg == "default" and (
out.get("placement_group_capture_child_tasks", None) is None
):
pg = PlacementGroup.empty()
if isinstance(pg, PlacementGroup):
check_placement_group_index(pg, bundle_index)
return out
python/ray/util/client/ray_client_helpers.py
+83
View File
@@ -0,0 +1,83 @@
import time
from contextlib import contextmanager
from typing import Any, Dict
import ray as real_ray
import ray.util.client.server.server as ray_client_server
from ray._common.network_utils import build_address, get_localhost_ip
from ray._private.client_mode_hook import disable_client_hook
from ray.job_config import JobConfig
from ray.util.client import ray
@contextmanager
def ray_start_client_server(metadata=None, ray_connect_handler=None, **kwargs):
with ray_start_client_server_pair(
metadata=metadata, ray_connect_handler=ray_connect_handler, **kwargs
) as pair:
client, server = pair
yield client
@contextmanager
def ray_start_client_server_for_address(address):
"""
Starts a Ray client server that initializes drivers at the specified address.
"""
def connect_handler(
job_config: JobConfig = None, **ray_init_kwargs: Dict[str, Any]
):
import ray
with disable_client_hook():
if not ray.is_initialized():
return ray.init(address, job_config=job_config, **ray_init_kwargs)
with ray_start_client_server(ray_connect_handler=connect_handler) as ray:
yield ray
@contextmanager
def ray_start_client_server_pair(metadata=None, ray_connect_handler=None, **kwargs):
ray._inside_client_test = True
with disable_client_hook():
assert not ray.is_initialized()
server = ray_client_server.serve(
get_localhost_ip(), 50051, ray_connect_handler=ray_connect_handler
)
ray.connect(build_address(get_localhost_ip(), 50051), metadata=metadata, **kwargs)
try:
yield ray, server
finally:
ray._inside_client_test = False
ray.disconnect()
server.stop(0)
del server
start = time.monotonic()
with disable_client_hook():
while ray.is_initialized():
time.sleep(1)
if time.monotonic() - start > 30:
raise RuntimeError("Failed to terminate Ray")
# Allow windows to close processes before moving on
time.sleep(3)
@contextmanager
def ray_start_cluster_client_server_pair(address):
ray._inside_client_test = True
def ray_connect_handler(job_config=None, **ray_init_kwargs):
real_ray.init(address=address)
server = ray_client_server.serve(
get_localhost_ip(), 50051, ray_connect_handler=ray_connect_handler
)
ray.connect(build_address(get_localhost_ip(), 50051))
try:
yield ray, server
finally:
ray._inside_client_test = False
ray.disconnect()
server.stop(0)
python/ray/util/client/runtime_context.py
+77
View File
@@ -0,0 +1,77 @@
from types import SimpleNamespace
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from ray import JobID, NodeID, WorkerID
from ray.runtime_context import RuntimeContext
class _ClientWorkerPropertyAPI:
"""Emulates the properties of the ray._private.worker object for the client"""
def __init__(self, worker):
assert worker is not None
self.worker = worker
def build_runtime_context(self) -> "RuntimeContext":
"""Creates a RuntimeContext backed by the properites of this API"""
# Defer the import of RuntimeContext until needed to avoid cycles
from ray.runtime_context import RuntimeContext
return RuntimeContext(self)
def _fetch_runtime_context(self):
import ray.core.generated.ray_client_pb2 as ray_client_pb2
return self.worker.get_cluster_info(
ray_client_pb2.ClusterInfoType.RUNTIME_CONTEXT
)
@property
def mode(self):
from ray._private.worker import SCRIPT_MODE
return SCRIPT_MODE
@property
def current_job_id(self) -> "JobID":
from ray import JobID
return JobID(self._fetch_runtime_context().job_id)
@property
def current_node_id(self) -> "NodeID":
from ray import NodeID
return NodeID(self._fetch_runtime_context().node_id)
@property
def worker_id(self) -> "WorkerID":
"""Binary worker id for the Ray Client server process (cluster driver worker)."""
from ray import WorkerID
return WorkerID(self._fetch_runtime_context().worker_id)
@property
def namespace(self) -> str:
return self._fetch_runtime_context().namespace
@property
def should_capture_child_tasks_in_placement_group(self) -> bool:
return self._fetch_runtime_context().capture_client_tasks
@property
def runtime_env(self) -> str:
return self._fetch_runtime_context().runtime_env
def check_connected(self) -> bool:
return self.worker.ping_server()
@property
def gcs_client(self) -> str:
return SimpleNamespace(address=self._fetch_runtime_context().gcs_address)
@property
def node(self):
"""Emulates the worker.node property for client mode"""
return SimpleNamespace(session_name=self._fetch_runtime_context().session_name)
python/ray/util/client/server/__init__.py
+1
View File
@@ -0,0 +1 @@
from ray.util.client.server.server import serve # noqa
python/ray/util/client/server/__main__.py
+4
View File
@@ -0,0 +1,4 @@
if __name__ == "__main__":
from ray.util.client.server.server import main
main()
python/ray/util/client/server/dataservicer.py
+415
View File
@@ -0,0 +1,415 @@
import logging
import sys
import time
from collections import defaultdict
from queue import Queue
from threading import Event, Lock, Thread
from typing import TYPE_CHECKING, Any, Dict, Iterator, Union
import grpc
import ray
import ray.core.generated.ray_client_pb2 as ray_client_pb2
import ray.core.generated.ray_client_pb2_grpc as ray_client_pb2_grpc
from ray._private.client_mode_hook import disable_client_hook
from ray.util.client.common import (
CLIENT_SERVER_MAX_THREADS,
OrderedResponseCache,
_propagate_error_in_context,
)
from ray.util.client.server.server_pickler import loads_from_client
from ray.util.debug import log_once
if TYPE_CHECKING:
from ray.util.client.server.server import RayletServicer
logger = logging.getLogger(__name__)
QUEUE_JOIN_SECONDS = 10
def _get_reconnecting_from_context(context: Any) -> bool:
"""
Get `reconnecting` from gRPC metadata, or False if missing.
"""
metadata = dict(context.invocation_metadata())
val = metadata.get("reconnecting")
if val is None or val not in ("True", "False"):
logger.error(
f'Client connecting with invalid value for "reconnecting": {val}, '
"This may be because you have a mismatched client and server "
"version."
)
return False
return val == "True"
def _should_cache(req: ray_client_pb2.DataRequest) -> bool:
"""
Returns True if the response should to the given request should be cached,
false otherwise. At the moment the only requests we do not cache are:
- asynchronous gets: These arrive out of order. Skipping caching here
is fine, since repeating an async get is idempotent
- acks: Repeating acks is idempotent
- clean up requests: Also idempotent, and client has likely already
wrapped up the data connection by this point.
- puts: We should only cache when we receive the final chunk, since
any earlier chunks won't generate a response
- tasks: We should only cache when we receive the final chunk,
since any earlier chunks won't generate a response
"""
req_type = req.WhichOneof("type")
if req_type == "get" and req.get.asynchronous:
return False
if req_type == "put":
return req.put.chunk_id == req.put.total_chunks - 1
if req_type == "task":
return req.task.chunk_id == req.task.total_chunks - 1
return req_type not in ("acknowledge", "connection_cleanup")
def fill_queue(
grpc_input_generator: Iterator[ray_client_pb2.DataRequest],
output_queue: "Queue[Union[ray_client_pb2.DataRequest, ray_client_pb2.DataResponse]]", # noqa: E501
) -> None:
"""
Pushes incoming requests to a shared output_queue.
"""
try:
for req in grpc_input_generator:
output_queue.put(req)
except grpc.RpcError as e:
logger.debug(
"closing dataservicer reader thread "
f"grpc error reading request_iterator: {e}"
)
finally:
# Set the sentinel value for the output_queue
output_queue.put(None)
class ChunkCollector:
"""
Helper class for collecting chunks from PutObject or ClientTask messages
"""
def __init__(self):
self.curr_req_id = None
self.last_seen_chunk_id = -1
self.data = bytearray()
def add_chunk(
self,
req: ray_client_pb2.DataRequest,
chunk: Union[ray_client_pb2.PutRequest, ray_client_pb2.ClientTask],
):
if self.curr_req_id is not None and self.curr_req_id != req.req_id:
raise RuntimeError(
"Expected to receive a chunk from request with id "
f"{self.curr_req_id}, but found {req.req_id} instead."
)
self.curr_req_id = req.req_id
next_chunk = self.last_seen_chunk_id + 1
if chunk.chunk_id < next_chunk:
# Repeated chunk, ignore
return
if chunk.chunk_id > next_chunk:
raise RuntimeError(
f"A chunk {chunk.chunk_id} of request {req.req_id} was "
"received out of order."
)
elif chunk.chunk_id == self.last_seen_chunk_id + 1:
self.data.extend(chunk.data)
self.last_seen_chunk_id = chunk.chunk_id
return chunk.chunk_id + 1 == chunk.total_chunks
def reset(self):
self.curr_req_id = None
self.last_seen_chunk_id = -1
self.data = bytearray()
class DataServicer(ray_client_pb2_grpc.RayletDataStreamerServicer):
def __init__(self, basic_service: "RayletServicer"):
self.basic_service = basic_service
self.clients_lock = Lock()
self.num_clients = 0 # guarded by self.clients_lock
# dictionary mapping client_id's to the last time they connected
self.client_last_seen: Dict[str, float] = {}
# dictionary mapping client_id's to their reconnect grace periods
self.reconnect_grace_periods: Dict[str, float] = {}
# dictionary mapping client_id's to their response cache
self.response_caches: Dict[str, OrderedResponseCache] = defaultdict(
OrderedResponseCache
)
# stopped event, useful for signals that the server is shut down
self.stopped = Event()
# Helper for collecting chunks from PutObject calls. Assumes that
# that put requests from different objects aren't interleaved.
self.put_request_chunk_collector = ChunkCollector()
# Helper for collecting chunks from ClientTask calls. Assumes that
# schedule requests from different remote calls aren't interleaved.
self.client_task_chunk_collector = ChunkCollector()
def Datapath(self, request_iterator, context):
start_time = time.time()
# set to True if client shuts down gracefully
cleanup_requested = False
metadata = dict(context.invocation_metadata())
client_id = metadata.get("client_id")
if client_id is None:
logger.error("Client connecting with no client_id")
return
logger.debug(f"New data connection from client {client_id}: ")
accepted_connection = self._init(client_id, context, start_time)
response_cache = self.response_caches[client_id]
# Set to False if client requests a reconnect grace period of 0
reconnect_enabled = True
if not accepted_connection:
return
try:
request_queue = Queue()
queue_filler_thread = Thread(
target=fill_queue, daemon=True, args=(request_iterator, request_queue)
)
queue_filler_thread.start()
"""For non `async get` requests, this loop yields immediately
For `async get` requests, this loop:
1) does not yield, it just continues
2) When the result is ready, it yields
"""
for req in iter(request_queue.get, None):
if isinstance(req, ray_client_pb2.DataResponse):
# Early shortcut if this is the result of an async get.
yield req
continue
assert isinstance(req, ray_client_pb2.DataRequest)
if _should_cache(req) and reconnect_enabled:
cached_resp = response_cache.check_cache(req.req_id)
if isinstance(cached_resp, Exception):
# Cache state is invalid, raise exception
raise cached_resp
if cached_resp is not None:
yield cached_resp
continue
resp = None
req_type = req.WhichOneof("type")
if req_type == "init":
resp_init = self.basic_service.Init(req.init)
resp = ray_client_pb2.DataResponse(
init=resp_init,
)
with self.clients_lock:
self.reconnect_grace_periods[
client_id
] = req.init.reconnect_grace_period
if req.init.reconnect_grace_period == 0:
reconnect_enabled = False
elif req_type == "get":
if req.get.asynchronous:
get_resp = self.basic_service._async_get_object(
req.get, client_id, req.req_id, request_queue
)
if get_resp is None:
# Skip sending a response for this request and
# continue to the next requst. The response for
# this request will be sent when the object is
# ready.
continue
else:
get_resp = self.basic_service._get_object(req.get, client_id)
resp = ray_client_pb2.DataResponse(get=get_resp)
elif req_type == "put":
if not self.put_request_chunk_collector.add_chunk(req, req.put):
# Put request still in progress
continue
put_resp = self.basic_service._put_object(
self.put_request_chunk_collector.data,
req.put.client_ref_id,
client_id,
)
self.put_request_chunk_collector.reset()
resp = ray_client_pb2.DataResponse(put=put_resp)
elif req_type == "release":
released = []
for rel_id in req.release.ids:
rel = self.basic_service.release(client_id, rel_id)
released.append(rel)
resp = ray_client_pb2.DataResponse(
release=ray_client_pb2.ReleaseResponse(ok=released)
)
elif req_type == "connection_info":
resp = ray_client_pb2.DataResponse(
connection_info=self._build_connection_response()
)
elif req_type == "prep_runtime_env":
with self.clients_lock:
resp_prep = self.basic_service.PrepRuntimeEnv(
req.prep_runtime_env
)
resp = ray_client_pb2.DataResponse(prep_runtime_env=resp_prep)
elif req_type == "connection_cleanup":
cleanup_requested = True
cleanup_resp = ray_client_pb2.ConnectionCleanupResponse()
resp = ray_client_pb2.DataResponse(connection_cleanup=cleanup_resp)
elif req_type == "acknowledge":
# Clean up acknowledged cache entries
response_cache.cleanup(req.acknowledge.req_id)
continue
elif req_type == "task":
with self.clients_lock:
task = req.task
if not self.client_task_chunk_collector.add_chunk(req, task):
# Not all serialized arguments have arrived
continue
arglist, kwargs = loads_from_client(
self.client_task_chunk_collector.data, self.basic_service
)
self.client_task_chunk_collector.reset()
resp_ticket = self.basic_service.Schedule(
req.task, arglist, kwargs, context
)
resp = ray_client_pb2.DataResponse(task_ticket=resp_ticket)
del arglist
del kwargs
elif req_type == "terminate":
with self.clients_lock:
response = self.basic_service.Terminate(req.terminate, context)
resp = ray_client_pb2.DataResponse(terminate=response)
elif req_type == "list_named_actors":
with self.clients_lock:
response = self.basic_service.ListNamedActors(
req.list_named_actors
)
resp = ray_client_pb2.DataResponse(list_named_actors=response)
else:
raise Exception(
f"Unreachable code: Request type "
f"{req_type} not handled in Datapath"
)
resp.req_id = req.req_id
if _should_cache(req) and reconnect_enabled:
response_cache.update_cache(req.req_id, resp)
yield resp
except Exception as e:
logger.exception("Error in data channel:")
recoverable = _propagate_error_in_context(e, context)
invalid_cache = response_cache.invalidate(e)
if not recoverable or invalid_cache:
context.set_code(grpc.StatusCode.FAILED_PRECONDITION)
# Connection isn't recoverable, skip cleanup
cleanup_requested = True
finally:
logger.debug(f"Stream is broken with client {client_id}")
queue_filler_thread.join(QUEUE_JOIN_SECONDS)
if queue_filler_thread.is_alive():
logger.error(
"Queue filler thread failed to join before timeout: {}".format(
QUEUE_JOIN_SECONDS
)
)
cleanup_delay = self.reconnect_grace_periods.get(client_id)
if not cleanup_requested and cleanup_delay is not None:
logger.debug(
"Cleanup wasn't requested, delaying cleanup by"
f"{cleanup_delay} seconds."
)
# Delay cleanup, since client may attempt a reconnect
# Wait on the "stopped" event in case the grpc server is
# stopped and we can clean up earlier.
self.stopped.wait(timeout=cleanup_delay)
else:
logger.debug("Cleanup was requested, cleaning up immediately.")
with self.clients_lock:
if client_id not in self.client_last_seen:
logger.debug("Connection already cleaned up.")
# Some other connection has already cleaned up this
# this client's session. This can happen if the client
# reconnects and then gracefully shut's down immediately.
return
last_seen = self.client_last_seen[client_id]
if last_seen > start_time:
# The client successfully reconnected and updated
# last seen some time during the grace period
logger.debug("Client reconnected, skipping cleanup")
return
# Either the client shut down gracefully, or the client
# failed to reconnect within the grace period. Clean up
# the connection.
self.basic_service.release_all(client_id)
del self.client_last_seen[client_id]
if client_id in self.reconnect_grace_periods:
del self.reconnect_grace_periods[client_id]
if client_id in self.response_caches:
del self.response_caches[client_id]
self.num_clients -= 1
logger.debug(
f"Removed client {client_id}, " f"remaining={self.num_clients}"
)
# It's important to keep the Ray shutdown
# within this locked context or else Ray could hang.
# NOTE: it is strange to start ray in server.py but shut it
# down here. Consider consolidating ray lifetime management.
with disable_client_hook():
if self.num_clients == 0:
logger.debug("Shutting down ray.")
ray.shutdown()
def _init(self, client_id: str, context: Any, start_time: float):
"""
Checks if resources allow for another client.
Returns a boolean indicating if initialization was successful.
"""
with self.clients_lock:
reconnecting = _get_reconnecting_from_context(context)
threshold = int(CLIENT_SERVER_MAX_THREADS / 2)
if self.num_clients >= threshold:
logger.warning(
f"[Data Servicer]: Num clients {self.num_clients} "
f"has reached the threshold {threshold}. "
f"Rejecting client: {client_id}. "
)
if log_once("client_threshold"):
logger.warning(
"You can configure the client connection "
"threshold by setting the "
"RAY_CLIENT_SERVER_MAX_THREADS env var "
f"(currently set to {CLIENT_SERVER_MAX_THREADS})."
)
context.set_code(grpc.StatusCode.RESOURCE_EXHAUSTED)
return False
if reconnecting and client_id not in self.client_last_seen:
# Client took too long to reconnect, session has been
# cleaned up.
context.set_code(grpc.StatusCode.NOT_FOUND)
context.set_details(
"Attempted to reconnect to a session that has already "
"been cleaned up."
)
return False
if client_id in self.client_last_seen:
logger.debug(f"Client {client_id} has reconnected.")
else:
self.num_clients += 1
logger.debug(
f"Accepted data connection from {client_id}. "
f"Total clients: {self.num_clients}"
)
self.client_last_seen[client_id] = start_time
return True
def _build_connection_response(self):
with self.clients_lock:
cur_num_clients = self.num_clients
return ray_client_pb2.ConnectionInfoResponse(
num_clients=cur_num_clients,
python_version="{}.{}.{}".format(
sys.version_info[0], sys.version_info[1], sys.version_info[2]
),
ray_version=ray.__version__,
ray_commit=ray.__commit__,
)
python/ray/util/client/server/logservicer.py
+125
View File
@@ -0,0 +1,125 @@
"""This file responds to log stream requests and forwards logs
with its handler.
"""
import io
import logging
import queue
import threading
import uuid
import grpc
import ray.core.generated.ray_client_pb2 as ray_client_pb2
import ray.core.generated.ray_client_pb2_grpc as ray_client_pb2_grpc
from ray._private.ray_logging import global_worker_stdstream_dispatcher
from ray._private.worker import print_worker_logs
from ray.util.client.common import CLIENT_SERVER_MAX_THREADS
logger = logging.getLogger(__name__)
class LogstreamHandler(logging.Handler):
def __init__(self, queue, level):
super().__init__()
self.queue = queue
self.level = level
def emit(self, record: logging.LogRecord):
logdata = ray_client_pb2.LogData()
logdata.msg = record.getMessage()
logdata.level = record.levelno
logdata.name = record.name
self.queue.put(logdata)
class StdStreamHandler:
def __init__(self, queue):
self.queue = queue
self.id = str(uuid.uuid4())
def handle(self, data):
logdata = ray_client_pb2.LogData()
logdata.level = -2 if data["is_err"] else -1
logdata.name = "stderr" if data["is_err"] else "stdout"
with io.StringIO() as file:
print_worker_logs(data, file)
logdata.msg = file.getvalue()
self.queue.put(logdata)
def register_global(self):
global_worker_stdstream_dispatcher.add_handler(self.id, self.handle)
def unregister_global(self):
global_worker_stdstream_dispatcher.remove_handler(self.id)
def log_status_change_thread(log_queue, request_iterator):
std_handler = StdStreamHandler(log_queue)
current_handler = None
root_logger = logging.getLogger("ray")
default_level = root_logger.getEffectiveLevel()
try:
for req in request_iterator:
if current_handler is not None:
root_logger.setLevel(default_level)
root_logger.removeHandler(current_handler)
std_handler.unregister_global()
if not req.enabled:
current_handler = None
continue
current_handler = LogstreamHandler(log_queue, req.loglevel)
std_handler.register_global()
root_logger.addHandler(current_handler)
root_logger.setLevel(req.loglevel)
except grpc.RpcError as e:
logger.debug(f"closing log thread " f"grpc error reading request_iterator: {e}")
finally:
if current_handler is not None:
root_logger.setLevel(default_level)
root_logger.removeHandler(current_handler)
std_handler.unregister_global()
log_queue.put(None)
class LogstreamServicer(ray_client_pb2_grpc.RayletLogStreamerServicer):
def __init__(self):
super().__init__()
self.num_clients = 0
self.client_lock = threading.Lock()
def Logstream(self, request_iterator, context):
initialized = False
with self.client_lock:
threshold = CLIENT_SERVER_MAX_THREADS / 2
if self.num_clients + 1 >= threshold:
context.set_code(grpc.StatusCode.RESOURCE_EXHAUSTED)
logger.warning(
f"Logstream: Num clients {self.num_clients} has reached "
f"the threshold {threshold}. Rejecting new connection."
)
return
self.num_clients += 1
initialized = True
logger.info(
"New logs connection established. " f"Total clients: {self.num_clients}"
)
log_queue = queue.Queue()
thread = threading.Thread(
target=log_status_change_thread,
args=(log_queue, request_iterator),
daemon=True,
)
thread.start()
try:
queue_iter = iter(log_queue.get, None)
for record in queue_iter:
if record is None:
break
yield record
except grpc.RpcError as e:
logger.debug(f"Closing log channel: {e}")
finally:
thread.join()
with self.client_lock:
if initialized:
self.num_clients -= 1
python/ray/util/client/server/proxier.py
+938
View File
@@ -0,0 +1,938 @@
import atexit
import json
import logging
import socket
import sys
import time
import traceback
import urllib
from concurrent import futures
from dataclasses import dataclass
from itertools import chain
from threading import Event, Lock, RLock, Thread
from typing import Callable, Dict, List, Optional, Tuple
from urllib.parse import urlparse, urlunparse
import grpc
import ray
import ray.core.generated.ray_client_pb2 as ray_client_pb2
import ray.core.generated.ray_client_pb2_grpc as ray_client_pb2_grpc
import ray.core.generated.runtime_env_agent_pb2 as runtime_env_agent_pb2
from ray._common.network_utils import (
build_address,
get_localhost_ip,
is_ipv6,
is_localhost,
)
from ray._common.tls_utils import add_port_to_grpc_server
from ray._common.utils import env_integer
from ray._private.authentication.http_token_authentication import (
format_authentication_http_error,
get_auth_headers_if_auth_enabled,
)
from ray._private.client_mode_hook import disable_client_hook
from ray._private.grpc_utils import init_grpc_channel
from ray._private.parameter import RayParams
from ray._private.runtime_env.context import RuntimeEnvContext
from ray._private.services import (
ProcessInfo,
get_node_with_retry,
start_ray_client_server,
)
from ray._private.utils import detect_fate_sharing_support
from ray._raylet import GcsClient
from ray.cloudpickle.compat import pickle
from ray.exceptions import AuthenticationError
from ray.job_config import JobConfig
from ray.util.client.common import (
CLIENT_SERVER_MAX_THREADS,
GRPC_OPTIONS,
ClientServerHandle,
_get_client_id_from_context,
_propagate_error_in_context,
)
from ray.util.client.server.dataservicer import _get_reconnecting_from_context
# Import psutil after ray so the packaged version is used.
import psutil
logger = logging.getLogger(__name__)
CHECK_PROCESS_INTERVAL_S = 30
MIN_SPECIFIC_SERVER_PORT = 23000
MAX_SPECIFIC_SERVER_PORT = 24000
CHECK_CHANNEL_TIMEOUT_S = env_integer("RAY_CLIENT_SERVER_CHECK_CHANNEL_TIMEOUT_S", 30)
LOGSTREAM_RETRIES = 5
LOGSTREAM_RETRY_INTERVAL_SEC = 2
@dataclass
class SpecificServer:
port: int
process_handle_future: futures.Future
channel: "grpc._channel.Channel"
def is_ready(self) -> bool:
"""Check if the server is ready or not (doesn't block)."""
return self.process_handle_future.done()
def wait_ready(self, timeout: Optional[float] = None) -> None:
"""
Wait for the server to actually start up.
"""
res = self.process_handle_future.result(timeout=timeout)
if res is None:
# This is only set to none when server creation specifically fails.
raise RuntimeError("Server startup failed.")
def poll(self) -> Optional[int]:
"""Check if the process has exited."""
try:
proc = self.process_handle_future.result(timeout=0.1)
if proc is not None:
return proc.process.poll()
except futures.TimeoutError:
return
def kill(self) -> None:
"""Try to send a KILL signal to the process."""
try:
proc = self.process_handle_future.result(timeout=0.1)
if proc is not None:
proc.process.kill()
except futures.TimeoutError:
# Server has not been started yet.
pass
def set_result(self, proc: Optional[ProcessInfo]) -> None:
"""Set the result of the internal future if it is currently unset."""
if not self.is_ready():
self.process_handle_future.set_result(proc)
def _match_running_client_server(command: List[str]) -> bool:
"""
Detects if the main process in the given command is the RayClient Server.
This works by ensuring that the command is of the form:
<py_executable> -m ray.util.client.server <args>
"""
flattened = " ".join(command)
return "-m ray.util.client.server" in flattened
class ProxyManager:
def __init__(
self,
address: Optional[str],
runtime_env_agent_address: str,
*,
session_dir: Optional[str] = None,
redis_username: Optional[str] = None,
redis_password: Optional[str] = None,
node_id: Optional[str] = None,
):
self.servers: Dict[str, SpecificServer] = dict()
self.server_lock = RLock()
self._address = address
self._redis_username = redis_username
self._redis_password = redis_password
self._free_ports: List[int] = list(
range(MIN_SPECIFIC_SERVER_PORT, MAX_SPECIFIC_SERVER_PORT)
)
if runtime_env_agent_address:
parsed = urlparse(runtime_env_agent_address)
# runtime env agent self-assigns a free port, fetch it from GCS
if parsed.port is None or parsed.port == 0:
if node_id is None:
raise ValueError(
"node_id is required when runtime_env_agent_address "
"has no port specified"
)
node_info = get_node_with_retry(address, node_id)
runtime_env_agent_address = urlunparse(
parsed._replace(
netloc=f"{parsed.hostname}:{node_info['runtime_env_agent_port']}"
)
)
self._runtime_env_agent_address = runtime_env_agent_address
self._check_thread = Thread(target=self._check_processes, daemon=True)
self._check_thread.start()
self.fate_share = bool(detect_fate_sharing_support())
self._node: Optional[ray._private.node.Node] = None
atexit.register(self._cleanup)
def _get_unused_port(self, family: int = socket.AF_INET) -> int:
"""
Search for a port in _free_ports that is unused.
"""
with self.server_lock:
num_ports = len(self._free_ports)
for _ in range(num_ports):
port = self._free_ports.pop(0)
s = socket.socket(family, socket.SOCK_STREAM)
try:
s.bind(("", port))
except OSError:
self._free_ports.append(port)
continue
finally:
s.close()
return port
raise RuntimeError("Unable to succeed in selecting a random port.")
@property
def address(self) -> str:
"""
Returns the provided Ray bootstrap address, or creates a new cluster.
"""
if self._address:
return self._address
# Start a new, locally scoped cluster.
connection_tuple = ray.init()
self._address = connection_tuple["address"]
self._session_dir = connection_tuple["session_dir"]
return self._address
@property
def node(self) -> ray._private.node.Node:
"""Gets a 'ray.Node' object for this node (the head node).
If it does not already exist, one is created using the bootstrap
address.
"""
if self._node:
return self._node
ray_params = RayParams(gcs_address=self.address)
self._node = ray._private.node.Node(
ray_params,
head=False,
shutdown_at_exit=False,
spawn_reaper=False,
connect_only=True,
)
return self._node
def create_specific_server(self, client_id: str) -> SpecificServer:
"""
Create, but not start a SpecificServer for a given client. This
method must be called once per client.
"""
with self.server_lock:
assert (
self.servers.get(client_id) is None
), f"Server already created for Client: {client_id}"
host = get_localhost_ip()
port = self._get_unused_port(
socket.AF_INET6 if is_ipv6(host) else socket.AF_INET
)
server = SpecificServer(
port=port,
process_handle_future=futures.Future(),
channel=init_grpc_channel(
build_address(host, port), options=GRPC_OPTIONS
),
)
self.servers[client_id] = server
return server
def _create_runtime_env(
self,
serialized_runtime_env: str,
runtime_env_config: str,
specific_server: SpecificServer,
):
"""Increase the runtime_env reference by sending an RPC to the agent.
Includes retry logic to handle the case when the agent is
temporarily unreachable (e.g., hasn't been started up yet).
"""
logger.info(
f"Increasing runtime env reference for "
f"ray_client_server_{specific_server.port}."
f"Serialized runtime env is {serialized_runtime_env}."
)
assert (
len(self._runtime_env_agent_address) > 0
), "runtime_env_agent_address not set"
create_env_request = runtime_env_agent_pb2.GetOrCreateRuntimeEnvRequest(
serialized_runtime_env=serialized_runtime_env,
runtime_env_config=runtime_env_config,
job_id=f"ray_client_server_{specific_server.port}".encode("utf-8"),
source_process="client_server",
)
retries = 0
max_retries = 5
wait_time_s = 0.5
last_exception = None
while retries <= max_retries:
try:
url = urllib.parse.urljoin(
self._runtime_env_agent_address, "/get_or_create_runtime_env"
)
data = create_env_request.SerializeToString()
headers = {"Content-Type": "application/octet-stream"}
headers.update(**get_auth_headers_if_auth_enabled(headers))
req = urllib.request.Request(
url, data=data, method="POST", headers=headers
)
response = urllib.request.urlopen(req, timeout=None)
response_data = response.read()
r = runtime_env_agent_pb2.GetOrCreateRuntimeEnvReply()
r.ParseFromString(response_data)
if r.status == runtime_env_agent_pb2.AgentRpcStatus.AGENT_RPC_STATUS_OK:
return r.serialized_runtime_env_context
elif (
r.status
== runtime_env_agent_pb2.AgentRpcStatus.AGENT_RPC_STATUS_FAILED
):
raise RuntimeError(
"Failed to create runtime_env for Ray client "
f"server, it is caused by:\n{r.error_message}"
)
else:
assert False, f"Unknown status: {r.status}."
except urllib.error.HTTPError as e:
body = ""
try:
body = e.read().decode("utf-8", "ignore")
except Exception:
body = e.reason if hasattr(e, "reason") else str(e)
formatted_error = format_authentication_http_error(e.code, body or "")
if formatted_error:
raise AuthenticationError(formatted_error) from e
# Treat non-auth HTTP errors like URLError (retry with backoff)
last_exception = e
logger.warning(
f"GetOrCreateRuntimeEnv request failed with HTTP {e.code}: {body or e}. "
f"Retrying after {wait_time_s}s. "
f"{max_retries-retries} retries remaining."
)
except urllib.error.URLError as e:
last_exception = e
logger.warning(
f"GetOrCreateRuntimeEnv request failed: {e}. "
f"Retrying after {wait_time_s}s. "
f"{max_retries-retries} retries remaining."
)
# Exponential backoff.
time.sleep(wait_time_s)
retries += 1
wait_time_s *= 2
raise TimeoutError(
f"GetOrCreateRuntimeEnv request failed after {max_retries} attempts."
f" Last exception: {last_exception}"
)
def start_specific_server(self, client_id: str, job_config: JobConfig) -> bool:
"""
Start up a RayClient Server for an incoming client to
communicate with. Returns whether creation was successful.
"""
specific_server = self._get_server_for_client(client_id)
assert specific_server, f"Server has not been created for: {client_id}"
output, error = self.node.get_log_file_handles(
f"ray_client_server_{specific_server.port}", unique=True
)
serialized_runtime_env = job_config._get_serialized_runtime_env()
runtime_env_config = job_config._get_proto_runtime_env_config()
if not serialized_runtime_env or serialized_runtime_env == "{}":
# TODO(edoakes): can we just remove this case and always send it
# to the agent?
serialized_runtime_env_context = RuntimeEnvContext().serialize()
else:
serialized_runtime_env_context = self._create_runtime_env(
serialized_runtime_env=serialized_runtime_env,
runtime_env_config=runtime_env_config,
specific_server=specific_server,
)
proc = start_ray_client_server(
self.address,
get_localhost_ip(),
specific_server.port,
stdout_file=output,
stderr_file=error,
fate_share=self.fate_share,
server_type="specific-server",
serialized_runtime_env_context=serialized_runtime_env_context,
redis_username=self._redis_username,
redis_password=self._redis_password,
)
# Wait for the process being run transitions from the shim process
# to the actual RayClient Server.
pid = proc.process.pid
if sys.platform != "win32":
psutil_proc = psutil.Process(pid)
else:
psutil_proc = None
# Don't use `psutil` on Win32
while psutil_proc is not None:
if proc.process.poll() is not None:
logger.error(f"SpecificServer startup failed for client: {client_id}")
break
cmd = psutil_proc.cmdline()
if _match_running_client_server(cmd):
break
logger.debug("Waiting for Process to reach the actual client server.")
time.sleep(0.5)
specific_server.set_result(proc)
logger.info(
f"SpecificServer started on port: {specific_server.port} "
f"with PID: {pid} for client: {client_id}"
)
return proc.process.poll() is None
def _get_server_for_client(self, client_id: str) -> Optional[SpecificServer]:
with self.server_lock:
client = self.servers.get(client_id)
if client is None:
logger.error(f"Unable to find channel for client: {client_id}")
return client
def has_channel(self, client_id: str) -> bool:
server = self._get_server_for_client(client_id)
if server is None:
return False
return server.is_ready()
def get_channel(
self,
client_id: str,
) -> Optional["grpc._channel.Channel"]:
"""
Find the gRPC Channel for the given client_id. This will block until
the server process has started.
"""
server = self._get_server_for_client(client_id)
if server is None:
return None
# Wait for the SpecificServer to become ready.
server.wait_ready()
try:
grpc.channel_ready_future(server.channel).result(
timeout=CHECK_CHANNEL_TIMEOUT_S
)
return server.channel
except grpc.FutureTimeoutError:
logger.exception(f"Timeout waiting for channel for {client_id}")
return None
def _check_processes(self):
"""
Keeps the internal servers dictionary up-to-date with running servers.
"""
while True:
with self.server_lock:
for client_id, specific_server in list(self.servers.items()):
if specific_server.poll() is not None:
logger.info(
f"Specific server {client_id} is no longer running"
f", freeing its port {specific_server.port}"
)
del self.servers[client_id]
# Port is available to use again.
self._free_ports.append(specific_server.port)
time.sleep(CHECK_PROCESS_INTERVAL_S)
def _cleanup(self) -> None:
"""
Forcibly kill all spawned RayClient Servers. This ensures cleanup
for platforms where fate sharing is not supported.
"""
for server in self.servers.values():
server.kill()
class RayletServicerProxy(ray_client_pb2_grpc.RayletDriverServicer):
def __init__(self, ray_connect_handler: Callable, proxy_manager: ProxyManager):
self.proxy_manager = proxy_manager
self.ray_connect_handler = ray_connect_handler
def _call_inner_function(
self, request, context, method: str
) -> Optional[ray_client_pb2_grpc.RayletDriverStub]:
client_id = _get_client_id_from_context(context)
chan = self.proxy_manager.get_channel(client_id)
if not chan:
logger.error(f"Channel for Client: {client_id} not found!")
context.set_code(grpc.StatusCode.NOT_FOUND)
return None
stub = ray_client_pb2_grpc.RayletDriverStub(chan)
try:
metadata = [("client_id", client_id)]
if context:
metadata = context.invocation_metadata()
return getattr(stub, method)(request, metadata=metadata)
except Exception as e:
# Error while proxying -- propagate the error's context to user
logger.exception(f"Proxying call to {method} failed!")
_propagate_error_in_context(e, context)
def _has_channel_for_request(self, context):
client_id = _get_client_id_from_context(context)
return self.proxy_manager.has_channel(client_id)
def Init(self, request, context=None) -> ray_client_pb2.InitResponse:
return self._call_inner_function(request, context, "Init")
def KVPut(self, request, context=None) -> ray_client_pb2.KVPutResponse:
"""Proxies internal_kv.put.
This is used by the working_dir code to upload to the GCS before
ray.init is called. In that case (if we don't have a server yet)
we directly make the internal KV call from the proxier.
Otherwise, we proxy the call to the downstream server as usual.
"""
if self._has_channel_for_request(context):
return self._call_inner_function(request, context, "KVPut")
with disable_client_hook():
already_exists = ray.experimental.internal_kv._internal_kv_put(
request.key, request.value, overwrite=request.overwrite
)
return ray_client_pb2.KVPutResponse(already_exists=already_exists)
def KVGet(self, request, context=None) -> ray_client_pb2.KVGetResponse:
"""Proxies internal_kv.get.
This is used by the working_dir code to upload to the GCS before
ray.init is called. In that case (if we don't have a server yet)
we directly make the internal KV call from the proxier.
Otherwise, we proxy the call to the downstream server as usual.
"""
if self._has_channel_for_request(context):
return self._call_inner_function(request, context, "KVGet")
with disable_client_hook():
value = ray.experimental.internal_kv._internal_kv_get(request.key)
return ray_client_pb2.KVGetResponse(value=value)
def KVDel(self, request, context=None) -> ray_client_pb2.KVDelResponse:
"""Proxies internal_kv.delete.
This is used by the working_dir code to upload to the GCS before
ray.init is called. In that case (if we don't have a server yet)
we directly make the internal KV call from the proxier.
Otherwise, we proxy the call to the downstream server as usual.
"""
if self._has_channel_for_request(context):
return self._call_inner_function(request, context, "KVDel")
with disable_client_hook():
ray.experimental.internal_kv._internal_kv_del(request.key)
return ray_client_pb2.KVDelResponse()
def KVList(self, request, context=None) -> ray_client_pb2.KVListResponse:
"""Proxies internal_kv.list.
This is used by the working_dir code to upload to the GCS before
ray.init is called. In that case (if we don't have a server yet)
we directly make the internal KV call from the proxier.
Otherwise, we proxy the call to the downstream server as usual.
"""
if self._has_channel_for_request(context):
return self._call_inner_function(request, context, "KVList")
with disable_client_hook():
keys = ray.experimental.internal_kv._internal_kv_list(request.prefix)
return ray_client_pb2.KVListResponse(keys=keys)
def KVExists(self, request, context=None) -> ray_client_pb2.KVExistsResponse:
"""Proxies internal_kv.exists.
This is used by the working_dir code to upload to the GCS before
ray.init is called. In that case (if we don't have a server yet)
we directly make the internal KV call from the proxier.
Otherwise, we proxy the call to the downstream server as usual.
"""
if self._has_channel_for_request(context):
return self._call_inner_function(request, context, "KVExists")
with disable_client_hook():
exists = ray.experimental.internal_kv._internal_kv_exists(request.key)
return ray_client_pb2.KVExistsResponse(exists=exists)
def PinRuntimeEnvURI(
self, request, context=None
) -> ray_client_pb2.ClientPinRuntimeEnvURIResponse:
"""Proxies internal_kv.pin_runtime_env_uri.
This is used by the working_dir code to upload to the GCS before
ray.init is called. In that case (if we don't have a server yet)
we directly make the internal KV call from the proxier.
Otherwise, we proxy the call to the downstream server as usual.
"""
if self._has_channel_for_request(context):
return self._call_inner_function(request, context, "PinRuntimeEnvURI")
with disable_client_hook():
ray.experimental.internal_kv._pin_runtime_env_uri(
request.uri, expiration_s=request.expiration_s
)
return ray_client_pb2.ClientPinRuntimeEnvURIResponse()
def ListNamedActors(
self, request, context=None
) -> ray_client_pb2.ClientListNamedActorsResponse:
return self._call_inner_function(request, context, "ListNamedActors")
def ClusterInfo(self, request, context=None) -> ray_client_pb2.ClusterInfoResponse:
# NOTE: We need to respond to the PING request here to allow the client
# to continue with connecting.
if request.type == ray_client_pb2.ClusterInfoType.PING:
resp = ray_client_pb2.ClusterInfoResponse(json=json.dumps({}))
return resp
return self._call_inner_function(request, context, "ClusterInfo")
def Terminate(self, req, context=None):
return self._call_inner_function(req, context, "Terminate")
def GetObject(self, request, context=None):
try:
yield from self._call_inner_function(request, context, "GetObject")
except Exception as e:
# Error while iterating over response from GetObject stream
logger.exception("Proxying call to GetObject failed!")
_propagate_error_in_context(e, context)
def PutObject(
self, request: ray_client_pb2.PutRequest, context=None
) -> ray_client_pb2.PutResponse:
return self._call_inner_function(request, context, "PutObject")
def WaitObject(self, request, context=None) -> ray_client_pb2.WaitResponse:
return self._call_inner_function(request, context, "WaitObject")
def Schedule(self, task, context=None) -> ray_client_pb2.ClientTaskTicket:
return self._call_inner_function(task, context, "Schedule")
def ray_client_server_env_prep(job_config: JobConfig) -> JobConfig:
return job_config
def prepare_runtime_init_req(
init_request: ray_client_pb2.DataRequest,
) -> Tuple[ray_client_pb2.DataRequest, JobConfig]:
"""
Extract JobConfig and possibly mutate InitRequest before it is passed to
the specific RayClient Server.
"""
init_type = init_request.WhichOneof("type")
assert init_type == "init", (
"Received initial message of type " f"{init_type}, not 'init'."
)
req = init_request.init
job_config = JobConfig()
if req.job_config:
job_config = pickle.loads(req.job_config)
new_job_config = ray_client_server_env_prep(job_config)
modified_init_req = ray_client_pb2.InitRequest(
job_config=pickle.dumps(new_job_config),
ray_init_kwargs=init_request.init.ray_init_kwargs,
reconnect_grace_period=init_request.init.reconnect_grace_period,
)
init_request.init.CopyFrom(modified_init_req)
return (init_request, new_job_config)
class RequestIteratorProxy:
def __init__(self, request_iterator):
self.request_iterator = request_iterator
def __iter__(self):
return self
def __next__(self):
try:
return next(self.request_iterator)
except grpc.RpcError as e:
# To stop proxying already CANCLLED request stream gracefully,
# we only translate the exact grpc.RpcError to StopIteration,
# not its subsclasses. ex: grpc._Rendezvous
# https://github.com/grpc/grpc/blob/v1.43.0/src/python/grpcio/grpc/_server.py#L353-L354
# This fixes the https://github.com/ray-project/ray/issues/23865
if type(e) is not grpc.RpcError:
raise e # re-raise other grpc exceptions
logger.exception(
"Stop iterating cancelled request stream with the following exception:"
)
raise StopIteration
class DataServicerProxy(ray_client_pb2_grpc.RayletDataStreamerServicer):
def __init__(self, proxy_manager: ProxyManager):
self.num_clients = 0
# dictionary mapping client_id's to the last time they connected
self.clients_last_seen: Dict[str, float] = {}
self.reconnect_grace_periods: Dict[str, float] = {}
self.clients_lock = Lock()
self.proxy_manager = proxy_manager
self.stopped = Event()
def modify_connection_info_resp(
self, init_resp: ray_client_pb2.DataResponse
) -> ray_client_pb2.DataResponse:
"""
Modify the `num_clients` returned the ConnectionInfoResponse because
individual SpecificServers only have **one** client.
"""
init_type = init_resp.WhichOneof("type")
if init_type != "connection_info":
return init_resp
modified_resp = ray_client_pb2.DataResponse()
modified_resp.CopyFrom(init_resp)
with self.clients_lock:
modified_resp.connection_info.num_clients = self.num_clients
return modified_resp
def Datapath(self, request_iterator, context):
request_iterator = RequestIteratorProxy(request_iterator)
cleanup_requested = False
start_time = time.time()
client_id = _get_client_id_from_context(context)
if client_id == "":
return
reconnecting = _get_reconnecting_from_context(context)
if reconnecting:
with self.clients_lock:
if client_id not in self.clients_last_seen:
# Client took too long to reconnect, session has already
# been cleaned up
context.set_code(grpc.StatusCode.NOT_FOUND)
context.set_details(
"Attempted to reconnect a session that has already "
"been cleaned up"
)
return
self.clients_last_seen[client_id] = start_time
server = self.proxy_manager._get_server_for_client(client_id)
channel = self.proxy_manager.get_channel(client_id)
# iterator doesn't need modification on reconnect
new_iter = request_iterator
else:
# Create Placeholder *before* reading the first request.
server = self.proxy_manager.create_specific_server(client_id)
with self.clients_lock:
self.clients_last_seen[client_id] = start_time
self.num_clients += 1
try:
if not reconnecting:
logger.info(f"New data connection from client {client_id}: ")
init_req = next(request_iterator)
with self.clients_lock:
self.reconnect_grace_periods[
client_id
] = init_req.init.reconnect_grace_period
try:
modified_init_req, job_config = prepare_runtime_init_req(init_req)
if not self.proxy_manager.start_specific_server(
client_id, job_config
):
logger.error(
f"Server startup failed for client: {client_id}, "
f"using JobConfig: {job_config}!"
)
raise RuntimeError(
"Starting Ray client server failed. See "
f"ray_client_server_{server.port}.err for "
"detailed logs."
)
channel = self.proxy_manager.get_channel(client_id)
if channel is None:
logger.error(f"Channel not found for {client_id}")
raise RuntimeError(
"Proxy failed to Connect to backend! Check "
"`ray_client_server.err` and "
f"`ray_client_server_{server.port}.err` on the "
"head node of the cluster for the relevant logs. "
"By default these are located at "
"/tmp/ray/session_latest/logs."
)
except Exception:
init_resp = ray_client_pb2.DataResponse(
init=ray_client_pb2.InitResponse(
ok=False, msg=traceback.format_exc()
)
)
init_resp.req_id = init_req.req_id
yield init_resp
return None
new_iter = chain([modified_init_req], request_iterator)
stub = ray_client_pb2_grpc.RayletDataStreamerStub(channel)
metadata = [("client_id", client_id), ("reconnecting", str(reconnecting))]
resp_stream = stub.Datapath(new_iter, metadata=metadata)
for resp in resp_stream:
resp_type = resp.WhichOneof("type")
if resp_type == "connection_cleanup":
# Specific server is skipping cleanup, proxier should too
cleanup_requested = True
yield self.modify_connection_info_resp(resp)
except Exception as e:
logger.exception("Proxying Datapath failed!")
# Propogate error through context
recoverable = _propagate_error_in_context(e, context)
if not recoverable:
# Client shouldn't attempt to recover, clean up connection
cleanup_requested = True
finally:
cleanup_delay = self.reconnect_grace_periods.get(client_id)
if not cleanup_requested and cleanup_delay is not None:
# Delay cleanup, since client may attempt a reconnect
# Wait on stopped event in case the server closes and we
# can clean up earlier
self.stopped.wait(timeout=cleanup_delay)
with self.clients_lock:
if client_id not in self.clients_last_seen:
logger.info(f"{client_id} not found. Skipping clean up.")
# Connection has already been cleaned up
return
last_seen = self.clients_last_seen[client_id]
logger.info(
f"{client_id} last started stream at {last_seen}. Current "
f"stream started at {start_time}."
)
if last_seen > start_time:
logger.info("Client reconnected. Skipping cleanup.")
# Client has reconnected, don't clean up
return
logger.debug(f"Client detached: {client_id}")
self.num_clients -= 1
del self.clients_last_seen[client_id]
if client_id in self.reconnect_grace_periods:
del self.reconnect_grace_periods[client_id]
server.set_result(None)
class LogstreamServicerProxy(ray_client_pb2_grpc.RayletLogStreamerServicer):
def __init__(self, proxy_manager: ProxyManager):
super().__init__()
self.proxy_manager = proxy_manager
def Logstream(self, request_iterator, context):
request_iterator = RequestIteratorProxy(request_iterator)
client_id = _get_client_id_from_context(context)
if client_id == "":
return
logger.debug(f"New logstream connection from client {client_id}: ")
channel = None
# We need to retry a few times because the LogClient *may* connect
# Before the DataClient has finished connecting.
for i in range(LOGSTREAM_RETRIES):
channel = self.proxy_manager.get_channel(client_id)
if channel is not None:
break
logger.warning(f"Retrying Logstream connection. {i+1} attempts failed.")
time.sleep(LOGSTREAM_RETRY_INTERVAL_SEC)
if channel is None:
context.set_code(grpc.StatusCode.NOT_FOUND)
context.set_details(
"Logstream proxy failed to connect. Channel for client "
f"{client_id} not found."
)
return None
stub = ray_client_pb2_grpc.RayletLogStreamerStub(channel)
resp_stream = stub.Logstream(
request_iterator, metadata=[("client_id", client_id)]
)
try:
for resp in resp_stream:
yield resp
except Exception:
logger.exception("Proxying Logstream failed!")
def serve_proxier(
host: str,
port: int,
gcs_address: Optional[str],
*,
redis_username: Optional[str] = None,
redis_password: Optional[str] = None,
session_dir: Optional[str] = None,
runtime_env_agent_address: Optional[str] = None,
node_id: Optional[str] = None,
):
# Initialize internal KV to be used to upload and download working_dir
# before calling ray.init within the RayletServicers.
# NOTE(edoakes): redis_address and redis_password should only be None in
# tests.
if gcs_address is not None:
gcs_cli = GcsClient(address=gcs_address)
ray.experimental.internal_kv._initialize_internal_kv(gcs_cli)
from ray._private.grpc_utils import create_grpc_server_with_interceptors
server = create_grpc_server_with_interceptors(
max_workers=CLIENT_SERVER_MAX_THREADS,
thread_name_prefix="ray_client_proxier",
options=GRPC_OPTIONS,
asynchronous=False,
)
proxy_manager = ProxyManager(
gcs_address,
session_dir=session_dir,
redis_username=redis_username,
redis_password=redis_password,
runtime_env_agent_address=runtime_env_agent_address,
node_id=node_id,
)
task_servicer = RayletServicerProxy(None, proxy_manager)
data_servicer = DataServicerProxy(proxy_manager)
logs_servicer = LogstreamServicerProxy(proxy_manager)
ray_client_pb2_grpc.add_RayletDriverServicer_to_server(task_servicer, server)
ray_client_pb2_grpc.add_RayletDataStreamerServicer_to_server(data_servicer, server)
ray_client_pb2_grpc.add_RayletLogStreamerServicer_to_server(logs_servicer, server)
if not is_localhost(host):
add_port_to_grpc_server(server, build_address(get_localhost_ip(), port))
add_port_to_grpc_server(server, build_address(host, port))
server.start()
return ClientServerHandle(
task_servicer=task_servicer,
data_servicer=data_servicer,
logs_servicer=logs_servicer,
grpc_server=server,
)
python/ray/util/client/server/server.py
+968
View File
@@ -0,0 +1,968 @@
import base64
import functools
import gc
import inspect
import json
import logging
import math
import os
import pickle
import queue
import threading
import time
from collections import defaultdict
from typing import Any, Callable, Dict, List, Optional, Set, Union
import grpc
import ray
import ray._private.state
import ray.core.generated.ray_client_pb2 as ray_client_pb2
import ray.core.generated.ray_client_pb2_grpc as ray_client_pb2_grpc
from ray import cloudpickle
from ray._common.network_utils import (
build_address,
get_all_interfaces_ip,
get_localhost_ip,
is_localhost,
)
from ray._common.tls_utils import add_port_to_grpc_server
from ray._private import ray_constants
from ray._private.client_mode_hook import disable_client_hook
from ray._private.ray_constants import env_integer
from ray._private.ray_logging import setup_logger
from ray._private.ray_logging.logging_config import LoggingConfig
from ray._private.services import canonicalize_bootstrap_address_or_die
from ray._raylet import GcsClient
from ray.job_config import JobConfig
from ray.util.client.common import (
CLIENT_SERVER_MAX_THREADS,
GRPC_OPTIONS,
OBJECT_TRANSFER_CHUNK_SIZE,
ClientServerHandle,
ResponseCache,
)
from ray.util.client.server.dataservicer import DataServicer
from ray.util.client.server.logservicer import LogstreamServicer
from ray.util.client.server.proxier import serve_proxier
from ray.util.client.server.server_pickler import dumps_from_server, loads_from_client
from ray.util.client.server.server_stubs import current_server
logger = logging.getLogger(__name__)
TIMEOUT_FOR_SPECIFIC_SERVER_S = env_integer("TIMEOUT_FOR_SPECIFIC_SERVER_S", 30)
def _use_response_cache(func):
"""
Decorator for gRPC stubs. Before calling the real stubs, checks if there's
an existing entry in the caches. If there is, then return the cached
entry. Otherwise, call the real function and use the real cache
"""
@functools.wraps(func)
def wrapper(self, request, context):
metadata = dict(context.invocation_metadata())
expected_ids = ("client_id", "thread_id", "req_id")
if any(i not in metadata for i in expected_ids):
# Missing IDs, skip caching and call underlying stub directly
return func(self, request, context)
# Get relevant IDs to check cache
client_id = metadata["client_id"]
thread_id = metadata["thread_id"]
req_id = int(metadata["req_id"])
# Check if response already cached
response_cache = self.response_caches[client_id]
cached_entry = response_cache.check_cache(thread_id, req_id)
if cached_entry is not None:
if isinstance(cached_entry, Exception):
# Original call errored, propagate error
context.set_code(grpc.StatusCode.FAILED_PRECONDITION)
context.set_details(str(cached_entry))
raise cached_entry
return cached_entry
try:
# Response wasn't cached, call underlying stub and cache result
resp = func(self, request, context)
except Exception as e:
# Unexpected error in underlying stub -- update cache and
# propagate to user through context
response_cache.update_cache(thread_id, req_id, e)
context.set_code(grpc.StatusCode.FAILED_PRECONDITION)
context.set_details(str(e))
raise
response_cache.update_cache(thread_id, req_id, resp)
return resp
return wrapper
class RayletServicer(ray_client_pb2_grpc.RayletDriverServicer):
def __init__(self, ray_connect_handler: Callable):
"""Construct a raylet service
Args:
ray_connect_handler: Function to connect to ray cluster
"""
# Stores client_id -> (ref_id -> ObjectRef)
self.object_refs: Dict[str, Dict[bytes, ray.ObjectRef]] = defaultdict(dict)
# Stores client_id -> (client_ref_id -> ref_id (in self.object_refs))
self.client_side_ref_map: Dict[str, Dict[bytes, bytes]] = defaultdict(dict)
self.function_refs = {}
self.actor_refs: Dict[bytes, ray.ActorHandle] = {}
self.actor_owners: Dict[str, Set[bytes]] = defaultdict(set)
self.registered_actor_classes = {}
self.named_actors = set()
self.state_lock = threading.Lock()
self.ray_connect_handler = ray_connect_handler
self.response_caches: Dict[str, ResponseCache] = defaultdict(ResponseCache)
def Init(
self, request: ray_client_pb2.InitRequest, context=None
) -> ray_client_pb2.InitResponse:
if request.job_config:
job_config = pickle.loads(request.job_config)
job_config._client_job = True
else:
job_config = None
current_job_config = None
with disable_client_hook():
if ray.is_initialized():
worker = ray._private.worker.global_worker
current_job_config = worker.core_worker.get_job_config()
else:
extra_kwargs = json.loads(request.ray_init_kwargs or "{}")
# Reconstruct LoggingConfig from dict after InitRequest ray_init_kwargs is parsed from JSON on the server.
if "logging_config" in extra_kwargs and isinstance(
extra_kwargs["logging_config"], dict
):
extra_kwargs["logging_config"] = LoggingConfig.from_dict(
extra_kwargs["logging_config"]
)
try:
self.ray_connect_handler(job_config, **extra_kwargs)
except Exception as e:
logger.exception("Running Ray Init failed:")
return ray_client_pb2.InitResponse(
ok=False,
msg=f"Call to `ray.init()` on the server failed with: {e}",
)
if job_config is None:
return ray_client_pb2.InitResponse(ok=True)
# NOTE(edoakes): this code should not be necessary anymore because we
# only allow a single client/job per server. There is an existing test
# that tests the behavior of multiple clients with the same job config
# connecting to one server (test_client_init.py::test_num_clients),
# so I'm leaving it here for now.
job_config = job_config._get_proto_job_config()
# If the server has been initialized, we need to compare whether the
# runtime env is compatible.
if current_job_config:
job_uris = set(job_config.runtime_env_info.uris.working_dir_uri)
job_uris.update(job_config.runtime_env_info.uris.py_modules_uris)
current_job_uris = set(
current_job_config.runtime_env_info.uris.working_dir_uri
)
current_job_uris.update(
current_job_config.runtime_env_info.uris.py_modules_uris
)
if job_uris != current_job_uris and len(job_uris) > 0:
return ray_client_pb2.InitResponse(
ok=False,
msg="Runtime environment doesn't match "
f"request one {job_config.runtime_env_info.uris} "
f"current one {current_job_config.runtime_env_info.uris}",
)
return ray_client_pb2.InitResponse(ok=True)
@_use_response_cache
def KVPut(self, request, context=None) -> ray_client_pb2.KVPutResponse:
try:
with disable_client_hook():
already_exists = ray.experimental.internal_kv._internal_kv_put(
request.key,
request.value,
overwrite=request.overwrite,
namespace=request.namespace,
)
except Exception as e:
return_exception_in_context(e, context)
already_exists = False
return ray_client_pb2.KVPutResponse(already_exists=already_exists)
def KVGet(self, request, context=None) -> ray_client_pb2.KVGetResponse:
try:
with disable_client_hook():
value = ray.experimental.internal_kv._internal_kv_get(
request.key, namespace=request.namespace
)
except Exception as e:
return_exception_in_context(e, context)
value = b""
return ray_client_pb2.KVGetResponse(value=value)
@_use_response_cache
def KVDel(self, request, context=None) -> ray_client_pb2.KVDelResponse:
try:
with disable_client_hook():
deleted_num = ray.experimental.internal_kv._internal_kv_del(
request.key,
del_by_prefix=request.del_by_prefix,
namespace=request.namespace,
)
except Exception as e:
return_exception_in_context(e, context)
deleted_num = 0
return ray_client_pb2.KVDelResponse(deleted_num=deleted_num)
def KVList(self, request, context=None) -> ray_client_pb2.KVListResponse:
try:
with disable_client_hook():
keys = ray.experimental.internal_kv._internal_kv_list(
request.prefix, namespace=request.namespace
)
except Exception as e:
return_exception_in_context(e, context)
keys = []
return ray_client_pb2.KVListResponse(keys=keys)
def KVExists(self, request, context=None) -> ray_client_pb2.KVExistsResponse:
try:
with disable_client_hook():
exists = ray.experimental.internal_kv._internal_kv_exists(
request.key, namespace=request.namespace
)
except Exception as e:
return_exception_in_context(e, context)
exists = False
return ray_client_pb2.KVExistsResponse(exists=exists)
def ListNamedActors(
self, request, context=None
) -> ray_client_pb2.ClientListNamedActorsResponse:
with disable_client_hook():
actors = ray.util.list_named_actors(all_namespaces=request.all_namespaces)
return ray_client_pb2.ClientListNamedActorsResponse(
actors_json=json.dumps(actors)
)
def ClusterInfo(self, request, context=None) -> ray_client_pb2.ClusterInfoResponse:
resp = ray_client_pb2.ClusterInfoResponse()
resp.type = request.type
if request.type == ray_client_pb2.ClusterInfoType.CLUSTER_RESOURCES:
with disable_client_hook():
resources = ray.cluster_resources()
# Normalize resources into floats
# (the function may return values that are ints)
float_resources = {k: float(v) for k, v in resources.items()}
resp.resource_table.CopyFrom(
ray_client_pb2.ClusterInfoResponse.ResourceTable(table=float_resources)
)
elif request.type == ray_client_pb2.ClusterInfoType.AVAILABLE_RESOURCES:
with disable_client_hook():
resources = ray.available_resources()
# Normalize resources into floats
# (the function may return values that are ints)
float_resources = {k: float(v) for k, v in resources.items()}
resp.resource_table.CopyFrom(
ray_client_pb2.ClusterInfoResponse.ResourceTable(table=float_resources)
)
elif request.type == ray_client_pb2.ClusterInfoType.RUNTIME_CONTEXT:
ctx = ray_client_pb2.ClusterInfoResponse.RuntimeContext()
with disable_client_hook():
rtc = ray.get_runtime_context()
ctx.job_id = ray._common.utils.hex_to_binary(rtc.get_job_id())
ctx.node_id = ray._common.utils.hex_to_binary(rtc.get_node_id())
ctx.worker_id = ray._common.utils.hex_to_binary(rtc.get_worker_id())
ctx.namespace = rtc.namespace
ctx.capture_client_tasks = (
rtc.should_capture_child_tasks_in_placement_group
)
ctx.gcs_address = rtc.gcs_address
ctx.runtime_env = rtc.get_runtime_env_string()
ctx.session_name = rtc.get_session_name()
resp.runtime_context.CopyFrom(ctx)
else:
with disable_client_hook():
resp.json = self._return_debug_cluster_info(request, context)
return resp
def _return_debug_cluster_info(self, request, context=None) -> str:
"""Handle ClusterInfo requests that only return a json blob."""
data = None
if request.type == ray_client_pb2.ClusterInfoType.NODES:
data = ray.nodes()
elif request.type == ray_client_pb2.ClusterInfoType.IS_INITIALIZED:
data = ray.is_initialized()
elif request.type == ray_client_pb2.ClusterInfoType.TIMELINE:
data = ray.timeline()
elif request.type == ray_client_pb2.ClusterInfoType.PING:
data = {}
elif request.type == ray_client_pb2.ClusterInfoType.DASHBOARD_URL:
data = {"dashboard_url": ray._private.worker.get_dashboard_url()}
else:
raise TypeError("Unsupported cluster info type")
return json.dumps(data)
def release(self, client_id: str, id: bytes) -> bool:
with self.state_lock:
if client_id in self.object_refs:
if id in self.object_refs[client_id]:
logger.debug(f"Releasing object {id.hex()} for {client_id}")
del self.object_refs[client_id][id]
return True
if client_id in self.actor_owners:
if id in self.actor_owners[client_id]:
logger.debug(f"Releasing actor {id.hex()} for {client_id}")
self.actor_owners[client_id].remove(id)
if self._can_remove_actor_ref(id):
logger.debug(f"Deleting reference to actor {id.hex()}")
del self.actor_refs[id]
return True
return False
def release_all(self, client_id):
with self.state_lock:
self._release_objects(client_id)
self._release_actors(client_id)
# NOTE: Try to actually dereference the object and actor refs.
# Otherwise dereferencing will happen later, which may run concurrently
# with ray.shutdown() and will crash the process. The crash is a bug
# that should be fixed eventually.
gc.collect()
def _can_remove_actor_ref(self, actor_id_bytes):
no_owner = not any(
actor_id_bytes in actor_list for actor_list in self.actor_owners.values()
)
return no_owner and actor_id_bytes not in self.named_actors
def _release_objects(self, client_id):
if client_id not in self.object_refs:
logger.debug(f"Releasing client with no references: {client_id}")
return
count = len(self.object_refs[client_id])
del self.object_refs[client_id]
if client_id in self.client_side_ref_map:
del self.client_side_ref_map[client_id]
if client_id in self.response_caches:
del self.response_caches[client_id]
logger.debug(f"Released all {count} objects for client {client_id}")
def _release_actors(self, client_id):
if client_id not in self.actor_owners:
logger.debug(f"Releasing client with no actors: {client_id}")
return
count = 0
actors_to_remove = self.actor_owners.pop(client_id)
for id_bytes in actors_to_remove:
count += 1
if self._can_remove_actor_ref(id_bytes):
logger.debug(f"Deleting reference to actor {id_bytes.hex()}")
del self.actor_refs[id_bytes]
logger.debug(f"Released all {count} actors for client: {client_id}")
@_use_response_cache
def Terminate(self, req, context=None):
if req.WhichOneof("terminate_type") == "task_object":
try:
object_ref = self.object_refs[req.client_id][req.task_object.id]
with disable_client_hook():
ray.cancel(
object_ref,
force=req.task_object.force,
recursive=req.task_object.recursive,
)
except Exception as e:
return_exception_in_context(e, context)
elif req.WhichOneof("terminate_type") == "actor":
try:
actor_ref = self.actor_refs[req.actor.id]
with disable_client_hook():
ray.kill(actor_ref, no_restart=req.actor.no_restart)
except Exception as e:
return_exception_in_context(e, context)
else:
raise RuntimeError(
"Client requested termination without providing a valid terminate_type"
)
return ray_client_pb2.TerminateResponse(ok=True)
def _async_get_object(
self,
request: ray_client_pb2.GetRequest,
client_id: str,
req_id: int,
result_queue: queue.Queue,
context=None,
) -> Optional[ray_client_pb2.GetResponse]:
"""Attempts to schedule a callback to push the GetResponse to the
main loop when the desired object is ready. If there is some failure
in scheduling, a GetResponse will be immediately returned.
"""
if len(request.ids) != 1:
raise ValueError(
f"Async get() must have exactly 1 Object ID. Actual: {request}"
)
rid = request.ids[0]
ref = self.object_refs[client_id].get(rid, None)
if not ref:
return ray_client_pb2.GetResponse(
valid=False,
error=cloudpickle.dumps(
ValueError(
f"ClientObjectRef with id {rid} not found for "
f"client {client_id}"
)
),
)
try:
logger.debug("async get: %s" % ref)
with disable_client_hook():
def send_get_response(result: Any) -> None:
"""Pushes GetResponses to the main DataPath loop to send
to the client. This is called when the object is ready
on the server side."""
try:
serialized = dumps_from_server(result, client_id, self)
total_size = len(serialized)
assert total_size > 0, "Serialized object cannot be zero bytes"
total_chunks = math.ceil(
total_size / OBJECT_TRANSFER_CHUNK_SIZE
)
for chunk_id in range(request.start_chunk_id, total_chunks):
start = chunk_id * OBJECT_TRANSFER_CHUNK_SIZE
end = min(
total_size, (chunk_id + 1) * OBJECT_TRANSFER_CHUNK_SIZE
)
get_resp = ray_client_pb2.GetResponse(
valid=True,
data=serialized[start:end],
chunk_id=chunk_id,
total_chunks=total_chunks,
total_size=total_size,
)
chunk_resp = ray_client_pb2.DataResponse(
get=get_resp, req_id=req_id
)
result_queue.put(chunk_resp)
except Exception as exc:
get_resp = ray_client_pb2.GetResponse(
valid=False, error=cloudpickle.dumps(exc)
)
resp = ray_client_pb2.DataResponse(get=get_resp, req_id=req_id)
result_queue.put(resp)
ref._on_completed(send_get_response)
return None
except Exception as e:
return ray_client_pb2.GetResponse(valid=False, error=cloudpickle.dumps(e))
def GetObject(self, request: ray_client_pb2.GetRequest, context):
metadata = dict(context.invocation_metadata())
client_id = metadata.get("client_id")
if client_id is None:
yield ray_client_pb2.GetResponse(
valid=False,
error=cloudpickle.dumps(
ValueError("client_id is not specified in request metadata")
),
)
else:
yield from self._get_object(request, client_id)
def _get_object(self, request: ray_client_pb2.GetRequest, client_id: str):
objectrefs = []
for rid in request.ids:
ref = self.object_refs[client_id].get(rid, None)
if ref:
objectrefs.append(ref)
else:
yield ray_client_pb2.GetResponse(
valid=False,
error=cloudpickle.dumps(
ValueError(
f"ClientObjectRef {rid} is not found for client {client_id}"
)
),
)
return
try:
logger.debug("get: %s" % objectrefs)
with disable_client_hook():
items = ray.get(objectrefs, timeout=request.timeout)
except Exception as e:
yield ray_client_pb2.GetResponse(valid=False, error=cloudpickle.dumps(e))
return
serialized = dumps_from_server(items, client_id, self)
total_size = len(serialized)
assert total_size > 0, "Serialized object cannot be zero bytes"
total_chunks = math.ceil(total_size / OBJECT_TRANSFER_CHUNK_SIZE)
for chunk_id in range(request.start_chunk_id, total_chunks):
start = chunk_id * OBJECT_TRANSFER_CHUNK_SIZE
end = min(total_size, (chunk_id + 1) * OBJECT_TRANSFER_CHUNK_SIZE)
yield ray_client_pb2.GetResponse(
valid=True,
data=serialized[start:end],
chunk_id=chunk_id,
total_chunks=total_chunks,
total_size=total_size,
)
def PutObject(
self, request: ray_client_pb2.PutRequest, context=None
) -> ray_client_pb2.PutResponse:
"""gRPC entrypoint for unary PutObject"""
return self._put_object(request.data, request.client_ref_id, "", context)
def _put_object(
self,
data: Union[bytes, bytearray],
client_ref_id: bytes,
client_id: str,
context: Optional[grpc.ServicerContext] = None,
) -> ray_client_pb2.PutResponse:
"""Put an object in the cluster with ray.put() via gRPC.
Args:
data: Pickled data. Can either be bytearray if this is called
from the dataservicer, or bytes if called from PutObject.
client_ref_id: The id associated with this object on the client.
client_id: The client who owns this data, for tracking when to
delete this reference.
context: gRPC context.
Returns:
A ``PutResponse`` containing the resulting object ref id, or an
error payload if the put failed.
"""
try:
obj = loads_from_client(data, self)
with disable_client_hook():
objectref = ray.put(obj)
except Exception as e:
logger.exception("Put failed:")
return ray_client_pb2.PutResponse(
id=b"", valid=False, error=cloudpickle.dumps(e)
)
self.object_refs[client_id][objectref.binary()] = objectref
if len(client_ref_id) > 0:
self.client_side_ref_map[client_id][client_ref_id] = objectref.binary()
logger.debug("put: %s" % objectref)
return ray_client_pb2.PutResponse(id=objectref.binary(), valid=True)
def WaitObject(self, request, context=None) -> ray_client_pb2.WaitResponse:
object_refs = []
for rid in request.object_ids:
if rid not in self.object_refs[request.client_id]:
raise Exception(
"Asking for a ref not associated with this client: %s" % str(rid)
)
object_refs.append(self.object_refs[request.client_id][rid])
num_returns = request.num_returns
timeout = request.timeout
try:
with disable_client_hook():
ready_object_refs, remaining_object_refs = ray.wait(
object_refs,
num_returns=num_returns,
timeout=timeout if timeout != -1 else None,
)
except Exception as e:
# TODO(ameer): improve exception messages.
logger.error(f"Exception {e}")
return ray_client_pb2.WaitResponse(valid=False)
logger.debug(
"wait: %s %s" % (str(ready_object_refs), str(remaining_object_refs))
)
ready_object_ids = [
ready_object_ref.binary() for ready_object_ref in ready_object_refs
]
remaining_object_ids = [
remaining_object_ref.binary()
for remaining_object_ref in remaining_object_refs
]
return ray_client_pb2.WaitResponse(
valid=True,
ready_object_ids=ready_object_ids,
remaining_object_ids=remaining_object_ids,
)
def Schedule(
self,
task: ray_client_pb2.ClientTask,
arglist: List[Any],
kwargs: Dict[str, Any],
context=None,
) -> ray_client_pb2.ClientTaskTicket:
logger.debug(
"schedule: %s %s"
% (task.name, ray_client_pb2.ClientTask.RemoteExecType.Name(task.type))
)
try:
with disable_client_hook():
if task.type == ray_client_pb2.ClientTask.FUNCTION:
result = self._schedule_function(task, arglist, kwargs, context)
elif task.type == ray_client_pb2.ClientTask.ACTOR:
result = self._schedule_actor(task, arglist, kwargs, context)
elif task.type == ray_client_pb2.ClientTask.METHOD:
result = self._schedule_method(task, arglist, kwargs, context)
elif task.type == ray_client_pb2.ClientTask.NAMED_ACTOR:
result = self._schedule_named_actor(task, context)
else:
raise NotImplementedError(
"Unimplemented Schedule task type: %s"
% ray_client_pb2.ClientTask.RemoteExecType.Name(task.type)
)
result.valid = True
return result
except Exception as e:
logger.debug("Caught schedule exception", exc_info=True)
return ray_client_pb2.ClientTaskTicket(
valid=False, error=cloudpickle.dumps(e)
)
def _schedule_method(
self,
task: ray_client_pb2.ClientTask,
arglist: List[Any],
kwargs: Dict[str, Any],
context=None,
) -> ray_client_pb2.ClientTaskTicket:
actor_handle = self.actor_refs.get(task.payload_id)
if actor_handle is None:
raise Exception("Can't run an actor the server doesn't have a handle for")
method = getattr(actor_handle, task.name)
opts = decode_options(task.options)
if opts is not None:
method = method.options(**opts)
output = method.remote(*arglist, **kwargs)
ids = self.unify_and_track_outputs(output, task.client_id)
return ray_client_pb2.ClientTaskTicket(return_ids=ids)
def _schedule_actor(
self,
task: ray_client_pb2.ClientTask,
arglist: List[Any],
kwargs: Dict[str, Any],
context=None,
) -> ray_client_pb2.ClientTaskTicket:
remote_class = self.lookup_or_register_actor(
task.payload_id, task.client_id, decode_options(task.baseline_options)
)
opts = decode_options(task.options)
if opts is not None:
remote_class = remote_class.options(**opts)
with current_server(self):
actor = remote_class.remote(*arglist, **kwargs)
self.actor_refs[actor._actor_id.binary()] = actor
self.actor_owners[task.client_id].add(actor._actor_id.binary())
return ray_client_pb2.ClientTaskTicket(return_ids=[actor._actor_id.binary()])
def _schedule_function(
self,
task: ray_client_pb2.ClientTask,
arglist: List[Any],
kwargs: Dict[str, Any],
context=None,
) -> ray_client_pb2.ClientTaskTicket:
remote_func = self.lookup_or_register_func(
task.payload_id, task.client_id, decode_options(task.baseline_options)
)
opts = decode_options(task.options)
if opts is not None:
remote_func = remote_func.options(**opts)
with current_server(self):
output = remote_func.remote(*arglist, **kwargs)
ids = self.unify_and_track_outputs(output, task.client_id)
return ray_client_pb2.ClientTaskTicket(return_ids=ids)
def _schedule_named_actor(
self, task: ray_client_pb2.ClientTask, context=None
) -> ray_client_pb2.ClientTaskTicket:
assert len(task.payload_id) == 0
# Convert empty string back to None.
actor = ray.get_actor(task.name, task.namespace or None)
bin_actor_id = actor._actor_id.binary()
if bin_actor_id not in self.actor_refs:
self.actor_refs[bin_actor_id] = actor
self.actor_owners[task.client_id].add(bin_actor_id)
self.named_actors.add(bin_actor_id)
return ray_client_pb2.ClientTaskTicket(return_ids=[actor._actor_id.binary()])
def lookup_or_register_func(
self, id: bytes, client_id: str, options: Optional[Dict]
) -> ray.remote_function.RemoteFunction:
with disable_client_hook():
if id not in self.function_refs:
funcref = self.object_refs[client_id][id]
func = ray.get(funcref)
if not inspect.isfunction(func):
raise Exception(
"Attempting to register function that isn't a function."
)
if options is None or len(options) == 0:
self.function_refs[id] = ray.remote(func)
else:
self.function_refs[id] = ray.remote(**options)(func)
return self.function_refs[id]
def lookup_or_register_actor(
self, id: bytes, client_id: str, options: Optional[Dict]
):
with disable_client_hook():
if id not in self.registered_actor_classes:
actor_class_ref = self.object_refs[client_id][id]
actor_class = ray.get(actor_class_ref)
if not inspect.isclass(actor_class):
raise Exception("Attempting to schedule actor that isn't a class.")
if options is None or len(options) == 0:
reg_class = ray.remote(actor_class)
else:
reg_class = ray.remote(**options)(actor_class)
self.registered_actor_classes[id] = reg_class
return self.registered_actor_classes[id]
def unify_and_track_outputs(self, output, client_id):
if output is None:
outputs = []
elif isinstance(output, list):
outputs = output
else:
outputs = [output]
for out in outputs:
if out.binary() in self.object_refs[client_id]:
logger.warning(f"Already saw object_ref {out}")
self.object_refs[client_id][out.binary()] = out
return [out.binary() for out in outputs]
def return_exception_in_context(err, context):
if context is not None:
context.set_details(encode_exception(err))
# Note: https://grpc.github.io/grpc/core/md_doc_statuscodes.html
# ABORTED used here since it should never be generated by the
# grpc lib -- this way we know the error was generated by ray logic
context.set_code(grpc.StatusCode.ABORTED)
def encode_exception(exception) -> str:
data = cloudpickle.dumps(exception)
return base64.standard_b64encode(data).decode()
def decode_options(options: ray_client_pb2.TaskOptions) -> Optional[Dict[str, Any]]:
if not options.pickled_options:
return None
opts = pickle.loads(options.pickled_options)
assert isinstance(opts, dict)
return opts
def serve(host: str, port: int, ray_connect_handler=None):
def default_connect_handler(
job_config: JobConfig = None, **ray_init_kwargs: Dict[str, Any]
):
with disable_client_hook():
if not ray.is_initialized():
return ray.init(job_config=job_config, **ray_init_kwargs)
from ray._private.grpc_utils import create_grpc_server_with_interceptors
ray_connect_handler = ray_connect_handler or default_connect_handler
server = create_grpc_server_with_interceptors(
max_workers=CLIENT_SERVER_MAX_THREADS,
thread_name_prefix="ray_client_server",
options=GRPC_OPTIONS,
asynchronous=False,
)
task_servicer = RayletServicer(ray_connect_handler)
data_servicer = DataServicer(task_servicer)
logs_servicer = LogstreamServicer()
ray_client_pb2_grpc.add_RayletDriverServicer_to_server(task_servicer, server)
ray_client_pb2_grpc.add_RayletDataStreamerServicer_to_server(data_servicer, server)
ray_client_pb2_grpc.add_RayletLogStreamerServicer_to_server(logs_servicer, server)
if not is_localhost(host):
add_port_to_grpc_server(server, build_address(get_localhost_ip(), port))
add_port_to_grpc_server(server, build_address(host, port))
current_handle = ClientServerHandle(
task_servicer=task_servicer,
data_servicer=data_servicer,
logs_servicer=logs_servicer,
grpc_server=server,
)
server.start()
return current_handle
def init_and_serve(host: str, port: int, *args, **kwargs):
with disable_client_hook():
# Disable client mode inside the worker's environment
info = ray.init(*args, **kwargs)
def ray_connect_handler(job_config=None, **ray_init_kwargs):
# Ray client will disconnect from ray when
# num_clients == 0.
if ray.is_initialized():
return info
else:
return ray.init(job_config=job_config, *args, **kwargs)
server_handle = serve(host, port, ray_connect_handler=ray_connect_handler)
return (server_handle, info)
def shutdown_with_server(server, _exiting_interpreter=False):
server.stop(1)
with disable_client_hook():
ray.shutdown(_exiting_interpreter=_exiting_interpreter)
def create_ray_handler(address, redis_password, redis_username=None):
def ray_connect_handler(job_config: JobConfig = None, **ray_init_kwargs):
if address:
if redis_password:
ray.init(
address=address,
_redis_username=redis_username,
_redis_password=redis_password,
job_config=job_config,
**ray_init_kwargs,
)
else:
ray.init(address=address, job_config=job_config, **ray_init_kwargs)
else:
ray.init(job_config=job_config, **ray_init_kwargs)
return ray_connect_handler
def try_create_gcs_client(address: Optional[str]) -> Optional[GcsClient]:
"""
Try to create a gcs client based on the command line args or by
autodetecting a running Ray cluster.
"""
address = canonicalize_bootstrap_address_or_die(address)
return GcsClient(address=address)
def main():
import argparse
parser = argparse.ArgumentParser()
parser.add_argument(
"--host",
type=str,
default=get_all_interfaces_ip(),
help="Host IP to bind to. Defaults to all interfaces (0.0.0.0/::).",
)
parser.add_argument("-p", "--port", type=int, default=10001, help="Port to bind to")
parser.add_argument(
"--mode",
type=str,
choices=["proxy", "legacy", "specific-server"],
default="proxy",
)
parser.add_argument(
"--address", required=False, type=str, help="Address to use to connect to Ray"
)
parser.add_argument(
"--redis-username",
required=False,
type=str,
help="username for connecting to Redis",
)
parser.add_argument(
"--runtime-env-agent-address",
required=False,
type=str,
default=None,
help="The port to use for connecting to the runtime_env_agent.",
)
parser.add_argument(
"--node-id",
required=False,
type=str,
default=None,
help="The hex ID of this node.",
)
args, _ = parser.parse_known_args()
redis_password = os.environ.get(ray_constants.RAY_REDIS_PASSWORD_ENV)
setup_logger(ray_constants.LOGGER_LEVEL, ray_constants.LOGGER_FORMAT)
ray_connect_handler = create_ray_handler(
args.address, redis_password, args.redis_username
)
hostport = build_address(args.host, args.port)
args_str = str(args)
logger.info(f"Starting Ray Client server on {hostport}, args {args_str}")
if args.mode == "proxy":
server = serve_proxier(
args.host,
args.port,
args.address,
redis_username=args.redis_username,
redis_password=redis_password,
runtime_env_agent_address=args.runtime_env_agent_address,
node_id=args.node_id,
)
else:
server = serve(args.host, args.port, ray_connect_handler)
try:
idle_checks_remaining = TIMEOUT_FOR_SPECIFIC_SERVER_S
while True:
health_report = {
"time": time.time(),
}
try:
if not ray.experimental.internal_kv._internal_kv_initialized():
gcs_client = try_create_gcs_client(args.address)
ray.experimental.internal_kv._initialize_internal_kv(gcs_client)
ray.experimental.internal_kv._internal_kv_put(
"ray_client_server",
json.dumps(health_report),
namespace=ray_constants.KV_NAMESPACE_HEALTHCHECK,
)
except Exception as e:
logger.error(
f"[{args.mode}] Failed to put health check on {args.address}"
)
logger.exception(e)
time.sleep(1)
if args.mode == "specific-server":
if server.data_servicer.num_clients > 0:
idle_checks_remaining = TIMEOUT_FOR_SPECIFIC_SERVER_S
else:
idle_checks_remaining -= 1
if idle_checks_remaining == 0:
raise KeyboardInterrupt()
if (
idle_checks_remaining % 5 == 0
and idle_checks_remaining != TIMEOUT_FOR_SPECIFIC_SERVER_S
):
logger.info(f"{idle_checks_remaining} idle checks before shutdown.")
except KeyboardInterrupt:
server.stop(0)
if __name__ == "__main__":
main()
python/ray/util/client/server/server_pickler.py
+124
View File
@@ -0,0 +1,124 @@
"""Implements the client side of the client/server pickling protocol.
These picklers are aware of the server internals and can find the
references held for the client within the server.
More discussion about the client/server pickling protocol can be found in:
ray/util/client/client_pickler.py
ServerPickler dumps ray objects from the server into the appropriate stubs.
ClientUnpickler loads stubs from the client and finds their associated handle
in the server instance.
"""
import io
from typing import TYPE_CHECKING, Any
import ray
import ray.cloudpickle as cloudpickle
from ray._private.client_mode_hook import disable_client_hook
from ray.util.client.client_pickler import PickleStub
from ray.util.client.server.server_stubs import (
ClientReferenceActor,
ClientReferenceFunction,
)
if TYPE_CHECKING:
from ray.util.client.server.server import RayletServicer
import pickle # noqa: F401
class ServerPickler(cloudpickle.CloudPickler):
def __init__(self, client_id: str, server: "RayletServicer", *args, **kwargs):
super().__init__(*args, **kwargs)
self.client_id = client_id
self.server = server
def persistent_id(self, obj):
if isinstance(obj, ray.ObjectRef):
obj_id = obj.binary()
if obj_id not in self.server.object_refs[self.client_id]:
# We're passing back a reference, probably inside a reference.
# Let's hold onto it.
self.server.object_refs[self.client_id][obj_id] = obj
return PickleStub(
type="Object",
client_id=self.client_id,
ref_id=obj_id,
name=None,
baseline_options=None,
)
elif isinstance(obj, ray.actor.ActorHandle):
actor_id = obj._actor_id.binary()
if actor_id not in self.server.actor_refs:
# We're passing back a handle, probably inside a reference.
self.server.actor_refs[actor_id] = obj
if actor_id not in self.server.actor_owners[self.client_id]:
self.server.actor_owners[self.client_id].add(actor_id)
return PickleStub(
type="Actor",
client_id=self.client_id,
ref_id=obj._actor_id.binary(),
name=None,
baseline_options=None,
)
return None
class ClientUnpickler(pickle.Unpickler):
def __init__(self, server, *args, **kwargs):
super().__init__(*args, **kwargs)
self.server = server
def persistent_load(self, pid):
assert isinstance(pid, PickleStub)
if pid.type == "Ray":
return ray
elif pid.type == "Object":
return self.server.object_refs[pid.client_id][pid.ref_id]
elif pid.type == "Actor":
return self.server.actor_refs[pid.ref_id]
elif pid.type == "RemoteFuncSelfReference":
return ClientReferenceFunction(pid.client_id, pid.ref_id)
elif pid.type == "RemoteFunc":
return self.server.lookup_or_register_func(
pid.ref_id, pid.client_id, pid.baseline_options
)
elif pid.type == "RemoteActorSelfReference":
return ClientReferenceActor(pid.client_id, pid.ref_id)
elif pid.type == "RemoteActor":
return self.server.lookup_or_register_actor(
pid.ref_id, pid.client_id, pid.baseline_options
)
elif pid.type == "RemoteMethod":
actor = self.server.actor_refs[pid.ref_id]
return getattr(actor, pid.name)
else:
raise NotImplementedError("Uncovered client data type")
def dumps_from_server(
obj: Any, client_id: str, server_instance: "RayletServicer", protocol=None
) -> bytes:
with io.BytesIO() as file:
sp = ServerPickler(client_id, server_instance, file, protocol=protocol)
sp.dump(obj)
return file.getvalue()
def loads_from_client(
data: bytes,
server_instance: "RayletServicer",
*,
fix_imports=True,
encoding="ASCII",
errors="strict"
) -> Any:
with disable_client_hook():
if isinstance(data, str):
raise TypeError("Can't load pickle from unicode string")
file = io.BytesIO(data)
return ClientUnpickler(
server_instance, file, fix_imports=fix_imports, encoding=encoding
).load()
python/ray/util/client/server/server_stubs.py
+66
View File
@@ -0,0 +1,66 @@
from abc import ABC, abstractmethod
from contextlib import contextmanager
_current_server = None
@contextmanager
def current_server(r):
global _current_server
remote = _current_server
_current_server = r
try:
yield
finally:
_current_server = remote
class ClientReferenceSentinel(ABC):
def __init__(self, client_id, id):
self.client_id = client_id
self.id = id
def __reduce__(self):
remote_obj = self.get_remote_obj()
if remote_obj is None:
return (self.__class__, (self.client_id, self.id))
return (identity, (remote_obj,))
@abstractmethod
def get_remote_obj(self):
pass
def get_real_ref_from_server(self):
global _current_server
if _current_server is None:
return None
client_map = _current_server.client_side_ref_map.get(self.client_id, None)
if client_map is None:
return None
return client_map.get(self.id, None)
class ClientReferenceActor(ClientReferenceSentinel):
def get_remote_obj(self):
global _current_server
real_ref_id = self.get_real_ref_from_server()
if real_ref_id is None:
return None
return _current_server.lookup_or_register_actor(
real_ref_id, self.client_id, None
)
class ClientReferenceFunction(ClientReferenceSentinel):
def get_remote_obj(self):
global _current_server
real_ref_id = self.get_real_ref_from_server()
if real_ref_id is None:
return None
return _current_server.lookup_or_register_func(
real_ref_id, self.client_id, None
)
def identity(x):
return x
python/ray/util/client/worker.py
+1005
View File
File diff suppressed because it is too large Load Diff