Buckets:

AsianPlayer
/

TriDrive

Files

xet

AsianPlayer/TriDrive / code /train /train_cwm3_oof.py

AsianPlayer

16 days ago

download

raw

18.8 kB

"""cwm3/train/train_cwm3_oof.py — OUT-OF-FOLD joint-model predictions for the warning head's TRAIN label rows (coordinator, 2026-09-10).
Problem: the warning head is trained on features the joint model produced on TRAIN rows it was fitted on (in-sample predictions), which
inflates the device-grouped CV of the flexible heads on the predicted-future features.  Fix: K device folds over the labelled TRAIN devices
(91 devices, 979 consensus + 1,250 legacy rows); for fold k the three base branches (V0/can4s, U0/kin_can_noK, E0/pooled — the same stage
recipes, `train/train_cwm3.py` code called as a library, no file edited) and then the joint model (`train/train_cwm3_joint.py` as a library,
same I1 recipe / variant) are trained on every train row EXCEPT the rows of the fold's devices, and the fold's label rows are predicted by
that model.  Merged: results/runs/<joint run>/pred_{regime}_train_labelrows_oof.npz (same layout as the in-sample pack + `fold` per row).
Devices without labels stay in every training set (they are never held out).  Test / val features stay those of the full models.
  python train_cwm3_oof.py --seed 0 [--variant all] [--folds 5] [--fold k] [--out-root results/runs/oof]
  2026-09-12 (GPU-baselines agent, roadmap E8 on the FINAL model; every flag opt-in, defaults = the command above):
  python train_cwm3_oof.py --seed 0 --variant all_minus_road2drv --bases V0:can4s_stats2,U0:kin_can_noK_clip8_ln,E0:pooled_P --kin-source causal_routes --kin-version v2 \
      --pack-version v2 --out-root results/runs/oof_final --target-run I1_all_minus_road2drv_s0_final --gate-each --gate-all-trainers --max-others 2
  (--bases: the three base variants (stage:variant); the fold bases inherit the variants' recipes incl. `kclip` / `u_in_ln` / `p_head` / `select_LF` and the data
  flags; base_cfgs = models.cwm3.base_cfg_summary so the joint's batch prep uses the base's clamp / side vector; --target-run: the run dir that receives
  pred_{regime}_train_labelrows_oof.npz; --gate-each: wait for the launch gates before EVERY training (not only per fold); --gate-all-trainers: count
  train_kp / pretrain processes too (project rule: < 4 trainers alive across agents).)
Runs inline (no lock), one fold at a time; ≈ 10 min of GPU for the three bases + ≈ 8 min for the joint fine-tune per fold.  Resumable per
stage (TRAIN_DONE / pred file present -> skipped).  The launcher waits for the run_queue_joint.sh gates (< 2 trainers, RAM, VRAM).
No openpilot DMS field is read.
"""
from __future__ import annotations

import os

os.environ.setdefault("OMP_NUM_THREADS", "4")
os.environ.setdefault("MKL_NUM_THREADS", "4")
import argparse                                             # noqa: E402
import importlib.util                                       # noqa: E402
import json                                                 # noqa: E402
import subprocess                                           # noqa: E402
import sys                                                  # noqa: E402
import time                                                 # noqa: E402
from pathlib import Path                                    # noqa: E402
from types import SimpleNamespace                           # noqa: E402

import numpy as np                                          # noqa: E402

CWM3 = Path(__file__).resolve().parents[1]
for p in (CWM3, CWM3 / "data", CWM3 / "common", CWM3 / "train", CWM3 / "eval", CWM3.parent / "s12_wm"):
    sys.path.insert(0, str(p))
import cwm3_dataset as ds                                   # noqa: E402
import batch as bt                                          # noqa: E402
import train_cwm3 as T                                      # noqa: E402
import train_cwm3_joint as J                                # noqa: E402

log = T.log
BASES = [("V0", "can4s"), ("U0", "kin_can_noK"), ("E0", "pooled")]
BASES_DEFAULT = "V0:can4s,U0:kin_can_noK,E0:pooled"
DFLAGS = {"pack_version": "v1", "kin_source": "cwm2", "kin_version": "v1"}          # module-level data flags (set by main from the opt-in CLI flags)


def label_table():
    spec = importlib.util.spec_from_file_location("cwm3_labels", CWM3 / "data" / "labels.py")
    LB = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(LB)
    return LB.train_label_table(sources=("consensus", "legacy_v1"))


def device_folds(tab, K):
    dev = np.array(sorted(tab.dongle.unique()))
    return {d: i % K for i, d in enumerate(dev)}


def gates_ok(min_avail_gb=26, min_vram_gb=12, max_others=1, pattern="train_cwm3"):     # measured joint-trainer VRAM 7.6 GB + 4 GB headroom (owner rule)
    """same gates as train/run_queue_joint.sh; this process is excluded from the trainer count (it is not training while it waits).
    max_others / pattern (2026-09-12, opt-in): the trainer-count limit and the process pattern (default = the original train_cwm3-only count <= 1)."""
    me = os.getpid()
    n = int(subprocess.run(["bash", "-c", f'pids=$(pgrep -f "^[^ ]*python[^ ]* [^ ]*({pattern})" | grep -vx {me} || true); n=0; for p in $pids; do pp=$(awk \'/^PPid/{{print $2}}\' /proc/$p/status 2>/dev/null); '
                                          'grep -qx -- "$pp" <<< "$pids" || n=$((n+1)); done; echo $n'], capture_output=True, text=True).stdout.strip() or 0)
    avail = int(open("/proc/meminfo").read().split("MemAvailable:")[1].split()[0]) // 1048576
    vram = int(subprocess.run(["nvidia-smi", "--query-gpu=memory.free", "--format=csv,noheader,nounits"], capture_output=True, text=True).stdout.split()[0]) // 1024
    return n <= max_others and avail >= min_avail_gb and vram >= min_vram_gb, (n, avail, vram)


GATE = {"each": False, "kw": {}}       # --gate-each: wait for gates_ok(**kw) before every training (default off = the original per-fold gate)


def wait_gates(log_out, what):
    if not GATE["each"]:
        return
    while True:
        ok, stat = gates_ok(**GATE["kw"])
        if ok:
            log(f"[oof] gates ok for {what} (others, availGB, vramGB) = {stat}", log_out)
            return
        time.sleep(60)


# ----------------------------------------------------------------------------- base branches (train_cwm3.py as a library)
def train_base_oof(stage, variant, seed, held, out, log_out):
    if (out / "TRAIN_DONE").exists() and (out / "BEST").exists():
        return
    out.mkdir(parents=True, exist_ok=True)
    cfg = T.stage_cfg(stage, variant)
    cfg["pack_version"], cfg["kin_source"], cfg["kin_version"], cfg["mask_routes"] = DFLAGS["pack_version"], DFLAGS["kin_source"], DFLAGS["kin_version"], ""
    cfg["stats_version"] = int(bt.STATS_VERSION)
    if cfg["branch"] == "U" and "kclip" in cfg["recipe"]:                     # recipes that pin the clamp (clip8 survivors) record it, as train_cwm3.main does; others keep the legacy path
        cfg["kclip"] = float(cfg["recipe"]["kclip"])
    args = SimpleNamespace(seed=seed, out=str(out), epochs=0, minutes=600, cpu=False, no_amp=False, synthetic=0, max_anchors=0, no_lock=True)
    (out / "config.json").write_text(json.dumps({"cfg": cfg, "args": vars(args), "oof": {"held_devices": sorted(held), "data_flags": DFLAGS}}, indent=1, default=str))
    wait_gates(log_out, f"base {stage}/{variant} s{seed}")
    bank = T.load_bank(args, cfg)
    tr, va = T.load_raw(args, cfg, "train"), T.load_raw(args, cfg, "val")
    keep = np.flatnonzero(~np.isin(tr["dongle"], list(held)))
    n_all = len(tr["t"])
    tr = ds.subset(tr, keep)
    st = bt.compute_stats(tr, T.target_stats(args))
    if not (out / "stats.npz").exists():
        np.savez(out / "stats.npz", **st)
    log(f"[oof] base {stage}/{variant} s{seed}: train rows {len(keep):,}/{n_all:,} after holding out {len(held)} devices", log_out)
    code = T.run_chunk(args, cfg, tr, va, st, bank)
    assert code == 0 and (out / "TRAIN_DONE").exists(), f"base {stage}/{variant} fold training failed ({code})"
    del tr, va, bank


# ----------------------------------------------------------------------------- joint fine-tune + fold label rows (train_cwm3_joint.py as a library)
def train_joint_oof(seed, variant, held, fold_rows, base_root, suffix, out, log_out, bases=None):
    if all((out / f"pred_{r}_train_labelrows.npz").exists() for r in fold_rows):
        return
    out.mkdir(parents=True, exist_ok=True)
    bases = bases or dict(BASES)
    cfg = J.stage_cfg("I1", variant, base_v=bases.get("V0"), base_e=bases.get("E0"), base_u=bases.get("U0"))
    cfg["pack_version"], cfg["kin_source"], cfg["kin_version"] = DFLAGS["pack_version"], DFLAGS["kin_source"], DFLAGS["kin_version"]
    cfg["stats_mode"], cfg["base_suffix_e"], cfg["run_suffix"], cfg["base_dirs"] = "merged", "", "", {}
    args = SimpleNamespace(stage="I1", variant=variant, seed=seed, out=str(out), epochs=0, minutes=600, cpu=False, no_amp=False, synthetic=0, max_anchors=0,
                           no_lock=True, no_label_rows=False, base_v=bases.get("V0", ""), base_e=bases.get("E0", ""), base_u=bases.get("U0", ""), regimes="", pred_splits="",
                           kin_source=DFLAGS["kin_source"], kin_version=DFLAGS["kin_version"], pack_version=DFLAGS["pack_version"], stats_mode="merged", run_suffix="", base_suffix_e="")
    # library monkey-patches (module-level, this process only): OOF base dirs, no PRED_DONE requirement, device exclusion, fold label rows only
    J.base_dirs = lambda cfg_, seed_: {b: base_root / f"{st_}_{va_}_s{seed_}{suffix}" for b, (st_, va_) in cfg_["base"].items()}
    J.base_info_numpy = _base_info_numpy_oof
    orig_load_raw = J.load_raw

    def load_raw_oof(a, c, split, subset=True, regimes=None):
        raw = orig_load_raw(a, c, split, subset, regimes)
        if split == "train" and subset:
            keep = np.flatnonzero(~np.isin(raw["dongle"], list(held)))
            raw = ds.subset(raw, keep)
            raw["regime_row"], raw["regime_id"] = raw["regime_row"], raw["regime_id"]
        return raw
    J.load_raw = load_raw_oof
    J.label_rows = lambda a, c: {r: np.asarray(v, np.int64) for r, v in fold_rows.items()}
    J.load_pred_raws = lambda a, c, lrows: _label_raws_only(a, c, lrows, orig_load_raw)
    info = J.base_info_numpy(cfg, seed)
    cfg["e_side"] = info["base_cfgs"]["E"].get("side", "vf256")
    sv = {b: c.get("stats_version") for b, c in info["base_cfgs"].items()}
    cfg["stats_version"], cfg["base_stats_versions"] = (int(min(v for v in sv.values() if v is not None)) if any(v is not None for v in sv.values()) else 1), sv
    (out / "config.json").write_text(json.dumps({"cfg": cfg, "args": vars(args), "oof": {"held_devices": sorted(held), "base_root": str(base_root), "suffix": suffix, "data_flags": DFLAGS}}, indent=1, default=str))
    wait_gates(log_out, f"joint {variant} s{seed}")
    st = {k: np.asarray(v, np.float32) for k, v in info["stats"].items()}
    st.update(J.aux_stats_from_files(cfg["regimes"]))
    np.savez(out / "stats.npz", **st)
    (out / "base_ckpts.json").write_text(json.dumps({"base_ckpts": info["base_ckpts"], "base_cfgs": info["base_cfgs"], "tok_source": cfg.get("tok_source")}, indent=1))
    bank = J.load_bank(args, cfg, info)
    if not (out / "TRAIN_DONE").exists():
        tr, va = J.load_raw(args, cfg, "train"), J.load_raw(args, cfg, "val")
        log(f"[oof] joint {variant} s{seed}: train rows {len(tr['t']):,} (fold devices held out)", log_out)
        code = J.run_chunk(args, cfg, tr, va, st, bank, info)
        assert code == 0, f"joint fold training failed ({code})"
        del tr, va
    raws = J.load_pred_raws(args, cfg, fold_rows)
    code = J.predict_many(args, cfg, raws, st, bank, info)
    assert code == 0
    (out / "PRED_DONE").write_text(json.dumps({"label_rows_only": True, "regimes": list(fold_rows)}))


def _base_info_numpy_oof(cfg, seed):
    from models.cwm3 import merge_stats, base_cfg_summary
    names, cfgs, stats, dirs = {}, {}, [], {}
    for b, d in J.base_dirs(cfg, seed).items():
        assert (d / "BEST").exists() and (d / "TRAIN_DONE").exists(), f"OOF base for {b} missing: {d}"
        c = json.loads((d / "config.json").read_text())["cfg"]
        names[b] = f"{d.name}:{(d / 'BEST').read_text().strip()}"
        cfgs[b] = base_cfg_summary(c)                       # 2026-09-12: kclip / u_in_ln / p_head / select_LF / side / stats_version reach the joint's batch prep (as train_cwm3_joint.base_info_numpy)
        st = {k: v for k, v in np.load(d / "stats.npz").items()}
        if "stats_version" in st and "stats_version" not in cfgs[b]:
            cfgs[b]["stats_version"] = int(st["stats_version"])
        stats.append(st)
        dirs[b] = str(d)
    return {"base_ckpts": names, "base_cfgs": cfgs, "stats": merge_stats(stats), "stats_mode": "merged", "base_dirs": dirs}


def _label_raws_only(args, cfg, lrows, orig_load_raw):
    raws = {}
    for regime, rows in lrows.items():
        if len(rows) and not (Path(args.out) / f"pred_{regime}_train_labelrows.npz").exists():
            full = orig_load_raw(args, cfg, "train", subset=False, regimes=[regime])
            sub = ds.subset(full, rows)
            sub["_label_rows"] = rows
            raws[(regime, "train")] = sub
            del full
    return raws


# ----------------------------------------------------------------------------- merge
def merge(seed, variant, K, out_root, target_run, log_out):
    for regime in ("manual", "engaged"):
        packs = [np.load(out_root / f"I1_{variant}_s{seed}_f{k}" / f"pred_{regime}_train_labelrows.npz", allow_pickle=True) for k in range(K)
                 if (out_root / f"I1_{variant}_s{seed}_f{k}" / f"pred_{regime}_train_labelrows.npz").exists()]
        if len(packs) < K:
            log(f"[oof] merge {regime}: only {len(packs)}/{K} folds present — skipped", log_out)
            continue
        rows = np.concatenate([p["rows"] for p in packs])
        order = np.argsort(rows, kind="mergesort")
        fold = np.concatenate([np.full(len(p["rows"]), k, np.int8) for k, p in enumerate(packs)])[order]
        out = {}
        for key in packs[0].files:
            a0 = packs[0][key]
            if a0.ndim >= 1 and a0.shape[0] == len(packs[0]["rows"]) and key != "rows":
                out[key] = np.concatenate([p[key] for p in packs])[order]
            elif key == "rows":
                out[key] = rows[order]
            else:
                out[key] = a0
        out["fold"] = fold
        out["oof"] = np.array(json.dumps({"K": K, "folds": [str(out_root / f"I1_{variant}_s{seed}_f{k}") for k in range(K)]}))
        target = CWM3 / "results" / "runs" / target_run / f"pred_{regime}_train_labelrows_oof.npz"
        if target.exists():
            log(f"[oof] {target} exists — not overwritten", log_out)
            continue
        np.savez_compressed(target, **out)
        log(f"[oof] merged {regime}: {len(rows)} label rows from {K} folds -> {target}", log_out)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--variant", default="all")
    ap.add_argument("--folds", type=int, default=5)
    ap.add_argument("--fold", type=int, default=-1, help="single fold (default: all)")
    ap.add_argument("--out-root", default=str(CWM3 / "results" / "runs" / "oof"))
    ap.add_argument("--no-gate", action="store_true")
    ap.add_argument("--bases", default=BASES_DEFAULT, help="stage:variant triple of the base branches (default = the original V0/can4s, U0/kin_can_noK, E0/pooled)")
    ap.add_argument("--kin-source", default="cwm2", choices=["cwm2", "causal_routes", "oldin_ctgt"], help="data layer kin source for the fold bases AND the joint (default cwm2 = original)")
    ap.add_argument("--kin-version", default="v1", choices=["v1", "v2"])
    ap.add_argument("--pack-version", default="v1", choices=["v1", "v2"])
    ap.add_argument("--target-run", default="", help="run dir under results/runs that receives pred_{regime}_train_labelrows_oof.npz (default I1_<variant>_s<seed>)")
    ap.add_argument("--gate-each", action="store_true", help="wait for the launch gates before EVERY training (default: once per fold, original behaviour)")
    ap.add_argument("--gate-all-trainers", action="store_true", help="count train_kp / pretrain processes in the trainer gate too (project rule: < 4 trainers alive)")
    ap.add_argument("--max-others", type=int, default=1, help="trainer-count gate: other trainers alive <= this (default 1 = original)")
    ap.add_argument("--min-avail-gb", type=int, default=26)
    ap.add_argument("--min-vram-gb", type=int, default=12)
    a = ap.parse_args()
    global BASES
    BASES = [tuple(x.split(":")) for x in a.bases.split(",")]
    assert [b[0] for b in BASES] == ["V0", "U0", "E0"], a.bases
    DFLAGS.update(pack_version=a.pack_version, kin_source=a.kin_source, kin_version=a.kin_version)
    ds.set_pack_version(a.pack_version)
    ds.set_kin_source(a.kin_source, a.kin_version)
    gate_kw = dict(min_avail_gb=a.min_avail_gb, min_vram_gb=a.min_vram_gb, max_others=a.max_others, pattern=("train_cwm3|train_kp|pretrain" if a.gate_all_trainers else "train_cwm3"))
    GATE["each"], GATE["kw"] = bool(a.gate_each and not a.no_gate), gate_kw
    t0 = time.time()
    out_root = Path(a.out_root)
    out_root.mkdir(parents=True, exist_ok=True)
    log_out = str(out_root)
    tab = label_table()
    folds = device_folds(tab, a.folds)
    (out_root / f"folds_s{a.seed}.json").write_text(json.dumps({"K": a.folds, "device_fold": folds, "n_label_rows": int(len(tab))}, indent=1))
    log(f"[oof] {len(folds)} labelled devices -> {a.folds} folds; {len(tab)} label rows; bases {BASES}; data flags {DFLAGS}; target run {a.target_run or f'I1_{a.variant}_s{a.seed}'}", log_out)
    for k in range(a.folds):
        if a.fold >= 0 and k != a.fold:
            continue
        held = {d for d, f in folds.items() if f == k}
        fold_rows = {r: np.sort(tab[(tab.regime == r) & tab.dongle.isin(held)].row.unique()).astype(np.int64) for r in ("manual", "engaged")}
        while not a.no_gate:
            ok, stat = gates_ok(**gate_kw)
            if ok:
                break
            time.sleep(60)
        log(f"[oof] === fold {k}: {len(held)} devices, label rows manual {len(fold_rows['manual'])} engaged {len(fold_rows['engaged'])} ({time.time() - t0:.0f}s)", log_out)
        for stage, variant in BASES:
            train_base_oof(stage, variant, a.seed, held, out_root / f"{stage}_{variant}_s{a.seed}_f{k}", log_out)
        train_joint_oof(a.seed, a.variant, held, fold_rows, out_root, f"_f{k}", out_root / f"I1_{a.variant}_s{a.seed}_f{k}", log_out, bases=dict(BASES))
    if a.fold < 0:
        merge(a.seed, a.variant, a.folds, out_root, a.target_run or f"I1_{a.variant}_s{a.seed}", log_out)
    log(f"[oof] DONE in {(time.time() - t0) / 60:.1f} min", log_out)


if __name__ == "__main__":
    main()

Xet Storage Details

Size:: 18.8 kB
Xet hash:: 6077934b0abb8473266c3516652f96bc0387e1cab2ed7f4102d0c20f9c44895a

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.

	"""cwm3/train/train_cwm3_oof.py — OUT-OF-FOLD joint-model predictions for the warning head's TRAIN label rows (coordinator, 2026-09-10).

	Problem: the warning head is trained on features the joint model produced on TRAIN rows it was fitted on (in-sample predictions), which
	inflates the device-grouped CV of the flexible heads on the predicted-future features. Fix: K device folds over the labelled TRAIN devices
	(91 devices, 979 consensus + 1,250 legacy rows); for fold k the three base branches (V0/can4s, U0/kin_can_noK, E0/pooled — the same stage
	recipes, `train/train_cwm3.py` code called as a library, no file edited) and then the joint model (`train/train_cwm3_joint.py` as a library,
	same I1 recipe / variant) are trained on every train row EXCEPT the rows of the fold's devices, and the fold's label rows are predicted by
	that model. Merged: results/runs/<joint run>/pred_{regime}_train_labelrows_oof.npz (same layout as the in-sample pack + `fold` per row).
	Devices without labels stay in every training set (they are never held out). Test / val features stay those of the full models.
	python train_cwm3_oof.py --seed 0 [--variant all] [--folds 5] [--fold k] [--out-root results/runs/oof]
	2026-09-12 (GPU-baselines agent, roadmap E8 on the FINAL model; every flag opt-in, defaults = the command above):
	python train_cwm3_oof.py --seed 0 --variant all_minus_road2drv --bases V0:can4s_stats2,U0:kin_can_noK_clip8_ln,E0:pooled_P --kin-source causal_routes --kin-version v2 \
	--pack-version v2 --out-root results/runs/oof_final --target-run I1_all_minus_road2drv_s0_final --gate-each --gate-all-trainers --max-others 2
	(--bases: the three base variants (stage:variant); the fold bases inherit the variants' recipes incl. `kclip` / `u_in_ln` / `p_head` / `select_LF` and the data
	flags; base_cfgs = models.cwm3.base_cfg_summary so the joint's batch prep uses the base's clamp / side vector; --target-run: the run dir that receives
	pred_{regime}_train_labelrows_oof.npz; --gate-each: wait for the launch gates before EVERY training (not only per fold); --gate-all-trainers: count
	train_kp / pretrain processes too (project rule: < 4 trainers alive across agents).)
	Runs inline (no lock), one fold at a time; ≈ 10 min of GPU for the three bases + ≈ 8 min for the joint fine-tune per fold. Resumable per
	stage (TRAIN_DONE / pred file present -> skipped). The launcher waits for the run_queue_joint.sh gates (< 2 trainers, RAM, VRAM).
	No openpilot DMS field is read.
	"""
	from __future__ import annotations

	import os

	os.environ.setdefault("OMP_NUM_THREADS", "4")
	os.environ.setdefault("MKL_NUM_THREADS", "4")
	import argparse # noqa: E402
	import importlib.util # noqa: E402
	import json # noqa: E402
	import subprocess # noqa: E402
	import sys # noqa: E402
	import time # noqa: E402
	from pathlib import Path # noqa: E402
	from types import SimpleNamespace # noqa: E402

	import numpy as np # noqa: E402

	CWM3 = Path(__file__).resolve().parents[1]
	for p in (CWM3, CWM3 / "data", CWM3 / "common", CWM3 / "train", CWM3 / "eval", CWM3.parent / "s12_wm"):
	sys.path.insert(0, str(p))
	import cwm3_dataset as ds # noqa: E402
	import batch as bt # noqa: E402
	import train_cwm3 as T # noqa: E402
	import train_cwm3_joint as J # noqa: E402

	log = T.log
	BASES = [("V0", "can4s"), ("U0", "kin_can_noK"), ("E0", "pooled")]
	BASES_DEFAULT = "V0:can4s,U0:kin_can_noK,E0:pooled"
	DFLAGS = {"pack_version": "v1", "kin_source": "cwm2", "kin_version": "v1"} # module-level data flags (set by main from the opt-in CLI flags)


	def label_table():
	spec = importlib.util.spec_from_file_location("cwm3_labels", CWM3 / "data" / "labels.py")
	LB = importlib.util.module_from_spec(spec)
	spec.loader.exec_module(LB)
	return LB.train_label_table(sources=("consensus", "legacy_v1"))


	def device_folds(tab, K):
	dev = np.array(sorted(tab.dongle.unique()))
	return {d: i % K for i, d in enumerate(dev)}


	def gates_ok(min_avail_gb=26, min_vram_gb=12, max_others=1, pattern="train_cwm3"): # measured joint-trainer VRAM 7.6 GB + 4 GB headroom (owner rule)
	"""same gates as train/run_queue_joint.sh; this process is excluded from the trainer count (it is not training while it waits).
	max_others / pattern (2026-09-12, opt-in): the trainer-count limit and the process pattern (default = the original train_cwm3-only count <= 1)."""
	me = os.getpid()
	n = int(subprocess.run(["bash", "-c", f'pids=$(pgrep -f "^[^ ]python[^ ] [^ ]*({pattern})" \| grep -vx {me} \|\| true); n=0; for p in $pids; do pp=$(awk \'/^PPid/{{print $2}}\' /proc/$p/status 2>/dev/null); '
	'grep -qx -- "$pp" <<< "$pids" \|\| n=$((n+1)); done; echo $n'], capture_output=True, text=True).stdout.strip() or 0)
	avail = int(open("/proc/meminfo").read().split("MemAvailable:")[1].split()[0]) // 1048576
	vram = int(subprocess.run(["nvidia-smi", "--query-gpu=memory.free", "--format=csv,noheader,nounits"], capture_output=True, text=True).stdout.split()[0]) // 1024
	return n <= max_others and avail >= min_avail_gb and vram >= min_vram_gb, (n, avail, vram)


	GATE = {"each": False, "kw": {}} # --gate-each: wait for gates_ok(**kw) before every training (default off = the original per-fold gate)


	def wait_gates(log_out, what):
	if not GATE["each"]:
	return
	while True:
	ok, stat = gates_ok(**GATE["kw"])
	if ok:
	log(f"[oof] gates ok for {what} (others, availGB, vramGB) = {stat}", log_out)
	return
	time.sleep(60)


	# ----------------------------------------------------------------------------- base branches (train_cwm3.py as a library)
	def train_base_oof(stage, variant, seed, held, out, log_out):
	if (out / "TRAIN_DONE").exists() and (out / "BEST").exists():
	return
	out.mkdir(parents=True, exist_ok=True)
	cfg = T.stage_cfg(stage, variant)
	cfg["pack_version"], cfg["kin_source"], cfg["kin_version"], cfg["mask_routes"] = DFLAGS["pack_version"], DFLAGS["kin_source"], DFLAGS["kin_version"], ""
	cfg["stats_version"] = int(bt.STATS_VERSION)
	if cfg["branch"] == "U" and "kclip" in cfg["recipe"]: # recipes that pin the clamp (clip8 survivors) record it, as train_cwm3.main does; others keep the legacy path
	cfg["kclip"] = float(cfg["recipe"]["kclip"])
	args = SimpleNamespace(seed=seed, out=str(out), epochs=0, minutes=600, cpu=False, no_amp=False, synthetic=0, max_anchors=0, no_lock=True)
	(out / "config.json").write_text(json.dumps({"cfg": cfg, "args": vars(args), "oof": {"held_devices": sorted(held), "data_flags": DFLAGS}}, indent=1, default=str))
	wait_gates(log_out, f"base {stage}/{variant} s{seed}")
	bank = T.load_bank(args, cfg)
	tr, va = T.load_raw(args, cfg, "train"), T.load_raw(args, cfg, "val")
	keep = np.flatnonzero(~np.isin(tr["dongle"], list(held)))
	n_all = len(tr["t"])
	tr = ds.subset(tr, keep)
	st = bt.compute_stats(tr, T.target_stats(args))
	if not (out / "stats.npz").exists():
	np.savez(out / "stats.npz", **st)
	log(f"[oof] base {stage}/{variant} s{seed}: train rows {len(keep):,}/{n_all:,} after holding out {len(held)} devices", log_out)
	code = T.run_chunk(args, cfg, tr, va, st, bank)
	assert code == 0 and (out / "TRAIN_DONE").exists(), f"base {stage}/{variant} fold training failed ({code})"
	del tr, va, bank


	# ----------------------------------------------------------------------------- joint fine-tune + fold label rows (train_cwm3_joint.py as a library)
	def train_joint_oof(seed, variant, held, fold_rows, base_root, suffix, out, log_out, bases=None):
	if all((out / f"pred_{r}_train_labelrows.npz").exists() for r in fold_rows):
	return
	out.mkdir(parents=True, exist_ok=True)
	bases = bases or dict(BASES)
	cfg = J.stage_cfg("I1", variant, base_v=bases.get("V0"), base_e=bases.get("E0"), base_u=bases.get("U0"))
	cfg["pack_version"], cfg["kin_source"], cfg["kin_version"] = DFLAGS["pack_version"], DFLAGS["kin_source"], DFLAGS["kin_version"]
	cfg["stats_mode"], cfg["base_suffix_e"], cfg["run_suffix"], cfg["base_dirs"] = "merged", "", "", {}
	args = SimpleNamespace(stage="I1", variant=variant, seed=seed, out=str(out), epochs=0, minutes=600, cpu=False, no_amp=False, synthetic=0, max_anchors=0,
	no_lock=True, no_label_rows=False, base_v=bases.get("V0", ""), base_e=bases.get("E0", ""), base_u=bases.get("U0", ""), regimes="", pred_splits="",
	kin_source=DFLAGS["kin_source"], kin_version=DFLAGS["kin_version"], pack_version=DFLAGS["pack_version"], stats_mode="merged", run_suffix="", base_suffix_e="")
	# library monkey-patches (module-level, this process only): OOF base dirs, no PRED_DONE requirement, device exclusion, fold label rows only
	J.base_dirs = lambda cfg_, seed_: {b: base_root / f"{st_}_{va_}_s{seed_}{suffix}" for b, (st_, va_) in cfg_["base"].items()}
	J.base_info_numpy = _base_info_numpy_oof
	orig_load_raw = J.load_raw

	def load_raw_oof(a, c, split, subset=True, regimes=None):
	raw = orig_load_raw(a, c, split, subset, regimes)
	if split == "train" and subset:
	keep = np.flatnonzero(~np.isin(raw["dongle"], list(held)))
	raw = ds.subset(raw, keep)
	raw["regime_row"], raw["regime_id"] = raw["regime_row"], raw["regime_id"]
	return raw
	J.load_raw = load_raw_oof
	J.label_rows = lambda a, c: {r: np.asarray(v, np.int64) for r, v in fold_rows.items()}
	J.load_pred_raws = lambda a, c, lrows: _label_raws_only(a, c, lrows, orig_load_raw)
	info = J.base_info_numpy(cfg, seed)
	cfg["e_side"] = info["base_cfgs"]["E"].get("side", "vf256")
	sv = {b: c.get("stats_version") for b, c in info["base_cfgs"].items()}
	cfg["stats_version"], cfg["base_stats_versions"] = (int(min(v for v in sv.values() if v is not None)) if any(v is not None for v in sv.values()) else 1), sv
	(out / "config.json").write_text(json.dumps({"cfg": cfg, "args": vars(args), "oof": {"held_devices": sorted(held), "base_root": str(base_root), "suffix": suffix, "data_flags": DFLAGS}}, indent=1, default=str))
	wait_gates(log_out, f"joint {variant} s{seed}")
	st = {k: np.asarray(v, np.float32) for k, v in info["stats"].items()}
	st.update(J.aux_stats_from_files(cfg["regimes"]))
	np.savez(out / "stats.npz", **st)
	(out / "base_ckpts.json").write_text(json.dumps({"base_ckpts": info["base_ckpts"], "base_cfgs": info["base_cfgs"], "tok_source": cfg.get("tok_source")}, indent=1))
	bank = J.load_bank(args, cfg, info)
	if not (out / "TRAIN_DONE").exists():
	tr, va = J.load_raw(args, cfg, "train"), J.load_raw(args, cfg, "val")
	log(f"[oof] joint {variant} s{seed}: train rows {len(tr['t']):,} (fold devices held out)", log_out)
	code = J.run_chunk(args, cfg, tr, va, st, bank, info)
	assert code == 0, f"joint fold training failed ({code})"
	del tr, va
	raws = J.load_pred_raws(args, cfg, fold_rows)
	code = J.predict_many(args, cfg, raws, st, bank, info)
	assert code == 0
	(out / "PRED_DONE").write_text(json.dumps({"label_rows_only": True, "regimes": list(fold_rows)}))


	def _base_info_numpy_oof(cfg, seed):
	from models.cwm3 import merge_stats, base_cfg_summary
	names, cfgs, stats, dirs = {}, {}, [], {}
	for b, d in J.base_dirs(cfg, seed).items():
	assert (d / "BEST").exists() and (d / "TRAIN_DONE").exists(), f"OOF base for {b} missing: {d}"
	c = json.loads((d / "config.json").read_text())["cfg"]
	names[b] = f"{d.name}:{(d / 'BEST').read_text().strip()}"
	cfgs[b] = base_cfg_summary(c) # 2026-09-12: kclip / u_in_ln / p_head / select_LF / side / stats_version reach the joint's batch prep (as train_cwm3_joint.base_info_numpy)
	st = {k: v for k, v in np.load(d / "stats.npz").items()}
	if "stats_version" in st and "stats_version" not in cfgs[b]:
	cfgs[b]["stats_version"] = int(st["stats_version"])
	stats.append(st)
	dirs[b] = str(d)
	return {"base_ckpts": names, "base_cfgs": cfgs, "stats": merge_stats(stats), "stats_mode": "merged", "base_dirs": dirs}


	def _label_raws_only(args, cfg, lrows, orig_load_raw):
	raws = {}
	for regime, rows in lrows.items():
	if len(rows) and not (Path(args.out) / f"pred_{regime}_train_labelrows.npz").exists():
	full = orig_load_raw(args, cfg, "train", subset=False, regimes=[regime])
	sub = ds.subset(full, rows)
	sub["_label_rows"] = rows
	raws[(regime, "train")] = sub
	del full
	return raws


	# ----------------------------------------------------------------------------- merge
	def merge(seed, variant, K, out_root, target_run, log_out):
	for regime in ("manual", "engaged"):
	packs = [np.load(out_root / f"I1_{variant}_s{seed}_f{k}" / f"pred_{regime}_train_labelrows.npz", allow_pickle=True) for k in range(K)
	if (out_root / f"I1_{variant}_s{seed}_f{k}" / f"pred_{regime}_train_labelrows.npz").exists()]
	if len(packs) < K:
	log(f"[oof] merge {regime}: only {len(packs)}/{K} folds present — skipped", log_out)
	continue
	rows = np.concatenate([p["rows"] for p in packs])
	order = np.argsort(rows, kind="mergesort")
	fold = np.concatenate([np.full(len(p["rows"]), k, np.int8) for k, p in enumerate(packs)])[order]
	out = {}
	for key in packs[0].files:
	a0 = packs[0][key]
	if a0.ndim >= 1 and a0.shape[0] == len(packs[0]["rows"]) and key != "rows":
	out[key] = np.concatenate([p[key] for p in packs])[order]
	elif key == "rows":
	out[key] = rows[order]
	else:
	out[key] = a0
	out["fold"] = fold
	out["oof"] = np.array(json.dumps({"K": K, "folds": [str(out_root / f"I1_{variant}_s{seed}_f{k}") for k in range(K)]}))
	target = CWM3 / "results" / "runs" / target_run / f"pred_{regime}_train_labelrows_oof.npz"
	if target.exists():
	log(f"[oof] {target} exists — not overwritten", log_out)
	continue
	np.savez_compressed(target, **out)
	log(f"[oof] merged {regime}: {len(rows)} label rows from {K} folds -> {target}", log_out)


	def main():
	ap = argparse.ArgumentParser()
	ap.add_argument("--seed", type=int, default=0)
	ap.add_argument("--variant", default="all")
	ap.add_argument("--folds", type=int, default=5)
	ap.add_argument("--fold", type=int, default=-1, help="single fold (default: all)")
	ap.add_argument("--out-root", default=str(CWM3 / "results" / "runs" / "oof"))
	ap.add_argument("--no-gate", action="store_true")
	ap.add_argument("--bases", default=BASES_DEFAULT, help="stage:variant triple of the base branches (default = the original V0/can4s, U0/kin_can_noK, E0/pooled)")
	ap.add_argument("--kin-source", default="cwm2", choices=["cwm2", "causal_routes", "oldin_ctgt"], help="data layer kin source for the fold bases AND the joint (default cwm2 = original)")
	ap.add_argument("--kin-version", default="v1", choices=["v1", "v2"])
	ap.add_argument("--pack-version", default="v1", choices=["v1", "v2"])
	ap.add_argument("--target-run", default="", help="run dir under results/runs that receives pred_{regime}_train_labelrows_oof.npz (default I1_<variant>_s<seed>)")
	ap.add_argument("--gate-each", action="store_true", help="wait for the launch gates before EVERY training (default: once per fold, original behaviour)")
	ap.add_argument("--gate-all-trainers", action="store_true", help="count train_kp / pretrain processes in the trainer gate too (project rule: < 4 trainers alive)")
	ap.add_argument("--max-others", type=int, default=1, help="trainer-count gate: other trainers alive <= this (default 1 = original)")
	ap.add_argument("--min-avail-gb", type=int, default=26)
	ap.add_argument("--min-vram-gb", type=int, default=12)
	a = ap.parse_args()
	global BASES
	BASES = [tuple(x.split(":")) for x in a.bases.split(",")]
	assert [b[0] for b in BASES] == ["V0", "U0", "E0"], a.bases
	DFLAGS.update(pack_version=a.pack_version, kin_source=a.kin_source, kin_version=a.kin_version)
	ds.set_pack_version(a.pack_version)
	ds.set_kin_source(a.kin_source, a.kin_version)
	gate_kw = dict(min_avail_gb=a.min_avail_gb, min_vram_gb=a.min_vram_gb, max_others=a.max_others, pattern=("train_cwm3\|train_kp\|pretrain" if a.gate_all_trainers else "train_cwm3"))
	GATE["each"], GATE["kw"] = bool(a.gate_each and not a.no_gate), gate_kw
	t0 = time.time()
	out_root = Path(a.out_root)
	out_root.mkdir(parents=True, exist_ok=True)
	log_out = str(out_root)
	tab = label_table()
	folds = device_folds(tab, a.folds)
	(out_root / f"folds_s{a.seed}.json").write_text(json.dumps({"K": a.folds, "device_fold": folds, "n_label_rows": int(len(tab))}, indent=1))
	log(f"[oof] {len(folds)} labelled devices -> {a.folds} folds; {len(tab)} label rows; bases {BASES}; data flags {DFLAGS}; target run {a.target_run or f'I1_{a.variant}_s{a.seed}'}", log_out)
	for k in range(a.folds):
	if a.fold >= 0 and k != a.fold:
	continue
	held = {d for d, f in folds.items() if f == k}
	fold_rows = {r: np.sort(tab[(tab.regime == r) & tab.dongle.isin(held)].row.unique()).astype(np.int64) for r in ("manual", "engaged")}
	while not a.no_gate:
	ok, stat = gates_ok(**gate_kw)
	if ok:
	break
	time.sleep(60)
	log(f"[oof] === fold {k}: {len(held)} devices, label rows manual {len(fold_rows['manual'])} engaged {len(fold_rows['engaged'])} ({time.time() - t0:.0f}s)", log_out)
	for stage, variant in BASES:
	train_base_oof(stage, variant, a.seed, held, out_root / f"{stage}_{variant}_s{a.seed}_f{k}", log_out)
	train_joint_oof(a.seed, a.variant, held, fold_rows, out_root, f"_f{k}", out_root / f"I1_{a.variant}_s{a.seed}_f{k}", log_out, bases=dict(BASES))
	if a.fold < 0:
	merge(a.seed, a.variant, a.folds, out_root, a.target_run or f"I1_{a.variant}_s{a.seed}", log_out)
	log(f"[oof] DONE in {(time.time() - t0) / 60:.1f} min", log_out)


	if __name__ == "__main__":
	main()