isolate 0.22.0__py3-none-any.whl

isolate/connections/grpc/definitions/common.proto

@@ -0,0 +1,49 @@
+syntax = "proto3";
+
+import "google/protobuf/timestamp.proto";
+
+message SerializedObject {
+    // The serialization method used to serialize the the raw_object. Must be
+    // present in the environment that is running the agent itself.
+    string method = 1;
+    // The Python object serialized with the method above.
+    bytes definition = 2;
+    // A flag indicating whether the given object was raised (e.g. an exception
+    // that was captured) or not.
+    bool was_it_raised = 3;
+    // The stringized version of the traceback, if it was raised.
+    optional string stringized_traceback = 4;
+}
+
+message PartialRunResult {
+    // A flag indicating whether the run has completed.
+    bool is_complete = 1;
+    // A list of logs collected during this partial execution. It does
+    // not include old logs.
+    repeated Log logs = 2;
+    // The result of the run, if it is complete.
+    optional SerializedObject result = 3;
+}
+
+message Log {
+    string message = 1;
+    LogSource source = 2;
+    LogLevel level = 3;
+    google.protobuf.Timestamp timestamp = 4;
+}
+
+enum LogSource {
+    BUILDER = 0;
+    BRIDGE = 1;
+    USER = 2;
+}
+
+enum LogLevel {
+    TRACE = 0;
+    DEBUG = 1;
+    INFO = 2;
+    WARNING = 3;
+    ERROR = 4;
+    STDOUT = 5;
+    STDERR = 6;
+}

isolate/connections/grpc/definitions/common_pb2.py

@@ -0,0 +1,35 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: common.proto
+# Protobuf Python Version: 4.25.1
+"""Generated protocol buffer code."""
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf.internal import builder as _builder
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+from google.protobuf import timestamp_pb2 as google_dot_protobuf_dot_timestamp__pb2
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x0c\x63ommon.proto\x1a\x1fgoogle/protobuf/timestamp.proto\"\x89\x01\n\x10SerializedObject\x12\x0e\n\x06method\x18\x01 \x01(\t\x12\x12\n\ndefinition\x18\x02 \x01(\x0c\x12\x15\n\rwas_it_raised\x18\x03 \x01(\x08\x12!\n\x14stringized_traceback\x18\x04 \x01(\tH\x00\x88\x01\x01\x42\x17\n\x15_stringized_traceback\"n\n\x10PartialRunResult\x12\x13\n\x0bis_complete\x18\x01 \x01(\x08\x12\x12\n\x04logs\x18\x02 \x03(\x0b\x32\x04.Log\x12&\n\x06result\x18\x03 \x01(\x0b\x32\x11.SerializedObjectH\x00\x88\x01\x01\x42\t\n\x07_result\"{\n\x03Log\x12\x0f\n\x07message\x18\x01 \x01(\t\x12\x1a\n\x06source\x18\x02 \x01(\x0e\x32\n.LogSource\x12\x18\n\x05level\x18\x03 \x01(\x0e\x32\t.LogLevel\x12-\n\ttimestamp\x18\x04 \x01(\x0b\x32\x1a.google.protobuf.Timestamp*.\n\tLogSource\x12\x0b\n\x07\x42UILDER\x10\x00\x12\n\n\x06\x42RIDGE\x10\x01\x12\x08\n\x04USER\x10\x02*Z\n\x08LogLevel\x12\t\n\x05TRACE\x10\x00\x12\t\n\x05\x44\x45\x42UG\x10\x01\x12\x08\n\x04INFO\x10\x02\x12\x0b\n\x07WARNING\x10\x03\x12\t\n\x05\x45RROR\x10\x04\x12\n\n\x06STDOUT\x10\x05\x12\n\n\x06STDERR\x10\x06\x62\x06proto3')
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'common_pb2', _globals)
+if _descriptor._USE_C_DESCRIPTORS == False:
+  DESCRIPTOR._options = None
+  _globals['_LOGSOURCE']._serialized_start=426
+  _globals['_LOGSOURCE']._serialized_end=472
+  _globals['_LOGLEVEL']._serialized_start=474
+  _globals['_LOGLEVEL']._serialized_end=564
+  _globals['_SERIALIZEDOBJECT']._serialized_start=50
+  _globals['_SERIALIZEDOBJECT']._serialized_end=187
+  _globals['_PARTIALRUNRESULT']._serialized_start=189
+  _globals['_PARTIALRUNRESULT']._serialized_end=299
+  _globals['_LOG']._serialized_start=301
+  _globals['_LOG']._serialized_end=424
+# @@protoc_insertion_point(module_scope)

isolate/connections/grpc/definitions/common_pb2.pyi

@@ -0,0 +1,152 @@
+"""
+@generated by mypy-protobuf.  Do not edit manually!
+isort:skip_file
+"""
+import builtins
+import collections.abc
+import google.protobuf.descriptor
+import google.protobuf.internal.containers
+import google.protobuf.internal.enum_type_wrapper
+import google.protobuf.message
+import google.protobuf.timestamp_pb2
+import sys
+import typing
+
+if sys.version_info >= (3, 10):
+    import typing as typing_extensions
+else:
+    import typing_extensions
+
+DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
+
+class _LogSource:
+    ValueType = typing.NewType("ValueType", builtins.int)
+    V: typing_extensions.TypeAlias = ValueType
+
+class _LogSourceEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[_LogSource.ValueType], builtins.type):
+    DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor
+    BUILDER: _LogSource.ValueType  # 0
+    BRIDGE: _LogSource.ValueType  # 1
+    USER: _LogSource.ValueType  # 2
+
+class LogSource(_LogSource, metaclass=_LogSourceEnumTypeWrapper): ...
+
+BUILDER: LogSource.ValueType  # 0
+BRIDGE: LogSource.ValueType  # 1
+USER: LogSource.ValueType  # 2
+global___LogSource = LogSource
+
+class _LogLevel:
+    ValueType = typing.NewType("ValueType", builtins.int)
+    V: typing_extensions.TypeAlias = ValueType
+
+class _LogLevelEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[_LogLevel.ValueType], builtins.type):
+    DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor
+    TRACE: _LogLevel.ValueType  # 0
+    DEBUG: _LogLevel.ValueType  # 1
+    INFO: _LogLevel.ValueType  # 2
+    WARNING: _LogLevel.ValueType  # 3
+    ERROR: _LogLevel.ValueType  # 4
+    STDOUT: _LogLevel.ValueType  # 5
+    STDERR: _LogLevel.ValueType  # 6
+
+class LogLevel(_LogLevel, metaclass=_LogLevelEnumTypeWrapper): ...
+
+TRACE: LogLevel.ValueType  # 0
+DEBUG: LogLevel.ValueType  # 1
+INFO: LogLevel.ValueType  # 2
+WARNING: LogLevel.ValueType  # 3
+ERROR: LogLevel.ValueType  # 4
+STDOUT: LogLevel.ValueType  # 5
+STDERR: LogLevel.ValueType  # 6
+global___LogLevel = LogLevel
+
+@typing_extensions.final
+class SerializedObject(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    METHOD_FIELD_NUMBER: builtins.int
+    DEFINITION_FIELD_NUMBER: builtins.int
+    WAS_IT_RAISED_FIELD_NUMBER: builtins.int
+    STRINGIZED_TRACEBACK_FIELD_NUMBER: builtins.int
+    method: builtins.str
+    """The serialization method used to serialize the the raw_object. Must be
+    present in the environment that is running the agent itself.
+    """
+    definition: builtins.bytes
+    """The Python object serialized with the method above."""
+    was_it_raised: builtins.bool
+    """A flag indicating whether the given object was raised (e.g. an exception
+    that was captured) or not.
+    """
+    stringized_traceback: builtins.str
+    """The stringized version of the traceback, if it was raised."""
+    def __init__(
+        self,
+        *,
+        method: builtins.str = ...,
+        definition: builtins.bytes = ...,
+        was_it_raised: builtins.bool = ...,
+        stringized_traceback: builtins.str | None = ...,
+    ) -> None: ...
+    def HasField(self, field_name: typing_extensions.Literal["_stringized_traceback", b"_stringized_traceback", "stringized_traceback", b"stringized_traceback"]) -> builtins.bool: ...
+    def ClearField(self, field_name: typing_extensions.Literal["_stringized_traceback", b"_stringized_traceback", "definition", b"definition", "method", b"method", "stringized_traceback", b"stringized_traceback", "was_it_raised", b"was_it_raised"]) -> None: ...
+    def WhichOneof(self, oneof_group: typing_extensions.Literal["_stringized_traceback", b"_stringized_traceback"]) -> typing_extensions.Literal["stringized_traceback"] | None: ...
+
+global___SerializedObject = SerializedObject
+
+@typing_extensions.final
+class PartialRunResult(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    IS_COMPLETE_FIELD_NUMBER: builtins.int
+    LOGS_FIELD_NUMBER: builtins.int
+    RESULT_FIELD_NUMBER: builtins.int
+    is_complete: builtins.bool
+    """A flag indicating whether the run has completed."""
+    @property
+    def logs(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___Log]:
+        """A list of logs collected during this partial execution. It does
+        not include old logs.
+        """
+    @property
+    def result(self) -> global___SerializedObject:
+        """The result of the run, if it is complete."""
+    def __init__(
+        self,
+        *,
+        is_complete: builtins.bool = ...,
+        logs: collections.abc.Iterable[global___Log] | None = ...,
+        result: global___SerializedObject | None = ...,
+    ) -> None: ...
+    def HasField(self, field_name: typing_extensions.Literal["_result", b"_result", "result", b"result"]) -> builtins.bool: ...
+    def ClearField(self, field_name: typing_extensions.Literal["_result", b"_result", "is_complete", b"is_complete", "logs", b"logs", "result", b"result"]) -> None: ...
+    def WhichOneof(self, oneof_group: typing_extensions.Literal["_result", b"_result"]) -> typing_extensions.Literal["result"] | None: ...
+
+global___PartialRunResult = PartialRunResult
+
+@typing_extensions.final
+class Log(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    MESSAGE_FIELD_NUMBER: builtins.int
+    SOURCE_FIELD_NUMBER: builtins.int
+    LEVEL_FIELD_NUMBER: builtins.int
+    TIMESTAMP_FIELD_NUMBER: builtins.int
+    message: builtins.str
+    source: global___LogSource.ValueType
+    level: global___LogLevel.ValueType
+    @property
+    def timestamp(self) -> google.protobuf.timestamp_pb2.Timestamp: ...
+    def __init__(
+        self,
+        *,
+        message: builtins.str = ...,
+        source: global___LogSource.ValueType = ...,
+        level: global___LogLevel.ValueType = ...,
+        timestamp: google.protobuf.timestamp_pb2.Timestamp | None = ...,
+    ) -> None: ...
+    def HasField(self, field_name: typing_extensions.Literal["timestamp", b"timestamp"]) -> builtins.bool: ...
+    def ClearField(self, field_name: typing_extensions.Literal["level", b"level", "message", b"message", "source", b"source", "timestamp", b"timestamp"]) -> None: ...
+
+global___Log = Log

isolate/connections/grpc/definitions/common_pb2_grpc.py

@@ -0,0 +1,4 @@
+# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!
+"""Client and server classes corresponding to protobuf-defined services."""
+import grpc
+

isolate/connections/grpc/interface.py

@@ -0,0 +1,71 @@
+"""A common gRPC interface for both the gRPC connection implementation
+and the Isolate Server to share."""
+
+import functools
+from typing import Any, Optional
+
+from isolate.common import timestamp
+from isolate.connections.common import load_serialized_object, serialize_object
+from isolate.connections.grpc import definitions
+from isolate.logs import Log, LogLevel, LogSource
+
+
+@functools.singledispatch
+def from_grpc(message: definitions.Message) -> Any:
+    """Materialize a gRPC message into a Python object."""
+    wrong_type = type(message).__name__
+    raise NotImplementedError(f"Can't convert {wrong_type} to a Python object.")
+
+
+@functools.singledispatch
+def to_grpc(obj: Any) -> definitions.Message:
+    """Convert a Python object into a gRPC message."""
+    wrong_type = type(obj).__name__
+    raise NotImplementedError(f"Cannot convert {wrong_type} to a gRPC message.")
+
+
+@from_grpc.register
+def _(message: definitions.SerializedObject) -> Any:
+    return load_serialized_object(
+        message.method,
+        message.definition,
+        was_it_raised=message.was_it_raised,
+        stringized_traceback=message.stringized_traceback,
+    )
+
+
+@from_grpc.register
+def _(message: definitions.Log) -> Log:
+    source = LogSource(definitions.LogSource.Name(message.source).lower())
+    level = LogLevel[definitions.LogLevel.Name(message.level).upper()]
+    return Log(
+        message=message.message,
+        source=source,
+        level=level,
+        timestamp=timestamp.to_datetime(message.timestamp),
+    )
+
+
+@to_grpc.register
+def _(obj: Log) -> definitions.Log:
+    return definitions.Log(
+        message=obj.message,
+        source=definitions.LogSource.Value(obj.source.name.upper()),
+        level=definitions.LogLevel.Value(obj.level.name.upper()),
+        timestamp=timestamp.from_datetime(obj.timestamp),
+    )
+
+
+def to_serialized_object(
+    obj: Any,
+    method: str,
+    was_it_raised: bool = False,
+    stringized_traceback: Optional[str] = None,
+) -> definitions.SerializedObject:
+    """Convert a Python object into a gRPC message."""
+    return definitions.SerializedObject(
+        method=method,
+        definition=serialize_object(method, obj),
+        was_it_raised=was_it_raised,
+        stringized_traceback=stringized_traceback,
+    )

isolate/connections/ipc/__init__.py

@@ -0,0 +1,5 @@
+from isolate.connections.ipc._base import (  # noqa: F401
+    IsolatedProcessConnection,
+    PythonExecutionBase,
+    PythonIPC,
+)

isolate/connections/ipc/_base.py

@@ -0,0 +1,225 @@
+from __future__ import annotations
+
+import base64
+import importlib
+import subprocess
+import time
+from contextlib import ExitStack, closing
+from dataclasses import dataclass
+from multiprocessing.connection import Connection, Listener
+from pathlib import Path
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    ContextManager,
+)
+
+from isolate.backends import (
+    BasicCallable,
+    CallResultType,
+    EnvironmentConnection,
+)
+from isolate.connections._local import PythonExecutionBase, agent_startup
+from isolate.connections.common import prepare_exc
+from isolate.connections.ipc import agent
+from isolate.logs import LogLevel, LogSource
+
+if TYPE_CHECKING:
+    # Somhow mypy can't figure out that `ConnectionWrapper`
+    # really exists.
+    class ConnectionWrapper(Connection):
+        def __init__(
+            self,
+            connection: Any,
+            loads: Callable[[bytes], Any],
+            dumps: Callable[[Any], bytes],
+        ) -> None: ...
+
+        def recv(self) -> Any: ...
+
+        def send(self, value: Any) -> None: ...
+
+        def close(self) -> None: ...
+
+else:
+    from multiprocessing.connection import ConnectionWrapper
+
+
+class AgentListener(Listener):
+    """A custom listener that can use any available serialization method
+    to communicate with the child process."""
+
+    def __init__(self, backend_name: str, *args: Any, **kwargs: Any) -> None:
+        self.serialization_backend = loadserialization_method(backend_name)
+        super().__init__(*args, **kwargs)
+
+    def accept(self) -> Connection:
+        return ConnectionWrapper(
+            super().accept(),
+            dumps=self.serialization_backend.dumps,
+            loads=self.serialization_backend.loads,
+        )
+
+
+def loadserialization_method(backend_name: str) -> Any:
+    # TODO(feat): This should probably throw a better error if the
+    # given backend does not exist.
+    return importlib.import_module(backend_name)
+
+
+def encode_service_address(address: tuple[str, int]) -> str:
+    host, port = address
+    return base64.b64encode(f"{host}:{port}".encode()).decode("utf-8")
+
+
+@dataclass
+class IsolatedProcessConnection(EnvironmentConnection):
+    """A generic IPC implementation for running the isolate backend
+    in a separated process.
+
+    Each implementation needs to define a start_process method to
+    spawn the agent."""
+
+    # The amount of seconds to wait before checking whether the
+    # isolated process has exited or not.
+    _DEFER_THRESHOLD = 0.25
+
+    def start_process(
+        self,
+        connection: AgentListener,
+        *args: Any,
+        **kwargs: Any,
+    ) -> ContextManager[subprocess.Popen]:
+        """Start the agent process."""
+        raise NotImplementedError
+
+    def run(  # type: ignore[return-value]
+        self,
+        executable: BasicCallable,
+        *args: Any,
+        **kwargs: Any,
+    ) -> CallResultType:  # type: ignore[type-var]
+        """Spawn an agent process using the given environment, run the given
+        `executable` in that process, and return the result object back."""
+
+        with ExitStack() as stack:
+            # IPC flow is the following:
+            #  1. [controller]: Create the socket server
+            #  2. [controller]: Spawn the call agent with the socket address
+            #  3.      [agent]: Connect to the socket server
+            #  4. [controller]: Accept the incoming connection request
+            #  5. [controller]: Send the executable over the established bridge
+            #  6.      [agent]: Receive the executable from the bridge
+            #  7.      [agent]: Execute the executable and once done send the result
+            #                   back
+            #  8. [controller]: Loop until either the isolated process exits or sends
+            #                   any data (will be interpreted as a tuple of two
+            #                   mutually exclusive objects, either a result object or
+            #                   an exception to be raised).
+            #
+
+            self.log("Starting the controller bridge.")
+            controller_service = stack.enter_context(
+                AgentListener(
+                    self.environment.settings.serialization_method,
+                    family="AF_INET",
+                )
+            )
+
+            self.log(
+                f"Controller server is listening at {controller_service.address}."
+                " Attempting to start the agent process."
+            )
+            assert not (args or kwargs), "run() should not receive any arguments."
+            isolated_process = stack.enter_context(
+                self.start_process(controller_service, *args, **kwargs)
+            )
+
+            # TODO(fix): this might hang if the agent process crashes before it can
+            # connect to the controller bridge.
+            self.log(
+                f"Awaiting agent process of {isolated_process.pid}"
+                " to establish a connection."
+            )
+            established_connection = stack.enter_context(
+                closing(controller_service.accept())
+            )
+
+            self.log("Bridge between controller and the agent has been established.")
+            established_connection.send(executable)
+
+            self.log("Executable has been sent, awaiting execution result.")
+            return self.poll_until_result(
+                isolated_process,
+                established_connection,
+            )
+
+    def poll_until_result(
+        self,
+        process: subprocess.Popen,
+        connection: Connection,
+    ) -> CallResultType:  # type: ignore[type-var]
+        """Take the given process, and poll until either it exits or returns
+        a result object."""
+
+        while not connection.poll():
+            # Normally, if we do connection.read() without having this loop
+            # it is going to block us indefinitely (even if the underlying
+            # process has crashed). We can use a combination of process.poll
+            # and connection.poll to check if the process is alive and has data
+            # to move forward.
+            if process.poll():
+                break
+
+            # For preventing busy waiting, we can sleep for a bit
+            # and let other threads run.
+            time.sleep(self._DEFER_THRESHOLD)
+            continue
+
+        if not connection.poll():
+            # If the process has exited but there is still no data, we
+            # can assume something terrible has happened.
+            raise OSError(
+                "The isolated process has exited unexpectedly with code "
+                f"'{process.poll()}' without sending any data back."
+            )
+
+        # TODO(fix): handle EOFError that might happen here (e.g. problematic
+        # serialization might cause it).
+        result, did_it_raise, stringized_traceback = connection.recv()
+
+        if did_it_raise:
+            raise prepare_exc(result, stringized_traceback=stringized_traceback)
+        else:
+            assert stringized_traceback is None
+            return result
+
+
+@dataclass
+class PythonIPC(PythonExecutionBase[AgentListener], IsolatedProcessConnection):
+    def get_python_cmd(
+        self,
+        executable: Path,
+        connection: AgentListener,
+        log_fd: int,
+    ) -> list[str | Path]:
+        assert isinstance(connection.address, tuple)
+        return [
+            executable,
+            agent_startup.__file__,
+            agent.__file__,
+            encode_service_address(connection.address),
+            # TODO(feat): we probably should check if the given backend is installed
+            # on the remote interpreter, otherwise it will fail without establishing
+            # the connection with the bridge.
+            "--serialization-backend",
+            self.environment.settings.serialization_method,
+            "--log-fd",
+            str(log_fd),
+        ]
+
+    def handle_agent_log(
+        self, line: str, *, level: LogLevel, source: LogSource
+    ) -> None:
+        self.log(line, level=level, source=source)