Repository: proxystore/aeris

Issue:
Pass exchange to `Launcher.launch()` rather than constructor
### Describe the Request

Currently, `Launcher` implementations are initialized with an exchange to use, which makes it unclear if the `Launcher` owns the exchange or not. The `Launcher` otherwise does not need the exchange for anything besides creating the `Agent` in `launch()`. This will be important for #25 .

### Sample Code

_No response_

Discussion:

Solution patch:
diff --git a/aeris/launcher/__init__.py b/aeris/launcher/__init__.py
index 7dfd8c8..388ac7b 100644
--- a/aeris/launcher/__init__.py
+++ b/aeris/launcher/__init__.py
@@ -5,6 +5,7 @@
 from typing import TypeVar
 
 from aeris.behavior import Behavior
+from aeris.exchange import Exchange
 from aeris.handle import RemoteHandle
 
 __all__ = ['Launcher']
@@ -23,11 +24,16 @@ def close(self) -> None:
         """Close the launcher and shutdown agents."""
         ...
 
-    def launch(self, behavior: BehaviorT) -> RemoteHandle[BehaviorT]:
+    def launch(
+        self,
+        behavior: BehaviorT,
+        exchange: Exchange,
+    ) -> RemoteHandle[BehaviorT]:
         """Launch a new agent with a specified behavior.
 
         Args:
             behavior: Behavior the agent should implement.
+            exchange: Exchange the agent will use for messaging.
 
         Returns:
             Handle to the agent.
diff --git a/aeris/launcher/executor.py b/aeris/launcher/executor.py
index e3438f4..7e00ec6 100644
--- a/aeris/launcher/executor.py
+++ b/aeris/launcher/executor.py
@@ -28,13 +28,11 @@ class ExecutorLauncher:
     """Launcher that wraps a [`concurrent.futures.Executor`][concurrent.futures.Executor].
 
     Args:
-        exchange: Exchange used for communication.
         executor: Executor used for launching agents. Note that this class
             takes ownership of the `executor`.
     """  # noqa: E501
 
-    def __init__(self, exchange: Exchange, executor: Executor) -> None:
-        self._exchange = exchange
+    def __init__(self, executor: Executor) -> None:
         self._executor = executor
         self._futures: dict[Future[None], AgentIdentifier] = {}
 
@@ -50,13 +48,10 @@ def __exit__(
         self.close()
 
     def __repr__(self) -> str:
-        return (
-            f'{type(self).__name__}'
-            f'(exchange={self._exchange!r}, executor={self._executor!r})'
-        )
+        return f'{type(self).__name__}(executor={self._executor!r})'
 
     def __str__(self) -> str:
-        return f'{type(self).__name__}<{self._exchange}; {self._executor}>'
+        return f'{type(self).__name__}<{self._executor}>'
 
     def _callback(self, future: Future[None]) -> None:
         aid = self._futures.pop(future)
@@ -76,21 +71,26 @@ def close(self) -> None:
         self._executor.shutdown(wait=True, cancel_futures=True)
         logger.info('Closed %s', self)
 
-    def launch(self, behavior: BehaviorT) -> RemoteHandle[BehaviorT]:
+    def launch(
+        self,
+        behavior: BehaviorT,
+        exchange: Exchange,
+    ) -> RemoteHandle[BehaviorT]:
         """Launch a new agent with a specified behavior.
 
         Args:
             behavior: Behavior the agent should implement.
+            exchange: Exchange the agent will use for messaging.
 
         Returns:
             Mailbox used to communicate with agent.
         """
-        aid = self._exchange.create_agent()
+        aid = exchange.create_agent()
 
         agent = Agent(
             behavior,
             aid=aid,
-            exchange=self._exchange,
+            exchange=exchange,
             close_exchange=True,
         )
         future = self._executor.submit(agent)
@@ -98,4 +98,4 @@ def launch(self, behavior: BehaviorT) -> RemoteHandle[BehaviorT]:
         self._futures[future] = aid
         logger.info('Launched agent with %s', agent)
 
-        return self._exchange.create_handle(aid)
+        return exchange.create_handle(aid)
diff --git a/aeris/launcher/thread.py b/aeris/launcher/thread.py
index 44de78f..011da27 100644
--- a/aeris/launcher/thread.py
+++ b/aeris/launcher/thread.py
@@ -39,9 +39,8 @@ class ThreadLauncher:
     scalability of agents.
     """
 
-    def __init__(self, exchange: Exchange) -> None:
+    def __init__(self) -> None:
         self._agents: dict[AgentIdentifier, _RunningAgent[Any]] = {}
-        self._exchange = exchange
 
     def __enter__(self) -> Self:
         return self
@@ -55,10 +54,10 @@ def __exit__(
         self.close()
 
     def __repr__(self) -> str:
-        return f'{type(self).__name__}(exchange={self._exchange!r})'
+        return f'{type(self).__name__}()'
 
     def __str__(self) -> str:
-        return f'{type(self).__name__}<{self._exchange}>'
+        return f'{type(self).__name__}<>'
 
     def close(self) -> None:
         """Close the launcher and shutdown agents."""
@@ -69,21 +68,31 @@ def close(self) -> None:
             self._agents[aid].thread.join()
         logger.info('Closed %s', self)
 
-    def launch(self, behavior: BehaviorT) -> RemoteHandle[BehaviorT]:
+    def launch(
+        self,
+        behavior: BehaviorT,
+        exchange: Exchange,
+    ) -> RemoteHandle[BehaviorT]:
         """Launch a new agent with a specified behavior.
 
         Args:
             behavior: Behavior the agent should implement.
+            exchange: Exchange the agent will use for messaging.
 
         Returns:
             Mailbox used to communicate with agent.
         """
-        aid = self._exchange.create_agent()
-
-        agent = Agent(behavior, aid=aid, exchange=self._exchange)
+        aid = exchange.create_agent()
+
+        agent = Agent(
+            behavior,
+            aid=aid,
+            exchange=exchange,
+            close_exchange=False,
+        )
         thread = threading.Thread(target=agent, name=f'{self}-{aid}')
         thread.start()
         self._agents[aid] = _RunningAgent(agent, thread)
         logger.info('Launched %s', agent)
 
-        return self._exchange.create_handle(aid)
+        return exchange.create_handle(aid)
diff --git a/examples/01-actor-client/run-01.py b/examples/01-actor-client/run-01.py
index cfda551..127b139 100644
--- a/examples/01-actor-client/run-01.py
+++ b/examples/01-actor-client/run-01.py
@@ -28,11 +28,9 @@ def get_count(self) -> int:
 def main() -> int:
     init_logging(logging.DEBUG)
 
-    behavior = Counter()
-    exchange = ThreadExchange()
-
-    with ThreadLauncher(exchange) as launcher:
-        agent = launcher.launch(behavior).bind_as_client()
+    with ThreadExchange() as exchange, ThreadLauncher() as launcher:
+        behavior = Counter()
+        agent = launcher.launch(behavior, exchange).bind_as_client()
 
         future: Future[int] = agent.action('get_count')
         assert future.result() == 0
diff --git a/examples/02-agent-loop/run-02.py b/examples/02-agent-loop/run-02.py
index 8b8ee41..7f07790 100644
--- a/examples/02-agent-loop/run-02.py
+++ b/examples/02-agent-loop/run-02.py
@@ -33,11 +33,9 @@ def get_count(self) -> int:
 def main() -> int:
     init_logging(logging.DEBUG)
 
-    behavior = Counter()
-    exchange = ThreadExchange()
-
-    with ThreadLauncher(exchange) as launcher:
-        agent = launcher.launch(behavior).bind_as_client()
+    with ThreadExchange() as exchange, ThreadLauncher() as launcher:
+        behavior = Counter()
+        agent = launcher.launch(behavior, exchange).bind_as_client()
 
         time.sleep(2)
 
diff --git a/examples/03-agent-agent/run-03.py b/examples/03-agent-agent/run-03.py
index 901c9e6..823a77e 100644
--- a/examples/03-agent-agent/run-03.py
+++ b/examples/03-agent-agent/run-03.py
@@ -42,11 +42,12 @@ def reverse(self, text: str) -> str:
 def main() -> int:
     init_logging(logging.DEBUG)
 
-    with ThreadLauncher(ThreadExchange()) as launcher:
-        lowerer = launcher.launch(Lowerer())
-        reverser = launcher.launch(Reverser())
+    with ThreadExchange() as exchange, ThreadLauncher() as launcher:
+        lowerer = launcher.launch(Lowerer(), exchange)
+        reverser = launcher.launch(Reverser(), exchange)
         coordinator = launcher.launch(
             Coordinator(lowerer, reverser),
+            exchange,
         ).bind_as_client()
 
         text = 'DEADBEEF'
diff --git a/examples/04-execution/run-04.py b/examples/04-execution/run-04.py
index 7b253ad..e3a056f 100644
--- a/examples/04-execution/run-04.py
+++ b/examples/04-execution/run-04.py
@@ -59,17 +59,15 @@ def main() -> int:
         # any concurrent.futures.Executor (here, a ProcessPoolExecutor) to
         # execute agents. The launcher is initialized to use the simple
         # exchange that we just set up for agent and client communication.
-        with ExecutorLauncher(
-            exchange=exchange,
-            executor=ProcessPoolExecutor(max_workers=3),
-        ) as launcher:
+        with ExecutorLauncher(ProcessPoolExecutor(max_workers=3)) as launcher:
             # Initialize and launch each of the three agents. The
             # bind_as_client method returns a client handle to the agent that
             # can be used to send commands to that agent.
-            lowerer = launcher.launch(Lowerer()).bind_as_client()
-            reverser = launcher.launch(Reverser()).bind_as_client()
+            lowerer = launcher.launch(Lowerer(), exchange).bind_as_client()
+            reverser = launcher.launch(Reverser(), exchange).bind_as_client()
             coordinator = launcher.launch(
                 Coordinator(lowerer, reverser),
+                exchange,
             ).bind_as_client()
 
             text = 'DEADBEEF'
diff --git a/testing/fixture.py b/testing/fixture.py
index 9236409..2b62537 100644
--- a/testing/fixture.py
+++ b/testing/fixture.py
@@ -6,7 +6,6 @@
 
 import pytest
 
-from aeris.exchange import Exchange
 from aeris.exchange.simple import SimpleServer
 from aeris.exchange.thread import ThreadExchange
 from aeris.launcher.thread import ThreadLauncher
@@ -16,14 +15,14 @@
 
 
 @pytest.fixture
-def exchange() -> Generator[Exchange]:
+def exchange() -> Generator[ThreadExchange]:
     with ThreadExchange() as exchange:
         yield exchange
 
 
 @pytest.fixture
-def launcher(exchange: Exchange) -> Generator[ThreadLauncher]:
-    with ThreadLauncher(exchange) as launcher:
+def launcher() -> Generator[ThreadLauncher]:
+    with ThreadLauncher() as launcher:
         yield launcher
 
 


Test patch:
diff --git a/tests/handle_test.py b/tests/handle_test.py
index cd72bad..ae1f062 100644
--- a/tests/handle_test.py
+++ b/tests/handle_test.py
@@ -9,6 +9,7 @@
 from aeris.exception import HandleClosedError
 from aeris.exception import HandleNotBoundError
 from aeris.exchange import Exchange
+from aeris.exchange.thread import ThreadExchange
 from aeris.handle import AgentRemoteHandle
 from aeris.handle import ClientRemoteHandle
 from aeris.handle import Handle
@@ -173,10 +174,12 @@ def test_client_remote_handle_bind(exchange: Exchange) -> None:
 
 
 def test_client_remote_handle_log_bad_response(
+    exchange: ThreadExchange,
     launcher: ThreadLauncher,
 ) -> None:
+    behavior = EmptyBehavior()
     handle: RemoteHandle[Any]
-    with launcher.launch(EmptyBehavior()) as handle:
+    with launcher.launch(behavior, exchange) as handle:
         client = handle.bind_as_client()
         assert client.hid is not None
         # Should log but not crash
@@ -206,8 +209,12 @@ def test_client_remote_handle_recv_thread_crash(exchange: Exchange) -> None:
         handle.close()
 
 
-def test_client_remote_handle_actions(launcher: ThreadLauncher) -> None:
-    with launcher.launch(CounterBehavior()).bind_as_client() as handle:
+def test_client_remote_handle_actions(
+    exchange: ThreadExchange,
+    launcher: ThreadLauncher,
+) -> None:
+    behavior = CounterBehavior()
+    with launcher.launch(behavior, exchange).bind_as_client() as handle:
         assert handle.ping() > 0
 
         add_future: Future[None] = handle.action('add', 1)
@@ -219,26 +226,36 @@ def test_client_remote_handle_actions(launcher: ThreadLauncher) -> None:
         handle.shutdown()
 
 
-def test_client_remote_handle_errors(launcher: ThreadLauncher) -> None:
-    with launcher.launch(ErrorBehavior()).bind_as_client() as handle:
+def test_client_remote_handle_errors(
+    exchange: ThreadExchange,
+    launcher: ThreadLauncher,
+) -> None:
+    behavior = ErrorBehavior()
+    with launcher.launch(behavior, exchange).bind_as_client() as handle:
         with pytest.raises(RuntimeError, match='This action always fails.'):
             handle.action('fails').result()
         with pytest.raises(AttributeError, match='null'):
             handle.action('null').result()
 
 
-def test_client_remote_handle_wait_futures(launcher: ThreadLauncher) -> None:
+def test_client_remote_handle_wait_futures(
+    exchange: ThreadExchange,
+    launcher: ThreadLauncher,
+) -> None:
     behavior = SleepBehavior()
-    handle = launcher.launch(behavior).bind_as_client()
+    handle = launcher.launch(behavior, exchange).bind_as_client()
 
     future: Future[None] = handle.action('sleep', TEST_SLEEP)
     handle.close(wait_futures=True)
     future.result(timeout=0)
 
 
-def test_client_remote_handle_cancel_futures(launcher: ThreadLauncher) -> None:
+def test_client_remote_handle_cancel_futures(
+    exchange: ThreadExchange,
+    launcher: ThreadLauncher,
+) -> None:
     behavior = SleepBehavior()
-    handle = launcher.launch(behavior).bind_as_client()
+    handle = launcher.launch(behavior, exchange).bind_as_client()
 
     future: Future[None] = handle.action('sleep', TEST_SLEEP)
     handle.close(wait_futures=False)
diff --git a/tests/launcher/executor_test.py b/tests/launcher/executor_test.py
index 9c52c1c..0a0ce7f 100644
--- a/tests/launcher/executor_test.py
+++ b/tests/launcher/executor_test.py
@@ -14,9 +14,9 @@
 from testing.constant import TEST_LOOP_SLEEP
 
 
-def test_protocol(exchange: Exchange) -> None:
+def test_protocol() -> None:
     executor = ThreadPoolExecutor(max_workers=1)
-    with ExecutorLauncher(exchange, executor) as launcher:
+    with ExecutorLauncher(executor) as launcher:
         assert isinstance(launcher, Launcher)
         assert isinstance(repr(launcher), str)
         assert isinstance(str(launcher), str)
@@ -25,9 +25,9 @@ def test_protocol(exchange: Exchange) -> None:
 def test_launch_agents_threads(exchange: Exchange) -> None:
     behavior = SleepBehavior(TEST_LOOP_SLEEP)
     executor = ThreadPoolExecutor(max_workers=2)
-    with ExecutorLauncher(exchange, executor) as launcher:
-        handle1 = launcher.launch(behavior).bind_as_client()
-        handle2 = launcher.launch(behavior).bind_as_client()
+    with ExecutorLauncher(executor) as launcher:
+        handle1 = launcher.launch(behavior, exchange).bind_as_client()
+        handle2 = launcher.launch(behavior, exchange).bind_as_client()
 
         time.sleep(5 * TEST_LOOP_SLEEP)
 
@@ -47,9 +47,9 @@ def test_launch_agents_processes(
     host, port = simple_exchange_server
 
     with SimpleExchange(host, port) as exchange:
-        with ExecutorLauncher(exchange, executor) as launcher:
-            handle1 = launcher.launch(behavior).bind_as_client()
-            handle2 = launcher.launch(behavior).bind_as_client()
+        with ExecutorLauncher(executor) as launcher:
+            handle1 = launcher.launch(behavior, exchange).bind_as_client()
+            handle2 = launcher.launch(behavior, exchange).bind_as_client()
 
             assert handle1.ping(timeout=TEST_CONNECTION_TIMEOUT) > 0
             assert handle2.ping(timeout=TEST_CONNECTION_TIMEOUT) > 0
diff --git a/tests/launcher/thread_test.py b/tests/launcher/thread_test.py
index 574e230..82f1a22 100644
--- a/tests/launcher/thread_test.py
+++ b/tests/launcher/thread_test.py
@@ -9,8 +9,8 @@
 from testing.constant import TEST_LOOP_SLEEP
 
 
-def test_protocol(exchange: Exchange) -> None:
-    with ThreadLauncher(exchange) as launcher:
+def test_protocol() -> None:
+    with ThreadLauncher() as launcher:
         assert isinstance(launcher, Launcher)
         assert isinstance(repr(launcher), str)
         assert isinstance(str(launcher), str)
@@ -18,9 +18,9 @@ def test_protocol(exchange: Exchange) -> None:
 
 def test_launch_agents(exchange: Exchange) -> None:
     behavior = SleepBehavior(TEST_LOOP_SLEEP)
-    with ThreadLauncher(exchange) as launcher:
-        handle1 = launcher.launch(behavior).bind_as_client()
-        handle2 = launcher.launch(behavior).bind_as_client()
+    with ThreadLauncher() as launcher:
+        handle1 = launcher.launch(behavior, exchange).bind_as_client()
+        handle2 = launcher.launch(behavior, exchange).bind_as_client()
 
         time.sleep(5 * TEST_LOOP_SLEEP)