diff --git a/src/lfx/src/lfx/execution/backends/in_process.py b/src/lfx/src/lfx/execution/backends/in_process.py index 37f327f089..9e645ce8a1 100644 --- a/src/lfx/src/lfx/execution/backends/in_process.py +++ b/src/lfx/src/lfx/execution/backends/in_process.py @@ -27,16 +27,32 @@ class InProcessExecutor(Executor): yield RunComplete(outputs=list(outputs)) return - async for result in graph.async_start( + # Hold a reference to the inner async iterator so we can guarantee + # finalization on consumer cancellation. ``async for`` does NOT call + # ``aclose()`` on its iterator when the loop is interrupted by an exception + # (e.g. ``GeneratorExit`` from this generator being closed). Without the + # explicit try/finally below, the inner generator's ``finally:`` block -- + # which is where event managers, connections, or vertex finalizers run -- + # would only execute on GC, leaking resources for an unpredictable window. + inner = graph.async_start( inputs=opts.get("initial_inputs"), max_iterations=opts.get("max_iterations"), config=opts.get("config"), event_manager=opts.get("event_manager"), reset_output_values=opts.get("reset_output_values", True), fallback_to_env_vars=opts.get("fallback_to_env_vars", False), - ): - yield StepResult(payload=result) - if isinstance(result, Finish): - break + ) + try: + async for result in inner: + yield StepResult(payload=result) + if isinstance(result, Finish): + break - yield RunComplete(outputs=[]) + yield RunComplete(outputs=[]) + finally: + # ``aclose`` is a no-op if the generator already completed. Async + # iterables that aren't generators (test stubs, custom iterators) may + # not implement aclose, so guard with hasattr. + aclose = getattr(inner, "aclose", None) + if aclose is not None: + await aclose() diff --git a/src/lfx/src/lfx/execution/coordinator.py b/src/lfx/src/lfx/execution/coordinator.py index 908b5ef5ec..ea2e219c45 100644 --- a/src/lfx/src/lfx/execution/coordinator.py +++ b/src/lfx/src/lfx/execution/coordinator.py @@ -1,4 +1,21 @@ -"""Coordinator: graph in, streaming step results out.""" +"""Coordinator: graph in, streaming step results out. + +The ``Coordinator`` is the single dispatch point between callers and executors. It +looks up the registered ``Executor`` for the configured ``kind``, partitions the +input graph into ``Unit`` value objects, and forwards each unit's event stream. + +Callers pick the view they want: + +- ``run()``: full seam stream (``StepResult`` + terminal ``RunComplete``). Use when + you need the terminal signal -- e.g. ``run_to_completion``. +- ``stream()``: payloads only. Drops the terminal ``RunComplete`` silently because + every existing payload-consumer in the codebase (``flow_executor``, the CLI, + ``Loop``, ``run/base``) processes events incrementally and has no use for the + terminal envelope. +- ``run_to_completion()``: drains the stream and returns ``RunComplete.outputs``. + Practically useful only with the in-process executor's legacy passthrough; see + the docstring on ``RunComplete.outputs`` for the full semantics. +""" from __future__ import annotations @@ -14,6 +31,8 @@ if TYPE_CHECKING: class Coordinator: + """Routes graph runs through a registered ``Executor`` of the configured kind.""" + def __init__( self, *, @@ -30,11 +49,30 @@ class Coordinator: inputs: list[dict[str, Any]], **runtime_options: Any, ) -> AsyncIterator[StepResult | RunComplete]: + """Stream the seam-typed event sequence for ``graph``. + + Yields the executor's ``StepResult`` items in order followed by the + terminal ``RunComplete``. If you only need payloads, use ``stream()``; + if you only need the final ``outputs`` list, use ``run_to_completion()``. + + ``runtime_options`` are forwarded as-is into ``Unit.runtime_options``; + consult the active executor's documentation for which keys it honors. + """ 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 + # Cascade cleanup: if our consumer aclose()s us, we aclose the executor + # stream, which lets InProcessExecutor's try/finally aclose the underlying + # graph generator. Without this, the executor stream is abandoned on + # consumer aclose and only finalizes when CPython gets around to GC. + inner = executor.execute(unit) + try: + async for item in inner: + yield item + finally: + aclose = getattr(inner, "aclose", None) + if aclose is not None: + await aclose() async def run_to_completion( self, @@ -43,6 +81,16 @@ class Coordinator: inputs: list[dict[str, Any]], **runtime_options: Any, ) -> list[Any]: + """Drain the stream and return ``RunComplete.outputs``. + + Raises ``RuntimeError`` if the executor's stream ends without a terminal + ``RunComplete`` -- a contract violation that every executor must avoid. + + See ``RunComplete.outputs`` for the per-executor semantics of the returned + list. In particular, the in-process executor's streaming path returns an + empty list here; callers that want final outputs from the streaming path + should use ``stream()`` and collect what they need from payloads. + """ async for item in self.run(graph, inputs=inputs, **runtime_options): if isinstance(item, RunComplete): return item.outputs @@ -56,6 +104,17 @@ class Coordinator: inputs: list[dict[str, Any]] | None = None, **runtime_options: Any, ) -> AsyncIterator[Any]: - async for item in self.run(graph, inputs=inputs or [], **runtime_options): - if isinstance(item, StepResult): - yield item.payload + """Yield ``StepResult.payload`` values only. + + The terminal ``RunComplete`` is intentionally NOT yielded -- this helper + exists for consumers that process events incrementally and have no use for + the end-of-run envelope. If you need to know when the run ends, use + ``run()`` directly. + """ + inner = self.run(graph, inputs=inputs or [], **runtime_options) + try: + async for item in inner: + if isinstance(item, StepResult): + yield item.payload + finally: + await inner.aclose() diff --git a/src/lfx/src/lfx/execution/executor.py b/src/lfx/src/lfx/execution/executor.py index 1546c1cf51..8965e17921 100644 --- a/src/lfx/src/lfx/execution/executor.py +++ b/src/lfx/src/lfx/execution/executor.py @@ -1,4 +1,23 @@ -"""Executor abstract base class.""" +"""Executor abstract base class. + +An ``Executor`` is the pluggable backend that turns a ``Unit`` into a stream of +``StepResult`` events terminated by a single ``RunComplete``. The ``Coordinator`` +holds at most one instance per ``kind`` and dispatches every ``run()`` against it, +so implementations must observe the lifecycle contract documented below. + +Authoring a new executor: + +1. Subclass ``Executor`` and set ``kind`` to a unique, stable string. +2. Implement ``execute`` as an ``async def`` generator. Yield zero or more + ``StepResult`` items, then exactly one ``RunComplete``. +3. Register through one of three paths (all equivalent at runtime): + - Explicit: ``get_service(ServiceType.EXECUTOR_SERVICE).register(MyExecutor())`` + - Plugin: ``[project.entry-points."lfx.executors"] my_kind = "pkg.mod:MyExecutor"`` + - Service replacement: ``@register_service(ServiceType.EXECUTOR_SERVICE)`` +4. Validate with the shared contract suite: subclass + ``tests.unit.execution.test_executor_contract.ExecutorContract`` in your tests + and provide the two fixtures. All universal seam guarantees are checked there. +""" from __future__ import annotations @@ -12,8 +31,46 @@ if TYPE_CHECKING: class Executor(ABC): + """Pluggable execution backend for the seam. + + Class attributes: + kind: Unique stable identifier used by ``ExecutorRegistry`` and by + ``settings.executor_kind``. Two executors with the same ``kind`` cannot + coexist in a single registry; entry-point discovery refuses collisions. + + Lifecycle contract: + - A single instance is shared across all runs of its ``kind``. The same + instance may receive concurrent ``execute()`` calls (e.g. parallel API + requests, ``Loop`` subgraphs). Implementations MUST be safe under + concurrent calls and MUST NOT carry per-run state on ``self``. + - ``execute()`` is called once per ``Unit``. Each call MUST return a fresh + async iterator; reusing/caching a generator across calls is a bug. + - The stream MUST emit zero or more ``StepResult`` followed by exactly one + terminal ``RunComplete``. ``Coordinator.run_to_completion`` raises if the + terminal item is missing. + - Exceptions raised inside the generator propagate to the consumer. The + generator's finalizer runs as usual, so any ``async with`` / cleanup code + will execute. Consumers MAY drop the iterator mid-stream (``aclose``); + implementations MUST tolerate this without hanging or leaking external + resources (subprocesses, sockets, locks, file handles). + - ``execute()`` MUST be tolerant of arbitrary keys in ``unit.runtime_options``: + ignore keys the executor doesn't recognize. Keys prefixed with ``_`` are + reserved for executor-internal flags; non-owning executors must ignore them. + """ + kind: ClassVar[str] @abstractmethod def execute(self, unit: Unit) -> AsyncIterator[StepResult | RunComplete]: - """Yield StepResult items, end with a RunComplete.""" + """Run ``unit`` and yield the seam-typed event stream. + + Args: + unit: The work item. ``unit.graph`` is executor-defined (an in-memory + ``Graph`` for in-process; whatever shape a remote/sandboxed executor + expects). ``unit.inputs`` and ``unit.runtime_options`` carry per-run + state. + + Yields: + ``StepResult`` items for mid-run events (executor-defined payloads), + terminated by exactly one ``RunComplete``. + """ diff --git a/src/lfx/src/lfx/execution/types.py b/src/lfx/src/lfx/execution/types.py index b24e0c0f9e..8151fe0d6f 100644 --- a/src/lfx/src/lfx/execution/types.py +++ b/src/lfx/src/lfx/execution/types.py @@ -1,4 +1,11 @@ -"""Value objects for the execution layer.""" +"""Value objects for the execution layer. + +These three dataclasses are the entire vocabulary of the seam: a ``Unit`` flows from +``Coordinator`` into ``Executor.execute()``, which yields a stream of ``StepResult`` +items terminated by a single ``RunComplete``. Every consumer downstream of the +coordinator (``run_to_completion``, ``stream``, ``flow_executor``, the CLI, ``Loop`` +subgraphs) reads only these types, so changes here ripple widely. Keep them small. +""" from __future__ import annotations @@ -8,6 +15,30 @@ from typing import Any @dataclass class Unit: + """A self-contained slice of work handed to an ``Executor``. + + Attributes: + graph: The graph (or graph-shaped object) the executor should run. Executors + decide what they can accept here -- the in-process executor wants a real + ``lfx.Graph``; a remote executor may want a serialized form. There is no + seam-level constraint other than "the executor knows what to do with it." + inputs: Per-run input dicts. Shape is consumer-defined. + runtime_options: A free-form bag of run-scoped options. Stringly typed on + purpose: the seam can't enumerate every option every executor might want + (event managers, session IDs, fallback flags, executor-specific tokens), + and pinning a schema here would force every change into this module. + Conventions worth knowing: + + - Keys starting with ``_`` are reserved for executor-internal flags + (e.g. ``_use_arun_legacy`` selects the legacy passthrough in the + in-process executor). Other executors MUST ignore unknown keys. + - Common keys recognized by the in-process executor: + ``initial_inputs``, ``max_iterations``, ``config``, ``event_manager``, + ``reset_output_values``, ``fallback_to_env_vars``, ``session_id``. + - Other executors are free to define their own keys; document them in + the executor module. + """ + graph: Any inputs: list[dict[str, Any]] = field(default_factory=list) runtime_options: dict[str, Any] = field(default_factory=dict) @@ -15,9 +46,54 @@ class Unit: @dataclass class StepResult: + """A single mid-run event from an executor. + + ``payload`` is intentionally untyped at the seam level. Each executor decides + what mid-run events look like for its backend: + + - ``InProcessExecutor`` yields ``Vertex`` / ``Finish`` instances from + ``Graph.async_start``. + - A remote executor (e.g. stepflow) might yield protobuf status events. + - A sandboxed executor might yield structured ``StepStarted`` / ``StepCompleted`` + records. + + Consumers that need a uniform shape across executors should normalize at the + consumer site (or wrap an executor with an adapter), not push a structured event + type down through the seam. Doing the latter would force every executor to + translate its native event vocabulary into a synthetic one and would mask + information the native vocabulary carries. + + ``Coordinator.stream()`` unwraps ``StepResult.payload`` for callers that want a + raw event stream; ``Coordinator.run()`` yields the full ``StepResult``/`` + RunComplete`` union for callers that want the seam shape. + """ + payload: Any @dataclass class RunComplete: + """Terminal item in an executor stream. + + ``outputs`` is populated only by executors / paths that have a meaningful notion + of "final outputs of the run." In particular: + + - The in-process executor's *legacy* path (selected via + ``runtime_options["_use_arun_legacy"]``) populates ``outputs`` with the + ``list[RunOutputs]`` returned by ``Graph._arun_legacy``. ``Graph.arun`` reads + this via ``Coordinator.run_to_completion``. + - The in-process executor's *streaming* path yields ``RunComplete(outputs=[])``; + consumers of the streaming path are expected to collect what they need from + ``StepResult.payload`` events along the way (this is what ``flow_executor``, + the CLI, ``Loop``, and ``run/base`` already do via ``Coordinator.stream``). + - Other executors MAY populate ``outputs`` with whatever terminal-status event + their backend provides (e.g. a stepflow ``RunCompletedEvent``). + + Bottom line: the *only* universal contract on ``outputs`` is that it is a list. + If you need final values across all executors, collect them from the stream; + if you specifically need the legacy ``list[RunOutputs]`` shape, use + ``Coordinator.run_to_completion`` with ``_use_arun_legacy=True`` against the + in-process executor. + """ + outputs: list[Any] diff --git a/src/lfx/tests/unit/execution/test_cancellation.py b/src/lfx/tests/unit/execution/test_cancellation.py new file mode 100644 index 0000000000..8bb04aceb8 --- /dev/null +++ b/src/lfx/tests/unit/execution/test_cancellation.py @@ -0,0 +1,187 @@ +"""Cancellation and cleanup tests for the execution seam. + +The contract suite covers the bare-minimum ``aclose()`` shape via +``test_consumer_cancellation_does_not_hang_or_leak``. This module exercises the +stricter behaviour the seam needs in production: explicit cleanup at the top of +the consumer chain MUST cascade all the way down to the underlying graph iterator. + +What we deliberately do NOT test here: + +- Implicit GC-based finalization (``async for ... break`` without ``aclosing``, + ``task.cancel()`` then immediate assert). CPython finalizes abandoned async + generators on a GC pass; asyncio runs the registered finalizer on event-loop + shutdown. Neither is observable inside a single test, and asserting either + would test CPython, not us. Consumers that need deterministic cleanup MUST + use ``contextlib.aclosing`` or call ``aclose()`` explicitly. +""" + +from __future__ import annotations + +import asyncio +from contextlib import aclosing + +import pytest +from lfx.execution.backends.in_process import InProcessExecutor +from lfx.execution.coordinator import Coordinator +from lfx.execution.registry import ExecutorRegistry +from lfx.execution.types import StepResult, Unit + + +class _InstrumentedGraph: + """Minimal graph stub whose ``async_start`` runs forever until closed. + + We track entry, items yielded, and finalizer invocation so tests can assert + the iterator was properly closed rather than abandoned. + """ + + def __init__(self) -> None: + self.entered: bool = False + self.finalized: bool = False + self.items_yielded: int = 0 + + async def async_start(self, **_kwargs): + self.entered = True + try: + for i in range(10_000): + self.items_yielded += 1 + yield f"item-{i}" + await asyncio.sleep(0) + finally: + self.finalized = True + + +def _unit(graph: _InstrumentedGraph) -> Unit: + return Unit(graph=graph, inputs=[], runtime_options={}) + + +@pytest.mark.asyncio +async def test_aclose_on_executor_finalizes_underlying_graph_iteration(): + """``aclose()`` on the executor's generator MUST run the underlying graph's finalizer. + + Without the explicit try/finally inside ``InProcessExecutor.execute``, the inner + generator is abandoned and the graph's ``finally:`` block (event managers, + connections, vertex teardown) only runs on a future GC pass. + """ + graph = _InstrumentedGraph() + executor = InProcessExecutor() + stream = executor.execute(_unit(graph)) + + first = await anext(stream) + assert isinstance(first, StepResult) + assert graph.entered + assert not graph.finalized + + await stream.aclose() + assert graph.finalized, "underlying async generator was not finalized on aclose()" + + +@pytest.mark.asyncio +async def test_aclosing_helper_chains_cleanup_through_executor(): + """``contextlib.aclosing`` MUST trigger the same cascade as explicit ``aclose``.""" + graph = _InstrumentedGraph() + executor = InProcessExecutor() + + async with aclosing(executor.execute(_unit(graph))) as stream: + async for _ in stream: + break + + assert graph.finalized + + +@pytest.mark.asyncio +async def test_coordinator_run_cascades_aclose_to_executor_and_graph(): + """Cleanup must propagate through every seam layer. + + Consumer aclose -> Coordinator.run -> Executor.execute -> graph.async_start. + A regression at any layer would leave the graph generator hanging. + """ + graph = _InstrumentedGraph() + registry = ExecutorRegistry() + registry.register(InProcessExecutor()) + coordinator = Coordinator(registry=registry) + + async with aclosing(coordinator.run(graph, inputs=[])) as stream: + first = await anext(stream) + assert first is not None + + assert graph.finalized, "Coordinator.run did not cascade aclose to the graph" + + +@pytest.mark.asyncio +async def test_coordinator_stream_cascades_aclose_to_graph(): + """``Coordinator.stream`` must cascade cleanup the same way ``Coordinator.run`` does.""" + graph = _InstrumentedGraph() + registry = ExecutorRegistry() + registry.register(InProcessExecutor()) + coordinator = Coordinator(registry=registry) + + async with aclosing(coordinator.stream(graph)) as stream: + async for _ in stream: + break + + assert graph.finalized + + +@pytest.mark.asyncio +async def test_cleanup_bounded_by_timeout(): + """Finalization must complete within a bounded window even on a slow CI box. + + A regression that swaps in a non-cancellable resource (blocking subprocess + wait, sync lock) would hang here rather than slow. + """ + graph = _InstrumentedGraph() + executor = InProcessExecutor() + + async def drive(): + gen = executor.execute(_unit(graph)) + await anext(gen) + await gen.aclose() + + await asyncio.wait_for(drive(), timeout=2.0) + assert graph.finalized + + +@pytest.mark.asyncio +async def test_exception_in_consumer_still_finalizes_graph(): + """An exception inside the consumer MUST still trigger cleanup via ``aclosing``.""" + graph = _InstrumentedGraph() + executor = InProcessExecutor() + + class _ConsumerError(RuntimeError): + pass + + async def consume_then_raise() -> None: + async with aclosing(executor.execute(_unit(graph))) as stream: + async for _ in stream: + msg = "downstream bug" + raise _ConsumerError(msg) + + with pytest.raises(_ConsumerError): + await consume_then_raise() + + assert graph.finalized + + +@pytest.mark.asyncio +async def test_concurrent_consumers_each_clean_up_independently(): + """Concurrent runs on one executor instance must clean up independently. + + Regression guard for any future change that introduces shared cleanup state + on the executor (which would couple unrelated consumers). + """ + graph_a = _InstrumentedGraph() + graph_b = _InstrumentedGraph() + executor = InProcessExecutor() + + async def drive(graph: _InstrumentedGraph, items_before_stop: int) -> None: + async with aclosing(executor.execute(_unit(graph))) as stream: + count = 0 + async for _ in stream: + count += 1 + if count >= items_before_stop: + break + + await asyncio.gather(drive(graph_a, 1), drive(graph_b, 3)) + + assert graph_a.finalized + assert graph_b.finalized diff --git a/src/lfx/tests/unit/execution/test_coordinator_e2e.py b/src/lfx/tests/unit/execution/test_coordinator_e2e.py new file mode 100644 index 0000000000..790d720136 --- /dev/null +++ b/src/lfx/tests/unit/execution/test_coordinator_e2e.py @@ -0,0 +1,156 @@ +"""End-to-end seam tests: a real Graph driven through Coordinator. + +Tests in ``test_coordinator.py`` check Coordinator behaviour at a behavioural level +but defer most of the heavy lifting to the executor or graph layer. This module +verifies the seam contract by running a real ``lfx.Graph`` through the full +``Coordinator -> ExecutorRegistry -> InProcessExecutor -> graph.async_start`` +chain and asserting the event stream looks like what consumers actually receive: + +- ``stream()`` payloads are vertex-shaped (not the seam envelopes). +- ``run()`` yields ``StepResult`` items in order then exactly one ``RunComplete``. +- ``stream()`` never leaks ``RunComplete`` through the payload channel. +- Switching the registered executor changes what runs (the dispatch is real, not + hard-wired to ``InProcessExecutor``). +- Lifecycle: registry get() returns the same instance for repeated calls. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import pytest +from lfx.execution.backends.in_process import InProcessExecutor +from lfx.execution.coordinator import Coordinator +from lfx.execution.executor import Executor +from lfx.execution.registry import ExecutorRegistry +from lfx.execution.types import RunComplete, StepResult, Unit + +if TYPE_CHECKING: + from collections.abc import AsyncIterator + + +def _make_coordinator() -> Coordinator: + registry = ExecutorRegistry() + registry.register(InProcessExecutor()) + return Coordinator(registry=registry) + + +@pytest.mark.asyncio +async def test_run_yields_step_results_then_terminal_run_complete(simple_graph): + coordinator = _make_coordinator() + items = [item async for item in coordinator.run(simple_graph, inputs=[{"input_value": "hi"}])] + + assert items, "real Graph drove zero events through Coordinator" + assert isinstance(items[-1], RunComplete), "stream must end with RunComplete" + assert sum(isinstance(i, RunComplete) for i in items) == 1, "exactly one RunComplete expected" + assert all(isinstance(i, StepResult) for i in items[:-1]), "pre-terminal items must all be StepResult" + + +@pytest.mark.asyncio +async def test_stream_yields_payloads_only_no_run_complete_leak(simple_graph): + """``Coordinator.stream`` MUST NOT yield ``RunComplete`` instances. + + It is the payload-only helper; the terminal envelope from the executor must + be stripped before reaching the consumer. + """ + coordinator = _make_coordinator() + payloads = [p async for p in coordinator.stream(simple_graph, initial_inputs=None)] + + assert payloads, "stream produced no payloads from a real graph" + assert not any(isinstance(p, RunComplete) for p in payloads), ( + "Coordinator.stream() leaked RunComplete into the payload channel" + ) + assert not any(isinstance(p, StepResult) for p in payloads), ( + "Coordinator.stream() leaked StepResult envelopes; payloads should be unwrapped" + ) + + +@pytest.mark.asyncio +async def test_payload_shape_is_vertex_or_finish_for_in_process(simple_graph): + """Pin the in-process payload contract. + + Changes to graph.async_start can't silently rename or restructure the + events consumers depend on without breaking this test. + """ + from lfx.graph.graph.constants import Finish + + coordinator = _make_coordinator() + payloads = [p async for p in coordinator.stream(simple_graph, initial_inputs=None)] + + vertex_ids = [getattr(p, "vertex", None) and p.vertex.id for p in payloads] + assert "chat_input" in vertex_ids + assert "chat_output" in vertex_ids + assert any(isinstance(p, Finish) for p in payloads), "expected a terminal Finish payload" + + +@pytest.mark.asyncio +async def test_dispatch_routes_to_configured_kind(simple_graph): + """Changing ``executor_kind`` changes where execution lands. + + Proves the registry lookup is the real dispatch, not a hard-coded reference. + """ + + class _Marker(Executor): + kind = "marker" + ran: bool = False + + async def execute(self, unit: Unit) -> AsyncIterator[StepResult | RunComplete]: # noqa: ARG002 + type(self).ran = True + yield StepResult(payload="sentinel") + yield RunComplete(outputs=["sentinel"]) + + registry = ExecutorRegistry() + registry.register(InProcessExecutor()) + registry.register(_Marker()) + coordinator = Coordinator(registry=registry, executor_kind="marker") + + items = [item async for item in coordinator.run(simple_graph, inputs=[])] + assert _Marker.ran is True + assert items[0].payload == "sentinel" + assert isinstance(items[-1], RunComplete) + assert items[-1].outputs == ["sentinel"] + + +@pytest.mark.asyncio +async def test_run_to_completion_returns_outputs_from_legacy_path(simple_graph): + """``run_to_completion`` with the legacy flag returns populated outputs. + + This is the only path through the seam where ``RunComplete.outputs`` carries + a meaningful value (see ``RunComplete`` docstring). + """ + coordinator = _make_coordinator() + outputs = await coordinator.run_to_completion( + simple_graph, + inputs=[{"input_value": "hi"}], + _use_arun_legacy=True, + ) + assert isinstance(outputs, list) + + +@pytest.mark.asyncio +async def test_registry_returns_same_executor_instance_across_runs(simple_graph): + """The seam shares one executor instance per kind across all runs. + + Two sequential ``coordinator.run`` calls MUST resolve to the same instance; + otherwise executor authors couldn't rely on instance-level setup. + """ + instances: list[Executor] = [] + real_inprocess = InProcessExecutor() + + class _Recorder(Executor): + kind = "recorder" + + async def execute(self, unit: Unit) -> AsyncIterator[StepResult | RunComplete]: + instances.append(self) + async for item in real_inprocess.execute(unit): + yield item + + registry = ExecutorRegistry() + registry.register(_Recorder()) + coordinator = Coordinator(registry=registry, executor_kind="recorder") + + [_ async for _ in coordinator.run(simple_graph, inputs=[{"input_value": "a"}])] + [_ async for _ in coordinator.run(simple_graph, inputs=[{"input_value": "b"}])] + + assert len(instances) == 2 + assert instances[0] is instances[1], "registry must return the same instance per kind" diff --git a/src/lfx/tests/unit/execution/test_executor_contract.py b/src/lfx/tests/unit/execution/test_executor_contract.py new file mode 100644 index 0000000000..d62befdc20 --- /dev/null +++ b/src/lfx/tests/unit/execution/test_executor_contract.py @@ -0,0 +1,228 @@ +"""Reusable contract test suite for any ``lfx.execution.Executor`` implementation. + +The seam has a small but load-bearing set of behaviors that any executor MUST satisfy +to be safely swappable behind ``Coordinator``. This module spells them out as a +parametrizable suite so: + +- The built-in ``InProcessExecutor`` is checked here against the same suite the + external executors (stepflow, future remote/sandbox, etc.) will be checked against. +- Authors of new executors can subclass ``ExecutorContract`` in their own project, + override the two fixtures, and inherit the full battery of behavioural assertions + rather than re-deriving them from the docstring of the ABC. + +The contract intentionally does NOT pin executor-specific event shapes -- only the +universal seam guarantees. If a future executor needs to add stricter checks for its +own payloads, those go in an executor-specific test file, not here. +""" + +from __future__ import annotations + +import asyncio +import inspect +from collections.abc import AsyncIterator, Callable +from typing import TYPE_CHECKING + +import pytest +from lfx.execution.types import RunComplete, StepResult, Unit + +if TYPE_CHECKING: + from lfx.execution.executor import Executor + + +class ExecutorContract: + """Base class of behavioural assertions every Executor must satisfy. + + Subclasses MUST override two fixtures: + + - ``contract_executor``: returns a freshly constructed ``Executor`` instance. + - ``contract_unit_factory``: returns a callable ``() -> Unit`` that produces a + fresh, runnable ``Unit`` each call. The unit should drive a minimal but real + execution -- not a stub -- so the test actually exercises the executor's + end-to-end path. + + Tests are designed to be independent: each test creates its own unit through + the factory so subclasses don't have to worry about state bleed between cases. + """ + + @pytest.fixture + def contract_executor(self) -> Executor: + msg = "Subclasses must override the contract_executor fixture" + raise NotImplementedError(msg) + + @pytest.fixture + def contract_unit_factory(self) -> Callable[[], Unit]: + msg = "Subclasses must override the contract_unit_factory fixture" + raise NotImplementedError(msg) + + # --- Universal seam guarantees ---------------------------------------------------- + + def test_kind_is_a_non_empty_string(self, contract_executor: Executor) -> None: + """``kind`` is a ClassVar on the ABC; it must be set and non-empty.""" + assert isinstance(contract_executor.kind, str) + assert contract_executor.kind, "Executor.kind must be a non-empty string" + + def test_execute_returns_async_iterator( + self, + contract_executor: Executor, + contract_unit_factory: Callable[[], Unit], + ) -> None: + """``execute(unit)`` MUST return an async iterator. + + The ABC types it as such; consumers (Coordinator.run, Coordinator.stream, + run_to_completion) rely on it. + """ + stream = contract_executor.execute(contract_unit_factory()) + assert isinstance(stream, AsyncIterator), ( + f"{type(contract_executor).__name__}.execute() returned a " + f"{type(stream).__name__}, expected AsyncIterator. Consumers can't iterate " + "this. Common cause: forgetting to make the method an `async def` " + "generator." + ) + + @pytest.mark.asyncio + async def test_stream_ends_with_exactly_one_run_complete( + self, + contract_executor: Executor, + contract_unit_factory: Callable[[], Unit], + ) -> None: + """Streams emit zero-or-more ``StepResult`` items terminated by one ``RunComplete``.""" + items = [item async for item in contract_executor.execute(contract_unit_factory())] + + assert items, "executor yielded nothing -- a RunComplete is mandatory" + assert isinstance(items[-1], RunComplete), ( + f"final item was {type(items[-1]).__name__}; the seam requires a terminal " + "RunComplete so run_to_completion() and downstream collectors can detect end-of-run." + ) + + run_completes = [i for i in items if isinstance(i, RunComplete)] + assert len(run_completes) == 1, ( + f"executor yielded {len(run_completes)} RunComplete items; the seam allows " + "at most one (and requires exactly one)." + ) + + non_terminal = items[:-1] + bad = [i for i in non_terminal if not isinstance(i, StepResult)] + assert not bad, ( + f"all pre-terminal items must be StepResult; got: {[type(i).__name__ for i in bad]}. " + "Coordinator.stream() relies on this to filter; mixing other types makes " + "the stream untyped at the seam boundary." + ) + + @pytest.mark.asyncio + async def test_executor_is_reusable_across_runs( + self, + contract_executor: Executor, + contract_unit_factory: Callable[[], Unit], + ) -> None: + """An Executor MUST be reusable: state that carries between runs is a bug. + + The registry returns the same instance for every ``coordinator.run()`` call, + so any internal state pinned to ``self`` will silently corrupt later runs. + """ + first = [item async for item in contract_executor.execute(contract_unit_factory())] + second = [item async for item in contract_executor.execute(contract_unit_factory())] + + assert isinstance(first[-1], RunComplete) + assert isinstance(second[-1], RunComplete) + # Same shape across runs proves no carry-over short-circuit. We don't assert + # equality of payloads because some executors include timestamps / run IDs. + assert len(first) >= 1 + assert len(second) >= 1 + + @pytest.mark.asyncio + async def test_concurrent_runs_on_one_instance_are_isolated( + self, + contract_executor: Executor, + contract_unit_factory: Callable[[], Unit], + ) -> None: + """Concurrent runs on one executor instance MUST be isolated. + + Coordinator may launch concurrent runs (Loop subgraphs, parallel API + requests). An executor that crosses state between them is broken. + """ + + async def drain() -> list[object]: + return [item async for item in contract_executor.execute(contract_unit_factory())] + + a, b = await asyncio.gather(drain(), drain()) + assert isinstance(a[-1], RunComplete) + assert isinstance(b[-1], RunComplete) + + @pytest.mark.asyncio + async def test_consumer_cancellation_does_not_hang_or_leak( + self, + contract_executor: Executor, + contract_unit_factory: Callable[[], Unit], + ) -> None: + """Consumer cancellation via ``aclose()`` MUST finalize within a bounded window. + + The executor's async generator is closed by the consumer; any internal + context managers, subprocesses, or sockets MUST be released. A strict + timeout surfaces regressions as test failures rather than CI hangs. + """ + gen = contract_executor.execute(contract_unit_factory()) + + try: + first = await asyncio.wait_for(anext(gen), timeout=10.0) + except StopAsyncIteration: + # Executor produced no items at all -- still a valid stream (just an empty one). + return + + assert isinstance(first, (StepResult, RunComplete)), ( + f"first yielded item was {type(first).__name__}; seam requires StepResult or RunComplete" + ) + + # Drop the iterator. If the executor holds a subprocess / network / lock, this + # is the moment of truth: aclose() must finalize without blocking forever. + await asyncio.wait_for(gen.aclose(), timeout=10.0) + + @pytest.mark.asyncio + async def test_execute_signature_takes_exactly_one_unit(self, contract_executor: Executor) -> None: + """Pin the ABC signature so executors don't drift away from a single ``Unit`` parameter. + + ``execute(*args, **kwargs)`` shapes bypass the ``Unit`` value object and + make the seam opaque to introspection. + """ + sig = inspect.signature(contract_executor.execute) + params = [p for p in sig.parameters.values() if p.name != "self"] + assert len(params) == 1, ( + f"{type(contract_executor).__name__}.execute must accept exactly one positional " + f"Unit argument; got signature: {sig}" + ) + + +class TestInProcessExecutorContract(ExecutorContract): + """Run the contract suite against the built-in InProcessExecutor. + + This is the reference implementation: if the contract suite is wrong, this is + where it will fail first. Stepflow / remote / sandbox executors are exercised + against the same suite in their own test files. + """ + + @pytest.fixture + def contract_executor(self) -> Executor: + from lfx.execution.backends.in_process import InProcessExecutor + + return InProcessExecutor() + + @pytest.fixture + def contract_unit_factory(self, simple_graph) -> Callable[[], Unit]: # noqa: ARG002 + # ``simple_graph`` is a pytest fixture from conftest.py. We don't capture it + # by closure across runs because Graph instances accumulate state; instead we + # rebuild on every call so each unit is fresh. + from lfx.components.input_output import ChatInput, ChatOutput + from lfx.graph import Graph + from lfx.schema.schema import InputValueRequest + + def factory() -> Unit: + ci = ChatInput(_id="chat_input") + ci.set(should_store_message=False) + co = ChatOutput(input_value="test", _id="chat_output") + co.set(sender_name=ci.message_response) + return Unit( + graph=Graph(ci, co), + inputs=[], + runtime_options={"initial_inputs": InputValueRequest(input_value="hi")}, + ) + + return factory