feat: HumanEval coding benchmark for TRINITY training + retry/timeout reliability fixes

.gitignore

@@ -1,3 +1,6 @@
+# machine-local Claude Code artifacts
+CLAUDE.local.md
+
 # caches / build
 __pycache__/
 *.pyc
@@ -25,3 +28,4 @@ results/*.out
 *.log
 conductor_*_out/
 conductor_toolscale_*/
+.tokensave

HANDOFF.md

@@ -0,0 +1,18 @@
+# HANDOFF: TRINITY materialized-weights loader
+
+Repo: /Users/russ/Projects/snowy2/OpenFugu (venv .venv with torch/transformers; `pip install safetensors huggingface_hub` into .venv if missing). Do NOT git commit.
+
+Problem: openfugu/mini.py loads a 19,456-dim CMA-ES vector (env FUGU_VECTOR, model_iter_60.npy) and reconstructs 9 SVF-adapted matrices + linear router head. That .npy no longer exists upstream. HF dataset `nshkrdotcom/trinity-coordinator-adapted-qwen3-0.6b` now ships MATERIALIZED weights:
+- checkpoints/0001_embedder.token_embedding.kernel.safetensors ... 0009_language_modeling_head.output.kernel.safetensors (embedder; decoder.blocks.26 self_attention query/key/value/output; ffn gate/intermediate/output; lm head)
+- router_head.safetensors (tensor `trinity.router_head.linear.weight`, [10,1024])
+- manifest.json: 7 agent_labels + 3 role_labels = 10 outputs, hidden 1024, opt layer [26]
+CAUTION: exported from Elixir/Bumblebee ("kernel" naming) — tensors may be TRANSPOSED vs PyTorch nn.Linear convention; verify by shape against the Qwen3-0.6B state dict and transpose where needed.
+
+Steps:
+1. Read openfugu/mini.py fully (FUGU_VECTOR load path, weight installation, head application) + docs/HOW_FUGU_IS_IMPLEMENTED.md (SVF + head sections).
+2. Download the 11 files via huggingface_hub.hf_hub_download(repo_type="dataset") into artifacts/materialized/ (~650MB).
+3. New file openfugu/materialized.py: `load_materialized(checkpoint_dir) -> (weights: dict[qwen_module_weight_name, torch.Tensor], head: torch.Tensor)`; map Bumblebee names → HF Qwen3 module paths as used by mini.py; handle dtype + transposition.
+4. Minimal hook in mini.py: env FUGU_CHECKPOINT_DIR set → install materialized weights + head instead of SVF reconstruction. Smallest possible diff.
+5. Verify: FUGU_MODEL=/Users/russ/.cache/huggingface/hub/models--Qwen--Qwen3-0.6B/snapshots/c1899de289a04d12100db370d81485cdf75e47ca FUGU_CHECKPOINT_DIR=$PWD/artifacts/materialized FUGU_FIXTURE=$PWD/artifacts/qwen_router_prompt_eval_cases.json .venv/bin/python -m openfugu.mini --self-test (check exact flag in mini.py). Target ≈95% agent / 100% role on the 37-case fixture; chance-level accuracy ⇒ suspect transposition/name mapping, iterate.
+
+Report: self-test numbers + mapping/transposition decisions.

openfugu/materialized.py

@@ -0,0 +1,94 @@
+import os
+import json
+import torch
+from safetensors.torch import load_file
+
+def load_materialized(checkpoint_dir: str) -> tuple[dict[str, torch.Tensor], torch.Tensor]:
+    """
+    Loads materialized weights and router head from the specified checkpoint directory.
+    Maps Bumblebee names -> HF Qwen3 PyTorch module paths.
+    Handles transposition for weights whose shapes do not match PyTorch convention.
+    """
+    manifest_path = os.path.join(checkpoint_dir, "manifest.json")
+    with open(manifest_path, "r") as f:
+        manifest = json.load(f)
+
+    # 1. Load the 9 selected tensors
+    weights = {}
+    selected_tensors = manifest.get("selected_tensors", [])
+
+    expected_shapes = {
+        "model.embed_tokens.weight": (151936, 1024),
+        "model.layers.26.self_attn.q_proj.weight": (2048, 1024),
+        "model.layers.26.self_attn.k_proj.weight": (1024, 1024),
+        "model.layers.26.self_attn.v_proj.weight": (1024, 1024),
+        "model.layers.26.self_attn.o_proj.weight": (1024, 2048),
+        "model.layers.26.mlp.gate_proj.weight": (3072, 1024),
+        "model.layers.26.mlp.up_proj.weight": (3072, 1024),
+        "model.layers.26.mlp.down_proj.weight": (1024, 3072),
+        "lm_head.weight": (151936, 1024),
+    }
+
+    for entry in selected_tensors:
+        chk_path = entry["checkpoint_path"]
+        abs_chk_path = os.path.join(checkpoint_dir, chk_path)
+        source_name = entry["source_name"]
+        path_key = entry["path"]
+
+        if not os.path.exists(abs_chk_path):
+            raise FileNotFoundError(f"Checkpoint file not found: {abs_chk_path}")
+
+        tensors_dict = load_file(abs_chk_path)
+        if path_key not in tensors_dict:
+            # fallback to checking if any key in tensors_dict contains the name
+            matching_keys = [k for k in tensors_dict if path_key in k]
+            if matching_keys:
+                tensor = tensors_dict[matching_keys[0]]
+            else:
+                raise KeyError(
+                    f"Tensor key {path_key!r} for {source_name!r} not found in "
+                    f"{abs_chk_path} (available keys: {sorted(tensors_dict)})"
+                )
+        else:
+            tensor = tensors_dict[path_key]
+
+        if source_name not in expected_shapes:
+            raise ValueError(f"No expected shape registered for tensor {source_name!r}")
+        expected_shape = expected_shapes[source_name]
+        if source_name in [
+            "model.layers.26.self_attn.k_proj.weight",
+            "model.layers.26.self_attn.v_proj.weight",
+        ]:
+            tensor = tensor.T.contiguous()
+        elif tensor.shape != expected_shape:
+            # Check if transposing makes shapes match
+            if tuple(tensor.shape[::-1]) == expected_shape:
+                tensor = tensor.T.contiguous()
+            else:
+                raise ValueError(
+                    f"Tensor {source_name} shape {tensor.shape} does not match expected shape {expected_shape}"
+                )
+
+        weights[source_name] = tensor
+
+    # 2. Load the router head
+    router_head_path = os.path.join(checkpoint_dir, "router_head.safetensors")
+    if not os.path.exists(router_head_path):
+        raise FileNotFoundError(f"Router head file not found: {router_head_path}")
+
+    router_head_dict = load_file(router_head_path)
+
+    # Try common keys
+    if "trinity_router_head" in router_head_dict:
+        head = router_head_dict["trinity_router_head"]
+    else:
+        head = list(router_head_dict.values())[0]
+
+    # Expected shape is (10, 1024)
+    if head.shape != (10, 1024):
+        if tuple(head.shape[::-1]) == (10, 1024):
+            head = head.T.contiguous()
+        else:
+            raise ValueError(f"Router head shape {head.shape} is not (10, 1024)")
+
+    return weights, head

openfugu/mini.py

@@ -105,17 +105,13 @@ class FuguRouter:
     which is what makes a routing decision ~one forward pass. [EXEC]
     """
 
-    def __init__(self, model_dir: str, vector_path: str, dtype="float32",
+    def __init__(self, model_dir: str, vector_path: str | None = None, dtype="float32",
                  device: str | None = None, seed: int | None = None):
         import torch
         from transformers import AutoModelForCausalLM, AutoTokenizer
         self.torch = torch
         self.rng = np.random.default_rng(seed)
 
-        vec = np.load(vector_path).astype(np.float64)
-        if vec.shape != (VEC_LEN,):
-            raise ValueError(f"router vector must be {VEC_LEN} floats, got {vec.shape}")
-
         self.tok = AutoTokenizer.from_pretrained(model_dir)
         # transformers >=5 uses dtype=, <5 uses torch_dtype= — support both
         td = getattr(torch, dtype)
@@ -127,9 +123,28 @@ def __init__(self, model_dir: str, vector_path: str, dtype="float32",
             self.model.to(device)
         self.device = next(self.model.parameters()).device
 
-        self._apply_svf(vec[:SVF_LEN])
-        # head: last 10240 -> (10, 1024) [EXEC]
-        self.head = torch.from_numpy(vec[SVF_LEN:].copy()).float().reshape(HEAD_ROWS, HIDDEN).to(self.device)
+        checkpoint_dir = os.environ.get("FUGU_CHECKPOINT_DIR")
+        if checkpoint_dir:
+            try:
+                from .materialized import load_materialized
+            except ImportError:
+                from materialized import load_materialized
+            weights, head = load_materialized(checkpoint_dir)
+            with torch.no_grad():
+                sd = self.model.state_dict()
+                for k, w in weights.items():
+                    sd[k].copy_(w.to(device=self.device, dtype=td))
+            self.svf_keys = list(weights.keys())
+            self.head = head.to(device=self.device, dtype=torch.float32)
+        else:
+            if not vector_path:
+                raise ValueError("vector_path must be provided if FUGU_CHECKPOINT_DIR is not set")
+            vec = np.load(vector_path).astype(np.float64)
+            if vec.shape != (VEC_LEN,):
+                raise ValueError(f"router vector must be {VEC_LEN} floats, got {vec.shape}")
+            self._apply_svf(vec[:SVF_LEN])
+            # head: last 10240 -> (10, 1024) [EXEC]
+            self.head = torch.from_numpy(vec[SVF_LEN:].copy()).float().reshape(HEAD_ROWS, HIDDEN).to(self.device)
 
     # SVF: scale only singular values, freeze U/V, energy-preserving renorm. [CODE]
     # Matrices consumed in state_dict order: embed_tokens, layer-26 {q,k,v,o,

train/train_trinity_coding.py

@@ -0,0 +1,214 @@
+#!/usr/bin/env python3
+# OpenFugu — Apache-2.0. Part of an independent, open reimplementation of
+# the Fugu orchestrator. NOT affiliated with Sakana AI. See NOTICE.
+# Reference: TRINITY (arXiv:2512.04695). Self-trains the coordinator via
+# sep-CMA-ES on REAL coding data (HumanEval) with a real worker pool — the
+# coding-benchmark variant of train_trinity_real.py. Original code.
+"""
+train_trinity_coding.py — TRINITY self-training on HumanEval coding benchmark.
+
+Same sep-CMA-ES loop as train_trinity_real.py, but for coding:
+  - tasks   : HumanEval problems (openai/openai_humaneval), split="test"
+  - features: a REAL Qwen3-0.6B penultimate hidden state of the prompt
+  - workers : a real pool via litellm (Novita), differing models
+  - reward  : real code execution (verifiable signal via Python subprocess execution)
+
+The coordinator (bias-free linear head over the hidden state) learns which
+worker to send each coding problem to, to maximize solved rate.
+"""
+from __future__ import annotations
+import argparse, os, re, subprocess, sys, tempfile, time
+import numpy as np
+
+HIDDEN = 1024
+HIDDEN_POS = -2
+
+
+def extract_completion(text: str) -> str:
+    """Extract candidate completion from fenced block or use raw response as-is."""
+    m = re.search(r"```python\n(.*?)\n?```", text, re.DOTALL)
+    if m:
+        return m.group(1)
+    m = re.search(r"```\n(.*?)\n?```", text, re.DOTALL)
+    if m:
+        return m.group(1)
+    return text
+
+
+class Backbone:
+    """Real Qwen3-0.6B -> penultimate hidden state of a question (the router feature)."""
+    def __init__(self, model_dir, device=None):
+        import torch
+        from transformers import AutoModelForCausalLM, AutoTokenizer
+        self.torch = torch
+        self.tok = AutoTokenizer.from_pretrained(model_dir)
+        try:
+            self.model = AutoModelForCausalLM.from_pretrained(model_dir, dtype=torch.float32).eval()
+        except TypeError:
+            self.model = AutoModelForCausalLM.from_pretrained(model_dir, torch_dtype=torch.float32).eval()
+        if device:
+            self.model.to(device)
+        self.device = next(self.model.parameters()).device
+        self._cache = {}
+
+    def feature(self, question: str) -> np.ndarray:
+        if question in self._cache:
+            return self._cache[question]
+        torch = self.torch
+        ids = self.tok(f"user: {question}", return_tensors="pt").to(self.device)
+        with torch.no_grad():
+            h = self.model.model(**ids).last_hidden_state[0, HIDDEN_POS, :]
+        v = h.float().cpu().numpy()
+        self._cache[question] = v
+        return v
+
+
+def route(head_vec, feat, n_workers):
+    """Bias-free linear head -> worker id (argmax), faithful to mini.py."""
+    W = head_vec.reshape(n_workers, HIDDEN)
+    return int(np.argmax(W @ feat))
+
+
+def main():
+    ap = argparse.ArgumentParser(description="TRINITY self-train on HumanEval (minimal).")
+    ap.add_argument("--model", default=os.environ.get("FUGU_MODEL", "Qwen/Qwen3-0.6B"))
+    ap.add_argument("--slot-models", required=True, help="csv of litellm worker ids (the pool)")
+    ap.add_argument("--n-train", type=int, default=40, help="HumanEval questions (kept small/cheap)")
+    ap.add_argument("--iters", type=int, default=25)
+    ap.add_argument("--sigma0", type=float, default=0.5)
+    ap.add_argument("--max-tokens", type=int, default=768)
+    ap.add_argument("--seed", type=int, default=42)
+    ap.add_argument("--out", default="trinity_humaneval.npy")
+    args = ap.parse_args()
+
+    import cma, litellm
+    from datasets import load_dataset
+
+    workers = args.slot_models.split(",")
+    n_workers = len(workers)
+    api_key = os.environ.get("FUGU_API_KEY") or os.environ.get("NOVITA_API_KEY") or os.environ.get("OPENAI_API_KEY")
+    api_base = os.environ.get("FUGU_BASE_URL") or os.environ.get("OPENAI_BASE_URL")
+
+    ds = load_dataset("openai/openai_humaneval", split="test")
+    tasks = list(ds)[:args.n_train]
+    print(f"[real-train] {len(tasks)} HumanEval tasks, {n_workers} workers: {workers}", flush=True)
+
+    bb = Backbone(args.model)
+    feats = [bb.feature(t["prompt"]) for t in tasks]  # cache real hidden states once
+    print(f"[real-train] cached {len(feats)} real Qwen3-0.6B features (dim {feats[0].shape[0]})", flush=True)
+
+    # worker call cache: (worker_id, task_id) -> solved? so CMA candidates reuse answers
+    solve_cache: dict = {}
+
+    def worker_solves(wid, task):
+        task_id = task["task_id"]
+        key = (wid, task_id)
+        if key in solve_cache:
+            return solve_cache[key]
+
+        prompt = task["prompt"]
+        test = task["test"]
+        entry_point = task["entry_point"]
+
+        worker_prompt = prompt + "\nComplete this Python function. Output ONLY the indented function body (do NOT repeat the `def` line or docstring), in a single ```python fenced code block, no explanation."
+
+        kw = dict(model=workers[wid],
+                  messages=[{"role": "user",
+                             "content": worker_prompt}],
+                  max_tokens=args.max_tokens, temperature=0.0, timeout=45)
+        if api_key: kw["api_key"] = api_key
+        if api_base: kw["api_base"] = api_base
+        ok = 0.0
+        call_succeeded = False
+        # Candidate code is untrusted model output, executed via subprocess (the whole
+        # point of HumanEval-style grading, not a shell/injection risk: no shell=True,
+        # args are a list). Give it an allowlisted env, not the parent's full env with
+        # secrets denylisted after the fact.
+        run_env = {k: os.environ[k] for k in ("PATH", "PYTHONPATH") if k in os.environ}
+        for attempt in range(5):
+            try:
+                out = litellm.completion(**kw).choices[0].message.content or ""
+                candidate = extract_completion(out)
+                call_succeeded = True
+
+                # Grading (the reward)
+                full_code = prompt + "\n" + candidate + "\n" + test + f"\ncheck({entry_point})\n"
+                tmp_file = tempfile.NamedTemporaryFile(suffix=".py", mode="w", delete=False, encoding="utf-8")
+                tmp_path = tmp_file.name
+                try:
+                    tmp_file.write(full_code)
+                    tmp_file.close()
+                    res = subprocess.run([sys.executable, tmp_path], capture_output=True, timeout=10, env=run_env)
+                    ok = 1.0 if res.returncode == 0 else 0.0
+                except subprocess.TimeoutExpired:
+                    ok = 0.0
+                except Exception as e:
+                    print(f"   [warn] execution error: {str(e)[:60]}", flush=True)
+                    ok = 0.0
+                finally:
+                    try:
+                        if os.path.exists(tmp_path):
+                            os.unlink(tmp_path)
+                    except Exception:
+                        pass
+                break
+            except litellm.RateLimitError:
+                wait = 2 ** attempt
+                print(f"   [rate-limit] worker {wid} 429, retry {attempt+1}/5 in {wait}s", flush=True)
+                time.sleep(wait)
+            except Exception as e:
+                print(f"   [warn] worker {wid} call failed: {str(e)[:60]}", flush=True)
+                break
+        if not call_succeeded:
+            # Transient/exhausted call failure, not a graded answer — don't poison the
+            # cache with a fake "unsolved" that would stick for the rest of the run.
+            print(f"   [call] worker {wid} task SKIPPED (call never succeeded)", flush=True)
+            return 0.0
+        solve_cache[key] = ok
+        print(f"   [call] worker {wid} task done -> {'OK' if ok else 'miss'} (cache={len(solve_cache)})", flush=True)
+        return ok
+
+    def fitness(head_vec):
+        tot = 0.0
+        for task, feat in zip(tasks, feats):
+            wid = route(head_vec, feat, n_workers)
+            tot += worker_solves(wid, task)
+        return tot / len(tasks)
+
+    # baseline: each worker alone + random (uses the same cache -> cheap)
+    rng = np.random.default_rng(args.seed)
+    per_worker = []
+    for w in range(n_workers):
+        per_worker.append(np.mean([worker_solves(w, t) for t in tasks]))
+    best_single = max(per_worker)
+    print("[baseline] per-worker solved rate: " +
+          ", ".join(f"{workers[w].split('/')[-1]}={per_worker[w]:.2f}" for w in range(n_workers)), flush=True)
+
+    # sep-CMA-ES over the head
+    dim = n_workers * HIDDEN
+    es = cma.CMAEvolutionStrategy(np.zeros(dim), args.sigma0,
+                                  {"seed": args.seed, "verbose": -9, "CMA_diagonal": True})
+    best_vec, best_fit = None, -1.0
+    for it in range(args.iters):
+        cands = es.ask()
+        fits = [fitness(c) for c in cands]
+        es.tell(cands, [-f for f in fits])
+        i = int(np.argmax(fits))
+        if fits[i] > best_fit:
+            best_fit, best_vec = fits[i], cands[i].copy()
+        print(f"[iter {it}] best_solved={best_fit:.3f}  "
+              f"(best single worker {best_single:.3f})  cache={len(solve_cache)}", flush=True)
+
+    np.save(args.out, best_vec)
+    print(f"\n[result] coordinator solved {best_fit:.3f} vs best single worker {best_single:.3f}")
+    print(f"[result] learned routing per task:")
+    for task, feat in zip(tasks, feats):
+        w = route(best_vec, feat, n_workers)
+        print(f"   -> {workers[w].split('/')[-1]:24s} | {task['prompt'][:50]}")
+    if best_fit >= best_single:
+        print("PASS — sep-CMA-ES self-trained TRINITY on REAL HumanEval, "
+              "coordinator >= best single worker")
+
+
+if __name__ == "__main__":
+    main()