wehub-resource-sync
609 lines
19 KiB
Python
609 lines
19 KiB
Python
import json
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
from uv import find_uv_bin
|
|
|
|
import ray
|
|
from ray._private.test_utils import (
|
|
format_web_url,
|
|
wait_until_server_available,
|
|
)
|
|
from ray.exceptions import RaySystemError
|
|
|
|
PYPROJECT_TOML = """
|
|
[project]
|
|
name = "test"
|
|
version = "0.1"
|
|
dependencies = [
|
|
"emoji",
|
|
]
|
|
requires-python = ">=3.9"
|
|
"""
|
|
|
|
|
|
@pytest.fixture(scope="function")
|
|
def tmp_working_dir():
|
|
"""A test fixture, which writes a pyproject.toml."""
|
|
with tempfile.TemporaryDirectory() as tmp_dir:
|
|
path = Path(tmp_dir)
|
|
|
|
script_file = path / "pyproject.toml"
|
|
with script_file.open(mode="w") as f:
|
|
f.write(PYPROJECT_TOML)
|
|
|
|
yield tmp_dir
|
|
|
|
|
|
@pytest.mark.skipif(sys.platform == "win32", reason="Not ported to Windows yet.")
|
|
def test_uv_run_simple(shutdown_only):
|
|
runtime_env = {
|
|
"py_executable": f"{find_uv_bin()} run --with emoji --no-project",
|
|
}
|
|
ray.init(runtime_env=runtime_env)
|
|
|
|
@ray.remote
|
|
def emojize():
|
|
import emoji
|
|
|
|
return emoji.emojize("Ray rocks :thumbs_up:")
|
|
|
|
assert ray.get(emojize.remote()) == "Ray rocks 👍"
|
|
|
|
|
|
@pytest.mark.skipif(sys.platform == "win32", reason="Not ported to Windows yet.")
|
|
def test_uv_run_pyproject(shutdown_only, tmp_working_dir):
|
|
tmp_dir = tmp_working_dir
|
|
|
|
ray.init(
|
|
runtime_env={
|
|
"working_dir": tmp_dir,
|
|
# We want to run in the system environment so the current installation of Ray can be found here
|
|
"py_executable": f"env PYTHONPATH={':'.join(sys.path)} {find_uv_bin()} run --python-preference=only-system",
|
|
}
|
|
)
|
|
|
|
@ray.remote
|
|
def emojize():
|
|
import emoji
|
|
|
|
return emoji.emojize("Ray rocks :thumbs_up:")
|
|
|
|
assert ray.get(emojize.remote()) == "Ray rocks 👍"
|
|
|
|
|
|
@pytest.mark.skipif(sys.platform == "win32", reason="Not ported to Windows yet.")
|
|
def test_uv_run_editable(shutdown_only, tmp_working_dir):
|
|
tmp_dir = tmp_working_dir
|
|
|
|
subprocess.run(
|
|
["git", "clone", "https://github.com/carpedm20/emoji/", "emoji_copy"],
|
|
cwd=tmp_dir,
|
|
)
|
|
|
|
subprocess.run(
|
|
["git", "reset", "--hard", "08c5cc4789d924ad4215e2fb2ee8f0b19a0d421f"],
|
|
cwd=Path(tmp_dir) / "emoji_copy",
|
|
)
|
|
|
|
subprocess.run(
|
|
[find_uv_bin(), "add", "--editable", "./emoji_copy"],
|
|
cwd=tmp_dir,
|
|
)
|
|
|
|
# Now edit the package
|
|
content = ""
|
|
with open(Path(tmp_dir) / "emoji_copy" / "emoji" / "core.py") as f:
|
|
content = f.read()
|
|
|
|
content = content.replace(
|
|
"return pattern.sub(replace, string)", 'return "The package was edited"'
|
|
)
|
|
|
|
with open(Path(tmp_dir) / "emoji_copy" / "emoji" / "core.py", "w") as f:
|
|
f.write(content)
|
|
|
|
ray.init(
|
|
runtime_env={
|
|
"working_dir": tmp_dir,
|
|
# We want to run in the system environment so the current installation of Ray can be found here
|
|
"py_executable": f"env PYTHONPATH={':'.join(sys.path)} {find_uv_bin()} run --python-preference=only-system",
|
|
}
|
|
)
|
|
|
|
@ray.remote
|
|
def emojize():
|
|
import emoji
|
|
|
|
return emoji.emojize("Ray rocks :thumbs_up:")
|
|
|
|
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():
|
|
import ray._private.runtime_env.uv_runtime_env_hook
|
|
|
|
def check_uv_run(
|
|
cmd, runtime_env, expected_output, subprocess_kwargs=None, expected_error=None
|
|
):
|
|
subprocess_kwargs = subprocess_kwargs or {}
|
|
env = dict(subprocess_kwargs.get("env") or os.environ)
|
|
env["UV_PYTHON"] = sys.executable
|
|
result = subprocess.run(
|
|
cmd + [json.dumps(runtime_env)],
|
|
capture_output=True,
|
|
**{**subprocess_kwargs, "env": env},
|
|
)
|
|
output = result.stdout.strip().decode()
|
|
if result.returncode != 0:
|
|
assert expected_error, result.stderr.decode()
|
|
assert expected_error in result.stderr.decode()
|
|
else:
|
|
assert json.loads(output) == expected_output
|
|
|
|
script = ray._private.runtime_env.uv_runtime_env_hook.__file__
|
|
|
|
check_uv_run(
|
|
cmd=[find_uv_bin(), "run", "--no-project", script],
|
|
runtime_env={},
|
|
expected_output={
|
|
"py_executable": f"{find_uv_bin()} run --no-project --python {sys.executable} python",
|
|
"working_dir": os.getcwd(),
|
|
},
|
|
)
|
|
check_uv_run(
|
|
cmd=[find_uv_bin(), "run", "--no-project", "--directory", "/tmp", script],
|
|
runtime_env={},
|
|
expected_output={
|
|
"py_executable": f"{find_uv_bin()} run --no-project --python {sys.executable} python",
|
|
"working_dir": os.path.realpath("/tmp"),
|
|
},
|
|
)
|
|
check_uv_run(
|
|
[find_uv_bin(), "run", "--no-project", script],
|
|
{"working_dir": "/some/path"},
|
|
{
|
|
"py_executable": f"{find_uv_bin()} run --no-project --python {sys.executable} python",
|
|
"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=[find_uv_bin(), "run", script],
|
|
runtime_env={},
|
|
expected_output={
|
|
"py_executable": f"{find_uv_bin()} run --python {sys.executable} python",
|
|
"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=[find_uv_bin(), "run", "--with-requirements", requirements, script],
|
|
runtime_env={},
|
|
expected_output={
|
|
"py_executable": f"{find_uv_bin()} run --with-requirements {requirements} --python {sys.executable} python",
|
|
"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=[find_uv_bin(), "run", script],
|
|
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=[
|
|
find_uv_bin(),
|
|
"run",
|
|
"--with-requirements",
|
|
tmp_dir / "requirements.txt",
|
|
script,
|
|
],
|
|
runtime_env={},
|
|
expected_output=None,
|
|
subprocess_kwargs={"cwd": tmp_dir / "cwd"},
|
|
expected_error="Make sure the requirements file is in the working directory.",
|
|
)
|
|
|
|
# Make sure the runtime environment hook gives the appropriate error message
|
|
# when combined with the 'pip' or 'uv' environment.
|
|
for runtime_env in [{"uv": ["emoji"]}, {"pip": ["emoji"]}]:
|
|
check_uv_run(
|
|
cmd=[find_uv_bin(), "run", "--no-project", script],
|
|
runtime_env=runtime_env,
|
|
expected_output=None,
|
|
expected_error="You are using the 'pip' or 'uv' runtime environments together with 'uv run'.",
|
|
)
|
|
|
|
# Check without uv run
|
|
subprocess.check_output([sys.executable, script, "{}"]).strip().decode() == "{}"
|
|
|
|
# Check in the case that there is one more level of subprocess indirection between
|
|
# the "uv run" process and the process that checks the environment
|
|
check_uv_run(
|
|
cmd=[find_uv_bin(), "run", "--no-project", script],
|
|
runtime_env={},
|
|
expected_output={
|
|
"py_executable": f"{find_uv_bin()} run --no-project --python {sys.executable} python",
|
|
"working_dir": os.getcwd(),
|
|
},
|
|
subprocess_kwargs={
|
|
"env": {**os.environ, "RAY_TEST_UV_ADD_SUBPROCESS_INDIRECTION": "1"}
|
|
},
|
|
)
|
|
|
|
# Check in the case that the script is started with multiprocessing spawn
|
|
check_uv_run(
|
|
cmd=[find_uv_bin(), "run", "--no-project", script],
|
|
runtime_env={},
|
|
expected_output={
|
|
"py_executable": f"{find_uv_bin()} run --no-project --python {sys.executable} python",
|
|
"working_dir": os.getcwd(),
|
|
},
|
|
subprocess_kwargs={
|
|
"env": {**os.environ, "RAY_TEST_UV_MULTIPROCESSING_SPAWN": "1"}
|
|
},
|
|
)
|
|
|
|
# Check in the case that a module is used for "uv run" (-m or --module)
|
|
check_uv_run(
|
|
cmd=[
|
|
find_uv_bin(),
|
|
"run",
|
|
"--no-project",
|
|
"-m",
|
|
"ray._private.runtime_env.uv_runtime_env_hook",
|
|
],
|
|
runtime_env={},
|
|
expected_output={
|
|
"py_executable": f"{find_uv_bin()} run --no-project --python {sys.executable} python",
|
|
"working_dir": os.getcwd(),
|
|
},
|
|
)
|
|
|
|
# Check in the case that a module is use for "uv run" and there is
|
|
# an argument immediately behind it
|
|
check_uv_run(
|
|
cmd=[
|
|
find_uv_bin(),
|
|
"run",
|
|
"--no-project",
|
|
"-m",
|
|
"ray._private.runtime_env.uv_runtime_env_hook",
|
|
"--extra-args",
|
|
],
|
|
runtime_env={},
|
|
expected_output={
|
|
"py_executable": f"{find_uv_bin()} run --no-project --python {sys.executable} python",
|
|
"working_dir": os.getcwd(),
|
|
},
|
|
)
|
|
|
|
|
|
def test_uv_run_parser():
|
|
from ray._private.runtime_env.uv_runtime_env_hook import (
|
|
_create_uv_run_parser,
|
|
_parse_args,
|
|
)
|
|
|
|
parser = _create_uv_run_parser()
|
|
|
|
options, command = _parse_args(parser, ["script.py"])
|
|
assert command == ["script.py"]
|
|
|
|
options, command = _parse_args(parser, ["--with", "requests", "example.py"])
|
|
assert options.with_packages == ["requests"]
|
|
assert command == ["example.py"]
|
|
|
|
options, command = _parse_args(parser, ["--python", "3.10", "example.py"])
|
|
assert options.python == "3.10"
|
|
assert command == ["example.py"]
|
|
|
|
options, command = _parse_args(
|
|
parser, ["--no-project", "script.py", "some", "args"]
|
|
)
|
|
assert options.no_project
|
|
assert command == ["script.py", "some", "args"]
|
|
|
|
options, command = _parse_args(
|
|
parser, ["--isolated", "-m", "module_name", "--extra-args"]
|
|
)
|
|
assert options.module == "module_name"
|
|
assert options.isolated
|
|
assert command == ["--extra-args"]
|
|
|
|
options, command = _parse_args(
|
|
parser,
|
|
[
|
|
"--isolated",
|
|
"--extra",
|
|
"vllm",
|
|
"-m",
|
|
"my_module.submodule",
|
|
"--model",
|
|
"Qwen/Qwen3-32B",
|
|
],
|
|
)
|
|
assert options.isolated
|
|
assert options.extras == ["vllm"]
|
|
assert options.module == "my_module.submodule"
|
|
assert command == ["--model", "Qwen/Qwen3-32B"]
|
|
|
|
|
|
@pytest.mark.skipif(sys.platform == "win32", reason="Not ported to Windows yet.")
|
|
def test_uv_run_runtime_env_hook_e2e(shutdown_only, temp_dir):
|
|
tmp_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_dir / "output.txt"}", "w") as out:
|
|
json.dump(ray.get(f.remote()), out)
|
|
"""
|
|
|
|
working_dir = tmp_dir / "working_dir"
|
|
working_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
script_file = working_dir / "script.py"
|
|
with open(script_file, "w") as f:
|
|
f.write(script)
|
|
f.close()
|
|
|
|
subprocess.run(
|
|
[
|
|
find_uv_bin(),
|
|
"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",
|
|
str(script_file),
|
|
],
|
|
stdin=subprocess.PIPE,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
env={
|
|
"PYTHONPATH": ":".join(sys.path),
|
|
"PATH": os.environ["PATH"],
|
|
},
|
|
cwd=working_dir,
|
|
check=True,
|
|
)
|
|
with open(tmp_dir / "output.txt") as f:
|
|
assert json.load(f) == {"working_dir_files": os.listdir(working_dir)}
|
|
|
|
|
|
@pytest.mark.skipif(sys.platform == "win32", reason="Not ported to Windows yet.")
|
|
@pytest.mark.parametrize(
|
|
"ray_start_cluster_head_with_env_vars",
|
|
[
|
|
{
|
|
"include_dashboard": True,
|
|
}
|
|
],
|
|
indirect=True,
|
|
)
|
|
def test_uv_run_runtime_env_hook_e2e_job(
|
|
ray_start_cluster_head_with_env_vars, temp_dir
|
|
):
|
|
cluster = ray_start_cluster_head_with_env_vars
|
|
assert wait_until_server_available(cluster.webui_url) is True
|
|
webui_url = format_web_url(cluster.webui_url)
|
|
|
|
tmp_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_dir / "output.txt"}", "w") as out:
|
|
json.dump(ray.get(f.remote()), out)
|
|
"""
|
|
|
|
working_dir = tmp_dir / "working_dir"
|
|
working_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
script_file = working_dir / "script.py"
|
|
with open(script_file, "w") as f:
|
|
f.write(script)
|
|
f.close()
|
|
|
|
requirements_file = working_dir / "requirements.txt"
|
|
with open(requirements_file, "w") as f:
|
|
f.write("emoji\n")
|
|
f.close()
|
|
|
|
# Test job submission
|
|
runtime_env_json = (
|
|
'{"env_vars": {"PYTHONPATH": "' + ":".join(sys.path) + '"}, "working_dir": "."}'
|
|
)
|
|
subprocess.run(
|
|
[
|
|
"ray",
|
|
"job",
|
|
"submit",
|
|
"--runtime-env-json",
|
|
runtime_env_json,
|
|
"--",
|
|
find_uv_bin(),
|
|
"run",
|
|
"--with-requirements",
|
|
str(requirements_file),
|
|
"--no-project",
|
|
str(script_file),
|
|
],
|
|
stdin=subprocess.PIPE,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
env={
|
|
"PATH": os.environ["PATH"],
|
|
"RAY_ADDRESS": webui_url,
|
|
},
|
|
cwd=working_dir,
|
|
check=True,
|
|
)
|
|
with open(tmp_dir / "output.txt") as f:
|
|
assert json.load(f) == {"working_dir_files": os.listdir(working_dir)}
|
|
|
|
|
|
@pytest.mark.skipif(sys.platform == "win32", reason="Not ported to Windows yet.")
|
|
@pytest.mark.parametrize(
|
|
"call_ray_start",
|
|
[{"cmd": "ray start --head --num-cpus=1 --ray-client-server-port=50052"}],
|
|
indirect=True,
|
|
)
|
|
def test_uv_run_ray_client_mode(call_ray_start, tmp_working_dir):
|
|
"""Test that UV environment works with Ray Client mode.
|
|
|
|
This verifies the fix for: https://github.com/ray-project/ray/issues/57991
|
|
|
|
The test verifies that when connecting via Ray Client, the client-side UV
|
|
hook correctly detects the UV environment and propagates the configuration
|
|
(py_executable, working_dir) to the cluster workers.
|
|
|
|
Architecture:
|
|
- Ray cluster runs in separate subprocess (started by call_ray_start fixture)
|
|
- Ray Client server is built-in (via --ray-client-server-port flag)
|
|
- This test process runs under 'uv run' and connects as a client
|
|
- UV hook detects UV environment in client process and propagates config
|
|
|
|
This is a cleaner approach than running cluster and client in the same process.
|
|
"""
|
|
from ray._private.runtime_env.uv_runtime_env_hook import _get_uv_run_cmdline
|
|
from ray.job_config import JobConfig
|
|
|
|
# Verify test is running under UV (prerequisite for this test)
|
|
uv_cmdline = _get_uv_run_cmdline()
|
|
if uv_cmdline is None:
|
|
pytest.skip("Test must run under 'uv run'")
|
|
|
|
tmp_dir = tmp_working_dir
|
|
|
|
# Connect as Ray Client from UV environment
|
|
# The UV hook will detect the 'uv run' process and set py_executable
|
|
ray.init(
|
|
"ray://localhost:50052", # Connect to the client server started by fixture
|
|
job_config=JobConfig(
|
|
runtime_env={
|
|
"working_dir": tmp_dir,
|
|
# Use system environment to find Ray installation
|
|
"env_vars": {"PYTHONPATH": ":".join(sys.path)},
|
|
}
|
|
),
|
|
)
|
|
|
|
try:
|
|
# Verify UV config was applied to runtime_env
|
|
context = ray.get_context()
|
|
conn_info = context.client_worker.connection_info()
|
|
runtime_env = conn_info.get("job_config", {}).get("runtime_env", {})
|
|
assert "py_executable" in runtime_env, "UV hook should set py_executable"
|
|
assert (
|
|
"uv run" in runtime_env["py_executable"]
|
|
), "py_executable should contain 'uv run'"
|
|
|
|
@ray.remote
|
|
def emojize():
|
|
import emoji
|
|
|
|
return emoji.emojize("Ray rocks :thumbs_up:")
|
|
|
|
# This should work because UV hook detected and propagated UV config
|
|
result = ray.get(emojize.remote())
|
|
assert (
|
|
result == "Ray rocks 👍"
|
|
), "UV should have installed emoji package on workers"
|
|
|
|
finally:
|
|
ray.shutdown()
|
|
|
|
|
|
@pytest.mark.skipif(sys.platform == "win32", reason="Not ported to Windows yet.")
|
|
def test_uv_run_worker_startup_failed(shutdown_only, monkeypatch):
|
|
"""Test that exception is raised when pop worker always fails
|
|
|
|
This test uses an invalid py_executable command that will fail during worker startup,
|
|
causing the worker registration to time out multiple times and eventually raise
|
|
RaySystemError.
|
|
"""
|
|
with monkeypatch.context() as m:
|
|
# Set environment variables to speed up the test (only for this test context)
|
|
# Reduce timeout to 3 seconds and max retries to 3
|
|
m.setenv("RAY_worker_register_timeout_seconds", "3")
|
|
m.setenv("RAY_pop_worker_max_retries", "3")
|
|
|
|
# Use an invalid command that will fail
|
|
runtime_env = {
|
|
"py_executable": f"{find_uv_bin()} run --with nonexistent-package-xyz123 --no-project",
|
|
}
|
|
|
|
ray.init(runtime_env=runtime_env)
|
|
|
|
@ray.remote
|
|
def failing_task():
|
|
return "should not succeed"
|
|
|
|
# The task should raise RaySystemError
|
|
with pytest.raises(RaySystemError) as exc_info:
|
|
ray.get(failing_task.remote(), timeout=30)
|
|
|
|
# Assert that the exception was actually raised and contains expected message
|
|
assert exc_info.type is RaySystemError
|
|
assert "Failed to startup worker" in str(exc_info.value)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-sv", __file__]))
|