Concurrency patterns for model calls

A model call spends almost all of its time waiting for a reply, so issuing them one after another leaves nearly all of your wall-clock on the table. This section is about running many at once without exhausting the provider's rate limit or your own memory: capping how many are in flight, collecting results as they finish rather than in order, and making sure one failure does not take the whole batch with it.

Requirement                                  What provides it
many slow calls overlapped                   asyncio.gather over coroutines
never more than N in flight                  asyncio.Semaphore(N), inside the task
one hung call cannot stall the job           asyncio.timeout around each call
one failure cannot lose the batch            return the exception as a value
a 429 does not kill the item                 retry with exponential backoff and jitter
a crash does not lose the work               append each result to disk as it lands
you can tell it apart from a hang            as_completed, print every N
you know what it cost before the bill        accumulate tokens per result

The shape you will write repeatedly#

Almost every eval harness, backfill and batch enrichment job has the same structure: many independent slow calls, a cap on how many run at once, a timeout on each, and a result set that tolerates some failures.

import asyncio

async def process_all(items, limit=8, timeout=60):
    sem = asyncio.Semaphore(limit)

    async def one(item):
        async with sem:
            try:
                async with asyncio.timeout(timeout):
                    return {"item": item, "ok": True, "result": await process(item)}
            except Exception as e:
                return {"item": item, "ok": False, "error": str(e)}

    return await asyncio.gather(*(one(i) for i in items))

Three details in that are load-bearing. The semaphore is acquired inside one, not around the gather, so a coroutine with no slot is parked rather than never created. The timeout is inside the semaphore, so a call that hangs gives its slot back after sixty seconds instead of holding it for ever. And the except returns rather than raises, which is what stops a single failure propagating out of gather and discarding every other result.

Retry the failures worth retrying#

Rate limits and transient server errors are the majority of failures in a large batch, and both succeed on a second attempt. Retry those and nothing else: a 400 for a malformed request will fail identically for ever, and retrying it three times wastes three times the latency to reach the same answer. Wait longer after each attempt, and add randomness, because otherwise every one of your eight workers fails at the same moment and retries at the same moment, which is the same burst again.

import random

RETRYABLE = (RateLimitError, APITimeoutError, InternalServerError)   # 429, 408, 5xx

async def with_retry(fn, *args, attempts=5, base=1.0):
    for attempt in range(attempts):
        try:
            return await fn(*args)
        except RETRYABLE:
            if attempt == attempts - 1:
                raise
            delay = base * (2 ** attempt)          # 1, 2, 4, 8 seconds
            await asyncio.sleep(delay * (0.5 + random.random()))   # jitter: 0.5x to 1.5x
    raise AssertionError("unreachable")

If the provider sends a retry-after header, obey it rather than your own schedule; it is the only number in the exchange that knows when the limit resets. And count retries separately from failures in whatever you print, because a job that succeeds after retrying four hundred times is telling you the concurrency limit is too high.

Progress on a long run#

A job that prints nothing for twenty minutes is indistinguishable from a hung one. as_completed lets you report as results arrive.

done = 0
for coro in asyncio.as_completed(tasks):
    result = await coro
    done += 1
    if done % 25 == 0:
        print(f"{done}/{len(tasks)}  cost so far ${total_cost:.2f}")

Print four things and you never have to guess what a running job is doing: how many are finished, how many failed, the elapsed time with a projection to the end, and the money spent so far. The projection is what makes you kill a job at minute two rather than discovering at minute ninety that it was never going to finish before the deadline.

start = time.perf_counter()
done = failed = 0
tokens_in = tokens_out = 0

for coro in asyncio.as_completed(tasks):
    r = await coro
    done += 1
    failed += not r["ok"]
    if r["ok"]:
        tokens_in += r["result"].usage.input_tokens
        tokens_out += r["result"].usage.output_tokens
    if done % 25 == 0:
        elapsed = time.perf_counter() - start
        rate = done / elapsed
        cost = tokens_in * 3e-6 + tokens_out * 15e-6
        remaining = (len(tasks) - done) / rate
        print(f"{done}/{len(tasks)}  {failed} failed  "
              f"{rate:.1f}/s  eta {remaining/60:.1f} min  ${cost:.2f}")

# 25/500  0 failed  4.1/s  eta 1.9 min  $0.41
# 50/500  2 failed  3.8/s  eta 2.0 min  $0.83

Make it resumable#

Long jobs get interrupted. Write each result as it completes rather than at the end, and skip anything already done on restart. This turns a crash from a lost afternoon into a lost minute.

done_ids = {json.loads(l)["id"] for l in open(out_path)} if exists(out_path) else set()
todo = [i for i in items if i.id not in done_ids]

with open(out_path, "a") as f:
    for coro in asyncio.as_completed([one(i) for i in todo]):
        result = await coro
        f.write(json.dumps(result) + "\n")
        f.flush()

Four things make this work and each is easy to leave out. The file is opened in append mode, so a restart adds rather than truncating. Each record carries its own id, because the order is completion order and position tells you nothing. flush is called after every line, or the last few thousand records sit in a buffer that a crash discards, which defeats the whole exercise. And the format is one JSON object per line, so a half-written final line from a kill in the middle of a write costs you exactly one record and the rest of the file still parses.

# Verified on 20 items, one of which always fails:
#   first run wrote 20 records in 0.07s
#   ok: 19  failed: 1
#   failures: [{'id': 7, 'ok': False, 'error': 'ValueError: bad item'}]
#   rerun processed: 0      -- every id was already in the file

# Skip the line a kill truncated rather than crashing on restart:
done_ids = set()
if exists(out_path):
    with open(out_path) as f:
        for line in f:
            try:
                done_ids.add(json.loads(line)["id"])
            except json.JSONDecodeError:
                pass        # a partial last line: it will simply be redone

Choosing the concurrency limit#

Too low and the job is slow. Too high and you are rate limited, which is slower still because of the backoff. Start at eight, watch for 429s, and adjust. The right number depends on your quota, not on your machine.

Providers usually limit two things at once: requests per minute and tokens per minute. A semaphore only caps requests, so a job sending short prompts can sit comfortably under the request limit and be throttled on tokens anyway. If that is happening, either lower the limit until it stops or track tokens yourself in a rolling window. The signal to watch is the proportion of calls that needed a retry: near zero means you can raise the limit, and consistently above about a tenth means you are already paying more in backoff than you gain in concurrency.

limit   wall clock   retries   what it means
    2       14m 20s         0      under-using the quota
    8        3m 40s         3      about right
   32        3m 10s       210      mostly backoff; barely faster, far noisier
  128        9m 50s     1,400+     slower than 8, and the provider notices

When async is not the answer#

Two other cases are worth recognising. If the work must happen in a fixed order, or each step needs the previous result, there is nothing to overlap and a plain loop is clearer. And if the job is large enough to run for hours, or needs to survive a machine restart, a queue with durable workers is the right shape rather than one script holding all the state in memory; the same batch of calls then restarts on a different machine without losing anything. Several providers also offer a batch endpoint that takes the whole job at a large discount and returns it within a day, which beats any amount of concurrency tuning when the answer is not needed now.