ai-resume/corpus/measure.py

"""Qdrant memory at 500 / 5,000 / 20,000 chunks.

`docs/phase-4-corpus.md` requires the measurement and `reference/stack.md` is
where the number is recorded. The corpus itself is small; the point of the
measurement is to show the *shape* of the growth so the 400 m cap is a known
quantity rather than a hope.

Method:

1. Create an empty collection and sample memory — the baseline.
2. For each size: recreate the collection, upsert that many chunks with payloads
   shaped exactly like production payloads (same fields, same text lengths),
   wait for the write to settle, then sample.

Memory is read from Qdrant's own Prometheus `app_process_resident_memory_bytes`,
which is the process that would be capped, not the container sandbox. When
`--container` is also given the container's counter is recorded next to it.

The measurement collection is dropped afterwards unless `--keep` is passed, and
it is never the production collection name.
"""

from __future__ import annotations

import hashlib
import json
import re
import statistics
import subprocess
import time
import uuid

from . import chunker, embeddings, loader, payloads
from .qdrant import Qdrant

# Qdrant 1.20 exposes memory_resident_bytes; older builds used the Prometheus
# default app_process_resident_memory_bytes. Accept either.
RSS_METRIC = re.compile(
    r"^(?:memory_resident_bytes|app_process_resident_memory_bytes)\s+([0-9.eE+]+)",
    re.MULTILINE,
)
INGEST_BATCH = 250
INGEST_SETTLE = 0.4
MEMUSAGE = re.compile(r"([0-9.]+)\s*([KMG]i?B)", re.IGNORECASE)
CAP_MIB = 400  # reference/stack.md memory budget for Qdrant


def measure(cfg, sizes=(500, 5000, 20000), container: str = "", keep: bool = False, log=None) -> dict:
    log = log or (lambda _msg: None)
    collection = f"{cfg.collection}_measure"
    provider = embeddings.build_provider(cfg)
    qdrant = Qdrant(cfg.qdrant_url, timeout=max(cfg.timeout, 60.0))
    if not qdrant.health():
        raise RuntimeError(f"Qdrant is not answering at {cfg.qdrant_url}")

    templates = _templates(cfg, provider)
    log(f"measure: {len(templates)} payload template(s), dim={cfg.dim}, collection={collection}")

    qdrant.drop_collection(collection)
    _create(qdrant, cfg, collection)
    time.sleep(1.0)
    baseline = _sample(qdrant, container)

    samples = []
    for size in sorted(int(value) for value in sizes):
        qdrant.drop_collection(collection)
        _create(qdrant, cfg, collection)
        points = _points(templates, size, provider)
        started = time.time()
        # Throttled ingest: the peak during a bulk write is what trips the
        # cgroup limit, not the settled size. One batch at a time, sampling as we
        # go, so the measurement reports the peak and not just the end state.
        sent = 0
        peak = baseline["rss_mib"]
        for start in range(0, len(points), INGEST_BATCH):
            sent += qdrant.upsert(collection, points[start : start + INGEST_BATCH], batch=INGEST_BATCH)
            time.sleep(INGEST_SETTLE)
            reading = _rss_mib(qdrant)
            if reading is not None:
                peak = max(peak, reading)
        elapsed = round(time.time() - started, 2)
        points_now = qdrant.point_count(collection)
        time.sleep(2.0)
        reading = _sample(qdrant, container)
        growth = reading["rss_mib"] - baseline["rss_mib"]
        samples.append(
            {
                "chunks_requested": size,
                "points": points_now,
                "sent": sent,
                "upsert_seconds": elapsed,
                "rss_mib": reading["rss_mib"],
                "peak_rss_mib": round(peak, 1),
                "container_mib": reading["container_mib"],
                "growth_mib": round(growth, 1),
                "bytes_per_chunk": int(growth * 1024 * 1024 / size) if size else None,
            }
        )
        log(
            f"measure: {size} chunk(s) -> {points_now} point(s), rss {reading['rss_mib']:.1f} MiB "
            f"(+{growth:.1f} MiB), upsert {elapsed}s"
        )

    if not keep:
        qdrant.drop_collection(collection)

    report = {
        "collection": collection,
        "dim": cfg.dim,
        "distance": "Cosine",
        "provider": cfg.provider,
        "model": cfg.model,
        "qdrant_version": qdrant.version(),
        "cap_mib": CAP_MIB,
        "baseline": baseline,
        "samples": samples,
        "within_cap": all(
            max(sample["rss_mib"], sample["peak_rss_mib"]) <= CAP_MIB for sample in samples
        ),
        "measured_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "kept": keep,
    }
    return report


def _create(qdrant: Qdrant, cfg, collection: str) -> None:
    """Same collection settings as production, so the memory reading means something."""
    qdrant.create_collection(
        collection, cfg.dim, vectors_on_disk=cfg.vectors_on_disk
    )


def _templates(cfg, provider) -> list:
    """Real payloads from content/, so field count and text length are realistic."""
    documents = loader.load_documents(cfg)
    out = []
    for doc in documents:
        for chunk in chunker.chunk_document(doc.body):
            payload = payloads.build_payload(doc, chunk, cfg)
            payload["id"] = "template"
            out.append(payload)
    if not out:
        raise RuntimeError("no content to build a measurement from")
    return out


def _points(templates: list, size: int, provider) -> list:
    """`size` points with unique ids, unique source paths and real payload shape."""
    points = []
    texts = []
    for index in range(size):
        template = templates[index % len(templates)]
        payload = dict(template)
        payload["id"] = str(uuid.UUID(hex=hashlib.sha256(f"measure:{index}".encode()).hexdigest()[:32]))
        payload["source_path"] = f"measure/{template['area']}/{index}.md"
        payload["chunk_index"] = index
        points.append({"id": payload["id"], "payload": payload})
        texts.append(f"{index} {payload['text']}")
    vectors = provider.embed(texts)
    for point, vector in zip(points, vectors):
        point["vector"] = vector
    return points


def _sample(qdrant: Qdrant, container: str) -> dict:
    """Three readings two seconds apart; report the median so a GC pause is not the number."""
    readings = []
    for attempt in range(3):
        readings.append(_rss_mib(qdrant))
        if attempt < 2:
            time.sleep(2.0)
    readings = [value for value in readings if value is not None]
    rss = round(statistics.median(readings), 1) if readings else 0.0
    return {"rss_mib": rss, "container_mib": _container_mib(container) if container else None}


def _rss_mib(qdrant: Qdrant):
    text = qdrant.metrics()
    if not text:
        return None
    match = RSS_METRIC.search(text)
    if not match:
        return None
    return float(match.group(1)) / (1024 * 1024)


def _container_mib(container: str):
    """Container memory from `docker stats`. None when docker or the container is absent."""
    try:
        completed = subprocess.run(
            ["docker", "stats", "--no-stream", "--format", "{{.MemUsage}}", container],
            capture_output=True,
            text=True,
            timeout=20,
            check=False,
        )
    except (OSError, subprocess.SubprocessError):
        return None
    if completed.returncode != 0:
        return None
    match = MEMUSAGE.search(completed.stdout)
    if not match:
        return None
    value = float(match.group(1))
    unit = match.group(2).upper()
    if unit.startswith("G"):
        value *= 1024
    elif unit.startswith("K"):
        value /= 1024
    return round(value, 1)


def render_text(report: dict) -> str:
    lines = [
        f"Qdrant {report['qdrant_version']} · collection {report['collection']} · "
        f"dim {report['dim']} cosine · provider {report['provider']}/{report['model']}",
        f"baseline (empty collection): {report['baseline']['rss_mib']:.1f} MiB",
        "",
        f"{'chunks':>8}  {'points':>8}  {'rss MiB':>8}  {'peak':>7}  {'growth':>7}  {'B/chunk':>8}  {'upsert s':>8}",
    ]
    for sample in report["samples"]:
        per_chunk = sample["bytes_per_chunk"] or 0
        lines.append(
            f"{sample['chunks_requested']:>8}  {sample['points']:>8}  {sample['rss_mib']:>8.1f}  "
            f"{sample['peak_rss_mib']:>7.1f}  {sample['growth_mib']:>+7.1f}  {per_chunk:>8}  "
            f"{sample['upsert_seconds']:>8}"
        )
    lines.append("")
    verdict = "within" if report["within_cap"] else "OVER"
    lines.append(f"cap {report['cap_mib']} MiB: {verdict} cap at every measured size")
    if report["kept"]:
        lines.append(f"collection {report['collection']} was kept for inspection")
    return "\n".join(lines)


if __name__ == "__main__":
    import sys

    from .config import Config

    cfg = Config.from_env(dim=int(sys.argv[1]) if len(sys.argv) > 1 else None)
    print(render_text(measure(cfg, container="", log=print)))