Add uv install

docs/api/python/image.rst

@@ -25,6 +25,7 @@ ImageSteupStepType
     .. autoattribute:: SET_ENV_VARS
     .. autoattribute:: PIP_INSTALL
     .. autoattribute:: CONDA_INSTALL
+    .. autoattribute:: UV_INSTALL
     .. autoattribute:: SYNC_PACKAGE
 
 ImageSetupStep

runhouse/resources/hardware/cluster.py

@@ -678,6 +678,7 @@ def _sync_image_to_cluster(self, parallel: bool = True):
         logger.info(f"Syncing default image {self.image} to cluster.")
 
         secrets_to_sync = []
+        uv_install = False
 
         for step in self.image.setup_steps:
             if step.step_type == ImageSetupStepType.SYNC_SECRETS:
@@ -687,6 +688,16 @@ def _sync_image_to_cluster(self, parallel: bool = True):
                 image_env_vars = _process_env_vars(step.kwargs.get("env_vars"))
                 env_vars.update(image_env_vars)
                 continue
+            elif step.step_type == ImageSetupStepType.UV_INSTALL and not uv_install:
+                self._run_setup_step(
+                    step=ImageSetupStep(
+                        step_type=ImageSetupStepType.PIP_INSTALL,
+                        reqs=["uv"],
+                        conda_env_name=self.conda_env_name,
+                    ),
+                    parallel=parallel,
+                )
+                uv_install = True
 
             self._run_setup_step(step, env_vars, parallel)
 
@@ -790,6 +801,26 @@ def pip_install(
             override_remote_version=override_remote_version,
         )
 
+    def uv_install(
+        self,
+        reqs: List[Union["Package", str]],
+        node: Optional[str] = None,
+        conda_env_name: Optional[str] = None,
+        force_sync_local: bool = False,
+    ):
+        from runhouse.resources.packages.package import Package
+
+        uv_packages = [
+            Package.from_string(f"uv:{req}") if isinstance(req, str) else req
+            for req in reqs
+        ]
+        self.install_packages(
+            reqs=uv_packages,
+            node=node,
+            conda_env_name=conda_env_name,
+            force_sync_local=force_sync_local,
+        )
+
     def conda_install(
         self,
         reqs: List[Union["Package", str]],

runhouse/resources/hardware/utils.py

@@ -192,6 +192,30 @@ def _do_setup_step_for_node(cluster, setup_step, node, env_vars):
             conda_env_name=setup_step.kwargs.get("conda_env_name"),
             node=node,
         )
+    elif setup_step.step_type == ImageSetupStepType.PIP_INSTALL:
+        cluster.pip_install(
+            setup_step.kwargs.get("reqs"),
+            conda_env_name=setup_step.kwargs.get("conda_env_name"),
+            node=node,
+        )
+    elif setup_step.step_type == ImageSetupStepType.UV_INSTALL:
+        cluster.uv_install(
+            setup_step.kwargs.get("reqs"),
+            conda_env_name=setup_step.kwargs.get("conda_env_name"),
+            node=node,
+        )
+    elif setup_step.step_type == ImageSetupStepType.CONDA_INSTALL:
+        cluster.conda_install(
+            setup_step.kwargs.get("reqs"),
+            conda_env_name=setup_step.kwargs.get("conda_env_name"),
+            node=node,
+        )
+    elif setup_step.step_type == ImageSetupStepType.SYNC_PACKAGE:
+        cluster.sync_package(
+            setup_step.kwargs.get("package"),
+            conda_env_name=setup_step.kwargs.get("conda_env_name"),
+            node=node,
+        )
     elif setup_step.step_type == ImageSetupStepType.CMD_RUN:
         return run_setup_command(
             cmd=setup_step.kwargs.get("command"),

runhouse/resources/images/image.py

@@ -12,6 +12,7 @@ class ImageSetupStepType(Enum):
     SYNC_SECRETS = "sync_secrets"
     SET_ENV_VARS = "set_env_vars"
     PIP_INSTALL = "pip_install"
+    UV_INSTALL = "uv_install"
     CONDA_INSTALL = "conda_install"
     SYNC_PACKAGE = "sync_package"
 
@@ -214,6 +215,26 @@ def pip_install(
         )
         return self
 
+    def uv_install(
+        self, reqs: List[Union["Package", str]], conda_env_name: Optional[str] = None
+    ):
+        """Uv pip install the given packages.
+
+        Args:
+            reqs (List[Package or str]): List of packages to uv pip install on cluster and env.
+            conda_env_name (str, optional): Name of conda env to install the package in, if relevant. If left empty,
+                defaults to base environment. (Default: ``None``)
+        """
+
+        self.setup_steps.append(
+            ImageSetupStep(
+                step_type=ImageSetupStepType.UV_INSTALL,
+                reqs=reqs,
+                conda_env_name=conda_env_name or self.conda_env_name,
+            )
+        )
+        return self
+
     def conda_install(
         self, reqs: List[Union["Package", str]], conda_env_name: Optional[str] = None
     ):

runhouse/resources/packages/package.py

@@ -30,7 +30,7 @@
 )
 
 
-INSTALL_METHODS = {"local", "pip", "conda"}
+INSTALL_METHODS = {"local", "pip", "uv", "conda"}
 
 logger = get_logger(__name__)
 
@@ -171,7 +171,10 @@ def _pip_install_cmd(
         self,
         conda_env_name: Optional[str] = None,
         cluster: "Cluster" = None,
+        uv: bool = None,
     ):
+        if uv is None:
+            uv = self.install_method == "uv"
         install_args = f" {self.install_args}" if self.install_args else ""
         install_extras = f"[{self.install_extras}]" if self.install_extras else ""
         if isinstance(self.install_target, InstallTarget):
@@ -198,9 +201,13 @@ def _pip_install_cmd(
             install_cmd = install_target + install_args
 
         install_cmd = f"pip install {self._install_cmd_for_torch(install_cmd, cluster)}"
-        install_cmd = self._prepend_python_executable(
-            install_cmd, cluster=cluster, conda_env_name=conda_env_name
-        )
+        if uv:
+            install_cmd = f"uv {install_cmd}"
+        else:
+            # uv doesn't need python executable
+            install_cmd = self._prepend_python_executable(
+                install_cmd, cluster=cluster, conda_env_name=conda_env_name
+            )
         install_cmd = self._prepend_env_command(
             install_cmd, conda_env_name=conda_env_name
         )
@@ -232,7 +239,7 @@ def _install_and_validate_output(
             if INSUFFICIENT_DISK_MSG in stdout or INSUFFICIENT_DISK_MSG in stderr:
                 raise InsufficientDiskError(command=install_cmd)
             raise RuntimeError(
-                f"Pip install {install_cmd} failed, check that the package exists and is available for your platform."
+                f"{self.install_method} install '{install_cmd}' failed, check that the package exists and is available for your platform."
             )
 
     def _install(
@@ -257,7 +264,7 @@ def _install(
         """
         logger.info(f"Installing {str(self)} with method {self.install_method}.")
 
-        if self.install_method == "pip":
+        if self.install_method in ["pip", "uv"]:
 
             # If this is a generic pip package, with no version pinned, we want to check if there is a version
             # already installed. If there is, and ``override_remote_version`` is not set to ``True``, then we ignore
@@ -302,7 +309,9 @@ def _install(
                         )
 
             install_cmd = self._pip_install_cmd(
-                conda_env_name=conda_env_name, cluster=cluster
+                conda_env_name=conda_env_name,
+                cluster=cluster,
+                uv=(self.install_method == "uv"),
             )
             logger.info(f"Running via install_method pip: {install_cmd}")
             self._install_and_validate_output(
@@ -533,7 +542,7 @@ def from_string(specifier: str, dryrun: bool = False):
                 else:
                     # We want to preferrably install this version of the package server-side
                     install_method == install_method or "pip"
-                    if install_method == "pip":
+                    if install_method in ["pip", "uv"]:
                         preferred_version = locally_installed_version
 
         # If install method is still not set, default to pip
@@ -544,7 +553,7 @@ def from_string(specifier: str, dryrun: bool = False):
             return Package(
                 install_target=target, install_method=install_method, dryrun=dryrun
             )
-        elif install_method in ["pip", "conda"]:
+        elif install_method in ["pip", "uv", "conda"]:
             return Package(
                 install_target=target,
                 install_args=args,

runhouse/servers/obj_store.py

@@ -2007,7 +2007,7 @@ async def ainstall_package_in_all_nodes_and_processes(
 
         logger.info(f"Installing {str(package)} with method {package.install_method}.")
 
-        if package.install_method == "pip":
+        if package.install_method in ["pip", "uv"]:
 
             # If this is a generic pip package, with no version pinned, we want to check if there is a version
             # already installed. If there is, then we ignore preferred version and leave the existing version.
@@ -2031,9 +2031,17 @@ async def ainstall_package_in_all_nodes_and_processes(
                             _path_to_sync_to_on_cluster=package.install_target,
                         )
 
-            install_cmd = package._pip_install_cmd(conda_env_name=conda_env_name)
-            logger.info(f"Running via install_method pip: {install_cmd}")
-            await self._install_packages_in_nodes_helper(install_cmd)
+            install_cmd = package._pip_install_cmd(
+                conda_env_name=conda_env_name, uv=(package.install_method == "uv")
+            )
+            logger.info(
+                f"Running via install_method {package.install_method}: {install_cmd}"
+            )
+            run_cmd_results = await self.arun_bash_command_on_all_nodes(install_cmd)
+            if any(run_cmd_result != 0 for run_cmd_result in run_cmd_results):
+                raise RuntimeError(
+                    f"Pip install {install_cmd} failed, check that the package exists and is available for your platform."
+                )
 
         elif package.install_method == "conda":
             install_cmd = package._conda_install_cmd(conda_env_name=conda_env_name)

tests/conftest.py

@@ -302,6 +302,7 @@ def event_loop():
     package,  # noqa: F401
     pip_package,  # noqa: F401
     s3_package,  # noqa: F401
+    uv_package,  # noqa: F401
 )
 
 from tests.fixtures.secret_fixtures import (

tests/fixtures/package_fixtures.py

@@ -26,6 +26,17 @@ def pip_package():
     return package
 
 
+@pytest.fixture(scope="session")
+def uv_package():
+    args = {
+        "install_target": "scipy",
+        "install_method": "uv",
+    }
+    package = rh.Package(**args)
+    init_args[id(package)] = args
+    return package
+
+
 @pytest.fixture(scope="session")
 def conda_package():
     args = {

tests/test_resources/test_data/test_package.py

@@ -23,6 +23,7 @@ class TestPackage(tests.test_resources.test_resource.TestResource):
 
     packages = [
         "pip_package",
+        "uv_package",
         "conda_package",
         "local_package",
     ]
@@ -62,6 +63,7 @@ def test_package_factory_and_properties(self, package):
         [
             "numpy",
             "pip:numpy",
+            "uv:numpy",
             "conda:numpy",
             "local:./",
         ],
@@ -73,6 +75,8 @@ def test_from_string(self, pkg_str):
 
         if package.install_method == "local":
             assert isinstance(package.install_target, InstallTarget)
+        else:
+            assert package.install_target == "numpy"
 
     # --------- test install command ---------
     @pytest.mark.level("unit")
@@ -82,6 +86,13 @@ def test_pip_install_cmd(self, pip_package):
             == f'{sys.executable} -m pip install "{pip_package.install_target}"'
         )
 
+    @pytest.mark.level("unit")
+    def test_uv_install_cmd(self, uv_package):
+        assert (
+            uv_package._pip_install_cmd()
+            == f'uv pip install "{uv_package.install_target}"'
+        )
+
     @pytest.mark.level("unit")
     def test_conda_install_cmd(self, conda_package):
         assert (