Repository: proxystore/aeris

Issue:
Catch `@loop` errors and shutdown agent
If a `@loop` method fails, the error is not propagated because it's in a separate thread. It would be helpful to cause this to shutdown the agent, so the launcher can see the error.

Discussion:

Solution patch:
diff --git a/aeris/agent.py b/aeris/agent.py
index e12a87d..f65d9b3 100644
--- a/aeris/agent.py
+++ b/aeris/agent.py
@@ -65,6 +65,10 @@ class Agent(Generic[BehaviorT]):
         An agent can only be run once. After `shutdown()` is called, later
         operations will raise a `RuntimeError`.
 
+    Note:
+        If any `@loop` method raises an error, the agent will be signaled
+        to shutdown.
+
     Args:
         behavior: Behavior that the agent will exhibit.
         agent_id: Identifier of this agent in a multi-agent system.
@@ -97,7 +101,7 @@ def __init__(
         self._action_pool: ThreadPoolExecutor | None = None
         self._action_futures: dict[ActionRequest, Future[None]] = {}
         self._loop_pool: ThreadPoolExecutor | None = None
-        self._loop_futures: set[Future[None]] = set()
+        self._loop_futures: dict[Future[None], str] = {}
 
         self._multiplexer = MailboxMultiplexer(
             self.agent_id,
@@ -198,6 +202,16 @@ def _request_handler(self, request: RequestMessage) -> None:
         else:
             raise AssertionError('Unreachable.')
 
+    def _loop_callback(self, future: Future[None]) -> None:
+        if future.exception() is not None:
+            name = self._loop_futures[future]
+            logger.warning(
+                'Error in loop %r (signaling shutdown): %r',
+                name,
+                future.exception(),
+            )
+            self.signal_shutdown()
+
     def action(self, action: str, args: Any, kwargs: Any) -> Any:
         """Invoke an action of the agent.
 
@@ -268,12 +282,13 @@ def start(self) -> None:
                 max_workers=len(self._loops) + 1,
             )
 
-            for method in self._loops.values():
+            for name, method in self._loops.items():
                 loop_future = self._loop_pool.submit(method, self._shutdown)
-                self._loop_futures.add(loop_future)
+                self._loop_futures[loop_future] = name
+                loop_future.add_done_callback(self._loop_callback)
 
             listener_future = self._loop_pool.submit(self._multiplexer.listen)
-            self._loop_futures.add(listener_future)
+            self._loop_futures[listener_future] = '_multiplexer.listen'
 
             self._state = _AgentState.RUNNING
 
@@ -297,11 +312,11 @@ def shutdown(self) -> None:
         Raises:
             Exception: Any exceptions raised inside threads.
         """
-        if self._state is _AgentState.SHUTDOWN:
-            return
-
-        logger.debug('Shutting down agent... (%s)', self.agent_id)
         with self._state_lock:
+            if self._state is _AgentState.SHUTDOWN:
+                return
+
+            logger.debug('Shutting down agent... (%s)', self.agent_id)
             self._state = _AgentState.TERMINTATING
             self._shutdown.set()
 
@@ -313,11 +328,9 @@ def shutdown(self) -> None:
             # the mailbox the multiplexer is listening to.
             self._multiplexer.close_mailbox()
 
-            # Wait on all the loops, raising any exceptions.
-            for future in self._loop_futures:
-                future.result()
+            # Shutdown the loop pool after waiting on the loops to exit.
             if self._loop_pool is not None:
-                self._loop_pool.shutdown()
+                self._loop_pool.shutdown(wait=True)
 
             if self.close_exchange:
                 self.exchange.close()
@@ -325,7 +338,11 @@ def shutdown(self) -> None:
             self.behavior.shutdown()
             self._state = _AgentState.SHUTDOWN
 
-        logger.info('Shutdown agent (%s)', self.agent_id)
+            logger.info('Shutdown agent (%s)', self.agent_id)
+
+            # Raise any exceptions from the loop threads as the final step.
+            for future in self._loop_futures:
+                future.result()
 
     def signal_shutdown(self) -> None:
         """Signal that the agent should exit.
diff --git a/aeris/manager.py b/aeris/manager.py
index 90ad951..0197413 100644
--- a/aeris/manager.py
+++ b/aeris/manager.py
@@ -69,7 +69,7 @@ def __init__(self, exchange: Exchange, launcher: Launcher) -> None:
         )
         self._listener_thread.start()
         logger.info(
-            'Initialized manager (%s; %s; %s',
+            'Initialized manager (%s; %s; %s)',
             self._mailbox_id,
             self._exchange,
             self._launcher,


Test patch:
diff --git a/tests/agent_test.py b/tests/agent_test.py
index 3a583fc..7e406d7 100644
--- a/tests/agent_test.py
+++ b/tests/agent_test.py
@@ -74,6 +74,24 @@ def test_agent_shutdown_without_start(exchange: Exchange) -> None:
     assert agent.behavior.shutdown_event.is_set()
 
 
+class LoopFailureBehavior(Behavior):
+    @loop
+    def bad(self, shutdown: threading.Event) -> None:
+        raise RuntimeError('Expected failure in loop.')
+
+
+def test_loop_failure_triggers_shutdown(exchange: Exchange) -> None:
+    agent_id = exchange.create_agent()
+    agent = Agent(LoopFailureBehavior(), agent_id=agent_id, exchange=exchange)
+
+    agent.start()
+    assert agent._shutdown.is_set()
+
+    # Still need to clean up resources...
+    with pytest.raises(RuntimeError, match='Expected failure in loop.'):
+        agent.shutdown()
+
+
 def test_agent_run_in_thread(exchange: Exchange) -> None:
     agent_id = exchange.create_agent()
     agent = Agent(SignalingBehavior(), agent_id=agent_id, exchange=exchange)