Merge pull request #1 from fal-ai/add-server

feat: Add isolate server

Author: isidentical

.github/workflows/test.yml

@@ -38,8 +38,8 @@ jobs:
 
     - name: Install dependencies
       run: |
-        python -m pip install pytest
-        python -m pip install -e .
+        python -m pip install -r dev-requirements.txt
+        python -m pip install -e ".[server]"
 
     - name: Test
       run: python -m pytest

dev-requirements.txt

@@ -0,0 +1,3 @@
+pytest
+cloudpickle>=2.2.0
+dill>=0.3.5.1

poetry.lock

@@ -6,12 +6,21 @@ authors = ["Features & Labels <[email protected]>"]
 
 [tool.poetry.dependencies]
 python = ">=3.7"
-virtualenv = "^20"
-cloudpickle = "^2.2.0"
-importlib-metadata = "^5.0.0"
+virtualenv = ">=20.4"
+importlib-metadata = ">=4.4"
+flask = { version = "*", optional = true }
+marshmallow = { version = "*", optional = true }
+
+[tool.poetry.extras]
+server = ["flask", "marshmallow"]
+
+[tool.poetry.plugins."isolate.backends"]
+"virtualenv" = "isolate.backends.virtual_env:VirtualPythonEnvironment"
+"conda" = "isolate.backends.conda:CondaEnvironment"
+"local" = "isolate.backends.local:LocalPythonEnvironment"
 
 [build-system]
-requires = ["poetry-core>=1.0.0"]
+requires = ["poetry-core>=1.1.0"]
 build-backend = "poetry.core.masonry.api"
 
 [tool.isort]
@@ -20,7 +29,3 @@ force_grid_wrap=0
 include_trailing_comma=true
 multi_line_output=3
 use_parentheses=true
-
-[tool.poetry.plugins."isolate.backends"]
-"virtualenv" = "isolate.backends.virtual_env:VirtualPythonEnvironment"
-"conda" = "isolate.backends.conda:CondaEnvironment"

pyproject.toml

@@ -1,4 +1,3 @@
-from isolate.managed import EnvironmentManager
 from isolate.registry import prepare_environment
 
 __version__ = "0.1.0"

src/isolate/__init__.py

@@ -27,13 +27,18 @@
     "EnvironmentConnection",
     "BaseEnvironment",
     "UserException",
+    "EnvironmentCreationError",
 ]
 
 ConnectionKeyType = TypeVar("ConnectionKeyType")
 CallResultType = TypeVar("CallResultType")
 BasicCallable = Callable[[], CallResultType]
 
 
+class EnvironmentCreationError(Exception):
+    """Raised when the environment cannot be created."""
+
+
 class BaseEnvironment(Generic[ConnectionKeyType]):
     """Represents a managed environment definition for an isolatation backend
     that can be used to run Python code with different set of dependencies."""

src/isolate/backends/_base.py

@@ -1,18 +1,16 @@
 from __future__ import annotations
 
 import functools
+import hashlib
 import importlib
 import os
 import shutil
 import sysconfig
 import threading
 from contextlib import contextmanager
-from functools import partial
+from functools import lru_cache
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Callable, Iterator, Optional, Tuple
-
-if TYPE_CHECKING:
-    from isolate.backends._base import BaseEnvironment
+from typing import Any, Callable, Iterator, Optional, Tuple
 
 
 @contextmanager
@@ -167,3 +165,12 @@ def run_serialized(serialization_backend_name: str, data: bytes) -> Any:
     serialization_backend = importlib.import_module(serialization_backend_name)
     executable = serialization_backend.loads(data)
     return executable()
+
+
+@lru_cache(maxsize=None)
+def sha256_digest_of(*unique_fields: str, _join_char: str = "\n") -> str:
+    """Return the SHA256 digest that corresponds to the combined version
+    of 'unique_fields. The order is preserved."""
+
+    inner_text = _join_char.join(unique_fields).encode()
+    return hashlib.sha256(inner_text).hexdigest()

src/isolate/backends/common.py

@@ -1,15 +1,19 @@
 from __future__ import annotations
 
-import hashlib
 import os
 import shutil
 import subprocess
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Any, ClassVar, Dict, List
 
-from isolate.backends import BaseEnvironment
-from isolate.backends.common import cache_static, logged_io, rmdir_on_fail
+from isolate.backends import BaseEnvironment, EnvironmentCreationError
+from isolate.backends.common import (
+    cache_static,
+    logged_io,
+    rmdir_on_fail,
+    sha256_digest_of,
+)
 from isolate.backends.connections import PythonIPC
 from isolate.backends.context import GLOBAL_CONTEXT, ContextType
 
@@ -38,7 +42,7 @@ def from_config(
 
     @property
     def key(self) -> str:
-        return hashlib.sha256(" ".join(self.packages).encode()).hexdigest()
+        return sha256_digest_of(*self.packages)
 
     def create(self) -> Path:
         path = self.context.get_cache_dir(self) / self.key
@@ -52,20 +56,25 @@ def create(self) -> Path:
                 self.log(f"Installing packages: {', '.join(self.packages)}")
 
             with logged_io(self.log) as (stdout, stderr):
-                subprocess.check_call(
-                    [
-                        conda_executable,
-                        "create",
-                        "--yes",
-                        # The environment will be created under $BASE_CACHE_DIR/conda
-                        # so that in the future we can reuse it.
-                        "--prefix",
-                        path,
-                        *self.packages,
-                    ],
-                    stdout=stdout,
-                    stderr=stderr,
-                )
+                try:
+                    subprocess.check_call(
+                        [
+                            conda_executable,
+                            "create",
+                            "--yes",
+                            # The environment will be created under $BASE_CACHE_DIR/conda
+                            # so that in the future we can reuse it.
+                            "--prefix",
+                            path,
+                            *self.packages,
+                        ],
+                        stdout=stdout,
+                        stderr=stderr,
+                    )
+                except subprocess.SubprocessError as exc:
+                    raise EnvironmentCreationError(
+                        "Failure during 'conda create'"
+                    ) from exc
 
         return path
 

src/isolate/backends/conda.py

@@ -1,5 +1,5 @@
 from isolate.backends.connections.ipc import (
-    DualPythonIPC,
+    ExtendedPythonIPC,
     IsolatedProcessConnection,
     PythonIPC,
 )

src/isolate/backends/connections/__init__.py

@@ -1,5 +1,5 @@
 from isolate.backends.connections.ipc._base import (
-    DualPythonIPC,
+    ExtendedPythonIPC,
     IsolatedProcessConnection,
     PythonIPC,
 )

src/isolate/backends/connections/ipc/__init__.py

@@ -4,11 +4,11 @@
 import subprocess
 import time
 from contextlib import ExitStack, closing, contextmanager
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from functools import partial
 from multiprocessing.connection import ConnectionWrapper, Listener
 from pathlib import Path
-from typing import Any, ContextManager, Iterator, List, Tuple, Union
+from typing import Any, ContextManager, Dict, Iterator, List, Tuple, Union
 
 from isolate.backends import (
     BasicCallable,
@@ -208,8 +208,8 @@ def start_process(
 
     def _get_python_env(self):
         return {
-            "PYTHONUNBUFFERED": "1",  # We want to stream the logs as they come.
             **os.environ,
+            "PYTHONUNBUFFERED": "1",  # We want to stream the logs as they come.
         }
 
     def _get_python_cmd(
@@ -247,26 +247,36 @@ def _parse_agent_and_log(self, line: str, level: LogLevel) -> None:
         self.log(line, level=level, source=source)
 
 
+# TODO: should we actually merge this with PythonIPC since it is
+# simple enough and interchangeable?
 @dataclass
-class DualPythonIPC(PythonIPC):
-    """A dual-environment Python IPC implementation that
-    can run the agent process in an environment with its
-    Python and also load the shared libraries from a different
-    one.
-
-    The user of DualPythonIPC must ensure that the Python versions from
-    both of these environments are the same. Using different versions is
-    an undefined behavior.
+class ExtendedPythonIPC(PythonIPC):
+    """A Python IPC implementation that can also inherit packages from
+    other environments (e.g. a virtual environment that has the core
+    requirements like `dill` can be inherited on a new environment).
+
+    The given extra_inheritance_paths should be a list of paths that
+    comply with the sysconfig, and it should be ordered in terms of
+    priority (e.g. the first path will be the most prioritized one,
+    right after the current environment). So if two environments have
+    conflicting versions of the same package, the first one present in
+    the inheritance chain will be used.
+
+    This works by including the `site-packages` directory of the
+    inherited environment in the `PYTHONPATH` when starting the
+    agent process.
     """
 
-    secondary_path: Path
+    extra_inheritance_paths: List[Path] = field(default_factory=list)
 
-    def _get_python_env(self):
-        # We are going to use the primary environment to run the Python
-        # interpreter, but at the same time we are going to inherit all
-        # the packages from the secondary environment.
-
-        # The search order is important, we want the primary path to
-        # take precedence.
-        python_path = python_path_for(self.environment_path, self.secondary_path)
-        return {"PYTHONPATH": python_path, **super()._get_python_env()}
+    def _get_python_env(self) -> Dict[str, str]:
+        env_variables = super()._get_python_env()
+
+        if self.extra_inheritance_paths:
+            # The order here should reflect the order of the inheritance
+            # where the actual environment already takes precedence.
+            python_path = python_path_for(
+                self.environment_path, *self.extra_inheritance_paths
+            )
+            env_variables["PYTHONPATH"] = python_path
+        return env_variables

src/isolate/backends/connections/ipc/_base.py

@@ -15,9 +15,14 @@
 # one being the actual result of the given callable, and the other one is a boolean flag
 # indicating whether the callable has raised an exception or not.
 
+# WARNING: Please do not import anything outside of standard library before
+# the call to load_pth_files(). After that point, imports to installed packages
+# are allowed.
+
 import base64
 import importlib
 import os
+import site
 import sys
 import time
 from argparse import ArgumentParser
@@ -26,6 +31,29 @@
 from typing import ContextManager, Tuple
 
 
+def load_pth_files() -> None:
+    """Each site dir in Python can contain some .pth files, which are
+    basically instructions that tell Python to load other stuff. This is
+    generally used for editable installations, and just setting PYTHONPATH
+    won't make them expand so we need manually process them. Luckily, site
+    module can simply take the list of new paths and recognize them.
+
+    https://docs.python.org/3/tutorial/modules.html#the-module-search-path
+    """
+    python_path = os.getenv("PYTHONPATH")
+    if python_path is None:
+        return None
+
+    # TODO: The order here is the same as the one that is used for generating the
+    # PYTHONPATH. The only problem that might occur is that, on a chain with
+    # 3 ore more nodes (A, B, C), if X is installed as an editable package to
+    # B and a normal package to C, then C might actually take precedence. This
+    # will need to be fixed once we are dealing with more than 2 nodes and editable
+    # packages.
+    for site_dir in python_path.split(os.pathsep):
+        site.addsitedir(site_dir)
+
+
 def decode_service_address(address: str) -> Tuple[str, int]:
     host, port = base64.b64decode(address).decode("utf-8").rsplit(":", 1)
     return host, int(port)
@@ -115,7 +143,7 @@ def _get_shell_bootstrap() -> str:
     return " ".join(
         f"{session_variable}={os.getenv(session_variable)}"
         for session_variable in [
-            # PYTHONPATH is customized by the Dual Environment IPC
+            # PYTHONPATH is customized by the Extended Environment IPC
             # system to make sure that the isolated process can
             # import stuff from the primary environment. Without this
             # the isolated process will not be able to run properly
@@ -158,4 +186,5 @@ def main() -> int:
 
 
 if __name__ == "__main__":
+    load_pth_files()
     sys.exit(main())

src/isolate/backends/connections/ipc/agent.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, ClassVar, Dict
+
+from isolate.backends import BaseEnvironment
+from isolate.backends.common import sha256_digest_of
+from isolate.backends.connections import PythonIPC
+from isolate.backends.context import GLOBAL_CONTEXT, ContextType
+
+
+@dataclass
+class LocalPythonEnvironment(BaseEnvironment[Path]):
+    BACKEND_NAME: ClassVar[str] = "local"
+
+    @classmethod
+    def from_config(
+        cls,
+        config: Dict[str, Any],
+        context: ContextType = GLOBAL_CONTEXT,
+    ) -> BaseEnvironment:
+        environment = cls()
+        environment.set_context(context)
+        return environment
+
+    @property
+    def key(self) -> str:
+        return sha256_digest_of(sys.exec_prefix)
+
+    def create(self) -> Path:
+        return Path(sys.exec_prefix)
+
+    def destroy(self, connection_key: Path) -> None:
+        raise NotImplementedError("LocalPythonEnvironment cannot be destroyed")
+
+    def exists(self) -> bool:
+        return True
+
+    def open_connection(self, connection_key: Path) -> PythonIPC:
+        return PythonIPC(self, connection_key)