Batching Embedding Requests

A nightly job re-embeds the document library: 50,000 files, one API call each. It finished in 41 minutes on the first run and started returning 429s on the third. The loop is correct and the design is wrong — every request pays round-trip latency, connection setup, and a slot in the provider's per-minute quota, and none of that buys anything.

Batching sends many strings per request:

from openai import OpenAI

client = OpenAI()
MODEL = 'text-embedding-3-small'

def chunks(items, n):
    for i in range(0, len(items), n):
        yield items[i:i + n]

def embed_batched(docs, n=100):
    out = []
    for batch in chunks(docs, n):
        resp = client.embeddings.create(model=MODEL, input=batch)
        out.extend(d.embedding for d in resp.data)   # order mirrors the batch
    return out

Input: 50,000 documents with n=100.

Output: 500 requests instead of 50,000, returning the same vectors in the same order.

out.extend(d.embedding for d in resp.data) is the line to defend: response entries mirror the batch, so position is the only link between a vector and its document — the same contract described in Embedding Models Explained.

Two caps decide batch size, and only one of them counts items. Providers limit entries per request and total tokens per request. A hundred short FAQ lines fit; a hundred whole documents do not, and the rejection comes back as a 400 against the entire batch. When document lengths vary, size batches by cumulative tokens rather than by count.

A workable sizing rule: take the provider's tokens-per-request limit, divide by the 95th percentile of your document token lengths, and use that as the item cap. It is approximate — a batch of short documents under-uses the request — but it keeps every request under the cap without measuring each document individually. Re-derive the cap after any chunking change, since chunking sets the length distribution the cap depends on.

Re-running the corpus after a transient failure re-embeds 49,900 documents that already succeeded — the expensive half of the job paid twice, to recover from a two-second throttle. Failures need an owner at each level instead.

import random, time
from openai import APIConnectionError, APIStatusError, RateLimitError

def embed_with_retry(batch, attempts=4):
    for attempt in range(attempts):
        try:
            return list(client.embeddings.create(model=MODEL, input=batch).data)
        except RateLimitError:
            time.sleep(2 ** attempt + random.random())   # full jitter; quota refills
        except APIConnectionError:
            time.sleep(2 ** attempt + random.random())   # socket died mid-flight
        except APIStatusError as err:
            if err.status_code and err.status_code >= 500:
                time.sleep(2 ** attempt + random.random())
            else:
                raise        # 400/401: the payload is wrong, retrying repeats it
    return None               # caller marks only this batch as pending

def embed_corpus(docs, n=100):
    done, failed = [], []
    for batch in chunks(docs, n):
        data = embed_with_retry(batch)
        if data is None:
            failed.extend(batch)          # 100 documents, not 50,000
        else:
            done.extend(zip(batch, (d.embedding for d in data)))
    return done, failed

Input: 50,000 documents where one batch of 100 meets a sustained throttle.

Output: 49900 embedded, 100 pending — the retry queue holds the failed slice, and the next run sends only those 100.

Three lines carry the policy:

• return None instead of raising. A batch that exhausts its attempts is data for the caller, not an exception that kills a backfill halfway through with 49,900 vectors already computed.
• The bare raise on 4xx. A 400 means one document in the batch is malformed or over the token cap; three identical requests will fail identically, so the batch gets bisected to find the bad input instead of burning attempts on it.
• The jitter term. Workers that slept the same duration wake together and recreate the 429 that throttled them; random.random() spreads the wake-ups across the window.

Batching and retries still re-embed text that never changed. On a library where three documents were edited yesterday, the third run should cost three documents, not fifty thousand. Store a fingerprint beside each vector and the run becomes arithmetic:

import hashlib

def fingerprint(text):
    return hashlib.sha256(text.encode('utf-8')).hexdigest()

def plan(docs, store):
    to_embed, reused = [], []
    for doc in docs:
        if store.get(doc.id, {}).get('hash') == fingerprint(doc.text):
            reused.append(doc.id)         # unchanged: keep the stored vector
        else:
            to_embed.append(doc)          # new, edited, or never seen
    return to_embed, reused

Input: 10,000 documents against a store holding 9,963 matching fingerprints.

Output: embed 37, reuse 9963 — 99.6% of the corpus costs no API call.

The job as a whole:

documents ──► fingerprint ──┬── hash matches stored ──► keep vector, no call
                            │
                            └── new or changed ──► token-sized batches
                                                          │
                                                    retry + jitter
                                                          │
                                              ┌───────────┴───────────┐
                                              ▼                       ▼
                                     upsert(id, hash, vector)   failed batch ─►
                                                                  pending queue

The upsert writes id, hash, and vector in one operation. Split them across two writes and a crash in between leaves a fresh hash pointing at an old vector — after which the skip logic reports a document as current when its index entry no longer matches its text.

Rate limits then shape the schedule rather than the code. Issuing two million chunks at full concurrency produces a 429 storm that no retry policy can outwait. Pace the job under the provider's published requests-per-minute and tokens-per-minute limits, run backfills off-peak, and chart the 429 rate against throughput — a rising rate at flat throughput is throttling, not an outage.