feat: Add UV support for Ray Client mode#60868
feat: Add UV support for Ray Client mode#60868preneond wants to merge 1 commit intoray-project:masterfrom
Conversation
There was a problem hiding this comment.
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.
python/ray/util/client/__init__.py
Outdated
| if job_config and job_config.runtime_env: | ||
| # Prefer job_config runtime_env if provided | ||
| runtime_env = runtime_env or job_config.runtime_env |
There was a problem hiding this comment.
The logic for determining the runtime_env has a potential bug and a misleading comment.
- The comment
Prefer job_config runtime_env if providedis incorrect. The code actually prefersruntime_envfromray_init_kwargsif it's truthy. - The expression
runtime_env or job_config.runtime_envis problematic whenray_init_kwargscontains an empty dictionary forruntime_env(e.g.,ray.init(runtime_env={})). An empty dictionary is a valid value that should take precedence overjob_config.runtime_env, but because it's falsy,job_config.runtime_envwill be used instead.
To fix this and align with standard Ray precedence rules (ray.init arguments > JobConfig), please make the logic more explicit.
| 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 |
python/ray/util/client/__init__.py
Outdated
| logger = logging.getLogger(__name__) | ||
|
|
||
|
|
||
| def _apply_uv_hook_for_client(runtime_env: Optional[Dict[str, Any]]) -> Dict[str, Any]: |
There was a problem hiding this comment.
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.
| 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]]: |
0fd3c0a to
fb20e97
Compare
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>
fb20e97 to
5b90e34
Compare
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_HOOKis 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 afterworking_dirupload. However, the UV hook is safe to run client-side before upload.Fixes #57991
Changes
Core Change:
python/ray/util/client/__init__.py_apply_uv_hook_for_client()function that detects UV environment on the client side before connecting to the clusterClientContext.connect()to apply the UV hook before the server connectionRAY_RUNTIME_ENV_HOOKremains disabled for Ray Client (safety preserved)How It Works
uv runin the local process tree using the existing_get_uv_run_cmdline()functionpy_executable = "uv run ..."in runtime_envworking_dir(withpyproject.toml/uv.lock) uploads to GCS as normalpy_executableand executeuv run default_worker.py, which installs dependencies from the uploaded lock filesThis avoids the issue from #26688 because the hook runs before the
working_dirupload (not after), so there's no conflict with GCS URIs.Tests:
python/ray/tests/test_client.pytest_ray_client_uv_hook_detection— verifies UV detection setspy_executablecorrectly, handles missing UV, and recovers from errorstest_ray_client_uv_no_detection_without_uv— verifies normal Ray Client operation is unaffectedtest_ray_client_uv_hook_with_existing_runtime_env— verifies user'sworking_dirandenv_varsare preservedTests:
python/ray/tests/test_runtime_env_uv_run.pytest_uv_run_ray_client_mode— end-to-end test verifying UV installs packages via Ray ClientRelated issue number
Fixes #57991
Checks
git commit -s) in this PR.scripts/format.shto lint the changes in this PR.test_client.pysuite: 46/46 passed)._apply_uv_hook_for_client()functiontest_env_hook_skipped_for_ray_client