isolate 0.22.0__py3-none-any.whl

isolate/connections/grpc/_base.py

@@ -0,0 +1,175 @@
+import os
+import socket
+import subprocess
+from contextlib import contextmanager
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, ContextManager, Iterator, List, Tuple, Union, cast
+
+import grpc
+
+from isolate.backends import (
+    BasicCallable,
+    CallResultType,
+    EnvironmentConnection,
+)
+from isolate.connections._local import PythonExecutionBase, agent_startup
+from isolate.connections.common import serialize_object
+from isolate.connections.grpc import agent, definitions
+from isolate.connections.grpc.configuration import get_default_options
+from isolate.connections.grpc.interface import from_grpc
+from isolate.logs import LogLevel, LogSource
+
+
+class AgentError(Exception):
+    """An internal problem caused by (most probably) the agent."""
+
+
+PROCESS_SHUTDOWN_TIMEOUT_SECONDS = float(
+    os.getenv("ISOLATE_SHUTDOWN_GRACE_PERIOD", "60")
+)
+
+
+@dataclass
+class GRPCExecutionBase(EnvironmentConnection):
+    """A customizable gRPC-based execution backend."""
+
+    def start_agent(self) -> ContextManager[Tuple[str, grpc.ChannelCredentials]]:
+        """Starts the gRPC agent and returns the address it is listening on and
+        the required credentials to connect to it."""
+        raise NotImplementedError
+
+    def abort_agent(self) -> None:
+        raise NotImplementedError
+
+    @contextmanager
+    def _establish_bridge(
+        self,
+        *,
+        max_wait_timeout: float = 20.0,
+    ) -> Iterator[definitions.AgentStub]:
+        with self.start_agent() as (address, credentials):
+            with grpc.secure_channel(
+                address,
+                credentials,
+                options=get_default_options(),
+            ) as channel:
+                channel_status = grpc.channel_ready_future(channel)
+                try:
+                    channel_status.result(timeout=max_wait_timeout)
+                except grpc.FutureTimeoutError:
+                    raise AgentError(
+                        "Couldn't connect to the gRPC server in the agent "
+                        f"(listening at {address}) in time."
+                    )
+                stub = definitions.AgentStub(channel)
+                stub._channel = channel  # type: ignore
+                yield stub
+
+    def run(
+        self,
+        executable: BasicCallable,
+        *args: Any,
+        **kwargs: Any,
+    ) -> CallResultType:  # type: ignore[type-var]
+        # Implementation details
+        # ======================
+        #
+        #  RPC Flow:
+        #  ---------
+        #  1. [controller]: Spawn the agent.
+        #  2.      [agent]: Start listening at the given address.
+        #  3. [controller]: Await *at most* max_wait_timeout seconds for the agent to
+        #                   be available if it doesn't do it until then,
+        #                   raise an AgentError.
+        #  4. [controller]: If the server is available, then establish the bridge and
+        #                   pass the 'function' as the input.
+        #  5.      [agent]: Receive the function, deserialize it, start the execution.
+        #  6. [controller]: Watch agent for logs (stdout/stderr), and as soon as they
+        #                   appear call the log handler.
+        #  7.      [agent]: Once the execution of the function is finished, send the
+        #                   result using the same serialization method.
+        #  8. [controller]: Receive the result back and return it.
+
+        method = self.environment.settings.serialization_method
+        function = definitions.SerializedObject(
+            method=method,
+            definition=serialize_object(method, executable),
+            was_it_raised=False,
+            stringized_traceback=None,
+        )
+        function_call = definitions.FunctionCall(
+            function=function,
+        )
+
+        with self._establish_bridge() as bridge:
+            for partial_result in bridge.Run(function_call):
+                for raw_log in partial_result.logs:
+                    log = from_grpc(raw_log)
+                    self.log(log.message, level=log.level, source=log.source)
+
+                if partial_result.is_complete:
+                    if not partial_result.result:
+                        raise AgentError(
+                            "The agent didn't return a result, but it should have."
+                        )
+
+                    return cast(CallResultType, from_grpc(partial_result.result))
+
+        raise AgentError(
+            "No result object was received from the agent "
+            "(it never set is_complete to True)."
+        )
+
+
+class LocalPythonGRPC(PythonExecutionBase[str], GRPCExecutionBase):
+    _process: Union[None, subprocess.Popen] = None
+
+    @contextmanager
+    def start_agent(self) -> Iterator[Tuple[str, grpc.ChannelCredentials]]:
+        def find_free_port() -> Tuple[str, int]:
+            """Find a free port in the system."""
+            with socket.socket() as _temp_socket:
+                _temp_socket.bind(("", 0))
+                return _temp_socket.getsockname()
+
+        host, port = find_free_port()
+        address = f"{host}:{port}"
+        self._process = None
+        try:
+            with self.start_process(address) as process:
+                self._process = process
+                yield address, grpc.local_channel_credentials()
+        finally:
+            self.abort_agent()
+
+    def abort_agent(self) -> None:
+        if self._process is not None:
+            try:
+                print("Terminating the agent process...")
+                self._process.terminate()
+                self._process.wait(timeout=PROCESS_SHUTDOWN_TIMEOUT_SECONDS)
+                print("Agent process shutdown gracefully")
+            except Exception as exc:
+                print(f"Failed to shutdown the agent process gracefully: {exc}")
+                self._process.kill()
+
+    def get_python_cmd(
+        self,
+        executable: Path,
+        connection: str,
+        log_fd: int,
+    ) -> List[Union[str, Path]]:
+        return [
+            executable,
+            agent_startup.__file__,
+            agent.__file__,
+            connection,
+            "--log-fd",
+            str(log_fd),
+        ]
+
+    def handle_agent_log(
+        self, line: str, *, level: LogLevel, source: LogSource
+    ) -> None:
+        self.log(line, level=level, source=source)

isolate/connections/grpc/agent.py

@@ -0,0 +1,284 @@
+# agent-requires: isolate[server]
+"""
+This file contains the implementation of the gRPC agent. The agent is a
+separate process that is responsible for running the user code in a
+sandboxed environment.
+
+This file is referenced by the latest version of the `isolate` package
+but then runs it in the context of the frozen agent built environment.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import os
+import signal
+import sys
+import traceback
+from argparse import ArgumentParser
+from concurrent import futures
+from dataclasses import dataclass
+from typing import (
+    Any,
+    AsyncIterator,
+    Iterable,
+)
+
+from grpc import StatusCode, aio, local_server_credentials
+
+from isolate.connections.grpc.definitions import PartialRunResult
+
+try:
+    from isolate import __version__ as agent_version
+except ImportError:
+    agent_version = "UNKNOWN"
+
+from isolate.backends.common import sha256_digest_of
+from isolate.connections.common import SerializationError, serialize_object
+from isolate.connections.grpc import definitions
+from isolate.connections.grpc.configuration import get_default_options
+from isolate.connections.grpc.interface import from_grpc
+
+
+@dataclass
+class AbortException(Exception):
+    message: str
+
+
+class AgentServicer(definitions.AgentServicer):
+    def __init__(self, log_fd: int | None = None):
+        super().__init__()
+
+        self._run_cache: dict[str, Any] = {}
+        self._log = sys.stdout if log_fd is None else os.fdopen(log_fd, "w")
+        self._thread_pool = futures.ThreadPoolExecutor(max_workers=1)
+
+        def handle_termination(*args):
+            self.log("Termination signal received, shutting down...")
+            signal.raise_signal(signal.SIGTERM)
+
+        signal.signal(signal.SIGINT, handle_termination)
+
+    async def Run(
+        self,
+        request: definitions.FunctionCall,
+        context: aio.ServicerContext,
+    ) -> AsyncIterator[PartialRunResult]:
+        self.log(f"A connection has been established: {context.peer()}!")
+        server_version = os.getenv("ISOLATE_SERVER_VERSION") or "unknown"
+        self.log(f"Isolate info: server {server_version}, agent {agent_version}")
+
+        extra_args = []
+        if request.HasField("setup_func"):
+            cache_key = sha256_digest_of(
+                request.setup_func.definition,
+                request.setup_func.method,
+            )
+            if cache_key not in self._run_cache:
+                try:
+                    (
+                        result,
+                        was_it_raised,
+                        stringized_tb,
+                    ) = await self.execute_function(
+                        request.setup_func,
+                        "setup",
+                    )
+
+                    if was_it_raised:
+                        self.log(
+                            "The setup function has thrown an error. Aborting the run."
+                        )
+                        yield self.send_object(
+                            request.setup_func.method,
+                            result,
+                            was_it_raised,
+                            stringized_tb,
+                        )
+                        raise AbortException("The setup function has thrown an error.")
+                except AbortException as exc:
+                    self.abort_with_msg(context, exc.message)
+                    return
+                else:
+                    assert not was_it_raised
+                    self._run_cache[cache_key] = result
+
+            extra_args.append(self._run_cache[cache_key])
+
+        try:
+            result, was_it_raised, stringized_tb = await self.execute_function(
+                request.function,
+                "function",
+                extra_args=extra_args,
+            )
+            yield self.send_object(
+                request.function.method,
+                result,
+                was_it_raised,
+                stringized_tb,
+            )
+        except AbortException as exc:
+            self.abort_with_msg(context, exc.message)
+            return
+
+    async def execute_function(
+        self,
+        function: definitions.SerializedObject,
+        function_kind: str,
+        *,
+        extra_args: Iterable[Any] = (),
+    ) -> tuple[Any, bool, str | None]:
+        if function.was_it_raised:
+            raise AbortException(
+                f"The {function_kind} function must be callable, "
+                "not a raised exception."
+            )
+
+        try:
+            # TODO: technically any sort of exception could be raised here, since
+            # depickling is basically involves code execution from the *user*.
+            function = from_grpc(function)
+        except SerializationError as exc:
+            str_tb = traceback.format_exc()
+            self.log(str_tb)
+            self.log(f"The {function_kind} function could not be deserialized.")
+            return exc, True, str_tb
+
+        if not callable(function):
+            raise AbortException(
+                f"The {function_kind} function must be callable, "
+                f"not {type(function).__name__}."
+            )
+
+        self.log(f"Starting the execution of the {function_kind} function.")
+
+        was_it_raised = False
+        stringized_tb = None
+        try:
+            # Newer fal SDK will mark async entrypoints with `_run_on_main_thread` so
+            # we execute on the main loop and can await the coroutine they return.
+            # Older fal SDK still call `asyncio.run(...)`.
+            # To avoid error "asyncio.run() cannot be called from a running event loop"
+            # and be backward compatible,
+            # we offload those unflagged functions to a thread pool.
+
+            if getattr(function, "_run_on_main_thread", False):
+                result = function(*extra_args)
+            else:
+                result = self._thread_pool.submit(function, *extra_args).result()
+
+            if asyncio.iscoroutine(result):
+                result = await result
+
+        except BaseException as exc:
+            result = exc
+            was_it_raised = True
+            num_frames = len(traceback.extract_stack()[:-5])
+            stringized_tb = "".join(traceback.format_exc(limit=-num_frames))
+
+        if not was_it_raised:
+            self.log(f"Completed the execution of the {function_kind} function.")
+        else:
+            self.log(
+                f"Completed the execution of the {function_kind} function"
+                f" with an error: {result}\nTraceback:\n{stringized_tb}"
+            )
+        return result, was_it_raised, stringized_tb
+
+    def send_object(
+        self,
+        serialization_method: str,
+        result: object,
+        was_it_raised: bool,
+        stringized_tb: str | None,
+    ) -> definitions.PartialRunResult:
+        try:
+            definition = serialize_object(serialization_method, result)
+        except SerializationError:
+            if stringized_tb:
+                print(stringized_tb, file=sys.stderr)
+            self.log(traceback.format_exc())
+            raise AbortException(
+                "Error while serializing the execution result "
+                f"(object of type {type(result)})."
+            )
+        except BaseException:
+            self.log(traceback.format_exc())
+            raise AbortException(
+                "An unexpected error occurred while serializing the result."
+            )
+
+        self.log("Sending the result.")
+        serialized_obj = definitions.SerializedObject(
+            method=serialization_method,
+            definition=definition,
+            was_it_raised=was_it_raised,
+            stringized_traceback=stringized_tb,
+        )
+        return definitions.PartialRunResult(
+            result=serialized_obj,
+            is_complete=True,
+            logs=[],
+        )
+
+    def log(self, message: str) -> None:
+        self._log.write(message + "\n")
+        self._log.flush()
+
+    def abort_with_msg(
+        self,
+        context: aio.ServicerContext,
+        message: str,
+        *,
+        code: StatusCode = StatusCode.INVALID_ARGUMENT,
+    ) -> None:
+        context.set_code(code)
+        context.set_details(message)
+        return None
+
+
+def create_server(address: str) -> aio.Server:
+    """Create a new (temporary) gRPC server listening on the given
+    address."""
+    # Use asyncio server so requests can run in the main thread and intercept signals
+    # There seems to be a weird bug with grpcio that makes subsequent requests fail with
+    # concurrent rpc limit exceeded if we set maximum_current_rpcs to 1. Setting it to 2
+    # fixes it, even though in practice, we only run one request at a time.
+    server = aio.server(
+        maximum_concurrent_rpcs=2,
+        options=get_default_options(),
+    )
+
+    # Local server credentials allow us to ensure that the
+    # connection is established by a local process.
+    server_credentials = local_server_credentials()
+    server.add_secure_port(address, server_credentials)
+    return server
+
+
+async def run_agent(address: str, log_fd: int | None = None) -> int:
+    """Run the agent servicer on the given address."""
+    server = create_server(address)
+    servicer = AgentServicer(log_fd=log_fd)
+
+    # This function just calls some methods on the server
+    # and register a generic handler for the bridge. It does
+    # not have any global side effects.
+    definitions.register_agent(servicer, server)
+
+    await server.start()
+    await server.wait_for_termination()
+    return 0
+
+
+async def main() -> int:
+    parser = ArgumentParser()
+    parser.add_argument("address", type=str)
+    parser.add_argument("--log-fd", type=int)
+
+    options = parser.parse_args()
+    return await run_agent(options.address, log_fd=options.log_fd)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

isolate/connections/grpc/configuration.py

@@ -0,0 +1,23 @@
+import ast
+import os
+
+_GRPC_OPTION_PREFIX = "ISOLATE_GRPC_CALL_"
+
+
+def get_default_options():
+    """Return the default list of GRPC call options (both for
+    server and client) which are set via environment variables.
+
+    Each environment variable starting with `ISOLATE_GRPC_CALL_`
+    will be converted to a GRPC option. The name of the option
+    will be the name of the environment variable, with the
+    `ISOLATE_GRPC_CALL_` prefix removed and converted to lowercase.
+    """
+
+    options = []
+    for raw_key, raw_value in os.environ.items():
+        if raw_key.startswith(_GRPC_OPTION_PREFIX):
+            field = raw_key[len(_GRPC_OPTION_PREFIX) :].lower()
+            value = ast.literal_eval(raw_value)
+            options.append((f"grpc.{field}", value))
+    return options

isolate/connections/grpc/definitions/__init__.py

@@ -0,0 +1,11 @@
+from google.protobuf.message import Message  # noqa: F401
+
+from isolate.connections.grpc.definitions.agent_pb2 import *  # noqa: F403
+from isolate.connections.grpc.definitions.agent_pb2_grpc import (  # noqa: F401
+    AgentServicer,
+    AgentStub,
+)
+from isolate.connections.grpc.definitions.agent_pb2_grpc import (  # noqa: F401
+    add_AgentServicer_to_server as register_agent,
+)
+from isolate.connections.grpc.definitions.common_pb2 import *  # noqa: F403

isolate/connections/grpc/definitions/agent.proto

@@ -0,0 +1,18 @@
+syntax = "proto3";
+
+import "common.proto";
+
+service Agent {
+    // Start running the given function, and stream results back.
+    rpc Run (FunctionCall) returns (stream PartialRunResult) {}
+}
+
+message FunctionCall {
+    // The function to execute and return the results to.
+    SerializedObject function = 1;
+    // Optionally the setup function which will be passed
+    // as the first argument to the given function. This
+    // has to be an idempotent step since the result for
+    // this executable will be cached.
+    optional SerializedObject setup_func = 2;
+}

isolate/connections/grpc/definitions/agent_pb2.py

@@ -0,0 +1,29 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: agent.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
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x0b\x61gent.proto\x1a\x0c\x63ommon.proto\"n\n\x0c\x46unctionCall\x12#\n\x08\x66unction\x18\x01 \x01(\x0b\x32\x11.SerializedObject\x12*\n\nsetup_func\x18\x02 \x01(\x0b\x32\x11.SerializedObjectH\x00\x88\x01\x01\x42\r\n\x0b_setup_func24\n\x05\x41gent\x12+\n\x03Run\x12\r.FunctionCall\x1a\x11.PartialRunResult\"\x00\x30\x01\x62\x06proto3')
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'agent_pb2', _globals)
+if _descriptor._USE_C_DESCRIPTORS == False:
+  DESCRIPTOR._options = None
+  _globals['_FUNCTIONCALL']._serialized_start=29
+  _globals['_FUNCTIONCALL']._serialized_end=139
+  _globals['_AGENT']._serialized_start=141
+  _globals['_AGENT']._serialized_end=193
+# @@protoc_insertion_point(module_scope)

isolate/connections/grpc/definitions/agent_pb2.pyi

@@ -0,0 +1,44 @@
+"""
+@generated by mypy-protobuf.  Do not edit manually!
+isort:skip_file
+"""
+import builtins
+from isolate.connections.grpc.definitions import common_pb2
+import google.protobuf.descriptor
+import google.protobuf.message
+import sys
+
+if sys.version_info >= (3, 8):
+    import typing as typing_extensions
+else:
+    import typing_extensions
+
+DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
+
+@typing_extensions.final
+class FunctionCall(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    FUNCTION_FIELD_NUMBER: builtins.int
+    SETUP_FUNC_FIELD_NUMBER: builtins.int
+    @property
+    def function(self) -> common_pb2.SerializedObject:
+        """The function to execute and return the results to."""
+    @property
+    def setup_func(self) -> common_pb2.SerializedObject:
+        """Optionally the setup function which will be passed
+        as the first argument to the given function. This
+        has to be an idempotent step since the result for
+        this executable will be cached.
+        """
+    def __init__(
+        self,
+        *,
+        function: common_pb2.SerializedObject | None = ...,
+        setup_func: common_pb2.SerializedObject | None = ...,
+    ) -> None: ...
+    def HasField(self, field_name: typing_extensions.Literal["_setup_func", b"_setup_func", "function", b"function", "setup_func", b"setup_func"]) -> builtins.bool: ...
+    def ClearField(self, field_name: typing_extensions.Literal["_setup_func", b"_setup_func", "function", b"function", "setup_func", b"setup_func"]) -> None: ...
+    def WhichOneof(self, oneof_group: typing_extensions.Literal["_setup_func", b"_setup_func"]) -> typing_extensions.Literal["setup_func"] | None: ...
+
+global___FunctionCall = FunctionCall

isolate/connections/grpc/definitions/agent_pb2_grpc.py

@@ -0,0 +1,68 @@
+# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!
+"""Client and server classes corresponding to protobuf-defined services."""
+import grpc
+
+from isolate.connections.grpc.definitions import agent_pb2 as agent__pb2
+from isolate.connections.grpc.definitions import common_pb2 as common__pb2
+
+
+class AgentStub(object):
+    """Missing associated documentation comment in .proto file."""
+
+    def __init__(self, channel):
+        """Constructor.
+
+        Args:
+            channel: A grpc.Channel.
+        """
+        self.Run = channel.unary_stream(
+                '/Agent/Run',
+                request_serializer=agent__pb2.FunctionCall.SerializeToString,
+                response_deserializer=common__pb2.PartialRunResult.FromString,
+                )
+
+
+class AgentServicer(object):
+    """Missing associated documentation comment in .proto file."""
+
+    def Run(self, request, context):
+        """Start running the given function, and stream results back.
+        """
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+
+def add_AgentServicer_to_server(servicer, server):
+    rpc_method_handlers = {
+            'Run': grpc.unary_stream_rpc_method_handler(
+                    servicer.Run,
+                    request_deserializer=agent__pb2.FunctionCall.FromString,
+                    response_serializer=common__pb2.PartialRunResult.SerializeToString,
+            ),
+    }
+    generic_handler = grpc.method_handlers_generic_handler(
+            'Agent', rpc_method_handlers)
+    server.add_generic_rpc_handlers((generic_handler,))
+
+
+ # This class is part of an EXPERIMENTAL API.
+class Agent(object):
+    """Missing associated documentation comment in .proto file."""
+
+    @staticmethod
+    def Run(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_stream(request, target, '/Agent/Run',
+            agent__pb2.FunctionCall.SerializeToString,
+            common__pb2.PartialRunResult.FromString,
+            options, channel_credentials,
+            insecure, call_credentials, compression, wait_for_ready, timeout, metadata)