diff --git a/src/lfx/src/lfx/execution/executor.py b/src/lfx/src/lfx/execution/executor.py new file mode 100644 index 0000000000..1546c1cf51 --- /dev/null +++ b/src/lfx/src/lfx/execution/executor.py @@ -0,0 +1,19 @@ +"""Executor abstract base class.""" + +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): + kind: ClassVar[str] + + @abstractmethod + def execute(self, unit: Unit) -> AsyncIterator[StepResult | RunComplete]: + """Yield StepResult items, end with a RunComplete.""" diff --git a/src/lfx/tests/unit/execution/test_executor.py b/src/lfx/tests/unit/execution/test_executor.py new file mode 100644 index 0000000000..3d226e8f01 --- /dev/null +++ b/src/lfx/tests/unit/execution/test_executor.py @@ -0,0 +1,37 @@ +import pytest +from lfx.execution.executor import Executor +from lfx.execution.types import RunComplete, StepResult, Unit + + +def test_executor_is_abstract(): + with pytest.raises(TypeError, match="abstract"): + Executor() + + +@pytest.mark.asyncio +async def test_concrete_executor_yields_steps_then_run_complete(): + class Toy(Executor): + kind = "toy" + + async def execute(self, unit): # noqa: ARG002 + yield StepResult(payload={"vertex_id": "a"}) + yield StepResult(payload={"vertex_id": "b"}) + yield RunComplete(outputs=["done"]) + + items = [item async for item in Toy().execute(Unit(graph=object(), inputs=[]))] + assert len(items) == 3 + assert isinstance(items[-1], RunComplete) + assert items[-1].outputs == ["done"] + + +@pytest.mark.asyncio +async def test_executor_must_yield_run_complete_last(): + class Toy(Executor): + kind = "toy" + + async def execute(self, unit): # noqa: ARG002 + yield StepResult(payload={}) + yield RunComplete(outputs=[]) + + items = [item async for item in Toy().execute(Unit(graph=object(), inputs=[]))] + assert isinstance(items[-1], RunComplete)