Orchestrating Multi-Stage AI Pipelines

Orchestration frameworks are mostly interchangeable. What is not interchangeable is whether your stages can be run twice without harm, because every queue you will ever use delivers at least once and sometimes more.

At-least-once is what you get

Exactly-once delivery is not available across a network. What is available is at-least-once delivery plus idempotent processing, which composes to exactly-once effect — and the second half is your job, not the queue’s.

The duplicates are not hypothetical. A worker that finishes its work and dies before acknowledging causes a redelivery. A visibility timeout that expires because a stage was slow causes a redelivery while the first attempt is still running. A deploy that restarts workers mid-flight causes a batch of them. In an AI pipeline every one of those costs money, because the duplicated work includes model calls.

There is a second reason to design for repetition that has nothing to do with queues: you will re-run things on purpose. A parser fix, a config change, a backfill, a partial recovery after an incident — each of those replays work that has already been done, and a pipeline whose stages are only safe to run once turns every one of them into a careful, manual, out-of-hours operation. Idempotency is what makes re-running boring, and boring is the property you want from the thing you do under pressure.

Making a stage idempotent

The mechanism is the same one that makes incremental processing work: derive a key from the content and the code, check before doing the expensive thing, and write under that key.

def embed_chunk(chunk_id: bytes, text: str, model: str, conn) -> None:
    # 1. Cheap existence check: the common case for a redelivery.
    if conn.execute("SELECT 1 FROM embeddings WHERE chunk_id=%s AND model=%s",
                    (chunk_id, model)).fetchone():
        return

    # 2. The expensive, billable call.
    vec = provider.embed(text, model=model)

    # 3. Write so that a concurrent duplicate is harmless, not an error.
    conn.execute(
        "INSERT INTO embeddings (chunk_id, model, dim, vec) "
        "VALUES (%s,%s,%s,%s) ON CONFLICT (chunk_id, model) DO NOTHING",
        (chunk_id, model, len(vec), vec))

Both halves are needed. The check at the top saves the money in the common case; the ON CONFLICT DO NOTHING handles the race where two workers pass the check simultaneously, which the check alone cannot prevent. A stage written this way can be run any number of times and the result is the same as running it once.

Where the expensive call itself has side effects — sending a notification, calling a stateful external API — the provider’s own idempotency key is the tool, and the general pattern is worth knowing separately.

Retry, but not for everything

Blanket retry is how a transient failure becomes a bill. Sort failures into three classes and treat them differently:

import random, time

RETRYABLE = {408, 429, 500, 502, 503, 504}

def with_retry(fn, attempts=5, base=1.0, cap=60.0):
    for i in range(attempts):
        try:
            return fn()
        except HTTPError as e:
            if e.status not in RETRYABLE or i == attempts - 1:
                raise                       # terminal, or out of attempts
            wait = float(e.headers.get("Retry-After", 0)) or \
                   random.uniform(0, min(cap, base * 2 ** i))
            time.sleep(wait)

Full jitter — a uniform draw from zero to the backoff bound, not the bound itself — is the detail that matters at scale. Without it, a thousand workers that failed together retry together, and the synchronised wave is what keeps the dependency down. The retry-safety rules are worth reading in full before you make retries automatic.

Partial reruns

The operation you will perform most often is “re-run stage 4 for the documents where it failed”, and the pipeline should make that a query rather than an argument. Two requirements.

Per-item state, not per-run state. A run that reports “failed” tells you nothing about which of its 40,000 items need attention. Record a row per item per stage with its status, attempt count, last error and the code version that ran.

-- what needs rerunning after a parser fix, as a query
SELECT document_id FROM stage_runs
WHERE stage = 'extract' AND status = 'failed'
  AND error_class = 'cid_leak'
  AND extractor_ver < '1.24.9';

Asset-shaped stages. Frameworks differ in how much they help here. Airflow models tasks; Dagster models assets, so “materialise this asset and everything downstream of it” is a first-class operation; Temporal models durable workflow executions, which is a good fit when a pipeline has to wait days for an external result. The abstraction you want for a document pipeline is the asset one, because “the chunks of document X under config Y” is a thing that either exists and is current or does not — and that is exactly what a content-addressed artefact store already gives you, framework or no framework.

The poison document

One document in the corpus will crash the parser, exhaust memory, or hit a timeout every single time. Delivered at-least-once, it is redelivered forever, and it takes a worker with it each time. Left alone, one file stops the pipeline.

Cap attempts per item and route the exhausted ones to a dead-letter queue with the full error. Then treat the DLQ as a work list with an owner, because the failure mode of dead-letter queues is that nobody looks at them: alert on its rate of growth rather than its depth, since depth is a number people learn to ignore and a change in rate is a new bug.

Guard resources per item as well as attempts. A memory limit per worker process, a wall-clock timeout per stage, and a maximum input size checked before the expensive call — a 900 MB PDF should be quarantined on the size check, not discovered when the box runs out of memory and the kernel kills something unrelated.

One nastier variant deserves naming: the item that does not crash but takes a hundred times longer than any other. It holds a worker for hours, the queue backs up behind it, and every dashboard says the pipeline is healthy because nothing has failed. A per-stage wall-clock timeout converts that silent stall into an ordinary failure with an error message, which is the only form in which anybody will notice it.