AareLC_train/tools/train_run.py

#!/usr/bin/env python3
"""Train one ledger run end to end: train → registry → gold measure → gate → promote.

    python tools/ledger.py new --dataset dataset/2026-10-02_alc --tag alc --base yolo26n-seg.pt
    python tools/train_run.py 2026-10-02_alc --device 0 2>&1 | tee runs/2026-10-02_alc/train.log

Every retraining starts from the base weights on the whole dataset (the cumulative
labelled pool), so runs stay comparable. The gate compares the gold measure with the
deployed model's own gold report (``models/registry/<active>/gold_report.json``):
a pass needs no drop beyond 0.005 in overall F1 and in the mask score (mean mask IoU
over all gold objects, unmatched = 0), an improvement in at least one of them, and no
drop beyond 0.02 in F1 for loop_all or pin. On a pass the
run is promoted on its own: ``models/active/production.json`` is rewritten. The
first run, with no baseline, only records its report as the future baseline.
"""

from __future__ import annotations

import argparse
import csv
import json
import shutil
import socket
import sys
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, Optional

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

from tools import ledger  # noqa: E402

# Recipe since 2026-09-29: no horizontal flip (the camera never shows the pin on the left) and
# full-resolution masks (mask_ratio 1), the combination that kept run A's detection with the
# tightest masks; before that fliplr=0.5 and mask_ratio=4, as in the legacy train_yolo_seg.py.
AUGMENT = dict(scale=0.7, translate=0.15, mosaic=0.5, close_mosaic=10, fliplr=0.0, flipud=0.3,
               shear=5.0, perspective=0.0, mask_ratio=1)
GATE_TOL = 0.005
CLASS_TOL = 0.02
# Since 2026-10-02 the gate knows how noisy the gold set is: when the baseline's predictions are
# at hand, a paired bootstrap over the gold frames decides. The candidate must beat the baseline
# on F1 or on the mask score in WIN_CLEAR of the draws, and must not lose the other one in more
# than 1 - WIN_FLOOR of them. The 0.005 point rule stays as the fallback and is always reported.
WIN_CLEAR = 0.90
WIN_FLOOR = 0.10
GUARDED_CLASSES = ("loop_all", "pin")


def dataset_dir(run: Dict[str, Any]) -> Path:
    name = str(run["dataset"]).split("/", 1)[-1]
    d = ROOT / "data" / "datasets" / name
    if not (d / "dataset.yaml").exists():
        raise SystemExit(f"dataset {run['dataset']} not found at {d} (git checkout the tag, git lfs pull)")
    return d


def resolved_yaml(run_dir: Path, ds: Path) -> Path:
    text = (ds / "dataset.yaml").read_text(encoding="utf-8").replace("path: .", f"path: {ds.resolve()}", 1)
    out = run_dir / "dataset.resolved.yaml"
    out.write_text(text, encoding="utf-8")
    return out


def base_weights_path(spec: str) -> str:
    p = Path(spec)
    if p.exists():
        return str(p)
    cand = ROOT / "models" / "base" / spec
    return str(cand) if cand.exists() else spec  # else an Ultralytics model name, downloaded


def results_summary(train_dir: Path) -> Dict[str, Any]:
    csv_path = train_dir / "results.csv"
    if not csv_path.exists():
        return {}
    rows = list(csv.DictReader(csv_path.open(encoding="utf-8")))
    if not rows:
        return {}
    last = {k.strip(): v for k, v in rows[-1].items()}
    out: Dict[str, Any] = {"epochs_run": len(rows)}
    for k, v in last.items():
        if "metrics/" in k:
            try:
                out[k.replace("metrics/", "")] = float(v)
            except ValueError:
                pass
    return out


def measure(weights: Path, run_dir: Path, names: list, conf: float = 0.25) -> Dict[str, Any]:
    from src.eval.evaluate import evaluate_dirs
    from src.eval.yolo_labels import PolygonInstance, write_label_file
    from ultralytics import YOLO

    gold = ROOT / "data" / "gold"
    img_dir, gold_labels = gold / "images", gold / "labels"
    if not img_dir.exists() or not gold_labels.exists():
        raise SystemExit("no gold set under data/gold (import it with tools/import_studio_export.py --gold)")
    pred_dir = run_dir / "gold_predictions"
    pred_dir.mkdir(parents=True, exist_ok=True)
    model = YOLO(str(weights))
    for img_path in sorted(p for p in img_dir.iterdir() if p.suffix.lower() in (".jpg", ".jpeg", ".png")):
        res = model.predict(str(img_path), conf=conf, verbose=False)[0]
        insts = []
        if res.masks is not None:
            h, w = res.orig_shape
            for poly, cls in zip(res.masks.xy, res.boxes.cls.tolist()):
                pts = [(float(x) / w, float(y) / h) for x, y in poly]
                if len(pts) >= 3:
                    insts.append(PolygonInstance(int(cls), pts))
        write_label_file(pred_dir / f"{img_path.stem}.txt", insts)
    return evaluate_dirs(pred_dir, gold_labels, names=names)


def mask_score(overall: Dict[str, Any]) -> float:
    """Mean mask IoU over every gold object, an unmatched object counting as 0.

    ``mean_mask_iou`` alone is averaged over matched objects only, so a model that
    finds only the easy objects scores higher than one that also finds the hard ones
    with looser masks; run 1 on 2026-09-28 lost the gate exactly that way. Weighting
    by recall removes that bias.
    """
    tp, fn = float(overall.get("tp", 0)), float(overall.get("fn", 0))
    return float(overall.get("mean_mask_iou", 0.0)) * (tp / (tp + fn) if tp + fn else 0.0)


def gate(report: Dict[str, Any], baseline: Optional[Path], paired: Optional[Dict[str, Any]] = None) -> Optional[Dict[str, Any]]:
    if not baseline or not Path(baseline).exists():
        return None
    base = json.loads(Path(baseline).read_text(encoding="utf-8"))
    f1_delta = float(report["overall"]["f1"]) - float(base["overall"]["f1"])
    iou_delta = mask_score(report["overall"]) - mask_score(base["overall"])
    per_class = report.get("per_class") or report.get("classes") or {}
    base_classes = base.get("per_class") or base.get("classes") or {}
    class_deltas = {}
    for c in GUARDED_CLASSES:
        a, b = per_class.get(c) or {}, base_classes.get(c) or {}
        if "f1" in a and "f1" in b:
            class_deltas[c] = float(a["f1"]) - float(b["f1"])
    guarded = all(d >= -CLASS_TOL for d in class_deltas.values())
    ok_point = f1_delta >= -GATE_TOL and iou_delta >= -GATE_TOL and (f1_delta > 0 or iou_delta > 0) and guarded
    out = {"baseline": str(baseline), "f1_delta": round(f1_delta, 4), "mask_iou_delta": round(iou_delta, 4),
           "mask_score": round(mask_score(report["overall"]), 4), "mask_score_baseline": round(mask_score(base["overall"]), 4),
           "matched_only_iou_delta": round(float(report["overall"]["mean_mask_iou"]) - float(base["overall"]["mean_mask_iou"]), 4),
           "class_f1_delta": {k: round(v, 4) for k, v in class_deltas.items()}, "pass_point_rule": bool(ok_point)}
    if paired and paired.get("frames"):
        f, m = paired["f1"], paired["mask_score"]
        improved = f["win"] >= WIN_CLEAR or m["win"] >= WIN_CLEAR
        not_worse = f["win"] >= WIN_FLOOR and m["win"] >= WIN_FLOOR
        out.update({"rule": "paired bootstrap", "bootstrap": paired, "pass": bool(improved and not_worse and guarded)})
    else:
        out.update({"rule": "point", "pass": bool(ok_point)})
    return out


def environment() -> Dict[str, Any]:
    """What a measurement depends on besides the weights: the library versions decide how masks are
    decoded and filtered, so two machines can score the same weights differently."""
    env: Dict[str, Any] = {"host": socket.gethostname()}
    for mod in ("ultralytics", "torch"):
        try:
            env[mod] = __import__(mod).__version__
        except Exception:
            env[mod] = None
    return env


def same_environment(report_env: Optional[Dict[str, Any]], here: Dict[str, Any]) -> bool:
    """Was a stored report measured with the libraries of this machine? Reports from before
    2026-10-02 carry no environment: they were all measured on the Spark."""
    if not report_env:
        return str(here.get("host", "")).startswith("mx-spark")
    return all(report_env.get(k) == here.get(k) for k in ("ultralytics", "torch"))


def baseline_here(baseline: Optional[Path], names: list, conf: float) -> tuple:
    """The baseline report to gate against, measured in THIS environment.

    On 2026-10-02 a cluster run was gated against a baseline measured on the Spark: the same
    weights scored F1 0.840 and mask score 0.706 on the cluster, 0.861 and 0.757 on the Spark,
    and the gate failed a model that passes it in a single environment. So when the stored
    baseline report comes from other library versions, the baseline weights are measured again
    here (predictions kept under runs/_measure_<version>, where the paired bootstrap finds
    them). Returns (report path or None, note or None); a baseline that cannot be measured here
    gives (None, why) and the caller refuses to promote.
    """
    if not baseline or not Path(baseline).exists():
        return baseline, None
    base = json.loads(Path(baseline).read_text(encoding="utf-8"))
    here = environment()
    if same_environment(base.get("environment"), here):
        return Path(baseline), None
    version = Path(baseline).parent.name
    weights = Path(baseline).parent / "weights.pt"
    if not weights.exists() or weights.stat().st_size < 1_000_000:      # absent, or a git-lfs pointer
        return None, (f"the baseline {version} was measured with {base.get('environment') or 'the Spark environment'}, this machine has "
                      f"{here}; its weights are not here to measure again (git lfs pull -I 'models/registry/{version}/weights.pt')")
    work = ROOT / "runs" / f"_measure_{version}"
    work.mkdir(parents=True, exist_ok=True)
    rep = measure(weights, work, names, conf=conf)
    rep["environment"] = here
    local = work / "gold_report.local.json"
    local.write_text(json.dumps(rep, indent=2) + "\n", encoding="utf-8")
    print(f"gate: baseline {version} measured again in this environment (ultralytics {here.get('ultralytics')}, torch {here.get('torch')}): "
          f"F1 {rep['overall']['f1']:.4f} (stored report: {base['overall']['f1']:.4f})")
    return local, f"baseline measured again here; stored report came from {base.get('environment') or 'the Spark'}"


def baseline_predictions(baseline: Optional[Path]) -> Optional[Path]:
    """The baseline model's predictions on the gold frames, when they were kept: the re-measure
    directory of its version, or the gold predictions of the run that produced it."""
    if not baseline:
        return None
    if Path(baseline).name == "gold_report.local.json":               # measured again here: predictions sit beside it
        d = Path(baseline).parent / "gold_predictions"
        return d if d.is_dir() else None
    version = Path(baseline).parent.name
    d = ROOT / "runs" / f"_measure_{version}" / "gold_predictions"
    if d.is_dir():
        return d
    for rj in sorted((ROOT / "runs").glob("*/run.json")):
        try:
            if json.loads(rj.read_text(encoding="utf-8")).get("version") == version and (rj.parent / "gold_predictions").is_dir():
                return rj.parent / "gold_predictions"
        except (OSError, ValueError):
            continue
    return None


def paired_with_baseline(run_dir: Path, baseline: Optional[Path]) -> Optional[Dict[str, Any]]:
    """Paired bootstrap of this run against the baseline on the gold frames both were measured on."""
    base_pred = baseline_predictions(baseline)
    if base_pred is None:
        return None
    from tools import gate_stats
    gold = ROOT / "data" / "gold" / "labels"
    cand = gate_stats.per_frame(run_dir / "gold_predictions", gold)
    base = gate_stats.per_frame(base_pred, gold)
    missing = len(cand) - len(set(cand) & set(base))
    if missing:
        print(f"gate: the baseline has no prediction for {missing} gold frames (gold grew since it was measured); pairing on the rest")
    return gate_stats.paired_bootstrap(cand, base, groups=gate_stats.load_groups(gold.parent)["mount"])


def active_version() -> Optional[str]:
    p = ROOT / "models" / "active" / "production.json"
    if not p.exists():
        return None
    return json.loads(p.read_text(encoding="utf-8")).get("version")


def promote(version: str, run_id: str, g: Dict[str, Any]) -> Path:
    p = ROOT / "models" / "active" / "production.json"
    p.parent.mkdir(parents=True, exist_ok=True)
    p.write_text(json.dumps({
        "version": version,
        "export": "weights.pt",   # engine built on the target GPU at deploy time
        "promoted_on": datetime.now().strftime("%Y-%m-%d"),
        "promoted_by": "gate",
        "run": run_id,
        "gate": g,
    }, indent=2) + "\n", encoding="utf-8")
    return p


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("run_id")
    ap.add_argument("--epochs", type=int, default=400)
    ap.add_argument("--patience", type=int, default=50)
    ap.add_argument("--imgsz", type=int, default=640)
    ap.add_argument("--batch", type=int, default=16)
    ap.add_argument("--device", default="0")
    ap.add_argument("--workers", type=int, default=8)
    ap.add_argument("--conf", type=float, default=0.25, help="confidence used for the gold measure")
    ap.add_argument("--baseline", default=None, help="gold report to gate against (default: the active model's)")
    ap.add_argument("--no-promote", action="store_true", help="gate, but leave production.json alone")
    ap.add_argument("--skip-train", action="store_true", help="reuse runs/<id>/ultralytics/train/weights/best.pt")
    ap.add_argument("--force", action="store_true", help="train again a run that already has a gate verdict")
    ap.add_argument("--set", nargs="*", default=[], metavar="KEY=VALUE",
                    help="override any Ultralytics train argument, e.g. fliplr=0 mask_ratio=1 overlap_mask=true; recorded in the ledger")
    args = ap.parse_args(argv)

    run = ledger.load(args.run_id)
    run_dir = ledger.run_path(args.run_id).parent
    # A gate verdict is an outcome, not a crash: the exit code 2 of a failed gate must not make a
    # wrapper's retry train the whole run a second time (cycle 4 lost two hours that way).
    if run.get("status") in ("gate_failed", "gate_passed", "promoted", "deployed") and not (args.skip_train or args.force):
        print(f"{args.run_id} is already {run['status']}: not training it again (--force to retrain, --skip-train to re-measure)")
        return 0
    ds = dataset_dir(run)
    data_yaml = resolved_yaml(run_dir, ds)
    names = [v for _, v in sorted(json.loads((ds / "source.json").read_text())["classes"].items(), key=lambda kv: int(kv[0]))] \
        if (ds / "source.json").exists() else ["loop_all", "pin", "crystal", "loop_face", "ice", "needle"]
    # A run may carry its own training arguments (ledger new --kv train_args=...): the cluster
    # job passes no --set, so they come from the ledger; a --set on the command line still wins.
    overrides: Dict[str, Any] = dict(run.get("train_args") or {})
    for kv in args.set or []:
        if "=" not in kv:
            raise SystemExit(f"--set expects KEY=VALUE, got {kv!r}")
        k, v = kv.split("=", 1)
        try:
            overrides[k] = json.loads(v.lower() if v.lower() in ("true", "false") else v)
        except json.JSONDecodeError:
            overrides[k] = v
    hyper = {"epochs": args.epochs, "patience": args.patience, "imgsz": args.imgsz, "batch": args.batch,
             "device": args.device, "overlap_mask": False, **AUGMENT, **overrides}
    ledger.set_status(args.run_id, "training", note="train_run.py", kv={
        "hyperparams": hyper, "host": socket.gethostname(), "started": datetime.now().isoformat(timespec="seconds")})

    best = run_dir / "ultralytics" / "train" / "weights" / "best.pt"
    if not args.skip_train:
        from ultralytics import YOLO
        model = YOLO(base_weights_path(run["base_weights"]))
        train_kwargs = dict(epochs=args.epochs, patience=args.patience, imgsz=args.imgsz, batch=args.batch,
                            device=args.device, workers=args.workers, overlap_mask=False, amp=False, **AUGMENT)
        train_kwargs.update(overrides)   # --set wins over the recipe
        model.train(data=str(data_yaml), project=str(run_dir / "ultralytics"), name="train", exist_ok=True, **train_kwargs)
    if not best.exists():
        ledger.set_status(args.run_id, "aborted", note=f"no best.pt at {best}")
        return 1

    family = Path(run["base_weights"]).stem.split("-")[0]
    tag = args.run_id.split("_", 1)[-1]
    version = f"{family}-seg-{tag}-{datetime.now().strftime('%Y-%m-%d')}"
    reg = ROOT / "models" / "registry" / version
    reg.mkdir(parents=True, exist_ok=True)
    shutil.copy2(best, reg / "weights.pt")

    report = measure(reg / "weights.pt", run_dir, names, conf=args.conf)
    report["environment"] = environment()
    baseline = Path(args.baseline) if args.baseline else None
    if baseline is None:
        av = active_version()
        if av and (ROOT / "models" / "registry" / av / "gold_report.json").exists():
            baseline = ROOT / "models" / "registry" / av / "gold_report.json"
    stored_baseline = baseline
    baseline, env_note = baseline_here(baseline, names, args.conf)
    paired = None
    try:
        paired = paired_with_baseline(run_dir, baseline)
    except Exception as e:  # the point rule still stands
        print(f"gate: paired bootstrap unavailable ({e}); falling back to the point rule")
    g = gate(report, baseline, paired)
    if g is not None:
        g["baseline"] = str(stored_baseline)
        if env_note:
            g["environment_note"] = env_note
    elif stored_baseline is not None and env_note:
        # there is a baseline, but not one this machine can be compared with: no verdict, no promotion
        g = {"baseline": str(stored_baseline), "pass": False, "rule": "none", "environment_note": env_note}
        print(f"gate: no verdict: {env_note}")
    report["gate"] = g
    (reg / "gold_report.json").write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")
    metrics = {"overall": report["overall"], "gate": g, "training": results_summary(run_dir / "ultralytics" / "train"),
               "version": version, "measured_at": datetime.now().isoformat(timespec="seconds")}
    (run_dir / "metrics.json").write_text(json.dumps(metrics, indent=2) + "\n", encoding="utf-8")
    (reg / "metadata.json").write_text(json.dumps({
        "model_id": version, "family": family, "task": "segmentation",
        "training_date": datetime.now().strftime("%Y-%m-%d"), "run": args.run_id,
        "dataset": run["dataset"], "base_weights": run["base_weights"], "hyperparams": hyper,
        "export_formats": ["pt"], "metrics": metrics["overall"], "gate": g,
    }, indent=2) + "\n", encoding="utf-8")
    ledger.record_metrics(args.run_id, run_dir / "metrics.json")
    ledger.set_status(args.run_id, "measured", kv={"version": version})

    if g is None:
        print(f"measured; no baseline, {version} is the baseline for the next run")
        return 0
    if g["pass"]:
        ledger.set_status(args.run_id, "gate_passed", note=json.dumps(g))
        if not args.no_promote:
            promote(version, args.run_id, g)
            ledger.set_status(args.run_id, "promoted", note="by the gate")
            print(f"PROMOTED {version}: {g}")
        return 0
    ledger.set_status(args.run_id, "gate_failed", note=json.dumps(g))
    print(f"gate failed: {g}")
    return 2


if __name__ == "__main__":
    raise SystemExit(main())