Lesson 2: Concurrency Control
Semaphores for bounded concurrency
gather over ten thousand documents launches ten thousand simultaneous requests. Every one of them opens a connection, and the provider sees ten thousand requests arrive at once. This is the rate limiting disaster from the opening problem statement.
A semaphore limits how many tasks may be inside a section at once.
import asyncio
async def embed_all(docs: list[Document], *, concurrency: int = 20) -> list[Embedding]:
"""Embed every document with at most `concurrency` calls in flight."""
semaphore = asyncio.Semaphore(concurrency)
async def bounded(doc: Document) -> Embedding:
async with semaphore:
return await embed_document(doc)
async with asyncio.TaskGroup() as group:
tasks = [group.create_task(bounded(doc)) for doc in docs]
return [task.result() for task in tasks]
async with semaphore acquires a slot, waiting if all are taken, and releases it on exit including under cancellation. Ten thousand tasks exist, and at most twenty are calling the provider at any moment.
Choosing the limit. Start from what your dependencies allow rather than from a number that feels right. The provider's concurrent request limit, the connection pool size from Module 7, and the rate limit in requests per minute all bound it. The pool size point is easy to miss: a semaphore of fifty against a database pool of ten means forty tasks queue on the pool, which the pool timeout then turns into failures.
Measure rather than guess. Raise concurrency until throughput stops improving, then stop, because beyond that point you are adding queueing and error rate without adding speed.
Separate limits per dependency. One semaphore for the embedding provider, another for the generation provider, another for the database. A single global limit means a slow database throttles your embedding calls for no reason.
Rate limiting: token buckets, RPM and TPM
A semaphore bounds simultaneous requests. It does not bound requests per unit time. Twenty concurrent requests each taking fifty milliseconds is four hundred requests per second, which will exceed most limits.
Providers typically enforce two limits at once: requests per minute and tokens per minute. Hitting either produces a 429. The token limit is the one people forget, and it is often the binding constraint for long documents, because a hundred requests carrying eight thousand tokens each is eight hundred thousand tokens.
A token bucket is the standard mechanism. A bucket holds capacity, refills at a steady rate, and each request removes an amount. When the bucket is empty, callers wait.
import asyncio
import time
class TokenBucket:
"""Rate limiter allowing bursts up to capacity, refilling steadily."""
def __init__(self, *, rate_per_second: float, capacity: float) -> None:
self._rate = rate_per_second
self._capacity = capacity
self._tokens = capacity
self._updated = time.monotonic()
self._lock = asyncio.Lock()
async def acquire(self, amount: float = 1.0) -> None:
"""Wait until `amount` tokens are available, then consume them."""
if amount > self._capacity:
raise ValueError(f"request of {amount} exceeds capacity {self._capacity}")
while True:
async with self._lock:
now = time.monotonic()
elapsed = now - self._updated
self._tokens = min(self._capacity, self._tokens + elapsed * self._rate)
self._updated = now
if self._tokens >= amount:
self._tokens -= amount
return
deficit = amount - self._tokens
wait = deficit / self._rate
await asyncio.sleep(wait)
Note that the sleep happens outside the lock. Holding a lock while sleeping would serialise every waiter behind the first one.
time.monotonic appears again for the reason Module 5 gave: a clock adjustment must not make elapsed time negative.
Use two buckets, one per limit:
class ProviderLimiter:
def __init__(self, *, rpm: int, tpm: int) -> None:
self._requests = TokenBucket(rate_per_second=rpm / 60, capacity=rpm / 6)
self._tokens = TokenBucket(rate_per_second=tpm / 60, capacity=tpm / 6)
async def acquire(self, estimated_tokens: int) -> None:
await self._requests.acquire(1)
await self._tokens.acquire(estimated_tokens)
The capacity settings allow a burst of roughly ten seconds' worth, which smooths ordinary variation without permitting a flood.
estimated_tokens must be estimated before the call, which requires counting input tokens and reserving for output. That is a later module's subject, and until then a conservative estimate is better than none.
Adaptive throttling on 429
Configured limits are guesses. They change, they differ per account, and they are sometimes lower than documented. A system that only obeys its configuration will keep hitting 429s.
Adaptive throttling reduces the rate when rejections occur and recovers slowly when they stop. The pattern is borrowed from network congestion control: back off quickly, recover gradually.
class AdaptiveLimiter:
"""Reduces concurrency on rate limiting, recovers slowly on success."""
def __init__(self, *, initial: int, minimum: int = 1, maximum: int = 50) -> None:
self._limit = initial
self._minimum = minimum
self._maximum = maximum
self._successes_since_change = 0
self._semaphore = asyncio.Semaphore(initial)
def record_rate_limited(self) -> None:
"""Halve the limit immediately."""
new_limit = max(self._minimum, self._limit // 2)
if new_limit != self._limit:
logger.warning("reducing concurrency %d to %d", self._limit, new_limit)
self._limit = new_limit
self._successes_since_change = 0
def record_success(self) -> None:
"""Increase by one after a run of clean successes."""
self._successes_since_change += 1
if self._successes_since_change >= 50 and self._limit < self._maximum:
self._limit += 1
self._successes_since_change = 0
Halving on failure and adding one on success is deliberate. Fast reduction stops the damage; slow recovery avoids oscillating between flooding and backing off.
This composes with, rather than replaces, the retry logic from Module 5. Retry handles the individual request that failed. Adaptive throttling changes the rate so that fewer requests fail in the first place. And Retry-After, when the provider sends it, still overrides your own calculation.
Batching strategies and optimal batch size
Many APIs accept several inputs per request. Batching reduces per-request overhead and consumes fewer entries from your requests-per-minute budget.
from itertools import batched
async def embed_corpus(docs: Iterable[Document], *, batch_size: int = 64) -> list[Embedding]:
"""Embed in batches, with bounded concurrency across batches."""
semaphore = asyncio.Semaphore(10)
async def embed_batch(batch: tuple[Document, ...]) -> list[Embedding]:
async with semaphore:
await limiter.acquire(estimate_tokens(batch))
vectors = await provider.embed([d.content for d in batch], model=MODEL)
if len(vectors) != len(batch):
raise OutputError(
f"expected {len(batch)} embeddings, received {len(vectors)}"
)
return [Embedding(d.id, v) for d, v in zip(batch, vectors, strict=True)]
async with asyncio.TaskGroup() as group:
tasks = [group.create_task(embed_batch(b)) for b in batched(docs, batch_size)]
return [e for task in tasks for e in task.result()]
The count check from Module 5 reappears, and strict=True on zip is the same defence in a different form: it raises if the lengths differ rather than silently truncating, which would misalign every subsequent pairing.
Choosing a batch size. Larger batches amortise overhead better and fail more expensively, since one failure loses the whole batch. They also increase per-request latency and token count, which can hit the token limit sooner.
The practical approach is to start around 32 to 64 for embeddings, measure throughput at several sizes, and take the point where improvement flattens. Keep batches well under any documented maximum, and keep the estimated token count of a batch under the per-request limit with room to spare.
Batching and retries interact badly. A batch of 64 that fails on one bad input fails entirely, and retrying resends 63 inputs that were fine. When failures are input-specific rather than transient, splitting the batch and retrying the halves isolates the offender in a few attempts rather than losing the batch.
Fan-out and fan-in
Many operations decompose into independent parts that can run at once and be combined at the end. This is fan-out and fan-in.
async def build_context(query: str) -> RetrievalContext:
"""Run retrieval, metadata lookup, and history fetch concurrently."""
async with asyncio.TaskGroup() as group:
chunks_task = group.create_task(retriever.search(query, top_k=50))
metadata_task = group.create_task(repository.get_query_metadata(query))
history_task = group.create_task(conversations.load(session_id))
chunks = chunks_task.result()
reranked = await reranker.rank(query, chunks) # depends on chunks
return RetrievalContext(
chunks=reranked[:10],
metadata=metadata_task.result(),
history=history_task.result(),
)
The three independent operations run concurrently, so total time is the slowest rather than the sum. The reranking step depends on retrieval, so it stays sequential, which is the shape of most real pipelines: a concurrent phase, then a dependent step, then perhaps another concurrent phase.
Two things to watch. Fan-out multiplies load on everything downstream, so a fan-out of five across a hundred concurrent requests is five hundred concurrent operations. And a single slow branch determines total latency, which is why timeouts per branch matter more here than in sequential code. Consider whether a branch is optional, in which case a short timeout and a fallback beats waiting for it.
Cancellation, timeouts, and graceful shutdown
asyncio.timeout is the modern way to bound an operation, available from Python 3.11.
try:
async with asyncio.timeout(30):
result = await slow_operation()
except TimeoutError:
logger.warning("operation exceeded 30 seconds")
It cancels everything inside the block when the deadline passes. This is a different level from the HTTP timeouts in Module 5, and both are needed: the HTTP timeout bounds one request, and this bounds an operation that may involve several.
Cancellation propagates downward. Cancelling a task cancels what it is awaiting, which cancels what that is awaiting. That is what makes the client disconnect handling in Lesson 5 work.
Graceful shutdown means finishing in-flight work rather than dropping it.
async def run_pipeline(docs: AsyncIterator[Document]) -> None:
shutting_down = asyncio.Event()
def request_shutdown() -> None:
logger.info("shutdown requested, finishing in-flight work")
shutting_down.set()
loop = asyncio.get_running_loop()
for sig in (signal.SIGTERM, signal.SIGINT):
loop.add_signal_handler(sig, request_shutdown)
async with asyncio.TaskGroup() as group:
async for doc in docs:
if shutting_down.is_set():
break
group.create_task(process(doc))
# the TaskGroup will not exit until in-flight tasks finish
The TaskGroup provides the important half: after the loop breaks, the block waits for tasks already started. Combined with the checkpointing from Module 4, a shutdown loses no work and a restart resumes from the last checkpoint.
loop.add_signal_handler is used rather than signal.signal because it schedules the handler on the loop rather than interrupting it, which is safer in async code. Note that it is not available on Windows, where the synchronous form from Module 4 is the fallback.