Blame for day26_code/best_of_n.py

"""best_of_n.py — day-25 deliverable: inference-time best-of-N + real-execution verifier.

Day-24 verdict: Qwen3-32B+LoRA hits a CAPABILITY CEILING — no training method moves the sealed held-out
(greedy@1). But the base HAS above-greedy solutions (harvest cracked ~75% with search). This raises the
EFFECTIVE solve rate WITHOUT touching weights: sample N candidates per problem at temperature, grade each
with the REAL test oracle, KEEP the first that passes. That is the honest "smarter version" — pass@1 via
search+verify, sidestepping the wall.

Measures, on the SEALED prereg holdout (deterministic-eval base): greedy@1 baseline, then best-of-N for
N in {1,2,4,8,16,32} — solve_rate = fraction of problems where >=1 of N sampled candidates passes the FULL
real test suite (verifier-selected = deployable: we'd return that candidate). Reports the pass@N curve.
NO training. Pure inference. This is the number we'd actually ship."""
import os, sys, json, subprocess, tempfile, hashlib
sys.path.insert(0,"/workspace/RSI")
from concurrent.futures import ThreadPoolExecutor
from src.utils.external_benchmarks import _try_load_from_datasets, _extract_code
from src.utils.vllm_backend import VLLMModelLoader

MODEL=os.environ.get("MODEL","/workspace/RSI/expanded_models/qwen3_32b")
BNB={"load_in_4bit":True,"bnb_4bit_compute_dtype":"bfloat16"}
PREREG="/workspace/RSI/outputs/prereg.json"
OUT=os.environ.get("BON_OUT","/workspace/RSI/outputs/best_of_n.jsonl")
NS=[int(x) for x in os.environ.get("N_LIST","1,2,4,8,16,32").split(",")]
NMAX=max(NS)
TEMP=float(os.environ.get("BON_TEMP","0.8")); TOPP=float(os.environ.get("BON_TOPP","0.95"))
MAXTOK=int(os.environ.get("MAXTOK","640")); VERIFY_TC=int(os.environ.get("VERIFY_TC","40")); VERIFY_TO=int(os.environ.get("VERIFY_TO","8"))
SET=os.environ.get("BON_SET","hard_holdout")  # which prereg set to score
_GW=int(os.environ.get("GRADE_WORKERS","32")); _POOL=ThreadPoolExecutor(max_workers=_GW)

def _norm(s): return "\n".join(l.rstrip() for l in str(s).strip().splitlines())
def _thash(it): return hashlib.md5((it.prompt+repr(it.meta.get("inputs"))+repr(it.meta.get("outputs"))).encode()).hexdigest()
def run_stdin(code,inp,t):
    # Day-25: a candidate that spawns its own child (or ignores SIGTERM) made subprocess.run's timeout
    # leave an uninterruptible orphan -> ThreadPool deadlocked at N=64 grading. Run in a NEW SESSION and
    # kill the whole PROCESS GROUP on timeout so no candidate can wedge the grader.
    import signal
    d=tempfile.mkdtemp(); pth=os.path.join(d,"s.py"); open(pth,"w").write(code)
    try:
        p=subprocess.Popen([sys.executable,pth],stdin=subprocess.PIPE,stdout=subprocess.PIPE,
                           stderr=subprocess.DEVNULL,text=True,start_new_session=True)
        try:
            out,_=p.communicate(input=str(inp),timeout=t); return out
        except subprocess.TimeoutExpired:
            try: os.killpg(os.getpgid(p.pid),signal.SIGKILL)
            except Exception: pass
            try: p.communicate(timeout=2)
            except Exception: pass
            return ""
        except Exception: return ""
    except Exception: return ""
    finally:
        try: os.remove(pth); os.rmdir(d)
        except: pass
def grade(code,item,maxtc=None,to=VERIFY_TO):
    ins=item.meta.get("inputs") or []; outs=item.meta.get("outputs") or []
    if not code or not ins: return 0.0
    tc=list(zip(ins,outs))
    if maxtc: step=max(1,len(tc)//maxtc); tc=tc[::step][:maxtc]
    res=list(_POOL.map(lambda io:_norm(run_stdin(code,io[0],to))==_norm(io[1]), tc))
    return sum(res)/max(1,len(res))

def main():
    cfg_seq=4096
    loader=VLLMModelLoader(model_path=MODEL,dtype="bfloat16",max_model_len=cfg_seq,gpu_memory_utilization=0.85,
        allow_remote_code=True,quantization_config=BNB,max_lora_rank=128,enable_chunked_prefill=False,
        enable_lora=True,enforce_eager=True)
    loader.load()
    PR=json.load(open(PREREG)); IDS=set(PR[SET])
    apps=[it for it in (_try_load_from_datasets("apps") or []) if not (it.meta.get("fn_name") or "").strip() and it.meta.get("inputs") and it.meta.get("outputs")]
    items=[it for it in apps if _thash(it) in IDS]
    print(f"[bon] {SET}: {len(items)} problems | N in {NS} | temp={TEMP} | verify full-suite@{VERIFY_TO}s",flush=True)

    # greedy@1 (deterministic) for reference
    g_outs=list(loader.generate_batch([it.prompt for it in items],max_new_tokens=MAXTOK,temperature=0.0,top_p=1.0))
    greedy_pass=[grade(_extract_code(o) or o,it,VERIFY_TC)>=0.999 for it,o in zip(items,g_outs)]
    print(f"[bon] greedy@1 solved = {sum(greedy_pass)}/{len(items)} ({sum(greedy_pass)/len(items):.4f})",flush=True)

    # sample NMAX candidates per problem, CHUNKED so the batch never explodes the KV cache.
    # (one items*NMAX call OOM-killed the vLLM engine at N=64: 3840 seqs x 640 tok. Day-25 fix.)
    CHUNK=int(os.environ.get("GEN_CHUNK","960"))  # max sequences per generate_batch call
    prompts=[it.prompt for it in items for _ in range(NMAX)]
    s_outs=[]
    for cs in range(0,len(prompts),CHUNK):
        s_outs.extend(loader.generate_batch(prompts[cs:cs+CHUNK],max_new_tokens=MAXTOK,temperature=TEMP,top_p=TOPP))
        print(f"[bon] gen {min(cs+CHUNK,len(prompts))}/{len(prompts)} candidates",flush=True)
    # per-problem ordered pass/fail of its NMAX samples
    per_prob=[]
    for i,it in enumerate(items):
        cand=s_outs[i*NMAX:(i+1)*NMAX]
        passes=[grade(_extract_code(o) or o,it,VERIFY_TC)>=0.999 for o in cand]
        per_prob.append(passes)
    # best-of-N curve: solved iff any of first N sampled candidates passes (verifier picks it)
    curve={}
    for N in NS:
        solved=sum(1 for passes in per_prob if any(passes[:N]))
        curve[N]=solved
        rec={"N":N,"set":SET,"solved":solved,"total":len(items),"solve_rate":round(solved/len(items),4),"temp":TEMP}
        open(OUT,"a").write(json.dumps(rec)+"\n")
        print(f"[bon] best-of-{N}: {solved}/{len(items)} ({solved/len(items):.4f})",flush=True)
    g=sum(greedy_pass); bN=curve[NMAX]
    print(f"[bon] LIFT greedy@1={g}/{len(items)} -> best-of-{NMAX}={bN}/{len(items)} (+{bN-g} = +{(bN-g)/len(items):.1%} absolute on sealed holdout)",flush=True)
    print("[bon] DONE",flush=True); sys.exit(0)

if __name__=="__main__":
    main()