feat(execution): add Coordinator with streaming run + run_to_completion

src/lfx/src/lfx/execution/coordinator.py

@@ -0,0 +1,51 @@
+"""Coordinator: graph in, streaming step results out."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from lfx.execution.partitioner import identity_partition
+from lfx.execution.types import RunComplete
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncIterator
+
+    from lfx.execution.registry import ExecutorRegistry
+    from lfx.execution.types import StepResult
+
+
+class Coordinator:
+    def __init__(
+        self,
+        *,
+        registry: ExecutorRegistry,
+        executor_kind: str = "in-process",
+    ) -> None:
+        self._registry = registry
+        self._executor_kind = executor_kind
+
+    async def run(
+        self,
+        graph: Any,
+        *,
+        inputs: list[dict[str, Any]],
+        **runtime_options: Any,
+    ) -> AsyncIterator[StepResult | RunComplete]:
+        units = identity_partition(graph, inputs=inputs, runtime_options=runtime_options)
+        executor = self._registry.get(self._executor_kind)
+        for unit in units:
+            async for item in executor.execute(unit):
+                yield item
+
+    async def run_to_completion(
+        self,
+        graph: Any,
+        *,
+        inputs: list[dict[str, Any]],
+        **runtime_options: Any,
+    ) -> list[Any]:
+        async for item in self.run(graph, inputs=inputs, **runtime_options):
+            if isinstance(item, RunComplete):
+                return item.outputs
+        msg = "Executor stream ended without a RunComplete"
+        raise RuntimeError(msg)

src/lfx/tests/unit/execution/test_coordinator.py

@@ -0,0 +1,62 @@
+import pytest
+from lfx.components.input_output import ChatInput, ChatOutput
+from lfx.execution.coordinator import Coordinator
+from lfx.execution.registry import ExecutorRegistry
+from lfx.execution.types import RunComplete, StepResult
+from lfx.graph import Graph
+
+
+def _simple_graph():
+    chat_input = ChatInput(_id="chat_input")
+    chat_input.set(should_store_message=False)
+    chat_output = ChatOutput(input_value="test", _id="chat_output")
+    chat_output.set(sender_name=chat_input.message_response)
+    return Graph(chat_input, chat_output)
+
+
+def _make_coordinator():
+    from lfx.execution.backends.in_process import InProcessExecutor
+
+    registry = ExecutorRegistry()
+    registry.register(InProcessExecutor())
+    return Coordinator(registry=registry)
+
+
+@pytest.mark.asyncio
+async def test_run_yields_streaming_step_results_ending_in_run_complete():
+    coordinator = _make_coordinator()
+    items = [item async for item in coordinator.run(_simple_graph(), inputs=[{"input_value": "hi"}])]
+    assert isinstance(items[-1], RunComplete)
+    assert all(isinstance(i, StepResult) for i in items[:-1])
+
+
+@pytest.mark.asyncio
+async def test_run_to_completion_drains_and_returns_outputs():
+    coordinator = _make_coordinator()
+    outputs = await coordinator.run_to_completion(_simple_graph(), inputs=[{"input_value": "hi"}])
+    assert isinstance(outputs, list)
+
+
+@pytest.mark.asyncio
+async def test_run_to_completion_propagates_executor_errors():
+    coordinator = _make_coordinator()
+    graph = _simple_graph()
+
+    async def boom(*args, **kwargs):  # noqa: ARG001
+        msg = "boom"
+        raise RuntimeError(msg)
+
+    graph._arun_legacy = boom
+
+    with pytest.raises(RuntimeError, match="boom"):
+        await coordinator.run_to_completion(graph, inputs=[{"input_value": "hi"}])
+
+
+@pytest.mark.asyncio
+async def test_run_with_no_executor_registered_raises():
+    from lfx.execution.registry import ExecutorNotFoundError
+
+    coordinator = Coordinator(registry=ExecutorRegistry())
+    with pytest.raises(ExecutorNotFoundError):
+        async for _ in coordinator.run(_simple_graph(), inputs=[]):
+            pass