Upload jobs/run_real_llm_standalone_v2.py

jobs/run_real_llm_standalone_v2.py

@@ -0,0 +1,460 @@
+"""
+Self-contained GPU job for real LLM code benchmark — V2 with Qwen chat templating.
+V1 failed because HumanEval sends raw Python stubs but Qwen-Coder expects chat formatting.
+This version wraps prompts in <|im_start|> markers.
+"""
+import json
+import os
+import random
+import re
+import subprocess
+import sys
+import tempfile
+import time
+from dataclasses import dataclass, field
+from enum import Enum
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
+
+from datasets import load_dataset
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+
+# ============== ORACLE (INLINE) ==============
+
+@dataclass
+class OracleResult:
+    raw_score: float
+    cost_adjusted_score: float
+    confidence: float
+    evidence: Dict[str, Any]
+    reason: str
+    failure_tags: List[str] = field(default_factory=list)
+    reward_value: float = 0.0
+
+
+class ImpactOracle:
+    def __init__(self, compute_penalty_rate=0.0001, gaming_penalty=2.0):
+        self.compute_penalty_rate = compute_penalty_rate
+        self.gaming_penalty = gaming_penalty
+
+    def score(self, mode, action, context, result, agent_id=""):
+        correctness = result.get("correctness", 0.0)
+        compute_cost = result.get("compute_cost", 0.0)
+        public_pass = result.get("public_pass", correctness)
+        hidden_pass = result.get("hidden_tests_pass", correctness)
+
+        failure_tags = []
+        if public_pass and not hidden_pass:
+            failure_tags.append("gaming_hidden_tests")
+
+        raw = correctness * 1.0 - compute_cost * self.compute_penalty_rate
+        if "gaming_hidden_tests" in failure_tags:
+            raw -= self.gaming_penalty
+
+        cost_adj = raw - compute_cost * self.compute_penalty_rate
+        return OracleResult(
+            raw_score=raw,
+            cost_adjusted_score=cost_adj,
+            confidence=result.get("confidence", correctness),
+            evidence={"correctness": correctness},
+            reason=f"correctness={correctness:.2f}, cost={compute_cost}",
+            failure_tags=failure_tags,
+            reward_value=cost_adj,
+        )
+
+
+# ============== LEDGER (INLINE) ==============
+
+@dataclass
+class LedgerEntry:
+    agent_id: str
+    task_id: str
+    action_id: str
+    earned_credit: float
+    spent_credit: float
+    decayed_credit: float
+    remaining_credit: float
+    reason: str
+    oracle_score: float
+    compute_cost: float
+    timestamp: float
+    capability_scope: str = "global"
+
+
+class CreditLedger:
+    def __init__(self, decay_lambda=0.05):
+        self.entries = []
+        self.balances = {}
+        self.decay_lambda = decay_lambda
+
+    def earn(self, agent_id, task_id, action_id, amount, oracle_score, compute_cost, reason, capability_scope="global"):
+        now = time.time()
+        self._apply_decay(agent_id, now, capability_scope)
+        current = self._get(agent_id, capability_scope)
+        new_bal = current + amount
+        self.entries.append(LedgerEntry(agent_id, task_id, action_id, amount, 0.0, 0.0, new_bal, reason, oracle_score, compute_cost, now, capability_scope))
+        self._set(agent_id, capability_scope, new_bal)
+
+    def spend(self, agent_id, task_id, action_id, amount, capability_scope="global", reason="spend"):
+        now = time.time()
+        self._apply_decay(agent_id, now, capability_scope)
+        current = self._get(agent_id, capability_scope)
+        if current < amount:
+            return False
+        new_bal = current - amount
+        self.entries.append(LedgerEntry(agent_id, task_id, action_id, 0.0, amount, 0.0, new_bal, reason, 0.0, 0.0, now, capability_scope))
+        self._set(agent_id, capability_scope, new_bal)
+        return True
+
+    def balance(self, agent_id, capability_scope="global"):
+        now = time.time()
+        self._apply_decay(agent_id, now, capability_scope)
+        return self._get(agent_id, capability_scope)
+
+    def _get(self, agent_id, cap):
+        return self.balances.get(agent_id, {}).get(cap, 0.0)
+
+    def _set(self, agent_id, cap, val):
+        if agent_id not in self.balances:
+            self.balances[agent_id] = {}
+        self.balances[agent_id][cap] = val
+
+    def _apply_decay(self, agent_id, now, cap):
+        current = self._get(agent_id, cap)
+        if current <= 0:
+            return
+        decayed = current * (1 - self.decay_lambda)
+        if decayed < current:
+            self.entries.append(LedgerEntry(agent_id, "decay", "decay", 0.0, 0.0, current - decayed, decayed, "credit_decay", 0.0, 0.0, now, cap))
+            self._set(agent_id, cap, decayed)
+
+
+# ============== BROKER (INLINE) ==============
+
+class Decision(Enum):
+    ALLOW = "allow"
+    DENY = "deny"
+    REQUIRE_APPROVAL = "require_approval"
+    DOWNGRADE = "downgrade"
+    ESCALATE = "escalate"
+    ASK_JUSTIFICATION = "ask_justification"
+
+
+@dataclass
+class ResourceDecision:
+    decision: Decision
+    reason: str
+    capability: str
+    downgrade_to: Optional[str] = None
+
+
+class ResourceBroker:
+    RESOURCE_RISK = {"model_call": "medium", "retrieval_call": "low", "verifier_call": "medium",
+                     "debate_turn": "low", "file_write": "high", "shell_execute": "high",
+                     "memory_write": "medium", "human_escalation": "high", "larger_model": "medium"}
+    DEFAULT_THRESHOLDS = {"low": 0.5, "medium": 2.0, "high": 5.0}
+
+    def __init__(self, thresholds=None, urgency_boost=0.5):
+        self.thresholds = thresholds or self.DEFAULT_THRESHOLDS.copy()
+        self.urgency_boost = urgency_boost
+        self.denial_history = {}
+
+    def request(self, capability, agent_id, credit_balance, task_state=None, risk_score=0.0, gaming_flags=None):
+        task_state = task_state or {}
+        gaming_flags = gaming_flags or []
+        risk_class = self.RESOURCE_RISK.get(capability, "medium")
+        threshold = self.thresholds.get(risk_class, 2.0)
+        urgency = task_state.get("urgency", 0.0)
+        adjusted = max(0.1, threshold - urgency * self.urgency_boost)
+
+        if gaming_flags:
+            return ResourceDecision(Decision.DENY, f"Gaming: {gaming_flags}", capability)
+        if risk_class == "high" and risk_score > 0.7:
+            return ResourceDecision(Decision.REQUIRE_APPROVAL, f"High risk {risk_score:.2f}", capability)
+        if credit_balance >= adjusted:
+            return ResourceDecision(Decision.ALLOW, f"Balance {credit_balance:.2f} >= {adjusted:.2f}", capability)
+        if credit_balance >= adjusted * 0.5:
+            if risk_class == "medium":
+                return ResourceDecision(Decision.DOWNGRADE, f"Downgrading from {capability}", capability, "retrieval_call")
+            return ResourceDecision(Decision.ASK_JUSTIFICATION, f"Justification required", capability)
+        denials = self.denial_history.get(agent_id, 0)
+        if denials > 3:
+            return ResourceDecision(Decision.ESCALATE, f"Denied {denials} times", capability)
+        self.denial_history[agent_id] = denials + 1
+        return ResourceDecision(Decision.DENY, f"Balance {credit_balance:.2f} < {adjusted:.2f}", capability)
+
+
+# ============== HELPERS ==============
+
+def extract_function_body(code, entry_point):
+    """Extract just the function body from generated code."""
+    pattern = rf'def\s+{re.escape(entry_point)}\s*\([^)]*\).*?(?=\n(?:def\s|\S|$))'
+    match = re.search(pattern, code, re.DOTALL)
+    if match:
+        return match.group(0)
+    # Fallback: take everything from 'def entry_point' to end
+    pattern2 = rf'def\s+{re.escape(entry_point)}\s*\([^)]*\).*'
+    match2 = re.search(pattern2, code, re.DOTALL)
+    return match2.group(0) if match2 else code
+
+
+def run_tests(code, test_code, timeout=5):
+    full = code + "\n\n" + test_code + "\n\ncheck()\n"
+    with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as f:
+        f.write(full)
+        tmp = f.name
+    try:
+        result = subprocess.run(['python', tmp], capture_output=True, text=True, timeout=timeout)
+        passed = result.returncode == 0
+        error = result.stderr if not passed else ""
+    except subprocess.TimeoutExpired:
+        passed = False
+        error = "Timeout"
+    except Exception as e:
+        passed = False
+        error = str(e)
+    finally:
+        os.unlink(tmp)
+    return passed, error
+
+
+# ============== QWEN CHAT TEMPLATE WRAPPER (V2 FIX) ==============
+
+QWEN_SYSTEM = (
+    "You are an expert Python programmer. "
+    "Complete the function definition. Return ONLY the completed function body — no explanations, no markdown formatting."
+)
+
+
+def wrap_prompt_chat(humaneval_prompt: str, tok: AutoTokenizer) -> str:
+    """
+    V2 FIX: Wrap raw HumanEval prompt as Qwen chat-formatted input.
+
+    HumanEval sends:
+        'def has_close_elements(numbers, threshold):\n    """ Check if...'
+
+    Qwen-Coder-Instruct expects:
+        <|im_start|>system
+        You are an expert Python programmer...
+        <|im_end|>
+        <|im_start|>user
+        def has_close_elements(numbers, threshold):
+            """ Check if...
+        <|im_end|>
+        <|im_start|>assistant
+    """
+    messages = [
+        {"role": "system", "content": QWEN_SYSTEM},
+        {"role": "user", "content": humaneval_prompt.strip()},
+    ]
+    # Use the tokenizer's chat template if available
+    if hasattr(tok, "apply_chat_template") and tok.chat_template:
+        formatted = tok.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True,  # adds <|im_start|>assistant at end
+        )
+        return formatted
+    else:
+        # Manual fallback for Qwen-style models
+        return (
+            f"<|im_start|>system\n{QWEN_SYSTEM}<|im_end|>\n"
+            f"<|im_start|>user\n{humaneval_prompt.strip()}<|im_end|>\n"
+            f"<|im_start|>assistant\n"
+        )
+
+
+# ============== BENCHMARK ==============
+
+class RealLLMBenchmarkV2:
+    def __init__(self, model_name="Qwen/Qwen2.5-Coder-0.5B-Instruct", n_problems=20, seed=42):
+        self.model_name = model_name
+        self.n_problems = n_problems
+        self.seed = seed
+        self.oracle = ImpactOracle()
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        print(f"Using device: {self.device}")
+
+    def load_problems(self):
+        ds = load_dataset("openai/openai_humaneval", split="test")
+        problems = []
+        for i, item in enumerate(ds):
+            if i >= self.n_problems:
+                break
+            problems.append({
+                "task_id": item["task_id"],
+                "prompt": item["prompt"],
+                "test": item["test"],
+                "entry_point": item["entry_point"],
+            })
+        return problems
+
+    def load_model(self):
+        print(f"Loading {self.model_name}...")
+        tok = AutoTokenizer.from_pretrained(self.model_name, trust_remote_code=True)
+        model = AutoModelForCausalLM.from_pretrained(
+            self.model_name,
+            trust_remote_code=True,
+            torch_dtype=torch.bfloat16 if self.device == "cuda" else torch.float32,
+            device_map="auto" if self.device == "cuda" else None,
+        )
+        if self.device == "cpu":
+            model = model.to("cpu").float()
+        print(f"Model loaded. Chat template: {bool(tok.chat_template)}")
+        return model, tok
+
+    def generate(self, model, tok, prompt_raw, max_new_tokens=256, temperature=0.7):
+        """Generate with chat template wrapping."""
+        chat_prompt = wrap_prompt_chat(prompt_raw, tok)
+        inputs = tok(chat_prompt, return_tensors="pt").to(model.device)
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=max_new_tokens,
+                temperature=temperature,
+                do_sample=True,
+                pad_token_id=tok.eos_token_id,
+            )
+        gen = tok.decode(outputs[0], skip_special_tokens=True)
+        # Strip the chat prompt prefix to get just the generated code
+        prompt_decoded = tok.decode(inputs.input_ids[0], skip_special_tokens=True)
+        code = gen[len(prompt_decoded):].strip()
+        return code
+
+    def run_baseline(self, problems, model, tok):
+        results = []
+        total_compute = 0
+        for problem in problems:
+            code = self.generate(model, tok, problem["prompt"], max_new_tokens=256)
+            tokens = len(tok.encode(code))
+            total_compute += tokens
+            func = extract_function_body(code, problem["entry_point"])
+            full = problem["prompt"] + func
+            passed, error = run_tests(full, problem["test"])
+            results.append({"task_id": problem["task_id"], "passed": passed, "tokens": tokens, "error": error})
+            print(f"  {problem['task_id']}: passed={passed}, tokens={tokens}")
+        return {
+            "accuracy": sum(1 for r in results if r["passed"]) / len(results),
+            "total_compute": total_compute,
+            "mean_tokens": total_compute / len(problems),
+            "results": results,
+        }
+
+    def run_occ(self, problems, model, tok):
+        ledger = CreditLedger(decay_lambda=0.02)
+        broker = ResourceBroker()
+        ledger.earn("code_agent", "seed", "seed", 25.0, 0.0, 0.0, "initial", "model_call")
+        results = []
+        total_compute = 0
+
+        for problem in problems:
+            budget_remaining = 2000
+            attempts = 0
+            passed = False
+            best_code = ""
+            best_score = -999
+
+            while budget_remaining > 100 and attempts < 3 and not passed:
+                attempts += 1
+                balance = ledger.balance("code_agent", "model_call")
+                dec = broker.request("model_call", "code_agent", balance,
+                                   task_state={"attempts": attempts, "budget_remaining": budget_remaining})
+                if dec.decision == Decision.DENY:
+                    break
+
+                temp = 0.5 if attempts == 1 else 0.8
+                max_tok = 128 if attempts == 1 else 256
+                code = self.generate(model, tok, problem["prompt"], max_new_tokens=max_tok, temperature=temp)
+                tokens = len(tok.encode(code))
+                budget_remaining -= tokens
+                total_compute += tokens
+
+                func = extract_function_body(code, problem["entry_point"])
+                full = problem["prompt"] + func
+                passed_now, error = run_tests(full, problem["test"])
+                score = 1.0 if passed_now else 0.0
+
+                oracle_res = self.oracle.score(
+                    mode="code",
+                    action={"attempt": attempts},
+                    context={},
+                    result={"correctness": score, "pass_at_k": score, "regression": False,
+                            "compute_cost": tokens, "public_pass": passed_now, "hidden_tests_pass": passed_now},
+                    agent_id="code_agent",
+                )
+
+                if oracle_res.raw_score > best_score:
+                    best_score = oracle_res.raw_score
+                    best_code = code
+                    passed = passed_now
+
+                if passed:
+                    ledger.earn("code_agent", problem["task_id"], f"att_{attempts}", 5.0, oracle_res.raw_score, tokens, "pass", "model_call")
+                else:
+                    ledger.spend("code_agent", problem["task_id"], f"att_{attempts}", 1.0, "model_call", reason="fail")
+
+                if attempts >= 2 and not passed:
+                    break
+
+            results.append({"task_id": problem["task_id"], "passed": passed, "attempts": attempts,
+                            "tokens_used": 2000 - budget_remaining, "best_score": best_score})
+            print(f"  {problem['task_id']}: passed={passed}, attempts={attempts}")
+
+        return {
+            "accuracy": sum(1 for r in results if r["passed"]) / len(results),
+            "total_compute": total_compute,
+            "mean_tokens": total_compute / len(problems),
+            "mean_attempts": sum(r["attempts"] for r in results) / len(results),
+            "results": results,
+        }
+
+    def run_all(self):
+        problems = self.load_problems()
+        print(f"Loaded {len(problems)} problems")
+        model, tok = self.load_model()
+
+        print("\n--- Baseline ---")
+        baseline = self.run_baseline(problems, model, tok)
+        print(f"Baseline: accuracy={baseline['accuracy']:.3f}, compute={baseline['total_compute']}")
+
+        print("\n--- OCC ---")
+        occ = self.run_occ(problems, model, tok)
+        print(f"OCC: accuracy={occ['accuracy']:.3f}, compute={occ['total_compute']}")
+
+        comp = {
+            "baseline_accuracy": baseline["accuracy"],
+            "occ_accuracy": occ["accuracy"],
+            "baseline_compute": baseline["total_compute"],
+            "occ_compute": occ["total_compute"],
+            "compute_reduction": 1.0 - (occ["total_compute"] / max(baseline["total_compute"], 1)),
+            "accuracy_delta": occ["accuracy"] - baseline["accuracy"],
+        }
+
+        return {"baseline_fixed": baseline, "occ_budget": occ, "comparison": comp}
+
+
+def main():
+    bench = RealLLMBenchmarkV2(n_problems=20, seed=42)
+    results = bench.run_all()
+
+    print("\n" + "=" * 60)
+    print("REAL LLM CODE BENCHMARK (V2 — Chat Template Fix)")
+    print("=" * 60)
+    comp = results["comparison"]
+    print(f"Baseline accuracy: {comp['baseline_accuracy']:.3f}")
+    print(f"OCC accuracy:      {comp['occ_accuracy']:.3f}")
+    print(f"Baseline compute:  {comp['baseline_compute']}")
+    print(f"OCC compute:       {comp['occ_compute']}")
+    print(f"Compute reduction: {comp['compute_reduction']:.1%}")
+    print(f"Accuracy delta:    {comp['accuracy_delta']:+.3f}")
+
+    out_dir = Path("/app/occ/reports")
+    out_dir.mkdir(parents=True, exist_ok=True)
+    with open(out_dir / "benchmark_code_real_llm_v2_results.json", "w") as f:
+        json.dump(results, f, indent=2, default=str)
+    print(f"\nSaved to {out_dir / 'benchmark_code_real_llm_v2_results.json'}")
+
+
+if __name__ == "__main__":
+    main()