day26_code/rest_em.py

"""rest_em.py — ReST-EM expert iteration (Singh et al. 2024, arXiv:2312.06585), the published method
whose RESULT is HELD-OUT TRANSFER — the exact thing SFT/GRPO/DPO failed at here.

THE load-bearing difference vs everything we tried: each Improve round trains a FRESH LoRA on the FROZEN
BASE — NEVER stacked on the prior adapter. Stacking (our cumulative ratchet) is what degraded held-out
(holdout_split_probe: train-on-distribution overfits + hurts untrained). ReST-EM trains-from-base each
round, so the policy improves only because the DATA improves (the generator gets better → more/diverse
verified-correct solutions), not because weights compound into a memorized mode.

Loop (ReST-EM E-step / M-step):
  GENERATE (E): sample K diverse (high-temp) per TRAIN problem with the CURRENT-BEST generator; verifier
                (real exec) filters to CORRECT-only; dedup near-identical (diversity > volume per Singh);
                cap per problem; ACCUMULATE into a growing pool.
  IMPROVE  (M): SFT a FRESH LoRA on the FROZEN BASE over the whole filtered pool (+ small base-distribution
                replay to preserve pass@k diversity / avoid mode-collapse).
  RATCHET     : eval greedy@1 on the SEALED held-out (deterministic ruler). Keep best adapter as next
                generator. Compounding = held-out greedy@1 rises round-over-round.

Outcome reward only (unit tests). Deterministic eval (VLLM_DETERMINISTIC=1). Qwen2.5-72B base."""
import os, sys, json, subprocess, tempfile, random, hashlib, pathlib, signal
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.config import SystemConfig, ModelConfig
from src.utils.vllm_backend import VLLMModelLoader
from src.trainer.custom_lora import CustomLoRATrainer
from src.generator.data_generator import TrainingSample

MODEL=os.environ.get("MODEL","/workspace/RSI/expanded_models/qwen72")
BNB={"load_in_4bit":True,"bnb_4bit_compute_dtype":"bfloat16"}
PREREG="/workspace/RSI/outputs/prereg.json"
OUT=os.environ.get("RESTEM_OUT","/workspace/RSI/outputs/rest_em_slope.jsonl")
ADAPT=os.environ.get("RESTEM_ADAPT","/workspace/RSI/outputs/rest_em_adapters")
POOLFILE=os.environ.get("RESTEM_POOL","/workspace/RSI/outputs/rest_em_pool.json")  # persisted, resume-safe
N_CYCLES=int(os.environ.get("N_CYCLES","12")); N_TRAIN=int(os.environ.get("N_TRAIN","64")); K=int(os.environ.get("K","6"))
TEMP=float(os.environ.get("RESTEM_TEMP","1.0")); MAXTOK=int(os.environ.get("MAXTOK","640"))
GEN_CHUNK=int(os.environ.get("GEN_CHUNK","720")); EVAL_TC=int(os.environ.get("EVAL_TC","40")); EVAL_TO=int(os.environ.get("EVAL_TO","8"))
RANK=int(os.environ.get("RESTEM_RANK","32")); LR=float(os.environ.get("RESTEM_LR","1e-5"))
EPOCHS=int(os.environ.get("RESTEM_EPOCHS","2")); STEPS=int(os.environ.get("RESTEM_STEPS","120"))
MAXPER=int(os.environ.get("MAXPER","3"))  # diverse correct solutions kept per problem (dedup)
POOL_CAP=int(os.environ.get("POOL_CAP","800")); N_REPLAY=int(os.environ.get("N_REPLAY","16"))
_GW=int(os.environ.get("GRADE_WORKERS","48")); _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 _chash(c): return hashlib.md5(_norm(c).encode()).hexdigest()
def run_stdin(code,inp,t):
    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 Exception: pass
def grade(code,item,maxtc=None,to=EVAL_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=SystemConfig(); cfg.model=ModelConfig(model_path=MODEL,dtype="bfloat16",quantization_config=BNB)
    cfg.trainer.use_rslora=False; cfg.trainer.lora_rank=RANK; cfg.trainer.lora_alpha=RANK
    cfg.trainer.learning_rate=LR; cfg.trainer.num_epochs=EPOCHS; cfg.trainer.max_steps_per_cycle=STEPS; cfg.trainer.training_mode="sft"
    for attr,val in (("gradient_accumulation_steps",1),("batch_size",2)):
        try: setattr(cfg.trainer,attr,val)
        except Exception: pass
    loader=VLLMModelLoader(model_path=MODEL,dtype="bfloat16",max_model_len=cfg.model.max_seq_length,
        gpu_memory_utilization=0.80,allow_remote_code=True,quantization_config=BNB,max_lora_rank=128,
        enable_chunked_prefill=False,enable_lora=True,enforce_eager=True)
    loader.load()
    gen_adapter=None  # the CURRENT-BEST generator (None=base for round 1)
    trainer=CustomLoRATrainer(cfg.trainer,loader); os.makedirs(ADAPT,exist_ok=True)
    def gen(ps,temp): return list(loader.generate_batch(ps,max_new_tokens=MAXTOK,temperature=temp,top_p=(1.0 if temp==0 else 0.95)))
    def gen_chunked(ps,temp):
        o=[]
        for c in range(0,len(ps),GEN_CHUNK): o.extend(gen(ps[c:c+GEN_CHUNK],temp))
        return o

    PR=json.load(open(PREREG)); HOLD=set(PR["hard_holdout"]); FRESH=set(PR.get("fresh_probe",[])); EASY=set(PR.get("easy",[]))
    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")]
    bh={_thash(it):it for it in apps}
    hold=[bh[t] for t in HOLD if t in bh]; easy=[bh[t] for t in EASY if t in bh]
    reserved=HOLD|FRESH|EASY; train=[it for it in apps if _thash(it) not in reserved]
    print(f"[restem] apps {len(apps)} | sealed holdout {len(hold)} | train {len(train)} | K={K} N_TRAIN={N_TRAIN} temp={TEMP} (FRESH-FROM-BASE each round)",flush=True)

    # small base-distribution replay anchor (easy verified-correct) to preserve diversity / avoid collapse
    replay=[]
    for it in easy:
        o=gen([it.prompt],0.0)[0]; c=_extract_code(o) or o
        if grade(c,it,EVAL_TC)>=0.999:
            replay.append(TrainingSample(prompt=it.prompt,response="```python\n"+c.strip()+"\n```",problem_id=_thash(it),verified=True,domain="apps"))
    print(f"[restem] base-replay anchor = {len(replay)} easy verified",flush=True)

    def holdout_greedy():
        outs=gen([it.prompt for it in hold],0.0)
        return sum(1 for it,o in zip(hold,outs) if grade(_extract_code(o) or o,it,EVAL_TC)>=0.999)
    base_solved=holdout_greedy(); best=base_solved; best_adapter=None
    print(f"[restem] BASELINE holdout greedy@1 = {base_solved}/{len(hold)} ({base_solved/len(hold):.4f})",flush=True)
    open(OUT,"a").write(json.dumps({"cycle":0,"holdout_solved":base_solved,"total":len(hold)})+"\n")

    pool={}  # problem_id -> list of (code_hash, TrainingSample) ; accumulates verified-correct DIVERSE solns
    if os.path.exists(POOLFILE):
        try:
            for r in json.load(open(POOLFILE)):
                pool.setdefault(r["pid"],[]).append((r["ch"],TrainingSample(prompt=r["prompt"],response=r["response"],problem_id=r["pid"],verified=True,domain="apps")))
            print(f"[restem] resumed pool: {sum(len(v) for v in pool.values())} solns / {len(pool)} problems",flush=True)
        except Exception as e: print(f"[restem] pool reload failed ({e})",flush=True)

    for c in range(1,N_CYCLES+1):
        random.seed(80000+c); batch=random.sample(train,min(N_TRAIN,len(train)))
        if gen_adapter and os.path.exists(gen_adapter): loader.set_lora_adapter(gen_adapter)
        else: loader.set_lora_adapter(None)
        # GENERATE (E-step): K diverse samples per problem with current-best generator
        prompts=[it.prompt for it in batch for _ in range(K)]
        outs=gen_chunked(prompts,TEMP)
        new=0
        for i,it in enumerate(batch):
            pid=_thash(it); have={h for h,_ in pool.get(pid,[])}
            for o in outs[i*K:(i+1)*K]:
                if len([1 for _ in pool.get(pid,[])])>=MAXPER: break
                code=_extract_code(o) or o; ch=_chash(code)
                if ch in have or not code.strip(): continue
                if grade(code,it,EVAL_TC)>=0.999:
                    pool.setdefault(pid,[]).append((ch,TrainingSample(prompt=it.prompt,response="```python\n"+code.strip()+"\n```",problem_id=pid,verified=True,domain="apps")))
                    have.add(ch); new+=1
        # persist pool
        flat=[{"pid":pid,"ch":ch,"prompt":s.prompt,"response":s.response} for pid,lst in pool.items() for ch,s in lst]
        json.dump(flat,open(POOLFILE,"w"))
        train_samples=[s for lst in pool.values() for _,s in lst][:POOL_CAP]
        n_prob=len(pool)
        print(f"[restem] c{c} E-step: +{new} new verified solns -> pool {len(train_samples)} solns / {n_prob} problems",flush=True)
        if len(train_samples)<8:
            print(f"[restem] c{c} pool<8 — skip improve",flush=True); open(OUT,"a").write(json.dumps({"cycle":c,"pool":len(train_samples),"skip":True})+"\n"); continue
        # IMPROVE (M-step): FRESH LoRA from BASE on the filtered pool (+ replay anchor)
        loader.swap_to_hf_for_training(); ckpt=None
        try:
            trainer.inject_lora()  # FRESH from base — the ReST-EM load-bearing choice (never load prior adapter)
            m=trainer.train(train_samples+replay[:N_REPLAY],c); ckpt=trainer.save_lora_weights(pathlib.Path(ADAPT),c)
        finally:
            loader.swap_to_vllm_after_training(adapter_path=str(ckpt) if ckpt else best_adapter)
        if ckpt is None: print(f"[restem] c{c} improve failed",flush=True); continue
        solved=holdout_greedy(); d=solved-base_solved
        rec={"cycle":c,"pool_solns":len(train_samples),"pool_probs":n_prob,"steps":getattr(m,"steps",0),"loss":round(getattr(m,"final_loss",0.0),4),
             "holdout_solved":solved,"total":len(hold),"d_vs_base":d,"vs_best":solved-best}
        open(OUT,"a").write(json.dumps(rec)+"\n"); print(f"[restem] c{c} {json.dumps(rec)}",flush=True)
        # RATCHET: best adapter becomes next generator
        if solved>=best:
            best=solved; best_adapter=str(ckpt); gen_adapter=str(ckpt)
            print(f"[restem] c{c} RATCHET UP best={best}/{len(hold)} ({best/len(hold):.4f}) gen<-c{c}",flush=True)
        else:
            print(f"[restem] c{c} {solved}<{best}; keep generating from best (c data still accumulates)",flush=True)
            if best_adapter: gen_adapter=best_adapter
    print(f"[restem] DONE base={base_solved} -> best={best}/{len(hold)} ({best/len(hold):.4f}) over {N_CYCLES} cycles",flush=True)
    sys.exit(0)

if __name__=="__main__":
    main()