async-conveyor 0.0.1__tar.gz

async_conveyor-0.0.1/PKG-INFO

@@ -0,0 +1,5 @@
+Metadata-Version: 2.3
+Name: async-conveyor
+Version: 0.0.1
+Requires-Dist: anyio>=4.15.1
+Requires-Python: >=3.11

async_conveyor-0.0.1/pyproject.toml

@@ -0,0 +1,43 @@
+[project]
+name = "async-conveyor"
+version = "0.0.1"
+requires-python = ">=3.11"
+dependencies = ["anyio>=4.15.1"]
+
+[build-system]
+requires = ["uv_build>=0.11,<0.12"]
+build-backend = "uv_build"
+
+[dependency-groups]
+dev = [
+    "pytest>=8.0",
+    "ruff>=0.6",
+    "mypy>=1.11",
+]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+addopts = "-q"
+
+[tool.ruff]
+line-length = 120
+target-version = "py311"
+
+[tool.ruff.lint]
+select = [
+    "E",
+    "F",
+    "W",
+    "I",
+    "UP",
+    "B",
+    "SIM",
+]
+
+[tool.mypy]
+python_version = "3.11"
+strict = true
+files = [
+    "src",
+    "tests",
+]

async_conveyor-0.0.1/pyproject.toml.orig

@@ -0,0 +1,34 @@
+[project]
+name = "async-conveyor"
+version = "0.0.1"
+requires-python = ">=3.11"
+dependencies = [
+    "anyio>=4.15.1",
+]
+
+[build-system]
+requires = ["uv_build>=0.11,<0.12"]
+build-backend = "uv_build"
+
+[dependency-groups]
+dev = [
+    "pytest>=8.0",
+    "ruff>=0.6",
+    "mypy>=1.11",
+]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+addopts = "-q"
+
+[tool.ruff]
+line-length = 120
+target-version = "py311"
+
+[tool.ruff.lint]
+select = ["E", "F", "W", "I", "UP", "B", "SIM"]
+
+[tool.mypy]
+python_version = "3.11"
+strict = true
+files = ["src", "tests"]

async_conveyor-0.0.1/src/async_conveyor/__init__.py

@@ -0,0 +1,20 @@
+"""Run a DAG of pipeline stages concurrently on anyio, with bounded queues between them."""
+
+from .conveyor import BaseAsyncConveyor
+from .stage import AnyStage, BaseAsyncPipelineStage, BasePipelineStage, BaseSyncPipelineStage
+from .telemetry import ItemStats, LoggingTelemetry, StageSummary, SummaryTelemetry, Telemetry
+from .topology import TopologyError
+
+__all__ = [
+    "AnyStage",
+    "BaseAsyncConveyor",
+    "BaseAsyncPipelineStage",
+    "BasePipelineStage",
+    "BaseSyncPipelineStage",
+    "ItemStats",
+    "LoggingTelemetry",
+    "StageSummary",
+    "SummaryTelemetry",
+    "Telemetry",
+    "TopologyError",
+]

async_conveyor-0.0.1/src/async_conveyor/conveyor.py

@@ -0,0 +1,118 @@
+import abc
+import time
+from collections.abc import AsyncIterator
+from contextlib import AbstractAsyncContextManager, AsyncExitStack, asynccontextmanager
+from typing import Any, TypeVar
+
+import anyio
+
+from .stage import AnyStage, BasePipelineStage
+from .telemetry import ItemStats, LoggingTelemetry, Telemetry
+from .topology import TopologyMixin
+
+T = TypeVar("T")
+Stage = BasePipelineStage[Any, Any, Any]
+
+
+class BaseAsyncConveyor(TopologyMixin, abc.ABC):
+    def __init__(self, topology: dict[AnyStage, list[AnyStage]], telemetry: Telemetry | None = None):
+        """Base class for the conveyor that can run multiple stages concurrently.
+        Telemetry - metrics, tracing and error hooks; defaults to LoggingTelemetry.
+        Topology - routing for stages. Supports broadcast fan-out, doesn't support fan-in.
+        Exactly one root stage (the one that is not listed as anyone's output) is required. It is called
+        exactly once with item=None; the end of its generator is the end of the stream.
+        Keys and values are stage classes. Example:
+            {
+                Producer: [Embed, Metadata],  # broadcast fan-out
+                Embed: [Writer],  # direct
+            }
+        Unsupported:
+            {
+                Producer: [Embed, Metadata],
+                Qwe: [Metadata],  # fan-in
+                Metadata: [Producer],  # cycle
+            }
+        """
+        self.topology = topology
+        self.telemetry: Telemetry = telemetry if telemetry is not None else LoggingTelemetry()
+        self.validate_topology()
+
+    async def run(self) -> None:
+        """Run until the root's stream is exhausted and every stage has drained. See the class docstring."""
+        stages = self._wire()
+        try:
+            async with AsyncExitStack() as stack:
+                contexts = {s: await stack.enter_async_context(_shielded(s.context())) for s in stages}
+                async with anyio.create_task_group() as tg:
+                    for s in stages:
+                        tg.start_soon(self._run_stage, s, contexts[s])  # noqa
+        except BaseExceptionGroup as group:  # each stage has its own task group, so groups nest; unwrap to the cause
+            exc: BaseException = group
+            while isinstance(exc, BaseExceptionGroup) and len(exc.exceptions) == 1:
+                exc = exc.exceptions[0]
+            raise exc from None
+
+    def _wire(self) -> list[Stage]:
+        instances = {cls: cls() for cls in self.stages}
+        for parent, children in self.topology.items():
+            for child in children:
+                send, receive = anyio.create_memory_object_stream[Any](max_buffer_size=parent.output_queue_size)
+                instances[parent].outputs[child.name] = send
+                instances[child].inbox = receive
+        return list(instances.values())
+
+    async def _run_stage(self, stage: Stage, context: Any) -> None:
+        async with AsyncExitStack() as stack:
+            for out in stage.outputs.values():
+                await stack.enter_async_context(out)  # closing the queues tells the children nothing more is coming
+            if stage.inbox is None:
+                await self._process(stage, 0, None, context)
+                return
+            async with stage.inbox, anyio.create_task_group() as tg:
+                for worker in range(stage.workers):
+                    tg.start_soon(self._worker, stage, worker, context)  # noqa
+
+    async def _worker(self, stage: Stage, worker: int, context: Any) -> None:
+        assert stage.inbox is not None
+        async for item in stage.inbox:
+            await self._process(stage, worker, item, context)
+
+    async def _process(self, stage: Stage, worker: int, item: Any, context: Any) -> None:
+        started = time.perf_counter()
+        stage.active_workers += 1
+        try:
+            async with self.telemetry.item_span(stage.name, worker):
+                emitted, blocked = await stage.run_one(item, context)
+            stats = ItemStats(
+                stage_name=stage.name,
+                worker=worker,
+                duration=time.perf_counter() - started,
+                emitted=emitted,
+                blocked=blocked,
+                in_queue_size=stage.in_queue_size,
+                out_queue_sizes=stage.out_queue_sizes,
+                active_workers=stage.active_workers,
+            )
+        except (anyio.BrokenResourceError, anyio.ClosedResourceError):
+            return  # a consumer went away: the conveyor is shutting down and the cause is reported elsewhere
+        except Exception as exc:
+            await self.telemetry.on_error(stage.name, worker, exc, item)
+            raise
+        finally:
+            stage.active_workers -= 1
+        await self.telemetry.on_item_done(stats)
+
+
+@asynccontextmanager
+async def _shielded(cm: AbstractAsyncContextManager[T]) -> AsyncIterator[T]:
+    """`async with cm`, but the exit survives cancellation, so resources do get closed."""
+    value = await cm.__aenter__()
+    try:
+        yield value
+    except BaseException as exc:
+        with anyio.CancelScope(shield=True):
+            if not await cm.__aexit__(type(exc), exc, exc.__traceback__):
+                raise
+    else:
+        with anyio.CancelScope(shield=True):
+            await cm.__aexit__(None, None, None)

async_conveyor-0.0.1/src/async_conveyor/stage.py

@@ -0,0 +1,108 @@
+import abc
+import time
+from collections.abc import AsyncGenerator, Generator
+from contextlib import AbstractAsyncContextManager, nullcontext
+from typing import Any, Generic, TypeVar
+
+import anyio
+from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
+
+InT = TypeVar("InT")
+OutT = TypeVar("OutT")
+ContextT = TypeVar("ContextT")
+
+
+class BasePipelineStage(abc.ABC, Generic[InT, OutT, ContextT]):
+    """One stage of the conveyor. Stages are used as classes, not instances: all config is class-level.
+    Retries and timeouts are the responsibility of `process` itself; an exception escaping `process` stops the conveyor.
+    Items must be treated as immutable: on broadcast fan-out the same object is handed to every consumer.
+    """
+
+    name: str = NotImplemented  # unique name for this stage
+    workers: int = NotImplemented  # items processed concurrently (coroutines for async, threads for sync). 1 for root
+    output_queue_size: int = NotImplemented  # OOM protection - block producer when an output queue holds this many
+
+    def __init__(self) -> None:
+        """Runtime state. The conveyor instantiates every stage class once per run; users never do."""
+        self.inbox: MemoryObjectReceiveStream[InT] | None = None  # None for the root
+        self.outputs: dict[str, MemoryObjectSendStream[OutT]] = {}  # child stage name -> queue into it
+        self.active_workers = 0
+
+    @property
+    def in_queue_size(self) -> int:
+        return self.inbox.statistics().current_buffer_used if self.inbox is not None else 0
+
+    @property
+    def out_queue_sizes(self) -> dict[str, int]:
+        return {name: out.statistics().current_buffer_used for name, out in self.outputs.items()}
+
+    async def send(self, out: OutT) -> float:
+        """Broadcast one item to every child. Returns seconds spent waiting for room in their queues."""
+        if not self.outputs:
+            return 0.0
+        started = time.perf_counter()
+        for queue in self.outputs.values():
+            await queue.send(out)
+        return time.perf_counter() - started
+
+    @abc.abstractmethod
+    async def run_one(self, item: InT, context: ContextT) -> tuple[int, float]:
+        """Run process() on one item and send everything it yields downstream.
+        Returns (items emitted, seconds blocked sending)."""
+
+    @classmethod
+    def context(cls) -> AbstractAsyncContextManager[ContextT]:
+        """Resources shared by every process() call of this stage: HTTP clients, DB pools, models.
+        Entered once on the event loop before the stage's workers start and exited after they finish,
+        also on error or cancellation. Default: no resources, process() gets None (declare ContextT as None).
+        Override with `@asynccontextmanager`, or return the client itself when it is an async context manager:
+
+            @classmethod
+            @asynccontextmanager
+            async def context(cls) -> AsyncIterator[httpx.AsyncClient]:
+                async with httpx.AsyncClient() as client:
+                    yield client
+
+        Plain state that needs no closing is `return nullcontext(value)`.
+        The value must be safe for `workers` concurrent process() calls. For a sync stage it is used
+        from threads, so it must be thread-safe and must not be an asyncio-only object
+        (use httpx.Client, not httpx.AsyncClient).
+        """
+        return nullcontext(None)  # type: ignore[arg-type]
+
+
+class BaseSyncPipelineStage(BasePipelineStage[InT, OutT, ContextT]):
+    @classmethod
+    @abc.abstractmethod
+    def process(cls, item: InT, context: ContextT) -> Generator[OutT, None, None]:
+        """Process one item. Runs in a thread, should be thread-safe.
+        The generator is lazy: it must be consumed entirely inside the thread."""
+
+    async def run_one(self, item: InT, context: ContextT) -> tuple[int, float]:
+        def consume() -> tuple[int, float]:  # the whole generator runs in this one thread
+            emitted, blocked = 0, 0.0
+            for out in self.process(item, context):
+                emitted += 1
+                blocked += anyio.from_thread.run(self.send, out)
+            return emitted, blocked
+
+        return await anyio.to_thread.run_sync(consume)
+
+
+class BaseAsyncPipelineStage(BasePipelineStage[InT, OutT, ContextT]):
+    @classmethod
+    @abc.abstractmethod
+    def process(cls, item: InT, context: ContextT) -> AsyncGenerator[OutT, None]:
+        """Process one item. Runs on the event loop.
+        Implement as an async generator function (`async def` with `yield`), not as a coroutine returning a generator.
+        """
+
+    async def run_one(self, item: InT, context: ContextT) -> tuple[int, float]:
+        emitted, blocked = 0, 0.0
+        async for out in self.process(item, context):
+            emitted += 1
+            blocked += await self.send(out)
+        return emitted, blocked
+
+
+AnyStage = type[BasePipelineStage[Any, Any, Any]]

async_conveyor-0.0.1/src/async_conveyor/telemetry.py

@@ -0,0 +1,105 @@
+import logging
+from collections import defaultdict
+from collections.abc import AsyncIterator
+from contextlib import AbstractAsyncContextManager, asynccontextmanager
+from dataclasses import asdict, dataclass
+from typing import Protocol
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class ItemStats:
+    """Outcome of processing one item by one worker, plus a snapshot of the stage's queues at that moment."""
+
+    stage_name: str
+    worker: int
+    duration: float  # seconds spent in process(), including time blocked on output
+    emitted: int  # items yielded by process()
+    blocked: float  # seconds spent waiting for output queues (backpressure)
+    in_queue_size: int  # items waiting for this stage
+    out_queue_sizes: dict[str, int]  # child stage name -> items waiting in the queue into it
+    active_workers: int  # workers of this stage busy right now, this one included
+
+
+class Telemetry(Protocol):
+    """Extension point for metrics, tracing and error reporting. Implement and pass to the conveyor."""
+
+    def item_span(self, stage_name: str, worker: int) -> AbstractAsyncContextManager[None]:
+        """Wraps one process() call. Open a span or transaction here."""
+        ...
+
+    async def on_item_done(self, stats: ItemStats) -> None: ...
+
+    async def on_error(self, stage_name: str, worker: int, exc: BaseException, item: object) -> None:
+        """Called before the exception propagates and stops the conveyor."""
+        ...
+
+
+class LoggingTelemetry:
+    """Default telemetry: writes everything to the `async_conveyor.telemetry` logger with structured `extra` fields."""
+
+    def __init__(self, log: logging.Logger = logger) -> None:
+        self.log = log
+
+    @asynccontextmanager
+    async def item_span(self, stage_name: str, worker: int) -> AsyncIterator[None]:
+        yield
+
+    async def on_item_done(self, stats: ItemStats) -> None:
+        self.log.debug(
+            "item done: stage=%s worker=%d duration=%.3fs emitted=%d blocked=%.3fs in=%d out=%s workers=%d",
+            stats.stage_name,
+            stats.worker,
+            stats.duration,
+            stats.emitted,
+            stats.blocked,
+            stats.in_queue_size,
+            stats.out_queue_sizes,
+            stats.active_workers,
+            extra=asdict(stats),
+        )
+
+    async def on_error(self, stage_name: str, worker: int, exc: BaseException, item: object) -> None:
+        self.log.error(
+            "stage=%s worker=%d failed on item %r",
+            stage_name,
+            worker,
+            item,
+            exc_info=exc,
+            extra={"stage_name": stage_name, "worker": worker},
+        )
+
+
+@dataclass
+class StageSummary:
+    """Per-stage totals over a run."""
+
+    items: int = 0  # process() calls completed
+    emitted: int = 0  # items yielded downstream
+    max_active_workers: int = 0  # peak concurrency reached
+    blocked_seconds: float = 0.0  # total time spent waiting on full output queues
+    max_in_queue: int = 0  # peak backlog waiting for this stage
+
+
+class SummaryTelemetry:
+    """Aggregates per-stage counters in memory. Read `.stages` and `.errors` after run(): scripts, notebooks, tests."""
+
+    def __init__(self) -> None:
+        self.stages: dict[str, StageSummary] = defaultdict(StageSummary)
+        self.errors: list[tuple[str, int, BaseException, object]] = []  # (stage_name, worker, exc, item)
+
+    @asynccontextmanager
+    async def item_span(self, stage_name: str, worker: int) -> AsyncIterator[None]:
+        yield
+
+    async def on_item_done(self, stats: ItemStats) -> None:
+        s = self.stages[stats.stage_name]
+        s.items += 1
+        s.emitted += stats.emitted
+        s.blocked_seconds += stats.blocked
+        s.max_active_workers = max(s.max_active_workers, stats.active_workers)
+        s.max_in_queue = max(s.max_in_queue, stats.in_queue_size)
+
+    async def on_error(self, stage_name: str, worker: int, exc: BaseException, item: object) -> None:
+        self.errors.append((stage_name, worker, exc, item))

async_conveyor-0.0.1/src/async_conveyor/topology.py

@@ -0,0 +1,119 @@
+import inspect
+
+from .stage import AnyStage, BaseAsyncPipelineStage, BasePipelineStage, BaseSyncPipelineStage
+
+
+class TopologyError(ValueError):
+    """Raised when the conveyor topology is invalid."""
+
+
+class TopologyMixin:
+    """Queries and validation over `self.topology`: a dict of stage class -> list of child stage classes."""
+
+    topology: dict[AnyStage, list[AnyStage]]
+
+    @property
+    def stages(self) -> list[AnyStage]:
+        """All stages in the topology, in first-seen order: keys first, then their outputs."""
+        seen: dict[AnyStage, None] = {}
+        for parent, children in self.topology.items():
+            seen.setdefault(parent)
+            for child in children:
+                seen.setdefault(child)
+        return list(seen)
+
+    @property
+    def root(self) -> AnyStage:
+        """The single stage that is nobody's output. It is the source of the conveyor."""
+        (root,) = self._roots()
+        return root
+
+    def _roots(self) -> list[AnyStage]:
+        children = {child for outputs in self.topology.values() for child in outputs}
+        return [stage for stage in self.stages if stage not in children]
+
+    def validate_topology(self) -> None:
+        """Raise TopologyError if the topology cannot be run."""
+        stages = self.stages
+        if not stages:
+            raise TopologyError("topology is empty")
+        for stage in stages:
+            self._validate_stage(stage)
+        self._validate_unique_names(stages)
+        self._validate_no_fan_in()
+        self._validate_no_cycles(stages)
+        self._validate_root()
+
+    @staticmethod
+    def _validate_stage(stage: AnyStage) -> None:
+        if not isinstance(stage, type) or not issubclass(stage, BasePipelineStage):
+            raise TopologyError(f"{stage!r} is not a stage class")
+        if not issubclass(stage, BaseSyncPipelineStage | BaseAsyncPipelineStage):
+            raise TopologyError(f"{stage.__name__} must subclass BaseSyncPipelineStage or BaseAsyncPipelineStage")
+        if inspect.isabstract(stage):
+            missing = ", ".join(sorted(stage.__abstractmethods__))
+            raise TopologyError(f"{stage.__name__} is abstract, missing: {missing}")
+        if issubclass(stage, BaseAsyncPipelineStage) and not inspect.isasyncgenfunction(stage.process):
+            raise TopologyError(
+                f"{stage.__name__}.process must be an async generator function (`async def` with `yield`)"
+            )
+        if issubclass(stage, BaseSyncPipelineStage) and not inspect.isgeneratorfunction(stage.process):
+            raise TopologyError(f"{stage.__name__}.process must be a generator function (`def` with `yield`)")
+        for attr in ("name", "workers", "output_queue_size"):
+            if getattr(stage, attr) is NotImplemented:
+                raise TopologyError(f"{stage.__name__}.{attr} is not set")
+        if stage.workers < 1:
+            raise TopologyError(f"{stage.__name__}.workers must be >= 1")
+        if stage.output_queue_size < 0:
+            raise TopologyError(f"{stage.__name__}.output_queue_size must be >= 0")
+
+    @staticmethod
+    def _validate_unique_names(stages: list[AnyStage]) -> None:
+        by_name: dict[str, AnyStage] = {}
+        for stage in stages:
+            other = by_name.setdefault(stage.name, stage)
+            if other is not stage:
+                raise TopologyError(f"duplicate stage name {stage.name!r}: {other.__name__} and {stage.__name__}")
+
+    def _validate_no_fan_in(self) -> None:
+        parents: dict[AnyStage, AnyStage] = {}
+        for parent, children in self.topology.items():
+            for child in children:
+                other = parents.setdefault(child, parent)
+                if other is not parent:
+                    raise TopologyError(
+                        f"fan-in is not supported: {child.__name__} is fed by both "
+                        f"{other.__name__} and {parent.__name__}"
+                    )
+                if children.count(child) > 1:
+                    raise TopologyError(f"{child.__name__} is listed twice in outputs of {parent.__name__}")
+
+    def _validate_no_cycles(self, stages: list[AnyStage]) -> None:
+        done: set[AnyStage] = set()
+        path: list[AnyStage] = []
+
+        def visit(stage: AnyStage) -> None:
+            if stage in done:
+                return
+            if stage in path:
+                cycle = [*path[path.index(stage) :], stage]
+                raise TopologyError("cycle: " + " -> ".join(s.__name__ for s in cycle))
+            path.append(stage)
+            for child in self.topology.get(stage, []):
+                visit(child)
+            path.pop()
+            done.add(stage)
+
+        for stage in stages:
+            visit(stage)
+
+    def _validate_root(self) -> None:
+        roots = self._roots()
+        if not roots:
+            raise TopologyError("topology has no root stage")
+        if len(roots) > 1:
+            names = ", ".join(root.__name__ for root in roots)
+            raise TopologyError(f"topology must have exactly one root stage, found {len(roots)}: {names}")
+        (root,) = roots
+        if root.workers != 1:
+            raise TopologyError(f"root stage {root.__name__} must have workers == 1, it is called once")