diff --git a/doc/source/ray-core/handling-dependencies.rst b/doc/source/ray-core/handling-dependencies.rst index 3a4f2915d1bf..6f0f9415755b 100644 --- a/doc/source/ray-core/handling-dependencies.rst +++ b/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 my_script.py + + This command sets the `py_executable` to `uv run ` 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 """"""""""""""""""" diff --git a/python/ray/_private/runtime_env/uv_runtime_env_hook.py b/python/ray/_private/runtime_env/uv_runtime_env_hook.py new file mode 100644 index 000000000000..2571d1d9dc6d --- /dev/null +++ b/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))) diff --git a/python/ray/tests/test_runtime_env_uv_run.py b/python/ray/tests/test_runtime_env_uv_run.py index 67f9ef65d611..8bdde29eda3e 100644 --- a/python/ray/tests/test_runtime_env_uv_run.py +++ b/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__]))