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.