AI-Engineering Context Manager Recipes

import torch


@torch.inference_mode()
def embed_batch(model: torch.nn.Module, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
    with torch.autocast(device_type=input_ids.device.type, dtype=torch.bfloat16):
        out = model(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state
    mask = attention_mask.unsqueeze(-1).to(out.dtype)
    return (out * mask).sum(1) / mask.sum(1).clamp(min=1)       # mean pooling

torch.no_grad()                  # BUG: creates the object, enters nothing (ctx_tool CTX012)
logits = model(x)

with torch.no_grad:              # BUG: the class, not an instance -> TypeError;
    ...                          #      always write torch.no_grad()

with torch.inference_mode():
    feats = backbone(x)
loss = head(feats).sum()
loss.backward()                  # RuntimeError: inference tensors cannot be saved for backward
                                 # -> use no_grad() when outputs feed training later

scaler = torch.amp.GradScaler("cuda")
for batch in loader:
    optimizer.zero_grad(set_to_none=True)
    with torch.autocast(device_type="cuda", dtype=torch.float16):
        loss = model(**batch).loss
    scaler.scale(loss).backward()          # backward OUTSIDE autocast
    scaler.step(optimizer)
    scaler.update()

from contextlib import contextmanager
from typing import Iterator

import torch


@contextmanager
def evaluating(model: torch.nn.Module) -> Iterator[torch.nn.Module]:
    """eval() + no_grad() for the block, restoring each submodule's original mode."""
    modes = {m: m.training for m in model.modules()}   # submodules may differ (frozen BN)
    model.eval()
    try:
        with torch.no_grad():
            yield model
    finally:
        for m, was_training in modes.items():
            m.train(was_training)


for step, batch in enumerate(train_loader):
    train_step(model, batch)
    if step % 500 == 0:
        with evaluating(model) as m:
            val_loss = sum(m(**b).loss.item() for b in val_loader) / len(val_loader)

import gc
from contextlib import contextmanager
from typing import Dict, Iterator

import torch


@contextmanager
def cuda_memory(label: str, device: int = 0) -> Iterator[Dict[str, float]]:
    """Report peak memory of the block; free cached blocks afterwards."""
    stats: Dict[str, float] = {}
    if not torch.cuda.is_available():
        yield stats
        return
    torch.cuda.synchronize(device)
    torch.cuda.reset_peak_memory_stats(device)
    start = torch.cuda.memory_allocated(device)
    try:
        yield stats
    finally:
        torch.cuda.synchronize(device)
        stats["peak_mib"] = (torch.cuda.max_memory_allocated(device) - start) / 2**20
        gc.collect()
        torch.cuda.empty_cache()
        print(f"[{label}] peak +{stats['peak_mib']:.0f} MiB")


with cuda_memory("generate bs=32") as mem:
    out = model.generate(**batch, max_new_tokens=256)

def generate_with_backoff(model, prompts, batch_size: int = 64):
    while batch_size >= 1:
        try:
            with torch.inference_mode():
                return [model.generate(**tok(prompts[i:i + batch_size])) for i in range(0, len(prompts), batch_size)]
        except torch.cuda.OutOfMemoryError:
            batch_size //= 2
            gc.collect()
            torch.cuda.empty_cache()     # the traceback has been released at this point
    raise RuntimeError("cannot fit even batch_size=1")

import time
from contextlib import contextmanager
from dataclasses import dataclass, field
from typing import Dict, Iterator


@dataclass
class Usage:
    prompt_tokens: int = 0
    completion_tokens: int = 0
    calls: int = 0
    seconds: float = 0.0
    cost_usd: float = 0.0

    def add(self, resp, price_in: float, price_out: float) -> None:
        u = resp.usage
        self.prompt_tokens += u.prompt_tokens
        self.completion_tokens += u.completion_tokens
        self.calls += 1
        self.cost_usd += u.prompt_tokens * price_in + u.completion_tokens * price_out


@contextmanager
def metered(budget_usd: float = float("inf")) -> Iterator[Usage]:
    """Collect token usage/cost for a block; raise if the budget was exceeded."""
    usage = Usage()
    t0 = time.perf_counter()
    try:
        yield usage
    finally:
        usage.seconds = time.perf_counter() - t0
    if usage.cost_usd > budget_usd:                     # only on success: don't mask errors
        raise RuntimeError(f"budget exceeded: ${usage.cost_usd:.4f} > ${budget_usd}")


with metered(budget_usd=0.50) as usage:
    for q in eval_questions:
        resp = client.chat.completions.create(model="gpt-4o-mini", messages=[{"role": "user", "content": q}])
        usage.add(resp, price_in=0.15e-6, price_out=0.60e-6)
print(f"{usage.calls} calls, {usage.prompt_tokens}+{usage.completion_tokens} tokens, "
      f"${usage.cost_usd:.4f}, {usage.seconds:.1f}s")

from opentelemetry import trace
from opentelemetry.trace import Status, StatusCode

tracer = trace.get_tracer("rag-service")


def answer(question: str) -> str:
    with tracer.start_as_current_span("rag.answer") as span:      # records exceptions by default
        span.set_attribute("question.length", len(question))
        with tracer.start_as_current_span("retrieve") as s:
            docs = retriever.search(question, k=8)
            s.set_attribute("retrieve.k", 8)
        with tracer.start_as_current_span("llm.generate") as s:
            resp = llm.complete(build_prompt(question, docs))
            s.set_attribute("gen_ai.usage.input_tokens", resp.usage.prompt_tokens)
            s.set_attribute("gen_ai.usage.output_tokens", resp.usage.completion_tokens)
        return resp.text

from unittest import mock

with mock.patch.dict("os.environ", {"OPENAI_BASE_URL": "http://localhost:8000/v1"}):
    client = OpenAI()        # picks up the local vLLM server

import asyncio
import random


async def call_with_retry(limiter, fn, *, attempts: int = 5):
    for attempt in range(attempts):
        async with limiter:                         # slot released before sleeping
            try:
                return await fn()
            except RateLimitError as exc:
                if attempt == attempts - 1:
                    raise
                delay = getattr(exc, "retry_after", None) or min(60, 2 ** attempt + random.random())
        await asyncio.sleep(delay)                  # back off OUTSIDE the semaphore

import threading

LLM_SLOTS = threading.BoundedSemaphore(8)

def classify(text: str) -> str:
    with LLM_SLOTS:
        return llm.complete(f"Classify: {text}")

from contextlib import asynccontextmanager
from typing import AsyncIterator

import gc

import httpx
import torch
from fastapi import FastAPI, Request
from sentence_transformers import SentenceTransformer
from starlette.concurrency import run_in_threadpool


@asynccontextmanager
async def lifespan(app: FastAPI) -> AsyncIterator[None]:
    app.state.embedder = SentenceTransformer("BAAI/bge-small-en-v1.5", device="cuda")
    app.state.http = httpx.AsyncClient(timeout=httpx.Timeout(60, connect=5))
    try:
        yield
    finally:
        await app.state.http.aclose()
        del app.state.embedder
        gc.collect()
        torch.cuda.empty_cache()


app = FastAPI(lifespan=lifespan)


@app.post("/embed")
async def embed(req: Request, texts: list[str]) -> list[list[float]]:
    model = req.app.state.embedder
    # encode() is blocking CPU/GPU work: run it off the event loop
    vectors = await run_in_threadpool(model.encode, texts, normalize_embeddings=True)
    return vectors.tolist()

from typing import AsyncIterator

from fastapi import Depends


async def get_db() -> AsyncIterator["AsyncSession"]:
    async with SessionLocal() as session:          # SQLAlchemy AsyncSession
        try:
            yield session
            await session.commit()
        except BaseException:
            await session.rollback()
            raise


@app.post("/documents")
async def add_document(doc: DocIn, db=Depends(get_db)):
    db.add(Document(**doc.model_dump()))

import sqlite3
from contextlib import closing, contextmanager
from typing import Iterator


@contextmanager
def transaction(conn: sqlite3.Connection) -> Iterator[sqlite3.Connection]:
    conn.execute("BEGIN")
    try:
        yield conn
    except BaseException:
        conn.rollback()
        raise
    else:
        conn.commit()


with closing(sqlite3.connect("rag.db", isolation_level=None)) as conn:   # autocommit; we manage BEGIN
    with transaction(conn):
        conn.execute("DELETE FROM chunks WHERE doc_id = ?", (doc_id,))
        conn.executemany("INSERT INTO chunks(doc_id, ord, text) VALUES (?, ?, ?)", rows)

@contextmanager
def staged_upsert(store, collection: str) -> Iterator[str]:
    """Write into a temp collection; swap alias on success, drop it on failure."""
    staging = f"{collection}__staging"
    store.create_collection(staging)
    try:
        yield staging
    except BaseException:
        store.delete_collection(staging)
        raise
    store.swap_alias(alias=collection, target=staging)

with atomic_write("checkpoints/step_1200.pt", mode="wb") as f:
    torch.save({"model": model.state_dict(), "optim": optimizer.state_dict(), "step": 1200}, f)

with atomic_write("index/manifest.json") as f:
    json.dump({"embedding_model": "bge-small-en-v1.5", "dim": 384, "count": n}, f)

rng = np.random.default_rng(1234)                 # no global state to restore
g = torch.Generator().manual_seed(1234)
loader = DataLoader(ds, shuffle=True, generator=g)

import os
import shutil
import subprocess
import sys
import tempfile
from contextlib import contextmanager
from pathlib import Path
from typing import Iterator


@contextmanager
def sandbox(files: dict[str, str], keep_on_error: bool = False) -> Iterator[Path]:
    """Materialise agent-written files in a fresh dir; delete afterwards."""
    root = Path(tempfile.mkdtemp(prefix="agent-"))
    ok = False
    try:
        for rel, content in files.items():
            p = (root / rel).resolve()
            if root.resolve() not in p.parents:            # block ../../ path traversal
                raise ValueError(f"path escapes sandbox: {rel}")
            p.parent.mkdir(parents=True, exist_ok=True)
            p.write_text(content, encoding="utf-8")
        yield root
        ok = True
    finally:
        if ok or not keep_on_error:
            shutil.rmtree(root, ignore_errors=True)


with sandbox({"main.py": code, "test_main.py": tests}) as root:
    proc = subprocess.run(
        [sys.executable, "-m", "pytest", "-q"], cwd=root,
        capture_output=True, text=True, timeout=60,
    )

@contextmanager
def frozen(module: torch.nn.Module) -> Iterator[torch.nn.Module]:
    flags = [p.requires_grad for p in module.parameters()]
    for p in module.parameters():
        p.requires_grad_(False)
    try:
        yield module
    finally:
        for p, flag in zip(module.parameters(), flags):
            p.requires_grad_(flag)


with frozen(model.encoder):
    train_head_for_one_epoch(model)

from contextlib import contextmanager
from dataclasses import replace
from typing import Any, Iterator


@contextmanager
def overridden(obj: Any, **changes: Any) -> Iterator[Any]:
    """Set attributes for the block and restore the originals (also for missing ones)."""
    missing = object()
    saved = {k: getattr(obj, k, missing) for k in changes}
    try:
        for k, v in changes.items():
            setattr(obj, k, v)
        yield obj
    finally:
        for k, v in saved.items():
            if v is missing:
                delattr(obj, k)
            else:
                setattr(obj, k, v)


with overridden(gen_config, temperature=0.0, max_new_tokens=64):
    score = run_eval(model, gen_config)