Runtime environment hooks to load runtime_env from uv run environment #50462

pcmoritz · 2025-02-12T01:47:14Z

Why are these changes needed?

Next step after #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

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

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 :(

kouroshHakha · 2025-02-12T16:20:27Z

python/ray/_private/runtime_env/__init__.py

+    # Extract the arguments of 'uv run' that are not arguments of the script
+    uv_run_args = cmdline[: len(cmdline) - len(sys.argv)]
+    runtime_env["py_executable"] = " ".join(uv_run_args)


how are you differntiating between uv run args vs. script args? Is this something uv is handling under the hood by not modifying sys.argv?

the sys.argv are only the script args (they get passed to the script, which is a subprocess of uv run, by uv run), and cmdline are the full args with both the uv run args and the script args. So uv run args are computed in the line cmdline[: len(cmdline) - len(sys.argv)] above by removing the script args (in sys.argv) from all the args (cmdline)

kouroshHakha · 2025-02-12T16:21:42Z

python/ray/_private/runtime_env/__init__.py

+    return runtime_env
+
+
+def uv_run_runtime_env_hook(runtime_env: Dict[str, Any]) -> Dict[str, Any]:


now where do we inject this hook? in snapshot_util?

tangent question: How do we deal with uv run on jobs?

For the jobs: It is actually very simple, you just submit the job with uv run <script> as the entry point :)

On the hook: Right now you need to define it manually but I think going forward we should just remove the hook and do this by default once it is ironed out a little more and works well :)

so we still need to define the env variable RUNTIME_ENV_HOOK (I forget its name)?

Yes, that's right (for now) :)

edoakes · 2025-02-13T00:05:26Z

python/ray/_private/runtime_env/uv_runtime_env_hook.py

+    # Parse known arguments of uv run that impact the runtime environment
+    uv_run_parser = argparse.ArgumentParser()
+    uv_run_parser.add_argument("--directory", nargs="?")
+    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")


how'd we come up with this list? does it need to be updated over time?

These are the features I'm aware of so far that need special handling (with one exception), since they interact with the file system. The exception is https://docs.astral.sh/uv/configuration/files/#env but it seems more advanced and we can add it if / when the need arises.

Most features just work out of the box since they can just passed through.

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

Co-authored-by: Edward Oakes <[email protected]> Signed-off-by: Philipp Moritz <[email protected]>

angelinalg

Just some style nits. Our style guide avoids using "please" just for conciseness.

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

python/ray/_private/runtime_env/uv_runtime_env_hook.py

python/ray/tests/test_runtime_env_uv_run.py

Co-authored-by: angelinalg <[email protected]> Signed-off-by: Philipp Moritz <[email protected]>

Signed-off-by: Philipp Moritz <[email protected]>

Co-authored-by: angelinalg <[email protected]> Signed-off-by: Philipp Moritz <[email protected]>

…y into runtime-env-override-hook

Builtin hook to load runtime_env from environment variable

5c9038d

pcmoritz requested review from edoakes and kouroshHakha February 12, 2025 01:47

pcmoritz changed the title ~~Runtime environment hook to load runtime_env from environment variable~~ [WIP] Runtime environment hook to load runtime_env from environment variable Feb 12, 2025

pcmoritz added the go add ONLY when ready to merge, run all tests label Feb 12, 2025

pcmoritz added 3 commits February 11, 2025 18:05

lint

b39bd70

update

1155f66

Include UV hook

b269c71

pcmoritz changed the title ~~[WIP] Runtime environment hook to load runtime_env from environment variable~~ [WIP] Runtime environment hooks to load runtime_env from UV or an environment override Feb 12, 2025

lint

e802452

kouroshHakha reviewed Feb 12, 2025

View reviewed changes

add test

ded47a9

kevin85421 mentioned this pull request Feb 12, 2025

Use uv in sample YAML files ray-project/kuberay#3039

Open

2 tasks

pcmoritz added 6 commits February 12, 2025 13:01

lint

7581bb5

fix lint

814ac7b

add tests

206fef1

remove

70f89ab

update

8a432ff

clean

b7cc9bb

pcmoritz changed the title ~~[WIP] Runtime environment hooks to load runtime_env from UV or an environment override~~ Runtime environment hooks to load runtime_env from UV or an environment override Feb 12, 2025

pcmoritz changed the title ~~Runtime environment hooks to load runtime_env from UV or an environment override~~ Runtime environment hooks to load runtime_env from uv run environment Feb 12, 2025

add docs

9af39b8

pcmoritz requested a review from a team as a code owner February 12, 2025 22:16

pcmoritz added 2 commits February 12, 2025 14:17

lint

11daae2

update

0c13c7b

edoakes approved these changes Feb 13, 2025

View reviewed changes

pcmoritz and others added 3 commits February 12, 2025 17:34

fix test

fc329d2

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

03bd04c

Co-authored-by: Edward Oakes <[email protected]> Signed-off-by: Philipp Moritz <[email protected]>

update docs

0ad1c5d

angelinalg approved these changes Feb 13, 2025

View reviewed changes

pcmoritz and others added 11 commits February 12, 2025 18:29

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

4dfe515

Co-authored-by: angelinalg <[email protected]> Signed-off-by: Philipp Moritz <[email protected]>

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

81792fa

Co-authored-by: angelinalg <[email protected]> Signed-off-by: Philipp Moritz <[email protected]>

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

d5a5af7

Co-authored-by: angelinalg <[email protected]> Signed-off-by: Philipp Moritz <[email protected]>

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

3de2b9b

Co-authored-by: angelinalg <[email protected]> Signed-off-by: Philipp Moritz <[email protected]>

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

74d6cca

Co-authored-by: angelinalg <[email protected]> Signed-off-by: Philipp Moritz <[email protected]>

Update handling-dependencies.rst

19505d1

Signed-off-by: Philipp Moritz <[email protected]>

Update python/ray/_private/runtime_env/uv_runtime_env_hook.py

2a1c2bc

Co-authored-by: angelinalg <[email protected]> Signed-off-by: Philipp Moritz <[email protected]>

Update python/ray/_private/runtime_env/uv_runtime_env_hook.py

f086647

Co-authored-by: angelinalg <[email protected]> Signed-off-by: Philipp Moritz <[email protected]>

Update python/ray/_private/runtime_env/uv_runtime_env_hook.py

891e229

Co-authored-by: angelinalg <[email protected]> Signed-off-by: Philipp Moritz <[email protected]>

Update python/ray/tests/test_runtime_env_uv_run.py

cec298c

Co-authored-by: angelinalg <[email protected]> Signed-off-by: Philipp Moritz <[email protected]>

Update python/ray/tests/test_runtime_env_uv_run.py

f1eceb5

Co-authored-by: angelinalg <[email protected]> Signed-off-by: Philipp Moritz <[email protected]>

kouroshHakha approved these changes Feb 13, 2025

View reviewed changes

pcmoritz added 4 commits February 12, 2025 21:21

fixes

7fb5fd3

Merge branch 'runtime-env-override-hook' of github.com:ray-project/ra…

bb77558

…y into runtime-env-override-hook

update

db30a14

handle both '--directory value' and '--directory=value'

c4ee8ed

pcmoritz enabled auto-merge (squash) February 13, 2025 06:50

pcmoritz merged commit c694360 into master Feb 13, 2025
6 checks passed

pcmoritz deleted the runtime-env-override-hook branch February 13, 2025 08:22

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Runtime environment hooks to load runtime_env from uv run environment #50462

Runtime environment hooks to load runtime_env from uv run environment #50462

pcmoritz commented Feb 12, 2025 •

edited

Loading

kouroshHakha Feb 12, 2025

pcmoritz Feb 12, 2025

kouroshHakha Feb 12, 2025

pcmoritz Feb 12, 2025 •

edited

Loading

kouroshHakha Feb 13, 2025

pcmoritz Feb 13, 2025

edoakes Feb 13, 2025

pcmoritz Feb 13, 2025

angelinalg left a comment

		return runtime_env


		def uv_run_runtime_env_hook(runtime_env: Dict[str, Any]) -> Dict[str, Any]:

Runtime environment hooks to load runtime_env from uv run environment #50462

Runtime environment hooks to load runtime_env from uv run environment #50462

Conversation

pcmoritz commented Feb 12, 2025 • edited Loading

Why are these changes needed?

Related issue number

Checks

kouroshHakha Feb 12, 2025

Choose a reason for hiding this comment

pcmoritz Feb 12, 2025

Choose a reason for hiding this comment

kouroshHakha Feb 12, 2025

Choose a reason for hiding this comment

pcmoritz Feb 12, 2025 • edited Loading

Choose a reason for hiding this comment

kouroshHakha Feb 13, 2025

Choose a reason for hiding this comment

pcmoritz Feb 13, 2025

Choose a reason for hiding this comment

edoakes Feb 13, 2025

Choose a reason for hiding this comment

pcmoritz Feb 13, 2025

Choose a reason for hiding this comment

angelinalg left a comment

Choose a reason for hiding this comment

pcmoritz commented Feb 12, 2025 •

edited

Loading

pcmoritz Feb 12, 2025 •

edited

Loading