AI Engineer Portal

Async Patterns for High-Throughput AI Data Pipelines

asyncio.gather with a semaphore handles most AI batch jobs. But as pipelines grow — multiple stages, mixed CPU and I/O work, producer-consumer relationships, downstream rate limits — the simple gather pattern breaks down. This article covers the patterns that scale.

When gather + semaphore is not enough

The basic pattern handles one stage well:

results = await asyncio.gather(
    *[process_one(item) for item in items],
    return_exceptions=True,
)

It breaks down when:

Stages have different rates. A document parser (CPU-bound, fast) feeds an embedding API (I/O-bound, slow). The parser can overwhelm the embedder if they are tightly coupled.
Output needs to stream to a downstream system. Waiting for all 10,000 items to complete before writing to the database wastes memory and delays output.
The pipeline has more than two stages. Gather is a fan-out, not a pipeline. For sequential stages, you need a different pattern.

Producer-consumer with asyncio.Queue

The asyncio.Queue pattern decouples stages and provides natural backpressure:

import asyncio
import logging

logger = logging.getLogger(__name__)


async def producer(items: list[dict], queue: asyncio.Queue) -> None:
    for item in items:
        await queue.put(item)  # blocks if queue is full (backpressure)
    await queue.put(None)  # sentinel: producer is done


async def consumer(
    queue: asyncio.Queue,
    output_queue: asyncio.Queue,
    client,
    semaphore: asyncio.Semaphore,
) -> None:
    while True:
        item = await queue.get()
        if item is None:
            await output_queue.put(None)  # forward sentinel
            queue.task_done()
            return
        async with semaphore:
            try:
                result = await asyncio.wait_for(client.process(item), timeout=30.0)
                await output_queue.put(result)
            except Exception as exc:
                logger.warning("consumer_error", extra={"id": item.get("id"), "error": str(exc)})
        queue.task_done()


async def writer(output_queue: asyncio.Queue, output_path: str) -> int:
    import json
    from pathlib import Path
    count = 0
    with Path(output_path).open("w") as f:
        while True:
            result = await output_queue.get()
            if result is None:
                return count
            f.write(json.dumps(result) + "\n")
            count += 1
            output_queue.task_done()

The maxsize parameter on asyncio.Queue(maxsize=N) creates backpressure: if the consumer is slower than the producer, await queue.put() blocks the producer until space is available. Without maxsize, the queue grows without bound and consumes all available memory.

Fan-out with per-stage semaphores

When a pipeline has multiple parallel stages with different rate limits:

async def run_pipeline(
    items: list[dict],
    embedding_client,
    llm_client,
    embed_concurrency: int = 10,
    llm_concurrency: int = 3,
) -> list[dict]:
    embed_sem = asyncio.Semaphore(embed_concurrency)
    llm_sem = asyncio.Semaphore(llm_concurrency)

    async def process_one(item: dict) -> dict:
        # Stage 1: embed (higher concurrency, cheaper API)
        async with embed_sem:
            embedding = await embedding_client.embed(item["text"])

        # Stage 2: generate (lower concurrency, expensive API)
        async with llm_sem:
            result = await llm_client.generate(
                prompt=build_prompt(item["text"], embedding)
            )
        return {**item, "embedding": embedding, "result": result}

    raw = await asyncio.gather(*[process_one(i) for i in items], return_exceptions=True)
    return [r for r in raw if not isinstance(r, Exception)]

Each stage has its own semaphore sized to that stage's rate limit. The embedding API might support 10 concurrent requests; the LLM API might only support 3. Both constraints are enforced without coupling.

asyncio vs multiprocessing: when to use each

The choice is determined by whether your bottleneck is I/O or CPU:

Bottleneck	Pattern	Reason
Provider API calls	`asyncio.gather` + `Semaphore`	I/O-bound: async is ideal
Embedding API calls	`asyncio.gather` + `Semaphore`	I/O-bound
Document parsing (text extraction)	`ProcessPoolExecutor`	CPU-bound: GIL prevents threading
Tokenization for large corpora	`ProcessPoolExecutor`	CPU-bound
Mixed: parse then embed	`to_thread` for parsing, `asyncio` for I/O	Blend both

import asyncio
from concurrent.futures import ProcessPoolExecutor


def parse_document(path: str) -> dict:
    # CPU-intensive: runs in a subprocess, no GIL contention
    text = extract_text_from_pdf(path)
    return {"path": path, "text": text}


async def pipeline(paths: list[str], embedding_client) -> list[dict]:
    sem = asyncio.Semaphore(8)
    loop = asyncio.get_event_loop()

    async def process_one(path: str) -> dict:
        # Run CPU work in a process pool
        parsed = await loop.run_in_executor(None, parse_document, path)
        # Run I/O work with asyncio
        async with sem:
            embedding = await embedding_client.embed(parsed["text"])
        return {**parsed, "embedding": embedding}

    results = await asyncio.gather(*[process_one(p) for p in paths], return_exceptions=True)
    return [r for r in results if not isinstance(r, Exception)]

loop.run_in_executor(None, fn, *args) runs a synchronous function in the default ThreadPoolExecutor. For CPU-intensive work, pass a ProcessPoolExecutor explicitly.

Monitoring pipeline health

For long-running pipelines, log progress periodically rather than at completion:

async def monitored_pipeline(items: list[dict], client) -> list[dict]:
    sem = asyncio.Semaphore(5)
    results = []
    processed = 0

    async def _one(item: dict) -> dict:
        nonlocal processed
        async with sem:
            result = await client.process(item)
            processed += 1
            if processed % 100 == 0:
                logger.info("pipeline_progress", extra={
                    "processed": processed,
                    "total": len(items),
                    "pct": round(100 * processed / len(items), 1),
                })
            return result

    raw = await asyncio.gather(*[_one(i) for i in items], return_exceptions=True)
    return [r for r in raw if not isinstance(r, Exception)]

Progress logging every N items is essential for long-running pipelines. Without it, a 10,000-item run with no output for 30 minutes looks like a hang.

The practical checklist for any async pipeline

Set maxsize on all queues — unbounded queues are memory leaks waiting for a slow consumer.
Set a timeout on every await that touches the network — use asyncio.wait_for.
Use return_exceptions=True in every gather — one failure should not cancel 999 others.
Size semaphores to provider rate limits, not throughput goals — the rate limit is a constraint, not a dial.
Log progress every N items for runs longer than 1 minute — progress visibility prevents false "it hung" diagnoses.

Your personal operating system for career transition.