Repository: proxystore/aeris

Issue:
Agent restarting and migration
Agents are currently one-time use, but eventually we want to support automatic restarting and potentially manual migration.

Discussion:

Solution patch:
diff --git a/aeris/agent.py b/aeris/agent.py
index 090468d..af869d5 100644
--- a/aeris/agent.py
+++ b/aeris/agent.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import dataclasses
 import enum
 import logging
 import sys
@@ -42,19 +43,29 @@ class _AgentState(enum.Enum):
     SHUTDOWN = 'shutdown'
 
 
+@dataclasses.dataclass(frozen=True)
+class AgentRunConfig:
+    """Agent run configuration."""
+
+    close_exchange_on_exit: bool = True
+    max_action_concurrency: int | None = None
+    terminate_on_error: bool = True
+    terminate_on_exit: bool = True
+
+
 # Helper for Agent.__reduce__ which cannot handle the keyword arguments
 # of the Agent constructor.
 def _agent_trampoline(
     behavior: BehaviorT,
     agent_id: AgentIdentifier,
     exchange: Exchange,
-    close_exchange: bool,
+    config: AgentRunConfig,
 ) -> Agent[BehaviorT]:
     return Agent(
         behavior,
         agent_id=agent_id,
         exchange=exchange,
-        close_exchange=close_exchange,
+        config=config,
     )
 
 
@@ -78,8 +89,7 @@ class Agent(Generic[BehaviorT]):
         agent_id: Identifier of this agent in a multi-agent system.
         exchange: Message exchange of multi-agent system. The agent will close
             the exchange when it finished running.
-        close_exchange: Close the `exchange` object when the agent is
-            shutdown.
+        config: Agent execution parameters.
     """
 
     def __init__(
@@ -88,17 +98,18 @@ def __init__(
         *,
         agent_id: AgentIdentifier,
         exchange: Exchange,
-        close_exchange: bool = False,
+        config: AgentRunConfig | None = None,
     ) -> None:
         self.agent_id = agent_id
         self.behavior = behavior
         self.exchange = exchange
-        self.close_exchange = close_exchange
+        self.config = config if config is not None else AgentRunConfig()
 
         self._actions = behavior.behavior_actions()
         self._loops = behavior.behavior_loops()
 
         self._shutdown = threading.Event()
+        self._expected_shutdown = True
         self._state_lock = threading.Lock()
         self._state = _AgentState.INITIALIZED
 
@@ -132,7 +143,7 @@ def __str__(self) -> str:
     def __reduce__(self) -> Any:
         return (
             _agent_trampoline,
-            (self.behavior, self.agent_id, self.exchange, self.close_exchange),
+            (self.behavior, self.agent_id, self.exchange, self.config),
         )
 
     def _bind_handle(self, handle: Handle[BehaviorT]) -> Handle[BehaviorT]:
@@ -208,7 +219,7 @@ def _loop_callback(self, future: Future[None]) -> None:
                 name,
                 future.exception(),
             )
-            self.signal_shutdown()
+            self.signal_shutdown(expected=False)
 
     def action(self, action: str, args: Any, kwargs: Any) -> Any:
         """Invoke an action of the agent.
@@ -280,7 +291,9 @@ def start(self) -> None:
             self._state = _AgentState.STARTING
             self._bind_handles()
             self.behavior.on_setup()
-            self._action_pool = ThreadPoolExecutor()
+            self._action_pool = ThreadPoolExecutor(
+                self.config.max_action_concurrency,
+            )
             self._loop_pool = ThreadPoolExecutor(
                 max_workers=len(self._loops) + 1,
             )
@@ -321,7 +334,8 @@ def shutdown(self) -> None:
                 return
 
             logger.debug(
-                'Shutting down agent... (%s; %s)',
+                'Shutting down agent... (expected: %s; %s; %s)',
+                self._expected_shutdown,
                 self.agent_id,
                 self.behavior,
             )
@@ -345,12 +359,26 @@ def shutdown(self) -> None:
             if self._loop_pool is not None:
                 self._loop_pool.shutdown(wait=True)
 
+            if (
+                self._expected_shutdown and not self.config.terminate_on_exit
+            ) or (
+                not self._expected_shutdown
+                and not self.config.terminate_on_error
+            ):
+                # TODO: This is a hack because we need to close the mailbox
+                # for the multiplexer listener thread to exit, but in some
+                # cases we don't actually want to close it permanently. This
+                # means there is a race where the mailbox is temporarily
+                # closed.
+                self.exchange.create_mailbox(self.agent_id)
+
+            self.behavior.on_shutdown()
+
             # Close the exchange last since the actions that finished
             # up may still need to use it to send replies.
-            if self.close_exchange:
+            if self.config.close_exchange_on_exit:
                 self.exchange.close()
 
-            self.behavior.on_shutdown()
             self._state = _AgentState.SHUTDOWN
 
             # Raise any exceptions from the loop threads as the final step.
@@ -362,13 +390,14 @@ def shutdown(self) -> None:
                 self.behavior,
             )
 
-    def signal_shutdown(self) -> None:
+    def signal_shutdown(self, expected: bool = True) -> None:
         """Signal that the agent should exit.
 
         If the agent has not started, this will cause the agent to immediately
         shutdown when next started. If the agent is shutdown, this has no
         effect.
         """
+        self._expected_shutdown = expected
         self._shutdown.set()
 
 
diff --git a/aeris/launcher/executor.py b/aeris/launcher/executor.py
index bb77c6f..a750c76 100644
--- a/aeris/launcher/executor.py
+++ b/aeris/launcher/executor.py
@@ -1,8 +1,9 @@
 from __future__ import annotations
 
-import concurrent.futures
+import dataclasses
 import logging
 import sys
+import threading
 from concurrent.futures import CancelledError
 from concurrent.futures import Executor
 from concurrent.futures import Future
@@ -16,6 +17,7 @@
     from typing_extensions import Self
 
 from aeris.agent import Agent
+from aeris.agent import AgentRunConfig
 from aeris.behavior import Behavior
 from aeris.exception import BadIdentifierError
 from aeris.exchange import Exchange
@@ -35,6 +37,17 @@ def _run_agent_on_worker(agent: Agent[Any]) -> None:
     agent.run()
 
 
+@dataclasses.dataclass
+class _ACB:
+    # Agent Control Block
+    agent_id: AgentIdentifier
+    behavior: Behavior
+    exchange: Exchange
+    done: threading.Event
+    future: Future[None] | None = None
+    launch_count: int = 0
+
+
 class ExecutorLauncher:
     """Launcher that wraps a [`concurrent.futures.Executor`][concurrent.futures.Executor].
 
@@ -46,6 +59,8 @@ class ExecutorLauncher:
             when the executor runs agents in separate processes, but should
             be `False` for the `ThreadPoolExecutor` to avoid closing
             shared exchange objects.
+        max_restarts: Maximum times to restart an agent if it exits with
+            an error.
     """  # noqa: E501
 
     def __init__(
@@ -53,11 +68,13 @@ def __init__(
         executor: Executor,
         *,
         close_exchange: bool = True,
+        max_restarts: int = 0,
     ) -> None:
         self._executor = executor
         self._close_exchange = close_exchange
-        self._future_to_id: dict[Future[None], AgentIdentifier] = {}
-        self._id_to_future: dict[AgentIdentifier, Future[None]] = {}
+        self._max_restarts = max_restarts
+        self._acbs: dict[AgentIdentifier, _ACB] = {}
+        self._future_to_acb: dict[Future[None], _ACB] = {}
 
     def __enter__(self) -> Self:
         return self
@@ -77,23 +94,68 @@ def __str__(self) -> str:
         return f'{type(self).__name__}<{type(self._executor).__name__}>'
 
     def _callback(self, future: Future[None]) -> None:
-        agent_id = self._future_to_id[future]
+        acb = self._future_to_acb.pop(future)
+        done = True
+
         try:
             future.result()
-            logger.debug('Completed agent future (%s)', agent_id)
+            logger.debug('Completed agent future (%s)', acb.agent_id)
         except CancelledError:  # pragma: no cover
-            logger.warning('Cancelled agent future (%s)', agent_id)
+            logger.warning('Cancelled agent future (%s)', acb.agent_id)
         except Exception:  # pragma: no cover
-            logger.exception('Received agent exception (%s)', agent_id)
+            logger.exception('Received agent exception (%s)', acb.agent_id)
+            if acb.launch_count < self._max_restarts + 1:
+                self._launch(acb.agent_id)
+                done = False
+
+        if done:
+            acb.done.set()
 
     def close(self) -> None:
         """Close the launcher and shutdown agents."""
         logger.debug('Waiting for agents to shutdown...')
-        for future in self._future_to_id.copy():
-            future.result()
+        for acb in self._acbs.values():
+            if acb.done.is_set() and acb.future is not None:
+                # Raise possible errors from agents so user is sure
+                # to see them.
+                acb.future.result()
         self._executor.shutdown(wait=True, cancel_futures=True)
         logger.debug('Closed launcher (%s)', self)
 
+    def _launch(self, agent_id: AgentIdentifier) -> None:
+        acb = self._acbs[agent_id]
+
+        agent = Agent(
+            acb.behavior,
+            agent_id=acb.agent_id,
+            exchange=acb.exchange,
+            config=AgentRunConfig(
+                close_exchange_on_exit=self._close_exchange,
+                terminate_on_error=acb.launch_count + 1 >= self._max_restarts,
+            ),
+        )
+        future = self._executor.submit(_run_agent_on_worker, agent)
+        acb.launch_count += 1
+        acb.future = future
+        self._future_to_acb[future] = acb
+        future.add_done_callback(self._callback)
+
+        if acb.launch_count == 1:
+            logger.debug(
+                'Launched agent (%s; %s)',
+                acb.agent_id,
+                acb.behavior,
+            )
+        else:
+            restarts = acb.launch_count - 1
+            logger.debug(
+                'Restarted agent (%d/%d retries; %s; %s)',
+                restarts,
+                self._max_restarts,
+                acb.agent_id,
+                acb.behavior,
+            )
+
     def launch(
         self,
         behavior: BehaviorT,
@@ -114,17 +176,9 @@ def launch(
         """
         agent_id = exchange.create_agent() if agent_id is None else agent_id
 
-        agent = Agent(
-            behavior,
-            agent_id=agent_id,
-            exchange=exchange,
-            close_exchange=self._close_exchange,
-        )
-        future = self._executor.submit(_run_agent_on_worker, agent)
-        future.add_done_callback(self._callback)
-        self._future_to_id[future] = agent_id
-        self._id_to_future[agent_id] = future
-        logger.debug('Launched agent (%s; %s)', agent_id, behavior)
+        acb = _ACB(agent_id, behavior, exchange, done=threading.Event())
+        self._acbs[agent_id] = acb
+        self._launch(agent_id)
 
         return exchange.create_handle(agent_id)
 
@@ -136,44 +190,54 @@ def running(self) -> set[AgentIdentifier]:
             launcher that have not completed yet.
         """
         running: set[AgentIdentifier] = set()
-        for future, agent_id in self._future_to_id.items():
-            if not future.done():
-                running.add(agent_id)
+        for acb in self._acbs.values():
+            if not acb.done.is_set():
+                running.add(acb.agent_id)
         return running
 
     def wait(
         self,
         agent_id: AgentIdentifier,
         *,
+        ignore_error: bool = False,
         timeout: float | None = None,
     ) -> None:
         """Wait for a launched agent to exit.
 
+        Note:
+            Calling `wait()` is only valid after `launch()` has succeeded.
+
         Args:
             agent_id: ID of launched agent.
+            ignore_error: Ignore any errors raised by the agent.
             timeout: Optional timeout in seconds to wait for agent.
 
         Raises:
             BadIdentifierError: If an agent with `agent_id` was not
                 launched by this launcher.
             TimeoutError: If `timeout` was exceeded while waiting for agent.
+            Exception: Any exception raised by the agent if
+                `ignore_error=False`.
         """
         try:
-            future = self._id_to_future[agent_id]
+            acb = self._acbs[agent_id]
         except KeyError:
             raise BadIdentifierError(agent_id) from None
 
-        if sys.version_info >= (3, 11):  # pragma: >=3.11 cover
-            try:
-                future.result(timeout=timeout)
-            except TimeoutError:
-                raise
-            except Exception:
-                pass
-        else:  # pragma: <3.11 cover
-            try:
-                future.result(timeout=timeout)
-            except concurrent.futures.TimeoutError:
-                raise TimeoutError() from None
-            except Exception:
-                pass
+        if not acb.done.wait(timeout):
+            raise TimeoutError(
+                f'Agent did not complete within {timeout}s timeout '
+                f'({acb.agent_id})',
+            )
+
+        # The only time _ACB.future is None is between constructing the _ACB
+        # in launch() and creating the future in _launch().
+        assert acb.future is not None
+        # _ACB.done event should only be set in callback of future so
+        # the future must be done.
+        assert acb.future.done()
+
+        if not ignore_error:
+            exc = acb.future.exception()
+            if exc is not None:
+                raise exc
diff --git a/aeris/launcher/thread.py b/aeris/launcher/thread.py
index 3a25048..edbf72c 100644
--- a/aeris/launcher/thread.py
+++ b/aeris/launcher/thread.py
@@ -15,6 +15,7 @@
     from typing_extensions import Self
 
 from aeris.agent import Agent
+from aeris.agent import AgentRunConfig
 from aeris.behavior import Behavior
 from aeris.exception import BadIdentifierError
 from aeris.exchange import Exchange
@@ -93,7 +94,7 @@ def launch(
             behavior,
             agent_id=agent_id,
             exchange=exchange,
-            close_exchange=False,
+            config=AgentRunConfig(close_exchange_on_exit=False),
         )
         thread = threading.Thread(target=agent, name=f'{self}-{agent_id}')
         thread.start()


Test patch:
diff --git a/tests/agent_test.py b/tests/agent_test.py
index 466c7d8..b501853 100644
--- a/tests/agent_test.py
+++ b/tests/agent_test.py
@@ -7,6 +7,7 @@
 import pytest
 
 from aeris.agent import Agent
+from aeris.agent import AgentRunConfig
 from aeris.behavior import action
 from aeris.behavior import Behavior
 from aeris.behavior import loop
@@ -65,6 +66,23 @@ def test_agent_start_shutdown(exchange: Exchange) -> None:
     assert agent.behavior.shutdown_event.is_set()
 
 
+def test_agent_shutdown_without_terminate(exchange: Exchange) -> None:
+    agent_id = exchange.create_agent()
+    agent = Agent(
+        SignalingBehavior(),
+        agent_id=agent_id,
+        exchange=exchange,
+        config=AgentRunConfig(
+            close_exchange_on_exit=False,
+            terminate_on_exit=False,
+        ),
+    )
+    agent.start()
+    agent.shutdown()
+    # Verify mailbox is open
+    exchange.send(agent_id, PingRequest(src=agent_id, dest=agent_id))
+
+
 def test_agent_shutdown_without_start(exchange: Exchange) -> None:
     agent_id = exchange.create_agent()
     agent = Agent(SignalingBehavior(), agent_id=agent_id, exchange=exchange)
diff --git a/tests/launcher/executor_test.py b/tests/launcher/executor_test.py
index 9a17cfd..9f8ca3e 100644
--- a/tests/launcher/executor_test.py
+++ b/tests/launcher/executor_test.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import multiprocessing
+import threading
 import time
 from concurrent.futures import ProcessPoolExecutor
 from concurrent.futures import ThreadPoolExecutor
@@ -8,6 +9,7 @@
 import pytest
 
 from aeris.behavior import Behavior
+from aeris.behavior import loop
 from aeris.exception import BadIdentifierError
 from aeris.exchange import Exchange
 from aeris.exchange.http import HttpExchange
@@ -95,17 +97,51 @@ def test_wait_timeout(exchange: Exchange) -> None:
 
 
 class FailOnStartupBehavior(Behavior):
+    def __init__(self, max_errors: int | None = None) -> None:
+        self.errors = 0
+        self.max_errors = max_errors
+
     def on_setup(self) -> None:
-        raise RuntimeError('Agent startup failed')
+        if self.max_errors is None or self.errors < self.max_errors:
+            self.errors += 1
+            raise RuntimeError('Agent startup failed')
+
+    @loop
+    def exit_after_startup(self, shutdown: threading.Event) -> None:
+        shutdown.set()
+
+
+def test_restart_on_error(exchange: Exchange) -> None:
+    behavior = FailOnStartupBehavior(max_errors=2)
+    # This test only works because we are using a ThreadPoolExecutor so
+    # the state inside FailOnStartupBehavior is shared.
+    executor = ThreadPoolExecutor(max_workers=1)
+    with ExecutorLauncher(
+        executor,
+        close_exchange=False,
+        max_restarts=3,
+    ) as launcher:
+        handle = launcher.launch(behavior, exchange).bind_as_client()
+        launcher.wait(handle.agent_id)
+        handle.close()
+        assert behavior.errors == 2  # noqa: PLR2004
 
 
-def test_wait_ignore_agent_errors(exchange: Exchange) -> None:
+@pytest.mark.parametrize('ignore_error', (True, False))
+def test_wait_ignore_agent_errors(
+    ignore_error: bool,
+    exchange: Exchange,
+) -> None:
     behavior = FailOnStartupBehavior()
     executor = ThreadPoolExecutor(max_workers=1)
     launcher = ExecutorLauncher(executor, close_exchange=False)
 
     handle = launcher.launch(behavior, exchange).bind_as_client()
-    launcher.wait(handle.agent_id)
+    if ignore_error:
+        launcher.wait(handle.agent_id, ignore_error=ignore_error)
+    else:
+        with pytest.raises(RuntimeError, match='Agent startup failed'):
+            launcher.wait(handle.agent_id, ignore_error=ignore_error)
     handle.close()
 
     with pytest.raises(RuntimeError, match='Agent startup failed'):