Hunt for data corruption in vllm

hunt.py

#!/usr/bin/env python3
"""
hunt.py — drive the compiled vLLM server with ragged concurrent load and probe
for the silent inductor NaN corruption (output collapses to token id 0 / '!!!!').

Exits 0 the moment corruption is confirmed (writes a JSON report); exits 2 on
timeout. Prints progress so it can be tailed live. Stdlib only.

Env knobs: HUNT_CONC (concurrency), HUNT_MINUTES (max run), HUNT_CANARY_SEC
(seconds between canary probes), HUNT_URL, HUNT_MODEL, HUNT_REPORT.
"""
import json, os, sys, time, threading, urllib.request, urllib.error, random
from collections import Counter
from concurrent.futures import ThreadPoolExecutor

URL   = os.environ.get("HUNT_URL", "http://localhost:8888/v1/chat/completions")
MODEL = os.environ.get("HUNT_MODEL", "qwen3.6-27b")
CONC  = int(os.environ.get("HUNT_CONC", "14"))
MINUTES = float(os.environ.get("HUNT_MINUTES", "45"))
CANARY_SEC = float(os.environ.get("HUNT_CANARY_SEC", "8"))
REPORT = os.environ.get("HUNT_REPORT", "/tmp/logs/hunt-report.json")

# Ragged prompt pool — varied lengths to force non-uniform batch tiles.
FILLERS = ["the history of", "a short note on", "in detail, discuss", "briefly,",
           "please describe", "explain like I am five", "compare and contrast",
           "list three facts about", "write one sentence about"]
TOPICS = ["the ocean", "transistors", "photosynthesis", "the Roman empire",
          "black holes", "coffee", "violins", "glaciers", "the printing press",
          "migratory birds", "quantum tunneling", "the water cycle", "volcanoes",
          "bicycles", "the immune system", "tea ceremonies", "sailing", "comets"]

sent = Counter()
lock = threading.Lock()
stop = threading.Event()


def post(messages, max_tokens, temperature):
    body = json.dumps({"model": MODEL, "messages": messages,
                       "max_tokens": max_tokens, "temperature": temperature,
                       "ignore_eos": True}).encode()
    req = urllib.request.Request(URL, data=body,
                                 headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=120) as r:
        d = json.loads(r.read())
    return d["choices"][0]["message"]["content"]


def make_prompt(rnd):
    # target a ragged token count up to ~260 to hit boundary tiles like the
    # benchmark that reliably tripped the bug (256-token prompts).
    target_words = rnd.choice([3, 8, 20, 40, 80, 120, 180, 210])
    words = [rnd.choice(FILLERS)]
    while len(words) < target_words:
        words.append(rnd.choice(TOPICS))
        words.append(rnd.choice(FILLERS))
    return " ".join(words[:target_words])


def load_worker():
    rnd = random.Random(threading.get_ident())
    while not stop.is_set():
        prompt = make_prompt(rnd)
        mt = rnd.choice([16, 24, 32, 48, 64, 96, 128, 160])  # ragged decode len
        try:
            post([{"role": "user", "content": prompt}], mt, 0.7)
            with lock:
                sent["ok"] += 1
        except Exception:
            with lock:
                sent["err"] += 1


def degenerate(text):
    """True if a canary answer shows the corruption signature."""
    t = (text or "").strip()
    if not t:
        return None  # inconclusive
    if "paris" in t.lower():
        return False
    # not the healthy answer — is it a degenerate repeat?
    if "!!!!" in t:
        return True
    if len(t) >= 4:
        top = Counter(t).most_common(1)[0][1]
        if top / len(t) > 0.5:
            return True
    return None  # unhealthy but not clearly degenerate — log & keep watching


def canary():
    return post([{"role": "user",
                  "content": "What is the capital of France? Answer in one word."}],
                20, 0.0)


def main():
    print(f"hunt: conc={CONC} minutes={MINUTES} canary_every={CANARY_SEC}s url={URL}",
          flush=True)
    deadline = time.time() + MINUTES * 60
    t0 = time.time()
    pool = ThreadPoolExecutor(max_workers=CONC)
    for _ in range(CONC):
        pool.submit(load_worker)

    last_canary = 0.0
    while time.time() < deadline:
        time.sleep(1.0)
        if time.time() - last_canary >= CANARY_SEC:
            last_canary = time.time()
            try:
                ans = canary()
            except Exception as e:
                print(f"[{time.time()-t0:6.0f}s] canary error: {e}", flush=True)
                continue
            deg = degenerate(ans)
            with lock:
                s = dict(sent)
            print(f"[{time.time()-t0:6.0f}s] sent={s} canary={ans!r} degenerate={deg}",
                  flush=True)
            if deg:
                # confirm persistence: 3 more canaries, engine now near-idle
                stop.set()
                time.sleep(3)
                confirms = []
                for _ in range(3):
                    try:
                        confirms.append(canary())
                    except Exception as e:
                        confirms.append(f"ERR:{e}")
                    time.sleep(1)
                report = {
                    "tripped": True,
                    "elapsed_sec": round(time.time() - t0, 1),
                    "requests_sent": s,
                    "trigger_canary": ans,
                    "confirm_canaries": confirms,
                    "confirmed_persistent": sum(
                        1 for c in confirms if degenerate(c)) >= 2,
                    "conc": CONC,
                }
                with open(REPORT, "w") as f:
                    json.dump(report, f, indent=2)
                print("=== TRIPPED ===", flush=True)
                print(json.dumps(report, indent=2), flush=True)
                pool.shutdown(wait=False)
                sys.exit(0)

    # timed out
    stop.set()
    with lock:
        s = dict(sent)
    report = {"tripped": False, "elapsed_sec": round(time.time() - t0, 1),
              "requests_sent": s, "conc": CONC}
    with open(REPORT, "w") as f:
        json.dump(report, f, indent=2)
    print("=== TIMEOUT, no corruption ===", flush=True)
    print(json.dumps(report, indent=2), flush=True)
    pool.shutdown(wait=False)
    sys.exit(2)


if __name__ == "__main__":
    main()