Merge pull request #104 from fal-ai/support-uv-resolver

feat: support UV as a pip replacement

.github/workflows/test.yml

@@ -51,7 +51,12 @@ jobs:
         python -m pip install -r dev-requirements.txt
         python -m pip install -e ".[build]"
 
+    - name: Install uv
+      if: ${{ matrix.python != '3.7' }}
+      run: |
+        python -m pip install uv
+
     - name: Test
       run: |
         export ISOLATE_PYENV_EXECUTABLE=pyenv/bin/pyenv
-        python -m pytest
+        python -m pytest -vvv

src/isolate/backends/common.py

@@ -238,3 +238,18 @@ def optional_import(module_name: str) -> ModuleType:
             f"accessing {module_name!r} import functionality. Please try: "
             f"'$ pip install \"isolate[build]\"' to install it."
         ) from exc
+
+
+@lru_cache(4)
+def get_executable(command: str, home: str | None = None) -> Path:
+    for path in [home, None]:
+        binary_path = shutil.which(command, path=path)
+        if binary_path is not None:
+            return Path(binary_path)
+    else:
+        # TODO: we should probably show some instructions on how you
+        # can install conda here.
+        raise FileNotFoundError(
+            f"Could not find {command} executable. If {command} executable is not available by default, please point isolate "
+            f" to the path where conda binary is available '{home}'."
+        )

src/isolate/backends/conda.py

@@ -14,6 +14,7 @@
 from isolate.backends import BaseEnvironment, EnvironmentCreationError
 from isolate.backends.common import (
     active_python,
+    get_executable,
     logged_io,
     optional_import,
     sha256_digest_of,
@@ -48,7 +49,6 @@ class CondaEnvironment(BaseEnvironment[Path]):
     _exec_home: Optional[str] = _ISOLATE_MAMBA_HOME
     _exec_command: Optional[str] = _MAMBA_COMMAND
 
-
     @classmethod
     def from_config(
         cls,
@@ -162,15 +162,17 @@ def destroy(self, connection_key: Path) -> None:
 
     def _run_create(self, env_path: str, env_name: str) -> None:
         if self._exec_command == "conda":
-            self._run_conda("env", "create", "--force", "--prefix", env_path, "-f", env_name)
+            self._run_conda(
+                "env", "create", "--force", "--prefix", env_path, "-f", env_name
+            )
         else:
             self._run_conda("env", "create", "--prefix", env_path, "-f", env_name)
 
     def _run_destroy(self, connection_key: str) -> None:
-        self._run_conda("remove","--yes","--all","--prefix", connection_key)
+        self._run_conda("remove", "--yes", "--all", "--prefix", connection_key)
 
     def _run_conda(self, *args: Any) -> None:
-        conda_executable = _get_executable(self._exec_command, self._exec_home)
+        conda_executable = get_executable(self._exec_command, self._exec_home)
         with logged_io(partial(self.log, level=LogLevel.INFO)) as (stdout, stderr):
             subprocess.check_call(
                 [conda_executable, *args],
@@ -186,21 +188,6 @@ def open_connection(self, connection_key: Path) -> PythonIPC:
         return PythonIPC(self, connection_key)
 
 
-@functools.lru_cache(1)
-def _get_executable(command: str, home: str | None = None) -> Path:
-    for path in [home, None]:
-        conda_path = shutil.which(command, path=path)
-        if conda_path is not None:
-            return Path(conda_path)
-    else:
-        # TODO: we should probably show some instructions on how you
-        # can install conda here.
-        raise FileNotFoundError(
-            "Could not find conda executable. If conda executable is not available by default, please point isolate "
-            " to the path where conda binary is available 'ISOLATE_CONDA_HOME'."
-        )
-
-
 def _depends_on(
     dependencies: List[Union[str, Dict[str, List[str]]]],
     package_name: str,

src/isolate/backends/virtualenv.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import os
+import shlex
 import shutil
 import subprocess
 from dataclasses import dataclass, field
@@ -11,6 +12,7 @@
 from isolate.backends import BaseEnvironment, EnvironmentCreationError
 from isolate.backends.common import (
     active_python,
+    get_executable,
     get_executable_path,
     logged_io,
     optional_import,
@@ -20,6 +22,9 @@
 from isolate.connections import PythonIPC
 from isolate.logs import LogLevel
 
+_UV_RESOLVER_EXECUTABLE = os.environ.get("ISOLATE_UV_EXE", "uv")
+_UV_RESOLVER_HOME = os.getenv("ISOLATE_UV_HOME")
+
 
 @dataclass
 class VirtualPythonEnvironment(BaseEnvironment[Path]):
@@ -30,6 +35,7 @@ class VirtualPythonEnvironment(BaseEnvironment[Path]):
     python_version: Optional[str] = None
     extra_index_urls: List[str] = field(default_factory=list)
     tags: List[str] = field(default_factory=list)
+    resolver: Optional[str] = None
 
     @classmethod
     def from_config(
@@ -39,6 +45,10 @@ def from_config(
     ) -> BaseEnvironment:
         environment = cls(**config)
         environment.apply_settings(settings)
+        if environment.resolver not in ("uv", None):
+            raise ValueError(
+                "Only 'uv' is supported as a resolver for virtualenv environments."
+            )
         return environment
 
     @property
@@ -49,13 +59,20 @@ def key(self) -> str:
         else:
             constraints = []
 
+        extras = []
+        if not self.resolver:
+            extras.append(f"resolver={self.resolver}")
+
         active_python_version = self.python_version or active_python()
         return sha256_digest_of(
             active_python_version,
             *self.requirements,
             *constraints,
             *self.extra_index_urls,
             *sorted(self.tags),
+            # This is backwards compatible with environments not using
+            # the 'resolver' field.
+            *extras,
         )
 
     def install_requirements(self, path: Path) -> None:
@@ -69,9 +86,22 @@ def install_requirements(self, path: Path) -> None:
             return None
 
         self.log(f"Installing requirements: {', '.join(self.requirements)}")
+        environ = os.environ.copy()
+
+        if self.resolver == "uv":
+            # Set VIRTUAL_ENV to the actual path of the environment since that is
+            # how uv discovers the environment. This is necessary when using uv
+            # as the resolver.
+            environ["VIRTUAL_ENV"] = str(path)
+            base_pip_cmd = [
+                get_executable(_UV_RESOLVER_EXECUTABLE, _UV_RESOLVER_HOME),
+                "pip",
+            ]
+        else:
+            base_pip_cmd = [get_executable_path(path, "pip")]
 
         pip_cmd: List[Union[str, os.PathLike]] = [
-            get_executable_path(path, "pip"),
+            *base_pip_cmd,  # type: ignore
             "install",
             *self.requirements,
         ]
@@ -87,6 +117,7 @@ def install_requirements(self, path: Path) -> None:
                     pip_cmd,
                     stdout=stdout,
                     stderr=stderr,
+                    env=environ,
                 )
             except subprocess.SubprocessError as exc:
                 raise EnvironmentCreationError(f"Failure during 'pip install': {exc}")

src/isolate/connections/common.py

@@ -2,8 +2,8 @@
 
 import importlib
 import os
-from dataclasses import dataclass
 from contextlib import contextmanager
+from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Iterator, Optional, cast
 
 from tblib import Traceback, TracebackParseError

src/isolate/server/server.py

@@ -134,7 +134,9 @@ def _allocate_new_agent(
                     agent.terminate()
 
         bound_context = ExitStack()
-        stub = bound_context.enter_context(connection._establish_bridge(max_wait_timeout=MAX_GRPC_WAIT_TIMEOUT))
+        stub = bound_context.enter_context(
+            connection._establish_bridge(max_wait_timeout=MAX_GRPC_WAIT_TIMEOUT)
+        )
         return RunnerAgent(stub, queue, bound_context)
 
     def _identify(self, connection: LocalPythonGRPC) -> Tuple[Any, ...]:

tests/test_backends.py

@@ -1,3 +1,4 @@
+import shlex
 import subprocess
 import sys
 import textwrap
@@ -10,8 +11,8 @@
 
 import isolate
 from isolate.backends import BaseEnvironment, EnvironmentCreationError
-from isolate.backends.common import sha256_digest_of
-from isolate.backends.conda import CondaEnvironment, _get_executable
+from isolate.backends.common import get_executable, sha256_digest_of
+from isolate.backends.conda import CondaEnvironment
 from isolate.backends.local import LocalPythonEnvironment
 from isolate.backends.pyenv import PyenvEnvironment, _get_pyenv_executable
 from isolate.backends.remote import IsolateServer
@@ -80,7 +81,9 @@ def test_create_generic_env_empty(self, tmp_path):
         with pytest.raises(ModuleNotFoundError):
             self.get_example_version(environment, connection_key)
 
-    @pytest.mark.skip(reason="This test fails on the 'both the original one and the duplicate one will be gone' section")
+    @pytest.mark.skip(
+        reason="This test fails on the 'both the original one and the duplicate one will be gone' section"
+    )
     def test_create_generic_env_cached(self, tmp_path, monkeypatch):
         environment_1 = self.get_project_environment(tmp_path, "old-example-project")
         environment_2 = self.get_project_environment(tmp_path, "new-example-project")
@@ -206,6 +209,12 @@ def test_custom_python_version(self, tmp_path):
             assert python_version.startswith(python_version)
 
 
+try:
+    UV_PATH = get_executable("uv")
+except FileNotFoundError:
+    UV_PATH = None
+
+
 class TestVirtualenv(GenericEnvironmentTests):
 
     backend_cls = VirtualPythonEnvironment
@@ -361,23 +370,40 @@ def test_tags_in_key(self, tmp_path, monkeypatch):
             "isolate.backends.pyenv._get_pyenv_executable", lambda: 1 / 0
         )
 
-        constraints = self.configs['old-example-project']
+        constraints = self.configs["old-example-project"]
         tagged = constraints.copy()
-        tagged['tags'] = ['tag1', 'tag2']
+        tagged["tags"] = ["tag1", "tag2"]
         tagged_environment = self.get_environment(tmp_path, tagged)
 
         no_tagged_environment = self.get_environment(tmp_path, constraints)
-        assert tagged_environment.key != no_tagged_environment.key, "Tagged environment should have different key"
+        assert (
+            tagged_environment.key != no_tagged_environment.key
+        ), "Tagged environment should have different key"
 
         tagged["tags"] = ["tag2", "tag1"]
         tagged_environment_2 = self.get_environment(tmp_path, tagged)
-        assert tagged_environment.key == tagged_environment_2.key, "Tag order should not matter"
+        assert (
+            tagged_environment.key == tagged_environment_2.key
+        ), "Tag order should not matter"
+
+    @pytest.mark.skipif(not UV_PATH, reason="uv is not available")
+    def test_try_using_uv(self, tmp_path):
+        environment = self.get_environment(
+            tmp_path,
+            {
+                "requirements": [f"pyjokes==0.5"],
+                "resolver": "uv",
+            },
+        )
+        connection_key = environment.create()
+        pyjokes_version = self.get_example_version(environment, connection_key)
+        assert pyjokes_version == "0.5.0"
 
 
 # Since mamba is an external dependency, we'll skip tests using it
 # if it is not installed.
 try:
-    _get_executable("micromamba")
+    get_executable("micromamba")
 except FileNotFoundError:
     IS_MAMBA_AVAILABLE = False
 else:
@@ -527,17 +553,22 @@ def test_add_pip_dependencies(self, tmp_path, configuration):
         assert "agent" in pip_dep  # And pip dependency is added
 
     def test_tags_in_key(self, tmp_path):
-        constraints = self.configs['old-example-project']
+        constraints = self.configs["old-example-project"]
         tagged = constraints.copy()
-        tagged['tags'] = ['tag1', 'tag2']
+        tagged["tags"] = ["tag1", "tag2"]
         tagged_environment = self.get_environment(tmp_path, tagged)
 
         no_tagged_environment = self.get_environment(tmp_path, constraints)
-        assert tagged_environment.key != no_tagged_environment.key, "Tagged environment should have different key"
+        assert (
+            tagged_environment.key != no_tagged_environment.key
+        ), "Tagged environment should have different key"
 
         tagged["tags"] = ["tag2", "tag1"]
         tagged_environment_2 = self.get_environment(tmp_path, tagged)
-        assert tagged_environment.key == tagged_environment_2.key, "Tag order should not matter"
+        assert (
+            tagged_environment.key == tagged_environment_2.key
+        ), "Tag order should not matter"
+
 
 def test_local_python_environment():
     """Since 'local' environment does not support installation of extra dependencies