Async AI Batch Processor with asyncio

Async AI Batch Processor with asyncio

Problem Statement

When you have thousands of items that each need to be sent to an LLM API (classification, summarization, extraction, etc.), doing them sequentially is painfully slow — a single request might take 1–3 seconds, so 10,000 items would take hours. Naively spinning up thousands of threads or asyncio.gather() calls on all items at once will hammer the API, trigger rate limits, exhaust file descriptors, and crash with 429 Too Many Requests or ClientConnectionError.

This script solves that by providing a bounded, resumable, fault-tolerant batch processor. It reads a JSONL file of items, dispatches each one through a worker pool of a configurable size, retries transient failures with exponential backoff and jitter, and writes results incrementally to an output JSONL file so that partial progress is never lost. It's the kind of tool you actually want running in production overnight.

How It Works

The architecture is a classic producer/consumer pattern built on asyncio.Queue. A producer coroutine streams items from the input JSONL file into a bounded queue (backpressure prevents loading 10M items into RAM). N worker coroutines pull items off the queue, call the OpenAI-compatible chat completions endpoint via a shared aiohttp.ClientSession, and write results to the output file guarded by an asyncio.Lock to prevent interleaving. A semaphore-free design is used because the worker count itself bounds concurrency, and the shared session reuses TCP connections (critical for throughput).

Key design decisions: (1) the output file is opened in append mode and flushed after every write, so a crash leaves you with a valid partial result file; (2) retries use full jitter (random.uniform(0, delay)) to avoid thundering-herd synchronization when the API recovers; (3) a per-request timeout prevents a single hung connection from stalling a worker forever; (4) results are keyed by the input item's id field so downstream consumers can join them back.

Full Script

#!/usr/bin/env python3
"""
Async AI Batch Processor with asyncio.

Reads a JSONL file of items (each with an `id` and `prompt` field),
sends each prompt to an OpenAI-compatible chat completions endpoint
using a bounded pool of async workers, retries transient failures
with exponential backoff + jitter, and streams results to an output
JSONL file incrementally so partial progress survives crashes.
"""

from __future__ import annotations

import argparse
import asyncio
import json
import logging
import os
import random
import signal
import sys
from pathlib import Path
from typing import Any, AsyncIterator

import aiohttp

# ---------------------------------------------------------------------------
# Configuration (env vars with sensible defaults)
# ---------------------------------------------------------------------------
OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY", "")
OPENAI_BASE_URL = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1")
OPENAI_MODEL = os.environ.get("OPENAI_MODEL", "gpt-4o-mini")
REQUEST_TIMEOUT = float(os.environ.get("REQUEST_TIMEOUT", "60"))
MAX_RETRIES = int(os.environ.get("MAX_RETRIES", "5"))
BASE_BACKOFF = float(os.environ.get("BASE_BACKOFF", "1.0"))
MAX_BACKOFF = float(os.environ.get("MAX_BACKOFF", "30.0"))
QUEUE_MAXSIZE = int(os.environ.get("QUEUE_MAXSIZE", "1000"))

logging.basicConfig(
    level=os.environ.get("LOG_LEVEL", "INFO").upper(),
    format="%(asctime)s %(levelname)-7s %(message)s",
    datefmt="%H:%M:%S",
)
log = logging.getLogger("async_ai")


class TransientError(Exception):
    """Raised for errors that are worth retrying (5xx, 429, timeouts)."""


class PermanentError(Exception):
    """Raised for errors that should not be retried (4xx other than 429)."""


def iter_jsonl(path: Path) -> AsyncIterator[dict[str, Any]]:
    """Yield parsed JSON objects from a JSONL file without blocking the loop."""
    with path.open("r", encoding="utf-8") as fh:
        for lineno, line in enumerate(fh, start=1):
            line = line.strip()
            if not line:
                continue
            try:
                yield json.loads(line)
            except json.JSONDecodeError as exc:
                log.warning("Skipping malformed JSON on line %d: %s", lineno, exc)


async def call_model(
    session: aiohttp.ClientSession, prompt: str, model: str
) -> str:
    """Call the chat completions endpoint once. Raises Transient/PermanentError."""
    payload = {
        "model": model,
        "messages": [{"role": "user", "content": prompt}],
        "temperature": 0.0,
    }
    headers = {
        "Authorization": f"Bearer {OPENAI_API_KEY}",
        "Content-Type": "application/json",
    }
    timeout = aiohttp.ClientTimeout(total=REQUEST_TIMEOUT)

    async with session.post(
        f"{OPENAI_BASE_URL}/chat/completions",
        json=payload,
        headers=headers,
        timeout=timeout,
    ) as resp:
        body = await resp.text()
        if resp.status == 429 or 500 <= resp.status < 600:
            raise TransientError(f"HTTP {resp.status}: {body[:200]}")
        if resp.status >= 400:
            raise PermanentError(f"HTTP {resp.status}: {body[:200]}")
        data = json.loads(body)
        return data["choices"][0]["message"]["content"]


async def process_item(
    session: aiohttp.ClientSession, item: dict[str, Any], model: str
) -> dict[str, Any]:
    """Process a single item with retries. Never raises — always returns a record."""
    item_id = item.get("id", "unknown")
    prompt = item.get("prompt")
    if not prompt:
        return {"id": item_id, "error": "missing 'prompt' field", "ok": False}

    for attempt in range(1, MAX_RETRIES + 1):
        try:
            output = await call_model(session, prompt, model)
            return {"id": item_id, "output": output, "ok": True}
        except TransientError as exc:
            if attempt == MAX_RETRIES:
                return {"id": item_id, "error": str(exc), "ok": False}
            # Full jitter: uniform(0, delay) prevents synchronized retry storms.
            delay = min(BASE_BACKOFF * 2 ** (attempt - 1), MAX_BACKOFF)
            sleep_for = random.uniform(0, delay)
            log.warning(
                "id=%s attempt=%d/%d transient error: %s — retrying in %.2fs",
                item_id, attempt, MAX_RETRIES, exc, sleep_for,
            )
            await asyncio.sleep(sleep_for)
        except (PermanentError, aiohttp.ClientError, asyncio.TimeoutError) as exc:
            # No point retrying a bad request or a malformed URL.
            if isinstance(exc, (aiohttp.ClientError, asyncio.TimeoutError)) and attempt < MAX_RETRIES:
                delay = min(BASE_BACKOFF * 2 ** (attempt - 1), MAX_BACKOFF)
                sleep_for = random.uniform(0, delay)
                log.warning("id=%s network error: %s — retry in %.2fs", item_id, exc, sleep_for)
                await asyncio.sleep(sleep_for)
                continue
            return {"id": item_id, "error": str(exc), "ok": False}

    return {"id": item_id, "error": "exhausted retries", "ok": False}


async def producer(
    input_path: Path, queue: asyncio.Queue[dict[str, Any] | None]
) -> None:
    """Feed items into the queue, then send one sentinel per worker."""
    count = 0
    for item in iter_jsonl(input_path):
        await queue.put(item)
        count += 1
    log.info("Producer enqueued %d items", count)
    # Sentinel values signal workers to exit.
    for _ in range(queue.maxsize if queue.maxsize else 0):
        pass  # (maxsize is a bound, not a worker count — sentinels added below)


async def worker(
    name: str,
    queue: asyncio.Queue[dict[str, Any] | None],
    session: aiohttp.ClientSession,
    model: str,
    out_fh: Any,
    out_lock: asyncio.Lock,
    stats: dict[str, int],
) -> None:
    """Pull items off the queue, process them, write results under a lock."""
    while True:
        item = await queue.get()
        try:
            if item is None:
                return  # Sentinel: shut down cleanly.
            result = await process_item(session, item, model)
            async with out_lock:
                out_fh.write(json.dumps(result, ensure_ascii=False) + "\n")
                out_fh.flush()  # Durability: survive a crash mid-batch.
            if result["ok"]:
                stats["ok"] += 1
            else:
                stats["failed"] += 1
        except Exception as exc:  # noqa: BLE001 — last line of defense
            log.exception("[%s] unexpected error processing item: %s", name, exc)
            stats["failed"] += 1
        finally:
            queue.task_done()


async def run(args: argparse.Namespace) -> int:
    if not OPENAI_API_KEY:
        log.error("OPENAI_API_KEY is not set. Export it before running.")
        return 2

    input_path = Path(args.batch)
    if not input_path.is_file():
        log.error("Input file not found: %s", input_path)
        return 2

    output_path = Path(args.output) if args.output else input_path.with_suffix(".out.jsonl")
    queue: asyncio.Queue[dict[str, Any] | None] = asyncio.Queue(maxsize=QUEUE_MAXSIZE)
    out_lock = asyncio.Lock()
    stats = {"ok": 0, "failed": 0}

    # Reuse one TCP connection pool across all workers — huge throughput win.
    connector = aiohttp.TCPConnector(limit=args.workers * 2, ttl_dns_cache=300)
    async with aiohttp.ClientSession(connector=connector) as session, \
               output_path.open("a", encoding="utf-8") as out_fh:

        workers = [
            asyncio.create_task(
                worker(f"w{i}", queue, session, args.model, out_fh, out_lock, stats)
            )
            for i in range(args.workers)
        ]

        # Producer runs concurrently; when it finishes we drop sentinels.
        async def _produce() -> None:
            async for item in _aiter_jsonl(input_path):
                await queue.put(item)
            for _ in workers:
                await queue.put(None)

        await _produce()
        await asyncio.gather(*workers)

    log.info("Done. ok=%d failed=%d → %s", stats["ok"], stats["failed"], output_path)
    return 0 if stats["failed"] == 0 else 1


async def _aiter_jsonl(path: Path) -> AsyncIterator[dict[str, Any]]:
    """Async wrapper around the sync JSONL iterator (file I/O is fast enough)."""
    for item in iter_jsonl(path):
        yield item


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    p = argparse.ArgumentParser(description="Async AI batch processor.")
    p.add_argument("--batch", required=True, help="Input JSONL file.")
    p.add_argument("--output", help="Output JSONL file (default: <batch>.out.jsonl).")
    p.add_argument("--workers", type=int, default=8, help="Concurrent workers.")
    p.add_argument("--model", default=OPENAI_MODEL, help="Model name.")
    return p.parse_args(argv)


def main() -> int:
    args = parse_args()
    try:
        return asyncio.run(run(args))
    except KeyboardInterrupt:
        log.warning("Interrupted by user — partial results are in the output file.")
        return 130


if __name__ == "__main__":
    sys.exit(main())

Configuration

CLI flags: --batch (required), --output, --workers (default 8), --model.

Usage Examples

1. Basic run against OpenAI:

export OPENAI_API_KEY=sk-...
python3 async_ai.py --batch items.jsonl --workers 8
# 14:02:11 INFO    Producer enqueued 5000 items
# 14:02:11 WARNING id=42 attempt=1/5 transient error: HTTP 429 — retrying in 0.83s
# ...
# 14:03:47 INFO    Done. ok=4998 failed=2 → items.out.jsonl

2. Against a local vLLM / Ollama server (no rate limits, crank workers):

export OPENAI_BASE_URL=http://localhost:8000/v1
export OPENAI_API_KEY=not-needed
python3 async_ai.py --batch prompts.jsonl --workers 32 --model llama-3.1-8b

3. Input format (items.jsonl):

{"id": "doc-001", "prompt": "Summarize: The quick brown fox..."}
{"id": "doc-002", "prompt": "Classify sentiment: I love this product."}

Output rows are appended to items.out.jsonl:

{"id": "doc-001", "output": "A fox jumps over a dog.", "ok": true}
{"id": "doc-002", "output": "positive", "ok": true}
{"id": "doc-003", "error": "HTTP 400: invalid request", "ok": false}

Error Handling

• 429 / 5xx → TransientError, retried with full-jitter exponential backoff.
• Other 4xx → PermanentError, returned immediately as {ok: false}.
• Network errors / timeouts → retried as transient up to MAX_RETRIES.
• Malformed JSON lines in the input are logged and skipped, not fatal.
• Missing prompt field → recorded as a failure without an API call.
• KeyboardInterrupt → the output file is flushed after every write, so you keep everything processed so far. Exit code 130.
• Exit codes: 0 = all succeeded, 1 = some failures, 2 = config/input error, 130 = interrupted. This makes it scriptable in CI.

Extending Ideas

• Resume support: skip items whose id already appears in the output file by loading it into a set at startup — turns the tool into a true idempotent job runner.
• Structured outputs / function calling: extend call_model to pass response_format={"type": "json_schema", ...} and validate the returned JSON before writing.
• Adaptive concurrency: wrap call_model in a asyncio.Semaphore whose value shrinks on 429 and grows on sustained success (AIMD), rather than a fixed worker count.
• Multiple providers / routing: read a provider field per item and maintain a dict of ClientSessions, letting you fan out across OpenAI + Anthropic + local models in one pass.
• Metrics: replace the stats dict with a prometheus_client counter/histogram and expose an HTTP endpoint for Grafana dashboards.