Allow uv version

Signed-off-by: dentiny <dentinyhao@gmail.com>

python/ray/_private/runtime_env/uv.py

@@ -79,14 +79,22 @@ async def _install_uv(
         virtualenv_path = virtualenv_utils.get_virtualenv_path(path)
         python = virtualenv_utils.get_virtualenv_python(path)
 
+        def _get_uv_exec_to_install():
+            """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",
+            _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)

python/ray/_private/runtime_env/validation.py

@@ -110,8 +110,7 @@ def parse_and_validate_conda(conda: Union[str, dict]) -> Union[str, dict]:
 # TODO(hjiang): More package installation options to implement:
 # 1. Allow users to pass in a local requirements.txt file, which relates to all
 # packages to install;
-# 2. Allow specific version of `uv` to use; as of now we only use default version.
-# 3. `pip_check` has different semantics for `uv` and `pip`, see
+# 2. `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]:
@@ -139,16 +138,21 @@ def parse_and_validate_uv(uv: Union[str, List[str], Dict]) -> Optional[Dict]:
     if 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 bool, "
+                f"got {type(uv['uv_version'])}"
+            )
 
         result = uv.copy()
         if not isinstance(uv["packages"], list):

python/ray/tests/test_runtime_env_uv.py

@@ -57,6 +57,19 @@ def f():
         ray.get(f.remote())
 
 
+# Specify uv version.
+def test_uv_with_version():
+    @ray.remote(
+        runtime_env={"uv": {"packages": ["requests==2.3.0"], "uv_version": "==0.4.0"}}
+    )
+    def f():
+        import requests
+
+        return requests.__version__
+
+    assert ray.get(f.remote()) == "2.3.0"
+
+
 if __name__ == "__main__":
     if os.environ.get("PARALLEL_CI"):
         sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))

python/ray/tests/unit/test_runtime_env_validation.py

@@ -54,6 +54,12 @@ def test_parse_and_validate_uv(self):
         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"}
+
 
 class TestValidatePip:
     def test_validate_pip_invalid_types(self):