code/scripts/benchmark_batch.py

#!/usr/bin/env python3
"""Measure warmed free-T4 throughput using repeated TRAINING prompts only.

Writes an incremental Drive report for the pinned base or a completed adapter.
These are capacity checks, not evidence about the frozen blind benchmark.
"""
from __future__ import annotations

import argparse
from datetime import datetime, timezone
import gc
import hashlib
import importlib.metadata
import json
import os
from pathlib import Path
import sys
import time

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src"))
from vibegame import BASE_MODEL
from generate import UnslothGenerator

BASE_REVISION = "4858886f896bd05db823557fc5cbc4ac28342af7"
BATCH_SIZES = [1, 4, 8, 16]
MEMORY_ESTIMATE_POLICY = "kv_full_context_v1"


def require_gpu(torch, pro_a100=False):
    expected = "A100" if pro_a100 else "T4"
    if (not torch.cuda.is_available() or torch.cuda.device_count() != 1
            or expected not in torch.cuda.get_device_name()):
        raise RuntimeError(f"Requires one already allocated {expected} CUDA GPU; no GPU is provisioned")


def capacity_prompts(adapter=None, training_jsonl=None):
    """The opt-in A100 path reads only the exact admitted training split."""
    path = training_jsonl or ROOT / "data/teacher_prompts.jsonl"
    if training_jsonl is not None:
        signature = json.loads((adapter / "run_manifest.json").read_text())["signature"]
        if hashlib.sha256(path.read_bytes()).hexdigest() != signature["train_sha256"]:
            raise ValueError("Capacity training file differs from the completed adapter's training split")
    rows = []
    with path.open() as handle:
        for line in handle:
            if line.strip():
                row = json.loads(line)
                rows.append({"id": row.get("id", f"training-row-{len(rows)}"), "instruction": row["instruction"]})
            if len(rows) == 4:
                break
    if len(rows) != 4 or any(not isinstance(row["instruction"], str) or not row["instruction"] for row in rows):
        raise ValueError("Four training prompts are required")
    return rows


def save_report(path: Path, report: dict) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(path.suffix + ".tmp")
    with temporary.open("w", encoding="utf-8") as handle:
        json.dump(report, handle, indent=2)
        handle.write("\n")
        handle.flush()
        os.fsync(handle.fileno())
    temporary.replace(path)


def select_candidate(measurements: list[dict], memory_limit_bytes: int) -> dict | None:
    candidates = [row for row in measurements if row.get("status") == "success"
                  and row["peak_allocated_bytes"] < memory_limit_bytes
                  and row["peak_reserved_bytes"] < memory_limit_bytes
                  and row.get("full_context_memory_estimate", {}).get("within_limit") is True]
    return max(candidates, key=lambda row: row["sum_tokens_per_second"]) if candidates else None


def estimate_full_context_memory(row: dict, config: dict, gpu_memory_bytes: int,
                                 memory_limit_bytes: int, cache_element_bytes: int = 2,
                                 growth_tokens: int = 512) -> dict:
    """Conservative KV-growth extrapolation, not a proof against CUDA OOM."""
    batch = row["batch_size"]
    inputs, outputs = row["per_prompt_input_tokens"], row["per_prompt_tokens"]
    if len(inputs) != batch or len(outputs) != batch or any(value <= 0 for value in inputs + outputs):
        raise ValueError("Memory estimate requires actual input/output lengths for every batch row")
    layers = int(config["num_hidden_layers"])
    heads = int(config["num_attention_heads"])
    kv_heads = int(config["num_key_value_heads"])
    head_dim = int(config.get("head_dim") or (int(config["hidden_size"]) // heads))
    if min(layers, heads, kv_heads, head_dim, cache_element_bytes) <= 0:
        raise ValueError("Invalid loaded model cache dimensions")
    growth_tokens = max(512, growth_tokens)
    bytes_per_sequence_token = 2 * layers * kv_heads * head_dim * cache_element_bytes
    # Left padding and completed rows remain in the same decoding batch until
    # the longest output finishes. The last emitted token need not be cached.
    observed_cache_tokens = max(inputs) + max(outputs) - 1
    target_tokens = 4096 + growth_tokens
    extra_cache_tokens = max(0, target_tokens - observed_cache_tokens)
    delta_kv_bytes = batch * bytes_per_sequence_token * extra_cache_tokens
    observed_peak = max(row["peak_allocated_bytes"], row["peak_reserved_bytes"])
    reserve_bytes = 2_000_000_000
    effective_limit = min(memory_limit_bytes, gpu_memory_bytes - reserve_bytes)
    estimated_peak = observed_peak + delta_kv_bytes
    return {"policy": MEMORY_ESTIMATE_POLICY, "assessed_batch_size": batch,
            "target_context_tokens": 4096, "cache_growth_allowance_tokens": growth_tokens,
            "cache_element_bytes": cache_element_bytes, "bytes_per_sequence_token": bytes_per_sequence_token,
            "model_dimensions": {"layers": layers, "attention_heads": heads, "kv_heads": kv_heads, "head_dim": head_dim},
            "observed_padded_prompt_tokens": max(inputs), "observed_max_generated_tokens": max(outputs),
            "observed_cache_tokens": observed_cache_tokens, "observed_peak_bytes": observed_peak,
            "additional_cache_tokens": extra_cache_tokens, "additional_kv_bytes": delta_kv_bytes,
            "estimated_peak_bytes": estimated_peak, "effective_limit_bytes": effective_limit,
            "reserve_bytes": reserve_bytes, "within_limit": estimated_peak < effective_limit,
            "limitations": "Estimate only, not a measured full-context peak or an OOM guarantee. Adds KV growth to max observed allocated/reserved memory, with an extra cache growth block. Other workspace growth, prefill shapes and allocator fragmentation are not modeled; at least 2 GB of physical GPU memory remains reserved."}


def full_check_reason(best: dict | None, speedup: float | None,
                      batch_sizes: list[int], short_only: bool) -> str | None:
    """A single requested size checks capacity without claiming comparative speed."""
    if short_only:
        return "short-only requested"
    if best is None:
        return "No successful short benchmark has the required memory headroom"
    if len(batch_sizes) == 1:
        return None
    if speedup is None or speedup < 1.05:
        return "No eligible batch improves aggregate throughput by at least 5% over batch one"
    return None


def measure(backend, prompts: list[str], token_limit: int, torch) -> dict:
    # Warm this batch shape separately; shape-specific JIT is excluded from timing.
    warmup_started = time.monotonic()
    backend.generate_many(prompts, 4)
    torch.cuda.synchronize()
    warmup_seconds = time.monotonic() - warmup_started
    gc.collect()
    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats()
    outer_started = time.monotonic()
    results = backend.generate_many(prompts, token_limit)
    outer_seconds = time.monotonic() - outer_started
    elapsed = max(row["generation_seconds"] for row in results)
    call_elapsed = max(row.get("generation_call_seconds", row["generation_seconds"]) for row in results)
    total_tokens = sum(row["tokens"] for row in results)
    if elapsed <= 0 or len(results) != len(prompts):
        raise RuntimeError("Invalid benchmark output count or elapsed time")
    return {"status": "success", "batch_size": len(prompts), "max_new_tokens": token_limit,
            "warmup_seconds": warmup_seconds, "generation_seconds": elapsed,
            "generation_call_seconds": call_elapsed,
            "end_to_end_call_seconds": outer_seconds, "sum_generated_tokens": total_tokens,
            "sum_tokens_per_second": total_tokens / elapsed,
            "sum_tokens_per_call_second": total_tokens / outer_seconds,
            "per_prompt_tokens": [row["tokens"] for row in results],
            "per_prompt_input_tokens": [row["prompt_tokens"] for row in results],
            "per_prompt_done_reason": [row["done_reason"] for row in results],
            "peak_allocated_bytes": torch.cuda.max_memory_allocated(),
            "peak_reserved_bytes": torch.cuda.max_memory_reserved(),
            "latency_note": "Every prompt waits the whole batch call, including tokenization and decoding. sum_tokens_per_second uses model.generate; sum_tokens_per_call_second uses the full outer call. Both are aggregate throughput, not request latency."}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("output", type=Path, help="JSON path under /content/drive/MyDrive")
    parser.add_argument("--adapter", type=Path, help="Completed local adapter; load its exact saved base revision and tokenizer")
    parser.add_argument("--batch-sizes", type=int, nargs="+", choices=[1, 2, 4, 8, 16, 32], default=BATCH_SIZES,
                        help="Ascending batch sizes to measure; pass one size for a direct capacity check")
    parser.add_argument("--short-tokens", type=int, choices=[128, 256], default=256)
    parser.add_argument("--memory-limit-gb", type=float, default=12.5,
                        help="Conservative decimal-GB threshold before the full check")
    parser.add_argument("--short-only", action="store_true", help="Skip automatic full-length check")
    parser.add_argument("--pro-a100", action="store_true", help="Explicit separate Pro A100 evaluation capacity check; never changes T4 defaults")
    parser.add_argument("--training-jsonl", type=Path, help="Exact admitted training split, required only for --pro-a100")
    args = parser.parse_args()
    output = args.output.expanduser().resolve()
    drive = Path("/content/drive/MyDrive")
    if not drive.is_dir() or not output.is_relative_to(drive.resolve()):
        parser.error("Benchmark evidence must persist under mounted /content/drive/MyDrive")
    if output.exists():
        parser.error("Report already exists; use a new path to preserve measurements")
    if not 1 <= args.memory_limit_gb <= 12.5:
        parser.error("memory-limit-gb must be within 1..12.5")
    if args.batch_sizes != sorted(set(args.batch_sizes)):
        parser.error("batch-sizes must be unique and ascending")
    if args.pro_a100:
        if not args.adapter or args.batch_sizes != [8] or args.short_only or not args.training_jsonl:
            parser.error("--pro-a100 requires a completed adapter, --batch-sizes 8, --training-jsonl, and the full check")
    elif args.training_jsonl is not None:
        parser.error("--training-jsonl is reserved for the opt-in --pro-a100 path")
    pinned_base, pinned_revision = BASE_MODEL, BASE_REVISION
    if args.adapter:
        args.adapter = args.adapter.expanduser().resolve(strict=True)
        if not (args.adapter / "TRAINING_COMPLETE.json").is_file():
            parser.error("Adapter capacity check requires TRAINING_COMPLETE.json from the finished training run")
        signature = json.loads((args.adapter / "run_manifest.json").read_text())["signature"]
        pinned_base, pinned_revision = signature["base_model"], signature["base_revision"]
    from unsloth import FastLanguageModel
    import torch
    from bitsandbytes.nn import Linear4bit
    require_gpu(torch, args.pro_a100)
    # Read exactly four teacher TRAINING prompts. Never open a blind-eval file.
    prompt_rows = capacity_prompts(args.adapter, args.training_jsonl)
    prompts = [row["instruction"] for row in prompt_rows]
    memory_limit = int(args.memory_limit_gb * 1_000_000_000)
    report = {"kind": ("trained_adapter" if args.adapter else "untrained_base") + "_warmed_batch_throughput_on_training_prompts_not_model_evaluation",
              "started_utc": datetime.now(timezone.utc).isoformat(), "status": "loading_adapter" if args.adapter else "loading_base",
              "base_model": pinned_base, "base_revision": pinned_revision,
              "adapter": str(args.adapter) if args.adapter else None,
              "gpu": torch.cuda.get_device_name(), "gpu_memory_bytes": torch.cuda.get_device_properties(0).total_memory,
              "batch_sizes": args.batch_sizes, "short_max_new_tokens": args.short_tokens,
              "full_max_new_tokens": 3072, "memory_limit_bytes": memory_limit,
              "memory_estimate_policy": MEMORY_ESTIMATE_POLICY,
              "training_prompt_ids": [row["id"] for row in prompt_rows],
              "training_prompts_sha256": hashlib.sha256(json.dumps(prompts, ensure_ascii=False).encode()).hexdigest(),
              "blind_prompts_read": False,
              "versions": {name: importlib.metadata.version(name) for name in ["torch", "unsloth", "transformers", "bitsandbytes"]},
              "measurements": [], "full_check": None,
              "limitations": ("Completed adapter capacity check. " if args.adapter else "Original base only; adapter still needs capacity check. ") + "A 3072-token cap can end at EOS earlier and is not a forced worst-case allocation test. No quality or gameplay conclusion."}
    if args.pro_a100:
        report["evaluation_profile"] = "pro-a100-opt-in"
        report["training_jsonl_sha256"] = hashlib.sha256(args.training_jsonl.read_bytes()).hexdigest()
    save_report(output, report)
    if args.adapter:
        backend = UnslothGenerator(str(args.adapter))
        report["generation"] = backend.metadata
    else:
        model, tokenizer = FastLanguageModel.from_pretrained(model_name=BASE_MODEL, revision=BASE_REVISION,
                        max_seq_length=4096, dtype=torch.float16, load_in_4bit=True,
                        trust_remote_code=False, use_exact_model_name=True)
        for module in model.modules():
            if isinstance(module, Linear4bit):
                module.compute_dtype = torch.float16
                module.compute_type_is_set = True
        FastLanguageModel.for_inference(model)
        backend = UnslothGenerator.__new__(UnslothGenerator)
        backend.model, backend.tokenizer, backend.torch = model, tokenizer, torch
    config = backend.model.config.to_dict()
    cache_element_bytes = backend.model.get_input_embeddings().weight.element_size()
    from unsloth.models.llama import KV_CACHE_INCREMENT
    cache_growth_tokens = max(512, int(KV_CACHE_INCREMENT))
    report["memory_estimate_cache_dtype_basis"] = str(backend.model.get_input_embeddings().weight.dtype)

    def add_memory_estimate(row):
        row["full_context_memory_estimate"] = estimate_full_context_memory(
            row, config, report["gpu_memory_bytes"], memory_limit,
            cache_element_bytes=cache_element_bytes, growth_tokens=cache_growth_tokens)

    report["status"] = "short_benchmarks"
    save_report(output, report)
    for size in args.batch_sizes:
        repeated = [prompts[index % 4] for index in range(size)]
        try:
            row = measure(backend, repeated, args.short_tokens, torch)
            add_memory_estimate(row)
        except Exception as exc:
            row = {"batch_size": size, "status": "failed", "error": type(exc).__name__ + ": " + str(exc)}
            gc.collect()
            torch.cuda.empty_cache()
        report["measurements"].append(row)
        save_report(output, report)
        print(json.dumps(row), flush=True)
        if row["status"] == "failed":
            # Avoid making an OOM worse by trying a still-larger batch.
            break
    best = select_candidate(report["measurements"], memory_limit)
    report["best_short_batch_size"] = best["batch_size"] if best else None
    baseline = next((row for row in report["measurements"] if row.get("status") == "success" and row["batch_size"] == 1), None)
    speedup = best["sum_tokens_per_second"] / baseline["sum_tokens_per_second"] if best and baseline else None
    report["short_speedup_vs_batch1"] = speedup
    skip_reason = full_check_reason(best, speedup, args.batch_sizes, args.short_only)
    if skip_reason is None:
        report["status"] = "full_length_check"
        save_report(output, report)
        size = best["batch_size"]
        try:
            full = measure(backend, [prompts[index % 4] for index in range(size)], 3072, torch)
            add_memory_estimate(full)
            full["within_memory_limit"] = full["peak_allocated_bytes"] < memory_limit and full["peak_reserved_bytes"] < memory_limit
            report["full_check"] = full
            report["recommended_batch_size"] = size if full["within_memory_limit"] and full["full_context_memory_estimate"]["within_limit"] else None
        except Exception as exc:
            report["full_check"] = {"batch_size": size, "status": "failed", "error": type(exc).__name__ + ": " + str(exc)}
            report["recommended_batch_size"] = None
        print(json.dumps({"full_check": report["full_check"]}), flush=True)
    else:
        report["full_check_skip_reason"] = skip_reason
        report["recommended_batch_size"] = None
    report["status"] = "complete"
    save_report(output, report)
    print(json.dumps({"report": str(output), "best_short_batch_size": report["best_short_batch_size"],
                      "recommended_batch_size": report["recommended_batch_size"]}), flush=True)


if __name__ == "__main__":
    main()