fix: propagate SIGTERM to agent (#184)

* feat: Allow agent to handle signals for graceful termination

* Add tests, and fix a few agent termination edge cases

* Fix flaky test

* Allow nested event loops

* Revert "Allow nested event loops"

This reverts commit 42d04ec.

* fix: log traceback if raised

* fix: handle coroutine function

* fix: test

* fix: test

* fix: run in a thread pool if _run_as_main_thread not set

* fix: add signal handler

* fix: redundant signal handler

* fix: terminate process before closing connection on abort

* fix: log sigterm

* fix: duplicated PROCESS_SHUTDOWN_TIMEOUT_SECONDS

* fix: set _run_as_main_thread on test

* Potential fix for code scanning alert no. 4: Binding a socket to all network interfaces

Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com>

* fix: logs test ignore warning

* docs: add comments

* refactor: rename run_as to run_on

* fix: add tests

---------

Co-authored-by: Sanjay Raveendran <sanjayprabhu@gmail.com>
Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com>

Author: 3 people

src/isolate/connections/grpc/_base.py

@@ -1,4 +1,6 @@
+import os
 import socket
+import subprocess
 from contextlib import contextmanager
 from dataclasses import dataclass
 from pathlib import Path
@@ -23,6 +25,11 @@ 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."""
@@ -32,6 +39,9 @@ def start_agent(self) -> ContextManager[Tuple[str, grpc.ChannelCredentials]]:
         the required credentials to connect to it."""
         raise NotImplementedError
 
+    def abort_agent(self) -> None:
+        raise NotImplementedError
+
     @contextmanager
     def _establish_bridge(
         self,
@@ -113,6 +123,8 @@ def run(
 
 
 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]:
@@ -123,14 +135,24 @@ def find_free_port() -> Tuple[str, int]:
 
         host, port = find_free_port()
         address = f"{host}:{port}"
-        process = None
+        self._process = None
         try:
             with self.start_process(address) as process:
+                self._process = process
                 yield address, grpc.local_channel_credentials()
         finally:
-            if process is not None:
-                # TODO: should we check the status code here?
-                process.terminate()
+            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,

src/isolate/connections/grpc/agent.py

@@ -10,20 +10,23 @@
 
 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,
-    Iterator,
 )
 
-import grpc
-from grpc import ServicerContext, StatusCode
+from grpc import StatusCode, aio, local_server_credentials
+
+from isolate.connections.grpc.definitions import PartialRunResult
 
 try:
     from isolate import __version__ as agent_version
@@ -48,12 +51,19 @@ def __init__(self, log_fd: int | None = None):
 
         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)
 
-    def Run(
+        signal.signal(signal.SIGINT, handle_termination)
+
+    async def Run(
         self,
         request: definitions.FunctionCall,
-        context: ServicerContext,
-    ) -> Iterator[definitions.PartialRunResult]:
+        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}")
@@ -70,7 +80,7 @@ def Run(
                         result,
                         was_it_raised,
                         stringized_tb,
-                    ) = self.execute_function(
+                    ) = await self.execute_function(
                         request.setup_func,
                         "setup",
                     )
@@ -87,15 +97,16 @@ def Run(
                         )
                         raise AbortException("The setup function has thrown an error.")
                 except AbortException as exc:
-                    return self.abort_with_msg(context, exc.message)
+                    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 = self.execute_function(
+            result, was_it_raised, stringized_tb = await self.execute_function(
                 request.function,
                 "function",
                 extra_args=extra_args,
@@ -107,9 +118,10 @@ def Run(
                 stringized_tb,
             )
         except AbortException as exc:
-            return self.abort_with_msg(context, exc.message)
+            self.abort_with_msg(context, exc.message)
+            return
 
-    def execute_function(
+    async def execute_function(
         self,
         function: definitions.SerializedObject,
         function_kind: str,
@@ -143,14 +155,34 @@ def execute_function(
         was_it_raised = False
         stringized_tb = None
         try:
-            result = function(*extra_args)
+            # 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))
 
-        self.log(f"Completed the execution of the {function_kind} function.")
+        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(
@@ -195,7 +227,7 @@ def log(self, message: str) -> None:
 
     def abort_with_msg(
         self,
-        context: ServicerContext,
+        context: aio.ServicerContext,
         message: str,
         *,
         code: StatusCode = StatusCode.INVALID_ARGUMENT,
@@ -205,23 +237,26 @@ def abort_with_msg(
         return None
 
 
-def create_server(address: str) -> grpc.Server:
+def create_server(address: str) -> aio.Server:
     """Create a new (temporary) gRPC server listening on the given
     address."""
-    server = grpc.server(
-        futures.ThreadPoolExecutor(max_workers=1),
-        maximum_concurrent_rpcs=1,
+    # 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 = grpc.local_server_credentials()
+    server_credentials = local_server_credentials()
     server.add_secure_port(address, server_credentials)
     return server
 
 
-def run_agent(address: str, log_fd: int | None = None) -> int:
+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)
@@ -231,19 +266,19 @@ def run_agent(address: str, log_fd: int | None = None) -> int:
     # not have any global side effects.
     definitions.register_agent(servicer, server)
 
-    server.start()
-    server.wait_for_termination()
+    await server.start()
+    await server.wait_for_termination()
     return 0
 
 
-def main() -> int:
+async def main() -> int:
     parser = ArgumentParser()
     parser.add_argument("address", type=str)
     parser.add_argument("--log-fd", type=int)
 
     options = parser.parse_args()
-    return run_agent(options.address, log_fd=options.log_fd)
+    return await run_agent(options.address, log_fd=options.log_fd)
 
 
 if __name__ == "__main__":
-    main()
+    asyncio.run(main())