mirror of
https://github.com/langflow-ai/langflow.git
synced 2026-07-23 15:09:37 +08:00
feat(lfx): introduce pluggable executor seam (#12981)
* feat(execution): scaffold execution package with test isolation fixture * feat(execution): add core types (Unit, StepResult, RunComplete) * feat(execution): add Executor abstract base class * feat(execution): add identity partition function * feat(execution): add executor registry (register/get only) * feat(execution): add InProcessExecutor (streaming + legacy passthrough) * feat(execution): add Coordinator with streaming run + run_to_completion * feat(execution): default singleton with in-process pre-registered + setter * feat(execution): route Graph.arun through coordinator (seam #1) * feat(execution): route flow_executor through coordinator (seam #2) * feat(execution): route CLI and run/base through coordinator (seam #3) * feat(execution): route Loop subgraph through coordinator (seam #4) * fix(execution): explicit legacy dispatch flag, forward fallback_to_env_vars, real-graph concurrency tests * refactor(execution): add Coordinator.stream helper, lift imports, replace em-dash * test(execution): shared simple_graph fixture, reset_default_coordinator helper, real-executor loop tests, behavior assertions * refactor(execution): wrap default singletons in _Defaults holder, drop noqas * feat(execution): make executor seam pluggable via lfx services Introduce an ExecutorService that owns the registry and the default coordinator, discovered through the standard lfx ServiceManager (parity with cache, database, storage, etc). Add an executor_kind setting so the default kind can be configured without touching code, and an lfx.executors entry-point group for third-party executors. Entry-point discovery refuses to overwrite an existing kind so an installed package cannot silently replace the built-in in-process executor; explicit replacement still works through ExecutorService.register(). The lfx.execution public API (get_default_coordinator, get_default_registry, set_default_coordinator, reset_default_coordinator) now resolves through the service manager but keeps the same surface, so existing call sites in Graph.arun, flow_executor, the CLI, run/base, and loop_utils are unchanged. Also surface root-cause tracebacks in lfx.services.deps.get_service: the helper still returns None on failure (callers like get_db_service rely on that), but it now logs the exception so init failures stop disappearing into the void. * test(execution): contract suite, cancellation guarantees, seam docstrings Robustness pass on the execution seam after using it to integrate an external executor end-to-end. Four things this pins down: 1. Reusable executor contract suite - tests/unit/execution/test_executor_contract.py - ExecutorContract with 7 universal seam guarantees: kind shape, execute() returns an AsyncIterator, stream ends with exactly one RunComplete, instance is reusable across runs, concurrent runs are isolated, consumer cancellation does not hang or leak, execute() signature is stable. - TestInProcessExecutorContract subclasses and provides fixtures. Downstream executors (stepflow, future remote/sandbox) get the same battery by subclassing and overriding two fixtures. 2. Cancellation/cleanup correctness - tests/unit/execution/test_cancellation.py - InProcessExecutor wraps graph.async_start() in try/finally with explicit aclose so consumer aclose cascades to the underlying graph generator. Previously the inner generator was abandoned and only finalized on a GC pass: a real leak for executors holding subprocesses or sockets. - Coordinator.run and Coordinator.stream propagate aclose the same way so the full chain (consumer -> coordinator -> executor -> graph) cleans up deterministically. 3. Documented seam contracts - Unit.runtime_options: free-form bag, "_"-prefixed keys reserved for executor-internal flags, common in-process keys listed. - StepResult.payload: untyped on purpose; each executor defines its own event vocabulary; consumers normalize at consumption site. - RunComplete.outputs: only the legacy in-process path populates it; streaming consumers should collect from StepResult.payload. - Executor ABC: lifecycle (shared instance, must be reusable, must tolerate concurrent execute() calls and consumer aclose). 4. Coordinator-routed E2E suite - tests/unit/execution/test_coordinator_e2e.py - Real Graph through Coordinator -> registry -> InProcessExecutor chain. - Asserts payload shape, RunComplete never leaks to stream() consumers, dispatch routes to configured kind, registry returns same instance. Full execution suite: 63 passing. flow_executor coordinator test: passing. * fix(execution): close async generators on early exit and harden seam tests - run_to_completion and Coordinator.stream wrap iteration in contextlib.aclosing so the underlying generator is finalized when the loop returns or raises early - execute_loop_body wraps its stream in aclosing for the same mid-iteration guarantee - drop the unused inputs param from Coordinator.stream; the streaming path reads initial_inputs from runtime_options, so inputs was a silent no-op, and document it - test_coordinator_uses_settings_executor_kind uses monkeypatch + a finalizer so a mid-test failure can't poison session settings - test_registry_returns_same_executor_instance_across_runs builds a fresh graph per run - test_set_default_coordinator_overrides_singleton restores the singleton afterward * fix(execution): register ExecutorService factory and update seam test mocks - register ExecutorServiceFactory in register_all_service_factories so EXECUTOR_SERVICE resolves at app startup; without it the seam raised NoFactoryRegisteredError across integration and mcp tests - test_base.py graph.async_start mocks now default the leading positional param (the seam calls async_start with inputs as a keyword), fixing the 'missing 1 required positional argument' failures - add executor_kind to the settings composition EXPECTED_FIELDS set * fix(execution): reject non-Executor entry-point plugins; doc corrections - ExecutorService discovery skips a loaded entry-point object that is not an Executor instance instead of registering it and failing only at execute() time - Coordinator.run docstring notes the single-Unit / single-RunComplete assumption that identity_partition currently guarantees - correct two stale docstrings: Unit's session_id is honored only on the legacy passthrough path, and ExecutorService.teardown no longer claims the service manager keeps the cached instance * fix(lfx): reset factory-registered flag on ServiceManager teardown teardown() emptied self.factories but left factory_registered set, so get_service()'s lazy re-registration was skipped and later lookups raised NoFactoryRegisteredError. Once Graph.arun started routing through the executor service, that surfaced as intermittent flow-execution 500s in the backend unit suite whenever a sibling test had torn the global manager down. * test(loop): accept executor kwargs in subgraph async_start mock The loop body now runs its subgraph through get_default_coordinator().stream, so InProcessExecutor calls async_start with inputs/max_iterations/config/ reset_output_values/fallback_to_env_vars alongside event_manager. The mock only accepted event_manager and raised TypeError on the extra kwargs. Widen it to **_kwargs while still capturing event_manager for the assertion.
This commit is contained in:
committed by
GitHub
parent
7dd634790a
commit
49fc1f862f
@ -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
|
||||
|
||||
@ -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()
|
||||
|
||||
|
||||
|
||||
@ -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
|
||||
@ -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": {}})
|
||||
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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()
|
||||
|
||||
64
src/lfx/src/lfx/execution/__init__.py
Normal file
64
src/lfx/src/lfx/execution/__init__.py
Normal file
@ -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",
|
||||
]
|
||||
0
src/lfx/src/lfx/execution/backends/__init__.py
Normal file
0
src/lfx/src/lfx/execution/backends/__init__.py
Normal file
58
src/lfx/src/lfx/execution/backends/in_process.py
Normal file
58
src/lfx/src/lfx/execution/backends/in_process.py
Normal file
@ -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()
|
||||
128
src/lfx/src/lfx/execution/coordinator.py
Normal file
128
src/lfx/src/lfx/execution/coordinator.py
Normal file
@ -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
|
||||
76
src/lfx/src/lfx/execution/executor.py
Normal file
76
src/lfx/src/lfx/execution/executor.py
Normal file
@ -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``.
|
||||
"""
|
||||
16
src/lfx/src/lfx/execution/partitioner.py
Normal file
16
src/lfx/src/lfx/execution/partitioner.py
Normal file
@ -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 {})]
|
||||
46
src/lfx/src/lfx/execution/registry.py
Normal file
46
src/lfx/src/lfx/execution/registry.py
Normal file
@ -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
|
||||
100
src/lfx/src/lfx/execution/types.py
Normal file
100
src/lfx/src/lfx/execution/types.py
Normal file
@ -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]
|
||||
@ -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]
|
||||
|
||||
@ -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:
|
||||
|
||||
@ -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
|
||||
|
||||
|
||||
|
||||
1
src/lfx/src/lfx/services/executor/__init__.py
Normal file
1
src/lfx/src/lfx/services/executor/__init__.py
Normal file
@ -0,0 +1 @@
|
||||
"""Executor service module."""
|
||||
24
src/lfx/src/lfx/services/executor/factory.py
Normal file
24
src/lfx/src/lfx/services/executor/factory.py
Normal file
@ -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)
|
||||
124
src/lfx/src/lfx/services/executor/service.py
Normal file
124
src/lfx/src/lfx/services/executor/service.py
Normal file
@ -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 = "<unknown>"
|
||||
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", "<unknown>")
|
||||
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()
|
||||
@ -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)
|
||||
|
||||
@ -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"
|
||||
|
||||
@ -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):
|
||||
|
||||
@ -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.
|
||||
|
||||
0
src/lfx/tests/unit/execution/__init__.py
Normal file
0
src/lfx/tests/unit/execution/__init__.py
Normal file
24
src/lfx/tests/unit/execution/conftest.py
Normal file
24
src/lfx/tests/unit/execution/conftest.py
Normal file
@ -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)
|
||||
187
src/lfx/tests/unit/execution/test_cancellation.py
Normal file
187
src/lfx/tests/unit/execution/test_cancellation.py
Normal file
@ -0,0 +1,187 @@
|
||||
"""Cancellation and cleanup tests for the execution seam.
|
||||
|
||||
The contract suite covers the bare-minimum ``aclose()`` shape via
|
||||
``test_consumer_cancellation_does_not_hang_or_leak``. This module exercises the
|
||||
stricter behaviour the seam needs in production: explicit cleanup at the top of
|
||||
the consumer chain MUST cascade all the way down to the underlying graph iterator.
|
||||
|
||||
What we deliberately do NOT test here:
|
||||
|
||||
- Implicit GC-based finalization (``async for ... break`` without ``aclosing``,
|
||||
``task.cancel()`` then immediate assert). CPython finalizes abandoned async
|
||||
generators on a GC pass; asyncio runs the registered finalizer on event-loop
|
||||
shutdown. Neither is observable inside a single test, and asserting either
|
||||
would test CPython, not us. Consumers that need deterministic cleanup MUST
|
||||
use ``contextlib.aclosing`` or call ``aclose()`` explicitly.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from contextlib import aclosing
|
||||
|
||||
import pytest
|
||||
from lfx.execution.backends.in_process import InProcessExecutor
|
||||
from lfx.execution.coordinator import Coordinator
|
||||
from lfx.execution.registry import ExecutorRegistry
|
||||
from lfx.execution.types import StepResult, Unit
|
||||
|
||||
|
||||
class _InstrumentedGraph:
|
||||
"""Minimal graph stub whose ``async_start`` runs forever until closed.
|
||||
|
||||
We track entry, items yielded, and finalizer invocation so tests can assert
|
||||
the iterator was properly closed rather than abandoned.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.entered: bool = False
|
||||
self.finalized: bool = False
|
||||
self.items_yielded: int = 0
|
||||
|
||||
async def async_start(self, **_kwargs):
|
||||
self.entered = True
|
||||
try:
|
||||
for i in range(10_000):
|
||||
self.items_yielded += 1
|
||||
yield f"item-{i}"
|
||||
await asyncio.sleep(0)
|
||||
finally:
|
||||
self.finalized = True
|
||||
|
||||
|
||||
def _unit(graph: _InstrumentedGraph) -> Unit:
|
||||
return Unit(graph=graph, inputs=[], runtime_options={})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aclose_on_executor_finalizes_underlying_graph_iteration():
|
||||
"""``aclose()`` on the executor's generator MUST run the underlying graph's finalizer.
|
||||
|
||||
Without the explicit try/finally inside ``InProcessExecutor.execute``, the inner
|
||||
generator is abandoned and the graph's ``finally:`` block (event managers,
|
||||
connections, vertex teardown) only runs on a future GC pass.
|
||||
"""
|
||||
graph = _InstrumentedGraph()
|
||||
executor = InProcessExecutor()
|
||||
stream = executor.execute(_unit(graph))
|
||||
|
||||
first = await anext(stream)
|
||||
assert isinstance(first, StepResult)
|
||||
assert graph.entered
|
||||
assert not graph.finalized
|
||||
|
||||
await stream.aclose()
|
||||
assert graph.finalized, "underlying async generator was not finalized on aclose()"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aclosing_helper_chains_cleanup_through_executor():
|
||||
"""``contextlib.aclosing`` MUST trigger the same cascade as explicit ``aclose``."""
|
||||
graph = _InstrumentedGraph()
|
||||
executor = InProcessExecutor()
|
||||
|
||||
async with aclosing(executor.execute(_unit(graph))) as stream:
|
||||
async for _ in stream:
|
||||
break
|
||||
|
||||
assert graph.finalized
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_coordinator_run_cascades_aclose_to_executor_and_graph():
|
||||
"""Cleanup must propagate through every seam layer.
|
||||
|
||||
Consumer aclose -> Coordinator.run -> Executor.execute -> graph.async_start.
|
||||
A regression at any layer would leave the graph generator hanging.
|
||||
"""
|
||||
graph = _InstrumentedGraph()
|
||||
registry = ExecutorRegistry()
|
||||
registry.register(InProcessExecutor())
|
||||
coordinator = Coordinator(registry=registry)
|
||||
|
||||
async with aclosing(coordinator.run(graph, inputs=[])) as stream:
|
||||
first = await anext(stream)
|
||||
assert first is not None
|
||||
|
||||
assert graph.finalized, "Coordinator.run did not cascade aclose to the graph"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_coordinator_stream_cascades_aclose_to_graph():
|
||||
"""``Coordinator.stream`` must cascade cleanup the same way ``Coordinator.run`` does."""
|
||||
graph = _InstrumentedGraph()
|
||||
registry = ExecutorRegistry()
|
||||
registry.register(InProcessExecutor())
|
||||
coordinator = Coordinator(registry=registry)
|
||||
|
||||
async with aclosing(coordinator.stream(graph)) as stream:
|
||||
async for _ in stream:
|
||||
break
|
||||
|
||||
assert graph.finalized
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cleanup_bounded_by_timeout():
|
||||
"""Finalization must complete within a bounded window even on a slow CI box.
|
||||
|
||||
A regression that swaps in a non-cancellable resource (blocking subprocess
|
||||
wait, sync lock) would hang here rather than slow.
|
||||
"""
|
||||
graph = _InstrumentedGraph()
|
||||
executor = InProcessExecutor()
|
||||
|
||||
async def drive():
|
||||
gen = executor.execute(_unit(graph))
|
||||
await anext(gen)
|
||||
await gen.aclose()
|
||||
|
||||
await asyncio.wait_for(drive(), timeout=2.0)
|
||||
assert graph.finalized
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_exception_in_consumer_still_finalizes_graph():
|
||||
"""An exception inside the consumer MUST still trigger cleanup via ``aclosing``."""
|
||||
graph = _InstrumentedGraph()
|
||||
executor = InProcessExecutor()
|
||||
|
||||
class _ConsumerError(RuntimeError):
|
||||
pass
|
||||
|
||||
async def consume_then_raise() -> None:
|
||||
async with aclosing(executor.execute(_unit(graph))) as stream:
|
||||
async for _ in stream:
|
||||
msg = "downstream bug"
|
||||
raise _ConsumerError(msg)
|
||||
|
||||
with pytest.raises(_ConsumerError):
|
||||
await consume_then_raise()
|
||||
|
||||
assert graph.finalized
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_consumers_each_clean_up_independently():
|
||||
"""Concurrent runs on one executor instance must clean up independently.
|
||||
|
||||
Regression guard for any future change that introduces shared cleanup state
|
||||
on the executor (which would couple unrelated consumers).
|
||||
"""
|
||||
graph_a = _InstrumentedGraph()
|
||||
graph_b = _InstrumentedGraph()
|
||||
executor = InProcessExecutor()
|
||||
|
||||
async def drive(graph: _InstrumentedGraph, items_before_stop: int) -> None:
|
||||
async with aclosing(executor.execute(_unit(graph))) as stream:
|
||||
count = 0
|
||||
async for _ in stream:
|
||||
count += 1
|
||||
if count >= items_before_stop:
|
||||
break
|
||||
|
||||
await asyncio.gather(drive(graph_a, 1), drive(graph_b, 3))
|
||||
|
||||
assert graph_a.finalized
|
||||
assert graph_b.finalized
|
||||
59
src/lfx/tests/unit/execution/test_coordinator.py
Normal file
59
src/lfx/tests/unit/execution/test_coordinator.py
Normal file
@ -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)
|
||||
174
src/lfx/tests/unit/execution/test_coordinator_e2e.py
Normal file
174
src/lfx/tests/unit/execution/test_coordinator_e2e.py
Normal file
@ -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"
|
||||
32
src/lfx/tests/unit/execution/test_default_registry.py
Normal file
32
src/lfx/tests/unit/execution/test_default_registry.py
Normal file
@ -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
|
||||
37
src/lfx/tests/unit/execution/test_executor.py
Normal file
37
src/lfx/tests/unit/execution/test_executor.py
Normal file
@ -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)
|
||||
228
src/lfx/tests/unit/execution/test_executor_contract.py
Normal file
228
src/lfx/tests/unit/execution/test_executor_contract.py
Normal file
@ -0,0 +1,228 @@
|
||||
"""Reusable contract test suite for any ``lfx.execution.Executor`` implementation.
|
||||
|
||||
The seam has a small but load-bearing set of behaviors that any executor MUST satisfy
|
||||
to be safely swappable behind ``Coordinator``. This module spells them out as a
|
||||
parametrizable suite so:
|
||||
|
||||
- The built-in ``InProcessExecutor`` is checked here against the same suite the
|
||||
external executors (stepflow, future remote/sandbox, etc.) will be checked against.
|
||||
- Authors of new executors can subclass ``ExecutorContract`` in their own project,
|
||||
override the two fixtures, and inherit the full battery of behavioural assertions
|
||||
rather than re-deriving them from the docstring of the ABC.
|
||||
|
||||
The contract intentionally does NOT pin executor-specific event shapes -- only the
|
||||
universal seam guarantees. If a future executor needs to add stricter checks for its
|
||||
own payloads, those go in an executor-specific test file, not here.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import inspect
|
||||
from collections.abc import AsyncIterator, Callable
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import pytest
|
||||
from lfx.execution.types import RunComplete, StepResult, Unit
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from lfx.execution.executor import Executor
|
||||
|
||||
|
||||
class ExecutorContract:
|
||||
"""Base class of behavioural assertions every Executor must satisfy.
|
||||
|
||||
Subclasses MUST override two fixtures:
|
||||
|
||||
- ``contract_executor``: returns a freshly constructed ``Executor`` instance.
|
||||
- ``contract_unit_factory``: returns a callable ``() -> Unit`` that produces a
|
||||
fresh, runnable ``Unit`` each call. The unit should drive a minimal but real
|
||||
execution -- not a stub -- so the test actually exercises the executor's
|
||||
end-to-end path.
|
||||
|
||||
Tests are designed to be independent: each test creates its own unit through
|
||||
the factory so subclasses don't have to worry about state bleed between cases.
|
||||
"""
|
||||
|
||||
@pytest.fixture
|
||||
def contract_executor(self) -> Executor:
|
||||
msg = "Subclasses must override the contract_executor fixture"
|
||||
raise NotImplementedError(msg)
|
||||
|
||||
@pytest.fixture
|
||||
def contract_unit_factory(self) -> Callable[[], Unit]:
|
||||
msg = "Subclasses must override the contract_unit_factory fixture"
|
||||
raise NotImplementedError(msg)
|
||||
|
||||
# --- Universal seam guarantees ----------------------------------------------------
|
||||
|
||||
def test_kind_is_a_non_empty_string(self, contract_executor: Executor) -> None:
|
||||
"""``kind`` is a ClassVar on the ABC; it must be set and non-empty."""
|
||||
assert isinstance(contract_executor.kind, str)
|
||||
assert contract_executor.kind, "Executor.kind must be a non-empty string"
|
||||
|
||||
def test_execute_returns_async_iterator(
|
||||
self,
|
||||
contract_executor: Executor,
|
||||
contract_unit_factory: Callable[[], Unit],
|
||||
) -> None:
|
||||
"""``execute(unit)`` MUST return an async iterator.
|
||||
|
||||
The ABC types it as such; consumers (Coordinator.run, Coordinator.stream,
|
||||
run_to_completion) rely on it.
|
||||
"""
|
||||
stream = contract_executor.execute(contract_unit_factory())
|
||||
assert isinstance(stream, AsyncIterator), (
|
||||
f"{type(contract_executor).__name__}.execute() returned a "
|
||||
f"{type(stream).__name__}, expected AsyncIterator. Consumers can't iterate "
|
||||
"this. Common cause: forgetting to make the method an `async def` "
|
||||
"generator."
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_ends_with_exactly_one_run_complete(
|
||||
self,
|
||||
contract_executor: Executor,
|
||||
contract_unit_factory: Callable[[], Unit],
|
||||
) -> None:
|
||||
"""Streams emit zero-or-more ``StepResult`` items terminated by one ``RunComplete``."""
|
||||
items = [item async for item in contract_executor.execute(contract_unit_factory())]
|
||||
|
||||
assert items, "executor yielded nothing -- a RunComplete is mandatory"
|
||||
assert isinstance(items[-1], RunComplete), (
|
||||
f"final item was {type(items[-1]).__name__}; the seam requires a terminal "
|
||||
"RunComplete so run_to_completion() and downstream collectors can detect end-of-run."
|
||||
)
|
||||
|
||||
run_completes = [i for i in items if isinstance(i, RunComplete)]
|
||||
assert len(run_completes) == 1, (
|
||||
f"executor yielded {len(run_completes)} RunComplete items; the seam allows "
|
||||
"at most one (and requires exactly one)."
|
||||
)
|
||||
|
||||
non_terminal = items[:-1]
|
||||
bad = [i for i in non_terminal if not isinstance(i, StepResult)]
|
||||
assert not bad, (
|
||||
f"all pre-terminal items must be StepResult; got: {[type(i).__name__ for i in bad]}. "
|
||||
"Coordinator.stream() relies on this to filter; mixing other types makes "
|
||||
"the stream untyped at the seam boundary."
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_executor_is_reusable_across_runs(
|
||||
self,
|
||||
contract_executor: Executor,
|
||||
contract_unit_factory: Callable[[], Unit],
|
||||
) -> None:
|
||||
"""An Executor MUST be reusable: state that carries between runs is a bug.
|
||||
|
||||
The registry returns the same instance for every ``coordinator.run()`` call,
|
||||
so any internal state pinned to ``self`` will silently corrupt later runs.
|
||||
"""
|
||||
first = [item async for item in contract_executor.execute(contract_unit_factory())]
|
||||
second = [item async for item in contract_executor.execute(contract_unit_factory())]
|
||||
|
||||
assert isinstance(first[-1], RunComplete)
|
||||
assert isinstance(second[-1], RunComplete)
|
||||
# Same shape across runs proves no carry-over short-circuit. We don't assert
|
||||
# equality of payloads because some executors include timestamps / run IDs.
|
||||
assert len(first) >= 1
|
||||
assert len(second) >= 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_runs_on_one_instance_are_isolated(
|
||||
self,
|
||||
contract_executor: Executor,
|
||||
contract_unit_factory: Callable[[], Unit],
|
||||
) -> None:
|
||||
"""Concurrent runs on one executor instance MUST be isolated.
|
||||
|
||||
Coordinator may launch concurrent runs (Loop subgraphs, parallel API
|
||||
requests). An executor that crosses state between them is broken.
|
||||
"""
|
||||
|
||||
async def drain() -> list[object]:
|
||||
return [item async for item in contract_executor.execute(contract_unit_factory())]
|
||||
|
||||
a, b = await asyncio.gather(drain(), drain())
|
||||
assert isinstance(a[-1], RunComplete)
|
||||
assert isinstance(b[-1], RunComplete)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_consumer_cancellation_does_not_hang_or_leak(
|
||||
self,
|
||||
contract_executor: Executor,
|
||||
contract_unit_factory: Callable[[], Unit],
|
||||
) -> None:
|
||||
"""Consumer cancellation via ``aclose()`` MUST finalize within a bounded window.
|
||||
|
||||
The executor's async generator is closed by the consumer; any internal
|
||||
context managers, subprocesses, or sockets MUST be released. A strict
|
||||
timeout surfaces regressions as test failures rather than CI hangs.
|
||||
"""
|
||||
gen = contract_executor.execute(contract_unit_factory())
|
||||
|
||||
try:
|
||||
first = await asyncio.wait_for(anext(gen), timeout=10.0)
|
||||
except StopAsyncIteration:
|
||||
# Executor produced no items at all -- still a valid stream (just an empty one).
|
||||
return
|
||||
|
||||
assert isinstance(first, (StepResult, RunComplete)), (
|
||||
f"first yielded item was {type(first).__name__}; seam requires StepResult or RunComplete"
|
||||
)
|
||||
|
||||
# Drop the iterator. If the executor holds a subprocess / network / lock, this
|
||||
# is the moment of truth: aclose() must finalize without blocking forever.
|
||||
await asyncio.wait_for(gen.aclose(), timeout=10.0)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_signature_takes_exactly_one_unit(self, contract_executor: Executor) -> None:
|
||||
"""Pin the ABC signature so executors don't drift away from a single ``Unit`` parameter.
|
||||
|
||||
``execute(*args, **kwargs)`` shapes bypass the ``Unit`` value object and
|
||||
make the seam opaque to introspection.
|
||||
"""
|
||||
sig = inspect.signature(contract_executor.execute)
|
||||
params = [p for p in sig.parameters.values() if p.name != "self"]
|
||||
assert len(params) == 1, (
|
||||
f"{type(contract_executor).__name__}.execute must accept exactly one positional "
|
||||
f"Unit argument; got signature: {sig}"
|
||||
)
|
||||
|
||||
|
||||
class TestInProcessExecutorContract(ExecutorContract):
|
||||
"""Run the contract suite against the built-in InProcessExecutor.
|
||||
|
||||
This is the reference implementation: if the contract suite is wrong, this is
|
||||
where it will fail first. Stepflow / remote / sandbox executors are exercised
|
||||
against the same suite in their own test files.
|
||||
"""
|
||||
|
||||
@pytest.fixture
|
||||
def contract_executor(self) -> Executor:
|
||||
from lfx.execution.backends.in_process import InProcessExecutor
|
||||
|
||||
return InProcessExecutor()
|
||||
|
||||
@pytest.fixture
|
||||
def contract_unit_factory(self, simple_graph) -> Callable[[], Unit]: # noqa: ARG002
|
||||
# ``simple_graph`` is a pytest fixture from conftest.py. We don't capture it
|
||||
# by closure across runs because Graph instances accumulate state; instead we
|
||||
# rebuild on every call so each unit is fresh.
|
||||
from lfx.components.input_output import ChatInput, ChatOutput
|
||||
from lfx.graph import Graph
|
||||
from lfx.schema.schema import InputValueRequest
|
||||
|
||||
def factory() -> Unit:
|
||||
ci = ChatInput(_id="chat_input")
|
||||
ci.set(should_store_message=False)
|
||||
co = ChatOutput(input_value="test", _id="chat_output")
|
||||
co.set(sender_name=ci.message_response)
|
||||
return Unit(
|
||||
graph=Graph(ci, co),
|
||||
inputs=[],
|
||||
runtime_options={"initial_inputs": InputValueRequest(input_value="hi")},
|
||||
)
|
||||
|
||||
return factory
|
||||
228
src/lfx/tests/unit/execution/test_executor_service.py
Normal file
228
src/lfx/tests/unit/execution/test_executor_service.py
Normal file
@ -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]
|
||||
@ -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"]
|
||||
146
src/lfx/tests/unit/execution/test_in_process_executor.py
Normal file
146
src/lfx/tests/unit/execution/test_in_process_executor.py
Normal file
@ -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"]
|
||||
17
src/lfx/tests/unit/execution/test_partitioner.py
Normal file
17
src/lfx/tests/unit/execution/test_partitioner.py
Normal file
@ -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 == {}
|
||||
47
src/lfx/tests/unit/execution/test_registry.py
Normal file
47
src/lfx/tests/unit/execution/test_registry.py
Normal file
@ -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")
|
||||
28
src/lfx/tests/unit/execution/test_types.py
Normal file
28
src/lfx/tests/unit/execution/test_types.py
Normal file
@ -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"}
|
||||
@ -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
|
||||
|
||||
@ -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",
|
||||
|
||||
@ -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."""
|
||||
|
||||
Reference in New Issue
Block a user