Add VX-One source code

908fd1a verified 5 days ago

22.3 kB

"""Loop de entrenamiento VX-One (spec F2 §6.3-§6.5, D11).
Hiperparámetros congelados:
- AdamW β=(0.9,0.999), eps=1e-8; weight decay 1e-2 SOLO en matrices de peso
  Linear 2D (excluye LN, biases, embeddings, corn_b y s_q).
- Schedule: warmup lineal 5% + coseno hasta 0; clip de norma 1.0.
- Early stopping: patience 4 sobre M_val = media de métricas de cabezas
  activas (AUROC noul, macro-F1 choice sin abstención, 1−MAE(ŝ,y)/3 score,
  macro-F1 ddos en subconjunto).
- AMP: bf16 en L4, fp16 en T4, off en CPU.
- Al terminar: se restaura el mejor epoch, se calibra cada cabeza en val
  (§5, NLL) y se escribe el artefacto self-contained (§8.2).
"""

from __future__ import annotations

import copy
import math
import os
import time
from pathlib import Path

import numpy as np
import pyarrow.parquet as pq
import torch
from sklearn.metrics import f1_score, roc_auc_score

from .calibration import (fit_choice_temperature, fit_corn_temperature,
                          fit_noul_temperature)
from .data import RunData
from .engine import SOURCE_KEY, TAU_ABS, save_artifact
from .losses import (KendallMultiTask, choice_loss, corn_expected, corn_loss,
                     corn_pmf, noul_loss)
from .network import VXOneNet

# Tabla congelada §6.3: {escala: {fuente: (batch_size, epochs)}}; lr por escala.
TRAIN_TABLE: dict[str, dict[str, tuple[int, int]]] = {
    "nano":  {"unsw": (256, 40),  "cic": (1024, 25), "ember": (2048, 12)},
    "small": {"unsw": (512, 30),  "cic": (2048, 20), "ember": (4096, 10)},
    "base":  {"unsw": (1024, 25), "cic": (4096, 15), "ember": (8192, 8)},
}
LR_TABLE = {"nano": 3e-3, "small": 1e-3, "base": 3e-4}
WARMUP_FRAC = 0.05
CLIP_NORM = 1.0
PATIENCE = 4
WEIGHT_DECAY = 1e-2
SEED = 42
EVAL_BS = 8192


def pick_amp(device: str, amp: str = "auto") -> str:
    """'off' | 'bf16' | 'fp16' — auto: bf16 si CC≥8 (L4), fp16 si T4, off CPU."""
    if amp != "auto":
        return amp
    if device.startswith("cuda"):
        cc = torch.cuda.get_device_capability(device)
        return "bf16" if cc[0] >= 8 else "fp16"
    return "off"


def _param_groups(net: VXOneNet) -> list[dict]:
    """wd solo en matrices de peso Linear 2D (D11): LN/biases/embeddings fuera
    (corn_b y s_q son 1D/0D y quedan excluidos por la regla de dims)."""
    decay = [p for n, p in net.named_parameters() if p.dim() == 2 and not n.startswith("emb.")]
    no_decay = [p for n, p in net.named_parameters() if not (p.dim() == 2 and not n.startswith("emb."))]
    return [{"params": decay, "weight_decay": WEIGHT_DECAY},
            {"params": no_decay, "weight_decay": 0.0}]


def _make_optimizer(net: VXOneNet, lr: float,
                    total_steps: int) -> tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]:
    opt = torch.optim.AdamW(_param_groups(net), lr=lr, betas=(0.9, 0.999), eps=1e-8)
    warmup = max(1, int(WARMUP_FRAC * total_steps))

    def lr_lambda(step: int) -> float:
        if step < warmup:
            return step / warmup
        p = (step - warmup) / max(1, total_steps - warmup)
        return 0.5 * (1.0 + math.cos(math.pi * min(1.0, p)))

    return opt, torch.optim.lr_scheduler.LambdaLR(opt, lr_lambda)


class _BatchLabels:
    """Vista ligera de labels de un batch (slices o shards)."""

    def __init__(self, y_fam: torch.Tensor, y_bin: torch.Tensor,
                 y_sev: torch.Tensor | None = None,
                 ddos_mask: torch.Tensor | None = None,
                 y_ddos: torch.Tensor | None = None) -> None:
        self.y_fam, self.y_bin, self.y_sev = y_fam, y_bin, y_sev
        self.ddos_mask, self.y_ddos = ddos_mask, y_ddos


def _batch_losses(net: VXOneNet, run: RunData, u, txt, lab: _BatchLabels):
    """Loss por cabeza + máscaras de actividad (§3.6). Los tensores ω viven
    en CPU (run es serializable); se mueven al device del batch aquí."""
    out = net(u, txt)
    dev = u.device
    fm = lab.y_fam >= 0  # sentinela -1 (A5-EMBER): fuera del CE de choice
    losses = {
        "noul": noul_loss(out["noul"], lab.y_bin),
        "choice": choice_loss(out["choice"][fm], lab.y_fam[fm], run.omega.to(dev))
        if bool(fm.any()) else torch.zeros((), device=dev),
    }
    active = {"noul": True, "choice": bool(fm.any())}
    if net.has_score:
        losses["score"] = corn_loss(out["corn_logits"], lab.y_sev)
        active["score"] = True
    if net.has_ddos:
        m = lab.ddos_mask
        if bool(m.any()):
            losses["ddos"] = choice_loss(out["ddos"][m], lab.y_ddos[m],
                                         run.omega_ddos.to(dev))
            active["ddos"] = True
        else:
            active["ddos"] = False
    head = run.meta.get("a2_head")
    if head:  # A2: una sola cabeza activa (loss + M_val)
        losses = {head: losses[head]}
        active = {q: q == head for q in active}
    return losses, active


@torch.no_grad()
def _forward_numpy(net: VXOneNet, u: torch.Tensor, txt, lab, device: str) -> dict[str, np.ndarray]:
    """Forward en batches → arrays numpy para métricas/calibración (UNSW/CIC)."""
    outs: dict[str, list[np.ndarray]] = {}
    for i in range(0, u.shape[0], EVAL_BS):
        ub = u[i:i + EVAL_BS].to(device)
        tb = {k: v[i:i + EVAL_BS].to(device) for k, v in (txt or {}).items()}
        o = net(ub, tb if tb else None)
        for k, v in o.items():
            outs.setdefault(k, []).append(v.float().cpu().numpy())
    flat = {k: np.concatenate(v) for k, v in outs.items()}
    res: dict[str, np.ndarray] = {
        "y_bin": lab.y_bin.numpy(),
        "y_fam": lab.y_fam.numpy(),
        "noul_logits": flat["noul"],
        "choice_logits": flat["choice"],
    }
    if "corn_logits" in flat:
        res["corn_logits"] = flat["corn_logits"]
        res["y_sev"] = lab.y_sev.numpy()
    if "ddos" in flat:
        m = lab.ddos_mask.numpy()
        res.update(ddos_logits=flat["ddos"][m], y_ddos=lab.y_ddos.numpy()[m])
    return res


@torch.no_grad()
def _eval_ember(net: VXOneNet, run: RunData, batcher, subset: str, device: str) -> dict[str, np.ndarray]:
    """Forward en streaming (val/test EMBER) → arrays numpy."""
    acc: dict[str, list[np.ndarray]] = {}
    y_bin_l, y_fam_l = [], []
    for u_v, y_fam_v, y_bin_v in batcher.eval_batches(subset, EVAL_BS):
        o = net(u_v.to(device), None)
        for k, v in o.items():
            acc.setdefault(k, []).append(v.float().cpu().numpy())
        y_bin_l.append(y_bin_v.numpy())
        y_fam_l.append(y_fam_v.numpy())
    flat = {k: np.concatenate(v) for k, v in acc.items()}
    return {"noul_logits": flat["noul"], "choice_logits": flat["choice"],
            "y_bin": np.concatenate(y_bin_l), "y_fam": np.concatenate(y_fam_l)}


def _m_val(o: dict[str, np.ndarray], names: tuple[str, ...]) -> tuple[dict[str, float], float]:
    """M_val = media de métricas de cabezas activas (§6.4). Las filas con
    y_fam = -1 (fuera del candidate set, D12) no participan en Choice."""
    p_mal = 1.0 / (1.0 + np.exp(-o["noul_logits"]))  # AUROC invariante a T
    fam_mask = o["y_fam"] >= 0
    fam_pred = np.argmax(o["choice_logits"], axis=1)
    m: dict[str, float] = {
        "noul": (roc_auc_score(o["y_bin"], p_mal)
                 if len(np.unique(o["y_bin"])) > 1 else 0.5),
        "choice": f1_score(o["y_fam"][fam_mask], fam_pred[fam_mask],
                           average="macro", zero_division=0),
    }
    if "corn_logits" in o:
        pi = corn_pmf(torch.tensor(o["corn_logits"]))[1].numpy()
        sev_hat = (pi * np.arange(4)).sum(axis=1)
        m["score"] = 1.0 - float(np.abs(sev_hat - o["y_sev"]).mean()) / 3.0
    if "y_ddos" in o:
        m["ddos"] = f1_score(o["y_ddos"], np.argmax(o["ddos_logits"], axis=1),
                             average="macro", zero_division=0)
    active = {k: m[k] for k in names if k in m}
    return active, float(np.mean(list(active.values())))


def train_run(run: RunData, out_path: str | Path, device: str = "cpu",
              amp: str = "auto", seed: int = SEED, epochs: int | None = None,
              batch_size: int | None = None, log=print,
              ckpt_path: str | Path | None = None,
              ckpt_every_steps: int = 2000, resume: bool = False,
              ablation: str | None = None) -> dict:
    """Entrena una corrida y escribe el artefacto §8.2. Devuelve el resumen.
    Ablations E6 (§7): `ablation` en {a1, a2_*, a3, a4, a5, a6} — un solo
    cambio por corrida sobre la referencia Small: a1 sin rama textual (vocab a
    None), a2 sin multi-task (a2_noul / a2_choice / a2_score: un modelo por
    cabeza, el merge de los tres dumps es la corrida a2), a3 sin calibración
    (T forzada a 1), a4 ω≡1 (posterior con el
    tilt del train), a5 unknown off (τ_abs→1 en abstención; EMBER: el
    candidate set ya se pasa sin unknown), a6 dropout 0.
    Checkpointing de recuperación (OOM/kill): cada `ckpt_every_steps` pasos y
    al cierre de cada epoch se escribe ATÓMICAMENTE `ckpt_path` con
    net+opt+sched+scaler+contadores de early-stop+rng. `resume=True` retoma
    desde ese checkpoint REINICIANDO el epoch en curso (la permutación es
    determinista seed+epoch) — pérdida máxima <1 epoch; las corridas que
    terminan bien borran el checkpoint. El determinismo exacto (rng fresco,
    seed fijo) se aplica a los reportes single-shot del paper; el resume es
    un mecanismo de crash-recovery."""
    torch.manual_seed(seed)
    src = {"unsw": "unsw", "cic": "cic", "ember2024_v1": "ember"}[run.source]
    bs_def, ep_def = TRAIN_TABLE[run.scale][src]
    bs = batch_size or bs_def
    n_epochs = epochs or ep_def
    lr = LR_TABLE[run.scale]
    amp_mode = pick_amp(device, amp)

    # ── ablations E6 ──────────────────────────────────────────────────
    if ablation == "a1":   # sin rama textual: el net no ve embeddings
        run.vocab = None
        for part in ("train_txt", "val_txt", "test_txt"):
            setattr(run, part, None)
    elif ablation == "a4":  # ω≡1: CE sin tilt, de-tilt identidad
        run.omega = torch.ones_like(run.omega)
        if run.omega_ddos is not None:
            run.omega_ddos = torch.ones_like(run.omega_ddos)
    elif ablation == "a5" and run.source == "ember2024_v1":
        # 10-way sin unknown: la clase sale del candidate set; sus filas
        # (benignos + sin familia + cola) quedan fuera del CE de choice vía
        # el sentinela -1 del imap y la máscara de _batch_losses
        from .losses import class_weights

        unk = run.candidates.index("unknown")
        counts = run.meta.get("train_counts", {})
        run.candidates = [c for c in run.candidates if c != "unknown"]
        run.omega = torch.tensor(class_weights([counts[c] for c in run.candidates]))
        # la red se construye desde run.heads (dataclass frozen): reemplazarlo
        import dataclasses
        run.heads = dataclasses.replace(run.heads,
                                        n_candidates=len(run.candidates))
        for part in ("train_labels", "val_labels", "test_labels"):
            lab = getattr(run, part, None)
            if lab is None or lab.y_fam is None:
                continue
            yf = lab.y_fam
            lab.y_fam = torch.where(yf == unk, torch.full_like(yf, -1),
                                    torch.where(yf > unk, yf - 1, yf))
    elif ablation == "a6":  # dropout 0
        import src.models.network as _net_mod
        _net_mod.DROPOUT = 0.0
    elif ablation in ("a2_noul", "a2_choice", "a2_score"):
        # A2 sin multi-task (§7): UN modelo por cabeza con el mismo
        # presupuesto total — las cabezas aportan <1% de params, la anchura
        # no requiere ajuste. Solo la cabeza elegida recibe pérdida y M_val;
        # las demás quedan sin gradiente y el dump combinado las excluye.
        run.meta["a2_head"] = ablation.removeprefix("a2_")

    log(f"[train] {run.source}/{run.scale}/{run.split_scheme} bs={bs} epochs={n_epochs} "
        f"lr={lr} amp={amp_mode} K={run.heads.n_candidates} device={device}"
        + (f" ablation={ablation}" if ablation else ""))

    net = VXOneNet(src, run.scale, run.heads,
                   d_num=len(run.feature_stage.impute_median),
                   vocab_sizes={k: v.size for k, v in (run.vocab or {}).items()})
    net.to(device)
    kendall = KendallMultiTask(run.heads.names).to(device)

    ember = run.source == "ember2024_v1"
    if ember:
        from .data import EmberBatcher
        batcher = EmberBatcher(run)
        rows = [pq.ParquetFile(str(p)).metadata.num_rows for p in run.ember_train_shards]
        n_train = sum(rows)
        steps_per_epoch = sum(r // bs for r in rows)
    else:
        n_train = run.train.shape[0]
        steps_per_epoch = math.ceil(n_train / bs)

    total_steps = n_epochs * steps_per_epoch
    opt, sched = _make_optimizer(net, lr, total_steps)
    scaler = torch.amp.GradScaler(device, enabled=amp_mode == "fp16")
    amp_dtype = {"bf16": torch.bfloat16, "fp16": torch.float16}.get(amp_mode)
    params = [p for g in opt.param_groups for p in g["params"]]

    # ── checkpointing de recuperación (atomic tmp→replace) ───────────
    def _save_ckpt(epoch_done: int, best_m: float, best_epoch: int,
                   bad: int, best_state) -> None:
        """epoch_done = nº de epochs COMPLETOS (resume re-playea el actual)."""
        if not ckpt_path:
            return
        tmp = f"{ckpt_path}.tmp"
        Path(tmp).parent.mkdir(parents=True, exist_ok=True)
        torch.save({
            "net": {k: v.cpu() for k, v in net.state_dict().items()},
            "opt": opt.state_dict(), "sched": sched.state_dict(),
            "scaler": scaler.state_dict() if amp_mode == "fp16" else None,
            "amp": amp_mode, "epoch_done": epoch_done, "step": step,
            "best_m": best_m, "best_epoch": best_epoch, "bad": bad,
            "best_state": ({k: v.cpu() for k, v in best_state.items()}
                           if best_state else None),
            "torch_rng": torch.get_rng_state(),
        }, tmp)
        os.replace(tmp, ckpt_path)

    start_epoch, step = 0, 0
    best_m, best_epoch, bad, best_state = -1.0, -1, 0, None
    if resume and ckpt_path and Path(ckpt_path).exists():
        ck = torch.load(ckpt_path, map_location="cpu", weights_only=False)
        net.load_state_dict(ck["net"])
        opt.load_state_dict(ck["opt"])
        for st in opt.state.values():  # AdamW: estados a device de los params
            for k, v in st.items():
                if torch.is_tensor(v):
                    st[k] = v.to(device)
        # sched fresco: el epoch en curso se REPLAYEA completo (permutación
        # determinista seed+epoch), así last_epoch vuelve a alinear con
        # e*steps_per_epoch y el coseno termina en total_steps como single-shot.
        _, sched = _make_optimizer(net, lr, total_steps)
        if amp_mode == "fp16" and ck.get("scaler"):
            scaler.load_state_dict(ck["scaler"])
        if ck.get("torch_rng") is not None:
            torch.set_rng_state(ck["torch_rng"])
        best_m, best_epoch, bad = ck["best_m"], ck["best_epoch"], ck["bad"]
        best_state = ck.get("best_state")
        start_epoch = ck["epoch_done"]
        step = start_epoch * steps_per_epoch
        log(f"[train] RESUME desde {ckpt_path}: epochs completados "
            f"{start_epoch}/{n_epochs} M_val={best_m:.4f} (best ep {best_epoch})")

    def run_epoch(epoch: int) -> None:
        nonlocal step
        net.train()
        if ember:
            batches = ((u.to(device), None,
                        _BatchLabels(yf.to(device), yb.to(device)))
                       for u, yf, yb in batcher.train_batches(epoch, bs))
        else:
            order = np.random.default_rng(seed + epoch).permutation(n_train)
            txt_tr = run.train_txt or {}
            label_tr = run.train_labels

            def _gen():
                for s in range(0, n_train, bs):
                    idx = torch.tensor(order[s:s + bs])
                    u_b = run.train[idx].to(device)
                    txt_b = {k: v[idx].to(device) for k, v in txt_tr.items()} or None
                    yield u_b, txt_b, _BatchLabels(
                        label_tr.y_fam[idx].to(device),
                        label_tr.y_bin[idx].to(device),
                        label_tr.y_sev[idx].to(device)
                        if label_tr.y_sev is not None else None,
                        label_tr.ddos_mask[idx].to(device)
                        if label_tr.ddos_mask is not None else None,
                        label_tr.y_ddos[idx].to(device)
                        if label_tr.y_ddos is not None else None)
            batches = _gen()
        for u_b, txt_b, lab_b in batches:
            with torch.autocast(device_type=device.split(":")[0], dtype=amp_dtype,
                                enabled=amp_dtype is not None and device.startswith("cuda")):
                losses, active = _batch_losses(net, run, u_b, txt_b, lab_b)
                total, _lambdas = kendall(losses, active)
            opt.zero_grad(set_to_none=True)
            if amp_mode == "fp16":
                scaler.scale(total).backward()
                scaler.unscale_(opt)
                torch.nn.utils.clip_grad_norm_(params, CLIP_NORM)
                scaler.step(opt)
                scaler.update()
            else:
                total.backward()
                torch.nn.utils.clip_grad_norm_(params, CLIP_NORM)
                opt.step()
            sched.step()
            step += 1
            # checkpoint por paso (cada N): el epoch en curso se replayea al
            # retomar, así que epoch_done = epoch (NO epoch+1)
            if ckpt_path and ckpt_every_steps and step % ckpt_every_steps == 0:
                _save_ckpt(epoch, best_m, best_epoch, bad, best_state)

    def val_outputs() -> dict[str, np.ndarray]:
        net.eval()
        if ember:
            return _eval_ember(net, run, batcher, "val", device)
        txt_va = {k: v.to(device) for k, v in (run.val_txt or {}).items()}
        return _forward_numpy(net, run.val.to(device), txt_va or None,
                              run.val_labels, device)

    t0 = time.time()
    for epoch in range(start_epoch, n_epochs):
        run_epoch(epoch)
        names = run.heads.names
        if run.meta.get("a2_head"):  # A2: early-stop solo sobre su cabeza
            names = (run.meta["a2_head"],)
        active_m, mv = _m_val(val_outputs(), names)
        marker = ""
        if mv > best_m:
            best_m, best_epoch, bad = mv, epoch, 0
            best_state = copy.deepcopy(net.state_dict())
            marker = " ←"
        else:
            bad += 1
        log(f"[train] epoch {epoch + 1}/{n_epochs} M_val={mv:.4f} "
            + " ".join(f"{k}={v:.3f}" for k, v in active_m.items())
            + f" bad={bad}{marker}")
        if ckpt_path:  # fin de epoch: epoch_done = epoch+1 (barato, seguro)
            _save_ckpt(epoch + 1, best_m, best_epoch, bad, best_state)
        if bad >= PATIENCE:
            log(f"[train] early stop (patience {PATIENCE}) en epoch {epoch + 1}")
            break

    # ── restaurar mejor epoch + calibración en val (§5) ──────────────
    net.load_state_dict(best_state)
    o = val_outputs()

    if ablation == "a3":  # sin calibración: T=1 en todas las cabezas
        temperatures = {"noul": 1.0}
        if "corn_logits" in o:
            temperatures["score"] = 1.0
        if "y_ddos" in o:
            temperatures["ddos"] = 1.0
        temperatures["choice"] = 1.0
    else:
        temperatures = {"noul": fit_noul_temperature(
            torch.tensor(o["noul_logits"]), torch.tensor(o["y_bin"]))}
        if "corn_logits" in o:
            temperatures["score"] = fit_corn_temperature(
                torch.tensor(o["corn_logits"]), torch.tensor(o["y_sev"]))
        if "y_ddos" in o:
            temperatures["ddos"] = fit_choice_temperature(
                torch.tensor(o["ddos_logits"]), torch.tensor(o["y_ddos"]), run.omega_ddos)
        # Choice: solo filas dentro del candidate set (y_fam ≥ 0, D12)
        fam_mask = o["y_fam"] >= 0
        temperatures["choice"] = fit_choice_temperature(
            torch.tensor(o["choice_logits"][fam_mask]),
            torch.tensor(o["y_fam"][fam_mask]), run.omega)

    run_id = f"{src}_{run.scale}_{run.split_scheme}"
    if ablation:
        run_id += f"_abl_{ablation}"
    # A5 (unknown off): τ_abs→1 ⇒ nadie abtiene (UNSW/CIC) o modo abstain
    # puro tras quitar la clase (EMBER); el candidate set sin unknown ya lo
    # dejó train_run en run.candidates.
    tau = 1.0 if ablation == "a5" else TAU_ABS
    save_artifact(out_path, net, run_meta={
        "source_key": SOURCE_KEY[src],
        "feature_stage": run.feature_stage,
        "vocab": run.vocab,
        "candidates": run.candidates,
        "omega": run.omega,
        "omega_ddos": run.omega_ddos,
        "ddos_candidates": run.ddos_candidates,
        "mode": "abstain" if ablation == "a5" else ("trained_unknown" if ember else "abstain"),
        "tau_abs": tau,
        "config": {"scale": run.scale, "split_scheme": run.split_scheme,
                   "batch_size": bs, "lr": lr, "amp": amp_mode, "seed": seed},
        "manifest": {
            "run_id": run_id, "created_at": time.strftime("%Y-%m-%dT%H:%M:%S"),
            "seed": seed, "best_epoch": best_epoch, "m_val": best_m,
            "n_train": n_train, "candidates": run.candidates,
            "train_counts": run.meta.get("train_counts", {}),
            "exclusion": run.meta.get("exclusion", {}),
            "temperatures": temperatures,
            "wall_time_s": round(time.time() - t0, 1),
        },
    }, temperatures=temperatures, d_num=len(run.feature_stage.impute_median))
    log(f"[train] artefacto → {out_path} (best_epoch={best_epoch} M_val={best_m:.4f})")
    if ckpt_path:  # corrida completa: el ckpt de recuperación ya no hace falta
        Path(ckpt_path).unlink(missing_ok=True)
    return {"run_id": run_id, "best_epoch": best_epoch, "m_val": best_m,
            "temperatures": temperatures, "wall_time_s": time.time() - t0}

	"""Loop de entrenamiento VX-One (spec F2 §6.3-§6.5, D11).

	Hiperparámetros congelados:
	- AdamW β=(0.9,0.999), eps=1e-8; weight decay 1e-2 SOLO en matrices de peso
	Linear 2D (excluye LN, biases, embeddings, corn_b y s_q).
	- Schedule: warmup lineal 5% + coseno hasta 0; clip de norma 1.0.
	- Early stopping: patience 4 sobre M_val = media de métricas de cabezas
	activas (AUROC noul, macro-F1 choice sin abstención, 1−MAE(ŝ,y)/3 score,
	macro-F1 ddos en subconjunto).
	- AMP: bf16 en L4, fp16 en T4, off en CPU.
	- Al terminar: se restaura el mejor epoch, se calibra cada cabeza en val
	(§5, NLL) y se escribe el artefacto self-contained (§8.2).
	"""

	from __future__ import annotations

	import copy
	import math
	import os
	import time
	from pathlib import Path

	import numpy as np
	import pyarrow.parquet as pq
	import torch
	from sklearn.metrics import f1_score, roc_auc_score

	from .calibration import (fit_choice_temperature, fit_corn_temperature,
	fit_noul_temperature)
	from .data import RunData
	from .engine import SOURCE_KEY, TAU_ABS, save_artifact
	from .losses import (KendallMultiTask, choice_loss, corn_expected, corn_loss,
	corn_pmf, noul_loss)
	from .network import VXOneNet

	# Tabla congelada §6.3: {escala: {fuente: (batch_size, epochs)}}; lr por escala.
	TRAIN_TABLE: dict[str, dict[str, tuple[int, int]]] = {
	"nano": {"unsw": (256, 40), "cic": (1024, 25), "ember": (2048, 12)},
	"small": {"unsw": (512, 30), "cic": (2048, 20), "ember": (4096, 10)},
	"base": {"unsw": (1024, 25), "cic": (4096, 15), "ember": (8192, 8)},
	}
	LR_TABLE = {"nano": 3e-3, "small": 1e-3, "base": 3e-4}
	WARMUP_FRAC = 0.05
	CLIP_NORM = 1.0
	PATIENCE = 4
	WEIGHT_DECAY = 1e-2
	SEED = 42
	EVAL_BS = 8192


	def pick_amp(device: str, amp: str = "auto") -> str:
	"""'off' \| 'bf16' \| 'fp16' — auto: bf16 si CC≥8 (L4), fp16 si T4, off CPU."""
	if amp != "auto":
	return amp
	if device.startswith("cuda"):
	cc = torch.cuda.get_device_capability(device)
	return "bf16" if cc[0] >= 8 else "fp16"
	return "off"


	def _param_groups(net: VXOneNet) -> list[dict]:
	"""wd solo en matrices de peso Linear 2D (D11): LN/biases/embeddings fuera
	(corn_b y s_q son 1D/0D y quedan excluidos por la regla de dims)."""
	decay = [p for n, p in net.named_parameters() if p.dim() == 2 and not n.startswith("emb.")]
	no_decay = [p for n, p in net.named_parameters() if not (p.dim() == 2 and not n.startswith("emb."))]
	return [{"params": decay, "weight_decay": WEIGHT_DECAY},
	{"params": no_decay, "weight_decay": 0.0}]


	def _make_optimizer(net: VXOneNet, lr: float,
	total_steps: int) -> tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]:
	opt = torch.optim.AdamW(_param_groups(net), lr=lr, betas=(0.9, 0.999), eps=1e-8)
	warmup = max(1, int(WARMUP_FRAC * total_steps))

	def lr_lambda(step: int) -> float:
	if step < warmup:
	return step / warmup
	p = (step - warmup) / max(1, total_steps - warmup)
	return 0.5 * (1.0 + math.cos(math.pi * min(1.0, p)))

	return opt, torch.optim.lr_scheduler.LambdaLR(opt, lr_lambda)


	class _BatchLabels:
	"""Vista ligera de labels de un batch (slices o shards)."""

	def __init__(self, y_fam: torch.Tensor, y_bin: torch.Tensor,
	y_sev: torch.Tensor \| None = None,
	ddos_mask: torch.Tensor \| None = None,
	y_ddos: torch.Tensor \| None = None) -> None:
	self.y_fam, self.y_bin, self.y_sev = y_fam, y_bin, y_sev
	self.ddos_mask, self.y_ddos = ddos_mask, y_ddos


	def _batch_losses(net: VXOneNet, run: RunData, u, txt, lab: _BatchLabels):
	"""Loss por cabeza + máscaras de actividad (§3.6). Los tensores ω viven
	en CPU (run es serializable); se mueven al device del batch aquí."""
	out = net(u, txt)
	dev = u.device
	fm = lab.y_fam >= 0 # sentinela -1 (A5-EMBER): fuera del CE de choice
	losses = {
	"noul": noul_loss(out["noul"], lab.y_bin),
	"choice": choice_loss(out["choice"][fm], lab.y_fam[fm], run.omega.to(dev))
	if bool(fm.any()) else torch.zeros((), device=dev),
	}
	active = {"noul": True, "choice": bool(fm.any())}
	if net.has_score:
	losses["score"] = corn_loss(out["corn_logits"], lab.y_sev)
	active["score"] = True
	if net.has_ddos:
	m = lab.ddos_mask
	if bool(m.any()):
	losses["ddos"] = choice_loss(out["ddos"][m], lab.y_ddos[m],
	run.omega_ddos.to(dev))
	active["ddos"] = True
	else:
	active["ddos"] = False
	head = run.meta.get("a2_head")
	if head: # A2: una sola cabeza activa (loss + M_val)
	losses = {head: losses[head]}
	active = {q: q == head for q in active}
	return losses, active


	@torch.no_grad()
	def _forward_numpy(net: VXOneNet, u: torch.Tensor, txt, lab, device: str) -> dict[str, np.ndarray]:
	"""Forward en batches → arrays numpy para métricas/calibración (UNSW/CIC)."""
	outs: dict[str, list[np.ndarray]] = {}
	for i in range(0, u.shape[0], EVAL_BS):
	ub = u[i:i + EVAL_BS].to(device)
	tb = {k: v[i:i + EVAL_BS].to(device) for k, v in (txt or {}).items()}
	o = net(ub, tb if tb else None)
	for k, v in o.items():
	outs.setdefault(k, []).append(v.float().cpu().numpy())
	flat = {k: np.concatenate(v) for k, v in outs.items()}
	res: dict[str, np.ndarray] = {
	"y_bin": lab.y_bin.numpy(),
	"y_fam": lab.y_fam.numpy(),
	"noul_logits": flat["noul"],
	"choice_logits": flat["choice"],
	}
	if "corn_logits" in flat:
	res["corn_logits"] = flat["corn_logits"]
	res["y_sev"] = lab.y_sev.numpy()
	if "ddos" in flat:
	m = lab.ddos_mask.numpy()
	res.update(ddos_logits=flat["ddos"][m], y_ddos=lab.y_ddos.numpy()[m])
	return res


	@torch.no_grad()
	def _eval_ember(net: VXOneNet, run: RunData, batcher, subset: str, device: str) -> dict[str, np.ndarray]:
	"""Forward en streaming (val/test EMBER) → arrays numpy."""
	acc: dict[str, list[np.ndarray]] = {}
	y_bin_l, y_fam_l = [], []
	for u_v, y_fam_v, y_bin_v in batcher.eval_batches(subset, EVAL_BS):
	o = net(u_v.to(device), None)
	for k, v in o.items():
	acc.setdefault(k, []).append(v.float().cpu().numpy())
	y_bin_l.append(y_bin_v.numpy())
	y_fam_l.append(y_fam_v.numpy())
	flat = {k: np.concatenate(v) for k, v in acc.items()}
	return {"noul_logits": flat["noul"], "choice_logits": flat["choice"],
	"y_bin": np.concatenate(y_bin_l), "y_fam": np.concatenate(y_fam_l)}


	def _m_val(o: dict[str, np.ndarray], names: tuple[str, ...]) -> tuple[dict[str, float], float]:
	"""M_val = media de métricas de cabezas activas (§6.4). Las filas con
	y_fam = -1 (fuera del candidate set, D12) no participan en Choice."""
	p_mal = 1.0 / (1.0 + np.exp(-o["noul_logits"])) # AUROC invariante a T
	fam_mask = o["y_fam"] >= 0
	fam_pred = np.argmax(o["choice_logits"], axis=1)
	m: dict[str, float] = {
	"noul": (roc_auc_score(o["y_bin"], p_mal)
	if len(np.unique(o["y_bin"])) > 1 else 0.5),
	"choice": f1_score(o["y_fam"][fam_mask], fam_pred[fam_mask],
	average="macro", zero_division=0),
	}
	if "corn_logits" in o:
	pi = corn_pmf(torch.tensor(o["corn_logits"]))[1].numpy()
	sev_hat = (pi * np.arange(4)).sum(axis=1)
	m["score"] = 1.0 - float(np.abs(sev_hat - o["y_sev"]).mean()) / 3.0
	if "y_ddos" in o:
	m["ddos"] = f1_score(o["y_ddos"], np.argmax(o["ddos_logits"], axis=1),
	average="macro", zero_division=0)
	active = {k: m[k] for k in names if k in m}
	return active, float(np.mean(list(active.values())))


	def train_run(run: RunData, out_path: str \| Path, device: str = "cpu",
	amp: str = "auto", seed: int = SEED, epochs: int \| None = None,
	batch_size: int \| None = None, log=print,
	ckpt_path: str \| Path \| None = None,
	ckpt_every_steps: int = 2000, resume: bool = False,
	ablation: str \| None = None) -> dict:
	"""Entrena una corrida y escribe el artefacto §8.2. Devuelve el resumen.

	Ablations E6 (§7): `ablation` en {a1, a2_*, a3, a4, a5, a6} — un solo
	cambio por corrida sobre la referencia Small: a1 sin rama textual (vocab a
	None), a2 sin multi-task (a2_noul / a2_choice / a2_score: un modelo por
	cabeza, el merge de los tres dumps es la corrida a2), a3 sin calibración
	(T forzada a 1), a4 ω≡1 (posterior con el
	tilt del train), a5 unknown off (τ_abs→1 en abstención; EMBER: el
	candidate set ya se pasa sin unknown), a6 dropout 0.

	Checkpointing de recuperación (OOM/kill): cada `ckpt_every_steps` pasos y
	al cierre de cada epoch se escribe ATÓMICAMENTE `ckpt_path` con
	net+opt+sched+scaler+contadores de early-stop+rng. `resume=True` retoma
	desde ese checkpoint REINICIANDO el epoch en curso (la permutación es
	determinista seed+epoch) — pérdida máxima <1 epoch; las corridas que
	terminan bien borran el checkpoint. El determinismo exacto (rng fresco,
	seed fijo) se aplica a los reportes single-shot del paper; el resume es
	un mecanismo de crash-recovery."""
	torch.manual_seed(seed)
	src = {"unsw": "unsw", "cic": "cic", "ember2024_v1": "ember"}[run.source]
	bs_def, ep_def = TRAIN_TABLE[run.scale][src]
	bs = batch_size or bs_def
	n_epochs = epochs or ep_def
	lr = LR_TABLE[run.scale]
	amp_mode = pick_amp(device, amp)

	# ── ablations E6 ──────────────────────────────────────────────────
	if ablation == "a1": # sin rama textual: el net no ve embeddings
	run.vocab = None
	for part in ("train_txt", "val_txt", "test_txt"):
	setattr(run, part, None)
	elif ablation == "a4": # ω≡1: CE sin tilt, de-tilt identidad
	run.omega = torch.ones_like(run.omega)
	if run.omega_ddos is not None:
	run.omega_ddos = torch.ones_like(run.omega_ddos)
	elif ablation == "a5" and run.source == "ember2024_v1":
	# 10-way sin unknown: la clase sale del candidate set; sus filas
	# (benignos + sin familia + cola) quedan fuera del CE de choice vía
	# el sentinela -1 del imap y la máscara de _batch_losses
	from .losses import class_weights

	unk = run.candidates.index("unknown")
	counts = run.meta.get("train_counts", {})
	run.candidates = [c for c in run.candidates if c != "unknown"]
	run.omega = torch.tensor(class_weights([counts[c] for c in run.candidates]))
	# la red se construye desde run.heads (dataclass frozen): reemplazarlo
	import dataclasses
	run.heads = dataclasses.replace(run.heads,
	n_candidates=len(run.candidates))
	for part in ("train_labels", "val_labels", "test_labels"):
	lab = getattr(run, part, None)
	if lab is None or lab.y_fam is None:
	continue
	yf = lab.y_fam
	lab.y_fam = torch.where(yf == unk, torch.full_like(yf, -1),
	torch.where(yf > unk, yf - 1, yf))
	elif ablation == "a6": # dropout 0
	import src.models.network as _net_mod
	_net_mod.DROPOUT = 0.0
	elif ablation in ("a2_noul", "a2_choice", "a2_score"):
	# A2 sin multi-task (§7): UN modelo por cabeza con el mismo
	# presupuesto total — las cabezas aportan <1% de params, la anchura
	# no requiere ajuste. Solo la cabeza elegida recibe pérdida y M_val;
	# las demás quedan sin gradiente y el dump combinado las excluye.
	run.meta["a2_head"] = ablation.removeprefix("a2_")

	log(f"[train] {run.source}/{run.scale}/{run.split_scheme} bs={bs} epochs={n_epochs} "
	f"lr={lr} amp={amp_mode} K={run.heads.n_candidates} device={device}"
	+ (f" ablation={ablation}" if ablation else ""))

	net = VXOneNet(src, run.scale, run.heads,
	d_num=len(run.feature_stage.impute_median),
	vocab_sizes={k: v.size for k, v in (run.vocab or {}).items()})
	net.to(device)
	kendall = KendallMultiTask(run.heads.names).to(device)

	ember = run.source == "ember2024_v1"
	if ember:
	from .data import EmberBatcher
	batcher = EmberBatcher(run)
	rows = [pq.ParquetFile(str(p)).metadata.num_rows for p in run.ember_train_shards]
	n_train = sum(rows)
	steps_per_epoch = sum(r // bs for r in rows)
	else:
	n_train = run.train.shape[0]
	steps_per_epoch = math.ceil(n_train / bs)

	total_steps = n_epochs * steps_per_epoch
	opt, sched = _make_optimizer(net, lr, total_steps)
	scaler = torch.amp.GradScaler(device, enabled=amp_mode == "fp16")
	amp_dtype = {"bf16": torch.bfloat16, "fp16": torch.float16}.get(amp_mode)
	params = [p for g in opt.param_groups for p in g["params"]]

	# ── checkpointing de recuperación (atomic tmp→replace) ───────────
	def _save_ckpt(epoch_done: int, best_m: float, best_epoch: int,
	bad: int, best_state) -> None:
	"""epoch_done = nº de epochs COMPLETOS (resume re-playea el actual)."""
	if not ckpt_path:
	return
	tmp = f"{ckpt_path}.tmp"
	Path(tmp).parent.mkdir(parents=True, exist_ok=True)
	torch.save({
	"net": {k: v.cpu() for k, v in net.state_dict().items()},
	"opt": opt.state_dict(), "sched": sched.state_dict(),
	"scaler": scaler.state_dict() if amp_mode == "fp16" else None,
	"amp": amp_mode, "epoch_done": epoch_done, "step": step,
	"best_m": best_m, "best_epoch": best_epoch, "bad": bad,
	"best_state": ({k: v.cpu() for k, v in best_state.items()}
	if best_state else None),
	"torch_rng": torch.get_rng_state(),
	}, tmp)
	os.replace(tmp, ckpt_path)

	start_epoch, step = 0, 0
	best_m, best_epoch, bad, best_state = -1.0, -1, 0, None
	if resume and ckpt_path and Path(ckpt_path).exists():
	ck = torch.load(ckpt_path, map_location="cpu", weights_only=False)
	net.load_state_dict(ck["net"])
	opt.load_state_dict(ck["opt"])
	for st in opt.state.values(): # AdamW: estados a device de los params
	for k, v in st.items():
	if torch.is_tensor(v):
	st[k] = v.to(device)
	# sched fresco: el epoch en curso se REPLAYEA completo (permutación
	# determinista seed+epoch), así last_epoch vuelve a alinear con
	# e*steps_per_epoch y el coseno termina en total_steps como single-shot.
	_, sched = _make_optimizer(net, lr, total_steps)
	if amp_mode == "fp16" and ck.get("scaler"):
	scaler.load_state_dict(ck["scaler"])
	if ck.get("torch_rng") is not None:
	torch.set_rng_state(ck["torch_rng"])
	best_m, best_epoch, bad = ck["best_m"], ck["best_epoch"], ck["bad"]
	best_state = ck.get("best_state")
	start_epoch = ck["epoch_done"]
	step = start_epoch * steps_per_epoch
	log(f"[train] RESUME desde {ckpt_path}: epochs completados "
	f"{start_epoch}/{n_epochs} M_val={best_m:.4f} (best ep {best_epoch})")

	def run_epoch(epoch: int) -> None:
	nonlocal step
	net.train()
	if ember:
	batches = ((u.to(device), None,
	_BatchLabels(yf.to(device), yb.to(device)))
	for u, yf, yb in batcher.train_batches(epoch, bs))
	else:
	order = np.random.default_rng(seed + epoch).permutation(n_train)
	txt_tr = run.train_txt or {}
	label_tr = run.train_labels

	def _gen():
	for s in range(0, n_train, bs):
	idx = torch.tensor(order[s:s + bs])
	u_b = run.train[idx].to(device)
	txt_b = {k: v[idx].to(device) for k, v in txt_tr.items()} or None
	yield u_b, txt_b, _BatchLabels(
	label_tr.y_fam[idx].to(device),
	label_tr.y_bin[idx].to(device),
	label_tr.y_sev[idx].to(device)
	if label_tr.y_sev is not None else None,
	label_tr.ddos_mask[idx].to(device)
	if label_tr.ddos_mask is not None else None,
	label_tr.y_ddos[idx].to(device)
	if label_tr.y_ddos is not None else None)
	batches = _gen()
	for u_b, txt_b, lab_b in batches:
	with torch.autocast(device_type=device.split(":")[0], dtype=amp_dtype,
	enabled=amp_dtype is not None and device.startswith("cuda")):
	losses, active = _batch_losses(net, run, u_b, txt_b, lab_b)
	total, _lambdas = kendall(losses, active)
	opt.zero_grad(set_to_none=True)
	if amp_mode == "fp16":
	scaler.scale(total).backward()
	scaler.unscale_(opt)
	torch.nn.utils.clip_grad_norm_(params, CLIP_NORM)
	scaler.step(opt)
	scaler.update()
	else:
	total.backward()
	torch.nn.utils.clip_grad_norm_(params, CLIP_NORM)
	opt.step()
	sched.step()
	step += 1
	# checkpoint por paso (cada N): el epoch en curso se replayea al
	# retomar, así que epoch_done = epoch (NO epoch+1)
	if ckpt_path and ckpt_every_steps and step % ckpt_every_steps == 0:
	_save_ckpt(epoch, best_m, best_epoch, bad, best_state)

	def val_outputs() -> dict[str, np.ndarray]:
	net.eval()
	if ember:
	return _eval_ember(net, run, batcher, "val", device)
	txt_va = {k: v.to(device) for k, v in (run.val_txt or {}).items()}
	return _forward_numpy(net, run.val.to(device), txt_va or None,
	run.val_labels, device)

	t0 = time.time()
	for epoch in range(start_epoch, n_epochs):
	run_epoch(epoch)
	names = run.heads.names
	if run.meta.get("a2_head"): # A2: early-stop solo sobre su cabeza
	names = (run.meta["a2_head"],)
	active_m, mv = _m_val(val_outputs(), names)
	marker = ""
	if mv > best_m:
	best_m, best_epoch, bad = mv, epoch, 0
	best_state = copy.deepcopy(net.state_dict())
	marker = " ←"
	else:
	bad += 1
	log(f"[train] epoch {epoch + 1}/{n_epochs} M_val={mv:.4f} "
	+ " ".join(f"{k}={v:.3f}" for k, v in active_m.items())
	+ f" bad={bad}{marker}")
	if ckpt_path: # fin de epoch: epoch_done = epoch+1 (barato, seguro)
	_save_ckpt(epoch + 1, best_m, best_epoch, bad, best_state)
	if bad >= PATIENCE:
	log(f"[train] early stop (patience {PATIENCE}) en epoch {epoch + 1}")
	break

	# ── restaurar mejor epoch + calibración en val (§5) ──────────────
	net.load_state_dict(best_state)
	o = val_outputs()

	if ablation == "a3": # sin calibración: T=1 en todas las cabezas
	temperatures = {"noul": 1.0}
	if "corn_logits" in o:
	temperatures["score"] = 1.0
	if "y_ddos" in o:
	temperatures["ddos"] = 1.0
	temperatures["choice"] = 1.0
	else:
	temperatures = {"noul": fit_noul_temperature(
	torch.tensor(o["noul_logits"]), torch.tensor(o["y_bin"]))}
	if "corn_logits" in o:
	temperatures["score"] = fit_corn_temperature(
	torch.tensor(o["corn_logits"]), torch.tensor(o["y_sev"]))
	if "y_ddos" in o:
	temperatures["ddos"] = fit_choice_temperature(
	torch.tensor(o["ddos_logits"]), torch.tensor(o["y_ddos"]), run.omega_ddos)
	# Choice: solo filas dentro del candidate set (y_fam ≥ 0, D12)
	fam_mask = o["y_fam"] >= 0
	temperatures["choice"] = fit_choice_temperature(
	torch.tensor(o["choice_logits"][fam_mask]),
	torch.tensor(o["y_fam"][fam_mask]), run.omega)

	run_id = f"{src}_{run.scale}_{run.split_scheme}"
	if ablation:
	run_id += f"_abl_{ablation}"
	# A5 (unknown off): τ_abs→1 ⇒ nadie abtiene (UNSW/CIC) o modo abstain
	# puro tras quitar la clase (EMBER); el candidate set sin unknown ya lo
	# dejó train_run en run.candidates.
	tau = 1.0 if ablation == "a5" else TAU_ABS
	save_artifact(out_path, net, run_meta={
	"source_key": SOURCE_KEY[src],
	"feature_stage": run.feature_stage,
	"vocab": run.vocab,
	"candidates": run.candidates,
	"omega": run.omega,
	"omega_ddos": run.omega_ddos,
	"ddos_candidates": run.ddos_candidates,
	"mode": "abstain" if ablation == "a5" else ("trained_unknown" if ember else "abstain"),
	"tau_abs": tau,
	"config": {"scale": run.scale, "split_scheme": run.split_scheme,
	"batch_size": bs, "lr": lr, "amp": amp_mode, "seed": seed},
	"manifest": {
	"run_id": run_id, "created_at": time.strftime("%Y-%m-%dT%H:%M:%S"),
	"seed": seed, "best_epoch": best_epoch, "m_val": best_m,
	"n_train": n_train, "candidates": run.candidates,
	"train_counts": run.meta.get("train_counts", {}),
	"exclusion": run.meta.get("exclusion", {}),
	"temperatures": temperatures,
	"wall_time_s": round(time.time() - t0, 1),
	},
	}, temperatures=temperatures, d_num=len(run.feature_stage.impute_median))
	log(f"[train] artefacto → {out_path} (best_epoch={best_epoch} M_val={best_m:.4f})")
	if ckpt_path: # corrida completa: el ckpt de recuperación ya no hace falta
	Path(ckpt_path).unlink(missing_ok=True)
	return {"run_id": run_id, "best_epoch": best_epoch, "m_val": best_m,
	"temperatures": temperatures, "wall_time_s": time.time() - t0}