[core] [6/N] Allow users to specify uv version to install (#48634)

Signed-off-by: dentiny <[email protected]>

python/ray/_private/runtime_env/uv.py

@@ -72,23 +72,32 @@ 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",
             "pip",
             "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.

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

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

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