Repository: proxystore/aeris

Issue:
Add a mechanism to wait on a launched agent
### Describe the Request

In a scenario where I have one agent that when done signals the end of the workflow, it would be helpful to wait on that agent exiting. For example, something like:
```python
class Launcher:
    def wait(self, agent_id: AgentIdentifer, *, timeout: float | None = None) -> None:
        ...
```
This could be re-exported on the `Manager` as well.

### Sample Code

_No response_

Discussion:

Solution patch:
diff --git a/aeris/handle.py b/aeris/handle.py
index 6c1f4a0..e61d15b 100644
--- a/aeris/handle.py
+++ b/aeris/handle.py
@@ -298,7 +298,10 @@ def action(
             Future to the result of the action.
 
         Raises:
-            HandleClosedError: if the handle was closed.
+            HandleClosedError: If the handle was closed.
+            MailboxClosedError: If the agent's mailbox was closed. This
+                typically indicates the agent shutdown for another reason
+                (it self terminated or via another handle).
         """
         if self.mailbox_id is None:
             # UnboundRemoteHandle overrides these methods and is the only
@@ -342,8 +345,11 @@ def ping(self, *, timeout: float | None = None) -> float:
             Round-trip time in seconds.
 
         Raises:
-            HandleClosedError: if the handle was closed.
-            TimeoutError: if the timeout is exceeded.
+            HandleClosedError: If the handle was closed.
+            MailboxClosedError: If the agent's mailbox was closed. This
+                typically indicates the agent shutdown for another reason
+                (it self terminated or via another handle).
+            TimeoutError: If the timeout is exceeded.
         """
         if self.mailbox_id is None:
             # UnboundRemoteHandle overrides these methods and is the only
@@ -380,7 +386,10 @@ def shutdown(self) -> None:
         This is non-blocking and will only send the message.
 
         Raises:
-            HandleClosedError: if the handle was closed.
+            HandleClosedError: If the handle was closed.
+            MailboxClosedError: If the agent's mailbox was closed. This
+                typically indicates the agent shutdown for another reason
+                (it self terminated or via another handle).
         """
         if self.mailbox_id is None:
             # UnboundRemoteHandle overrides these methods and is the only
diff --git a/aeris/launcher/__init__.py b/aeris/launcher/__init__.py
index 8cfbb22..2d8514f 100644
--- a/aeris/launcher/__init__.py
+++ b/aeris/launcher/__init__.py
@@ -44,3 +44,31 @@ def launch(
             Handle to the agent.
         """
         ...
+
+    def running(self) -> set[AgentIdentifier]:
+        """Get a set of IDs for all running agents.
+
+        Returns:
+            Set of agent IDs corresponding to all agents launched by this \
+            launcher that have not completed yet.
+        """
+        ...
+
+    def wait(
+        self,
+        agent_id: AgentIdentifier,
+        *,
+        timeout: float | None = None,
+    ) -> None:
+        """Wait for a launched agent to exit.
+
+        Args:
+            agent_id: ID of launched 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.
+        """
+        ...
diff --git a/aeris/launcher/executor.py b/aeris/launcher/executor.py
index ee8d041..54c6b54 100644
--- a/aeris/launcher/executor.py
+++ b/aeris/launcher/executor.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import concurrent.futures
 import logging
 import sys
 from concurrent.futures import CancelledError
@@ -15,6 +16,7 @@
 
 from aeris.agent import Agent
 from aeris.behavior import Behavior
+from aeris.exception import BadIdentifierError
 from aeris.exchange import Exchange
 from aeris.handle import RemoteHandle
 from aeris.identifier import AgentIdentifier
@@ -34,7 +36,8 @@ class ExecutorLauncher:
 
     def __init__(self, executor: Executor) -> None:
         self._executor = executor
-        self._futures: dict[Future[None], AgentIdentifier] = {}
+        self._future_to_id: dict[Future[None], AgentIdentifier] = {}
+        self._id_to_future: dict[AgentIdentifier, Future[None]] = {}
 
     def __enter__(self) -> Self:
         return self
@@ -54,7 +57,7 @@ def __str__(self) -> str:
         return f'{type(self).__name__}<{type(self._executor).__name__}>'
 
     def _callback(self, future: Future[None]) -> None:
-        agent_id = self._futures.pop(future)
+        agent_id = self._future_to_id[future]
         try:
             future.result()
             logger.debug('Completed agent future (%s)', agent_id)
@@ -66,8 +69,8 @@ def _callback(self, future: Future[None]) -> None:
     def close(self) -> None:
         """Close the launcher and shutdown agents."""
         logger.debug('Waiting for agents to shutdown...')
-        for fut in self._futures.copy():
-            fut.result()
+        for future in self._future_to_id.copy():
+            future.result()
         self._executor.shutdown(wait=True, cancel_futures=True)
         logger.debug('Closed launcher (%s)', self)
 
@@ -99,7 +102,58 @@ def launch(
         )
         future = self._executor.submit(agent)
         future.add_done_callback(self._callback)
-        self._futures[future] = agent_id
+        self._future_to_id[future] = agent_id
+        self._id_to_future[agent_id] = future
         logger.debug('Launched agent (%s; %s)', agent_id, behavior)
 
         return exchange.create_handle(agent_id)
+
+    def running(self) -> set[AgentIdentifier]:
+        """Get a set of IDs for all running agents.
+
+        Returns:
+            Set of agent IDs corresponding to all agents launched by this \
+            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)
+        return running
+
+    def wait(
+        self,
+        agent_id: AgentIdentifier,
+        *,
+        timeout: float | None = None,
+    ) -> None:
+        """Wait for a launched agent to exit.
+
+        Args:
+            agent_id: ID of launched 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.
+        """
+        try:
+            future = self._id_to_future[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
diff --git a/aeris/launcher/thread.py b/aeris/launcher/thread.py
index f064e4c..3a25048 100644
--- a/aeris/launcher/thread.py
+++ b/aeris/launcher/thread.py
@@ -16,6 +16,7 @@
 
 from aeris.agent import Agent
 from aeris.behavior import Behavior
+from aeris.exception import BadIdentifierError
 from aeris.exchange import Exchange
 from aeris.handle import RemoteHandle
 from aeris.identifier import AgentIdentifier
@@ -100,3 +101,42 @@ def launch(
         logger.debug('Launched agent (%s; %s)', agent_id, behavior)
 
         return exchange.create_handle(agent_id)
+
+    def running(self) -> set[AgentIdentifier]:
+        """Get a set of IDs for all running agents.
+
+        Returns:
+            Set of agent IDs corresponding to all agents launched by this \
+            launcher that have not completed yet.
+        """
+        running: set[AgentIdentifier] = set()
+        for agent_id, agent in self._agents.items():
+            if agent.thread.is_alive():
+                running.add(agent_id)
+        return running
+
+    def wait(
+        self,
+        agent_id: AgentIdentifier,
+        *,
+        timeout: float | None = None,
+    ) -> None:
+        """Wait for a launched agent to exit.
+
+        Args:
+            agent_id: ID of launched 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.
+        """
+        try:
+            agent = self._agents[agent_id]
+        except KeyError:
+            raise BadIdentifierError(agent_id) from None
+
+        agent.thread.join(timeout=timeout)
+        if agent.thread.is_alive():
+            raise TimeoutError()
diff --git a/aeris/manager.py b/aeris/manager.py
index 6ff11f7..90ad951 100644
--- a/aeris/manager.py
+++ b/aeris/manager.py
@@ -1,9 +1,11 @@
 from __future__ import annotations
 
+import contextlib
 import logging
 import sys
 import threading
 from types import TracebackType
+from typing import Any
 from typing import TypeVar
 
 if sys.version_info >= (3, 11):  # pragma: >=3.11 cover
@@ -12,8 +14,10 @@
     from typing_extensions import Self
 
 from aeris.behavior import Behavior
+from aeris.exception import BadIdentifierError
+from aeris.exception import MailboxClosedError
 from aeris.exchange import Exchange
-from aeris.handle import RemoteHandle
+from aeris.handle import BoundRemoteHandle
 from aeris.identifier import AgentIdentifier
 from aeris.identifier import ClientIdentifier
 from aeris.launcher import Launcher
@@ -58,6 +62,7 @@ def __init__(self, exchange: Exchange, launcher: Launcher) -> None:
             self._exchange,
             self._handle_request,
         )
+        self._handles: dict[AgentIdentifier, BoundRemoteHandle[Any]] = {}
         self._listener_thread = threading.Thread(
             target=self._multiplexer.listen,
             name=f'multiplexer-{self.mailbox_id.uid}-listener',
@@ -119,15 +124,17 @@ def _handle_request(self, request: RequestMessage) -> None:
     def close(self) -> None:
         """Close the manager and cleanup resources.
 
-        1. Call shutdown on all launched agents.
+        1. Call shutdown on all running agents.
         1. Close all handles created by the manager.
         1. Close the mailbox associated with the manager.
         1. Close the exchange.
         1. Close the launcher.
         """
-        for handle in self._multiplexer.bound_handles.values():
-            handle.shutdown()
-        logger.debug('Instructing managed agents to shutdown')
+        for agent_id in self.launcher.running():
+            handle = self._handles[agent_id]
+            with contextlib.suppress(MailboxClosedError):
+                handle.shutdown()
+        logger.debug('Instructed managed agents to shutdown')
         self._multiplexer.close_bound_handles()
         self._multiplexer.close_mailbox()
         self._listener_thread.join()
@@ -140,7 +147,7 @@ def launch(
         behavior: BehaviorT,
         *,
         agent_id: AgentIdentifier | None = None,
-    ) -> RemoteHandle[BehaviorT]:
+    ) -> BoundRemoteHandle[BehaviorT]:
         """Launch a new agent with a specified behavior.
 
         Note:
@@ -162,5 +169,55 @@ def launch(
         )
         logger.info('Launched agent (%s; %s)', unbound.agent_id, behavior)
         bound = self._multiplexer.bind(unbound)
+        self._handles[bound.agent_id] = bound
         logger.debug('Bound agent handle to manager (%s)', bound)
         return bound
+
+    def shutdown(
+        self,
+        agent_id: AgentIdentifier,
+        *,
+        blocking: bool = True,
+        timeout: float | None = None,
+    ) -> None:
+        """Shutdown a launched agent.
+
+        Args:
+            agent_id: ID of launched agent.
+            blocking: Wait for the agent to exit before returning.
+            timeout: Optional timeout is seconds when `blocking=True`.
+
+        Raises:
+            BadIdentifierError: If an agent with `agent_id` was not
+                launched by this launcher.
+            TimeoutError: If `timeout` was exceeded while blocking for agent.
+        """
+        try:
+            handle = self._handles[agent_id]
+        except KeyError:
+            raise BadIdentifierError(agent_id) from None
+
+        with contextlib.suppress(MailboxClosedError):
+            handle.shutdown()
+
+        if blocking:
+            self.wait(agent_id, timeout=timeout)
+
+    def wait(
+        self,
+        agent_id: AgentIdentifier,
+        *,
+        timeout: float | None = None,
+    ) -> None:
+        """Wait for a launched agent to exit.
+
+        Args:
+            agent_id: ID of launched 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.
+        """
+        self.launcher.wait(agent_id, timeout=timeout)


Test patch:
diff --git a/tests/launcher/executor_test.py b/tests/launcher/executor_test.py
index 0a0ce7f..77f392e 100644
--- a/tests/launcher/executor_test.py
+++ b/tests/launcher/executor_test.py
@@ -5,6 +5,10 @@
 from concurrent.futures import ProcessPoolExecutor
 from concurrent.futures import ThreadPoolExecutor
 
+import pytest
+
+from aeris.behavior import Behavior
+from aeris.exception import BadIdentifierError
 from aeris.exchange import Exchange
 from aeris.exchange.simple import SimpleExchange
 from aeris.launcher import Launcher
@@ -29,6 +33,8 @@ def test_launch_agents_threads(exchange: Exchange) -> None:
         handle1 = launcher.launch(behavior, exchange).bind_as_client()
         handle2 = launcher.launch(behavior, exchange).bind_as_client()
 
+        assert len(launcher.running()) == 2  # noqa: PLR2004
+
         time.sleep(5 * TEST_LOOP_SLEEP)
 
         handle1.shutdown()
@@ -37,6 +43,11 @@ def test_launch_agents_threads(exchange: Exchange) -> None:
         handle1.close()
         handle2.close()
 
+        launcher.wait(handle1.agent_id)
+        launcher.wait(handle2.agent_id)
+
+        assert len(launcher.running()) == 0
+
 
 def test_launch_agents_processes(
     simple_exchange_server: tuple[str, int],
@@ -59,3 +70,44 @@ def test_launch_agents_processes(
 
             handle1.close()
             handle2.close()
+
+
+def test_wait_bad_identifier(exchange: Exchange) -> None:
+    executor = ThreadPoolExecutor(max_workers=1)
+    with ExecutorLauncher(executor) as launcher:
+        agent_id = exchange.create_agent()
+
+        with pytest.raises(BadIdentifierError):
+            launcher.wait(agent_id)
+
+
+def test_wait_timeout(exchange: Exchange) -> None:
+    behavior = SleepBehavior(TEST_LOOP_SLEEP)
+    executor = ThreadPoolExecutor(max_workers=1)
+    with ExecutorLauncher(executor) as launcher:
+        handle = launcher.launch(behavior, exchange).bind_as_client()
+
+        with pytest.raises(TimeoutError):
+            launcher.wait(handle.agent_id, timeout=TEST_LOOP_SLEEP)
+
+        handle.shutdown()
+        handle.close()
+
+
+class FailOnStartupBehavior(Behavior):
+    def setup(self) -> None:
+        raise RuntimeError('Agent startup failed')
+
+
+def test_wait_ignore_agent_errors(exchange: Exchange) -> None:
+    behavior = FailOnStartupBehavior()
+    executor = ThreadPoolExecutor(max_workers=1)
+    launcher = ExecutorLauncher(executor)
+
+    handle = launcher.launch(behavior, exchange).bind_as_client()
+    launcher.wait(handle.agent_id)
+    handle.close()
+
+    with pytest.raises(RuntimeError, match='Agent startup failed'):
+        launcher.close()
+    executor.shutdown()
diff --git a/tests/launcher/thread_test.py b/tests/launcher/thread_test.py
index 82f1a22..d44dded 100644
--- a/tests/launcher/thread_test.py
+++ b/tests/launcher/thread_test.py
@@ -2,6 +2,9 @@
 
 import time
 
+import pytest
+
+from aeris.exception import BadIdentifierError
 from aeris.exchange import Exchange
 from aeris.launcher import Launcher
 from aeris.launcher.thread import ThreadLauncher
@@ -22,6 +25,8 @@ def test_launch_agents(exchange: Exchange) -> None:
         handle1 = launcher.launch(behavior, exchange).bind_as_client()
         handle2 = launcher.launch(behavior, exchange).bind_as_client()
 
+        assert len(launcher.running()) == 2  # noqa: PLR2004
+
         time.sleep(5 * TEST_LOOP_SLEEP)
 
         handle1.shutdown()
@@ -29,3 +34,28 @@ def test_launch_agents(exchange: Exchange) -> None:
 
         handle1.close()
         handle2.close()
+
+        launcher.wait(handle1.agent_id)
+        launcher.wait(handle2.agent_id)
+
+        assert len(launcher.running()) == 0
+
+
+def test_wait_bad_identifier(exchange: Exchange) -> None:
+    with ThreadLauncher() as launcher:
+        agent_id = exchange.create_agent()
+
+        with pytest.raises(BadIdentifierError):
+            launcher.wait(agent_id)
+
+
+def test_wait_timeout(exchange: Exchange) -> None:
+    behavior = SleepBehavior(TEST_LOOP_SLEEP)
+    with ThreadLauncher() as launcher:
+        handle = launcher.launch(behavior, exchange).bind_as_client()
+
+        with pytest.raises(TimeoutError):
+            launcher.wait(handle.agent_id, timeout=TEST_LOOP_SLEEP)
+
+        handle.shutdown()
+        handle.close()
diff --git a/tests/manager_test.py b/tests/manager_test.py
index c4dbe7c..e7e4c71 100644
--- a/tests/manager_test.py
+++ b/tests/manager_test.py
@@ -2,6 +2,9 @@
 
 import time
 
+import pytest
+
+from aeris.exception import BadIdentifierError
 from aeris.exchange.thread import ThreadExchange
 from aeris.launcher.thread import ThreadLauncher
 from aeris.manager import Manager
@@ -9,6 +12,7 @@
 from aeris.message import PingResponse
 from testing.behavior import SleepBehavior
 from testing.constant import TEST_LOOP_SLEEP
+from testing.constant import TEST_THREAD_JOIN_TIMEOUT
 
 
 def test_protocol() -> None:
@@ -43,3 +47,59 @@ def test_reply_to_requests_with_error() -> None:
         response = manager.exchange.recv(client_id)
         assert isinstance(response, PingResponse)
         assert isinstance(response.exception, TypeError)
+
+
+def test_wait_bad_identifier(exchange: ThreadExchange) -> None:
+    with Manager(
+        exchange=ThreadExchange(),
+        launcher=ThreadLauncher(),
+    ) as manager:
+        agent_id = manager.exchange.create_agent()
+
+        with pytest.raises(BadIdentifierError):
+            manager.wait(agent_id)
+
+
+def test_wait_timeout(exchange: ThreadExchange) -> None:
+    behavior = SleepBehavior(TEST_LOOP_SLEEP)
+    with Manager(
+        exchange=ThreadExchange(),
+        launcher=ThreadLauncher(),
+    ) as manager:
+        handle = manager.launch(behavior)
+
+        with pytest.raises(TimeoutError):
+            manager.wait(handle.agent_id, timeout=TEST_LOOP_SLEEP)
+
+
+def test_shutdown_bad_identifier(exchange: ThreadExchange) -> None:
+    with Manager(
+        exchange=ThreadExchange(),
+        launcher=ThreadLauncher(),
+    ) as manager:
+        agent_id = manager.exchange.create_agent()
+
+        with pytest.raises(BadIdentifierError):
+            manager.shutdown(agent_id)
+
+
+def test_shutdown_nonblocking(exchange: ThreadExchange) -> None:
+    behavior = SleepBehavior(TEST_LOOP_SLEEP)
+    with Manager(
+        exchange=ThreadExchange(),
+        launcher=ThreadLauncher(),
+    ) as manager:
+        handle = manager.launch(behavior)
+        manager.shutdown(handle.agent_id, blocking=False)
+        manager.wait(handle.agent_id, timeout=TEST_THREAD_JOIN_TIMEOUT)
+
+
+def test_shutdown_blocking(exchange: ThreadExchange) -> None:
+    behavior = SleepBehavior(TEST_LOOP_SLEEP)
+    with Manager(
+        exchange=ThreadExchange(),
+        launcher=ThreadLauncher(),
+    ) as manager:
+        handle = manager.launch(behavior)
+        manager.shutdown(handle.agent_id, blocking=True)
+        manager.wait(handle.agent_id, timeout=TEST_LOOP_SLEEP)