diff --git a/src/backend/base/langflow/agentic/services/flow_executor.py b/src/backend/base/langflow/agentic/services/flow_executor.py index 3499239ef5..4e1772508c 100644 --- a/src/backend/base/langflow/agentic/services/flow_executor.py +++ b/src/backend/base/langflow/agentic/services/flow_executor.py @@ -12,6 +12,7 @@ from typing import TYPE_CHECKING, Any from fastapi import HTTPException from lfx.cli.script_loader import extract_structured_result from lfx.events.event_manager import EventManager, create_default_event_manager +from lfx.execution import get_default_coordinator from lfx.log.logger import logger from lfx.schema.schema import InputValueRequest from lfx.utils.flow_validation import CustomComponentValidationError @@ -59,7 +60,12 @@ async def _run_graph_with_events( graph.prepare() inputs = InputValueRequest(input_value=input_value) if input_value else None - results = [result async for result in graph.async_start(inputs=inputs, event_manager=event_manager)] + results = [ + payload + async for payload in get_default_coordinator().stream( + graph, initial_inputs=inputs, event_manager=event_manager + ) + ] execution_result.result = extract_structured_result(results) except Exception as e: # noqa: BLE001 execution_result.error = e @@ -131,7 +137,7 @@ async def execute_flow_file( graph.prepare() inputs = InputValueRequest(input_value=input_value) if input_value else None - results = [result async for result in graph.async_start(inputs=inputs)] + results = [payload async for payload in get_default_coordinator().stream(graph, initial_inputs=inputs)] flow_result = extract_structured_result(results) except HTTPException: raise diff --git a/src/backend/base/langflow/services/utils.py b/src/backend/base/langflow/services/utils.py index 5b8752e9d4..14b623a65e 100644 --- a/src/backend/base/langflow/services/utils.py +++ b/src/backend/base/langflow/services/utils.py @@ -487,6 +487,7 @@ def register_all_service_factories() -> None: from lfx.services.schema import ServiceType service_manager = get_service_manager() + from lfx.services.executor import factory as executor_factory from lfx.services.mcp_composer import factory as mcp_composer_factory from lfx.services.settings import factory as settings_factory @@ -542,6 +543,7 @@ def register_all_service_factories() -> None: ) service_manager.register_factory(authorization_factory.AuthorizationServiceFactory()) service_manager.register_factory(mcp_composer_factory.MCPComposerServiceFactory()) + service_manager.register_factory(executor_factory.ExecutorServiceFactory()) service_manager.set_factory_registered() diff --git a/src/backend/tests/unit/agentic/services/test_flow_executor_uses_coordinator.py b/src/backend/tests/unit/agentic/services/test_flow_executor_uses_coordinator.py new file mode 100644 index 0000000000..288b2fe6ca --- /dev/null +++ b/src/backend/tests/unit/agentic/services/test_flow_executor_uses_coordinator.py @@ -0,0 +1,78 @@ +"""Verify flow_executor.py routes through the executor coordinator.""" + +from __future__ import annotations + +import asyncio +from typing import Any + +import pytest +from lfx.execution import ( + Coordinator, + ExecutorRegistry, + reset_default_coordinator, + set_default_coordinator, +) +from lfx.execution.executor import Executor +from lfx.execution.types import RunComplete, StepResult + + +class _Recording(Executor): + kind = "in-process" + + def __init__(self) -> None: + self.units_seen: list[Any] = [] + + async def execute(self, unit): + self.units_seen.append(unit) + yield StepResult(payload=object()) + yield RunComplete(outputs=[]) + + +@pytest.fixture(autouse=True) +def _reset_singleton(): + reset_default_coordinator() + yield + reset_default_coordinator() + + +@pytest.mark.asyncio +async def test_run_graph_uses_default_coordinator(): + from dataclasses import dataclass, field + + from langflow.agentic.services.flow_executor import _run_graph_with_events + from langflow.agentic.services.flow_types import FlowExecutionResult + + recording = _Recording() + registry = ExecutorRegistry() + registry.register(recording) + set_default_coordinator(Coordinator(registry=registry)) + + @dataclass + class _FakeGraph: + flow_id: str | None = None + flow_name: str | None = None + user_id: str | None = None + session_id: str | None = None + context: dict = field(default_factory=dict) + + def prepare(self): + pass + + queue: asyncio.Queue = asyncio.Queue() + execution_result = FlowExecutionResult() + + from lfx.events.event_manager import create_default_event_manager + + em = create_default_event_manager(asyncio.Queue()) + await _run_graph_with_events( + graph=_FakeGraph(), + input_value="hi", + global_variables=None, + user_id=None, + session_id=None, + event_manager=em, + event_queue=queue, + execution_result=execution_result, + ) + + assert len(recording.units_seen) == 1 diff --git a/src/backend/tests/unit/components/flow_controls/test_loop.py b/src/backend/tests/unit/components/flow_controls/test_loop.py index 1cc1fb219a..11034395d1 100644 --- a/src/backend/tests/unit/components/flow_controls/test_loop.py +++ b/src/backend/tests/unit/components/flow_controls/test_loop.py @@ -383,7 +383,7 @@ class TestLoopComponentSubgraphExecution: event_manager_received = [] # Create an async generator that yields mock results - async def mock_async_start(event_manager=None): + async def mock_async_start(event_manager=None, **_kwargs): event_manager_received.append(event_manager) yield MagicMock(valid=True, result_dict={"outputs": {}}) diff --git a/src/lfx/src/lfx/base/flow_controls/loop_utils.py b/src/lfx/src/lfx/base/flow_controls/loop_utils.py index 575ededb97..77c63399f7 100644 --- a/src/lfx/src/lfx/base/flow_controls/loop_utils.py +++ b/src/lfx/src/lfx/base/flow_controls/loop_utils.py @@ -1,8 +1,10 @@ """Utility functions for loop component execution.""" from collections import deque +from contextlib import aclosing from typing import TYPE_CHECKING +from lfx.execution import get_default_coordinator from lfx.schema.data import Data if TYPE_CHECKING: @@ -266,12 +268,17 @@ async def execute_loop_body( # Execute subgraph and collect results # Pass event_manager so UI receives events from subgraph execution results = [] - async for result in iteration_subgraph.async_start(event_manager=event_manager): - results.append(result) - # Stop all on error (as per design decision) - if hasattr(result, "valid") and not result.valid: - msg = f"Error in loop iteration: {result}" - raise RuntimeError(msg) + # aclosing guarantees the stream is finalized even when we raise mid-iteration + # on an invalid result, so the underlying subgraph generator's cleanup runs. + async with aclosing( + get_default_coordinator().stream(iteration_subgraph, event_manager=event_manager) + ) as stream: + async for result in stream: + results.append(result) + # Stop all on error (as per design decision) + if hasattr(result, "valid") and not result.valid: + msg = f"Error in loop iteration: {result}" + raise RuntimeError(msg) # Extract output from final result output = extract_loop_output(results, end_vertex_id) diff --git a/src/lfx/src/lfx/cli/common.py b/src/lfx/src/lfx/cli/common.py index 5500984f1a..fe11590664 100644 --- a/src/lfx/src/lfx/cli/common.py +++ b/src/lfx/src/lfx/cli/common.py @@ -29,6 +29,7 @@ from lfx.cli.script_loader import ( find_graph_variable, load_graph_from_script, ) +from lfx.execution import get_default_coordinator from lfx.load import load_flow_from_json from lfx.run._defaults import apply_run_defaults, resolve_fallback_to_env_vars from lfx.schema.schema import InputValueRequest @@ -360,7 +361,12 @@ async def execute_graph_with_capture( try: sys.stdout = captured_stdout sys.stderr = captured_stderr - results = [result async for result in graph.async_start(inputs, fallback_to_env_vars=fallback_to_env_vars)] + results = [ + payload + async for payload in get_default_coordinator().stream( + graph, initial_inputs=inputs, fallback_to_env_vars=fallback_to_env_vars + ) + ] except Exception as exc: # Capture any error output that was written to stderr error_output = captured_stderr.getvalue() diff --git a/src/lfx/src/lfx/execution/__init__.py b/src/lfx/src/lfx/execution/__init__.py new file mode 100644 index 0000000000..7ee48a03aa --- /dev/null +++ b/src/lfx/src/lfx/execution/__init__.py @@ -0,0 +1,64 @@ +"""Public entry points for the execution layer. + +The execution coordinator and registry are owned by the :class:`ExecutorService` +(``lfx.services.executor``). The helpers below are thin convenience wrappers that +resolve the service through the standard service manager so callers do not need to +know about the service plumbing. +""" + +from __future__ import annotations + +from lfx.execution.backends.in_process import InProcessExecutor +from lfx.execution.coordinator import Coordinator +from lfx.execution.executor import Executor +from lfx.execution.partitioner import identity_partition +from lfx.execution.registry import ExecutorNotFoundError, ExecutorRegistry +from lfx.execution.types import RunComplete, StepResult, Unit + + +def _get_executor_service(): + from lfx.services.deps import get_service + from lfx.services.schema import ServiceType + + service = get_service(ServiceType.EXECUTOR_SERVICE) + if service is None: + msg = "ExecutorService is not available; check earlier logs for the underlying init failure." + raise RuntimeError(msg) + return service + + +def get_default_registry() -> ExecutorRegistry: + return _get_executor_service().registry + + +def get_default_coordinator() -> Coordinator: + return _get_executor_service().coordinator + + +def set_default_coordinator(coordinator: Coordinator) -> None: + _get_executor_service().set_coordinator(coordinator) + + +def reset_default_coordinator() -> None: + """Drop the executor service so the next access rebuilds registry + coordinator. For tests.""" + from lfx.services.manager import get_service_manager + from lfx.services.schema import ServiceType + + get_service_manager().update(ServiceType.EXECUTOR_SERVICE) + + +__all__ = [ + "Coordinator", + "Executor", + "ExecutorNotFoundError", + "ExecutorRegistry", + "InProcessExecutor", + "RunComplete", + "StepResult", + "Unit", + "get_default_coordinator", + "get_default_registry", + "identity_partition", + "reset_default_coordinator", + "set_default_coordinator", +] diff --git a/src/lfx/src/lfx/execution/backends/__init__.py b/src/lfx/src/lfx/execution/backends/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/src/lfx/src/lfx/execution/backends/in_process.py b/src/lfx/src/lfx/execution/backends/in_process.py new file mode 100644 index 0000000000..9e645ce8a1 --- /dev/null +++ b/src/lfx/src/lfx/execution/backends/in_process.py @@ -0,0 +1,58 @@ +"""In-process executor: wraps today's graph runner.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from lfx.execution.executor import Executor +from lfx.execution.types import RunComplete, StepResult +from lfx.graph.graph.constants import Finish + +if TYPE_CHECKING: + from collections.abc import AsyncIterator + + from lfx.execution.types import Unit + + +class InProcessExecutor(Executor): + kind = "in-process" + + async def execute(self, unit: Unit) -> AsyncIterator[StepResult | RunComplete]: + graph = unit.graph + opts = unit.runtime_options + + if opts.get("_use_arun_legacy") and hasattr(graph, "_arun_legacy"): + legacy_kwargs = {k: v for k, v in opts.items() if not k.startswith("_")} + outputs = await graph._arun_legacy(inputs=unit.inputs, **legacy_kwargs) # noqa: SLF001 + yield RunComplete(outputs=list(outputs)) + return + + # 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), + ) + try: + async for result in inner: + yield StepResult(payload=result) + if isinstance(result, Finish): + break + + 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 new file mode 100644 index 0000000000..7f5f15e376 --- /dev/null +++ b/src/lfx/src/lfx/execution/coordinator.py @@ -0,0 +1,128 @@ +"""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 + +from contextlib import aclosing +from typing import TYPE_CHECKING, Any + +from lfx.execution.partitioner import identity_partition +from lfx.execution.types import RunComplete, StepResult + +if TYPE_CHECKING: + from collections.abc import AsyncIterator + + from lfx.execution.registry import ExecutorRegistry + + +class Coordinator: + """Routes graph runs through a registered ``Executor`` of the configured kind.""" + + 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]: + """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. + + ``identity_partition`` currently yields exactly one ``Unit``, so this emits a + single terminal ``RunComplete``. A future partitioner returning multiple units + would emit one ``RunComplete`` per unit; consumers relying on a single terminal + envelope (notably ``run_to_completion``) would need updating first. + """ + units = identity_partition(graph, inputs=inputs, runtime_options=runtime_options) + executor = self._registry.get(self._executor_kind) + for unit in units: + # 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, + graph: Any, + *, + 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 with aclosing(self.run(graph, inputs=inputs, **runtime_options)) as stream: + async for item in stream: + if isinstance(item, RunComplete): + return item.outputs + msg = "Executor stream ended without a RunComplete" + raise RuntimeError(msg) + + async def stream( + self, + graph: Any, + **runtime_options: Any, + ) -> AsyncIterator[Any]: + """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. + + Note there is no ``inputs`` parameter here. The in-process executor's + streaming path reads its inputs from ``runtime_options["initial_inputs"]``, + not from the seam-level ``inputs`` list (which only the legacy passthrough + in ``run_to_completion`` consumes). Pass ``initial_inputs=...`` to feed a + streaming run. + """ + async with aclosing(self.run(graph, inputs=[], **runtime_options)) as inner: + async for item in inner: + if isinstance(item, StepResult): + yield item.payload diff --git a/src/lfx/src/lfx/execution/executor.py b/src/lfx/src/lfx/execution/executor.py new file mode 100644 index 0000000000..8965e17921 --- /dev/null +++ b/src/lfx/src/lfx/execution/executor.py @@ -0,0 +1,76 @@ +"""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 + +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, ClassVar + +if TYPE_CHECKING: + from collections.abc import AsyncIterator + + from lfx.execution.types import RunComplete, StepResult, Unit + + +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]: + """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/partitioner.py b/src/lfx/src/lfx/execution/partitioner.py new file mode 100644 index 0000000000..3322506698 --- /dev/null +++ b/src/lfx/src/lfx/execution/partitioner.py @@ -0,0 +1,16 @@ +"""Identity partitioner: one unit per graph.""" + +from __future__ import annotations + +from typing import Any + +from lfx.execution.types import Unit + + +def identity_partition( + graph: Any, + *, + inputs: list[dict[str, Any]], + runtime_options: dict[str, Any] | None = None, +) -> list[Unit]: + return [Unit(graph=graph, inputs=inputs, runtime_options=runtime_options or {})] diff --git a/src/lfx/src/lfx/execution/registry.py b/src/lfx/src/lfx/execution/registry.py new file mode 100644 index 0000000000..e9d7332085 --- /dev/null +++ b/src/lfx/src/lfx/execution/registry.py @@ -0,0 +1,46 @@ +"""Executor registry.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from lfx.execution.executor import Executor + + +class ExecutorNotFoundError(LookupError): + pass + + +class ExecutorKindCollisionError(ValueError): + """Raised when registering an executor whose kind is already registered and replace=False.""" + + +class ExecutorRegistry: + def __init__(self) -> None: + self._by_kind: dict[str, Executor] = {} + + def register(self, executor: Executor, *, replace: bool = True) -> None: + """Register an executor by its kind. + + Args: + executor: The executor to register. + replace: If True (default), overwrite any existing executor with the same kind. + If False, raise :class:`ExecutorKindCollisionError` instead. Discovery paths + pass ``replace=False`` so an installed package cannot silently replace a + pre-registered executor (notably the built-in ``in-process``). + """ + if not replace and executor.kind in self._by_kind: + msg = f"Executor kind={executor.kind!r} is already registered" + raise ExecutorKindCollisionError(msg) + self._by_kind[executor.kind] = executor + + def has(self, kind: str) -> bool: + return kind in self._by_kind + + def get(self, kind: str) -> Executor: + try: + return self._by_kind[kind] + except KeyError as exc: + msg = f"No executor registered for kind={kind!r}" + raise ExecutorNotFoundError(msg) from exc diff --git a/src/lfx/src/lfx/execution/types.py b/src/lfx/src/lfx/execution/types.py new file mode 100644 index 0000000000..d894d79641 --- /dev/null +++ b/src/lfx/src/lfx/execution/types.py @@ -0,0 +1,100 @@ +"""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 + +from dataclasses import dataclass, field +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`` is + honored only by the legacy passthrough path; the streaming path drops it. + - 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) + + +@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/src/lfx/graph/graph/base.py b/src/lfx/src/lfx/graph/graph/base.py index c3a12449b8..1400650f57 100644 --- a/src/lfx/src/lfx/graph/graph/base.py +++ b/src/lfx/src/lfx/graph/graph/base.py @@ -878,25 +878,34 @@ class Graph: fallback_to_env_vars: bool = False, event_manager: EventManager | None = None, ) -> list[RunOutputs]: - """Runs the graph with the given inputs. + """Runs the graph with the given inputs via the configured executor.""" + from lfx.execution import get_default_coordinator - Args: - inputs (list[Dict[str, str]]): The input values for the graph. - inputs_components (Optional[list[list[str]]], optional): Components to run for the inputs. Defaults to None. - types (Optional[list[Optional[InputType]]], optional): The types of the inputs. Defaults to None. - outputs (Optional[list[str]], optional): The outputs to retrieve from the graph. Defaults to None. - session_id (Optional[str], optional): The session ID for the graph. Defaults to None. - stream (bool, optional): Whether to stream the results or not. Defaults to False. - fallback_to_env_vars (bool, optional): Whether to fallback to environment variables. Defaults to False. - event_manager (EventManager | None): The event manager for the graph. + return await get_default_coordinator().run_to_completion( + self, + inputs=inputs, + inputs_components=inputs_components, + types=types, + outputs=outputs, + session_id=session_id, + stream=stream, + fallback_to_env_vars=fallback_to_env_vars, + event_manager=event_manager, + _use_arun_legacy=True, + ) - Returns: - List[RunOutputs]: The outputs of the graph. - """ - # inputs is {"message": "Hello, world!"} - # we need to go through self.inputs and update the self.raw_params - # of the vertices that are inputs - # if the value is a list, we need to run multiple times + async def _arun_legacy( + self, + inputs: list[dict[str, str]], + *, + inputs_components: list[list[str]] | None = None, + types: list[InputType | None] | None = None, + outputs: list[str] | None = None, + session_id: str | None = None, + stream: bool = False, + fallback_to_env_vars: bool = False, + event_manager: EventManager | None = None, + ) -> list[RunOutputs]: vertex_outputs = [] if not isinstance(inputs, list): inputs = [inputs] diff --git a/src/lfx/src/lfx/run/base.py b/src/lfx/src/lfx/run/base.py index e8755611e1..662beb9d27 100644 --- a/src/lfx/src/lfx/run/base.py +++ b/src/lfx/src/lfx/run/base.py @@ -15,6 +15,7 @@ from lfx.cli.script_loader import ( load_graph_from_script, ) from lfx.cli.validation import validate_global_variables_for_env +from lfx.execution import get_default_coordinator from lfx.log.logger import logger from lfx.run._defaults import apply_run_defaults, resolve_fallback_to_env_vars, validate_provided_id from lfx.schema.schema import InputValueRequest @@ -502,8 +503,11 @@ async def run_flow( # fall through to os.environ on miss instead of erroring the build). fallback_to_env_vars = resolve_fallback_to_env_vars() - async for result in graph.async_start( - inputs, event_manager=event_manager, fallback_to_env_vars=fallback_to_env_vars + async for result in get_default_coordinator().stream( + graph, + initial_inputs=inputs, + event_manager=event_manager, + fallback_to_env_vars=fallback_to_env_vars, ): result_count += 1 if verbosity > 0: diff --git a/src/lfx/src/lfx/services/deps.py b/src/lfx/src/lfx/services/deps.py index aa0cab4478..25a3601904 100644 --- a/src/lfx/src/lfx/services/deps.py +++ b/src/lfx/src/lfx/services/deps.py @@ -61,6 +61,10 @@ def get_service(service_type: ServiceType, default=None): try: return service_manager.get(service_type, default) except Exception: # noqa: BLE001 + # Preserve the traceback in logs so callers seeing a None return have something to grep + # for. Returning None remains the contract because several callers (e.g. get_db_service) + # treat absence as "not configured" and substitute a noop implementation. + logger.exception("Failed to resolve service %s", service_type) return None diff --git a/src/lfx/src/lfx/services/executor/__init__.py b/src/lfx/src/lfx/services/executor/__init__.py new file mode 100644 index 0000000000..57a6cf6e87 --- /dev/null +++ b/src/lfx/src/lfx/services/executor/__init__.py @@ -0,0 +1 @@ +"""Executor service module.""" diff --git a/src/lfx/src/lfx/services/executor/factory.py b/src/lfx/src/lfx/services/executor/factory.py new file mode 100644 index 0000000000..b3d02fd7ab --- /dev/null +++ b/src/lfx/src/lfx/services/executor/factory.py @@ -0,0 +1,24 @@ +"""Factory for the executor service.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from lfx.services.executor.service import ExecutorService +from lfx.services.factory import ServiceFactory +from lfx.services.schema import ServiceType + +if TYPE_CHECKING: + from lfx.services.settings.service import SettingsService + + +class ExecutorServiceFactory(ServiceFactory): + """Factory for creating ExecutorService instances.""" + + def __init__(self) -> None: + super().__init__() + self.service_class = ExecutorService + self.dependencies = [ServiceType.SETTINGS_SERVICE] + + def create(self, settings_service: SettingsService) -> ExecutorService: + return ExecutorService(settings_service=settings_service) diff --git a/src/lfx/src/lfx/services/executor/service.py b/src/lfx/src/lfx/services/executor/service.py new file mode 100644 index 0000000000..2f8e683156 --- /dev/null +++ b/src/lfx/src/lfx/services/executor/service.py @@ -0,0 +1,124 @@ +"""Executor service: owns the executor registry and the default coordinator.""" + +from __future__ import annotations + +from importlib.metadata import entry_points +from typing import TYPE_CHECKING + +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 ExecutorKindCollisionError, ExecutorRegistry +from lfx.log.logger import logger +from lfx.services.base import Service + +if TYPE_CHECKING: + from lfx.services.settings.service import SettingsService + + +ENTRY_POINT_GROUP = "lfx.executors" + + +class ExecutorService(Service): + """Service that holds the executor registry and the default coordinator.""" + + name = "executor_service" + + def __init__(self, settings_service: SettingsService) -> None: + super().__init__() + self._settings_service = settings_service + self._registry = ExecutorRegistry() + self._coordinator: Coordinator | None = None + self._populate_registry() + + def register(self, executor: Executor) -> None: + """Register an executor; replaces any prior registration for the same kind. + + Invalidates the cached coordinator so the next access rebuilds against the new + registry contents. Direct mutation of the registry through other paths bypasses + this invalidation -- always go through this method to register or replace. + """ + self._registry.register(executor) + self._coordinator = None + + def has(self, kind: str) -> bool: + """Return True if an executor with the given kind is registered.""" + return self._registry.has(kind) + + def get(self, kind: str) -> Executor: + """Return the executor registered under the given kind.""" + return self._registry.get(kind) + + @property + def registry(self) -> ExecutorRegistry: + """Return the underlying registry. + + Read access only. Mutating the registry directly (e.g. ``service.registry.register(...)``) + bypasses the cached-coordinator invalidation done by :meth:`register`; always go through + :meth:`register` for mutations. + """ + return self._registry + + @property + def coordinator(self) -> Coordinator: + if self._coordinator is None: + kind = self._settings_service.settings.executor_kind + self._coordinator = Coordinator(registry=self._registry, executor_kind=kind) + return self._coordinator + + def set_coordinator(self, coordinator: Coordinator) -> None: + """Override the default coordinator (test/extension hook).""" + self._coordinator = coordinator + + def _populate_registry(self) -> None: + """Register the built-in executor and run plugin discovery.""" + self._registry.register(InProcessExecutor()) + self._discover_entry_points() + + def _discover_entry_points(self) -> None: + try: + eps = entry_points(group=ENTRY_POINT_GROUP) + except Exception: # noqa: BLE001 + logger.exception( + "Failed to enumerate entry points for group=%r. Executor plugin discovery skipped.", + ENTRY_POINT_GROUP, + ) + return + for ep in eps: + kind: str = "" + try: + obj = ep.load() + executor = obj() if isinstance(obj, type) else obj + if not isinstance(executor, Executor): + logger.warning( + "Skipping entry point %r: loaded object of type %r is not an Executor.", + ep.name, + type(executor).__name__, + ) + continue + kind = getattr(executor, "kind", "") + self._registry.register(executor, replace=False) + logger.debug("Loaded executor from entry point: %s -> %s", ep.name, kind) + except ExecutorKindCollisionError: + logger.warning( + "Skipping entry point %r: an executor with kind=%r is already registered. " + "Entry-point discovery cannot replace existing executors; register the " + "executor explicitly via ExecutorService.register() if replacement is intended.", + ep.name, + kind, + ) + except Exception: # noqa: BLE001 + # Broad catch is intentional: a single broken plugin must not break discovery + # of others. The full traceback is preserved via logger.exception. + logger.exception("Failed to load executor entry point %s", ep.name) + + async def teardown(self) -> None: + """Reset the service to its post-init state. + + Rebuilds the registry (built-in InProcessExecutor + plugin discovery) and clears + the cached coordinator, so the service stays usable if it is accessed again before + the service manager evicts it. + """ + self._coordinator = None + self._registry = ExecutorRegistry() + self._populate_registry() diff --git a/src/lfx/src/lfx/services/manager.py b/src/lfx/src/lfx/services/manager.py index 4c206a6722..86b16f303f 100644 --- a/src/lfx/src/lfx/services/manager.py +++ b/src/lfx/src/lfx/services/manager.py @@ -314,6 +314,12 @@ class ServiceManager: self.services = {} self.factories = {} + # ``teardown`` empties the factory registry, so the "registered" flag has + # to drop too: get_service() re-registers factories only when + # are_factories_registered() is False. Leaving it set after a teardown + # makes the next lookup skip re-registration and raise + # NoFactoryRegisteredError. + self.factory_registered = False if raise_on_error and errors: names = ", ".join(name for name, _ in errors) diff --git a/src/lfx/src/lfx/services/schema.py b/src/lfx/src/lfx/services/schema.py index 47a53ba257..1bba5de093 100644 --- a/src/lfx/src/lfx/services/schema.py +++ b/src/lfx/src/lfx/services/schema.py @@ -27,3 +27,4 @@ class ServiceType(str, Enum): FLOW_EVENTS_SERVICE = "flow_events_service" TELEMETRY_WRITER_SERVICE = "telemetry_writer_service" EXTENSION_EVENTS_SERVICE = "extension_events_service" + EXECUTOR_SERVICE = "executor_service" diff --git a/src/lfx/src/lfx/services/settings/groups/runtime.py b/src/lfx/src/lfx/services/settings/groups/runtime.py index d83b1d5b25..43d25baa00 100644 --- a/src/lfx/src/lfx/services/settings/groups/runtime.py +++ b/src/lfx/src/lfx/services/settings/groups/runtime.py @@ -90,6 +90,14 @@ class RuntimeSettings(BaseModel): celery_enabled: bool = False + executor_kind: str = "in-process" + """The default executor kind used by the execution coordinator. + + Must match the `kind` of an Executor registered with the executor service. The built-in + `in-process` executor runs graphs in the current process; third-party executors registered + via the `lfx.executors` entry-point group can be selected by setting this to their kind. + """ + @field_validator("event_delivery", mode="before") @classmethod def set_event_delivery(cls, value, info): diff --git a/src/lfx/tests/unit/components/flow_controls/test_loop_events.py b/src/lfx/tests/unit/components/flow_controls/test_loop_events.py index 5431770a95..5e8f324df9 100644 --- a/src/lfx/tests/unit/components/flow_controls/test_loop_events.py +++ b/src/lfx/tests/unit/components/flow_controls/test_loop_events.py @@ -77,79 +77,107 @@ class TestEventManagerPropagation: @pytest.mark.asyncio async def test_event_manager_passed_to_subgraph_async_start(self): - """Test that event_manager is passed to subgraph's async_start method.""" - mock_event_manager = MagicMock() - received_event_manager = None - - # Create a mock subgraph that captures the event_manager - async def mock_async_start(event_manager=None): - nonlocal received_event_manager - received_event_manager = event_manager - yield MagicMock(valid=True, result_dict=MagicMock(outputs={})) - - def create_mock_subgraph(_vertex_ids): - mock_subgraph = MagicMock() - mock_subgraph._vertices = [] - mock_subgraph.prepare = MagicMock() - mock_subgraph.async_start = mock_async_start - mock_subgraph.get_vertex = MagicMock(return_value=MagicMock(custom_component=MagicMock())) - return mock_subgraph - - mock_graph = MagicMock() - mock_graph.create_subgraph = create_subgraph_context_manager_mock(create_mock_subgraph) - - data_list = [Data(text="item1")] - - await execute_loop_body( - graph=mock_graph, - data_list=data_list, - loop_body_vertex_ids={"vertex1"}, - start_vertex_id="vertex1", - start_edge=MagicMock(target_handle=MagicMock(field_name="data")), - end_vertex_id="vertex1", - event_manager=mock_event_manager, + """The loop body's runner must forward event_manager into the executor's runtime_options.""" + from lfx.execution import ( + Coordinator, + ExecutorRegistry, + reset_default_coordinator, + set_default_coordinator, ) + from lfx.execution.executor import Executor + from lfx.execution.types import RunComplete - # Verify event_manager was passed to async_start - assert received_event_manager is mock_event_manager + try: + mock_event_manager = MagicMock() + seen: list = [] + + class _Recorder(Executor): + kind = "in-process" + + async def execute(self, unit): + seen.append(unit.runtime_options.get("event_manager")) + yield RunComplete(outputs=[]) + + registry = ExecutorRegistry() + registry.register(_Recorder()) + set_default_coordinator(Coordinator(registry=registry)) + + def create_mock_subgraph(_vertex_ids): + mock_subgraph = MagicMock() + mock_subgraph._vertices = [] + mock_subgraph.prepare = MagicMock() + mock_subgraph.get_vertex = MagicMock(return_value=MagicMock(custom_component=MagicMock())) + return mock_subgraph + + mock_graph = MagicMock() + mock_graph.create_subgraph = create_subgraph_context_manager_mock(create_mock_subgraph) + + await execute_loop_body( + graph=mock_graph, + data_list=[Data(text="item1")], + loop_body_vertex_ids={"vertex1"}, + start_vertex_id="vertex1", + start_edge=MagicMock(target_handle=MagicMock(field_name="data")), + end_vertex_id="vertex1", + event_manager=mock_event_manager, + ) + + assert len(seen) == 1 + assert seen[0] is mock_event_manager + finally: + reset_default_coordinator() @pytest.mark.asyncio async def test_event_manager_passed_for_each_iteration(self): - """Test that event_manager is passed to async_start for each loop iteration.""" - mock_event_manager = MagicMock() - event_manager_calls = [] - - async def mock_async_start(event_manager=None): - event_manager_calls.append(event_manager) - yield MagicMock(valid=True, result_dict=MagicMock(outputs={})) - - def create_mock_subgraph(_vertex_ids): - mock_subgraph = MagicMock() - mock_subgraph._vertices = [] - mock_subgraph.prepare = MagicMock() - mock_subgraph.async_start = mock_async_start - mock_subgraph.get_vertex = MagicMock(return_value=MagicMock(custom_component=MagicMock())) - return mock_subgraph - - mock_graph = MagicMock() - mock_graph.create_subgraph = create_subgraph_context_manager_mock(create_mock_subgraph) - - # 3 items = 3 iterations - data_list = [Data(text="item1"), Data(text="item2"), Data(text="item3")] - - await execute_loop_body( - graph=mock_graph, - data_list=data_list, - loop_body_vertex_ids={"vertex1"}, - start_vertex_id="vertex1", - start_edge=MagicMock(target_handle=MagicMock(field_name="data")), - end_vertex_id="vertex1", - event_manager=mock_event_manager, + """Every loop iteration runs through the executor with the same event_manager.""" + from lfx.execution import ( + Coordinator, + ExecutorRegistry, + reset_default_coordinator, + set_default_coordinator, ) + from lfx.execution.executor import Executor + from lfx.execution.types import RunComplete - # Verify event_manager was passed for each iteration - assert len(event_manager_calls) == 3 - assert all(em is mock_event_manager for em in event_manager_calls) + try: + mock_event_manager = MagicMock() + seen: list = [] + + class _Recorder(Executor): + kind = "in-process" + + async def execute(self, unit): + seen.append(unit.runtime_options.get("event_manager")) + yield RunComplete(outputs=[]) + + registry = ExecutorRegistry() + registry.register(_Recorder()) + set_default_coordinator(Coordinator(registry=registry)) + + def create_mock_subgraph(_vertex_ids): + mock_subgraph = MagicMock() + mock_subgraph._vertices = [] + mock_subgraph.prepare = MagicMock() + mock_subgraph.get_vertex = MagicMock(return_value=MagicMock(custom_component=MagicMock())) + return mock_subgraph + + mock_graph = MagicMock() + mock_graph.create_subgraph = create_subgraph_context_manager_mock(create_mock_subgraph) + + await execute_loop_body( + graph=mock_graph, + data_list=[Data(text="item1"), Data(text="item2"), Data(text="item3")], + loop_body_vertex_ids={"vertex1"}, + start_vertex_id="vertex1", + start_edge=MagicMock(target_handle=MagicMock(field_name="data")), + end_vertex_id="vertex1", + event_manager=mock_event_manager, + ) + + assert len(seen) == 3 + assert all(em is mock_event_manager for em in seen) + finally: + reset_default_coordinator() def test_subgraph_preserves_vertex_ids(self): """Test that subgraph vertices maintain original IDs. diff --git a/src/lfx/tests/unit/execution/__init__.py b/src/lfx/tests/unit/execution/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/src/lfx/tests/unit/execution/conftest.py b/src/lfx/tests/unit/execution/conftest.py new file mode 100644 index 0000000000..b4cb974351 --- /dev/null +++ b/src/lfx/tests/unit/execution/conftest.py @@ -0,0 +1,24 @@ +"""Shared fixtures for execution tests.""" + +from __future__ import annotations + +import pytest +from lfx.components.input_output import ChatInput, ChatOutput +from lfx.execution import reset_default_coordinator +from lfx.graph import Graph + + +@pytest.fixture(autouse=True) +def _reset_execution_singletons(): + reset_default_coordinator() + yield + reset_default_coordinator() + + +@pytest.fixture +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) 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.py b/src/lfx/tests/unit/execution/test_coordinator.py new file mode 100644 index 0000000000..9ecc100958 --- /dev/null +++ b/src/lfx/tests/unit/execution/test_coordinator.py @@ -0,0 +1,59 @@ +import pytest +from lfx.execution.coordinator import Coordinator +from lfx.execution.registry import ExecutorRegistry +from lfx.execution.types import RunComplete, StepResult + + +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(simple_graph): + 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(simple_graph): + 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(simple_graph): + coordinator = _make_coordinator() + + async def boom(*args, **kwargs): # noqa: ARG001 + msg = "boom" + raise RuntimeError(msg) + + simple_graph._arun_legacy = boom + + with pytest.raises(RuntimeError, match="boom"): + await coordinator.run_to_completion(simple_graph, inputs=[{"input_value": "hi"}], _use_arun_legacy=True) + + +@pytest.mark.asyncio +async def test_run_with_no_executor_registered_raises(simple_graph): + from lfx.execution.registry import ExecutorNotFoundError + + coordinator = Coordinator(registry=ExecutorRegistry()) + with pytest.raises(ExecutorNotFoundError): + async for _ in coordinator.run(simple_graph, inputs=[]): + pass + + +@pytest.mark.asyncio +async def test_stream_helper_yields_payloads_directly(simple_graph): + coordinator = _make_coordinator() + payloads = [p async for p in coordinator.stream(simple_graph, initial_inputs=None)] + assert payloads + assert any(hasattr(p, "vertex") for p in payloads) 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..9ac4aec15b --- /dev/null +++ b/src/lfx/tests/unit/execution/test_coordinator_e2e.py @@ -0,0 +1,174 @@ +"""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) + + +def _build_simple_graph(): + """Build a fresh single-use graph (mirror of the ``simple_graph`` fixture). + + Tests that drive more than one run need a new graph each time, since a Graph + can't be re-driven after it finishes. + """ + from lfx.components.input_output import ChatInput, ChatOutput + from lfx.graph import 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) + + +@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(): + """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") + + # Build a fresh graph per run; a Graph is single-use once driven to completion, + # so reusing one across both runs would couple the second run to the first's state. + [_ async for _ in coordinator.run(_build_simple_graph(), inputs=[{"input_value": "a"}])] + [_ async for _ in coordinator.run(_build_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_default_registry.py b/src/lfx/tests/unit/execution/test_default_registry.py new file mode 100644 index 0000000000..35892b638d --- /dev/null +++ b/src/lfx/tests/unit/execution/test_default_registry.py @@ -0,0 +1,32 @@ +from lfx.execution import ( + Coordinator, + get_default_coordinator, + get_default_registry, + set_default_coordinator, +) + + +def test_default_registry_has_in_process(): + assert get_default_registry().get("in-process").kind == "in-process" + + +def test_default_coordinator_uses_default_registry(): + c = get_default_coordinator() + assert isinstance(c, Coordinator) + + +def test_set_default_coordinator_overrides_singleton(): + original = get_default_coordinator() + custom = Coordinator(registry=get_default_registry()) + set_default_coordinator(custom) + try: + assert get_default_coordinator() is custom + finally: + # Restore the singleton so this override doesn't bleed into later tests. + set_default_coordinator(original) + + +def test_default_coordinator_is_idempotent_within_a_test(): + a = get_default_coordinator() + b = get_default_coordinator() + assert a is b diff --git a/src/lfx/tests/unit/execution/test_executor.py b/src/lfx/tests/unit/execution/test_executor.py new file mode 100644 index 0000000000..3d226e8f01 --- /dev/null +++ b/src/lfx/tests/unit/execution/test_executor.py @@ -0,0 +1,37 @@ +import pytest +from lfx.execution.executor import Executor +from lfx.execution.types import RunComplete, StepResult, Unit + + +def test_executor_is_abstract(): + with pytest.raises(TypeError, match="abstract"): + Executor() + + +@pytest.mark.asyncio +async def test_concrete_executor_yields_steps_then_run_complete(): + class Toy(Executor): + kind = "toy" + + async def execute(self, unit): # noqa: ARG002 + yield StepResult(payload={"vertex_id": "a"}) + yield StepResult(payload={"vertex_id": "b"}) + yield RunComplete(outputs=["done"]) + + items = [item async for item in Toy().execute(Unit(graph=object(), inputs=[]))] + assert len(items) == 3 + assert isinstance(items[-1], RunComplete) + assert items[-1].outputs == ["done"] + + +@pytest.mark.asyncio +async def test_executor_must_yield_run_complete_last(): + class Toy(Executor): + kind = "toy" + + async def execute(self, unit): # noqa: ARG002 + yield StepResult(payload={}) + yield RunComplete(outputs=[]) + + items = [item async for item in Toy().execute(Unit(graph=object(), inputs=[]))] + assert isinstance(items[-1], RunComplete) 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 diff --git a/src/lfx/tests/unit/execution/test_executor_service.py b/src/lfx/tests/unit/execution/test_executor_service.py new file mode 100644 index 0000000000..14cc47e8f9 --- /dev/null +++ b/src/lfx/tests/unit/execution/test_executor_service.py @@ -0,0 +1,228 @@ +"""Service-level contract tests for the ExecutorService.""" + +from __future__ import annotations + +import pytest +from lfx.execution.executor import Executor +from lfx.execution.types import RunComplete +from lfx.services.deps import get_service +from lfx.services.executor.service import ExecutorService +from lfx.services.manager import get_service_manager +from lfx.services.schema import ServiceType + + +def _get_service() -> ExecutorService: + service = get_service(ServiceType.EXECUTOR_SERVICE) + assert service is not None + return service + + +def test_executor_service_is_resolvable_via_service_manager(): + service = _get_service() + assert isinstance(service, ExecutorService) + + +def test_executor_service_preregisters_in_process(): + service = _get_service() + assert service.has("in-process") + assert service.get("in-process").kind == "in-process" + + +@pytest.mark.asyncio +async def test_coordinator_uses_settings_executor_kind(monkeypatch, request): + """Behavioral assertion: coordinator dispatches to whichever kind the settings select.""" + settings_service = get_service(ServiceType.SETTINGS_SERVICE) + # monkeypatch restores the original kind even if an assertion below fails, so a + # mid-test failure can't poison the session-wide settings for later tests. The + # finalizer rebuilds the executor service after the kind reverts, so the service + # built from the patched kind doesn't linger either. + monkeypatch.setattr(settings_service.settings, "executor_kind", "kind-from-settings") + request.addfinalizer(lambda: get_service_manager().update(ServiceType.EXECUTOR_SERVICE)) + + seen: list[str] = [] + + class Sentinel(Executor): + kind = "kind-from-settings" + + async def execute(self, unit): # noqa: ARG002 + seen.append(self.kind) + yield RunComplete(outputs=["ok"]) + + get_service_manager().update(ServiceType.EXECUTOR_SERVICE) + service = _get_service() + service.register(Sentinel()) + + outputs = await service.coordinator.run_to_completion(graph=object(), inputs=[]) + + assert seen == ["kind-from-settings"] + assert outputs == ["ok"] + + +def test_register_replaces_and_invalidates_cached_coordinator(): + service = _get_service() + first = service.coordinator + + class Other(Executor): + kind = "in-process" + + async def execute(self, unit): # noqa: ARG002 + yield RunComplete(outputs=[]) + + service.register(Other()) + assert service.coordinator is not first + assert isinstance(service.get("in-process"), Other) + + +@pytest.mark.asyncio +async def test_teardown_preserves_builtin_in_process(): + """teardown() must leave the service usable: the built-in must still be registered. + + ServiceManager.teardown() does not evict the cached service instance, so any code + that resolves the service after teardown (e.g. background tasks racing app shutdown) + must still get a working coordinator. + """ + service = _get_service() + + custom_kinds_before = service.has("in-process") + assert custom_kinds_before + + await service.teardown() + + assert service.has("in-process") + assert service.coordinator is not None # rebuilds lazily over the fresh registry + + +def test_entry_point_cannot_replace_builtin_in_process(monkeypatch): + """A package exposing an executor with kind="in-process" must not silently replace the built-in.""" + from lfx.execution.backends.in_process import InProcessExecutor + from lfx.services.executor import service as service_module + + class Hijacker(Executor): + kind = "in-process" + + async def execute(self, unit): # noqa: ARG002 + yield RunComplete(outputs=[]) + + class _FakeEntryPoint: + name = "hijacker" + + def load(self): + return Hijacker + + monkeypatch.setattr( + service_module, + "entry_points", + lambda group: [_FakeEntryPoint()] if group == service_module.ENTRY_POINT_GROUP else [], + ) + + settings_service = get_service(ServiceType.SETTINGS_SERVICE) + fresh = ExecutorService(settings_service=settings_service) + + assert isinstance(fresh.get("in-process"), InProcessExecutor) + + +def test_entry_point_with_unique_kind_is_registered(monkeypatch): + """Entry points with non-colliding kinds are still loaded.""" + from lfx.services.executor import service as service_module + + class RemoteStub(Executor): + kind = "remote-stub" + + async def execute(self, unit): # noqa: ARG002 + yield RunComplete(outputs=[]) + + class _FakeEntryPoint: + name = "remote_stub" + + def load(self): + return RemoteStub + + monkeypatch.setattr( + service_module, + "entry_points", + lambda group: [_FakeEntryPoint()] if group == service_module.ENTRY_POINT_GROUP else [], + ) + + settings_service = get_service(ServiceType.SETTINGS_SERVICE) + fresh = ExecutorService(settings_service=settings_service) + + assert fresh.has("remote-stub") + assert isinstance(fresh.get("remote-stub"), RemoteStub) + + +def test_entry_point_load_failure_does_not_break_discovery(monkeypatch): + """A single broken entry point must not prevent the built-in or other plugins from loading.""" + from lfx.execution.backends.in_process import InProcessExecutor + from lfx.services.executor import service as service_module + + class GoodStub(Executor): + kind = "good-stub" + + async def execute(self, unit): # noqa: ARG002 + yield RunComplete(outputs=[]) + + class _BrokenEntryPoint: + name = "broken" + + def load(self): + msg = "simulated import failure" + raise ImportError(msg) + + class _GoodEntryPoint: + name = "good_stub" + + def load(self): + return GoodStub + + monkeypatch.setattr( + service_module, + "entry_points", + lambda group: [_BrokenEntryPoint(), _GoodEntryPoint()] if group == service_module.ENTRY_POINT_GROUP else [], + ) + + settings_service = get_service(ServiceType.SETTINGS_SERVICE) + fresh = ExecutorService(settings_service=settings_service) + + assert isinstance(fresh.get("in-process"), InProcessExecutor) + assert isinstance(fresh.get("good-stub"), GoodStub) + + +def test_entry_point_returning_non_executor_is_skipped(monkeypatch): + """An entry point whose loaded object is not an Executor must be skipped, not registered.""" + from lfx.execution.backends.in_process import InProcessExecutor + from lfx.services.executor import service as service_module + + class NotAnExecutor: + # Carries a `kind` (so it would survive registry.register's attribute read) but is + # not an Executor. Without the isinstance guard it would register and only blow up + # later at execute() time; discovery must reject it up front. + kind = "not-really" + + class _BadEntryPoint: + name = "bad" + + def load(self): + return NotAnExecutor + + monkeypatch.setattr( + service_module, + "entry_points", + lambda group: [_BadEntryPoint()] if group == service_module.ENTRY_POINT_GROUP else [], + ) + + settings_service = get_service(ServiceType.SETTINGS_SERVICE) + fresh = ExecutorService(settings_service=settings_service) + + # The bad plugin is rejected outright... + assert not fresh.has("not-really") + # ...and the built-in executor still works. + assert isinstance(fresh.get("in-process"), InProcessExecutor) + + +def test_factory_declares_settings_dependency(): + """Pin the factory contract: ExecutorService depends on settings_service only.""" + from lfx.services.executor.factory import ExecutorServiceFactory + + factory = ExecutorServiceFactory() + assert factory.service_class is ExecutorService + assert factory.dependencies == [ServiceType.SETTINGS_SERVICE] diff --git a/src/lfx/tests/unit/execution/test_graph_arun_routes_through_coordinator.py b/src/lfx/tests/unit/execution/test_graph_arun_routes_through_coordinator.py new file mode 100644 index 0000000000..648648186c --- /dev/null +++ b/src/lfx/tests/unit/execution/test_graph_arun_routes_through_coordinator.py @@ -0,0 +1,70 @@ +import asyncio + +import pytest +from lfx.components.input_output import ChatInput, ChatOutput +from lfx.execution import ( + Coordinator, + ExecutorRegistry, + set_default_coordinator, +) +from lfx.execution.executor import Executor +from lfx.execution.types import RunComplete +from lfx.graph import Graph +from lfx.graph.schema import RunOutputs + + +@pytest.mark.asyncio +async def test_arun_dispatches_through_coordinator(simple_graph): + units_seen: list[object] = [] + + class Recording(Executor): + kind = "in-process" + + async def execute(self, unit): + units_seen.append(unit) + yield RunComplete(outputs=[]) + + registry = ExecutorRegistry() + registry.register(Recording()) + set_default_coordinator(Coordinator(registry=registry)) + + await simple_graph.arun(inputs=[{"input_value": "hi"}]) + assert len(units_seen) == 1 + seen = units_seen[0] + assert seen.inputs == [{"input_value": "hi"}] + assert seen.runtime_options.get("_use_arun_legacy") is True + + +@pytest.mark.asyncio +async def test_arun_returns_run_outputs_shape(simple_graph): + outputs = await simple_graph.arun(inputs=[{"input_value": "hi"}]) + assert isinstance(outputs, list) + assert len(outputs) == 1 + assert isinstance(outputs[0], RunOutputs) + assert outputs[0].inputs == {"input_value": "hi"} + + +@pytest.mark.asyncio +async def test_arun_concurrent_on_separate_graphs_keep_options_isolated(): + """Concurrent arun on separate Graph instances keeps kwargs isolated.""" + seen: list[str] = [] + + async def patched(*, inputs, **kwargs): # noqa: ARG001 + seen.append(kwargs.get("session_id") or "") + await asyncio.sleep(0.01) + return [] + + def _build(): + 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 Graph(ci, co) + + async def run_with(session_id): + graph = _build() + graph._arun_legacy = patched + await graph.arun(inputs=[{"input_value": session_id}], session_id=session_id) + + await asyncio.gather(run_with("A"), run_with("B")) + assert sorted(seen) == ["A", "B"] diff --git a/src/lfx/tests/unit/execution/test_in_process_executor.py b/src/lfx/tests/unit/execution/test_in_process_executor.py new file mode 100644 index 0000000000..5f2b6674f5 --- /dev/null +++ b/src/lfx/tests/unit/execution/test_in_process_executor.py @@ -0,0 +1,146 @@ +import asyncio + +import pytest +from lfx.events.event_manager import EventManager +from lfx.execution.backends.in_process import InProcessExecutor +from lfx.execution.types import RunComplete, StepResult, Unit + + +def test_in_process_executor_kind(): + assert InProcessExecutor().kind == "in-process" + + +@pytest.mark.asyncio +async def test_streams_step_results_then_run_complete_for_real_graph(simple_graph): + """Streaming path emits StepResult per vertex build plus a Finish step, then RunComplete.""" + from lfx.graph.graph.constants import Finish + from lfx.schema.schema import InputValueRequest + + unit = Unit( + graph=simple_graph, + inputs=[], + runtime_options={"initial_inputs": InputValueRequest(input_value="hello world")}, + ) + + items = [item async for item in InProcessExecutor().execute(unit)] + step_payloads = [i.payload for i in items if isinstance(i, StepResult)] + + assert isinstance(items[-1], RunComplete) + assert all(isinstance(i, StepResult) for i in items[:-1]) + vertex_ids = [getattr(p, "vertex", None) and p.vertex.id for p in step_payloads] + assert "chat_input" in vertex_ids + assert "chat_output" in vertex_ids + assert any(isinstance(p, Finish) for p in step_payloads) + + +@pytest.mark.asyncio +async def test_run_complete_outputs_match_arun_shape(simple_graph): + """Legacy path: dict inputs + flag yields list[RunOutputs] from _arun_legacy.""" + from lfx.graph.schema import RunOutputs + + unit = Unit( + graph=simple_graph, + inputs=[{"input_value": "hello world"}], + runtime_options={"_use_arun_legacy": True}, + ) + + items = [item async for item in InProcessExecutor().execute(unit)] + rc = items[-1] + assert isinstance(rc, RunComplete) + assert len(rc.outputs) == 1 + assert isinstance(rc.outputs[0], RunOutputs) + assert rc.outputs[0].inputs == {"input_value": "hello world"} + + +@pytest.mark.asyncio +async def test_event_manager_in_runtime_options_is_used(simple_graph): + """Streaming path forwards event_manager into async_start so vertex events fire.""" + from lfx.schema.schema import InputValueRequest + + queue: asyncio.Queue = asyncio.Queue() + em = EventManager(queue=queue) + + seen: list[str] = [] + + def record(*, manager, event_type, data): # noqa: ARG001 + seen.append(f"{event_type}:{data.get('id') if isinstance(data, dict) else ''}") + + em.register_event("on_build_start", "build_start", record) + + unit = Unit( + graph=simple_graph, + inputs=[], + runtime_options={ + "event_manager": em, + "initial_inputs": InputValueRequest(input_value="hi"), + }, + ) + [_ async for _ in InProcessExecutor().execute(unit)] + + assert any(s.startswith("build_start:chat_input") for s in seen) + assert any(s.startswith("build_start:chat_output") for s in seen) + + +@pytest.mark.asyncio +async def test_propagates_graph_errors(simple_graph): + async def boom(*args, **kwargs): # noqa: ARG001 + msg = "boom" + raise RuntimeError(msg) + + simple_graph._arun_legacy = boom + + unit = Unit(graph=simple_graph, inputs=[{}], runtime_options={"_use_arun_legacy": True}) + with pytest.raises(RuntimeError, match="boom"): + async for _ in InProcessExecutor().execute(unit): + pass + + +@pytest.mark.asyncio +async def test_fallback_to_env_vars_forwarded_in_streaming_path(): + """Regression: the streaming path must forward fallback_to_env_vars to async_start.""" + + class _Stub: + seen_fallback: bool | None = None + + async def async_start(self, **kwargs): + _Stub.seen_fallback = kwargs.get("fallback_to_env_vars") + return + yield + + unit = Unit(graph=_Stub(), inputs=[], runtime_options={"fallback_to_env_vars": True}) + [_ async for _ in InProcessExecutor().execute(unit)] + assert _Stub.seen_fallback is True + + +@pytest.mark.asyncio +async def test_concurrent_runs_on_separate_graphs_keep_runtime_options_isolated(): + """Separate Graph instances stay isolated through the executor layer.""" + from lfx.components.input_output import ChatInput, ChatOutput + from lfx.graph import Graph + + seen: list[str | None] = [] + + async def patched_arun_legacy(*, inputs, **kwargs): # noqa: ARG001 + seen.append(kwargs.get("session_id")) + await asyncio.sleep(0.01) + return [] + + def _build(): + 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) + g = Graph(ci, co) + g._arun_legacy = patched_arun_legacy + return g + + async def run_with(session_id): + unit = Unit( + graph=_build(), + inputs=[{}], + runtime_options={"session_id": session_id, "_use_arun_legacy": True}, + ) + [_ async for _ in InProcessExecutor().execute(unit)] + + await asyncio.gather(run_with("A"), run_with("B")) + assert sorted(seen) == ["A", "B"] diff --git a/src/lfx/tests/unit/execution/test_partitioner.py b/src/lfx/tests/unit/execution/test_partitioner.py new file mode 100644 index 0000000000..c29cc9ef9e --- /dev/null +++ b/src/lfx/tests/unit/execution/test_partitioner.py @@ -0,0 +1,17 @@ +from lfx.execution.partitioner import identity_partition +from lfx.execution.types import Unit + + +def test_identity_partition_returns_one_unit(simple_graph): + units = identity_partition(simple_graph, inputs=[{"input_value": "hi"}], runtime_options={"session_id": "s1"}) + assert len(units) == 1 + [unit] = units + assert isinstance(unit, Unit) + assert unit.graph is simple_graph + assert unit.inputs == [{"input_value": "hi"}] + assert unit.runtime_options == {"session_id": "s1"} + + +def test_identity_partition_default_runtime_options(simple_graph): + units = identity_partition(simple_graph, inputs=[]) + assert units[0].runtime_options == {} diff --git a/src/lfx/tests/unit/execution/test_registry.py b/src/lfx/tests/unit/execution/test_registry.py new file mode 100644 index 0000000000..421776ada5 --- /dev/null +++ b/src/lfx/tests/unit/execution/test_registry.py @@ -0,0 +1,47 @@ +import pytest +from lfx.execution.executor import Executor +from lfx.execution.registry import ExecutorKindCollisionError, ExecutorNotFoundError, ExecutorRegistry + + +class _Stub(Executor): + kind = "stub" + + async def execute(self, unit): # noqa: ARG002 + return + yield + + +def test_register_and_get_by_kind(): + registry = ExecutorRegistry() + registry.register(_Stub()) + assert registry.get("stub").kind == "stub" + + +def test_get_unknown_raises(): + with pytest.raises(ExecutorNotFoundError, match="nope"): + ExecutorRegistry().get("nope") + + +def test_register_replaces_same_kind(): + registry = ExecutorRegistry() + a = _Stub() + b = _Stub() + registry.register(a) + registry.register(b) + assert registry.get("stub") is b + + +def test_register_replace_false_raises_on_collision(): + registry = ExecutorRegistry() + first = _Stub() + registry.register(first) + with pytest.raises(ExecutorKindCollisionError, match="stub"): + registry.register(_Stub(), replace=False) + assert registry.get("stub") is first + + +def test_has_reflects_registration(): + registry = ExecutorRegistry() + assert not registry.has("stub") + registry.register(_Stub()) + assert registry.has("stub") diff --git a/src/lfx/tests/unit/execution/test_types.py b/src/lfx/tests/unit/execution/test_types.py new file mode 100644 index 0000000000..522b8eb1bc --- /dev/null +++ b/src/lfx/tests/unit/execution/test_types.py @@ -0,0 +1,28 @@ +from lfx.execution.types import RunComplete, StepResult, Unit + + +def test_unit_carries_graph_inputs_and_runtime_options(): + sentinel = object() + unit = Unit( + graph=sentinel, + inputs=[{"input_value": "hi"}], + runtime_options={"event_manager": None, "session_id": "s1"}, + ) + assert unit.graph is sentinel + assert unit.inputs == [{"input_value": "hi"}] + assert unit.runtime_options == {"event_manager": None, "session_id": "s1"} + + +def test_unit_runtime_options_defaults_to_empty_dict(): + unit = Unit(graph=object(), inputs=[]) + assert unit.runtime_options == {} + + +def test_run_complete_carries_outputs(): + rc = RunComplete(outputs=["a", "b"]) + assert rc.outputs == ["a", "b"] + + +def test_step_result_carries_payload(): + sr = StepResult(payload={"vertex_id": "x"}) + assert sr.payload == {"vertex_id": "x"} diff --git a/src/lfx/tests/unit/run/test_base.py b/src/lfx/tests/unit/run/test_base.py index fda94111be..e9b9e0e9c3 100644 --- a/src/lfx/tests/unit/run/test_base.py +++ b/src/lfx/tests/unit/run/test_base.py @@ -180,7 +180,7 @@ class TestRunFlowJsonInput: mock_graph.edges = [] mock_graph.prepare = MagicMock() - async def mock_async_start(_inputs, **_kwargs): + async def mock_async_start(_inputs=None, **_kwargs): yield mock_graph.async_start = mock_async_start @@ -308,7 +308,7 @@ class TestRunFlowSessionId: graph.edges = [] graph.prepare = MagicMock() - async def _async_start(_inputs, **_kwargs): + async def _async_start(_inputs=None, **_kwargs): yield graph.async_start = _async_start @@ -469,7 +469,7 @@ class TestRunFlowSessionIdPropagation: graph.edges = [] graph.prepare = MagicMock() - async def _async_start(_inputs, **_kwargs): + async def _async_start(_inputs=None, **_kwargs): yield graph.async_start = _async_start @@ -574,7 +574,7 @@ class TestRunFlowUserId: graph.edges = [] graph.prepare = MagicMock() - async def _async_start(_inputs, **_kwargs): + async def _async_start(_inputs=None, **_kwargs): yield graph.async_start = _async_start @@ -671,7 +671,7 @@ class TestRunFlowFallbackToEnvVars: graph.edges = [] graph.prepare = MagicMock() - async def _async_start(_inputs, **kwargs): + async def _async_start(_inputs=None, **kwargs): captured.update(kwargs) yield @@ -743,7 +743,7 @@ class TestRunFlowGlobalVariables: mock_graph.edges = [] mock_graph.prepare = MagicMock() - async def mock_async_start(_inputs, **_kwargs): + async def mock_async_start(_inputs=None, **_kwargs): yield mock_graph.async_start = mock_async_start @@ -784,7 +784,7 @@ graph = Graph(chat_input, chat_output) mock_graph.edges = [] mock_graph.prepare = MagicMock() - async def mock_async_start(_inputs, **_kwargs): + async def mock_async_start(_inputs=None, **_kwargs): yield mock_graph.async_start = mock_async_start @@ -825,7 +825,7 @@ class TestRunFlowOutputFormats: mock_graph.edges = [] mock_graph.prepare = MagicMock() - async def mock_async_start(_inputs, **_kwargs): + async def mock_async_start(_inputs=None, **_kwargs): yield mock_graph.async_start = mock_async_start @@ -938,7 +938,7 @@ class TestRunFlowTiming: mock_result.vertex.display_name = "TestComponent" mock_result.vertex.id = "test-id-123" - async def mock_async_start(_inputs, **_kwargs): + async def mock_async_start(_inputs=None, **_kwargs): yield mock_result mock_graph.async_start = mock_async_start @@ -1133,7 +1133,7 @@ class TestRunFlowVariableValidation: mock_graph.edges = [] mock_graph.prepare = MagicMock() - async def mock_async_start(_inputs, **_kwargs): + async def mock_async_start(_inputs=None, **_kwargs): yield mock_graph.async_start = mock_async_start @@ -1169,7 +1169,7 @@ class TestRunFlowInputValueHandling: mock_graph.edges = [] mock_graph.prepare = MagicMock() - async def mock_async_start(_inputs, **_kwargs): + async def mock_async_start(_inputs=None, **_kwargs): yield mock_graph.async_start = mock_async_start @@ -1207,7 +1207,7 @@ class TestRunFlowInputValueHandling: mock_graph.edges = [] mock_graph.prepare = MagicMock() - async def mock_async_start(_inputs, **_kwargs): + async def mock_async_start(_inputs=None, **_kwargs): yield mock_graph.async_start = mock_async_start @@ -1264,7 +1264,7 @@ class TestRunFlowJsonFileExecution: mock_graph.edges = [] mock_graph.prepare = MagicMock() - async def mock_async_start(_inputs, **_kwargs): + async def mock_async_start(_inputs=None, **_kwargs): yield mock_graph.async_start = mock_async_start @@ -1364,7 +1364,7 @@ class TestRunFlowExecutionErrors: mock_graph.edges = [] mock_graph.prepare = MagicMock() - async def failing_async_start(_inputs, **_kwargs): + async def failing_async_start(_inputs=None, **_kwargs): msg = "Execution failed" raise ValueError(msg) yield # Required to make it an async generator @@ -1662,7 +1662,7 @@ class TestUpgradeFlowOption: mock_graph.edges = [] mock_graph.prepare = MagicMock() - async def _async_start(_inputs, **_kwargs): + async def _async_start(_inputs=None, **_kwargs): yield mock_graph.async_start = _async_start @@ -1703,7 +1703,7 @@ class TestUpgradeFlowOption: mock_graph.edges = [] mock_graph.prepare = MagicMock() - async def _async_start(_inputs, **_kwargs): + async def _async_start(_inputs=None, **_kwargs): yield mock_graph.async_start = _async_start @@ -1740,7 +1740,7 @@ class TestUpgradeFlowOption: mock_graph.edges = [] mock_graph.prepare = MagicMock() - async def _async_start(_inputs, **_kwargs): + async def _async_start(_inputs=None, **_kwargs): yield mock_graph.async_start = _async_start @@ -1783,7 +1783,7 @@ class TestUpgradeFlowOption: mock_graph.edges = [] mock_graph.prepare = MagicMock() - async def _async_start(_inputs, **_kwargs): + async def _async_start(_inputs=None, **_kwargs): yield mock_graph.async_start = _async_start @@ -1822,7 +1822,7 @@ class TestUpgradeFlowOption: mock_graph.edges = [] mock_graph.prepare = MagicMock() - async def _async_start(_inputs, **_kwargs): + async def _async_start(_inputs=None, **_kwargs): yield mock_graph.async_start = _async_start diff --git a/src/lfx/tests/unit/services/settings/test_settings_composition.py b/src/lfx/tests/unit/services/settings/test_settings_composition.py index 3da6e6027e..b6cce368c3 100644 --- a/src/lfx/tests/unit/services/settings/test_settings_composition.py +++ b/src/lfx/tests/unit/services/settings/test_settings_composition.py @@ -178,6 +178,7 @@ EXPECTED_FIELDS = { "redis_queue_polling_stale_threshold_s", "redis_queue_polling_watchdog_interval_s", "max_ingestion_timeout_secs", + "executor_kind", # UiSettings "embedded_mode", "hide_getting_started_progress", diff --git a/src/lfx/tests/unit/services/test_service_manager.py b/src/lfx/tests/unit/services/test_service_manager.py index 17008719f5..f2f906e9a0 100644 --- a/src/lfx/tests/unit/services/test_service_manager.py +++ b/src/lfx/tests/unit/services/test_service_manager.py @@ -359,6 +359,27 @@ class TestTeardown: # Services should be cleared assert ServiceType.STORAGE_SERVICE not in service_manager.services + @pytest.mark.asyncio + async def test_teardown_clears_factories_registered_flag(self, service_manager): + """teardown() must clear the factories-registered flag, not just the dict. + + get_service() (both lfx and langflow) re-registers factories only when + are_factories_registered() returns False. teardown() empties + self.factories, so if it leaves the flag set, every later lookup skips + re-registration and raises NoFactoryRegisteredError. That surfaced as + flow-execution 500s once Graph.arun routed through the executor service + and a sibling test had torn the global manager down. + """ + service_manager.register_service_class(ServiceType.STORAGE_SERVICE, LocalStorageService) + service_manager.get(ServiceType.STORAGE_SERVICE) + service_manager.set_factory_registered() + assert service_manager.are_factories_registered() is True + + await service_manager.teardown() + + assert service_manager.factories == {} + assert service_manager.are_factories_registered() is False + class TestConfigDirectorySource: """Tests for config_dir parameter with real services."""