isolate 0.22.0__py3-none-any.whl

isolate/connections/ipc/agent.py

@@ -0,0 +1,205 @@
+# This file defines an "isolate" agent for inter-process communication over
+# sockets. It is spawned by the controller process with a single argument (a
+# base64 encoded server address) and expected to go through the following procedures:
+#   1. Decode the given address
+#   2. Create a connection to the transmission bridge using the address
+#   3. Receive a callable object from the bridge
+#   4. Execute the callable object
+#   5. Send the result back to the bridge
+#   6. Exit
+#
+# Up until to point 4, the agent process has no way of transmitting information
+# to the controller so it should use the stderr/stdout channels appropriately. After
+# the executable is received, the controller process will switch to the listening mode
+# and wait for agent to return something. The expected object is a tuple of two objects
+# 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.
+
+from __future__ import annotations
+
+import base64
+import importlib
+import os
+import sys
+import time
+import traceback
+from argparse import ArgumentParser
+from contextlib import closing
+from multiprocessing.connection import Client
+from typing import TYPE_CHECKING, Any, Callable, ContextManager
+
+if TYPE_CHECKING:
+    # Somhow mypy can't figure out that `ConnectionWrapper`
+    # really exists.
+    class ConnectionWrapper:
+        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
+
+
+def decode_service_address(address: str) -> tuple[str, int]:
+    host, port = base64.b64decode(address).decode("utf-8").rsplit(":", 1)
+    return host, int(port)
+
+
+def child_connection(
+    serialization_method: str, address: tuple[str, int]
+) -> ContextManager[ConnectionWrapper]:
+    serialization_backend = importlib.import_module(serialization_method)
+    return closing(
+        ConnectionWrapper(
+            Client(address),
+            loads=serialization_backend.loads,
+            dumps=serialization_backend.dumps,
+        )
+    )
+
+
+IS_DEBUG_MODE = os.getenv("ISOLATE_ENABLE_DEBUGGING") == "1"
+DEBUG_TIMEOUT = 60 * 15
+
+
+def run_client(
+    serialization_method: str,
+    address: tuple[str, int],
+    *,
+    with_pdb: bool = False,
+    log_fd: int | None = None,
+) -> None:
+    # Debug Mode
+    # ==========
+    #
+    # Isolated processes are really tricky to debug properly
+    # so we want to have a smooth way into the process and see
+    # what is really going on in the case of errors.
+    #
+    # For using the debug mode, you first need to set ISOLATE_ENABLE_DEBUGGING
+    # environment variable to "1" from your controller process. This will
+    # make the isolated process hang at the initialization, and make it print
+    # the instructions to connect to the controller process.
+    #
+    # On a separate shell (while letting the the controller process hang), you can
+    # execute the given command to drop into the PDB (Python Debugger). With that
+    # you can observe each step of the connection and run process.
+
+    if with_pdb:
+        # This condition will only be activated if we want to
+        # debug the isolated process by passing the --with-pdb
+        # flag when executing the binary.
+        import pdb
+
+        pdb.set_trace()
+
+    if log_fd is None:
+        _log = sys.stdout
+    else:
+        _log = os.fdopen(log_fd, "w")
+
+    def log(_msg):
+        _log.write(_msg)
+        _log.flush()
+
+    log(f"Trying to create a connection to {address}")
+    # TODO(feat): this should probably run in a loop instead of
+    # receiving a single function and then exitting immediately.
+    with child_connection(serialization_method, address) as connection:
+        log(f"Created child connection to {address}")
+        callable = connection.recv()
+        log(f"Received the callable at {address}")
+
+        result = None
+        did_it_raise = False
+        stringized_tb = None
+        try:
+            result = callable()
+        except BaseException as exc:
+            result = exc
+            did_it_raise = True
+            num_frames = len(traceback.extract_stack()[:-4])
+            stringized_tb = "".join(traceback.format_exc(limit=-num_frames))
+        finally:
+            try:
+                connection.send((result, did_it_raise, stringized_tb))
+            except BaseException:
+                if did_it_raise:
+                    # If we can't even send it through the connection
+                    # still try to dump it to the stderr as the last
+                    # resort.
+                    assert isinstance(result, BaseException)
+                    traceback.print_exception(
+                        type(result),
+                        result,
+                        result.__traceback__,
+                    )
+                raise
+
+
+def _get_shell_bootstrap() -> str:
+    # Return a string that contains environment variables that
+    # might be used during isolated hook's execution.
+    return " ".join(
+        f"{session_variable}={os.getenv(session_variable)}"
+        for session_variable in [
+            # 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
+            # on the newly created debug session.
+            "PYTHONPATH",
+        ]
+        if session_variable in os.environ
+    )
+
+
+def main() -> int:
+    parser = ArgumentParser()
+    parser.add_argument("listen_at")
+    parser.add_argument("--with-pdb", action="store_true", default=False)
+    parser.add_argument("--serialization-backend", default="pickle")
+    parser.add_argument("--log-fd", type=int)
+
+    options = parser.parse_args()
+    if IS_DEBUG_MODE:
+        assert not options.with_pdb, "--with-pdb can't be used in the debug mode"
+        message = "=" * 60
+        message += "\n" * 3
+        message += (
+            "Debug mode successfully activated. "
+            "You can start your debugging session with the following command:\n"
+        )
+        message += (
+            f"    $ {_get_shell_bootstrap()}\\\n     "
+            f"{sys.executable} {os.path.abspath(__file__)} "
+            f"--serialization-backend {options.serialization_backend} "
+            f"--with-pdb {options.listen_at}"
+        )
+        message += "\n" * 3
+        message += "=" * 60
+        print(message)
+        time.sleep(DEBUG_TIMEOUT)
+
+    serialization_method = options.serialization_backend
+    address = decode_service_address(options.listen_at)
+    run_client(
+        serialization_method,
+        address,
+        with_pdb=options.with_pdb,
+        log_fd=options.log_fd,
+    )
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

isolate/logger.py

@@ -0,0 +1,53 @@
+import json
+import os
+from datetime import datetime, timezone
+from typing import Dict
+
+from isolate.logs import LogLevel, LogSource
+
+
+# NOTE: we probably should've created a proper `logging.getLogger` here,
+# but it handling `source` would be not trivial, so we are better off
+# just keeping it simple for now.
+class IsolateLogger:
+    extra_labels: Dict[str, str] = {}
+
+    def __init__(self, log_labels: Dict[str, str]):
+        self.log_labels = log_labels
+
+    def log(self, level: LogLevel, message: str, source: LogSource) -> None:
+        record = {
+            # Set the timestamp from source so we can be sure no buffering or
+            # latency is affecting the timestamp.
+            "logged_at": datetime.now(tz=timezone.utc).isoformat(),
+            "isolate_source": source.name,
+            "level": level.name,
+            "message": message,
+            **self.log_labels,
+            **self.extra_labels,
+        }
+        print(json.dumps(record))
+
+    @classmethod
+    def with_env_expanded(cls, labels: Dict[str, str]) -> "IsolateLogger":
+        for key, value in labels.items():
+            if value.startswith("$"):
+                expanded = os.getenv(value[1:])
+            else:
+                expanded = value
+            if expanded is not None:
+                labels[key] = expanded
+
+        return cls(labels)
+
+    @classmethod
+    def from_env(cls) -> "IsolateLogger":
+        _labels: Dict[str, str] = {}
+        raw = os.getenv("ISOLATE_LOG_LABELS")
+        if raw:
+            try:
+                _labels = json.loads(raw)
+            except json.JSONDecodeError:
+                print("Failed to parse ISOLATE_LOG_LABELS")
+
+        return cls.with_env_expanded(labels=_labels)

isolate/logs.py

@@ -0,0 +1,76 @@
+from __future__ import annotations
+
+import tempfile
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from enum import Enum
+from functools import total_ordering
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from isolate.backends import BaseEnvironment
+
+_SYSTEM_TEMP_DIR = Path(tempfile.gettempdir())
+
+
+class LogSource(str, Enum):
+    """Represents where the log orinates from."""
+
+    # During the environment creation process (e.g. if the environment
+    # is already created/cached, then no logs from this source will be
+    # emitted).
+    BUILDER = "builder"
+
+    # During the environment execution process (from the server<->agent
+    # communication, mostly for debugging purposes).
+    BRIDGE = "bridge"
+
+    # From the user script itself (e.g. a print() call in the given
+    # function). The stream will be attached as level (stdout or stderr)
+    USER = "user"
+
+
+@total_ordering
+class LogLevel(Enum):
+    """Represents the log level."""
+
+    TRACE = 0
+    DEBUG = 10
+    INFO = 20
+    WARNING = 30
+    ERROR = 40
+
+    # For user scripts
+    STDOUT = 100
+    STDERR = 110
+
+    def __lt__(self, other: LogLevel) -> bool:
+        if self.__class__ is other.__class__:
+            return self.value < other.value
+        return NotImplemented
+
+    def __str__(self) -> str:
+        return self.name.lower()
+
+
+@dataclass
+class Log:
+    """A structured log message with an option source and level."""
+
+    message: str
+    source: LogSource
+    level: LogLevel = LogLevel.INFO
+    bound_env: BaseEnvironment | None = field(default=None, repr=False)
+    timestamp: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
+
+    def __str__(self) -> str:
+        parts = [self.timestamp.strftime("%m/%d/%Y %H:%M:%S")]
+        if self.bound_env:
+            parts.append(f"[{self.bound_env.key[:6]}]")
+        else:
+            parts.append("[global]")
+
+        parts.append(f"[{self.source}]".ljust(10))
+        parts.append(f"[{self.level}]".ljust(10))
+        return " ".join(parts) + self.message

isolate/registry.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+import sys
+from typing import TYPE_CHECKING, Any
+
+if sys.version_info >= (3, 10):
+    import importlib.metadata as importlib_metadata
+else:
+    import importlib_metadata
+
+if TYPE_CHECKING:
+    from isolate.backends import BaseEnvironment
+
+# Any new environments can register themselves during package installation
+# time by simply adding an entry point to the `isolate.environment` group.
+_ENTRY_POINT = "isolate.backends"
+
+_ENTRY_POINTS: dict[str, importlib_metadata.EntryPoint] = {}
+_ENVIRONMENTS: dict[str, type[BaseEnvironment]] = {}
+
+
+def _reload_registry() -> None:
+    entry_points = importlib_metadata.entry_points()
+    _ENTRY_POINTS.update(
+        {
+            # We are not immediately loading the backend class here
+            # since it might cause importing modules that we won't be
+            # using at all.
+            entry_point.name: entry_point
+            for entry_point in entry_points.select(group=_ENTRY_POINT)
+        }
+    )
+
+
+_reload_registry()
+
+
+def prepare_environment(
+    kind: str,
+    **kwargs: Any,
+) -> BaseEnvironment:
+    """Get the environment for the given `kind` with the given `config`."""
+    from isolate.backends.settings import DEFAULT_SETTINGS
+
+    if kind not in _ENVIRONMENTS:
+        entry_point = _ENTRY_POINTS.get(kind)
+        if entry_point is None:
+            raise ValueError(f"Unknown environment: '{kind}'")
+
+        _ENVIRONMENTS[kind] = entry_point.load()
+
+    settings = kwargs.pop("context", DEFAULT_SETTINGS)
+    return _ENVIRONMENTS[kind].from_config(config=kwargs, settings=settings)

isolate/server/__init__.py

@@ -0,0 +1 @@
+from isolate.server.server import BridgeManager, IsolateServicer  # noqa: F401

isolate/server/definitions/__init__.py

@@ -0,0 +1,13 @@
+from google.protobuf.json_format import MessageToDict as struct_to_dict  # noqa: F401
+from google.protobuf.struct_pb2 import Struct  # noqa: F401
+
+# Inherit everything from the gRPC connection handler.
+from isolate.connections.grpc.definitions import *  # noqa: F403
+from isolate.server.definitions.server_pb2 import *  # noqa: F403
+from isolate.server.definitions.server_pb2_grpc import (  # noqa: F401
+    IsolateServicer,
+    IsolateStub,
+)
+from isolate.server.definitions.server_pb2_grpc import (  # noqa: F401
+    add_IsolateServicer_to_server as register_isolate,
+)

isolate/server/definitions/server.proto

@@ -0,0 +1,80 @@
+syntax = "proto3";
+
+import "common.proto";
+import "google/protobuf/struct.proto";
+
+service Isolate {
+    // Run the given function on the specified environment. Streams logs
+    // and the result originating from that function.
+    rpc Run (BoundFunction) returns (stream PartialRunResult) {}
+
+    // Submit a function to be run without waiting for results.
+    rpc Submit (SubmitRequest) returns (SubmitResponse) {}
+
+    // Set the metadata for a task.
+    rpc SetMetadata (SetMetadataRequest) returns (SetMetadataResponse) {}
+
+    // List running tasks
+    rpc List (ListRequest) returns (ListResponse) {}
+
+    // Cancel a running task
+    rpc Cancel (CancelRequest) returns (CancelResponse) {}
+}
+
+message BoundFunction {
+    repeated EnvironmentDefinition environments = 1;
+    SerializedObject function = 2;
+    optional SerializedObject setup_func = 3;
+    bool stream_logs = 4;
+}
+
+message EnvironmentDefinition {
+    // Kind of the isolate environment.
+    string kind = 1;
+    // A free-form definition of environment properties.
+    google.protobuf.Struct configuration = 2;
+    // Whether to force-create this environment or not.
+    bool force = 3;
+}
+
+message SubmitRequest {
+    // The function to run.
+    BoundFunction function = 1;
+    // Task metadata.
+    TaskMetadata metadata = 2;
+}
+
+message TaskMetadata {
+    // Labels to attach to the logs.
+    map<string, string> logger_labels = 1;
+}
+
+message SubmitResponse {
+    string task_id = 1;
+}
+
+message SetMetadataRequest{
+    string task_id = 1;
+    TaskMetadata metadata = 2;
+}
+
+message SetMetadataResponse {
+}
+
+message ListRequest {
+}
+
+message TaskInfo {
+    string task_id = 1;
+}
+
+message ListResponse {
+    repeated TaskInfo tasks = 1;
+}
+
+message CancelRequest {
+    string task_id = 1;
+}
+
+message CancelResponse {
+}

isolate/server/definitions/server_pb2.py

@@ -0,0 +1,56 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: server.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 isolate.connections.grpc.definitions import common_pb2 as common__pb2
+from google.protobuf import struct_pb2 as google_dot_protobuf_dot_struct__pb2
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x0cserver.proto\x1a\x0c\x63ommon.proto\x1a\x1cgoogle/protobuf/struct.proto\"\xb2\x01\n\rBoundFunction\x12,\n\x0c\x65nvironments\x18\x01 \x03(\x0b\x32\x16.EnvironmentDefinition\x12#\n\x08\x66unction\x18\x02 \x01(\x0b\x32\x11.SerializedObject\x12*\n\nsetup_func\x18\x03 \x01(\x0b\x32\x11.SerializedObjectH\x00\x88\x01\x01\x12\x13\n\x0bstream_logs\x18\x04 \x01(\x08\x42\r\n\x0b_setup_func\"d\n\x15\x45nvironmentDefinition\x12\x0c\n\x04kind\x18\x01 \x01(\t\x12.\n\rconfiguration\x18\x02 \x01(\x0b\x32\x17.google.protobuf.Struct\x12\r\n\x05\x66orce\x18\x03 \x01(\x08\"R\n\rSubmitRequest\x12 \n\x08\x66unction\x18\x01 \x01(\x0b\x32\x0e.BoundFunction\x12\x1f\n\x08metadata\x18\x02 \x01(\x0b\x32\r.TaskMetadata\"{\n\x0cTaskMetadata\x12\x36\n\rlogger_labels\x18\x01 \x03(\x0b\x32\x1f.TaskMetadata.LoggerLabelsEntry\x1a\x33\n\x11LoggerLabelsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t:\x02\x38\x01\"!\n\x0eSubmitResponse\x12\x0f\n\x07task_id\x18\x01 \x01(\t\"F\n\x12SetMetadataRequest\x12\x0f\n\x07task_id\x18\x01 \x01(\t\x12\x1f\n\x08metadata\x18\x02 \x01(\x0b\x32\r.TaskMetadata\"\x15\n\x13SetMetadataResponse\"\r\n\x0bListRequest\"\x1b\n\x08TaskInfo\x12\x0f\n\x07task_id\x18\x01 \x01(\t\"(\n\x0cListResponse\x12\x18\n\x05tasks\x18\x01 \x03(\x0b\x32\t.TaskInfo\" \n\rCancelRequest\x12\x0f\n\x07task_id\x18\x01 \x01(\t\"\x10\n\x0e\x43\x61ncelResponse2\xf4\x01\n\x07Isolate\x12,\n\x03Run\x12\x0e.BoundFunction\x1a\x11.PartialRunResult\"\x00\x30\x01\x12+\n\x06Submit\x12\x0e.SubmitRequest\x1a\x0f.SubmitResponse\"\x00\x12:\n\x0bSetMetadata\x12\x13.SetMetadataRequest\x1a\x14.SetMetadataResponse\"\x00\x12%\n\x04List\x12\x0c.ListRequest\x1a\r.ListResponse\"\x00\x12+\n\x06\x43\x61ncel\x12\x0e.CancelRequest\x1a\x0f.CancelResponse\"\x00\x62\x06proto3')
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'server_pb2', _globals)
+if _descriptor._USE_C_DESCRIPTORS == False:
+  DESCRIPTOR._options = None
+  _globals['_TASKMETADATA_LOGGERLABELSENTRY']._options = None
+  _globals['_TASKMETADATA_LOGGERLABELSENTRY']._serialized_options = b'8\001'
+  _globals['_BOUNDFUNCTION']._serialized_start=61
+  _globals['_BOUNDFUNCTION']._serialized_end=239
+  _globals['_ENVIRONMENTDEFINITION']._serialized_start=241
+  _globals['_ENVIRONMENTDEFINITION']._serialized_end=341
+  _globals['_SUBMITREQUEST']._serialized_start=343
+  _globals['_SUBMITREQUEST']._serialized_end=425
+  _globals['_TASKMETADATA']._serialized_start=427
+  _globals['_TASKMETADATA']._serialized_end=550
+  _globals['_TASKMETADATA_LOGGERLABELSENTRY']._serialized_start=499
+  _globals['_TASKMETADATA_LOGGERLABELSENTRY']._serialized_end=550
+  _globals['_SUBMITRESPONSE']._serialized_start=552
+  _globals['_SUBMITRESPONSE']._serialized_end=585
+  _globals['_SETMETADATAREQUEST']._serialized_start=587
+  _globals['_SETMETADATAREQUEST']._serialized_end=657
+  _globals['_SETMETADATARESPONSE']._serialized_start=659
+  _globals['_SETMETADATARESPONSE']._serialized_end=680
+  _globals['_LISTREQUEST']._serialized_start=682
+  _globals['_LISTREQUEST']._serialized_end=695
+  _globals['_TASKINFO']._serialized_start=697
+  _globals['_TASKINFO']._serialized_end=724
+  _globals['_LISTRESPONSE']._serialized_start=726
+  _globals['_LISTRESPONSE']._serialized_end=766
+  _globals['_CANCELREQUEST']._serialized_start=768
+  _globals['_CANCELREQUEST']._serialized_end=800
+  _globals['_CANCELRESPONSE']._serialized_start=802
+  _globals['_CANCELRESPONSE']._serialized_end=818
+  _globals['_ISOLATE']._serialized_start=821
+  _globals['_ISOLATE']._serialized_end=1065
+# @@protoc_insertion_point(module_scope)