isolate 0.22.0__py3-none-any.whl

isolate/backends/virtualenv.py

@@ -0,0 +1,204 @@
+from __future__ import annotations
+
+import os
+import shutil
+import subprocess
+from dataclasses import dataclass, field
+from functools import partial
+from pathlib import Path
+from typing import Any, ClassVar
+
+from isolate.backends import BaseEnvironment, EnvironmentCreationError
+from isolate.backends.common import (
+    active_python,
+    get_executable,
+    get_executable_path,
+    logged_io,
+    optional_import,
+    sha256_digest_of,
+)
+from isolate.backends.settings import DEFAULT_SETTINGS, IsolateSettings
+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]):
+    BACKEND_NAME: ClassVar[str] = "virtualenv"
+
+    requirements: list[str] = field(default_factory=list)
+    constraints_file: os.PathLike | None = None
+    python_version: str | None = None
+    extra_index_urls: list[str] = field(default_factory=list)
+    tags: list[str] = field(default_factory=list)
+    resolver: str | None = None
+
+    @classmethod
+    def from_config(
+        cls,
+        config: dict[str, Any],
+        settings: IsolateSettings = DEFAULT_SETTINGS,
+    ) -> 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
+    def key(self) -> str:
+        if self.constraints_file is not None:
+            with open(self.constraints_file) as stream:
+                constraints = stream.read().splitlines()
+        else:
+            constraints = []
+
+        extras = []
+        if self.resolver is not None:
+            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:
+        """Install the requirements of this environment using 'pip' to the
+        given virtualenv path.
+
+        If there are any constraint files specified, they will be also passed to
+        the package resolver.
+        """
+        if not self.requirements:
+            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[str | os.PathLike] = [
+            *base_pip_cmd,  # type: ignore
+            "install",
+            *self.requirements,
+        ]
+        if self.constraints_file:
+            pip_cmd.extend(["-c", self.constraints_file])
+
+        for extra_index_url in self.extra_index_urls:
+            pip_cmd.extend(["--extra-index-url", extra_index_url])
+
+        with logged_io(partial(self.log, level=LogLevel.INFO)) as (stdout, stderr, _):
+            try:
+                subprocess.check_call(
+                    pip_cmd,
+                    stdout=stdout,
+                    stderr=stderr,
+                    env=environ,
+                )
+            except subprocess.SubprocessError as exc:
+                raise EnvironmentCreationError(f"Failure during 'pip install': {exc}")
+
+    def _install_python_through_pyenv(self) -> str:
+        from isolate.backends.pyenv import PyenvEnvironment
+
+        self.log(
+            f"Requested Python version of {self.python_version} is not available "
+            "in the system, attempting to install it through pyenv."
+        )
+
+        pyenv = PyenvEnvironment.from_config(
+            {"python_version": self.python_version},
+            settings=self.settings,
+        )
+        return str(get_executable_path(pyenv.create(), "python"))
+
+    def _decide_python(self) -> str:
+        from isolate.backends.pyenv import _get_pyenv_executable
+
+        builtin_discovery = optional_import("virtualenv.discovery.builtin")
+        interpreter = builtin_discovery.get_interpreter(self.python_version, ())
+        if interpreter is not None:
+            return interpreter.executable
+
+        try:
+            _get_pyenv_executable()
+        except Exception:
+            raise EnvironmentCreationError(
+                f"Python {self.python_version} is not available in your "
+                "system. Please install it first."
+            ) from None
+        else:
+            return self._install_python_through_pyenv()
+
+    def create(self, *, force: bool = False) -> Path:
+        virtualenv = optional_import("virtualenv")
+
+        venv_path = self.settings.cache_dir_for(self)
+        completion_marker = self.settings.completion_marker_for(venv_path)
+        with self.settings.cache_lock_for(venv_path):
+            if not force:
+                is_cached = venv_path.exists()
+                if self.settings.strict_cache:
+                    is_cached &= completion_marker.exists()
+
+                if is_cached:
+                    return venv_path
+
+            self.log(f"Creating the environment at '{venv_path}'")
+
+            args = [str(venv_path)]
+            if self.python_version:
+                args.append(f"--python={self._decide_python()}")
+
+            try:
+                # This is not an official API, so it can throw anything at us.
+                virtualenv.cli_run(args)
+            except (SystemExit, RuntimeError, OSError) as exc:
+                raise EnvironmentCreationError(
+                    f"Failed to create the environment at '{venv_path}': {exc}"
+                )
+
+            self.install_requirements(venv_path)
+            completion_marker.touch()
+
+        self.log(f"New environment cached at '{venv_path}'")
+        return venv_path
+
+    def destroy(self, connection_key: Path) -> None:
+        with self.settings.cache_lock_for(connection_key):
+            # It might be destroyed already (when we are awaiting
+            # for the lock to be released).
+            if not connection_key.exists():
+                return
+
+            shutil.rmtree(connection_key)
+
+    def exists(self) -> bool:
+        path = self.settings.cache_dir_for(self)
+        return path.exists()
+
+    def open_connection(self, connection_key: Path) -> PythonIPC:
+        return PythonIPC(self, connection_key)

isolate/common/timestamp.py

@@ -0,0 +1,15 @@
+from __future__ import annotations
+
+from datetime import datetime, timezone
+
+from google.protobuf.timestamp_pb2 import Timestamp
+
+
+def from_datetime(time: datetime) -> Timestamp:
+    timestamp = Timestamp()
+    timestamp.FromDatetime(time)
+    return timestamp
+
+
+def to_datetime(timestamp: Timestamp) -> datetime:
+    return timestamp.ToDatetime(tzinfo=timezone.utc)

isolate/connections/__init__.py

@@ -0,0 +1,21 @@
+import importlib
+
+from isolate.connections.ipc import IsolatedProcessConnection, PythonIPC  # noqa: F401
+
+
+def __getattr__(name):
+    if name == "LocalPythonGRPC":
+        extra = "grpc"
+        module_name = "isolate.connections.grpc"
+    else:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+    try:
+        module = importlib.import_module(module_name)
+    except ImportError:
+        raise AttributeError(
+            f"For using {name!r} you need to install isolate with {extra!r} support."
+            f'\n    $ pip install "isolate[{extra}]"'
+        )
+
+    return getattr(module, name)

isolate/connections/_local/__init__.py

@@ -0,0 +1,2 @@
+from isolate.connections._local import agent_startup  # noqa: F401
+from isolate.connections._local._base import PythonExecutionBase  # noqa: F401

isolate/connections/_local/_base.py

@@ -0,0 +1,190 @@
+from __future__ import annotations
+
+import os
+import subprocess
+import sysconfig
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from functools import partial
+from pathlib import Path
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Generic,
+    Iterator,
+    TypeVar,
+)
+
+from isolate import __version__ as isolate_version
+from isolate.backends.common import active_python, get_executable_path, logged_io
+from isolate.connections.common import AGENT_SIGNATURE
+from isolate.logs import LogLevel, LogSource
+
+if TYPE_CHECKING:
+    from isolate.backends import BaseEnvironment
+
+ConnectionType = TypeVar("ConnectionType")
+
+
+def binary_path_for(*search_paths: Path) -> str:
+    """Return the binary search path for the given 'search_paths'.
+    It will be a combination of the 'bin/' folders in them and
+    the existing PATH environment variable."""
+
+    paths = []
+    for search_path in search_paths:
+        path = sysconfig.get_path("scripts", vars={"base": search_path})
+        paths.append(path)
+        # Some distributions (conda) might include both a 'bin' and
+        # a 'scripts' folder.
+
+        auxilary_binary_path = search_path / "bin"
+        if path != auxilary_binary_path and auxilary_binary_path.exists():
+            paths.append(str(auxilary_binary_path))
+
+    if "PATH" in os.environ:
+        paths.append(os.environ["PATH"])
+
+    return os.pathsep.join(paths)
+
+
+def python_path_for(*search_paths: Path) -> str:
+    """Return the PYTHONPATH for the library paths residing
+    in the given 'search_paths'. The order of the paths is
+    preserved."""
+    assert len(search_paths) >= 1
+    lib_paths = []
+    for search_path in search_paths:
+        # sysconfig defines the schema of the directories under
+        # any comforming Python installation (like venv, conda, etc.).
+        #
+        # Be aware that Debian's system installation does not
+        # comform sysconfig.
+        raw_glob_expr = sysconfig.get_path(
+            "purelib",
+            vars={
+                "base": search_path,
+                "python_version": "*",
+                "py_version_short": "*",
+                "py_version_nodot": "*",
+            },
+        )
+        relative_glob_expr = Path(raw_glob_expr).relative_to(search_path).as_posix()
+
+        # Try to find expand the Python version in the path. This is
+        # necessary for supporting multiple Python versions in the same
+        # environment.
+        for file in search_path.glob(relative_glob_expr):
+            lib_paths.append(str(file))
+
+    return os.pathsep.join(lib_paths)
+
+
+@dataclass
+class PythonExecutionBase(Generic[ConnectionType]):
+    """A generic Python execution implementation that can trigger a new process
+    and start watching stdout/stderr for the logs. The environment_path must be
+    the base directory of a new Python environment (structure that complies with
+    sysconfig). Python binary inside that environment will be used to run the
+    agent process.
+
+    If set, extra_inheritance_paths allows extending the custom package search
+    system with additional environments. As an example, the current environment_path
+    might point to an environment with numpy and the extra_inheritance_paths might
+    point to an environment with pandas. In this case, the agent process will have
+    access to both numpy and pandas. The order is important here, as the first
+    path in the list will be the first one to be looked up (so if there is multiple
+    versions of the same package in different environments, the one in the first
+    path will take precedence). Dependency resolution and compatibility must be
+    handled by the user."""
+
+    environment: BaseEnvironment
+    environment_path: Path
+    extra_inheritance_paths: list[Path] = field(default_factory=list)
+
+    @contextmanager
+    def start_process(
+        self,
+        connection: ConnectionType,
+        *args: Any,
+        **kwargs: Any,
+    ) -> Iterator[subprocess.Popen]:
+        """Start the agent process with the Python binary available inside the
+        bound environment."""
+
+        python_version = getattr(self.environment, "python_version", active_python())
+        try:
+            # prefer the specific version if available.
+            python_executable = get_executable_path(
+                self.environment_path, f"python{python_version}"
+            )
+        except FileNotFoundError:
+            # fallback to the generic binary.
+            python_executable = get_executable_path(self.environment_path, "python")
+
+        with logged_io(
+            partial(
+                self.handle_agent_log, source=LogSource.USER, level=LogLevel.STDOUT
+            ),
+            partial(
+                self.handle_agent_log, source=LogSource.USER, level=LogLevel.STDERR
+            ),
+            partial(
+                self.handle_agent_log, source=LogSource.BRIDGE, level=LogLevel.TRACE
+            ),
+        ) as (stdout, stderr, log_fd):
+            yield subprocess.Popen(
+                self.get_python_cmd(python_executable, connection, log_fd),
+                env=self.get_env_vars(),
+                stdout=stdout,
+                stderr=stderr,
+                pass_fds=(log_fd,),
+                text=True,
+            )
+
+    def get_env_vars(self) -> dict[str, str]:
+        """Return the environment variables to run the agent process with. By default
+        PYTHONUNBUFFERED is set to 1 to ensure the prints to stdout/stderr are reflect
+        immediately (so that we can seamlessly transfer logs)."""
+
+        custom_vars = {}
+        custom_vars["ISOLATE_SERVER_VERSION"] = isolate_version
+        custom_vars[AGENT_SIGNATURE] = "1"
+        custom_vars["PYTHONUNBUFFERED"] = "1"
+
+        # NOTE: we don't have to manually set PYTHONPATH here if we are
+        # using a single environment since python will automatically
+        # use the proper path.
+        if self.extra_inheritance_paths:
+            # The order here should reflect the order of the inheritance
+            # where the actual environment already takes precedence.
+            custom_vars["PYTHONPATH"] = python_path_for(
+                self.environment_path, *self.extra_inheritance_paths
+            )
+
+        # But the PATH must be always set since it will be not be
+        # automatically set by Python (think of this as ./venv/bin/activate)
+        custom_vars["PATH"] = binary_path_for(
+            self.environment_path, *self.extra_inheritance_paths
+        )
+
+        return {
+            **os.environ,
+            **custom_vars,
+        }
+
+    def get_python_cmd(
+        self,
+        executable: Path,
+        connection: ConnectionType,
+        log_fd: int,
+    ) -> list[str | Path]:
+        """Return the command to run the agent process with."""
+        raise NotImplementedError
+
+    def handle_agent_log(
+        self, line: str, *, level: LogLevel, source: LogSource
+    ) -> None:
+        """Handle a log line emitted by the agent process. The level will be either
+        STDOUT or STDERR."""
+        raise NotImplementedError

isolate/connections/_local/agent_startup.py

@@ -0,0 +1,53 @@
+"""
+Agent process execution wrapper for handling extended PYTHONPATH.
+"""
+
+import os
+import runpy
+import site
+import sys
+import traceback
+
+
+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):
+        try:
+            site.addsitedir(site_dir)
+        except Exception:
+            # NOTE: there could be .pth files that are model weights and not
+            # python path configuration files.
+            traceback.print_exc()
+            print(f"Error adding site directory {site_dir}, skipping...")
+
+
+def main():
+    real_agent, *real_arguments = sys.argv[1:]
+
+    load_pth_files()
+    # TODO(feat): implement a check to parse "agent-requires" line and see if
+    # all the dependencies are installed.
+    sys.argv = [real_agent] + real_arguments
+    runpy.run_path(real_agent, run_name="__main__")
+
+
+if __name__ == "__main__":
+    load_pth_files()
+    main()

isolate/connections/common.py

@@ -0,0 +1,121 @@
+from __future__ import annotations
+
+import importlib
+import os
+from contextlib import contextmanager
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Iterator, cast
+
+from tblib import Traceback, TracebackParseError
+from tblib.pickling_support import install as tblib_install
+
+if TYPE_CHECKING:
+    from typing import Protocol
+
+    class SerializationBackend(Protocol):
+        def loads(self, data: bytes) -> Any: ...
+
+        def dumps(self, obj: Any) -> bytes: ...
+
+
+AGENT_SIGNATURE = "IS_ISOLATE_AGENT"
+
+
+@dataclass
+class SerializationError(Exception):
+    """An error that happened during the serialization process."""
+
+    message: str
+
+
+# NOTE: tblib's install() will search for BaseException subclasses,
+# so we have to call it after the SerializationError is defined.
+tblib_install()
+
+
+@contextmanager
+def _step(message: str) -> Iterator[None]:
+    """A context manager to capture every expression
+    underneath it and if any of them fails for any reason
+    then it will raise a SerializationError with the
+    given message."""
+
+    try:
+        yield
+    except BaseException as exception:
+        raise SerializationError("Error while " + message) from exception
+
+
+def as_serialization_method(backend: Any) -> SerializationBackend:
+    """Ensures that the given backend has loads/dumps methods, and returns
+    it as is (also convinces type checkers that the given object satisfies
+    the serialization protocol)."""
+
+    if not hasattr(backend, "loads") or not hasattr(backend, "dumps"):
+        raise TypeError(
+            f"The given serialization backend ({backend.__name__}) does "
+            "not have one of the required methods (loads/dumps)."
+        )
+
+    return cast("SerializationBackend", backend)
+
+
+def load_serialized_object(
+    serialization_method: str,
+    raw_object: bytes,
+    *,
+    was_it_raised: bool = False,
+    stringized_traceback: str | None = None,
+) -> Any:
+    """Load the given serialized object using the given serialization method. If
+    anything fails, then a SerializationError will be raised. If the was_it_raised
+    flag is set to true, then the given object will be raised as an exception (instead
+    of being returned)."""
+
+    with _step(f"preparing the serialization backend ({serialization_method})"):
+        serialization_backend = as_serialization_method(
+            importlib.import_module(serialization_method)
+        )
+
+    with _step("deserializing the given object"):
+        result = serialization_backend.loads(raw_object)
+
+    if was_it_raised:
+        raise prepare_exc(result, stringized_traceback=stringized_traceback)
+    else:
+        return result
+
+
+def serialize_object(serialization_method: str, object: Any) -> bytes:
+    """Serialize the given object using the given serialization method. If
+    anything fails, then a SerializationError will be raised."""
+
+    with _step(f"preparing the serialization backend ({serialization_method})"):
+        serialization_backend = as_serialization_method(
+            importlib.import_module(serialization_method)
+        )
+
+    with _step("serializing the given object"):
+        return serialization_backend.dumps(object)
+
+
+def is_agent() -> bool:
+    """Returns true if the current process is an isolate agent."""
+    return os.environ.get(AGENT_SIGNATURE) == "1"
+
+
+def prepare_exc(
+    exc: BaseException,
+    *,
+    stringized_traceback: str | None = None,
+) -> BaseException:
+    if stringized_traceback:
+        try:
+            traceback = Traceback.from_string(stringized_traceback).as_traceback()
+        except TracebackParseError:
+            traceback = None
+    else:
+        traceback = None
+
+    exc.__traceback__ = traceback
+    return exc

isolate/connections/grpc/__init__.py

@@ -0,0 +1 @@
+from isolate.connections.grpc._base import AgentError, LocalPythonGRPC  # noqa: F401