Asyncio for AI Engineering

from __future__ import annotations

import asyncio
import random
from dataclasses import dataclass, field


class TransientLLMError(Exception):
    def __init__(self, msg: str, retry_after: float | None = None) -> None:
        super().__init__(msg)
        self.retry_after = retry_after


class PermanentLLMError(Exception):
    pass


@dataclass
class Stats:
    calls: int = 0
    retries: int = 0
    failures: int = 0
    peak_in_flight: int = 0
    in_flight: int = field(default=0, repr=False)


class FakeTransport:
    """Stands in for an SDK: ~10% of calls get a 429, prompts containing 'bad' are rejected."""

    def __init__(self, seed: int = 7) -> None:
        self.rng = random.Random(seed)

    async def complete(self, prompt: str) -> str:
        await asyncio.sleep(self.rng.uniform(0.001, 0.005))
        if "bad" in prompt:
            raise PermanentLLMError("400 invalid request")
        if self.rng.random() < 0.1:
            raise TransientLLMError("429 rate limited", retry_after=0.001)
        return f"answer to {prompt}"


class LLMClient:
    def __init__(self, transport, *, max_concurrency: int = 16, attempt_timeout: float = 60.0,
                 attempts: int = 5) -> None:
        self.transport = transport
        self.sem = asyncio.Semaphore(max_concurrency)
        self.attempt_timeout = attempt_timeout
        self.attempts = attempts
        self.stats = Stats()

    async def complete(self, prompt: str) -> str:
        for attempt in range(self.attempts):
            try:
                async with self.sem:
                    self.stats.calls += 1
                    self.stats.in_flight += 1
                    self.stats.peak_in_flight = max(self.stats.peak_in_flight, self.stats.in_flight)
                    try:
                        async with asyncio.timeout(self.attempt_timeout):
                            return await self.transport.complete(prompt)
                    finally:
                        self.stats.in_flight -= 1
            except (TransientLLMError, TimeoutError) as err:
                if attempt == self.attempts - 1:
                    self.stats.failures += 1
                    raise
                self.stats.retries += 1
                delay = getattr(err, "retry_after", None) or random.uniform(0, min(30.0, 0.5 * 2 ** attempt))
                await asyncio.sleep(delay)          # sleep OUTSIDE the semaphore: do not hold a slot
            except PermanentLLMError:
                self.stats.failures += 1
                raise
        raise AssertionError("unreachable")


async def main() -> None:
    client = LLMClient(FakeTransport(), max_concurrency=8)
    prompts = [f"q{i}" for i in range(200)] + ["bad prompt"]
    results = await asyncio.gather(*(client.complete(p) for p in prompts), return_exceptions=True)
    ok = sum(not isinstance(r, BaseException) for r in results)
    print(ok, "ok;", "peak in flight:", client.stats.peak_in_flight, "; retried:", client.stats.retries > 0)


asyncio.run(main())
# Output: 200 ok; peak in flight: 8 ; retried: True

import asyncio
import json
from contextlib import aclosing


async def fake_sse_lines(text: str):
    """What `response.aiter_lines()` yields for an SSE chat stream."""
    for word in text.split():
        await asyncio.sleep(0)
        yield "data: " + json.dumps({"delta": word + " "})
        yield ""
    yield "data: [DONE]"


async def stream_tokens(lines):
    async with aclosing(lines) as it:              # closes the HTTP stream even on early exit
        async for line in it:
            if not line.startswith("data: "):
                continue                           # blank separators, comments, event: lines
            payload = line[len("data: "):]
            if payload == "[DONE]":
                return
            yield json.loads(payload)["delta"]


async def main() -> None:
    out = []
    async with aclosing(stream_tokens(fake_sse_lines("retrieval augmented generation works"))) as tokens:
        async for tok in tokens:
            out.append(tok)
            if "generation" in tok:
                break                              # early stop: both generators are closed
    print("".join(out).strip())


asyncio.run(main())
# Output: retrieval augmented generation

import asyncio
import json
from collections.abc import AsyncIterator

from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse

app = FastAPI()


async def generate(prompt: str) -> AsyncIterator[str]:
    for tok in ("Hello", " from", " the", " model"):
        await asyncio.sleep(0.05)
        yield tok


@app.get("/chat")
async def chat(request: Request, q: str) -> StreamingResponse:
    async def events() -> AsyncIterator[str]:
        try:
            async for tok in generate(q):
                if await request.is_disconnected():     # cheap check; stop paying for tokens
                    break
                yield f"data: {json.dumps({'token': tok})}\n\n"
            yield "data: [DONE]\n\n"
        except asyncio.CancelledError:
            # Raised into the generator when the server cancels the response (client gone, shutdown).
            raise
        finally:
            pass                                       # close upstream stream / record usage here

    return StreamingResponse(events(), media_type="text/event-stream",
                             headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"})

from __future__ import annotations

import asyncio


class BatchTooLarge(Exception):
    pass


class BadInput(Exception):
    pass


async def embed_batch(texts: list[str]) -> list[list[float]]:
    await asyncio.sleep(0.001)
    if any(t == "<binary garbage>" for t in texts):
        raise BadInput("input 400: invalid characters")
    if sum(len(t) for t in texts) > 200:
        raise BatchTooLarge("413: batch exceeds token limit")
    return [[float(len(t))] for t in texts]


def make_batches(texts: list[str], max_items: int, max_chars: int) -> list[list[int]]:
    batches, cur, size = [], [], 0
    for i, t in enumerate(texts):
        if cur and (len(cur) >= max_items or size + len(t) > max_chars):
            batches.append(cur)
            cur, size = [], 0
        cur.append(i)
        size += len(t)
    if cur:
        batches.append(cur)
    return batches


async def embed_all(texts: list[str], max_items: int = 8, max_chars: int = 400,
                    concurrency: int = 4) -> tuple[list[list[float] | None], dict[int, str]]:
    vectors: list[list[float] | None] = [None] * len(texts)
    errors: dict[int, str] = {}
    sem = asyncio.Semaphore(concurrency)

    async def run(idx: list[int]) -> None:
        try:
            async with sem:
                result = await embed_batch([texts[i] for i in idx])
        except (BatchTooLarge, BadInput) as err:
            if len(idx) == 1:
                errors[idx[0]] = str(err)               # isolated the bad input
                return
            mid = len(idx) // 2                          # halve and retry both parts
            await asyncio.gather(run(idx[:mid]), run(idx[mid:]))
            return
        for i, vec in zip(idx, result):
            vectors[i] = vec

    await asyncio.gather(*(run(b) for b in make_batches(texts, max_items, max_chars)))
    return vectors, errors


async def main() -> None:
    texts = [f"chunk {i} " * (i % 5 + 1) for i in range(30)]
    texts[13] = "<binary garbage>"
    vectors, errors = await embed_all(texts)
    print(sum(v is not None for v in vectors), "embedded;", errors)


asyncio.run(main())
# Output: 29 embedded; {13: 'input 400: invalid characters'}

import asyncio
import json


async def search_docs(query: str) -> str:
    await asyncio.sleep(0.01)
    return json.dumps({"hits": [f"doc about {query}"]})


async def get_weather(city: str) -> str:
    await asyncio.sleep(1.0)                      # a hung upstream
    return "sunny"


async def run_sql(query: str) -> str:
    raise PermissionError("read-only role cannot run DELETE")


TOOLS = {"search_docs": search_docs, "get_weather": get_weather, "run_sql": run_sql}


async def execute_tool_call(call: dict, timeout: float = 0.1) -> dict:
    fn = TOOLS.get(call["name"])
    try:
        if fn is None:
            raise LookupError(f"unknown tool {call['name']!r}")
        async with asyncio.timeout(timeout):
            content = await fn(**json.loads(call["arguments"]))
        return {"tool_call_id": call["id"], "content": content}
    except Exception as err:                      # CancelledError is not an Exception: it propagates
        detail = str(err) or f"no result within {timeout}s"
        return {"tool_call_id": call["id"], "content": f"ERROR {type(err).__name__}: {detail}", "is_error": True}


async def main() -> None:
    calls = [
        {"id": "c1", "name": "search_docs", "arguments": json.dumps({"query": "refund policy"})},
        {"id": "c2", "name": "get_weather", "arguments": json.dumps({"city": "Paris"})},
        {"id": "c3", "name": "run_sql", "arguments": json.dumps({"query": "DELETE FROM users"})},
    ]
    async with asyncio.TaskGroup() as tg:
        tasks = [tg.create_task(execute_tool_call(c), name=f"tool-{c['name']}") for c in calls]
    for t in tasks:
        r = t.result()
        print(r["tool_call_id"], r["content"])


asyncio.run(main())
# Output:
# c1 {"hits": ["doc about refund policy"]}
# c2 ERROR TimeoutError: no result within 0.1s
# c3 ERROR PermissionError: read-only role cannot run DELETE

import asyncio
import json
import os
import tempfile


async def model_answer(prompt: str) -> str:
    await asyncio.sleep(0.001)
    return prompt[::-1]


async def judge(prompt: str, answer: str) -> float:
    await asyncio.sleep(0.001)
    return 1.0 if answer == prompt[::-1] else 0.0


def _read_rows(path: str) -> list[dict]:
    if not os.path.exists(path):
        return []
    with open(path, encoding="utf-8") as fh:
        return [json.loads(line) for line in fh if line.strip()]


def _append(path: str, line: str) -> None:
    with open(path, "a", encoding="utf-8") as fh:
        fh.write(line)


async def run_eval(cases: list[dict], out_path: str, concurrency: int = 16) -> float:
    done_ids = {row["id"] for row in await asyncio.to_thread(_read_rows, out_path)}   # resume
    todo = [c for c in cases if c["id"] not in done_ids]
    sem = asyncio.Semaphore(concurrency)
    write_lock = asyncio.Lock()

    async def one(case: dict) -> None:
        async with sem:
            answer = await model_answer(case["prompt"])
            score = await judge(case["prompt"], answer)
        line = json.dumps({"id": case["id"], "score": score}) + "\n"
        async with write_lock:                                   # one writer at a time
            await asyncio.to_thread(_append, out_path, line)    # file I/O off the loop

    await asyncio.gather(*(one(c) for c in todo))
    scores = [row["score"] for row in await asyncio.to_thread(_read_rows, out_path)]
    return sum(scores) / len(scores)


async def main() -> None:
    cases = [{"id": i, "prompt": f"prompt-{i}"} for i in range(100)]
    with tempfile.TemporaryDirectory() as d:
        path = os.path.join(d, "results.jsonl")
        await run_eval(cases[:40], path)              # "crashed" after 40
        mean = await run_eval(cases, path)            # resume: only the remaining 60 run
        rows = await asyncio.to_thread(_read_rows, path)
        print(len(rows), "rows, mean score", mean)


asyncio.run(main())
# Output: 100 rows, mean score 1.0

from contextlib import asynccontextmanager

import httpx
from fastapi import Depends, FastAPI, Request


@asynccontextmanager
async def lifespan(app: FastAPI):
    app.state.llm_http = httpx.AsyncClient(
        base_url="https://llm.internal",
        timeout=httpx.Timeout(120, connect=5, pool=10),
        limits=httpx.Limits(max_connections=64, max_keepalive_connections=32),
    )
    try:
        yield
    finally:
        await app.state.llm_http.aclose()


app = FastAPI(lifespan=lifespan)


def llm_http(request: Request) -> httpx.AsyncClient:
    return request.app.state.llm_http


@app.post("/v1/answer")
async def answer(payload: dict, http: httpx.AsyncClient = Depends(llm_http)) -> dict:
    resp = await http.post("/generate", json=payload)
    resp.raise_for_status()
    return resp.json()

async def upsert_stage(queue: asyncio.Queue, store, batch_size: int = 256) -> None:
    batch = []
    while True:
        item = await queue.get()
        try:
            if item is None:                            # sentinel: flush and stop
                if batch:
                    await store.upsert(batch)
                return
            batch.append(item)
            if len(batch) >= batch_size:
                await store.upsert(batch)
                batch = []
        finally:
            queue.task_done()

import asyncio
from concurrent.futures import ProcessPoolExecutor


def chunk_document(text: str, size: int = 400, overlap: int = 50) -> list[str]:
    step = size - overlap
    return [text[i:i + size] for i in range(0, max(len(text) - overlap, 1), step)]


async def chunk_all(texts: list[str], pool: ProcessPoolExecutor) -> list[list[str]]:
    loop = asyncio.get_running_loop()
    return await asyncio.gather(*(loop.run_in_executor(pool, chunk_document, t) for t in texts))


async def main() -> None:
    with ProcessPoolExecutor(max_workers=2) as pool:
        chunks = await chunk_all(["x" * 1000, "y" * 120], pool)
    print([len(c) for c in chunks])


if __name__ == "__main__":
    asyncio.run(main())
# Output: [3, 1]

import asyncio


class DynamicBatcher:
    def __init__(self, infer_batch, max_batch: int = 16, max_wait: float = 0.005) -> None:
        self.infer_batch = infer_batch            # sync function: list[input] -> list[output]
        self.max_batch, self.max_wait = max_batch, max_wait
        self.queue: asyncio.Queue = asyncio.Queue(maxsize=max_batch * 8)
        self.batches_run = 0

    async def submit(self, x):
        fut = asyncio.get_running_loop().create_future()
        await self.queue.put((x, fut))            # backpressure when the GPU falls behind
        return await fut

    async def run(self) -> None:
        loop = asyncio.get_running_loop()
        while True:
            items = [await self.queue.get()]
            deadline = loop.time() + self.max_wait
            while len(items) < self.max_batch:
                timeout = deadline - loop.time()
                if timeout <= 0:
                    break
                try:
                    async with asyncio.timeout(timeout):
                        items.append(await self.queue.get())
                except TimeoutError:
                    break
            inputs = [x for x, _ in items]
            try:
                outputs = await asyncio.to_thread(self.infer_batch, inputs)   # GPU work off the loop
            except Exception as err:
                for _, fut in items:
                    if not fut.done():
                        fut.set_exception(err)
            else:
                for (_, fut), out in zip(items, outputs):
                    if not fut.done():                                     # caller may have been cancelled
                        fut.set_result(out)
            self.batches_run += 1


def fake_model(batch: list[str]) -> list[int]:
    return [len(s) for s in batch]


async def main() -> None:
    batcher = DynamicBatcher(fake_model, max_batch=16)
    worker = asyncio.create_task(batcher.run())
    results = await asyncio.gather(*(batcher.submit("x" * i) for i in range(64)))
    worker.cancel()
    try:
        await worker
    except asyncio.CancelledError:
        pass
    print(results[:5], "batches:", batcher.batches_run)


asyncio.run(main())
# Output: [0, 1, 2, 3, 4] batches: 4