Skip to content

feat: Add UV support for Ray Client mode#60868

Draft
preneond wants to merge 1 commit intoray-project:masterfrom
preneond:feature/uv-ray-client-support
Draft

feat: Add UV support for Ray Client mode#60868
preneond wants to merge 1 commit intoray-project:masterfrom
preneond:feature/uv-ray-client-support

Conversation

@preneond
Copy link
Contributor

@preneond preneond commented Feb 9, 2026

Why are these changes needed?

UV package manager integration currently fails when using Ray Client connections via ray.init("ray://..."). While UV works correctly with the Ray Jobs API (ray job submit -- uv run script.py), it is completely non-functional for interactive Ray Client workflows.

The root cause is that RAY_RUNTIME_ENV_HOOK is explicitly skipped for Ray Client mode (introduced in #26688), which also prevents the built-in UV hook from running. The original skip was necessary because server-side hooks could overwrite GCS URIs with local paths after working_dir upload. However, the UV hook is safe to run client-side before upload.

Fixes #57991

Changes

Core Change: python/ray/util/client/__init__.py

  • Added _apply_uv_hook_for_client() function that detects UV environment on the client side before connecting to the cluster
  • Modified ClientContext.connect() to apply the UV hook before the server connection
  • Custom RAY_RUNTIME_ENV_HOOK remains disabled for Ray Client (safety preserved)

How It Works

  1. Client-side: Detects uv run in the local process tree using the existing _get_uv_run_cmdline() function
  2. Client-side: Extracts UV args and sets py_executable = "uv run ..." in runtime_env
  3. Upload: working_dir (with pyproject.toml/uv.lock) uploads to GCS as normal
  4. Cluster: Workers receive py_executable and execute uv run default_worker.py, which installs dependencies from the uploaded lock files

This avoids the issue from #26688 because the hook runs before the working_dir upload (not after), so there's no conflict with GCS URIs.

Tests: python/ray/tests/test_client.py

  • test_ray_client_uv_hook_detection — verifies UV detection sets py_executable correctly, handles missing UV, and recovers from errors
  • test_ray_client_uv_no_detection_without_uv — verifies normal Ray Client operation is unaffected
  • test_ray_client_uv_hook_with_existing_runtime_env — verifies user's working_dir and env_vars are preserved

Tests: python/ray/tests/test_runtime_env_uv_run.py

  • test_uv_run_ray_client_mode — end-to-end test verifying UV installs packages via Ray Client

Related issue number

Fixes #57991

Checks

  • I've signed off every commit (by using the -s flag, i.e., git commit -s) in this PR.
  • [ x] I've run scripts/format.sh to lint the changes in this PR.
  • I've included any doc changes needed for https://docs.ray.io/en/master/.
  • I've added any new APIs to the API Reference.
  • I've made sure the tests are passing (full test_client.py suite: 46/46 passed).
  • Testing Strategy
    • Unit tests for _apply_uv_hook_for_client() function
    • Integration tests with Ray Client server
    • Backward compatibility with existing test_env_hook_skipped_for_ray_client

@preneond preneond requested a review from a team as a code owner February 9, 2026 15:41
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request adds support for UV package manager in Ray Client mode, which was previously non-functional. The core change involves adding a client-side hook to detect and apply the UV environment before connecting to the cluster. The implementation is well-tested with new unit and integration tests. I've found a couple of issues: an incorrect type hint and a potential bug in how runtime_env is resolved from different sources. My review includes suggestions to fix these.

Comment on lines 121 to 123
if job_config and job_config.runtime_env:
# Prefer job_config runtime_env if provided
runtime_env = runtime_env or job_config.runtime_env
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

The logic for determining the runtime_env has a potential bug and a misleading comment.

  1. The comment Prefer job_config runtime_env if provided is incorrect. The code actually prefers runtime_env from ray_init_kwargs if it's truthy.
  2. The expression runtime_env or job_config.runtime_env is problematic when ray_init_kwargs contains an empty dictionary for runtime_env (e.g., ray.init(runtime_env={})). An empty dictionary is a valid value that should take precedence over job_config.runtime_env, but because it's falsy, job_config.runtime_env will be used instead.

To fix this and align with standard Ray precedence rules (ray.init arguments > JobConfig), please make the logic more explicit.

Suggested change
if job_config and job_config.runtime_env:
# Prefer job_config runtime_env if provided
runtime_env = runtime_env or job_config.runtime_env
if runtime_env is None and job_config and job_config.runtime_env:
# If runtime_env is not in ray.init(), use the one from JobConfig.
runtime_env = job_config.runtime_env

logger = logging.getLogger(__name__)


def _apply_uv_hook_for_client(runtime_env: Optional[Dict[str, Any]]) -> Dict[str, Any]:
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The return type hint for this function is Dict[str, Any], but it can return None if the input runtime_env is None and the UV hook is not applied. This is confirmed by test_ray_client_uv_hook_detection (Test 3). Please update the type hint to Optional[Dict[str, Any]] for correctness and to aid static analysis.

Suggested change
def _apply_uv_hook_for_client(runtime_env: Optional[Dict[str, Any]]) -> Dict[str, Any]:
def _apply_uv_hook_for_client(runtime_env: Optional[Dict[str, Any]]) -> Optional[Dict[str, Any]]:

@preneond preneond force-pushed the feature/uv-ray-client-support branch from 0fd3c0a to fb20e97 Compare February 9, 2026 15:43
@preneond preneond marked this pull request as draft February 9, 2026 15:45
This commit enables UV package manager integration for Ray Client
connections, fixing issue ray-project#57991.

Previously, the RAY_RUNTIME_ENV_HOOK was disabled for Ray Client mode,
preventing UV from detecting and installing dependencies from
pyproject.toml and uv.lock files. This caused ModuleNotFoundError
when using Ray Client with UV environments.

Changes:
- Add _apply_uv_hook_for_client() to detect UV on client side
- Modify ClientContext.connect() to apply UV hook before connection
- UV hook now runs client-side where 'uv run' process exists
- Server-side hooks remain disabled for safety

The fix detects 'uv run' on the client machine and propagates the
UV configuration (py_executable, working_dir) to cluster workers,
allowing them to install dependencies via UV.

Tests added:
- test_ray_client_uv_detection: Verify UV detection and propagation
- test_ray_client_uv_no_detection_without_uv: Verify normal operation
- test_ray_client_uv_with_job_config: Test with JobConfig
- test_uv_run_ray_client_mode: End-to-end UV installation test

Fixes ray-project#57991

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
Signed-off-by: Ondrej Prenek <ondra.prenek@gmail.com>
@preneond preneond force-pushed the feature/uv-ray-client-support branch from fb20e97 to 5b90e34 Compare February 9, 2026 15:46
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

uv integration doesn't work with Ray Client (gRPC)

1 participant