isolate 0.14.0__tar.gz → 0.14.2__tar.gz

{isolate-0.14.0 → isolate-0.14.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: isolate
-Version: 0.14.0
+Version: 0.14.2
 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.14.0 → isolate-0.14.2}/src/isolate/_isolate_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.14.0'
-__version_tuple__ = version_tuple = (0, 14, 0)
+__version__ = version = '0.14.2'
+__version_tuple__ = version_tuple = (0, 14, 2)

{isolate-0.14.0 → isolate-0.14.2}/src/isolate/backends/common.py

@@ -117,6 +117,7 @@ def _io_observer(
                 return  # Nothing to read
 
             for line in raw_data.splitlines():
+                # TODO: parse the lines to include `extra={...}` added by the logger?
                 hook(line)
 
     def _reader():

{isolate-0.14.0 → isolate-0.14.2}/src/isolate/logger.py

@@ -9,6 +9,8 @@ from isolate.logs import LogLevel, LogSource
 # 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
 
@@ -18,6 +20,7 @@ class IsolateLogger:
             "level": level.name,
             "message": message,
             **self.log_labels,
+            **self.extra_labels,
         }
         print(json.dumps(record))
 

{isolate-0.14.0 → isolate-0.14.2}/src/isolate/server/server.py

@@ -36,6 +36,7 @@ from isolate.server.health_server import HealthServicer
 from isolate.server.interface import from_grpc, to_grpc
 
 EMPTY_MESSAGE_INTERVAL = float(os.getenv("ISOLATE_EMPTY_MESSAGE_INTERVAL", "600"))
+SKIP_EMPTY_LOGS = os.getenv("ISOLATE_SKIP_EMPTY_LOGS") == "1"
 MAX_GRPC_WAIT_TIMEOUT = float(os.getenv("ISOLATE_MAX_GRPC_WAIT_TIMEOUT", "10.0"))
 
 # Whether to inherit all the packages from the current environment or not.
@@ -366,9 +367,7 @@ class IsolateServicer(definitions.IsolateServicer):
 
         task = self.background_tasks[request.task_id]
 
-        task.logger = IsolateLogger.with_env_expanded(
-            dict(request.metadata.logger_labels)
-        )
+        task.logger.extra_labels = dict(request.metadata.logger_labels)
 
         return definitions.SetMetadataResponse()
 
@@ -378,13 +377,19 @@ class IsolateServicer(definitions.IsolateServicer):
         context: ServicerContext,
     ) -> Iterator[definitions.PartialRunResult]:
         try:
-            yield from self._run_task(RunTask(request=request))
+            # HACK: we can support only one task at a time for Run
+            # TODO: move away from this when we use submit for env-aware tasks
+            task = RunTask(request=request)
+            self.background_tasks["RUN"] = task
+            yield from self._run_task(task)
         except GRPCException as exc:
             return self.abort_with_msg(
                 exc.message,
                 context,
                 code=exc.code,
             )
+        finally:
+            self.background_tasks.pop("RUN", None)
 
     def List(
         self,
@@ -464,7 +469,8 @@ class IsolateServicer(definitions.IsolateServicer):
         return None
 
     def cancel_tasks(self):
-        for task in self.background_tasks.values():
+        tasks_copy = self.background_tasks.copy()
+        for task in tasks_copy.values():
             task.cancel()
 
 
@@ -484,8 +490,9 @@ class LogHandler:
     task: RunTask
 
     def handle(self, log: Log) -> None:
-        self.task.logger.log(log.level, log.message, source=log.source)
-        self._add_log_to_queue(log)
+        if not SKIP_EMPTY_LOGS or log.message.strip():
+            self.task.logger.log(log.level, log.message, source=log.source)
+            self._add_log_to_queue(log)
 
     def _add_log_to_queue(self, log: Log) -> None:
         grpc_log = cast(definitions.Log, to_grpc(log))
@@ -534,6 +541,16 @@ class SingleTaskInterceptor(ServerBoundInterceptor):
     """Sets server to terminate after the first Submit/Run task."""
 
     _done: bool = False
+    _task_id: str | None = None
+
+    def __init__(self):
+        def terminate(request: Any, context: grpc.ServicerContext) -> Any:
+            context.abort(
+                grpc.StatusCode.RESOURCE_EXHAUSTED,
+                "Server has already served one Run/Submit task.",
+            )
+
+        self._terminator = grpc.unary_unary_rpc_method_handler(terminate)
 
     def intercept_service(self, continuation, handler_call_details):
         handler = continuation(handler_call_details)
@@ -542,29 +559,62 @@ class SingleTaskInterceptor(ServerBoundInterceptor):
         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:
+        if not is_new_task:
             # Let other requests like List/Cancel/etc pass through
-            return continuation(handler_call_details)
+            return handler
+
+        if self._done:
+            # Fail the request if the server has already served or is serving
+            # a Run/Submit task.
+            return self._terminator
+
+        self._done = True
 
         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:
+            def _wrapper(request: Any, context: grpc.ServicerContext) -> Any:
+                def termination() -> None:
+                    if is_run:
+                        print("Stopping server since run is finished")
+                        # Stop the server after the Run task is finished
+                        self.server.stop(grace=0.1)
+
+                    elif is_submit:
+                        # Wait until the task_id is assigned
+                        while self._task_id is None:
                             time.sleep(0.1)
-                    self.server.stop(grace=0.1)
 
-                context.add_callback(_stop)
-                return method_impl(request, context)
+                        # Get the task from the background tasks
+                        task = self.servicer.background_tasks.get(self._task_id)
+
+                        if task is not None:
+                            # Wait until the task future is assigned
+                            tries = 0
+                            while task.future is None:
+                                time.sleep(0.1)
+                                tries += 1
+                                if tries > 100:
+                                    raise RuntimeError(
+                                        "Task future was not assigned in time."
+                                    )
+
+                            def _stop(*args):
+                                # Small sleep to make sure the cancellation is processed
+                                time.sleep(0.1)
+                                print("Stopping server since the task is finished")
+                                self.server.stop(grace=0.1)
+
+                            # Add a callback which will stop the server
+                            # after the task is finished
+                            task.future.add_done_callback(_stop)
+
+                context.add_callback(termination)
+                res = method_impl(request, context)
+
+                if is_submit:
+                    self._task_id = cast(definitions.SubmitResponse, res).task_id
+
+                return res
 
             return _wrapper
 
@@ -598,7 +648,7 @@ def main(argv: list[str] | None = None) -> None:
     server = grpc.server(
         futures.ThreadPoolExecutor(max_workers=options.num_workers),
         options=get_default_options(),
-        interceptors=interceptors,
+        interceptors=interceptors,  # type: ignore
     )
 
     for interceptor in interceptors:

{isolate-0.14.0 → isolate-0.14.2}/src/isolate.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: isolate
-Version: 0.14.0
+Version: 0.14.2
 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.14.0 → isolate-0.14.2}/tests/test_server.py

@@ -5,7 +5,7 @@ from contextlib import contextmanager
 from dataclasses import dataclass
 from functools import partial
 from pathlib import Path
-from typing import Any, List, Optional, cast
+from typing import Any, Iterator, List, Optional, cast
 
 import grpc
 import pytest
@@ -15,7 +15,12 @@ from isolate.logs import Log, LogLevel, LogSource
 from isolate.server import definitions, health
 from isolate.server.health_server import HealthServicer
 from isolate.server.interface import from_grpc, to_serialized_object
-from isolate.server.server import BridgeManager, IsolateServicer
+from isolate.server.server import (
+    BridgeManager,
+    IsolateServicer,
+    ServerBoundInterceptor,
+    SingleTaskInterceptor,
+)
 
 REPO_DIR = Path(__file__).parent.parent
 assert (
@@ -34,14 +39,32 @@ class Stubs:
     health_stub: health.HealthStub
 
 
+@pytest.fixture
+def interceptors():
+    return []
+
+
 @contextmanager
-def make_server(tmp_path):
+def make_server(
+    tmp_path: Path, interceptors: Optional[List[ServerBoundInterceptor]] = None
+) -> Iterator[Stubs]:
+    interceptors = interceptors or []
     server = grpc.server(
-        futures.ThreadPoolExecutor(max_workers=1), options=get_default_options()
+        futures.ThreadPoolExecutor(max_workers=1),
+        options=get_default_options(),
+        interceptors=interceptors,  # type: ignore
     )
+
+    for interceptor in interceptors:
+        interceptor.register_server(server)
+
     test_settings = IsolateSettings(cache_dir=tmp_path / "cache")
     with BridgeManager() as bridge:
         servicer = IsolateServicer(bridge, test_settings)
+
+        for interceptor in interceptors:
+            interceptor.register_servicer(servicer)
+
         definitions.register_isolate(servicer, server)
         health.register_health(HealthServicer(), server)
         host, port = "localhost", server.add_insecure_port("[::]:0")
@@ -69,14 +92,14 @@ def make_server(tmp_path):
 
 
 @pytest.fixture
-def stub(tmp_path):
-    with make_server(tmp_path) as stubs:
+def stub(tmp_path, interceptors):
+    with make_server(tmp_path, interceptors) as stubs:
         yield stubs.isolate_stub
 
 
 @pytest.fixture
-def health_stub(tmp_path):
-    with make_server(tmp_path) as stubs:
+def health_stub(tmp_path, interceptors):
+    with make_server(tmp_path, interceptors) as stubs:
         yield stubs.health_stub
 
 
@@ -719,3 +742,79 @@ def test_server_submit_server(
     stub.Cancel(definitions.CancelRequest(task_id=task_id))
 
     assert not list(stub.List(definitions.ListRequest()).tasks)
+
+
+@pytest.mark.parametrize(
+    "interceptors",
+    [
+        [SingleTaskInterceptor()],
+    ],
+)
+def test_server_single_use_submit(
+    stub: definitions.IsolateStub,
+    monkeypatch: Any,
+) -> None:
+    import time
+
+    inherit_from_local(monkeypatch)
+
+    request = definitions.SubmitRequest(function=prepare_request(myserver))
+    task_id = stub.Submit(request).task_id
+
+    tasks = [task.task_id for task in stub.List(definitions.ListRequest()).tasks]
+    assert task_id in tasks
+
+    # Now try to Submit again
+    with pytest.raises(grpc.RpcError) as exc_info:
+        stub.Submit(request)
+    assert exc_info.value.code() == grpc.StatusCode.RESOURCE_EXHAUSTED
+
+    # And try to Run a task
+    with pytest.raises(grpc.RpcError) as exc_info:
+        run_request(stub, prepare_request(myserver))
+    assert exc_info.value.code() == grpc.StatusCode.RESOURCE_EXHAUSTED
+
+    stub.Cancel(definitions.CancelRequest(task_id=task_id))
+    time.sleep(1)
+
+    with pytest.raises(grpc.RpcError) as exc_info:
+        stub.List(definitions.ListRequest())
+
+    # Server should be shutting down
+    assert exc_info.value.code() == grpc.StatusCode.UNAVAILABLE
+
+
+@pytest.mark.parametrize(
+    "interceptors",
+    [
+        [SingleTaskInterceptor()],
+    ],
+)
+def test_server_single_use_run(
+    stub: definitions.IsolateStub,
+    monkeypatch: Any,
+) -> None:
+    import time
+
+    inherit_from_local(monkeypatch)
+
+    run_function(stub, check_machine)
+    time.sleep(1)
+
+    # Now try to Submit again
+    with pytest.raises(grpc.RpcError) as exc_info:
+        submit_request = definitions.SubmitRequest(function=prepare_request(myserver))
+        stub.Submit(submit_request)
+
+    assert exc_info.value.code() == grpc.StatusCode.UNAVAILABLE
+
+    # And try to Run a task
+    with pytest.raises(grpc.RpcError) as exc_info:
+        run_function(stub, check_machine)
+
+    assert exc_info.value.code() == grpc.StatusCode.UNAVAILABLE
+
+    with pytest.raises(grpc.RpcError) as exc_info:
+        stub.List(definitions.ListRequest())
+
+    assert exc_info.value.code() == grpc.StatusCode.UNAVAILABLE