#!/usr/bin/env python3 """geo_spectrum_run.py — churn the GEO free-form probes across a SPECTRUM of local models on abzu. Brendan: "Churn it for a spectrum across all local models ... run many times and produce the raw output." Each model: load (from /mnt/cache NVMe if staged, else NFS), run every question × N samples through the VERBATIM GEO baseline-answer + cold + post-hoc criteria probes, append raw JSONL. While model N runs, the NEXT models are staged NFS->/mnt/cache in the background so switching never waits on the network; finished models are evicted (keep <= CACHE_KEEP large models on the 1.8T cache). Raw output: one JSONL line per (model, query_id, sample) with the answer + both criteria self-reports. Resumable: skips (model, query_id, sample) triples already present in the output. python3 geo_spectrum_run.py --models 01-ai/Yi-1.5-34B-Chat,tencent/Hunyuan-7B-Instruct \ --questions /tmp/jlens_lab_harness/data/lab_questions_20.csv --samples 1 --out /mnt/cache/geo_raw.jsonl """ from __future__ import annotations import argparse, csv, glob, json, os, shutil, subprocess, sys, threading, time os.environ.setdefault("HF_HOME", "/mnt/assets/models/huggingface") os.environ.setdefault("HF_HUB_OFFLINE", "1") os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") sys.path.insert(0, "/tmp/jlens_lab_harness") import torch from transformers import AutoModelForCausalLM, AutoTokenizer from geo.common import prompts, criteria HUB_NFS = "/mnt/assets/models/huggingface/hub" CACHE = "/mnt/cache/hf/hub" # NVMe staging mirror (same hub layout) CACHE_KEEP = 3 # max large models kept on the 1.8T NVMe def _hub_dir(model): return "models--" + model.replace("/", "--") def snapshot_path(root, model): for c in sorted(glob.glob(f"{root}/{_hub_dir(model)}/snapshots/*/")): if c.rstrip("/").endswith(".partial"): continue if os.path.exists(os.path.join(c, "config.json")) and ( glob.glob(os.path.join(c, "*.safetensors")) or glob.glob(os.path.join(c, "*.bin"))): return c return None _stage_lock = threading.Lock() _staged = set() def stage_to_cache(model): """Copy a model's snapshot NFS->NVMe. Serialized + deduped: only one cp at a time, each model staged at most once (the prefetch windows overlap, so without this they race on the temp dir).""" with _stage_lock: if model in _staged or snapshot_path(CACHE, model): _staged.add(model); return CACHE src_snap = snapshot_path(HUB_NFS, model) if not src_snap: return None dst_snap = src_snap.replace(HUB_NFS, CACHE).rstrip("/") tmp = dst_snap + ".partial" shutil.rmtree(tmp, ignore_errors=True) os.makedirs(os.path.dirname(dst_snap), exist_ok=True) subprocess.run(["cp", "-rL", src_snap.rstrip("/"), tmp], check=True) shutil.rmtree(dst_snap, ignore_errors=True) os.replace(tmp, dst_snap) _staged.add(model) return CACHE def evict_from_cache(model): d = f"{CACHE}/{_hub_dir(model)}" if os.path.isdir(d): shutil.rmtree(d, ignore_errors=True) def cache_models_present(): return [os.path.basename(p) for p in glob.glob(f"{CACHE}/models--*")] def _merge_system_into_user(messages): """Some templates (Gemma-2, etc.) reject a system role. Fold the system content into the first user turn — same VERBATIM content, only the delivery structure changes.""" if not messages or messages[0]["role"] != "system": return messages sys_txt = messages[0]["content"]; out = []; injected = False for m in messages[1:]: if m["role"] == "user" and not injected: out.append({"role": "user", "content": sys_txt + "\n\n" + m["content"]}); injected = True else: out.append(m) return out def gen(model, tok, messages, max_new, temperature=0.7): try: enc = tok.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt", return_dict=True) except Exception as e: if "system" in str(e).lower(): enc = tok.apply_chat_template(_merge_system_into_user(messages), add_generation_prompt=True, return_tensors="pt", return_dict=True) else: raise enc = {k: v.to(model.device) for k, v in enc.items()} n_in = enc["input_ids"].shape[1] with torch.no_grad(): out = model.generate(**enc, max_new_tokens=max_new, do_sample=temperature > 0, temperature=max(temperature, 1e-5), top_p=0.95, pad_token_id=(tok.pad_token_id or tok.eos_token_id)) return tok.decode(out[0, n_in:], skip_special_tokens=True).strip() def cold_user(query): t = criteria.COLD_CRITERIA_USER_TEMPLATE return t.format(query=query) if "{query}" in t else t + f'\n\nQuestion: "{query}"' def parse(raw): try: return criteria.parse_criteria(raw) except Exception as e: return {"_parse_error": str(e), "_raw": raw[:1500]} def run_model(model_id, run_label, questions, samples, out_f, done, max_new, trust, device, locale="en-AU"): load_root = CACHE if snapshot_path(CACHE, model_id) else HUB_NFS path = snapshot_path(load_root, model_id) # Neither staging root has this model: hand the HF id straight to transformers and let it # resolve from HF_HOME (or download). Without this the resolver returns None, None reaches # from_pretrained, and the error is the memorable "None is not a local folder" -- which reads # like a bad model id rather than "this host has no NFS mount". The sparks have neither # /mnt/assets nor a pre-staged cache, so every probe there died on it. src = "NVMe" if load_root == CACHE else "NFS" if path is None: path, src = model_id, "hub" print(f"[load] {model_id} from {src}", file=sys.stderr, flush=True) t0 = time.time() tok = AutoTokenizer.from_pretrained(path) # --device auto: shard across every GPU and spill the tail to CPU RAM — how a 206G giant # (GLM-4.5-Air, Phase 7) generates on 2x97GB. Slower per token for the spilled layers; fine # for 20 questions. dm = "auto" if device == "auto" else {"": device} # auto mode caps each GPU at 78GiB: the first giant OOM'd because device_map=auto packs the # cards to the brim and leaves nothing for activations + KV cache. The remainder spills to CPU. mm = ({i: "78GiB" for i in range(torch.cuda.device_count())} | {"cpu": "220GiB"}) if device == "auto" else None model = AutoModelForCausalLM.from_pretrained(path, dtype=torch.bfloat16, device_map=dm, max_memory=mm, attn_implementation="sdpa", trust_remote_code=trust) model.eval() print(f"[loaded] {model_id} in {time.time()-t0:.0f}s on {model.device}", file=sys.stderr, flush=True) for q in questions: qid, query = q["query_id"], q["query"] for s in range(samples): key = (run_label, qid, s) if key in done: continue answer = gen(model, tok, [ {"role": "system", "content": prompts.BASE_SYSTEM_PROMPT}, {"role": "user", "content": prompts.BASE_USER_TEMPLATE.format(locale=locale, query=query)}], max_new) cold = gen(model, tok, [ {"role": "system", "content": criteria.COLD_CRITERIA_SYSTEM_PROMPT}, {"role": "user", "content": cold_user(query)}], max_new, temperature=0.0) posthoc = gen(model, tok, [ {"role": "system", "content": prompts.BASE_SYSTEM_PROMPT}, {"role": "user", "content": prompts.BASE_USER_TEMPLATE.format(locale=locale, query=query)}, {"role": "assistant", "content": answer}, {"role": "user", "content": criteria.POSTHOC_CRITERIA_FOLLOWUP}], max_new, temperature=0.0) rec = {"model": model_id, "run_label": run_label, "query_id": qid, "sample": s, "expected_criteria": q.get("expected_criteria", ""), "intent": q.get("intent", ""), "product_category": q.get("product_category", ""), "answer_text": answer, "criteria_cold": parse(cold), "criteria_posthoc": parse(posthoc), "ts": time.time()} out_f.write(json.dumps(rec) + "\n"); out_f.flush() print(f" [{run_label}] {qid} s{s}: ans={len(answer)}c " f"cold={len(rec['criteria_cold']) if isinstance(rec['criteria_cold'],list) else 'ERR'} " f"posthoc={len(rec['criteria_posthoc']) if isinstance(rec['criteria_posthoc'],list) else 'ERR'}", file=sys.stderr, flush=True) del model, tok torch.cuda.empty_cache() def main(): ap = argparse.ArgumentParser() ap.add_argument("--models", required=True, help="comma-separated HF ids, the spectrum in order") ap.add_argument("--questions", required=True); ap.add_argument("--samples", type=int, default=1) ap.add_argument("--label-suffix", default="", help="append to run_label, e.g. -r2 for a C0 replicate") ap.add_argument("--out", required=True); ap.add_argument("--no-prefetch", action="store_true") ap.add_argument("--max-new", type=int, default=768); ap.add_argument("--trust", action="store_true") ap.add_argument("--device", default="cuda:0") a = ap.parse_args() models = [m.strip() for m in a.models.split(",") if m.strip()] questions = list(csv.DictReader(open(a.questions))) os.makedirs(CACHE, exist_ok=True) # resume: read already-done (run_label, query_id, sample) done = set() if os.path.exists(a.out): for line in open(a.out): try: d = json.loads(line); done.add((d["run_label"], d["query_id"], d["sample"])) except Exception: pass print(f"[spectrum] {len(models)} models, {len(questions)} questions x {a.samples} " f"samples; {len(done)} already done", file=sys.stderr, flush=True) def label(m): return m.split("/")[-1].lower() + a.label_suffix out_f = open(a.out, "a") for i, m in enumerate(models): # prefetch: stage the NEXT model(s) to NVMe in the background while this one runs prefetch_threads = [] if not a.no_prefetch: for nxt in models[i+1:i+3]: th = threading.Thread(target=lambda mm=nxt: (print(f"[prefetch] staging {mm}", file=sys.stderr, flush=True), stage_to_cache(mm)), daemon=True) th.start(); prefetch_threads.append(th) try: run_model(m, label(m), questions, a.samples, out_f, done, a.max_new, a.trust, a.device) except Exception as e: print(f"[ERROR] {m}: {type(e).__name__}: {e}", file=sys.stderr, flush=True) # evict models beyond the keep-window to bound the NVMe present = cache_models_present() keep = {_hub_dir(x) for x in models[i:i+CACHE_KEEP]} for hubname in present: if hubname not in keep: shutil.rmtree(f"{CACHE}/{hubname}", ignore_errors=True) for th in prefetch_threads: th.join(timeout=1) out_f.close() print("[spectrum] done", file=sys.stderr, flush=True) if __name__ == "__main__": main()