diff --git a/src/ezmsg/core/backend.py b/src/ezmsg/core/backend.py index 8b150798..084282ec 100644 --- a/src/ezmsg/core/backend.py +++ b/src/ezmsg/core/backend.py @@ -321,9 +321,10 @@ def __init__( graph_address: AddressType | None = None, force_single_process: bool = False, profiler_log_name: str | None = None, + auto_start: bool | None = None, **components_kwargs: Component, ) -> None: - + components = either_dict_or_kwargs(components, components_kwargs, "GraphRunner") if components is None: raise ValueError("Must supply at least one component to run") @@ -336,6 +337,7 @@ def __init__( self._graph_address = graph_address self._force_single_process = force_single_process self._profiler_log_name = profiler_log_name + self._auto_start = auto_start self._execution_context = None self._graph_context = None @@ -690,7 +692,9 @@ def _initialize(self, force_single_process: bool, wait_for_ready: bool) -> bool: try: async def create_graph_context() -> GraphContext: - return await GraphContext(self._graph_address).__aenter__() + return await GraphContext( + self._graph_address, auto_start=self._auto_start + ).__aenter__() graph_context = asyncio.run_coroutine_threadsafe( create_graph_context(), self._loop @@ -912,6 +916,7 @@ def run( graph_address: AddressType | None = None, force_single_process: bool = False, profiler_log_name: str | None = None, + auto_start: bool | None = None, **components_kwargs: Component, ) -> None: """ @@ -942,6 +947,11 @@ def run( :type graph_address: AddressType | None :param force_single_process: Whether to force all components into a single process :type force_single_process: bool + :param auto_start: Whether to spawn a GraphServer if none is reachable at ``graph_address``. + If ``None`` (default), a server is auto-started only when no explicit ``graph_address`` + is provided and no environment override is set. Pass ``True`` to force auto-start when + binding to a specific address (e.g. an ephemeral port chosen by the parent process). + :type auto_start: bool | None :param components_kwargs: Additional components specified as keyword arguments :type components_kwargs: Component @@ -968,6 +978,7 @@ def run( graph_address=graph_address, force_single_process=force_single_process, profiler_log_name=profiler_log_name, + auto_start=auto_start, ) runner.run_blocking() diff --git a/tests/test_run.py b/tests/test_run.py index 323bc4ba..6923c4ab 100644 --- a/tests/test_run.py +++ b/tests/test_run.py @@ -129,3 +129,27 @@ def test_run_collection(passthrough_settings, num_messages): for line in lines: results.append(json.loads(line)) assert len(results) == num_messages + + +def test_run_auto_start_with_explicit_address(unused_tcp_port): + # When ``graph_address`` is provided, GraphContext defaults to ``auto_start=False`` + # and refuses to spawn a server. ``auto_start=True`` overrides that, which lets a + # parent process pick an ephemeral port and hand it to a child that owns the graph. + num_messages = 3 + with get_test_fn() as test_filename: + comps = { + "SIMPLE_PUB": MessageGenerator(num_msgs=num_messages), + "SIMPLE_SUB": MessageReceiver(num_msgs=num_messages, output_fn=test_filename), + } + conns = ((comps["SIMPLE_PUB"].OUTPUT, comps["SIMPLE_SUB"].INPUT),) + + ez.run( + components=comps, + connections=conns, + graph_address=("127.0.0.1", unused_tcp_port), + auto_start=True, + ) + + with open(test_filename, "r") as file: + results = [json.loads(line) for line in file.readlines()] + assert len(results) == num_messages