feat: add agent idle timeout (#189)

* feat: add agent idle timeout

* feat: handle SIGCHLD in parent server

* fix: use events and default idle_timeout 0

* fix: add more tests

* fix: prevent blocking event loop

* fix: test

* fix: test fd

* fix: flaky test_server_single_use_submit

Author: hanchchch

pyproject.toml

@@ -48,6 +48,7 @@ test = [
     "cloudpickle>=2.2.0",
     "dill>=0.3.5.1",
     "flaky",
+    "psutil>=7.2.1",
 ]
 dev = [
     "isolate[test]",

src/isolate/connections/grpc/_base.py

@@ -153,6 +153,12 @@ def abort_agent(self) -> None:
             except Exception as exc:
                 print(f"Failed to shutdown the agent process gracefully: {exc}")
                 self._process.kill()
+            self._process = None
+
+    def is_alive(self) -> bool:
+        if self._process is None:
+            return False
+        return self._process.poll() is None
 
     def get_python_cmd(
         self,

src/isolate/connections/grpc/agent.py

@@ -39,6 +39,8 @@
 from isolate.connections.grpc.configuration import get_default_options
 from isolate.connections.grpc.interface import from_grpc
 
+IDLE_TIMEOUT_SECONDS = int(os.getenv("ISOLATE_AGENT_IDLE_TIMEOUT_SECONDS", "0"))
+
 
 @dataclass
 class AbortException(Exception):
@@ -52,17 +54,64 @@ 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)
+        self._idle_timeout_seconds = IDLE_TIMEOUT_SECONDS
+        self._is_running = asyncio.Event()
+        self._is_idle = asyncio.Event()
+        self._is_idle.set()
 
-        def handle_termination(*args):
-            self.log("Termination signal received, shutting down...")
+        def handle_sigint(*args):
+            self.log("SIGINT signal received, shutting down...")
             signal.raise_signal(signal.SIGTERM)
 
-        signal.signal(signal.SIGINT, handle_termination)
+        signal.signal(signal.SIGINT, handle_sigint)
+
+    async def wait_for_idle_timeout(self) -> None:
+        while True:
+            # print(f"Hello, world! {self._idle_timeout_seconds}")
+            # wait for the agent to be idle
+            await self._is_idle.wait()
+
+            # idle timeout disabled
+            if self._idle_timeout_seconds <= 0:
+                # prevent blocking the event loop
+                await asyncio.sleep(0.1)
+                continue
+
+            try:
+                # wait for the agent to be running
+                await asyncio.wait_for(
+                    self._is_running.wait(), timeout=self._idle_timeout_seconds
+                )
+            except asyncio.TimeoutError:
+                self.log(
+                    f"Idle timeout {self._idle_timeout_seconds} seconds exceeded, shutting down..."  # noqa: E501
+                )
+                # This kills the agent itself, however it will remain as a zombie state
+                # unless the parent process (server) properly handles the SIGCHLD.
+                signal.raise_signal(signal.SIGTERM)
+                break
+            except asyncio.CancelledError:
+                # Cancelled when the server is shutting down
+                break
 
     async def Run(
         self,
         request: definitions.FunctionCall,
         context: aio.ServicerContext,
+    ) -> AsyncIterator[PartialRunResult]:
+        self._is_idle.clear()
+        self._is_running.set()
+        try:
+            async for result in self._Run(request, context):
+                yield result
+        finally:
+            self._is_running.clear()
+            self._is_idle.set()
+
+    async def _Run(
+        self,
+        request: definitions.FunctionCall,
+        context: aio.ServicerContext,
     ) -> AsyncIterator[PartialRunResult]:
         self.log(f"A connection has been established: {context.peer()}!")
         server_version = os.getenv("ISOLATE_SERVER_VERSION") or "unknown"
@@ -267,7 +316,18 @@ async def run_agent(address: str, log_fd: int | None = None) -> int:
     definitions.register_agent(servicer, server)
 
     await server.start()
-    await server.wait_for_termination()
+
+    _, pending = await asyncio.wait(
+        [
+            asyncio.create_task(server.wait_for_termination()),
+            asyncio.create_task(servicer.wait_for_idle_timeout()),
+        ],
+        return_when=asyncio.FIRST_COMPLETED,
+    )
+    for task in pending:
+        print(f"Cancelling task: {task}")
+        task.cancel()
+
     return 0
 
 

src/isolate/server/server.py

@@ -182,6 +182,14 @@ def __exit__(self, *exc_info: Any) -> None:
             for agent in agents:
                 agent.terminate()
 
+    def abort_unreachable_agents(self) -> None:
+        for agents in self._agents.values():
+            for agent in agents:
+                connection = agent._connection
+                if connection is not None and not connection.is_alive():
+                    connection.abort_agent()
+                    # maybe restart the agent?
+
 
 @dataclass
 class RunTask:
@@ -650,7 +658,7 @@ def termination() -> None:
 
                             def _stop(*args):
                                 # Small sleep to make sure the cancellation is processed
-                                time.sleep(0.1)
+                                time.sleep(0.3)
                                 print("Stopping server since the task is finished")
                                 self.servicer.shutdown()
                                 self.server.stop(grace=0.1)
@@ -720,8 +728,13 @@ def handle_termination(*args):
             servicer.shutdown()
             server.stop(grace=0.1)
 
+        def handle_child_termination(*args):
+            print("Child termination signal received, aborting unreachable agents...")
+            bridge_manager.abort_unreachable_agents()
+
         signal.signal(signal.SIGINT, handle_termination)
         signal.signal(signal.SIGTERM, handle_termination)
+        signal.signal(signal.SIGCHLD, handle_child_termination)
 
         server.add_insecure_port(f"[::]:{options.port}")
         print(f"Started listening at {options.host}:{options.port}")

tests/test_shutdown.py

@@ -7,10 +7,14 @@
 import sys
 import threading
 import time
+from typing import Iterator, Tuple
 from unittest.mock import Mock
 
 import grpc
+import psutil
 import pytest
+from isolate.connections.grpc.definitions import FunctionCall
+from isolate.connections.grpc.definitions.agent_pb2_grpc import AgentStub
 from isolate.server.definitions.server_pb2 import BoundFunction, EnvironmentDefinition
 from isolate.server.definitions.server_pb2_grpc import IsolateStub
 from isolate.server.interface import to_serialized_object
@@ -44,13 +48,23 @@ def servicer():
 
 
 @pytest.fixture
-def isolate_server_subprocess(monkeypatch):
+def single_use():
+    return True
+
+
+@pytest.fixture
+def idle_timeout_seconds():
+    return 0
+
+
+@pytest.fixture
+def isolate_server_subprocess(
+    single_use: bool, idle_timeout_seconds: int
+) -> Iterator[Tuple[subprocess.Popen, int]]:
     """Set up a gRPC server with the IsolateServicer for testing."""
     # Find a free port
     import socket
 
-    monkeypatch.setenv("ISOLATE_SHUTDOWN_GRACE_PERIOD", "2")
-
     # Bind only to the loopback interface to avoid exposing the socket on all interfaces
     with socket.socket() as s:
         s.bind(("127.0.0.1", 0))
@@ -61,10 +75,14 @@ def isolate_server_subprocess(monkeypatch):
             sys.executable,
             "-m",
             "isolate.server.server",
-            "--single-use",
+            *(["--single-use"] if single_use else []),
             "--port",
             str(port),
-        ]
+        ],
+        env={
+            "ISOLATE_SHUTDOWN_GRACE_PERIOD": "2",
+            "ISOLATE_AGENT_IDLE_TIMEOUT_SECONDS": str(idle_timeout_seconds),
+        },
     )
 
     time.sleep(5)  # Wait for server to start
@@ -76,7 +94,50 @@ def isolate_server_subprocess(monkeypatch):
         process.wait(timeout=10)
 
 
-def consume_responses(responses):
+@pytest.fixture
+def isolate_agent_subprocess(
+    idle_timeout_seconds: int,
+) -> Iterator[Tuple[subprocess.Popen, int]]:
+    """Set up a gRPC server with the IsolateServicer for testing."""
+    # Find a free port
+    import socket
+
+    # Bind only to the loopback interface to avoid exposing the socket on all interfaces
+    with socket.socket() as s:
+        s.bind(("127.0.0.1", 0))
+        port = s.getsockname()[1]
+
+    # Use /dev/null for log output since pytest may capture stdout
+    # (making fileno() fail with "Bad file descriptor")
+    log_file = open(os.devnull, "w")
+
+    process = subprocess.Popen(
+        [
+            sys.executable,
+            "-m",
+            "isolate.connections.grpc.agent",
+            f"localhost:{port}",
+            "--log-fd",
+            str(log_file.fileno()),
+        ],
+        env={
+            "ISOLATE_AGENT_IDLE_TIMEOUT_SECONDS": str(idle_timeout_seconds),
+        },
+        pass_fds=(log_file.fileno(),),
+    )
+
+    time.sleep(1)  # Wait for server to start
+    try:
+        yield process, port
+    finally:
+        # Cleanup
+        if process.poll() is None:
+            process.terminate()
+            process.wait(timeout=10)
+        log_file.close()
+
+
+def consume_responses(responses: Iterator, wait: bool = False) -> None:
     def _consume():
         try:
             for response in responses:
@@ -87,6 +148,8 @@ def _consume():
 
     response_thread = threading.Thread(target=_consume, daemon=True)
     response_thread.start()
+    if wait:
+        response_thread.join()
 
 
 def test_shutdown_with_terminate(servicer):
@@ -180,5 +243,103 @@ def handle_term(signum, frame):
     ), "Function should have received SIGTERM and created the file"
 
 
-if __name__ == "__main__":
-    pytest.main([__file__, "-v"])
+@pytest.mark.parametrize(
+    "idle_timeout_seconds",
+    [0, 2],
+)
+def test_idle_timeout_no_request(isolate_agent_subprocess, idle_timeout_seconds):
+    process, port = isolate_agent_subprocess
+
+    p = psutil.Process(process.pid)
+    for _ in range(10):
+        if p.is_running():
+            break
+        time.sleep(1)
+    else:
+        assert False, "Agent should be running"
+
+    # Wait for the idle timeout to trigger
+    try:
+        p.wait(timeout=5)
+        terminated = True
+    except psutil.TimeoutExpired:
+        terminated = False
+
+    if idle_timeout_seconds == 0:
+        assert not terminated, "Agent should not have terminated"
+    else:
+        assert terminated, "Agent should have terminated after idle timeout"
+
+
+@pytest.mark.parametrize(
+    "idle_timeout_seconds",
+    [0, 2],
+)
+def test_idle_timeout(isolate_agent_subprocess, idle_timeout_seconds):
+    process, port = isolate_agent_subprocess
+
+    p = psutil.Process(process.pid)
+    for _ in range(10):
+        if p.is_running():
+            break
+        time.sleep(1)
+    else:
+        assert False, "Agent should be running"
+
+    channel = grpc.insecure_channel(f"localhost:{port}")
+    stub = AgentStub(channel)
+
+    def fn():
+        import time
+
+        time.sleep(3)  # longer than the idle timeout
+        print("Function finished")
+
+    responses = stub.Run(FunctionCall(function=to_serialized_object(fn, method="dill")))
+    consume_responses(responses, wait=True)
+
+    # Wait for the idle timeout to trigger
+    try:
+        p.wait(timeout=5)
+        terminated = True
+    except psutil.TimeoutExpired:
+        terminated = False
+
+    if idle_timeout_seconds == 0:
+        assert not terminated, "Agent should not have terminated"
+    else:
+        assert terminated, "Agent should have terminated after idle timeout"
+
+
+@pytest.mark.parametrize(
+    ["single_use", "idle_timeout_seconds"],
+    [(False, 2)],  # to prevent the server from shutting down automatically
+)
+def test_idle_timeout_server_handle(isolate_server_subprocess):
+    process, port = isolate_server_subprocess
+    channel = grpc.insecure_channel(f"localhost:{port}")
+    stub = IsolateStub(channel)
+
+    def fn():
+        import time
+
+        time.sleep(5)  # longer than the idle timeout
+        print("Function finished")
+
+    responses = stub.Run(create_run_request(fn))
+    consume_responses(responses, wait=True)
+
+    # Send the first request to start the agent
+    p = psutil.Process(process.pid)
+    assert len(p.children()) == 1, "Server should have one agent process"
+
+    # Wait for the idle timeout to trigger
+    time.sleep(3)
+    assert (
+        len(p.children()) == 0
+    ), "Agent process should have terminated after idle timeout"
+
+    # Server should be able to handle a new request after the idle timeout
+    responses = stub.Run(create_run_request(fn))
+    consume_responses(responses, wait=True)
+    assert len(p.children()) == 1, "Server should have one agent process"