isolate 0.13.8__py3-none-any.whl → 0.13.9__py3-none-any.whl

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.9'
+__version_tuple__ = version_tuple = (0, 13, 9)

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.dist-info → isolate-0.13.9.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: isolate
-Version: 0.13.8
+Version: 0.13.9
 Summary: Managed isolated environments for Python
 Author-email: Features & Labels <hello@fal.ai>
 Project-URL: Issues, https://github.com/fal-ai/isolate/issues
@@ -8,23 +8,23 @@ Project-URL: Source, https://github.com/fal-ai/isolate
 Requires-Python: >=3.8
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: grpcio ==1.64.0
+Requires-Dist: grpcio==1.64.0
 Requires-Dist: protobuf
-Requires-Dist: tblib >=1.7.0
+Requires-Dist: tblib>=1.7.0
 Requires-Dist: platformdirs
-Requires-Dist: importlib-metadata >=4.4 ; python_version < "3.10"
+Requires-Dist: importlib-metadata>=4.4; python_version < "3.10"
 Provides-Extra: build
-Requires-Dist: virtualenv >=20.4 ; extra == 'build'
-Requires-Dist: PyYAML >=6.0 ; extra == 'build'
+Requires-Dist: virtualenv>=20.4; extra == "build"
+Requires-Dist: PyYAML>=6.0; extra == "build"
 Provides-Extra: dev
-Requires-Dist: isolate[test] ; extra == 'dev'
-Requires-Dist: grpcio-tools ==1.64.0 ; extra == 'dev'
+Requires-Dist: isolate[test]; extra == "dev"
+Requires-Dist: grpcio-tools==1.64.0; extra == "dev"
 Provides-Extra: test
-Requires-Dist: isolate[build] ; extra == 'test'
-Requires-Dist: pytest ; extra == 'test'
-Requires-Dist: cloudpickle >=2.2.0 ; extra == 'test'
-Requires-Dist: dill >=0.3.5.1 ; extra == 'test'
-Requires-Dist: pytest-rerunfailures ; extra == 'test'
+Requires-Dist: isolate[build]; extra == "test"
+Requires-Dist: pytest; extra == "test"
+Requires-Dist: cloudpickle>=2.2.0; extra == "test"
+Requires-Dist: dill>=0.3.5.1; extra == "test"
+Requires-Dist: pytest-rerunfailures; extra == "test"
 
 # Isolate
 

{isolate-0.13.8.dist-info → isolate-0.13.9.dist-info}/RECORD

@@ -1,5 +1,5 @@
 isolate/__init__.py,sha256=uXOKnONs7sXgARNgElwr4_A1sKoA6ACHVEvs3IDiX1M,127
-isolate/_isolate_version.py,sha256=hnyeLxwVkL3Abdy189hhIoA4c9gB_O3-7cct0Y_GGj8,413
+isolate/_isolate_version.py,sha256=I63RKqTcPf9RTCjzPcwnvKN-SWCoNcNIRzHaQnVkYig,413
 isolate/_version.py,sha256=05pXvy-yr5t3I1m9JMn42Ilzpg7fa8IB2J8a3G7t1cU,274
 isolate/logger.py,sha256=SehnK6rPx-HDqQHJ3sKWqhDGmD3fTDGBIkjnNQYnFJU,453
 isolate/logs.py,sha256=R_AHUVYD18z_PhtK_mDWi9Gch79CxmwHY09hUDShtwg,2079
@@ -42,7 +42,7 @@ isolate/connections/ipc/agent.py,sha256=hGlL4x78FhRvMZ4DkVh3dk-EmWQqxHW4LIipgyOk
 isolate/server/__init__.py,sha256=7R3GuWmxuqe0q28rVqETJN9OCrP_-Svjv9h0NR1GFL0,79
 isolate/server/health_server.py,sha256=yN7F1Q28DdX8-Zk3gef7XcQEE25XwlHwzV5GBM75aQM,1249
 isolate/server/interface.py,sha256=nGbjdxrN0p9m1LNdeds8NIoJOwPYW2NM6ktmbhfG4_s,687
-isolate/server/server.py,sha256=D02YtYvGv3K3bYV5DFkb1o5VJpl3RQy37xyXGOqWeQQ,16560
+isolate/server/server.py,sha256=9abptODXQIdswNBQqEYY3yRGz6Jr4NFVs15FiftfRdw,20125
 isolate/server/definitions/__init__.py,sha256=f_Q3pdjMuZrjgNlbM60btFKiB1Vg8cnVyKEbp0RmU0A,572
 isolate/server/definitions/server.proto,sha256=GIou0c_RP4sVDxjdYQzUtSptOfJtvbgsU4MGhL9hsMQ,1372
 isolate/server/definitions/server_pb2.py,sha256=SdbigRwRd1pU1yFwvyrA4vkmwQHPelPkQT7dtggKUlk,3194
@@ -53,9 +53,9 @@ isolate/server/health/health.proto,sha256=wE2_QD0OQAblKkEBG7sALLXEOj1mOLKG-FbC4t
 isolate/server/health/health_pb2.py,sha256=onOdP3M4Tpqhqs2PlGcyfoKe2VVKUEDx5ALeRcObb9A,1899
 isolate/server/health/health_pb2.pyi,sha256=CPyvxvDzra-1d-mBsukaJnscMUDBaqSACvo9LiXlFzo,2416
 isolate/server/health/health_pb2_grpc.py,sha256=XgsULrnRBmYIqvKr8eI7bqs6NIea5A0kkqdOOc2JHBY,5303
-isolate-0.13.8.dist-info/LICENSE,sha256=427vuyirL5scgBLqA9UWcdnxKrtSGc0u_JfUupk6lAA,11359
-isolate-0.13.8.dist-info/METADATA,sha256=nHwIIZULtQ-heNKVZd4IRw89JuKGff-XR-jngc4sBVI,3209
-isolate-0.13.8.dist-info/WHEEL,sha256=y4mX-SOX4fYIkonsAGA5N0Oy-8_gI4FXw5HNI1xqvWg,91
-isolate-0.13.8.dist-info/entry_points.txt,sha256=s3prh2EERaVCbL8R45tfY5WFPZ1TsYOsz305YR7s-Pc,360
-isolate-0.13.8.dist-info/top_level.txt,sha256=W9QJBHcq5WXRkbOXf25bvftzFsOZZN4n1DAatdroZrs,8
-isolate-0.13.8.dist-info/RECORD,,
+isolate-0.13.9.dist-info/LICENSE,sha256=427vuyirL5scgBLqA9UWcdnxKrtSGc0u_JfUupk6lAA,11359
+isolate-0.13.9.dist-info/METADATA,sha256=nSnYPN5KxPVH23bG5jW0ffaaQZXCoTG-pCIHeOrViB4,3191
+isolate-0.13.9.dist-info/WHEEL,sha256=R0nc6qTxuoLk7ShA2_Y-UWkN8ZdfDBG2B6Eqpz2WXbs,91
+isolate-0.13.9.dist-info/entry_points.txt,sha256=s3prh2EERaVCbL8R45tfY5WFPZ1TsYOsz305YR7s-Pc,360
+isolate-0.13.9.dist-info/top_level.txt,sha256=W9QJBHcq5WXRkbOXf25bvftzFsOZZN4n1DAatdroZrs,8
+isolate-0.13.9.dist-info/RECORD,,

{isolate-0.13.8.dist-info → isolate-0.13.9.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (70.2.0)
+Generator: setuptools (72.1.0)
 Root-Is-Purelib: true
 Tag: py3-none-any