Runtime environment hooks to load runtime_env from uv run environment (…

…ray-project#50462)

<!-- Thank you for your contribution! Please review
https://github.com/ray-project/ray/blob/master/CONTRIBUTING.rst before
opening a pull request. -->

<!-- Please add a reviewer to the assignee section when you create a PR.
If you don't have the access to it, we will shortly find a reviewer and
assign them to your PR. -->

## Why are these changes needed?

Next step after ray-project#50160 to make it
more convenient to use UV with Ray.

This is a useful runtime environment hook for mirroring the environment
of `uv run` to the workers (currently the args to uv run and the
working_dir). This is useful because it will allow people to intuitively
use `uv run` in a distributed application with the same behavior as for
a single python process.

This only modifies the environment if the driver was run with `uv run`
and could conceivably become the default for drivers run with uv run.

This is currently a developer API as implied by the fact that it is in
the `_private` namespace. It is currently for experimentation and can
needs to be opted in via

```shell
export RAY_RUNTIME_ENV_HOOK=ray._private.runtime_env.uv_runtime_env_hook.hook
```

If it works well, we might make it the default in the `uv run` case.

## Related issue number

<!-- For example: "Closes ray-project#1234" -->

## Checks

- [ ] I've signed off every commit(by using the -s flag, i.e., `git
commit -s`) in this PR.
- [ ] 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. For example, if I
added a
method in Tune, I've added it in `doc/source/tune/api/` under the
           corresponding `.rst` file.
- [ ] I've made sure the tests are passing. Note that there might be a
few flaky tests, see the recent failures at https://flakey-tests.ray.io/
- Testing Strategy
   - [ ] Unit tests
   - [ ] Release tests
   - [ ] This PR is not tested :(

---------

Signed-off-by: Philipp Moritz <pcmoritz@gmail.com>
Co-authored-by: Edward Oakes <ed.nmi.oakes@gmail.com>
Co-authored-by: angelinalg <122562471+angelinalg@users.noreply.github.com>

doc/source/ray-core/handling-dependencies.rst

@@ -373,6 +373,29 @@ file and use `uv run --with-requirements requirements.txt` for your `py_executab
 or use the `--with` flag to specify individual requirements.
 
 
+.. tip::
+
+  In order to make this functionality available in a convenient way without having
+  to specify `py_executable` in the runtime environment, you can use the following
+  runtime environment hook:
+
+  .. code-block:: sh
+
+    export RAY_RUNTIME_ENV_HOOK=ray._private.runtime_env.uv_runtime_env_hook.hook
+
+  Run your driver with the following command:
+
+  .. code-block:: sh
+
+    uv run <args> my_script.py
+
+  This command sets the `py_executable` to `uv run <args>` and also sets the
+  `working_dir` to the same working directory that the driver is using, either
+  the current working directory or the `--directory` in `uv run`.
+  Note that this hook is experimental, in the future the Ray team might make
+  this behavior the default.
+
+
 Library Development
 """""""""""""""""""
 

python/ray/_private/runtime_env/uv_runtime_env_hook.py

@@ -0,0 +1,99 @@
+import argparse
+import os
+from pathlib import Path
+import sys
+from typing import Any, Dict, Optional
+
+import psutil
+
+
+def hook(runtime_env: Optional[Dict[str, Any]]) -> Dict[str, Any]:
+    """Hook that detects if the driver is run in 'uv run' and sets the runtime environment accordingly."""
+
+    runtime_env = runtime_env or {}
+
+    parent = psutil.Process().parent()
+    cmdline = parent.cmdline()
+    if os.path.basename(cmdline[0]) != "uv" or cmdline[1] != "run":
+        # This means the driver was not run with 'uv run' -- in this case
+        # we leave the runtime environment unchanged
+        return runtime_env
+
+    # Extract the arguments of 'uv run' that are not arguments of the script
+    uv_run_args = cmdline[: len(cmdline) - len(sys.argv)]
+
+    # Remove the "--directory" argument since it has already been taken into
+    # account when setting the current working directory of the current process
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--directory", nargs="?")
+    _, remaining_uv_run_args = parser.parse_known_args(uv_run_args)
+
+    runtime_env["py_executable"] = " ".join(remaining_uv_run_args)
+
+    # If the user specified a working_dir, we always honor it, otherwise
+    # use the same working_dir that uv run would use
+    if "working_dir" not in runtime_env:
+        runtime_env["working_dir"] = os.getcwd()
+
+    # In the last part of the function we do some error checking that should catch
+    # the most common cases of how things are different in Ray, i.e. not the whole
+    # file system will be available on the workers, only the working_dir.
+
+    # First parse the arguments we need to check
+    uv_run_parser = argparse.ArgumentParser()
+    uv_run_parser.add_argument("--with-requirements", nargs="?")
+    uv_run_parser.add_argument("--project", nargs="?")
+    uv_run_parser.add_argument("--no-project", action="store_true")
+    known_args, _ = uv_run_parser.parse_known_args(uv_run_args)
+
+    working_dir = Path(runtime_env["working_dir"]).resolve()
+
+    # Check if the requirements.txt file is in the working_dir
+    if known_args.with_requirements and not Path(
+        known_args.with_requirements
+    ).resolve().is_relative_to(working_dir):
+        raise RuntimeError(
+            f"You specified --with-requirements={known_args.with_requirements} but "
+            f"the requirements file is not in the working_dir {runtime_env['working_dir']}, "
+            "so the workers will not have access to the file. Make sure "
+            "the requirements file is in the working directory. "
+            "You can do so by specifying --directory in 'uv run', by changing the current "
+            "working directory before running 'uv run', or by using the 'working_dir' "
+            "parameter of the runtime_environment."
+        )
+
+    # Check if the pyproject.toml file is in the working_dir
+    pyproject = None
+    if known_args.no_project:
+        pyproject = None
+    elif known_args.project:
+        pyproject = Path(known_args.project)
+    else:
+        # Walk up the directory tree until pyproject.toml is found
+        current_path = Path.cwd().resolve()
+        while current_path != current_path.parent:
+            if (current_path / "pyproject.toml").exists():
+                pyproject = Path(current_path / "pyproject.toml")
+                break
+            current_path = current_path.parent
+
+    if pyproject and not pyproject.resolve().is_relative_to(working_dir):
+        raise RuntimeError(
+            f"Your {pyproject.resolve()} is not in the working_dir {runtime_env['working_dir']}, "
+            "so the workers will not have access to the file. Make sure "
+            "the pyproject.toml file is in the working directory. "
+            "You can do so by specifying --directory in 'uv run', by changing the current "
+            "working directory before running 'uv run', or by using the 'working_dir' "
+            "parameter of the runtime_environment."
+        )
+
+    return runtime_env
+
+
+if __name__ == "__main__":
+    import json
+
+    # This is used for unit testing if the runtime_env_hook picks up the
+    # right settings.
+    runtime_env = json.loads(sys.argv[1])
+    print(json.dumps(hook(runtime_env)))

python/ray/tests/test_runtime_env_uv_run.py

@@ -1,5 +1,6 @@
-# End-to-end tests for using "uv run" with py_executable
+# End-to-end tests for using "uv run"
 
+import json
 import os
 from pathlib import Path
 import pytest
@@ -143,6 +144,170 @@ def emojize():
     assert ray.get(emojize.remote()) == "The package was edited"
 
 
+@pytest.mark.skipif(sys.platform == "win32", reason="Not ported to Windows yet.")
+def test_uv_run_runtime_env_hook(with_uv):
+
+    import ray._private.runtime_env.uv_runtime_env_hook
+
+    uv = with_uv
+
+    def check_uv_run(
+        cmd, runtime_env, expected_output, subprocess_kwargs=None, expected_error=None
+    ):
+        result = subprocess.run(
+            cmd
+            + [ray._private.runtime_env.uv_runtime_env_hook.__file__]
+            + [json.dumps(runtime_env)],
+            capture_output=True,
+            **(subprocess_kwargs if subprocess_kwargs else {}),
+        )
+        output = result.stdout.strip().decode()
+        if result.returncode != 0:
+            assert expected_error
+            assert expected_error in result.stderr.decode()
+        else:
+            assert json.loads(output) == expected_output
+
+    check_uv_run(
+        cmd=[uv, "run", "--no-project"],
+        runtime_env={},
+        expected_output={
+            "py_executable": f"{uv} run --no-project",
+            "working_dir": os.getcwd(),
+        },
+    )
+    check_uv_run(
+        cmd=[uv, "run", "--no-project", "--directory", "/tmp"],
+        runtime_env={},
+        expected_output={
+            "py_executable": f"{uv} run --no-project",
+            "working_dir": os.path.realpath("/tmp"),
+        },
+    )
+    check_uv_run(
+        [uv, "run", "--no-project"],
+        {"working_dir": "/some/path"},
+        {"py_executable": f"{uv} run --no-project", "working_dir": "/some/path"},
+    )
+
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        tmp_dir = Path(tmp_dir).resolve()
+        with open(tmp_dir / "pyproject.toml", "w") as file:
+            file.write("[project]\n")
+            file.write('name = "test"\n')
+            file.write('version = "0.1"\n')
+            file.write('dependencies = ["psutil"]\n')
+        check_uv_run(
+            cmd=[uv, "run"],
+            runtime_env={},
+            expected_output={"py_executable": f"{uv} run", "working_dir": f"{tmp_dir}"},
+            subprocess_kwargs={"cwd": tmp_dir},
+        )
+
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        tmp_dir = Path(tmp_dir).resolve()
+        os.makedirs(tmp_dir / "cwd")
+        requirements = tmp_dir / "requirements.txt"
+        with open(requirements, "w") as file:
+            file.write("psutil\n")
+        check_uv_run(
+            cmd=[uv, "run", "--with-requirements", requirements],
+            runtime_env={},
+            expected_output={
+                "py_executable": f"{uv} run --with-requirements {requirements}",
+                "working_dir": f"{tmp_dir}",
+            },
+            subprocess_kwargs={"cwd": tmp_dir},
+        )
+
+    # Check things fail if there is a pyproject.toml upstream of the current working directory
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        tmp_dir = Path(tmp_dir).resolve()
+        os.makedirs(tmp_dir / "cwd")
+        with open(tmp_dir / "pyproject.toml", "w") as file:
+            file.write("[project]\n")
+            file.write('name = "test"\n')
+            file.write('version = "0.1"\n')
+            file.write('dependencies = ["psutil"]\n')
+        check_uv_run(
+            cmd=[uv, "run"],
+            runtime_env={},
+            expected_output=None,
+            subprocess_kwargs={"cwd": tmp_dir / "cwd"},
+            expected_error="Make sure the pyproject.toml file is in the working directory.",
+        )
+
+    # Check things fail if there is a requirements.txt upstream to the current working directory
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        tmp_dir = Path(tmp_dir).resolve()
+        os.makedirs(tmp_dir / "cwd")
+        with open(tmp_dir / "requirements.txt", "w") as file:
+            file.write("psutil\n")
+        check_uv_run(
+            cmd=[uv, "run", "--with-requirements", tmp_dir / "requirements.txt"],
+            runtime_env={},
+            expected_output=None,
+            subprocess_kwargs={"cwd": tmp_dir / "cwd"},
+            expected_error="Make sure the requirements file is in the working directory.",
+        )
+
+    # Check without uv run
+    subprocess.check_output(
+        [sys.executable, ray._private.runtime_env.uv_runtime_env_hook.__file__, "{}"]
+    ).strip().decode() == "{}"
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="Not ported to Windows yet.")
+def test_uv_run_runtime_env_hook_e2e(shutdown_only, with_uv, temp_dir):
+
+    uv = with_uv
+    tmp_out_dir = Path(temp_dir)
+
+    script = f"""
+import json
+import ray
+import os
+
+@ray.remote
+def f():
+    import emoji
+    return {{"working_dir_files": os.listdir(os.getcwd())}}
+
+with open("{tmp_out_dir / "output.txt"}", "w") as out:
+    json.dump(ray.get(f.remote()), out)
+"""
+
+    with tempfile.NamedTemporaryFile("w", suffix=".py", delete=False) as f:
+        f.write(script)
+        f.close()
+        subprocess.run(
+            [
+                uv,
+                "run",
+                # We want to run in the system environment so the current installation of Ray can be found here
+                "--python-preference=only-system",
+                "--with",
+                "emoji",
+                "--no-project",
+                f.name,
+            ],
+            stdin=subprocess.PIPE,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            env={
+                "RAY_RUNTIME_ENV_HOOK": "ray._private.runtime_env.uv_runtime_env_hook.hook",
+                "PYTHONPATH": ":".join(sys.path),
+                "PATH": os.environ["PATH"],
+            },
+            cwd=os.path.dirname(uv),
+            check=True,
+        )
+        with open(tmp_out_dir / "output.txt") as f:
+            assert json.load(f) == {
+                "working_dir_files": os.listdir(os.path.dirname(uv))
+            }
+
+
 if __name__ == "__main__":
     if os.environ.get("PARALLEL_CI"):
         sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))