day26_code/soar_loop.py

"""soar_loop.py — the compounding RSI loop: Search -> Oracle -> Amortize -> Ratchet (iterative DPO).

Goal: >=1% COMPOUNDING per <=20min cycle on the SEALED held-out greedy@1 ruler. Method is NEITHER plain
SFT NOR GRPO (both shown not to compound on 32B). This is ITERATIVE DPO EXPERT-ITERATION:

  SEARCH   : best-of-K sampling per TRAIN problem (disjoint from sealed holdout), temp 0.9 (proven sweet spot)
  ORACLE   : real-execution verifier labels each candidate pass/fail (non-gameable ground truth)
  AMORTIZE : build PreferencePairs (verified-pass = chosen, fail = rejected) and DPO-update the policy,
             KL-anchored (dpo_beta) so it shifts toward winners without collapsing — the on-policy
             preference signal SFT/GRPO lacked
  RATCHET  : eval greedy@1 on the SEALED held-out (deterministic ruler); keep the adapter only if it does
             not regress; track best. Compounding = held-out greedy@1 rises >=1% relative each cycle.

Runs on the strongest base we have (Qwen2.5-72B-4bit). Cumulative LoRA across cycles. vLLM for gen/eval,
swap to HF for DPO, swap back. Deterministic eval (VLLM_DETERMINISTIC=1) so 1% is readable above noise."""
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 PreferencePair

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("SOAR_OUT","/workspace/RSI/outputs/soar_loop_slope.jsonl")
ADAPT=os.environ.get("SOAR_ADAPT","/workspace/RSI/outputs/soar_loop_adapters")
N_CYCLES=int(os.environ.get("N_CYCLES","10")); N_TRAIN=int(os.environ.get("N_TRAIN","48")); K=int(os.environ.get("K","6"))
TEMP=float(os.environ.get("SOAR_TEMP","0.9")); 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("SOAR_RANK","32")); LR=float(os.environ.get("SOAR_LR","5e-6")); DPO_BETA=float(os.environ.get("DPO_BETA","0.1"))
MAXPAIRS_PER_PROB=int(os.environ.get("MAXPAIRS_PER_PROB","2")); _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 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=int(os.environ.get("SOAR_EPOCHS","1"))
    cfg.trainer.max_steps_per_cycle=int(os.environ.get("SOAR_STEPS","60")); cfg.trainer.training_mode="dpo"; cfg.trainer.dpo_beta=DPO_BETA
    for attr,val in (("gradient_accumulation_steps",1),("batch_size",1)):
        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()
    cur_adapter=os.environ.get("SOAR_RESUME","") or None
    if cur_adapter and os.path.exists(cur_adapter): loader.set_lora_adapter(cur_adapter)
    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]
    reserved=HOLD|FRESH|EASY; train=[it for it in apps if _thash(it) not in reserved]
    print(f"[soar] apps {len(apps)} | sealed holdout {len(hold)} | train pool {len(train)} | K={K} N_TRAIN={N_TRAIN} temp={TEMP} dpo_beta={DPO_BETA}",flush=True)

    def holdout_greedy():  # deterministic greedy@1 on sealed holdout = the compounding ruler
        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
    print(f"[soar] 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")

    for c in range(1,N_CYCLES+1):
        random.seed(70000+c); batch=random.sample(train,min(N_TRAIN,len(train)))
        # SEARCH: K samples per train problem
        prompts=[it.prompt for it in batch for _ in range(K)]
        outs=gen_chunked(prompts,TEMP)
        # ORACLE + build preference pairs
        pairs=[]; n_pass=0; n_solvable=0
        for i,it in enumerate(batch):
            cand=outs[i*K:(i+1)*K]
            passes=[]; fails=[]
            for o in cand:
                code=_extract_code(o) or o
                (passes if grade(code,it,EVAL_TC)>=0.999 else fails).append(o)
            n_pass+=len(passes)
            if passes and fails:
                n_solvable+=1
                for j in range(min(MAXPAIRS_PER_PROB,len(passes),len(fails))):
                    pairs.append(PreferencePair(prompt=it.prompt,chosen_response=passes[j],rejected_response=fails[j],domain="apps"))
        print(f"[soar] c{c} search: {n_pass} verified passes over {len(batch)} probs; {n_solvable} yielded pairs; {len(pairs)} DPO pairs",flush=True)
        if len(pairs)<4:
            print(f"[soar] c{c} <4 pairs — skip train this cycle (search too sparse)",flush=True)
            open(OUT,"a").write(json.dumps({"cycle":c,"pairs":len(pairs),"skip":True})+"\n"); continue
        # AMORTIZE: DPO update (KL-anchored)
        loader.swap_to_hf_for_training(); ckpt=None
        try:
            if cur_adapter and os.path.exists(cur_adapter): trainer.load_lora_weights(cur_adapter)
            else: trainer.inject_lora()
            m=trainer.train([],c,preference_pairs=pairs); ckpt=trainer.save_lora_weights(pathlib.Path(ADAPT),c)
        finally:
            loader.swap_to_vllm_after_training(adapter_path=str(ckpt) if ckpt else cur_adapter)
        if ckpt is None: print(f"[soar] c{c} DPO failed",flush=True); continue
        # RATCHET: held-out greedy@1
        solved=holdout_greedy(); margin=getattr(m,"avg_reward_margin",0.0)
        d=solved-base_solved; rel=(solved-best)/max(1,best)
        rec={"cycle":c,"pairs":len(pairs),"steps":getattr(m,"steps",0),"reward_margin":round(margin,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"[soar] c{c} {json.dumps(rec)}",flush=True)
        if solved>=best:
            best=solved; prev=cur_adapter; cur_adapter=str(ckpt)
            print(f"[soar] c{c} RATCHET UP -> keep (best={best}/{len(hold)})",flush=True)
        else:
            loader.swap_to_vllm_after_training(adapter_path=cur_adapter)  # revert to best
            print(f"[soar] c{c} regressed ({solved}<{best}) -> revert, retry next cycle",flush=True)
    print(f"[soar] DONE base={base_solved} -> best={best}/{len(hold)} over {N_CYCLES} cycles",flush=True)
    sys.exit(0)

if __name__=="__main__":
    main()