From 7a970513f739b62f15c3b8dfa0b1e27e1c644d47 Mon Sep 17 00:00:00 2001 From: Advay Monga Date: Thu, 8 Oct 2026 01:08:57 -0400 Subject: [PATCH 1/7] profile bundle: window-only measurements, corpus workloads, op/kernel tables - warmup sized to the window (default = window requests, fired together, random ids at the window's prompt lengths so no prefix-cache hits); settings recorded in meta - meta: profiler_on + profiler config, GPU snapshot at window start and end, timeline events written/dropped, per-request outcomes (errors, empty, short, full) and token counts; empty responses count as failed - memory.json: CUDA peaks reset at window start; window_peak and lifetime_peak named apart - stats.json: scheduler stats at window start/end plus counter deltas - gpu: throttle reasons (clocks_event_reasons / clocks_throttle_reasons fallback, null if unsupported), application clocks; gpu.csv gets an epoch_s column - ops.json from key_averages, kernels.json for CUDA kernels by name - lab.profile takes --corpus-class/--workload/--prompts/--synthetic (one required); the profile tool samples a seen corpus class outside the jail (lab.corpus.sample) Co-Authored-By: Claude Opus 5.5 --- CLAUDE.md | 2 +- lab/README.md | 27 ++++++--- lab/bundle.py | 45 +++++++++++++-- lab/corpus.py | 10 ++++ lab/gpu.py | 67 ++++++++++++++++++++--- lab/profile.py | 112 +++++++++++++++++++++++++++++++------- lab/tools.py | 30 +++++++--- tests/test_lab_corpus.py | 13 +++++ tests/test_lab_gpu.py | 82 ++++++++++++++++++++++++++++ tests/test_lab_profile.py | 87 ++++++++++++++++++++++++++++- 10 files changed, 424 insertions(+), 51 deletions(-) create mode 100644 tests/test_lab_gpu.py diff --git a/CLAUDE.md b/CLAUDE.md index 4b76509..dc9de87 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -34,7 +34,7 @@ pytest/ruff config). Each folder with commands has its own `python -m`. (append-only JSONL, `lab/ledger/` gitignored), `budget.py`, `engine.py` (where the engine repo and its python are), `safety/` (write surfaces, srt jail, grader: lint/tests run with the engine's python in the jail), `profile.py` + `bundle.py` + `gpu.py` (profiler harness, staged into the pristine tree under `_harness/` so the - jail never opens this repo, which holds `.env` and the ledger), `canary.py` (honest regressions patched into + jail never opens this repo, which holds `.env` and the ledger; the tool samples its seen-split corpus workload out here, `corpus.sample`, and stages it as `_harness/workload.json`; window-only memory peaks and scheduler counters, warmup sized to the window, op/kernel tables), `canary.py` (honest regressions patched into the engine's process), `vm.py` + `providers/` + `vm-setup.sh` (one GPU VM on Verda/Nebius/Crusoe; pushes both trees, engine to `LAB_VM_DIR`, this repo to `LAB_VM_ENV_DIR`; `run --env` runs here; `run` stops the VM when done unless `--keep`, and every session is capped by `LAB_VM_MAX_MINUTES` via `timeout` + a local watchdog + in-VM diff --git a/lab/README.md b/lab/README.md index 57df716..bcc8e5e 100644 --- a/lab/README.md +++ b/lab/README.md @@ -32,18 +32,29 @@ copy it, change the engine section (and the reference if the model changes), and ## Tools - python -m lab.profile --backend custom-mps --requests 8 --max-tokens 32 --out lab/runs + python -m lab.profile --backend custom-mps --corpus-class steady_interactive --requests 8 --max-tokens 32 --out lab/runs The engine configuration comes from the same env vars the server reads (`MAX_BATCH_SIZE`, `PREFILL_MODE`, ...), so a profile measures what is served; `--backend` and `--model` override -`BACKEND` and `MODEL_NAME`. Writes one bundle directory per run, raw files only: +`BACKEND` and `MODEL_NAME`. The workload is explicit, one of: `--corpus-class NAME` (`--split`, `--seed`, +`--requests`: sampled through the verified loader, chat-templated with the engine's tokenizer, thinking off +unless `--enable-thinking`), `--workload FILE` (the same, pre-sampled by `lab.corpus.sample`; the `profile` +tool uses it because the jail holds no corpus), `--prompts FILE` or `--synthetic` (random ids of `--prompt-len`). +Output length is `--max-tokens` for every request. Before the window, `--warmup` requests (default: as many as +the window) at the window's prompt lengths with random ids are fired together, so compiles and graph captures for +the window's batch sizes happen outside it. Writes one bundle directory per run, raw files only: events.jsonl engine event timeline: scheduler decisions and phases, per step trace.json torch.profiler chrome trace; phase ranges carry the step id - memory.json device allocator stats and peak host RSS - gpu.csv nvidia-smi samples at 100 ms (CUDA hosts only) - stats.json the scheduler's own counters at the end of the run - meta.json git sha, torch, device, clock state, engine settings, workload hash, window + ops.json prof.key_averages() by self device then self CPU time: top 200 and totals (us) + kernels.json CUDA kernels by name: count, total, mean, max (us); only when there was CUDA activity + memory.json window peak (reset at window start) and lifetime peak (CUDA), allocator snapshots at window + start and end, peak host RSS (lifetime) + gpu.csv nvidia-smi samples at 100 ms with an `epoch_s` column, throttle reasons when the driver has them (CUDA hosts only) + stats.json the scheduler's counters at window start and end, and their difference + meta.json git sha, torch, device, engine settings, workload source and hash, window, warmup, profiler config + (`profiler_on`: every number was taken under the profiler), GPU clock state and throttle reasons at + window start and end, timeline events written/dropped, per-request outcomes and token counts ## Ledger @@ -79,7 +90,7 @@ with the referee's rights, snapshot the workspace and write the ledger on every | tool | does | |---|---| | `test` | lint and the fast suite on a pristine two-commit copy of the workspace, jailed | -| `profile` | `lab.profile` on the pristine copy, jailed; the bundle goes into the ledger as a blob | +| `profile` | `lab.profile` on the pristine copy, jailed, on `args.corpus_class` (default `steady_interactive`, seen split, `args.seed`) or `args.synthetic`; the bundle goes into the ledger as a blob | | `ledger` | read records (this run and earlier ones) | | `budget` | dollars left | | `restore` | workspace back to a snapshot id (`base` resets) | @@ -122,7 +133,7 @@ commands either way: python -m lab.vm start # create (first time) or start, wait for ssh python -m lab.vm setup # lab/vm-setup.sh: venv, CUDA torch, Nsight, counter and clock checks python -m lab.vm run --env --fetch lab/runs -- \ - env BACKEND=custom-cuda python -m lab.profile --requests 8 # --env: run in this repo's tree + env BACKEND=custom-cuda python -m lab.profile --corpus-class steady_interactive --requests 8 # --env: run in this repo's tree python -m lab.vm run -- python scripts/gpu_tests/checks.py # default: run in the engine's tree python -m lab.vm stop # ends all billing (Verda: deletes the VM and its disk) diff --git a/lab/bundle.py b/lab/bundle.py index 28b36b1..bf6cae4 100644 --- a/lab/bundle.py +++ b/lab/bundle.py @@ -16,10 +16,9 @@ from typing import Any import torch +from torch.autograd import DeviceType -from lab import gpu - -FILES = ("events.jsonl", "trace.json", "memory.json", "stats.json", "meta.json") +FILES = ("events.jsonl", "trace.json", "ops.json", "memory.json", "stats.json", "meta.json") def new_dir(root: str | Path) -> Path: @@ -61,11 +60,49 @@ def device_memory(device: str) -> dict: return out +def device_peaks(device: str) -> dict | None: + """Allocator high-water marks since the last reset; CUDA only (MPS keeps no peak).""" + if not (device.startswith("cuda") and torch.cuda.is_available()): + return None + s = torch.cuda.memory_stats() + return {"allocated_bytes": s.get("allocated_bytes.all.peak", 0), "reserved_bytes": s.get("reserved_bytes.all.peak", 0)} + + +def reset_peaks(device: str) -> None: + if device.startswith("cuda") and torch.cuda.is_available(): + torch.cuda.reset_peak_memory_stats() + + +def ops_table(prof: Any, top: int = 200) -> dict: + """prof.key_averages() as rows (times in us), sorted by self device time then self CPU time; top N plus totals.""" + rows = [{"name": a.key, "count": a.count, "self_cpu_us": a.self_cpu_time_total, "cpu_us": a.cpu_time_total, + "self_device_us": a.self_device_time_total, "device_us": a.device_time_total} + for a in prof.key_averages()] + rows.sort(key=lambda r: (r["self_device_us"], r["self_cpu_us"]), reverse=True) + return {"sort": "self_device_us, self_cpu_us", "ops_total": len(rows), "top": rows[:top], + "totals": {"self_cpu_us": sum(r["self_cpu_us"] for r in rows), + "self_device_us": sum(r["self_device_us"] for r in rows)}} + + +def kernels_table(events: Any) -> dict | None: + """CUDA kernels aggregated by name (times in us), by total time; None when the trace has no CUDA activity.""" + agg: dict[str, list[float]] = {} + for e in events: + if e.device_type == DeviceType.CUDA: + agg.setdefault(e.name, []).append(e.time_range.elapsed_us()) + if not agg: + return None + rows = [{"name": k, "count": len(v), "total_us": sum(v), "mean_us": sum(v) / len(v), "max_us": max(v)} + for k, v in agg.items()] + rows.sort(key=lambda r: r["total_us"], reverse=True) + return {"kernels": rows, "total_us": sum(r["total_us"] for r in rows), "launches": sum(r["count"] for r in rows)} + + def meta(device: str, settings: Any, prompts: list[list[int]], max_tokens: int, t0: float, t1: float, **extra: Any) -> dict: """A number is a fact about a config, so the engine settings travel with every bundle.""" return {"git_sha": git_sha(), "torch": torch.__version__, "python": sys.version.split()[0], - "platform": platform.platform(), "device": device, "gpu": gpu.query(), + "platform": platform.platform(), "device": device, "settings": asdict(settings), "workload_hash": workload_hash(prompts, max_tokens), "requests": len(prompts), "max_tokens": max_tokens, "window": {"start": t0, "end": t1, "wall_s": t1 - t0}, diff --git a/lab/corpus.py b/lab/corpus.py index 0f3c207..33196e5 100644 --- a/lab/corpus.py +++ b/lab/corpus.py @@ -13,6 +13,7 @@ import hashlib import json +import random from dataclasses import asdict, dataclass, field from pathlib import Path from typing import Any @@ -186,3 +187,12 @@ def load_trace(class_name: str, split: str, if class_name not in m.classes: raise CorpusError(f"unknown workload class {class_name!r}; have {sorted(m.classes)}") return m, read_trace(corpus_dir / m.classes[class_name].trace_file(split)) + + +def sample(class_name: str, split: str, seed: int, n: int, corpus_dir: Path = CORPUS_DIR) -> dict[str, Any]: + """`n` requests of one (class, split) drawn without replacement by `seed`: chat messages plus their provenance.""" + m, reqs = load_trace(class_name, split, corpus_dir) + lines = random.Random(seed).sample(range(len(reqs)), n) + return {"source": {"kind": "corpus", "class": class_name, "split": split, "seed": seed, + "corpus_version": m.corpus_version, "lines": lines}, + "messages": [reqs[i].chat_messages() for i in lines]} diff --git a/lab/gpu.py b/lab/gpu.py index 7871979..8e48977 100644 --- a/lab/gpu.py +++ b/lab/gpu.py @@ -2,8 +2,12 @@ from __future__ import annotations +import csv +import functools +import io import shutil import subprocess +from datetime import datetime from pathlib import Path QUERY = ("name", "driver_version", "clocks.sm", "clocks.mem", "clocks.max.sm", "clocks.max.mem", @@ -11,36 +15,79 @@ "memory.used", "memory.total") SAMPLE = ("timestamp", "utilization.gpu", "utilization.memory", "clocks.sm", "clocks.mem", "power.draw", "temperature.gpu", "memory.used") +REASONS = ("clocks_event_reasons.active", "clocks_throttle_reasons.active") # the field's name on newer, older drivers +APP_CLOCKS = ("clocks.applications.graphics", "clocks.applications.memory") # locked clocks are not queryable here +MISSING = ("[N/A]", "[Not Supported]", "N/A") def available() -> bool: return shutil.which("nvidia-smi") is not None +def _smi(fields: tuple[str, ...]) -> str | None: + out = subprocess.run(["nvidia-smi", f"--query-gpu={','.join(fields)}", "--format=csv,noheader,nounits", + "-i", "0"], capture_output=True, text=True) + return out.stdout if out.returncode == 0 else None + + +@functools.lru_cache(maxsize=None) +def supported(field: str) -> bool: + """Whether this driver accepts `field`: one unknown field fails the whole query.""" + return available() and _smi((field,)) is not None + + +def reasons_field() -> str | None: + return next((f for f in REASONS if supported(f)), None) + + def query() -> dict | None: - """One snapshot of the first GPU: identity, clocks and limits. The clock state that travels with a number.""" + """One snapshot of the first GPU: identity, clocks, limits, throttle reasons; None for what the driver lacks.""" if not available(): return None - out = subprocess.run(["nvidia-smi", f"--query-gpu={','.join(QUERY)}", "--format=csv,noheader,nounits", - "-i", "0"], capture_output=True, text=True) - if out.returncode != 0: + reasons = reasons_field() + extra = tuple(f for f in ((reasons,) if reasons else ()) + APP_CLOCKS if supported(f)) + out = _smi(QUERY + extra) + if out is None: return None - return dict(zip(QUERY, [v.strip() for v in out.stdout.strip().split(",")])) + got = {k: (None if v in MISSING else v) + for k, v in zip(QUERY + extra, [v.strip() for v in out.strip().split(",")])} + snap = {k: got.get(k) for k in QUERY + APP_CLOCKS} + snap[REASONS[0]] = got.get(reasons) if reasons else None + return snap + + +def add_epoch(raw: str) -> str: + """nvidia-smi's local-time `timestamp` column, prefixed with `epoch_s` so samples line up with timeline events.""" + buf = io.StringIO() + w = csv.writer(buf, lineterminator="\n") + for i, row in enumerate(r for r in csv.reader(io.StringIO(raw), skipinitialspace=True) if r): + if i == 0: + w.writerow(["epoch_s", *row]) + continue + try: + epoch = f"{datetime.strptime(row[0], '%Y/%m/%d %H:%M:%S.%f').timestamp():.3f}" + except ValueError: + epoch = "" + w.writerow([epoch, *row]) + return buf.getvalue() class Sampler: - """nvidia-smi sampling into a CSV at `interval_ms`, as a subprocess so it costs the engine nothing.""" + """nvidia-smi sampling at `interval_ms`, as a subprocess so it costs the engine nothing; `epoch_s` added on stop.""" def __init__(self, path: Path, interval_ms: int = 100): self.path, self.interval_ms, self._proc = Path(path), interval_ms, None + self._raw = self.path.with_suffix(".smi.csv") def start(self) -> bool: if not available(): return False - self._file = open(self.path, "w") + reasons = reasons_field() + self._file = open(self._raw, "w") self._proc = subprocess.Popen( - ["nvidia-smi", f"--query-gpu={','.join(SAMPLE)}", "--format=csv,nounits", - "-lms", str(self.interval_ms), "-i", "0"], stdout=self._file, stderr=subprocess.DEVNULL) + ["nvidia-smi", f"--query-gpu={','.join(SAMPLE + ((reasons,) if reasons else ()))}", + "--format=csv,nounits", "-lms", str(self.interval_ms), "-i", "0"], + stdout=self._file, stderr=subprocess.DEVNULL) return True def stop(self) -> None: @@ -56,3 +103,5 @@ def stop(self) -> None: finally: self._file.close() self._proc = None + self.path.write_text(add_epoch(self._raw.read_text())) + self._raw.unlink() diff --git a/lab/profile.py b/lab/profile.py index baedf49..3b47391 100644 --- a/lab/profile.py +++ b/lab/profile.py @@ -24,24 +24,35 @@ logger = logging.getLogger(__name__) +COUNTERS = ("total_admitted", "total_completed", "total_rejected", "total_preempted", "total_expired", + "total_iteration_errors", "decode_steps", "kv_admit_blocked", "prefill_chunks_processed") + + def synthetic_prompts(n: int, length: int, vocab: int = 1000, seed: int = 0) -> list[list[int]]: rng = random.Random(seed) return [[rng.randrange(1, vocab) for _ in range(length)] for _ in range(n)] -def _profiler(device: str) -> torch.profiler.profile: +def warmup_prompts(prompts: list[list[int]], n: int, vocab: int = 1000, seed: int = 1) -> list[list[int]]: + """Random ids at the window's prompt lengths: the same shapes get compiled, nothing lands in the prefix cache.""" + rng = random.Random(seed) + return [[rng.randrange(1, vocab) for _ in prompts[i % len(prompts)]] for i in range(n)] + + +def _profiler(device: str) -> tuple[torch.profiler.profile, dict]: """All threads: by default the profiler records only the thread that opened it, and the scheduler has its own.""" acts = [ProfilerActivity.CPU] if device.startswith("cuda"): acts.append(ProfilerActivity.CUDA) + config = {"activities": [a.name for a in acts], "all_threads": True} try: from torch._C._profiler import _ExperimentalConfig - config = _ExperimentalConfig(profile_all_threads=True) + exp = _ExperimentalConfig(profile_all_threads=True) except (ImportError, TypeError): # older torch: the trace will miss the scheduler thread logger.warning("torch %s cannot profile all threads; trace covers the main thread only", torch.__version__) - return torch.profiler.profile(activities=acts) - return torch.profiler.profile(activities=acts, experimental_config=config) + return torch.profiler.profile(activities=acts), {**config, "all_threads": False} + return torch.profiler.profile(activities=acts, experimental_config=exp), config async def _fire(sched: ContinuousBatchScheduler, prompts: list[list[int]], max_tokens: int, @@ -52,34 +63,75 @@ async def _fire(sched: ContinuousBatchScheduler, prompts: list[list[int]], max_t return await asyncio.gather(*(sched.submit(r) for r in reqs), return_exceptions=True) +def stats_delta(start: dict, end: dict) -> dict: + """What the window itself did: counter differences, wave sizes included.""" + d = {k: end[k] - start[k] for k in COUNTERS if k in start and k in end} + waves = {k: v - start.get("wave_sizes", {}).get(k, 0) for k, v in end.get("wave_sizes", {}).items()} + d["wave_sizes"] = {k: v for k, v in waves.items() if v} + return d + + +def outcomes(results: list, max_tokens: int) -> dict: + """Per-request outcome counts and output lengths; `short` stopped before max_tokens (EOS or cut: not recorded).""" + toks = [len(r) for r in results if not isinstance(r, BaseException)] + return {"errors": len(results) - len(toks), "empty": toks.count(0), + "short": sum(0 < n < max_tokens for n in toks), "full": toks.count(max_tokens), + "tokens_out": ({"total": sum(toks), "min": min(toks), "mean": sum(toks) / len(toks), "max": max(toks)} + if toks else None)} + + async def run(backend: InferenceBackend, settings: Settings, prompts: list[list[int]], - max_tokens: int, out: Path, warmup: int = 1) -> Path: - """Profile `prompts` through the served scheduler config on `backend`; return the bundle directory.""" + max_tokens: int, out: Path, warmup: int | None = None, source: dict | None = None) -> Path: + """Profile `prompts` through the served scheduler config on `backend`; return the bundle directory. + `warmup` requests (default: as many as the window) are fired together first, outside the window.""" out = bundle.new_dir(out) device = backend.device_str + warmup = len(prompts) if warmup is None else warmup timeline = Timeline(out) sched = build_scheduler(backend, settings, timeline=timeline) sched.start() sampler = gpu.Sampler(out / "gpu.csv") + prof, prof_config = _profiler(device) try: + w0 = time.time() if warmup: - await _fire(sched, prompts[:warmup], max_tokens, "warmup") + await _fire(sched, warmup_prompts(prompts, warmup), max_tokens, "warmup") + w1 = time.time() + peak_before = bundle.device_peaks(device) + bundle.reset_peaks(device) + mem0, stats0, gpu0 = bundle.device_memory(device), sched.stats(), gpu.query() timeline.event("profile_window", state="begin") sampler.start() t0 = time.time() - with _profiler(device) as prof: + with prof: results = await _fire(sched, prompts, max_tokens, "profile") t1 = time.time() timeline.event("profile_window", state="end") + mem1, stats1, gpu1 = bundle.device_memory(device), sched.stats(), gpu.query() + peak_window = bundle.device_peaks(device) finally: sampler.stop() await sched.stop() prof.export_chrome_trace(str(out / "trace.json")) - bundle.write_json(out / "memory.json", bundle.device_memory(device)) - bundle.write_json(out / "stats.json", sched.stats()) - failed = sum(isinstance(r, BaseException) for r in results) - bundle.write_json(out / "meta.json", bundle.meta(device, settings, prompts, max_tokens, t0, t1, - failed=failed, warmup=warmup)) + bundle.write_json(out / "ops.json", bundle.ops_table(prof)) + kernels = bundle.kernels_table(prof.events()) + if kernels is not None: + bundle.write_json(out / "kernels.json", kernels) + lifetime = (None if peak_window is None else + {k: max(v, (peak_before or {}).get(k, 0)) for k, v in peak_window.items()}) + bundle.write_json(out / "memory.json", { + "device": device, "window_peak": peak_window, "lifetime_peak": lifetime, + "host_rss_peak_lifetime_bytes": mem1["peak_host_rss_bytes"], + "at_window_start": mem0, "at_window_end": mem1}) + bundle.write_json(out / "stats.json", {"window_start": stats0, "window_end": stats1, + "window_delta": stats_delta(stats0, stats1)}) + done = outcomes(results, max_tokens) + bundle.write_json(out / "meta.json", bundle.meta( + device, settings, prompts, max_tokens, t0, t1, failed=done["errors"] + done["empty"], outcomes=done, + warmup={"requests": warmup, "concurrent": True, "prompts": "random ids at the window's prompt lengths", + "max_tokens": max_tokens, "wall_s": w1 - w0}, + profiler_on=True, profiler=prof_config, gpu={"window_start": gpu0, "window_end": gpu1}, + timeline=timeline.stats(), workload=source)) return out @@ -87,11 +139,19 @@ def main(argv: list[str] | None = None) -> int: ap = argparse.ArgumentParser(description=__doc__.split("(")[0]) ap.add_argument("--backend", help="override BACKEND, e.g. custom-mps") ap.add_argument("--model", help="override MODEL_NAME") - ap.add_argument("--prompts", help="JSON list of prompt strings; synthetic token ids when absent") + src = ap.add_mutually_exclusive_group(required=True) + src.add_argument("--corpus-class", help="prompts drawn from this corpus class (verified loader)") + src.add_argument("--workload", help="JSON from lab.corpus.sample, for a tree without the corpus (the profile tool)") + src.add_argument("--prompts", help="JSON list of prompt strings") + src.add_argument("--synthetic", action="store_true", help="random token ids of --prompt-len") + ap.add_argument("--split", default="seen", choices=("seen", "heldout")) + ap.add_argument("--corpus-dir", help="default: this repo's corpus/") + ap.add_argument("--enable-thinking", action="store_true", help="chat template thinking on, for corpus prompts") + ap.add_argument("--seed", type=int, default=0) ap.add_argument("--requests", type=int, default=8) ap.add_argument("--prompt-len", type=int, default=64) ap.add_argument("--max-tokens", type=int, default=32) - ap.add_argument("--warmup", type=int, default=1) + ap.add_argument("--warmup", type=int, help="warmup requests, fired together; default: as many as the window") ap.add_argument("--out", default="lab/runs") args = ap.parse_args(argv) @@ -100,14 +160,26 @@ def main(argv: list[str] | None = None) -> int: overrides = {k: v for k, v in (("backend_name", args.backend), ("model_name", args.model)) if v} settings = dataclasses.replace(settings, **overrides) backend, _ = build_backend(settings) - if args.prompts: + if args.synthetic: + prompts = synthetic_prompts(args.requests, args.prompt_len, seed=args.seed) + source = {"kind": "synthetic", "prompt_len": args.prompt_len, "seed": args.seed} + else: from inference_server.tokenizer import Tokenizer tok = Tokenizer(settings.model_name, settings.context_window) - prompts = [tok.encode_chat(t) for t in json.loads(Path(args.prompts).read_text())] - else: - prompts = synthetic_prompts(args.requests, args.prompt_len) + if args.prompts: + prompts = [tok.encode_chat(t) for t in json.loads(Path(args.prompts).read_text())] + source = {"kind": "prompts", "file": args.prompts} + else: + if args.workload: + w = json.loads(Path(args.workload).read_text()) + else: + from lab import corpus + w = corpus.sample(args.corpus_class, args.split, args.seed, args.requests, + Path(args.corpus_dir) if args.corpus_dir else corpus.CORPUS_DIR) + prompts = [tok.encode_messages(m, thinking=args.enable_thinking) for m in w["messages"]] + source = {**w["source"], "thinking": args.enable_thinking} out = asyncio.run(run(backend, settings, prompts, args.max_tokens, Path(args.out), - warmup=args.warmup)) + warmup=args.warmup, source=source)) print(out) return 0 diff --git a/lab/tools.py b/lab/tools.py index f4b8ffd..90d9c75 100644 --- a/lab/tools.py +++ b/lab/tools.py @@ -9,7 +9,7 @@ from pathlib import Path from typing import Any -from lab import engine, ledger +from lab import corpus, engine, ledger, target from lab.agent import ToolSpec from lab.evaltools import EvalTools from lab.safety import grader @@ -73,10 +73,20 @@ def profile(self, args: dict) -> str: snap = self._audited("profile", args) tree = self._pristine() out = tree / "lab" / "runs" - stage_harness(tree) - argv = [engine.python(), "-m", "lab.profile", "--requests", str(args.get("requests", 8)), - "--prompt-len", str(args.get("prompt_len", 64)), "--max-tokens", str(args.get("max_tokens", 32)), - "--out", str(out)] + harness = stage_harness(tree) + n = int(args.get("requests", 8)) + argv = [engine.python(), "-m", "lab.profile", "--requests", str(n), + "--max-tokens", str(args.get("max_tokens", 32)), "--out", str(out)] + if args.get("synthetic"): + argv += ["--synthetic", "--prompt-len", str(args.get("prompt_len", 64))] + else: # sampled out here: the jail holds no corpus, and the verified loader needs the held-out split too + t = target.load() + w = corpus.sample(args.get("corpus_class") or DEFAULT_CORPUS_CLASS, "seen", int(args.get("seed", 0)), n, + t.corpus_dir) + (harness / "workload.json").write_text(json.dumps(w)) + argv += ["--workload", str(harness / "workload.json")] + if t.chat_kwargs.get("enable_thinking"): + argv.append("--enable-thinking") t0 = time.monotonic() proc = (self.s.profile_runner or default_profile_runner)(tree, argv) bundles = sorted(out.glob("*")) if out.exists() else [] @@ -124,9 +134,12 @@ def specs(self) -> list[ToolSpec]: return [ ToolSpec("test", "Lint and run the test suite on a pristine copy of your current workspace, " "inside the referee's jail. Returns pass/fail and the tail of the output.", obj, self.test), - ToolSpec("profile", "Run the engine on a synthetic workload under the profiler and store the raw " - "bundle (event timeline, chrome trace, memory, provenance) in the ledger.", - {"type": "object", "properties": {"requests": {"type": "integer"}, "prompt_len": {"type": "integer"}, + ToolSpec("profile", "Run the engine under the profiler on requests sampled from a seen corpus class (or " + "synthetic token ids) and store the raw bundle (event timeline, chrome trace, op and kernel " + "tables, memory, provenance) in the ledger.", + {"type": "object", "properties": {"corpus_class": {"type": "string"}, "seed": {"type": "integer"}, + "synthetic": {"type": "boolean"}, "requests": {"type": "integer"}, + "prompt_len": {"type": "integer"}, "max_tokens": {"type": "integer"}}}, self.profile), ToolSpec("ledger", "Read past records: every tool call in this and earlier runs, with snapshots and results.", {"type": "object", "properties": {"kind": {"type": "string"}, "session": {"type": "string"}, @@ -154,6 +167,7 @@ def specs(self) -> list[ToolSpec]: HARNESS = ("__init__.py", "profile.py", "bundle.py", "gpu.py") # all lab.profile imports from the environment HARNESS_DIR = "_harness" +DEFAULT_CORPUS_CLASS = "steady_interactive" # targets name no corpus class of their own def stage_harness(tree: Path) -> Path: diff --git a/tests/test_lab_corpus.py b/tests/test_lab_corpus.py index 76da4d6..d347ce5 100644 --- a/tests/test_lab_corpus.py +++ b/tests/test_lab_corpus.py @@ -24,6 +24,7 @@ load_manifest, load_trace, read_trace, + sample, write_trace, ) @@ -226,6 +227,18 @@ def test_missing_trace_is_refused(tmp_path): load_manifest(tmp_path) +def test_sample_is_seeded_and_carries_its_provenance(tmp_path): + m = _mini_corpus(tmp_path) + write_trace(tmp_path / "cold_start" / "seen.jsonl", [TraceRequest(0.0, f"s{i}", 0, f"p{i}", 64) for i in range(5)]) + m = build_manifest(m.classes, tmp_path) + (tmp_path / "manifest.json").write_text(json.dumps(m.to_dict(), indent=2) + "\n") + w = sample("cold_start", "seen", 3, 4, tmp_path) + assert w == sample("cold_start", "seen", 3, 4, tmp_path) and len(set(w["source"]["lines"])) == 4 + assert w["source"] == {"kind": "corpus", "class": "cold_start", "split": "seen", "seed": 3, + "corpus_version": m.corpus_version, "lines": w["source"]["lines"]} + assert w["messages"][0] == [{"role": "user", "content": f"p{w['source']['lines'][0]}"}] + + def test_trace_round_trip(tmp_path): reqs = [TraceRequest(0.0, "s-0", 0, "hello", 16), TraceRequest(1.5, "s-0", 1, "hello\n\nhi!\n\nmore", 16, expected_output_hash="ab", diff --git a/tests/test_lab_gpu.py b/tests/test_lab_gpu.py new file mode 100644 index 0000000..51fdc3c --- /dev/null +++ b/tests/test_lab_gpu.py @@ -0,0 +1,82 @@ +"""lab.gpu against a fake nvidia-smi: field fallbacks, unsupported fields as None, epoch column on the samples.""" + +from __future__ import annotations + +import csv +import io +import os +import stat +import sys +import time +from datetime import datetime + +import pytest + +from lab import gpu + +FAKE = """#!{python} +import sys, time +args = sys.argv[1:] +fields = next(a for a in args if a.startswith("--query-gpu=")).split("=", 1)[1].split(",") +if any(f in {bad!r} for f in fields): + sys.exit(6) +vals = {{"clocks_throttle_reasons.active": "0x0000000000000004", "clocks.applications.memory": "[N/A]", + "timestamp": "2026/10/08 01:02:03.250"}} +row = ", ".join(vals.get(f, "1") for f in fields) +if "-lms" in args: + print(", ".join(fields), flush=True) + while True: + print(row, flush=True) + time.sleep(0.01) +print(row) +""" + + +@pytest.fixture +def fake_smi(tmp_path, monkeypatch): + def install(bad=("clocks_event_reasons.active",)): + p = tmp_path / "nvidia-smi" + p.write_text(FAKE.format(python=sys.executable, bad=set(bad))) + p.chmod(p.stat().st_mode | stat.S_IEXEC) + monkeypatch.setenv("PATH", f"{tmp_path}{os.pathsep}{os.environ['PATH']}") + gpu.supported.cache_clear() + yield install + gpu.supported.cache_clear() + + +def test_no_nvidia_smi_is_none(monkeypatch): + monkeypatch.setattr(gpu.shutil, "which", lambda _: None) + assert gpu.query() is None and gpu.reasons_field() is None + assert gpu.Sampler("x.csv").start() is False + + +def test_query_falls_back_to_the_older_throttle_field(fake_smi): + fake_smi() + q = gpu.query() + assert q["clocks_event_reasons.active"] == "0x0000000000000004" and q["name"] == "1" + assert q["clocks.applications.graphics"] == "1" and q["clocks.applications.memory"] is None + + +def test_query_without_any_reasons_field_reports_none(fake_smi): + fake_smi(bad=gpu.REASONS + ("clocks.applications.graphics",)) + q = gpu.query() + assert q["clocks_event_reasons.active"] is None and q["clocks.applications.graphics"] is None + assert q["clocks.sm"] == "1" + + +def test_sampler_writes_epoch_seconds(fake_smi, tmp_path): + fake_smi() + s = gpu.Sampler(tmp_path / "gpu.csv") + assert s.start() + time.sleep(0.3) + s.stop() + rows = list(csv.DictReader(io.StringIO((tmp_path / "gpu.csv").read_text()), skipinitialspace=True)) + assert rows and not (tmp_path / "gpu.smi.csv").exists() + want = datetime(2026, 10, 8, 1, 2, 3, 250000).timestamp() + assert float(rows[0]["epoch_s"]) == pytest.approx(want, abs=1e-3) + assert rows[0]["clocks_throttle_reasons.active"] == "0x0000000000000004" + + +def test_unparseable_timestamp_leaves_epoch_empty(): + out = gpu.add_epoch("timestamp, x\nnot a time, 1\n") + assert out.splitlines() == ["epoch_s,timestamp,x", ",not a time,1"] diff --git a/tests/test_lab_profile.py b/tests/test_lab_profile.py index adabea0..25dff27 100644 --- a/tests/test_lab_profile.py +++ b/tests/test_lab_profile.py @@ -22,6 +22,11 @@ async def test_bundle_is_complete(tmp_path, stub_backend_cls): meta = json.loads((out / "meta.json").read_text()) assert meta["requests"] == 4 and meta["max_tokens"] == 3 and meta["failed"] == 0 + assert meta["outcomes"]["full"] == 4 and meta["outcomes"]["tokens_out"]["total"] == 12 + assert meta["warmup"]["requests"] == 1 and meta["profiler_on"] is True + assert meta["profiler"]["activities"] == ["CPU"] and "all_threads" in meta["profiler"] + assert set(meta["gpu"]) == {"window_start", "window_end"} + assert meta["timeline"]["events_dropped"] == 0 and meta["timeline"]["events_written"] > 0 assert meta["workload_hash"] == bundle.workload_hash(prompts, 3) assert meta["device"] == "cpu" and meta["window"]["wall_s"] > 0 and meta["torch"] assert meta["settings"]["max_batch_size"] == 4 @@ -37,7 +42,87 @@ async def test_bundle_is_complete(tmp_path, stub_backend_cls): finished = [e for e in events if e["kind"] == "finish"] assert len(finished) == 5 # 1 warmup + 4 profiled stats = json.loads((out / "stats.json").read_text()) - assert stats["total_completed"] == 5 and stats["timeline"]["events_dropped"] == 0 + assert stats["window_start"]["total_completed"] == 1 and stats["window_end"]["total_completed"] == 5 + assert stats["window_delta"]["total_completed"] == 4 and stats["window_delta"]["total_admitted"] == 4 + + ops = json.loads((out / "ops.json").read_text()) + assert ops["top"] and {"name", "count", "self_cpu_us", "self_device_us"} <= set(ops["top"][0]) + assert not (out / "kernels.json").exists() # no CUDA activity on CPU + mem = json.loads((out / "memory.json").read_text()) + assert mem["window_peak"] is None and mem["host_rss_peak_lifetime_bytes"] > 0 + + +@pytest.mark.asyncio +async def test_warmup_defaults_to_the_window_and_peaks_reset_after_it(tmp_path, stub_backend_cls, monkeypatch): + calls = [] + peaks = iter([{"allocated_bytes": 900, "reserved_bytes": 1000}, {"allocated_bytes": 500, "reserved_bytes": 1200}]) + monkeypatch.setattr(bundle, "device_peaks", lambda d: calls.append("peak") or next(peaks)) + monkeypatch.setattr(bundle, "reset_peaks", lambda d: calls.append("reset")) + prompts = profile.synthetic_prompts(3, 8) + out = await profile.run(stub_backend_cls(), Settings(max_batch_size=4), prompts, max_tokens=2, out=tmp_path) + assert calls == ["peak", "reset", "peak"] + mem = json.loads((out / "memory.json").read_text()) + assert mem["window_peak"] == {"allocated_bytes": 500, "reserved_bytes": 1200} + assert mem["lifetime_peak"] == {"allocated_bytes": 900, "reserved_bytes": 1200} + meta = json.loads((out / "meta.json").read_text()) + assert meta["warmup"]["requests"] == 3 + events = [json.loads(l) for l in (out / "events.jsonl").read_text().splitlines()] + begin = next(i for i, e in enumerate(events) if e["kind"] == "profile_window") + assert sum(e["kind"] == "finish" for e in events[:begin]) == 3 + + +def test_warmup_prompts_match_lengths_not_contents(): + prompts = [[5] * 3, [7] * 9] + w = profile.warmup_prompts(prompts, 3) + assert [len(p) for p in w] == [3, 9, 3] and not any(p in prompts for p in w) + + +def test_outcomes_count_empty_as_failed_and_short_separately(): + o = profile.outcomes([[1, 2, 3], [], [1], RuntimeError("x")], max_tokens=3) + assert (o["errors"], o["empty"], o["short"], o["full"]) == (1, 1, 1, 1) + assert o["tokens_out"] == {"total": 4, "min": 0, "mean": 4 / 3, "max": 3} + + +def test_stats_delta_counts_only_the_window(): + s0 = {"total_completed": 2, "decode_steps": 10, "active_size": 1, "wave_sizes": {1: 2}} + s1 = {"total_completed": 6, "decode_steps": 25, "active_size": 0, "wave_sizes": {1: 2, 4: 1}} + assert profile.stats_delta(s0, s1) == {"total_completed": 4, "decode_steps": 15, "wave_sizes": {4: 1}} + + +def test_kernels_table_aggregates_cuda_events_by_name(): + from types import SimpleNamespace as NS + from torch.autograd import DeviceType + + def ev(name, us, dev=DeviceType.CUDA): + return NS(name=name, device_type=dev, time_range=NS(elapsed_us=lambda: us)) + assert bundle.kernels_table([ev("aten::mm", 50, DeviceType.CPU)]) is None + k = bundle.kernels_table([ev("gemm", 10), ev("gemm", 30), ev("softmax", 5), ev("aten::mm", 50, DeviceType.CPU)]) + assert k["launches"] == 3 and k["total_us"] == 45 + assert k["kernels"][0] == {"name": "gemm", "count": 2, "total_us": 40, "mean_us": 20, "max_us": 30} + + +def test_main_profiles_a_presampled_corpus_workload(tmp_path, stub_backend_cls, monkeypatch): + import inference_server.tokenizer as tokmod + + class Tok: + def __init__(self, *a): pass + def encode_messages(self, msgs, thinking=True): + return [len(m["content"]) for m in msgs] + [int(thinking)] + monkeypatch.setattr(tokmod, "Tokenizer", Tok) + monkeypatch.setattr(profile, "build_backend", lambda s: (stub_backend_cls(), None)) + w = {"source": {"kind": "corpus", "class": "steady_interactive", "split": "seen", "seed": 0, + "corpus_version": "v", "lines": [0, 1]}, + "messages": [[{"role": "user", "content": "hi"}], [{"role": "user", "content": "hello"}]]} + (tmp_path / "w.json").write_text(json.dumps(w)) + assert profile.main(["--workload", str(tmp_path / "w.json"), "--max-tokens", "2", "--out", str(tmp_path / "o")]) == 0 + meta = json.loads(next((tmp_path / "o").glob("*/meta.json")).read_text()) + assert meta["workload"] == {**w["source"], "thinking": False} and meta["requests"] == 2 + assert meta["workload_hash"] == bundle.workload_hash([[2, 0], [5, 0]], 2) + + +def test_main_needs_an_explicit_workload_source(): + with pytest.raises(SystemExit): + profile.main([]) def test_synthetic_prompts_are_deterministic(): From bec43b91390f280d2521778d5f22964b87960e44 Mon Sep 17 00:00:00 2001 From: Advay Monga Date: Thu, 8 Oct 2026 01:10:16 -0400 Subject: [PATCH 2/7] Passive data on every bench, equiv and submit Each served run keeps device samples (DCGM profiling fields if dcgmi, else nvidia-smi with throttle reasons; epoch-stamped), engine-native telemetry from the target's optional [engine.telemetry] env ({dir} in the served tree), every client row without text, and the serve log. Bench and equiv store them as ledger blobs under result.artifacts; submit keeps them only in /heldout-private/, never in the ledger. lab/artifacts.py loads a record's artifacts and joins client rows to engine rows by trace id. Co-Authored-By: Claude Opus 5.5 --- CLAUDE.md | 8 +++- lab/LEDGER.md | 27 ++++++++++++- lab/README.md | 4 +- lab/artifacts.py | 72 +++++++++++++++++++++++++++++++++++ lab/evaltools.py | 67 +++++++++++++++++++++++++------- lab/gpu.py | 58 +++++++++++++++++++++++++++- lab/serve.py | 41 ++++++++++++++++++-- lab/target.py | 3 +- regimes/suite.py | 19 +++++++-- targets/inference-server.toml | 3 ++ tests/fake_engine.py | 7 +++- tests/test_lab_artifacts.py | 35 +++++++++++++++++ tests/test_lab_evaltools.py | 41 +++++++++++++++++++- 13 files changed, 357 insertions(+), 28 deletions(-) create mode 100644 lab/artifacts.py create mode 100644 tests/test_lab_artifacts.py diff --git a/CLAUDE.md b/CLAUDE.md index 4b76509..aaca951 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -20,7 +20,7 @@ Top-level packages, run from the repo root (no install step; `pyproject.toml` li pytest/ruff config). Each folder with commands has its own `python -m`. - `targets/.toml` + `lab/target.py` — **the target spec** (since 2026-10-06): model + chat kwargs, engine (repo, python, launch command with `{port}`, health path, env, logprobs `api` vllm|sglang|none, write surface, - add-only globs, test and lint commands), reference server + its outputs dir, MLPerf latency limits, correctness + add-only globs, test and lint commands, optional `[engine.telemetry] env` with `{dir}`), reference server + its outputs dir, MLPerf latency limits, correctness tasks and policy, corpus dir. Selected by `LAB_TARGET` or `--target`; `target.load()` raises `NoTarget` otherwise. Nothing defaults to the user's engine or to Qwen3; `targets/inference-server.toml` is the committed example and what `tests/conftest.py` points at. Surfaces, grader commands, `lab/engine.py`, and the `correctness`, @@ -30,7 +30,11 @@ pytest/ruff config). Each folder with commands has its own `python -m`. `tools.py` (metered tools: test, profile, ledger, budget, restore, note) + `evaltools.py` (bench → `regimes` on the seen split; equiv → `correctness` candidate vs the target's reference; submit → held-out full tier vs the base commit, one aggregate per regime through the Thresholdout guard; all three serve the pristine tree via - `serve.py`, jailed with local binding allowed), `workspace.py` (exported engine copy + snapshots), `ledger.py` + `serve.py`, jailed with local binding allowed; every serve also keeps its passive data, `artifacts.py`: device + samples via `gpu.DeviceSampler` (DCGM if `dcgmi`, else nvidia-smi), engine telemetry files from `{dir}`, all + client rows, the serve log; bench/equiv as ledger blobs under `result.artifacts`, submit only in + `/heldout-private/`, never in the ledger; `artifacts.joined` joins client rows to engine rows by trace id), + `workspace.py` (exported engine copy + snapshots), `ledger.py` (append-only JSONL, `lab/ledger/` gitignored), `budget.py`, `engine.py` (where the engine repo and its python are), `safety/` (write surfaces, srt jail, grader: lint/tests run with the engine's python in the jail), `profile.py` + `bundle.py` + `gpu.py` (profiler harness, staged into the pristine tree under `_harness/` so the diff --git a/lab/LEDGER.md b/lab/LEDGER.md index 3b53022..5b6aba3 100644 --- a/lab/LEDGER.md +++ b/lab/LEDGER.md @@ -11,8 +11,8 @@ on write. Every record: `id`, `at`, `schema` (set by the writer), `kind`, `run`, | `session` | the loop | provider, model, turns, seconds, status (`stop`/`continue`), error, `violation`, `claim.note` | | `test` | `test` | `lint`, `tests` (bools), returncode, output tails, `changed`, `scratch_left_out` | | `profile` | `profile` | `bundle` (ledger blob of the profiler output), `workspace_copy`, seconds | -| `bench` | `bench` | `config: {split: seen, tier}`; `verdict: ok`, `metrics: {regime: {objective, value, better, valid, invalid_reasons}}`, `regimes` (full per-regime results), `ready_s` | -| `equiv` | `equiv` | `config: {split: seen, tier}`; `verdict: pass\|fail\|inconclusive`, `passed`, `reasons`, `gates`, `metrics` (accuracy with CI, length, consistency, divergence, per-task flips, unanswered), `thresholds` | +| `bench` | `bench` | `config: {split: seen, tier}`; `verdict: ok`, `metrics: {regime: {objective, value, better, valid, invalid_reasons}}`, `regimes` (full per-regime results), `ready_s`, `artifacts` | +| `equiv` | `equiv` | `config: {split: seen, tier}`; `verdict: pass\|fail\|inconclusive`, `passed`, `reasons`, `gates`, `metrics` (accuracy with CI, length, consistency, divergence, per-task flips, unanswered), `thresholds`, `artifacts` | | `submit` | `submit` | **held-out shape**: no `result`; top-level `config: {split: heldout, tier: full}` and `metrics: {regime: {base, new, delta_pct, band_pct, verdict}}`, scalars only. The ledger refuses anything else that mentions the held-out split. | | `finding` | `lab.ledger seed` | a `knowledge/` entry, whole, as `claim`; `source`, `content_sha` | | `note` | `note`, `restore`, refusals | free text, or `{verdict: refused, reason}` | @@ -21,3 +21,26 @@ A tool that could not run writes `{verdict: refused | error, reason}` and the ag crashes the session. The integrity verdict (`feedback/lab_verdict.py`) reads only these records: a run is valid iff a completed `submit` exists whose snapshot also has a completed seen-split `bench`, a passing `test`, a passing `equiv` (tier full) and no failing `equiv`, and no record reports a write-surface violation. + +## Passive data (`artifacts`) + +Every bench and equiv (including one that errored after serving) keeps the cheap data from the same served run as +ledger blobs, referenced from `result.artifacts: {device, telemetry, client_rows, serve_log}` (blob paths under the +ledger root, `telemetry` None when the target has no `[engine.telemetry]` or the engine wrote nothing). The tool's +reply to the agent never includes them. `lab/artifacts.py` loads them back (`load`) and joins client rows to engine +rows by trace id (`joined`: client `X-Trace-Id` is `bg-`). + +- `device/`: `meta.json` (`available`; when true `source` dcgm|nvidia-smi, `coarse` (nvidia-smi utilization is + coarse, DCGM profiling fields are not), `fields`, `interval_ms`, `exited` if the sampler died) and `samples.csv`, + each raw tool line prefixed with epoch seconds. Covers the whole served window, launch to teardown. +- `telemetry/`: whatever the engine wrote into `{dir}` (inference-server: `.sqlite` with a `requests` table, + `events.jsonl`). +- `client_rows`: JSONL, one per measured request, without generated text. Bench: the regimes' rows with `regime`, + `call` (which run_open/run_closed call), `call_started_at` (epoch) and run-clock times; warmups are not kept. + Equiv: one per gate item (`id`, `completion_tokens`, `finish_reason`, `error`). +- `serve_log`: this run's part of the serve log. + +**Submit keeps none of this in the ledger.** Its held-out record stays in the strict held-out shape above, with no +field pointing at passive data, and no blob is written for it. The same data for the held-out serve goes to +`/heldout-private/-