diff --git a/TRACKER.md b/TRACKER.md index 67c4a4a..aa989d1 100644 --- a/TRACKER.md +++ b/TRACKER.md @@ -10,7 +10,7 @@ **Last updated:** 2026-05-24 **Current phase:** Phase 1 — Ingestion + Knowledge Store -**Next action:** Phase 1 Step 1.1d — Pipeline + Batcher primitives in `rag-core` (async DAG with bounded queues; DataLoader-pattern Batcher middleware coalescing concurrent SPI calls) +**Next action:** Phase 1 Step 1.1e — Cache SPI split (`EmbeddingCache` / `RetrievalCache` / `AnswerCache`) + hot-path discipline + async telemetry path with drop-on-overflow counter > **Refactor window (Steps 1.1a–1.1f):** Before resuming the connectors framework (1.2), we insert a six-step refactor that locks in architecture + optimization decisions which are very expensive to retrofit later (PolicyEngine PDP, RequestContext-threaded SPIs, split Retrieval/Index backends, bulk + streaming + ID-only methods, Pipeline + Batcher primitives, three-way cache split, hot-path discipline). See [docs/adr/ADR-0005…0009] and [docs/architecture/policy-engine.md], [request-context.md], [caching.md], [performance.md]. @@ -71,7 +71,7 @@ | 1.1a | Core type & SPI refactor | ✅ | `build/phase-1/step-1.1a-core-type-spi-refactor` | [#44](https://github.com/officialCodeWork/AgentContextOS/pull/44) | `RequestContext` frozen model threaded through every SPI; `tenant_id` + `acl_labels` typed required on `Chunk`/`Embedding` (not metadata dict); `trust_level` on `Chunk` for prompt-injection defense; `dtype` on `Embedding` (float32/int8/binary); `BlobRef` for lazy chunk text; `QueryPlan` + `ChunkRef` + `Cost` + `PlanNode` types; typed `StageEvent`. `tests/contract/spi_signature.py` linter (RequestContext-first); rag-backends (`PgVectorStore`, `QdrantVectorStore`, `RedisCache`, `S3Storage`, `LocalFileStorage`) migrated; conformance + integration tests updated; `Budget.spend()` for agent-loop sub-turn budgets; schemas regenerated. ADR-0005 / ADR-0007 / ADR-0008 / ADR-0009 referenced. | | 1.1b | SPI split — Retrieval/Index, bulk + streaming + ID-only | ✅ | `build/phase-1/step-1.1b-spi-split-retrieval-index` | [#45](https://github.com/officialCodeWork/AgentContextOS/pull/45) | Split `VectorStore`/`KeywordStore`/`GraphStore` into `*RetrievalBackend` (read) + `*IndexBackend` (write) composite ABCs. `retrieve_ids` returns `list[ChunkRef]`; `hydrate` lives on the retrieval side (keyword full-Chunk; vector pass-through). Bulk: `bulk_index`/`bulk_delete` (+ graph bulk node/edge variants); streaming: `stream_index` async-iterator default that batches into `bulk_index`. `Embedder` split into single `embed` + canonical `bulk_embed`. New `IndexHint` + `WriteVolume` types passed to writes (per ADR-0009). Noop impls, `PgVectorStore`, `QdrantVectorStore` migrated; conformance + integration tests updated; `spi_signature.py` extended to enforce the split. Schemas regenerated (`IndexHint.json`). | | 1.1c | PolicyEngine package | ✅ | `build/phase-1/step-1.1c-policy-engine-package` | [#46](https://github.com/officialCodeWork/AgentContextOS/pull/46) | New `packages/policy/` (`rag-policy` v0.1.0): `PolicyEngine` SPI + `NoopPolicyEngine` (always-ALLOW with tenant-scoped `filter_pushdown`); `PolicyDecision` enum (read_chunk / ingest_doc / egress_text / quota_check / rate_limit / execute_plan); `PolicyResult` (allow/deny/transform) with predicate helpers; `QuotaSubject` / `RateLimitSubject`; `FilterExpr` mini-language (Eq / AnyIn / And / Or / Not / TrueExpr) returned by `filter_pushdown`; `PolicyWriter` facade mirroring `AuditWriter` and emitting `policy.decision` structured logs via `rag-observability`. Coverage linter `tests/policy/coverage.py` greps for governance-relevant SPI calls without adjacent `PolicyEngine`/`PolicyWriter` consultation, with file allowlist that consumers shrink as they wire the PDP in. Workspace + pytest pythonpath updated. 20 conformance tests added. ADR-0005 finalized. | -| 1.1d | Pipeline + Batcher primitives | ⏳ | 1.1a | — | `Pipeline` primitive in `rag-core`: async DAG with bounded queues, per-stage worker counts, backpressure (used by Step 1.10 write path). `Batcher[Req, Resp]` middleware (DataLoader pattern) coalescing concurrent SPI calls into batched provider calls; sits under Embedder/Reranker SPIs. | +| 1.1d | Pipeline + Batcher primitives | 🚧 | `build/phase-1/step-1.1d-pipeline-batcher-primitives` | — | `Pipeline` primitive in `rag-core` (`rag_core.pipeline`): async DAG with bounded queues, per-stage worker counts, backpressure, joiner-based shutdown handling per-stage worker-count differences, `on_error="fail"`/`"skip"` policies, optional `error_handler`. `Batcher[Req, Resp]` middleware (`rag_core.batcher`, DataLoader pattern) coalescing concurrent `load(req)` calls into batched `batch_fn(reqs)` invocations; size + time triggers (`max_batch_size`, `max_wait_ms`); per-request `asyncio.Future` for isolated error propagation; `flush()` for deterministic shutdown. PEP-695-typed generics. 22 unit tests (8 Batcher + 14 Pipeline). `docs/architecture/pipeline-batcher.md` + `rag-core` reference + `performance.md` backlinks. `rag-core` 0.5.0 → 0.6.0. | | 1.1e | Cache SPI split + perf discipline + async telemetry | ⏳ | 1.1a | — | Split `Cache` into `EmbeddingCache` (key model_id+text_hash), `RetrievalCache` (plan_hash+corpus_version), `AnswerCache` (plan_hash+corpus_version+policy_version). Each has distinct invalidation. Hot-path convention doc (Pydantic at SPI boundary, `model_construct`/msgspec inside). Async telemetry path with bounded buffer + drop-on-overflow counter. | | 1.1f | ADRs 0005–0009 + reviewer checklist | ⏳ | 1.1a–1.1e | — | ADR-0005 (PolicyEngine PDP), ADR-0006 (two-stage reranker default), ADR-0007 (tiered storage with BlobRef), ADR-0008 (cost-aware planner replacing reactive fallback), ADR-0009 (vector index strategy + quantization). Reviewer checklist in `docs/architecture/performance.md`. | | 1.2 | Connectors framework | ⏳ | 1.1a–1.1f | — | `Connector` SPI implementation (now receives `RequestContext` and `ConnectorState` watermark); built-in: filesystem, S3, GCS; crawler base class | @@ -217,6 +217,7 @@ | [#46](https://github.com/officialCodeWork/AgentContextOS/pull/46) | feat(policy): rag-policy package — PolicyEngine PDP + coverage linter (Step 1.1c) | `build/phase-1/step-1.1c-policy-engine-package` | ✅ Merged | 2026-05-24 | | [#47](https://github.com/officialCodeWork/AgentContextOS/pull/47) | chore(schemas): track generated JSON Schemas + add drift CI gate | `chore/track-generated-schemas-drift-gate` | ✅ Merged | 2026-05-24 | | [#48](https://github.com/officialCodeWork/AgentContextOS/pull/48) | chore(ci): lint scripts/ + ignore S603/S607 in scripts/** | `chore/lint-ignore-subprocess-in-scripts` | ✅ Merged | 2026-05-24 | +| [#49](https://github.com/officialCodeWork/AgentContextOS/pull/49) | chore(tracker): sync PR links for steps 1.1a/1.1c + log #46–#48 | `chore/tracker-sync-pr44-48` | ✅ Merged | 2026-05-24 | --- diff --git a/docs/README.md b/docs/README.md index 4d73d51..3e5df2b 100644 --- a/docs/README.md +++ b/docs/README.md @@ -10,6 +10,7 @@ | [policy-engine.md](architecture/policy-engine.md) | `PolicyEngine` (PDP) — single decision point for ACL, PII, quotas, redaction; replaces scattered governance checks | | [caching.md](architecture/caching.md) | Three-cache split: `EmbeddingCache`, `RetrievalCache`, `AnswerCache` — distinct invalidation rules | | [performance.md](architecture/performance.md) | Hot-path discipline, per-SPI p99 budgets, async telemetry, reviewer checklist | +| [pipeline-batcher.md](architecture/pipeline-batcher.md) | `Pipeline` (async DAG, bounded queues, per-stage workers) + `Batcher` (DataLoader-pattern coalescing) primitives — Step 1.1d | | [eval-skeleton.md](architecture/eval-skeleton.md) | Eval framework architecture: golden-set schema, metric functions, RAGAS spike, `ragctl eval` CLI, extension points | | [iac.md](architecture/iac.md) | IaC overview: Terraform module design, Helm chart structure, dev/prod environments, extension points | | [storage-backends.md](architecture/storage-backends.md) | Storage backend architecture: PgVector, Qdrant, Redis, S3/MinIO, tenant isolation, integration test strategy | @@ -20,7 +21,7 @@ |------|-------------| | [ragctl.md](reference/ragctl.md) | Full `ragctl` command reference — public usage, internals, extension points | | [backends.md](reference/backends.md) | `rag-backends` reference — PgVectorStore, QdrantVectorStore, RedisCache, S3Storage, LocalFileStorage | -| [rag-core.md](reference/rag-core.md) | `rag-core` type surface — `RequestContext`, `Budget`, `BlobRef`, `QueryPlan`, `ChunkRef`, `StageEvent` | +| [rag-core.md](reference/rag-core.md) | `rag-core` type surface — `RequestContext`, `Budget`, `BlobRef`, `QueryPlan`, `ChunkRef`, `StageEvent`, `Pipeline`, `Batcher` | | [rag-policy.md](reference/rag-policy.md) | `rag-policy` reference — `PolicyEngine`, `PolicyWriter`, `PolicyDecision`, `PolicyResult`, `FilterExpr` | ## guides/ diff --git a/docs/architecture/performance.md b/docs/architecture/performance.md index 0f085a7..f3d0aa8 100644 --- a/docs/architecture/performance.md +++ b/docs/architecture/performance.md @@ -107,5 +107,6 @@ Apply to every PR that touches an SPI, a hot path, or an SPI consumer: - [request-context.md](request-context.md) — the per-request envelope. - [policy-engine.md](policy-engine.md) — governance PDP. - [caching.md](caching.md) — three caches. +- [pipeline-batcher.md](pipeline-batcher.md) — `Pipeline` and `Batcher` primitives (bounded queues + DataLoader-pattern coalescing). - [ADR-0006](../adr/ADR-0006-two-stage-rerank.md), [ADR-0008](../adr/ADR-0008-cost-aware-planner.md), [ADR-0009](../adr/ADR-0009-vector-index-strategy.md) — performance-relevant decisions. - TRACKER.md Steps 1.1a–1.1f (refactor window), 4.6 (latency tuning), 7.1 (load testing). diff --git a/docs/architecture/pipeline-batcher.md b/docs/architecture/pipeline-batcher.md new file mode 100644 index 0000000..5c3ca76 --- /dev/null +++ b/docs/architecture/pipeline-batcher.md @@ -0,0 +1,158 @@ +# Pipeline + Batcher primitives + +## Overview + +`rag-core` ships two cross-cutting concurrency primitives used throughout +the ingest and retrieval paths: + +- **`Pipeline`** — a linear async DAG with bounded queues and per-stage + worker counts. Lives in `rag_core.pipeline`. Drives the Step 1.10 write + path (connector → parse → chunk → enrich → PII → embed → store) and any + other multi-stage processing. +- **`Batcher[Req, Resp]`** — a DataLoader-pattern middleware that coalesces + concurrent `load(req)` calls into a single `batch_fn(reqs)` invocation. + Lives in `rag_core.batcher`. Sits under every Embedder / Reranker / LLM + adapter that talks to a billed external API, so that N in-flight concurrent + calls collapse to one provider call. + +They are introduced in Step 1.1d and locked in ahead of Step 1.2 because +both shapes show up everywhere — adding them once, correctly, with the +right backpressure and error semantics is far cheaper than retrofitting +ten consumers later. + +## Usage + +### `Pipeline` + +Each stage is an `async fn(item) -> result`. The result type is overloaded: + +| Return value | Semantics | +|-------------------|----------------------------------------------------| +| a single object | forwarded to the next stage as one item | +| `list[...]` | each element forwarded separately (fan-out) | +| `None` | dropped (filter) | + +Build with `add_stage`, then drive with `run(source)`: + +```python +from rag_core import Pipeline + +async def parse(doc): return await parse_async(doc) +async def chunk(parsed): return chunker.split(parsed) # → list[Chunk] +async def embed(chunk): return await embedder.embed(ctx, chunk.text, chunk.id) + +pipe = ( + Pipeline() + .add_stage("parse", parse, workers=2) + .add_stage("chunk", chunk, workers=2) + .add_stage("embed", embed, workers=8) # downstream wider than upstream +) + +async for embedding in pipe.run(source_documents()): + await vector_index.append(ctx, embedding) +``` + +Bounded-queue backpressure is automatic — if the consumer of `pipe.run` +stops pulling, the embed workers fill their input queue, the chunk workers +fill theirs, and so on back to the source iterator. + +#### Error policy + +- `on_error="fail"` (default) — the first stage exception aborts the whole + pipeline. The exception surfaces at the consumer as `PipelineError` with + the original error in `__cause__`. Matches the fail-fast convention. +- `on_error="skip"` — the failing item is dropped; the pipeline keeps + draining. Useful for best-effort backfills. + +In either mode, an optional `error_handler(stage_name, item, exc)` +callback fires for every failure — wire it up to structured logging or +metrics. Error handlers must not raise; exceptions from the handler are +swallowed to keep the pipeline draining. + +### `Batcher` + +```python +from rag_core import Batcher + +async def call_provider(reqs: list[str]) -> list[list[float]]: + return await openai.embeddings.create(input=reqs, model="...") + +embed_batcher: Batcher[str, list[float]] = Batcher( + call_provider, + max_batch_size=64, + max_wait_ms=5, +) + +# Inside an Embedder.embed implementation: +vec = await embed_batcher.load(text) +``` + +A pending batch flushes when *either* `max_batch_size` is reached *or* +`max_wait_ms` elapses since the first request landed. Default window is +5 ms — long enough to coalesce a burst, short enough not to add visible +latency to single requests. + +`flush()` forces an immediate flush of whatever is pending. Use at +shutdown or in deterministic tests. + +## Internals + +### Pipeline shutdown — joiners and sentinels + +Each stage transition is governed by a *joiner* task that awaits all of +the upstream stage's workers and then injects exactly `next_stage.workers` +`_DONE` sentinels into the next queue. This lets per-stage worker counts +differ without losing or leaking sentinels. The feeder injects sentinels +for the first stage. + +When a stage worker hits an exception under `on_error="fail"`: + +1. The worker calls `error_handler` (if set). +2. It pushes a `_Failure(PipelineError(...))` marker into its output queue + and sets the run's abort event. +3. The marker propagates through downstream stages (workers forward it on + sight and exit). At the output queue, the consumer pulls the marker + and `raise`s. +4. Other workers on the failing stage drain their queues without invoking + the stage function (the abort event short-circuits them) until they + see a `_DONE` sentinel and exit. +5. Joiners complete normally; the consumer's `try/finally` in `run` + cancels any tasks that are still alive (e.g. the feeder waiting on a + full queue). + +This shape lets `aclose()` on the async iterator returned by `run` perform +clean shutdown if the consumer breaks early — no stranded tasks, no +leaked workers. + +### Batcher — single-flight flush task + +A pending batch holds at most one timer task. When a request arrives: + +- If the batch is empty, the request kicks off the timer. +- If the batch reaches `max_batch_size`, the current pending slot is + rotated and the batch is flushed synchronously by the caller. The + outstanding timer is cancelled. + +`asyncio.Future` per request decouples error handling: an exception from +`batch_fn` is set on every future in the batch, never on a single one. + +## Extension points + +- **Pipeline observation.** Wire an `error_handler` to emit structured + logs or `StageEvent` records; the gateway publishes `StageEvent`s over + SSE per `docs/architecture/RAG-Platform-HLD.md`. +- **Pipeline shape.** Linear-only by design. A second primitive + (`Workflow`) can be added later if a real consumer needs a non-linear + DAG; do not generalize this one preemptively. +- **Batcher under SPIs.** Every adapter for a billed external API should + wrap its provider call in a `Batcher`. The reviewer checklist in + [performance.md](performance.md) flags missing batchers in code review. + +## Related + +- [performance.md](performance.md) — hot-path discipline, p99 budgets, + reviewer checklist (where the Batcher rule is enforced). +- [request-context.md](request-context.md) — the envelope every stage + receives via closure over `ctx`. +- TRACKER.md Step 1.1d (this work) and Step 1.10 (Pipeline's first big + consumer — the ingest write path). diff --git a/docs/reference/rag-core.md b/docs/reference/rag-core.md index 1ddc1e2..45ded06 100644 --- a/docs/reference/rag-core.md +++ b/docs/reference/rag-core.md @@ -221,6 +221,22 @@ Removing or renaming a field is a major change and requires an ADR. --- +## Concurrency primitives + +`rag-core` also ships two cross-cutting concurrency primitives used by +ingest and retrieval consumers (added in Step 1.1d): + +| Name | Module | Purpose | +|---|---|---| +| `Pipeline`, `Stage` | `rag_core.pipeline` | Async DAG with bounded queues + per-stage worker counts; drives the ingest write path | +| `Batcher[Req, Resp]` | `rag_core.batcher` | DataLoader-pattern middleware; coalesces concurrent SPI calls into batched provider calls | + +Full design notes, usage, and error semantics live in +[docs/architecture/pipeline-batcher.md](../architecture/pipeline-batcher.md). +The reviewer checklist in +[docs/architecture/performance.md](../architecture/performance.md) calls +out where each must be used. + ## Related - [docs/architecture/request-context.md](../architecture/request-context.md) — design rationale. diff --git a/packages/core/pyproject.toml b/packages/core/pyproject.toml index 7e75c5f..38eb17c 100644 --- a/packages/core/pyproject.toml +++ b/packages/core/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "rag-core" -version = "0.3.0" +version = "0.6.0" description = "AgentContextOS — core domain types, error hierarchy, and plugin SPIs" readme = "README.md" requires-python = ">=3.12" diff --git a/packages/core/src/rag_core/__init__.py b/packages/core/src/rag_core/__init__.py index 1c0d0fc..0aafb3c 100644 --- a/packages/core/src/rag_core/__init__.py +++ b/packages/core/src/rag_core/__init__.py @@ -5,6 +5,7 @@ Eval domain types added in Step 0.8. """ +from rag_core.batcher import Batcher from rag_core.errors import ( ACLDeniedError, AuthError, @@ -46,6 +47,7 @@ ) from rag_core.logging import configure_logging, get_logger, reset_logging, set_log_context from rag_core.metrics import SpiMetrics +from rag_core.pipeline import Pipeline, Stage from rag_core.spi import ( LLM, OCR, @@ -122,7 +124,7 @@ WriteVolume, ) -__version__ = "0.5.0" +__version__ = "0.6.0" __all__ = [ # eval @@ -208,6 +210,10 @@ "SchemaValidationError", "TenantIsolationError", "TimeoutError", + # pipeline + batcher primitives + "Batcher", + "Pipeline", + "Stage", # telemetry + metrics "SpiMetrics", "configure_telemetry", diff --git a/packages/core/src/rag_core/batcher.py b/packages/core/src/rag_core/batcher.py new file mode 100644 index 0000000..2476f4b --- /dev/null +++ b/packages/core/src/rag_core/batcher.py @@ -0,0 +1,168 @@ +"""Batcher — DataLoader-pattern middleware for coalescing concurrent SPI calls. + +A ``Batcher[Req, Resp]`` accepts individual ``load(req)`` calls and groups +those issued within a short window (or until a max batch size is reached) +into a single ``batch_fn(reqs)`` invocation against the underlying provider. + +This is the substrate for the convention captured in +``docs/architecture/performance.md``: every Embedder / Reranker / LLM +adapter that talks to a billed external API sits under a Batcher so that +N concurrent in-flight requests collapse to one provider call. + +Design notes +------------ + +* **Per-batch isolation.** Each batch has its own ``asyncio.Future`` per + request. An error returned for one request never poisons the others; + set ``return_exceptions=True`` semantics if you need partial success at + the caller layer. +* **Trigger.** A pending batch is flushed when *either* ``max_batch_size`` + is reached *or* ``max_wait_ms`` elapses since the first request landed + in the current batch. Use small windows (1–10 ms) — the goal is to + coalesce truly concurrent fan-out, not to wait for trickle traffic. +* **Single-flight wait task.** Exactly one background task is alive per + pending batch; it cancels itself if the size trigger fires first. +* **Backpressure.** Calls are not rate-limited here — that is the + provider adapter's responsibility (or a wrapping ``Pipeline`` stage). + The Batcher only batches. +""" + +from __future__ import annotations + +import asyncio +from collections.abc import Awaitable, Callable +from dataclasses import dataclass, field + +type BatchFn[Req, Resp] = Callable[[list[Req]], Awaitable[list[Resp]]] + + +@dataclass +class _PendingItem[Req, Resp]: + req: Req + future: asyncio.Future[Resp] + + +@dataclass +class _PendingBatch[Req, Resp]: + items: list[_PendingItem[Req, Resp]] = field(default_factory=list) + flush_task: asyncio.Task[None] | None = None + flushing: bool = False + + +class Batcher[Req, Resp]: + """Coalesce concurrent ``load`` calls into batched provider invocations. + + Parameters + ---------- + batch_fn: + Async callable invoked once per flushed batch. Must return one + response per request, in the same order. + max_batch_size: + Maximum number of requests in a single ``batch_fn`` call. When + reached, the batch flushes immediately without waiting for the + time window. + max_wait_ms: + Maximum time (milliseconds) to wait for additional requests after + the first one lands in the pending batch. Defaults to ``5`` — + long enough to coalesce a burst, short enough not to add visible + latency. + """ + + def __init__( + self, + batch_fn: BatchFn[Req, Resp], + *, + max_batch_size: int = 64, + max_wait_ms: float = 5.0, + ) -> None: + if max_batch_size < 1: + raise ValueError("max_batch_size must be >= 1") + if max_wait_ms < 0: + raise ValueError("max_wait_ms must be >= 0") + + self._batch_fn = batch_fn + self._max_batch_size = max_batch_size + self._max_wait_s = max_wait_ms / 1000.0 + self._pending: _PendingBatch[Req, Resp] = _PendingBatch() + self._lock = asyncio.Lock() + + async def load(self, req: Req) -> Resp: + """Submit ``req`` and wait for the response from the next batch.""" + future: asyncio.Future[Resp] = asyncio.get_running_loop().create_future() + await self._enqueue(req, future) + return await future + + async def _enqueue(self, req: Req, future: asyncio.Future[Resp]) -> None: + flush_now: _PendingBatch[Req, Resp] | None = None + async with self._lock: + batch = self._pending + batch.items.append(_PendingItem(req=req, future=future)) + + if len(batch.items) >= self._max_batch_size: + # Size trigger — rotate the pending slot and flush this batch + # synchronously below, outside the lock. + self._pending = _PendingBatch() + if batch.flush_task is not None: + batch.flush_task.cancel() + batch.flushing = True + flush_now = batch + elif batch.flush_task is None: + batch.flush_task = asyncio.create_task(self._wait_then_flush(batch)) + + if flush_now is not None: + await self._run_batch(flush_now) + + async def _wait_then_flush(self, batch: _PendingBatch[Req, Resp]) -> None: + try: + await asyncio.sleep(self._max_wait_s) + except asyncio.CancelledError: + return + + async with self._lock: + if batch.flushing or not batch.items: + return + # Rotate the pending slot only if it still points at this batch + # (the size trigger may have already swapped it). + if self._pending is batch: + self._pending = _PendingBatch() + batch.flushing = True + + await self._run_batch(batch) + + async def _run_batch(self, batch: _PendingBatch[Req, Resp]) -> None: + reqs = [item.req for item in batch.items] + try: + results = await self._batch_fn(reqs) + except BaseException as exc: # noqa: BLE001 — propagate to every waiter + for item in batch.items: + if not item.future.done(): + item.future.set_exception(exc) + return + + if len(results) != len(reqs): + err = ValueError(f"batch_fn returned {len(results)} responses for {len(reqs)} requests") + for item in batch.items: + if not item.future.done(): + item.future.set_exception(err) + return + + for item, result in zip(batch.items, results, strict=True): + if not item.future.done(): + item.future.set_result(result) + + async def flush(self) -> None: + """Force any pending batch to flush immediately. + + Useful in tests and at shutdown. Safe to call when nothing is + pending — returns instantly. + """ + async with self._lock: + batch = self._pending + if not batch.items or batch.flushing: + return + self._pending = _PendingBatch() + if batch.flush_task is not None: + batch.flush_task.cancel() + batch.flushing = True + + await self._run_batch(batch) diff --git a/packages/core/src/rag_core/pipeline.py b/packages/core/src/rag_core/pipeline.py new file mode 100644 index 0000000..532e142 --- /dev/null +++ b/packages/core/src/rag_core/pipeline.py @@ -0,0 +1,359 @@ +"""Pipeline — async DAG with bounded queues and per-stage worker counts. + +The ``Pipeline`` primitive is the substrate for the ingest write path +(Step 1.10) and any other multi-stage processing the system runs end-to-end. +Each stage is an async function transforming one input item into zero or more +output items; stages are connected by bounded ``asyncio.Queue``s, which +enforce backpressure: a slow downstream consumer pauses every upstream +producer. + +The shape is intentionally linear (not a free-form DAG) — every real ingest +pipeline we have planned is a chain (connector → parse → chunk → enrich → +PII → embed → store). A linear chain composes cleanly with the per-stage +worker count to express the only fan-out the platform needs at this layer +(per-stage concurrency, e.g., 8 embedder workers downstream of a single +parse worker). Free-form DAG support can be added later if a real consumer +needs it; doing so now would be speculative. + +Design notes +------------ + +* **Bounded queues.** Every stage's input queue has a finite ``maxsize``. + ``asyncio.Queue()`` defaults to unbounded — see + ``docs/architecture/performance.md`` for why that is banned. +* **Per-stage workers.** Each stage runs ``workers`` coroutines pulling from + its input queue concurrently. Tune downstream stages (e.g. embedding) + higher than upstream stages. +* **Fan-out via list returns.** A stage that returns ``list[T]`` enqueues + each element separately; a stage that returns ``None`` drops the item. + This expresses chunking (one document → many chunks) without a separate + fan-out primitive. +* **Backpressure on push.** The feeder consumes the async source with + ``await queue.put(...)``; when downstream backs up, the source pauses. +* **Errors.** A stage exception is delivered to ``error_handler``; the + default policy is ``"fail"`` (abort the pipeline and re-raise at the + consumer) — matching the fail-fast convention used elsewhere in the + codebase. ``"skip"`` drops the failing item and continues. +* **Shutdown.** Each stage transition is governed by a *joiner* task that + waits for all upstream workers to complete, then injects exactly the + right number of sentinels into the next queue. This lets stage worker + counts differ without losing or leaking sentinels. +""" + +from __future__ import annotations + +import asyncio +import contextlib +from collections.abc import AsyncIterator, Awaitable, Callable, Iterable +from dataclasses import dataclass, field +from typing import Any, Literal + +from rag_core.errors import PipelineError + +type StageFn = Callable[[Any], Awaitable[Any | list[Any] | None]] +"""Signature for a stage function: takes one item, returns one, many, or none. + +Typed as ``Any`` because a pipeline's stage chain has heterogeneous types +that the type system cannot track across the dynamic ``add_stage`` builder. +Call sites should keep their per-stage callables narrowly typed. +""" + +type ErrorPolicy = Literal["fail", "skip"] +type ErrorHandler = Callable[[str, Any, BaseException], Awaitable[None] | None] + + +class _Done: + """Sentinel pushed through queues to signal end-of-stream.""" + + __slots__ = () + + +_DONE = _Done() + + +@dataclass +class _Failure: + """Marker pushed into the output queue when the pipeline aborts.""" + + error: BaseException + + +@dataclass +class Stage: + """One stage in the pipeline. + + Attributes + ---------- + name: + Human-readable identifier used in error messages and telemetry. + fn: + Async callable applied to each input item. Returns either: + * a single item (forwarded to the next stage), + * a ``list`` (each element forwarded separately — fan-out), or + * ``None`` (drop the item). + workers: + Number of concurrent coroutines running this stage. Defaults to 1. + queue_size: + Max number of items the stage's *input* queue holds before producers + block. Defaults to ``workers * 2`` (minimum 2) if left as 0. + """ + + name: str + fn: StageFn + workers: int = 1 + queue_size: int = 0 + + def __post_init__(self) -> None: + if self.workers < 1: + raise ValueError(f"stage {self.name!r}: workers must be >= 1") + if self.queue_size < 0: + raise ValueError(f"stage {self.name!r}: queue_size must be >= 0") + if self.queue_size == 0: + self.queue_size = max(self.workers * 2, 2) + + +class Pipeline[T, U]: + """Composable async DAG with bounded queues and per-stage workers. + + Build with ``add_stage`` (or pass ``stages=`` to the constructor) and + drive with ``run(source)``. The output of the final stage is yielded + from ``run`` as an async iterator. + + Example + ------- + >>> async def parse(doc): return parsed + >>> async def chunk(parsed): return [c1, c2, c3] # fan-out + >>> async def embed(chunk): return embedded + >>> + >>> pipe = ( + ... Pipeline() + ... .add_stage("parse", parse, workers=2) + ... .add_stage("chunk", chunk, workers=2) + ... .add_stage("embed", embed, workers=8) + ... ) + >>> async for emb in pipe.run(source_docs()): + ... await store.put(emb) + """ + + def __init__( + self, + stages: Iterable[Stage] | None = None, + *, + on_error: ErrorPolicy = "fail", + error_handler: ErrorHandler | None = None, + ) -> None: + self._stages: list[Stage] = list(stages) if stages else [] + self._on_error: ErrorPolicy = on_error + self._error_handler: ErrorHandler | None = error_handler + + def add_stage( + self, + name: str, + fn: StageFn, + *, + workers: int = 1, + queue_size: int = 0, + ) -> Pipeline[T, U]: + """Append a stage. Returns ``self`` for chaining.""" + if any(s.name == name for s in self._stages): + raise ValueError(f"duplicate stage name: {name!r}") + self._stages.append(Stage(name=name, fn=fn, workers=workers, queue_size=queue_size)) + return self + + @property + def stages(self) -> tuple[Stage, ...]: + return tuple(self._stages) + + async def run(self, source: AsyncIterator[Any]) -> AsyncIterator[Any]: + """Drive the pipeline from ``source``; yield outputs of the final stage. + + Bounded-queue backpressure means that if a downstream consumer of + ``run``'s output stops pulling, every upstream stage (and the source) + eventually pauses. + """ + if not self._stages: + async for item in source: + yield item + return + + run = _PipelineRun( + stages=self._stages, + on_error=self._on_error, + error_handler=self._error_handler, + ) + try: + async for item in run.execute(source): + yield item + finally: + await run.cancel_all() + + +@dataclass +class _PipelineRun: + stages: list[Stage] + on_error: ErrorPolicy + error_handler: ErrorHandler | None + _all_tasks: list[asyncio.Task[Any]] = field(default_factory=list) + _abort_event: asyncio.Event = field(default_factory=asyncio.Event) + + async def execute(self, source: AsyncIterator[Any]) -> AsyncIterator[Any]: + queues: list[asyncio.Queue[Any]] = [ + asyncio.Queue(maxsize=stage.queue_size) for stage in self.stages + ] + output: asyncio.Queue[Any] = asyncio.Queue(maxsize=max(self.stages[-1].workers * 2, 2)) + + worker_tasks_by_stage: list[list[asyncio.Task[None]]] = [] + for idx, stage in enumerate(self.stages): + next_q: asyncio.Queue[Any] = queues[idx + 1] if idx + 1 < len(self.stages) else output + tasks = [ + asyncio.create_task( + self._worker(stage=stage, in_q=queues[idx], out_q=next_q), + name=f"pipeline.{stage.name}.{w}", + ) + for w in range(stage.workers) + ] + worker_tasks_by_stage.append(tasks) + self._all_tasks.extend(tasks) + + feeder = asyncio.create_task( + self._feed(source, queues[0], self.stages[0].workers), + name="pipeline.feeder", + ) + self._all_tasks.append(feeder) + + # One joiner per stage: waits for that stage's workers to finish, + # then pushes exactly the right number of sentinels downstream. + for idx, stage in enumerate(self.stages): + next_q = queues[idx + 1] if idx + 1 < len(self.stages) else output + next_count = self.stages[idx + 1].workers if idx + 1 < len(self.stages) else 1 + joiner = asyncio.create_task( + self._join(worker_tasks_by_stage[idx], next_q, next_count), + name=f"pipeline.join.{stage.name}", + ) + self._all_tasks.append(joiner) + + while True: + item = await output.get() + if isinstance(item, _Done): + return + if isinstance(item, _Failure): + raise item.error + yield item + + async def _feed( + self, + source: AsyncIterator[Any], + first_queue: asyncio.Queue[Any], + worker_count: int, + ) -> None: + # CancelledError MUST propagate — do not catch BaseException here. + try: + async for item in source: + if self._abort_event.is_set(): + return + await first_queue.put(item) + except Exception as exc: + self._abort_event.set() + wrapped = PipelineError(f"source failed: {exc}") + wrapped.__cause__ = exc + # Push the failure straight through so the consumer raises promptly + # even if no worker has produced anything yet. Use put_nowait so + # we don't block if the queue happens to be full at this moment; + # the next consumer of the queue will see it. + with contextlib.suppress(asyncio.QueueFull): + first_queue.put_nowait(_Failure(wrapped)) + finally: + # Send one sentinel per worker so they exit cleanly. If we are + # being cancelled, put_nowait avoids blocking; the workers will + # be cancelled directly by ``cancel_all`` anyway. + for _ in range(worker_count): + try: + first_queue.put_nowait(_DONE) + except asyncio.QueueFull: + # Make room by awaiting, but if we're already cancelled, + # bail out and let cancel_all clean up. + if self._abort_event.is_set(): + break + try: + await first_queue.put(_DONE) + except asyncio.CancelledError: + break + + async def _worker( + self, + *, + stage: Stage, + in_q: asyncio.Queue[Any], + out_q: asyncio.Queue[Any], + ) -> None: + while True: + item = await in_q.get() + if isinstance(item, _Done): + return + if isinstance(item, _Failure): + # Forward and exit so downstream sees the failure quickly. + await out_q.put(item) + self._abort_event.set() + return + if self._abort_event.is_set(): + continue # drain remaining items quickly until sentinel + + try: + result = await stage.fn(item) + except Exception as exc: + # CancelledError is intentionally not caught — it must + # propagate so cancellation tears the pipeline down cleanly. + await self._invoke_error_handler(stage.name, item, exc) + if self.on_error == "fail": + wrapped = PipelineError( + f"stage {stage.name!r} failed: {exc}", + stage=stage.name, + ) + wrapped.__cause__ = exc + await out_q.put(_Failure(wrapped)) + self._abort_event.set() + return + continue + + if result is None: + continue + if isinstance(result, list): + for sub in result: + await out_q.put(sub) + else: + await out_q.put(result) + + async def _join( + self, + worker_tasks: list[asyncio.Task[None]], + next_q: asyncio.Queue[Any], + next_count: int, + ) -> None: + # Wait for every worker on this stage to finish. Exceptions are + # already surfaced by the workers via _Failure markers; we just + # need to know everyone has stopped so we can close the next stage. + await asyncio.gather(*worker_tasks, return_exceptions=True) + for _ in range(next_count): + await next_q.put(_DONE) + + async def _invoke_error_handler(self, stage_name: str, item: Any, exc: BaseException) -> None: + if self.error_handler is None: + return + try: + maybe = self.error_handler(stage_name, item, exc) + if asyncio.iscoroutine(maybe): + await maybe + except Exception: # noqa: S110, BLE001 + # Error handler must not crash the pipeline. CancelledError is + # intentionally allowed to propagate (only Exception is caught). + pass + + async def cancel_all(self) -> None: + self._abort_event.set() + for t in self._all_tasks: + if not t.done(): + t.cancel() + # gather with return_exceptions absorbs each child's CancelledError / + # Exception so we never raise from cleanup. + if self._all_tasks: + await asyncio.gather(*self._all_tasks, return_exceptions=True) diff --git a/packages/core/tests/test_batcher.py b/packages/core/tests/test_batcher.py new file mode 100644 index 0000000..caaa7e1 --- /dev/null +++ b/packages/core/tests/test_batcher.py @@ -0,0 +1,132 @@ +"""Unit tests for ``rag_core.batcher.Batcher``.""" + +from __future__ import annotations + +import asyncio + +import pytest +from rag_core.batcher import Batcher + + +async def test_coalesces_concurrent_loads_into_one_batch() -> None: + calls: list[list[int]] = [] + + async def batch_fn(reqs: list[int]) -> list[int]: + calls.append(list(reqs)) + return [r * 2 for r in reqs] + + batcher: Batcher[int, int] = Batcher(batch_fn, max_batch_size=10, max_wait_ms=20) + + results = await asyncio.gather(*(batcher.load(i) for i in range(5))) + + assert results == [0, 2, 4, 6, 8] + assert calls == [[0, 1, 2, 3, 4]] + + +async def test_size_trigger_flushes_without_waiting() -> None: + calls: list[list[int]] = [] + + async def batch_fn(reqs: list[int]) -> list[int]: + calls.append(list(reqs)) + return [r + 100 for r in reqs] + + # max_wait_ms is long; we rely on size trigger to flush. + batcher: Batcher[int, int] = Batcher(batch_fn, max_batch_size=3, max_wait_ms=5_000) + + results = await asyncio.gather(*(batcher.load(i) for i in range(3))) + + assert results == [100, 101, 102] + assert calls == [[0, 1, 2]] + + +async def test_size_trigger_splits_across_batches() -> None: + calls: list[list[int]] = [] + + async def batch_fn(reqs: list[int]) -> list[int]: + calls.append(list(reqs)) + return list(reqs) + + batcher: Batcher[int, int] = Batcher(batch_fn, max_batch_size=2, max_wait_ms=5) + + results = await asyncio.gather(*(batcher.load(i) for i in range(5))) + + assert results == [0, 1, 2, 3, 4] + # First two batches flushed by size, last batch flushed by timer. + assert len(calls) == 3 + assert sum(len(c) for c in calls) == 5 + + +async def test_time_trigger_flushes_partial_batch() -> None: + calls: list[list[int]] = [] + + async def batch_fn(reqs: list[int]) -> list[int]: + calls.append(list(reqs)) + return list(reqs) + + batcher: Batcher[int, int] = Batcher(batch_fn, max_batch_size=100, max_wait_ms=10) + + result = await batcher.load(42) + + assert result == 42 + assert calls == [[42]] + + +async def test_batch_fn_exception_propagates_to_every_caller() -> None: + class Boom(RuntimeError): + pass + + async def batch_fn(reqs: list[int]) -> list[int]: + raise Boom("provider down") + + batcher: Batcher[int, int] = Batcher(batch_fn, max_batch_size=10, max_wait_ms=5) + + awaitables = [batcher.load(i) for i in range(4)] + results = await asyncio.gather(*awaitables, return_exceptions=True) + + assert len(results) == 4 + assert all(isinstance(r, Boom) for r in results) + + +async def test_length_mismatch_is_reported_to_callers() -> None: + async def batch_fn(reqs: list[int]) -> list[int]: + return reqs[:-1] # one too few + + batcher: Batcher[int, int] = Batcher(batch_fn, max_batch_size=10, max_wait_ms=5) + + results = await asyncio.gather( + *(batcher.load(i) for i in range(3)), + return_exceptions=True, + ) + + assert all(isinstance(r, ValueError) for r in results) + + +async def test_flush_drains_pending_batch_immediately() -> None: + calls: list[list[int]] = [] + seen = asyncio.Event() + + async def batch_fn(reqs: list[int]) -> list[int]: + calls.append(list(reqs)) + seen.set() + return list(reqs) + + batcher: Batcher[int, int] = Batcher(batch_fn, max_batch_size=100, max_wait_ms=10_000) + + task = asyncio.create_task(batcher.load(7)) + await asyncio.sleep(0) # let load enqueue + await batcher.flush() + result = await task + + assert result == 7 + assert calls == [[7]] + assert seen.is_set() + + +async def test_invalid_construction_args() -> None: + async def batch_fn(reqs: list[int]) -> list[int]: + return list(reqs) + + with pytest.raises(ValueError): + Batcher(batch_fn, max_batch_size=0) + with pytest.raises(ValueError): + Batcher(batch_fn, max_wait_ms=-1) diff --git a/packages/core/tests/test_pipeline.py b/packages/core/tests/test_pipeline.py new file mode 100644 index 0000000..d1132d2 --- /dev/null +++ b/packages/core/tests/test_pipeline.py @@ -0,0 +1,243 @@ +"""Unit tests for ``rag_core.pipeline.Pipeline``.""" + +from __future__ import annotations + +import asyncio +from collections.abc import AsyncIterator + +import pytest +from rag_core.errors import PipelineError +from rag_core.pipeline import Pipeline, Stage + + +async def _aiter(items: list[int]) -> AsyncIterator[int]: + for it in items: + yield it + + +async def test_empty_pipeline_passes_items_through() -> None: + pipe: Pipeline[int, int] = Pipeline() + + out: list[int] = [] + async for item in pipe.run(_aiter([1, 2, 3])): + out.append(item) + + assert out == [1, 2, 3] + + +async def test_single_stage_one_to_one() -> None: + async def double(x: int) -> int: + return x * 2 + + pipe: Pipeline[int, int] = Pipeline().add_stage("double", double) + + out = [v async for v in pipe.run(_aiter([1, 2, 3]))] + assert out == [2, 4, 6] + + +async def test_stage_returning_none_drops_item() -> None: + async def even_only(x: int) -> int | None: + return x if x % 2 == 0 else None + + pipe: Pipeline[int, int] = Pipeline().add_stage("filter", even_only) + + out = [v async for v in pipe.run(_aiter([1, 2, 3, 4, 5]))] + assert out == [2, 4] + + +async def test_stage_returning_list_fans_out() -> None: + async def explode(x: int) -> list[int]: + return [x, x, x] + + pipe: Pipeline[int, int] = Pipeline().add_stage("explode", explode) + + out = [v async for v in pipe.run(_aiter([1, 2]))] + assert out == [1, 1, 1, 2, 2, 2] + + +async def test_multi_stage_chain_preserves_throughput() -> None: + async def parse(x: int) -> int: + return x + 1 + + async def chunk(x: int) -> list[int]: + return [x, x + 100] + + async def embed(x: int) -> int: + return x * -1 + + pipe: Pipeline[int, int] = ( + Pipeline() + .add_stage("parse", parse, workers=2) + .add_stage("chunk", chunk, workers=2) + .add_stage("embed", embed, workers=4) + ) + + out = sorted([v async for v in pipe.run(_aiter([1, 2, 3]))]) + # parse: 2,3,4 → chunk: (2,102),(3,103),(4,104) → embed: -2,-102,-3,-103,-4,-104 + assert out == sorted([-2, -102, -3, -103, -4, -104]) + + +async def test_per_stage_workers_run_concurrently() -> None: + in_flight = 0 + peak = 0 + lock = asyncio.Lock() + + async def slow(x: int) -> int: + nonlocal in_flight, peak + async with lock: + in_flight += 1 + peak = max(peak, in_flight) + await asyncio.sleep(0.02) + async with lock: + in_flight -= 1 + return x + + pipe: Pipeline[int, int] = Pipeline().add_stage("slow", slow, workers=4) + + out = [v async for v in pipe.run(_aiter(list(range(8))))] + + assert sorted(out) == list(range(8)) + assert peak >= 2, f"expected concurrent workers; peak was {peak}" + + +async def test_bounded_queue_creates_backpressure() -> None: + produced: list[int] = [] + + async def source() -> AsyncIterator[int]: + for i in range(20): + produced.append(i) + yield i + + sink_received = 0 + + async def slow_sink(x: int) -> int: + nonlocal sink_received + await asyncio.sleep(0.005) + sink_received += 1 + return x + + pipe: Pipeline[int, int] = Pipeline().add_stage("sink", slow_sink, workers=1, queue_size=2) + + iterator = pipe.run(source()) + + # Pull just one item, then sleep so the source can run. + first = await iterator.__anext__() + assert first == 0 + await asyncio.sleep(0.02) + + # The source should be backpressured — it cannot have run far ahead. + # With queue_size=2, output_size=2, 1 in-flight worker, and 1 already + # consumed, the steady-state ceiling is roughly 1+2+1+2+(1 feeder + # in-flight) = 7. Anything well below 20 (the full source) proves the + # backpressure is real — we use 10 to give scheduling jitter slack. + assert len(produced) < 10, f"backpressure failed; produced={produced}" + + # Drain the rest. + rest = [v async for v in iterator] + assert sorted([first, *rest]) == list(range(20)) + + +async def test_error_policy_fail_aborts_pipeline() -> None: + handled: list[tuple[str, int, str]] = [] + + async def explode(x: int) -> int: + if x == 3: + raise RuntimeError("nope") + return x + + async def handler(stage: str, item: int, exc: BaseException) -> None: + handled.append((stage, item, str(exc))) + + pipe: Pipeline[int, int] = Pipeline(on_error="fail", error_handler=handler).add_stage( + "explode", explode + ) + + with pytest.raises(PipelineError) as ei: + async for _ in pipe.run(_aiter([1, 2, 3, 4, 5])): + pass + + assert "explode" in str(ei.value) + assert ("explode", 3, "nope") in handled + + +async def test_error_policy_skip_drops_failing_item() -> None: + async def explode(x: int) -> int: + if x == 3: + raise RuntimeError("nope") + return x + + pipe: Pipeline[int, int] = Pipeline(on_error="skip").add_stage("explode", explode) + + out = [v async for v in pipe.run(_aiter([1, 2, 3, 4, 5]))] + assert sorted(out) == [1, 2, 4, 5] + + +async def test_source_exception_surfaces_as_pipeline_error() -> None: + async def bad_source() -> AsyncIterator[int]: + yield 1 + raise RuntimeError("source died") + + async def passthrough(x: int) -> int: + return x + + pipe: Pipeline[int, int] = Pipeline().add_stage("noop", passthrough) + + with pytest.raises(PipelineError) as ei: + async for _ in pipe.run(bad_source()): + pass + assert "source failed" in str(ei.value) + + +async def test_duplicate_stage_name_rejected() -> None: + async def fn(x: int) -> int: + return x + + pipe: Pipeline[int, int] = Pipeline().add_stage("a", fn) + with pytest.raises(ValueError, match="duplicate stage"): + pipe.add_stage("a", fn) + + +async def test_stage_validates_arguments() -> None: + async def fn(x: int) -> int: + return x + + with pytest.raises(ValueError): + Stage(name="bad", fn=fn, workers=0) + with pytest.raises(ValueError): + Stage(name="bad", fn=fn, queue_size=-1) + + +async def test_stage_default_queue_size_scales_with_workers() -> None: + async def fn(x: int) -> int: + return x + + s = Stage(name="x", fn=fn, workers=5) + assert s.queue_size == 10 # workers * 2 + + s2 = Stage(name="y", fn=fn) # workers=1, queue=2 minimum + assert s2.queue_size == 2 + + +async def test_consumer_early_break_does_not_leak_tasks() -> None: + processed = 0 + + async def fn(x: int) -> int: + nonlocal processed + processed += 1 + await asyncio.sleep(0) + return x + + pipe: Pipeline[int, int] = Pipeline().add_stage("fn", fn, workers=2, queue_size=2) + + iterator = pipe.run(_aiter(list(range(100)))) + first = await iterator.__anext__() + assert first == 0 + + # Closing the async generator should cancel everything. + await iterator.aclose() # type: ignore[attr-defined] + + # Give any leftover tasks a chance to be cleaned up. + await asyncio.sleep(0.01) + + # No assertion on `processed` value beyond "didn't process all 100". + assert processed < 100 diff --git a/uv.lock b/uv.lock index cc711f1..b35d026 100644 --- a/uv.lock +++ b/uv.lock @@ -3247,7 +3247,7 @@ provides-extras = ["eval"] [[package]] name = "rag-core" -version = "0.3.0" +version = "0.6.0" source = { editable = "packages/core" } dependencies = [ { name = "opentelemetry-api" },