The black-box probe runner. System-role folding for templates that reject one (Gemma-2), NVMe→NFS→hub fallback, label-suffix for replicates, and the PRODUCED-NOTHING guard family.
raw: geo_spectrum_run.py · annotate via ../marginalia/code-geo_spectrum_run-py.json
#!/usr/bin/env python3 class="s">"""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 class="s">""" from __future__ import annotations import argparse, csv, glob, json, os, shutil, subprocess, sys, threading, time os.environ.setdefault(class="s">"HF_HOME", class="s">"/mnt/assets/models/huggingface") os.environ.setdefault(class="s">"HF_HUB_OFFLINE", class="s">"1") os.environ.setdefault(class="s">"TOKENIZERS_PARALLELISM", class="s">"false") sys.path.insert(0, class="s">"/tmp/jlens_lab_harness") import torch from transformers import AutoModelForCausalLM, AutoTokenizer from geo.common import prompts, criteria HUB_NFS = class="s">"/mnt/assets/models/huggingface/hub" CACHE = class="s">"/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 class="s">"models--" + model.replace(class="s">"/", class="s">"--") def snapshot_path(root, model): for c in sorted(glob.glob(fclass="s">"{root}/{_hub_dir(model)}/snapshots/*/")): if c.rstrip(class="s">"/").endswith(class="s">".partial"): continue if os.path.exists(os.path.join(c, class="s">"config.json")) and ( glob.glob(os.path.join(c, class="s">"*.safetensors")) or glob.glob(os.path.join(c, class="s">"*.bin"))): return c return None _stage_lock = threading.Lock() _staged = set() def stage_to_cache(model): class="s">"""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).class="s">""" 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(class="s">"/") tmp = dst_snap + class="s">".partial" shutil.rmtree(tmp, ignore_errors=True) os.makedirs(os.path.dirname(dst_snap), exist_ok=True) subprocess.run([class="s">"cp", class="s">"-rL", src_snap.rstrip(class="s">"/"), 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 = fclass="s">"{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(fclass="s">"{CACHE}/models--*")] def _merge_system_into_user(messages): class="s">"""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.class="s">""" if not messages or messages[0][class="s">"role"] != class="s">"system": return messages sys_txt = messages[0][class="s">"content"]; out = []; injected = False for m in messages[1:]: if m[class="s">"role"] == class="s">"user" and not injected: out.append({class="s">"role": class="s">"user", class="s">"content": sys_txt + class="s">"\n\n" + m[class="s">"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=class="s">"pt", return_dict=True) except Exception as e: if class="s">"system" in str(e).lower(): enc = tok.apply_chat_template(_merge_system_into_user(messages), add_generation_prompt=True, return_tensors=class="s">"pt", return_dict=True) else: raise enc = {k: v.to(model.device) for k, v in enc.items()} n_in = enc[class="s">"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 class="s">"{query}" in t else t + fclass="s">'\n\nQuestion: "{query}"' def parse(raw): try: return criteria.parse_criteria(raw) except Exception as e: return {class="s">"_parse_error": str(e), class="s">"_raw": raw[:1500]} def run_model(model_id, run_label, questions, samples, out_f, done, max_new, trust, device, locale=class="s">"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 = class="s">"NVMe" if load_root == CACHE else class="s">"NFS" if path is None: path, src = model_id, class="s">"hub" print(fclass="s">"[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 = class="s">"auto" if device == class="s">"auto" else {class="s">"": 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: class="s">"78GiB" for i in range(torch.cuda.device_count())} | {class="s">"cpu": class="s">"220GiB"}) if device == class="s">"auto" else None model = AutoModelForCausalLM.from_pretrained(path, dtype=torch.bfloat16, device_map=dm, max_memory=mm, attn_implementation=class="s">"sdpa", trust_remote_code=trust) model.eval() print(fclass="s">"[loaded] {model_id} in {time.time()-t0:.0f}s on {model.device}", file=sys.stderr, flush=True) for q in questions: qid, query = q[class="s">"query_id"], q[class="s">"query"] for s in range(samples): key = (run_label, qid, s) if key in done: continue answer = gen(model, tok, [ {class="s">"role": class="s">"system", class="s">"content": prompts.BASE_SYSTEM_PROMPT}, {class="s">"role": class="s">"user", class="s">"content": prompts.BASE_USER_TEMPLATE.format(locale=locale, query=query)}], max_new) cold = gen(model, tok, [ {class="s">"role": class="s">"system", class="s">"content": criteria.COLD_CRITERIA_SYSTEM_PROMPT}, {class="s">"role": class="s">"user", class="s">"content": cold_user(query)}], max_new, temperature=0.0) posthoc = gen(model, tok, [ {class="s">"role": class="s">"system", class="s">"content": prompts.BASE_SYSTEM_PROMPT}, {class="s">"role": class="s">"user", class="s">"content": prompts.BASE_USER_TEMPLATE.format(locale=locale, query=query)}, {class="s">"role": class="s">"assistant", class="s">"content": answer}, {class="s">"role": class="s">"user", class="s">"content": criteria.POSTHOC_CRITERIA_FOLLOWUP}], max_new, temperature=0.0) rec = {class="s">"model": model_id, class="s">"run_label": run_label, class="s">"query_id": qid, class="s">"sample": s, class="s">"expected_criteria": q.get(class="s">"expected_criteria", class="s">""), class="s">"intent": q.get(class="s">"intent", class="s">""), class="s">"product_category": q.get(class="s">"product_category", class="s">""), class="s">"answer_text": answer, class="s">"criteria_cold": parse(cold), class="s">"criteria_posthoc": parse(posthoc), class="s">"ts": time.time()} out_f.write(json.dumps(rec) + class="s">"\n"); out_f.flush() print(fclass="s">" [{run_label}] {qid} s{s}: ans={len(answer)}c " fclass="s">"cold={len(rec['criteria_cold']) if isinstance(rec['criteria_cold'],list) else 'ERR'} " fclass="s">"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(class="s">"--models", required=True, help=class="s">"comma-separated HF ids, the spectrum in order") ap.add_argument(class="s">"--questions", required=True); ap.add_argument(class="s">"--samples", type=int, default=1) ap.add_argument(class="s">"--label-suffix", default=class="s">"", help=class="s">"append to run_label, e.g. -r2 for a C0 replicate") ap.add_argument(class="s">"--out", required=True); ap.add_argument(class="s">"--no-prefetch", action=class="s">"store_true") ap.add_argument(class="s">"--max-new", type=int, default=768); ap.add_argument(class="s">"--trust", action=class="s">"store_true") ap.add_argument(class="s">"--device", default=class="s">"cuda:0") a = ap.parse_args() models = [m.strip() for m in a.models.split(class="s">",") 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[class="s">"run_label"], d[class="s">"query_id"], d[class="s">"sample"])) except Exception: pass print(fclass="s">"[spectrum] {len(models)} models, {len(questions)} questions x {a.samples} " fclass="s">"samples; {len(done)} already done", file=sys.stderr, flush=True) def label(m): return m.split(class="s">"/")[-1].lower() + a.label_suffix out_f = open(a.out, class="s">"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(fclass="s">"[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(fclass="s">"[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(fclass="s">"{CACHE}/{hubname}", ignore_errors=True) for th in prefetch_threads: th.join(timeout=1) out_f.close() print(class="s">"[spectrum] done", file=sys.stderr, flush=True) if __name__ == class="s">"__main__": main()