diff --git a/python/ray/_private/runtime_env/uv.py b/python/ray/_private/runtime_env/uv.py index 345c37b6c30b..270fb1ad891e 100644 --- a/python/ray/_private/runtime_env/uv.py +++ b/python/ray/_private/runtime_env/uv.py @@ -72,15 +72,23 @@ def __init__( self._uv_env = os.environ.copy() self._uv_env.update(self._runtime_env.env_vars()) - # TODO(hjiang): Check `uv` existence before installation, so we don't blindly - # install. async def _install_uv( self, path: str, cwd: str, pip_env: dict, logger: logging.Logger ): - """Before package install, make sure `uv` is installed.""" + """Before package install, make sure the required version `uv` (if specifieds) + is installed. + """ virtualenv_path = virtualenv_utils.get_virtualenv_path(path) python = virtualenv_utils.get_virtualenv_python(path) + def _get_uv_exec_to_install() -> str: + """Get `uv` executable with version to install.""" + uv_version = self._uv_config.get("uv_version", None) + if uv_version: + return f"uv{uv_version}" + # Use default version. + return "uv" + uv_install_cmd = [ python, "-m", @@ -88,7 +96,8 @@ async def _install_uv( "install", "--disable-pip-version-check", "--no-cache-dir", - "uv", + "--force-reinstall", + _get_uv_exec_to_install(), ] logger.info("Installing package uv to %s", virtualenv_path) await check_output_cmd(uv_install_cmd, logger=logger, cwd=cwd, env=pip_env) @@ -131,7 +140,12 @@ async def _install_uv_packages( uv_exists = await self._check_uv_existence(python, cwd, pip_env, logger) # Install uv, which acts as the default package manager. - if not uv_exists: + # + # TODO(hjiang): If `uv` in virtual env perfectly matches the version users + # require, we don't need to install also. It requires a different + # implementation to execute and check existence. Here we take the simpliest + # implementation, always reinstall the required version. + if (not uv_exists) or (self._uv_config.get("uv_version", None) is not None): await self._install_uv(path, cwd, pip_env, logger) # Avoid blocking the event loop. diff --git a/python/ray/_private/runtime_env/validation.py b/python/ray/_private/runtime_env/validation.py index dfd1abfe50d7..85e9c5f60afb 100644 --- a/python/ray/_private/runtime_env/validation.py +++ b/python/ray/_private/runtime_env/validation.py @@ -116,8 +116,7 @@ def parse_and_validate_conda(conda: Union[str, dict]) -> Union[str, dict]: # TODO(hjiang): More package installation options to implement: -# 1. Allow specific version of `uv` to use; as of now we only use default version. -# 2. `pip_check` has different semantics for `uv` and `pip`, see +# 1. `pip_check` has different semantics for `uv` and `pip`, see # https://github.com/astral-sh/uv/pull/2544/files, consider whether we need to support # it; or simply ignore the field when people come from `pip`. def parse_and_validate_uv(uv: Union[str, List[str], Dict]) -> Optional[Dict]: @@ -149,16 +148,21 @@ def parse_and_validate_uv(uv: Union[str, List[str], Dict]) -> Optional[Dict]: elif isinstance(uv, list) and all(isinstance(dep, str) for dep in uv): result = dict(packages=uv) elif isinstance(uv, dict): - if set(uv.keys()) - {"packages"}: + if set(uv.keys()) - {"packages", "uv_version"}: raise ValueError( "runtime_env['uv'] can only have these fields: " - "packages, but got: " + "packages and uv_version, but got: " f"{list(uv.keys())}" ) if "packages" not in uv: raise ValueError( f"runtime_env['uv'] must include field 'packages', but got {uv}" ) + if "uv_version" in uv and not isinstance(uv["uv_version"], str): + raise TypeError( + "runtime_env['uv']['uv_version'] must be of type str, " + f"got {type(uv['uv_version'])}" + ) result = uv.copy() if not isinstance(uv["packages"], list): diff --git a/python/ray/tests/test_runtime_env_uv.py b/python/ray/tests/test_runtime_env_uv.py index 215d85645566..bfcb4b57d060 100644 --- a/python/ray/tests/test_runtime_env_uv.py +++ b/python/ray/tests/test_runtime_env_uv.py @@ -69,6 +69,21 @@ def f(): ray.get(f.remote()) +# Specify uv version and check. +def test_uv_with_version_and_check(shutdown_only): + @ray.remote( + runtime_env={"uv": {"packages": ["requests==2.3.0"], "uv_version": "==0.4.0"}} + ) + def f(): + import pkg_resources + import requests + + assert pkg_resources.get_distribution("uv").version == "0.4.0" + assert requests.__version__ == "2.3.0" + + ray.get(f.remote()) + + # Package installation via requirements file. def test_package_install_with_requirements(shutdown_only, tmp_working_dir): requirements_file = tmp_working_dir diff --git a/python/ray/tests/unit/test_runtime_env_validation.py b/python/ray/tests/unit/test_runtime_env_validation.py index 564f35ca7258..27f7c61dee29 100644 --- a/python/ray/tests/unit/test_runtime_env_validation.py +++ b/python/ray/tests/unit/test_runtime_env_validation.py @@ -50,6 +50,12 @@ def test_parse_and_validate_uv(self, test_directory): with pytest.raises(ValueError): result = validation.parse_and_validate_uv({"random_key": "random_value"}) + # Valid case w/ uv version. + result = validation.parse_and_validate_uv( + {"packages": ["tensorflow"], "uv_version": "==0.4.30"} + ) + assert result == {"packages": ["tensorflow"], "uv_version": "==0.4.30"} + # Valid requirement files. _, requirements_file = test_directory requirements_file = requirements_file.resolve()