mirror of
https://github.com/langflow-ai/langflow.git
synced 2026-07-23 23:13:58 +08:00
test(execution): contract suite, cancellation guarantees, seam docstrings
Robustness pass on the execution seam after using it to integrate an external
executor end-to-end. Four things this pins down:
1. Reusable executor contract suite
- tests/unit/execution/test_executor_contract.py
- ExecutorContract with 7 universal seam guarantees: kind shape, execute()
returns an AsyncIterator, stream ends with exactly one RunComplete,
instance is reusable across runs, concurrent runs are isolated, consumer
cancellation does not hang or leak, execute() signature is stable.
- TestInProcessExecutorContract subclasses and provides fixtures. Downstream
executors (stepflow, future remote/sandbox) get the same battery by
subclassing and overriding two fixtures.
2. Cancellation/cleanup correctness
- tests/unit/execution/test_cancellation.py
- InProcessExecutor wraps graph.async_start() in try/finally with explicit
aclose so consumer aclose cascades to the underlying graph generator.
Previously the inner generator was abandoned and only finalized on a GC
pass: a real leak for executors holding subprocesses or sockets.
- Coordinator.run and Coordinator.stream propagate aclose the same way so
the full chain (consumer -> coordinator -> executor -> graph) cleans up
deterministically.
3. Documented seam contracts
- Unit.runtime_options: free-form bag, "_"-prefixed keys reserved for
executor-internal flags, common in-process keys listed.
- StepResult.payload: untyped on purpose; each executor defines its own
event vocabulary; consumers normalize at consumption site.
- RunComplete.outputs: only the legacy in-process path populates it;
streaming consumers should collect from StepResult.payload.
- Executor ABC: lifecycle (shared instance, must be reusable, must tolerate
concurrent execute() calls and consumer aclose).
4. Coordinator-routed E2E suite
- tests/unit/execution/test_coordinator_e2e.py
- Real Graph through Coordinator -> registry -> InProcessExecutor chain.
- Asserts payload shape, RunComplete never leaks to stream() consumers,
dispatch routes to configured kind, registry returns same instance.
Full execution suite: 63 passing. flow_executor coordinator test: passing.
This commit is contained in:
@ -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()
|
||||
|
||||
@ -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()
|
||||
|
||||
@ -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``.
|
||||
"""
|
||||
|
||||
@ -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]
|
||||
|
||||
187
src/lfx/tests/unit/execution/test_cancellation.py
Normal file
187
src/lfx/tests/unit/execution/test_cancellation.py
Normal file
@ -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
|
||||
156
src/lfx/tests/unit/execution/test_coordinator_e2e.py
Normal file
156
src/lfx/tests/unit/execution/test_coordinator_e2e.py
Normal file
@ -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"
|
||||
228
src/lfx/tests/unit/execution/test_executor_contract.py
Normal file
228
src/lfx/tests/unit/execution/test_executor_contract.py
Normal file
@ -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
|
||||
Reference in New Issue
Block a user