chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -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.
|
||||
Reference in New Issue
Block a user