isolate 0.13.8__tar.gz → 0.13.10__tar.gz

{isolate-0.13.8 → isolate-0.13.10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: isolate
-Version: 0.13.8
+Version: 0.13.10
 Summary: Managed isolated environments for Python
 Author-email: Features & Labels <hello@fal.ai>
 Project-URL: Issues, https://github.com/fal-ai/isolate/issues

{isolate-0.13.8 → isolate-0.13.10}/src/isolate/_isolate_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.13.8'
-__version_tuple__ = version_tuple = (0, 13, 8)
+__version__ = version = '0.13.10'
+__version_tuple__ = version_tuple = (0, 13, 10)

{isolate-0.13.8 → isolate-0.13.10}/src/isolate/logger.py

@@ -1,15 +1,29 @@
 import json
+import os
 
 
 # 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:
+    def __init__(self):
+        self.log_labels = {}
+        raw = os.getenv("ISOLATE_LOG_LABELS")
+        if raw:
+            labels = json.loads(raw)
+            for key, value in labels.items():
+                if value.startswith("$"):
+                    expanded = os.getenv(value[1:])
+                else:
+                    expanded = value
+                self.log_labels[key] = expanded
+
     def log(self, level, message, source):
         record = {
             "isolate_source": source.name,
             "level": level.name,
             "message": message,
+            **self.log_labels,
         }
         print(json.dumps(record))
 

{isolate-0.13.8 → isolate-0.13.10}/src/isolate/server/server.py

@@ -1,10 +1,12 @@
 from __future__ import annotations
 
+import functools
 import os
 import threading
 import time
 import traceback
 import uuid
+from argparse import ArgumentParser
 from collections import defaultdict
 from concurrent import futures
 from concurrent.futures import ThreadPoolExecutor
@@ -16,6 +18,7 @@ from typing import Any, Callable, Iterator, cast
 
 import grpc
 from grpc import ServicerContext, StatusCode
+from grpc.experimental import wrap_server_method_handler
 
 from isolate.backends import (
     EnvironmentCreationError,
@@ -465,13 +468,120 @@ class LogHandler:
         self.messages.put_nowait(grpc_result)
 
 
-def main() -> None:
+@dataclass
+class ServerBoundInterceptor(grpc.ServerInterceptor):
+    _server: grpc.Server | None = None
+    _servicer: IsolateServicer | None = None
+
+    def register_server(self, server: grpc.Server) -> None:
+        if self._server is not None:
+            raise RuntimeError("A server is already bound to this interceptor.")
+
+        self._server = server
+
+    @property
+    def server(self) -> grpc.Server:
+        if self._server is None:
+            raise RuntimeError("No server was bound to this interceptor.")
+
+        return self._server
+
+    def register_servicer(self, servicer: IsolateServicer) -> None:
+        if self._servicer is not None:
+            raise RuntimeError("A servicer is already bound to this interceptor.")
+
+        self._servicer = servicer
+
+    @property
+    def servicer(self) -> IsolateServicer:
+        if self._servicer is None:
+            raise RuntimeError("No servicer was bound to this interceptor.")
+
+        return self._servicer
+
+
+@dataclass
+class SingleTaskInterceptor(ServerBoundInterceptor):
+    """Sets server to terminate after the first Submit/Run task."""
+
+    _done: bool = False
+
+    def intercept_service(self, continuation, handler_call_details):
+        handler = continuation(handler_call_details)
+
+        is_submit = handler_call_details.method == "/Isolate/Submit"
+        is_run = handler_call_details.method == "/Isolate/Run"
+        is_new_task = is_submit or is_run
+
+        if is_new_task and self._done:
+            raise grpc.RpcError(
+                grpc.StatusCode.UNAVAILABLE,
+                "Server has already served one Run/Submit task.",
+            )
+        elif is_new_task:
+            self._done = True
+        else:
+            # Let other requests like List/Cancel/etc pass through
+            return continuation(handler_call_details)
+
+        def wrapper(method_impl):
+            @functools.wraps(method_impl)
+            def _wrapper(request, context):
+                def _stop():
+                    if is_submit:
+                        # Wait for the task to finish
+                        while self.server.servicer.background_tasks:
+                            time.sleep(0.1)
+                    self.server.stop(grace=0.1)
+
+                context.add_callback(_stop)
+                return method_impl(request, context)
+
+            return _wrapper
+
+        return wrap_server_method_handler(wrapper, handler)
+
+
+def main(argv: list[str] | None = None) -> None:
+    parser = ArgumentParser()
+    parser.add_argument("--host", default="0.0.0.0")
+    parser.add_argument("--port", type=int, default=50001)
+    parser.add_argument(
+        "--single-use",
+        action="store_true",
+        help="Terminate the server after the first Run or Submit task is completed.",
+    )
+    parser.add_argument(
+        "--num-workers",
+        type=int,
+        default=MAX_THREADS,
+        help="Number of worker threads to use for the gRPC server.",
+    )
+
+    options = parser.parse_args(argv)
+    if options.num_workers is None:
+        options.num_workers = 1 if options.single_use else os.cpu_count()
+
+    interceptors: list[ServerBoundInterceptor] = []
+    if options.single_use:
+        interceptors.append(SingleTaskInterceptor())
+
     server = grpc.server(
-        futures.ThreadPoolExecutor(max_workers=MAX_THREADS),
+        futures.ThreadPoolExecutor(max_workers=options.num_workers),
         options=get_default_options(),
+        interceptors=interceptors,
     )
+
+    for interceptor in interceptors:
+        interceptor.register_server(server)
+
     with BridgeManager() as bridge_manager:
-        definitions.register_isolate(IsolateServicer(bridge_manager), server)
+        servicer = IsolateServicer(bridge_manager)
+
+        for interceptor in interceptors:
+            interceptor.register_servicer(servicer)
+
+        definitions.register_isolate(servicer, server)
         health.register_health(HealthServicer(), server)
 
         server.add_insecure_port("[::]:50001")

{isolate-0.13.8 → isolate-0.13.10}/src/isolate.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: isolate
-Version: 0.13.8
+Version: 0.13.10
 Summary: Managed isolated environments for Python
 Author-email: Features & Labels <hello@fal.ai>
 Project-URL: Issues, https://github.com/fal-ai/isolate/issues

{isolate-0.13.8 → isolate-0.13.10}/src/isolate.egg-info/SOURCES.txt

@@ -73,6 +73,7 @@ tests/test_concurrency.py
 tests/test_connections.py
 tests/test_isolate.py
 tests/test_log.py
+tests/test_logger.py
 tests/test_serialization.py
 tests/test_server.py
 tools/Dockerfile

isolate-0.13.10/tests/test_logger.py

@@ -0,0 +1,22 @@
+import json
+
+from isolate.logger import IsolateLogger
+
+
+def test_logger(monkeypatch):
+    labels = {
+        "foo": "$MYENVVAR1",
+        "bar": "$MYENVVAR2",
+        "baz": "baz",
+        "qux": "qux",
+    }
+    monkeypatch.setenv("MYENVVAR1", "myenvvar1")
+    monkeypatch.setenv("MYENVVAR2", "myenvvar2")
+    monkeypatch.setenv("ISOLATE_LOG_LABELS", json.dumps(labels))
+    logger = IsolateLogger()
+    assert logger.log_labels == {
+        "foo": "myenvvar1",
+        "bar": "myenvvar2",
+        "baz": "baz",
+        "qux": "qux",
+    }