isolate 0.14.0__py3-none-any.whl → 0.14.2__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.14.0'
-__version_tuple__ = version_tuple = (0, 14, 0)
+__version__ = version = '0.14.2'
+__version_tuple__ = version_tuple = (0, 14, 2)

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

@@ -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.dist-info → isolate-0.14.2.dist-info}/RECORD

@@ -1,13 +1,13 @@
 isolate/__init__.py,sha256=uXOKnONs7sXgARNgElwr4_A1sKoA6ACHVEvs3IDiX1M,127
-isolate/_isolate_version.py,sha256=9XbM6DRatFmjNkeSNxF4a6enpQeo8nPyAy9RTPPwPpY,413
+isolate/_isolate_version.py,sha256=rlRT2sdhOGuXEsCM8RZA3pLtxwYV_87Qc0n3HT-mx2U,413
 isolate/_version.py,sha256=05pXvy-yr5t3I1m9JMn42Ilzpg7fa8IB2J8a3G7t1cU,274
-isolate/logger.py,sha256=d9mZyTOtplMZeQSNgRNiXhOPwJ0be2B1nGGaQx09C3g,1450
+isolate/logger.py,sha256=gN_HzbvNkfloXw5iom6M7jL8e8VOgrMSi1Wp5I1eoX0,1522
 isolate/logs.py,sha256=R_AHUVYD18z_PhtK_mDWi9Gch79CxmwHY09hUDShtwg,2079
 isolate/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 isolate/registry.py,sha256=hpzv4HI7iihG5I7i5r8Pb257ibhEKY18xQcG-w1-BgI,1590
 isolate/backends/__init__.py,sha256=LLrSM7UnDFW8tIT5oYlE1wVJrxKcaj_v7cFwvTjQTlc,119
 isolate/backends/_base.py,sha256=Kt5pkhDzXZblq4vxaM3DQTo-Bj-7pIRZFlqJR7OfejY,4112
-isolate/backends/common.py,sha256=Zx0HXnBX_jlOLpFNJzY4ue8NcW-kktqo_WZOJmPSjvI,8481
+isolate/backends/common.py,sha256=GzmgzIkCY90elMH_7lpj-LwRNu-GLia2j4PrS6zNHwA,8567
 isolate/backends/conda.py,sha256=S5q5bdY787AMTck2iMGtwu-LHMH4a1qCIjNHDKTkqok,7649
 isolate/backends/container.py,sha256=MCQJbcmQvRUS-tTgTW_pKYBMKwSJO2KZsLeaBMXpPC0,1645
 isolate/backends/local.py,sha256=woxe4dmXuEHxWKsGNndoRA1_sP6yG-dg6tlFZni0mZc,1360
@@ -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=THjWTLFxy9Y9atlHtenjtN7ltSGOLAVUAj9lDzl0qUw,21197
+isolate/server/server.py,sha256=41Y_3dtUmeJR8RuQF2VMC5x_kM_k0w1hyzPFl5uW5Do,23477
 isolate/server/definitions/__init__.py,sha256=f_Q3pdjMuZrjgNlbM60btFKiB1Vg8cnVyKEbp0RmU0A,572
 isolate/server/definitions/server.proto,sha256=UihlFbYG8fbwm0IUfKDRH1vNkopzP3C-wplXXcAO1c8,1761
 isolate/server/definitions/server_pb2.py,sha256=TYBJC_z_dNf2J6FgHukY9mDyc3WBt9EPQOVFd_ayQNc,4304
@@ -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=AK-DPCpJzoYhU6DydD856c0Ywx84x6k-Cs4m6HpNv5A,2459
 isolate/server/health/health_pb2_grpc.py,sha256=XgsULrnRBmYIqvKr8eI7bqs6NIea5A0kkqdOOc2JHBY,5303
-isolate-0.14.0.dist-info/LICENSE,sha256=427vuyirL5scgBLqA9UWcdnxKrtSGc0u_JfUupk6lAA,11359
-isolate-0.14.0.dist-info/METADATA,sha256=STY6fLw2onLEsBbyyseZBKbWamPsAYrps5b9PvXC_MM,3191
-isolate-0.14.0.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
-isolate-0.14.0.dist-info/entry_points.txt,sha256=s3prh2EERaVCbL8R45tfY5WFPZ1TsYOsz305YR7s-Pc,360
-isolate-0.14.0.dist-info/top_level.txt,sha256=W9QJBHcq5WXRkbOXf25bvftzFsOZZN4n1DAatdroZrs,8
-isolate-0.14.0.dist-info/RECORD,,
+isolate-0.14.2.dist-info/LICENSE,sha256=427vuyirL5scgBLqA9UWcdnxKrtSGc0u_JfUupk6lAA,11359
+isolate-0.14.2.dist-info/METADATA,sha256=kFsnmpSczOcMgcM6idLXEuVDLWyOD8oG7MR2hV2x8KI,3191
+isolate-0.14.2.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
+isolate-0.14.2.dist-info/entry_points.txt,sha256=s3prh2EERaVCbL8R45tfY5WFPZ1TsYOsz305YR7s-Pc,360
+isolate-0.14.2.dist-info/top_level.txt,sha256=W9QJBHcq5WXRkbOXf25bvftzFsOZZN4n1DAatdroZrs,8
+isolate-0.14.2.dist-info/RECORD,,