scripts/train.py

#!/usr/bin/env python
"""Training entry point for the main StapleBridge model (Full Exact-SB).

This is the orchestrator that produced the released checkpoint:

* full 4020-lead training each epoch;
* full validation on every validation lead (111), every epoch;
* a fixed 10 epochs, no early stopping (all 10 always run);
* two checkpoints maintained independently, each updated after that epoch's
  validation pass:
    - ``best_kl.pt``: minimum ``q_star_vs_q_theta_kl`` -- the selection rule;
    - ``best_pv.pt``: maximum ``mean_delta_penetrance_vs_original_lead``,
      recorded for monitoring only and not used to select the released model;
* per-epoch logging of both metrics with their running bests.

``checkpoints/staplebridge_seed42_best.pt`` is the ``best_kl.pt`` of this run:
the epoch minimising ``q_star_vs_q_theta_kl`` on the validation split. Model,
loss, Exact-SB, property scoring, decoding and every other training setting are
read from the config.
"""
from __future__ import annotations

import argparse
import gc
import json
import random
import sys
import time
from pathlib import Path
from typing import Any

import numpy as np
import torch
import yaml

PACKAGE_ROOT = Path(__file__).resolve().parents[1]
if str(PACKAGE_ROOT) not in sys.path:
    sys.path.insert(0, str(PACKAGE_ROOT))

from staplebridge.data.dataset import load_leads  # noqa: E402
from staplebridge.hydrocarbon.exact_sb_cache import build_cache_from_config  # noqa: E402
from staplebridge.hydrocarbon.plan_control import (  # noqa: E402
    HydrocarbonPlanControlConfig, build_hydrocarbon_plan_head,
)
from staplebridge.hydrocarbon.property_energy import (  # noqa: E402
    HydrocarbonPropertyEnergyConfig,
    HydrocarbonPropertyScorer,
    required_original_lead_properties,
)
from staplebridge.hydrocarbon.tokenizer import tokenize_sequence  # noqa: E402
from staplebridge.training.main_loop import (  # noqa: E402
    train_enabled_epoch, validate_enabled,
)
from staplebridge.training.records import write_json, write_jsonl  # noqa: E402
from staplebridge.training.stack import (  # noqa: E402
    build_energy, build_models, build_predictor, build_stack, load_config,
    seed_everything, select_leads,
)

FULL_TRAIN_N = 4020
SEED = 42
EPOCHS = 10
KL_KEY = "q_star_vs_q_theta_kl"
PV_KEY = "mean_delta_penetrance_vs_original_lead"


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--config", type=Path,
                        default=PACKAGE_ROOT / "configs/staplebridge_main.yaml")
    parser.add_argument("--out-dir", type=Path,
                        default=PACKAGE_ROOT / "outputs/main_seed42")
    parser.add_argument("--resume", type=Path, default=None,
                        help="Checkpoint to resume from (e.g. checkpoints/epoch_008.pt). "
                             "Restores model/optimizer/RNG so remaining epochs are identical "
                             "to an uninterrupted run; appends to the existing metrics/log.")
    return parser.parse_args()


def require(condition: bool, message: str) -> None:
    if not condition:
        raise SystemExit(message)


def rng_payload() -> dict[str, Any]:
    return {
        "python_random_state": random.getstate(),
        "numpy_random_state": np.random.get_state(),
        "torch_rng_state": torch.get_rng_state(),
        "cuda_rng_state_all": torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None,
    }


def build_original_cache(config: dict[str, Any], leads: list[Any], out_dir: Path) -> dict:
    """Compute the configured unedited-lead property cache once."""
    path = out_dir / "original_linear_cache.jsonl"
    wrapper, _ = build_predictor(config)
    cache: dict[tuple[Any, ...], dict[str, Any]] = {}
    scorer = HydrocarbonPropertyScorer(wrapper, original_linear_cache=cache)
    property_cfg = HydrocarbonPropertyEnergyConfig.from_dict(
        (((config.get("hydrocarbon") or {}).get("terminal_energy") or {}).get("property"))
    )
    properties = required_original_lead_properties(property_cfg)
    started = time.perf_counter()
    if (
        property_cfg.enable_developability_constraints
        or property_cfg.enable_halflife_preservation
        or property_cfg.enable_joint_perm_halflife_support
    ):
        scorer.prefetch(
            properties,
            [
                scorer.original_linear_smiles(tokenize_sequence(lead.linear_sequence))
                for lead in leads
            ],
        )
    for index, lead in enumerate(leads):
        scorer.score_original_linear(
            tokenize_sequence(lead.linear_sequence),
            lead_key=str(lead.example_id),
            properties=properties,
        )
        if (index + 1) % 512 == 0 or index + 1 == len(leads):
            print(f"[original baseline] {index + 1}/{len(leads)}", flush=True)
    rows = [{"lead_key": k[0], "tokens": list(k[1]), "scores": v} for k, v in cache.items()]
    write_jsonl(path, rows)
    del scorer, wrapper
    import gc
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    print(f"[original baseline] built {len(cache)} entries in {time.perf_counter() - started:.1f}s", flush=True)
    return cache


def load_original_cache(path: Path) -> dict[tuple[Any, ...], dict[str, Any]]:
    cache: dict[tuple[Any, ...], dict[str, Any]] = {}
    with path.open(encoding="utf-8") as handle:
        for line in handle:
            row = json.loads(line)
            cache[(row["lead_key"], tuple(row["tokens"]))] = dict(row["scores"])
    return cache


def save_checkpoint(path: Path, payload: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_suffix(path.suffix + ".tmp")
    torch.save(payload, tmp)
    tmp.replace(path)


def main() -> None:
    args = parse_args()
    config = load_config(args.config.resolve())
    out_dir = args.out_dir.resolve()
    resuming = args.resume is not None

    if resuming:
        require(out_dir.is_dir(), f"resume requires existing output dir: {out_dir}")
        require(args.resume.is_file(), f"resume checkpoint not found: {args.resume}")
    else:
        # Never overwrite prior outputs: require a fresh/empty directory.
        if out_dir.exists():
            leftovers = [p for p in out_dir.iterdir()]
            require(not leftovers, f"output directory not empty (refusing to overwrite): {out_dir}")
        out_dir.mkdir(parents=True, exist_ok=True)

    require(int(config.get("train_n", -1)) == FULL_TRAIN_N, "train_n must be 4020")
    require(int(config["training"]["epochs"]) == EPOCHS, f"epochs must be {EPOCHS}")
    device = torch.device(str(config["training"]["device"]))
    if device.type == "cuda":
        require(torch.cuda.is_available(), f"CUDA unavailable for {device}")
        torch.cuda.set_device(device)

    (out_dir / "resolved_config.yaml").write_text(yaml.safe_dump(config, sort_keys=False))
    seed_everything(SEED)

    # ---- data: full train (4020) + FULL valid (all 111, not the 32 subset) --
    train_leads = select_leads(
        Path(config["data"]["root"]) / config["data"]["train_file"],
        FULL_TRAIN_N, int(config["data"]["max_lead_length"]),
    )
    valid_cap = int(config["validation"]["max_lead_length"])
    valid_leads = [
        lead for lead in load_leads(Path(config["data"]["root"]) / config["data"]["valid_file"])
        if len(lead.linear_sequence) <= valid_cap
    ]
    print(f"[data] train={len(train_leads)} full_valid={len(valid_leads)} (cap len<={valid_cap})", flush=True)

    # ---- original-lead penetrance cache (train + valid) for delta-PV --------
    original_path = out_dir / "original_linear_cache.jsonl"
    if resuming and original_path.is_file():
        original_cache = load_original_cache(original_path)
        print(f"[original baseline] reused {len(original_cache)} entries from {original_path}", flush=True)
    else:
        original_cache = build_original_cache(config, train_leads + valid_leads, out_dir)

    # ---- model + energy stack (identical construction to the standard run) --
    stack = build_stack(config, SEED)
    wrapper, _ = build_predictor(config)
    scorer = HydrocarbonPropertyScorer(wrapper, original_linear_cache=original_cache)
    energy_fn = build_energy(config, stack, scorer)
    policy, value, kernel, optimizer = build_models(config, stack, device)
    head = build_hydrocarbon_plan_head(
        config, int(config["model"]["emb_dim"]), device,
        esm2_prior=stack["reference_priors"].peptide,
    )
    optimizer.add_param_group({"params": list(head.parameters())})
    parameters = list(policy.parameters()) + list(value.parameters()) + list(head.parameters())
    plan_cfg = HydrocarbonPlanControlConfig.from_config(config)
    plan_rng = random.Random(SEED)

    # ---- reuse the existing persistent q* cache (fingerprint unchanged) -----
    exact_sb_cache = build_cache_from_config(config, catalog=stack["catalog"], repo_root=PACKAGE_ROOT)
    if exact_sb_cache.enabled:
        print(f"[exact-sb cache] {json.dumps(exact_sb_cache.describe(), ensure_ascii=False)}", flush=True)

    def make_payload(epoch: int) -> dict[str, Any]:
        return {
            "epoch": epoch, "config": config, "plan_control_enabled": True,
            "policy_state_dict": policy.state_dict(), "value_state_dict": value.state_dict(),
            "plan_head_state_dict": head.state_dict(), "optimizer_state_dict": optimizer.state_dict(),
            "plan_rng_state": plan_rng.getstate(), **rng_payload(),
        }

    metrics_path = out_dir / "metrics.jsonl"
    history: list[dict[str, Any]] = []
    best_kl = {"value": float("inf"), "epoch": None}
    best_pv = {"value": float("-inf"), "epoch": None}
    start_epoch = 0

    if resuming:
        ckpt = torch.load(args.resume, map_location=device, weights_only=False)
        policy.load_state_dict(ckpt["policy_state_dict"])
        value.load_state_dict(ckpt["value_state_dict"])
        head.load_state_dict(ckpt["plan_head_state_dict"])
        optimizer.load_state_dict(ckpt["optimizer_state_dict"])
        plan_rng.setstate(ckpt["plan_rng_state"])
        random.setstate(ckpt["python_random_state"])
        np.random.set_state(ckpt["numpy_random_state"])
        torch.set_rng_state(ckpt["torch_rng_state"].cpu())
        if torch.cuda.is_available() and ckpt.get("cuda_rng_state_all") is not None:
            torch.cuda.set_rng_state_all([s.cpu() for s in ckpt["cuda_rng_state_all"]])
        start_epoch = int(ckpt["epoch"])
        # Rebuild running bests + history from the persisted per-epoch metrics so
        # best_kl/best_pv provenance carries across the resume boundary.
        for line in metrics_path.read_text().splitlines():
            rec = json.loads(line)
            if int(rec["epoch"]) <= start_epoch:
                history.append(rec)
        if history:
            last = history[-1]
            best_kl = {"value": float(last["best_kl_value"]), "epoch": int(last["best_kl_epoch"])}
            best_pv = {"value": float(last["best_pv_value"]), "epoch": int(last["best_pv_epoch"])}
        print(
            f"[resume] from {args.resume} completed_epoch={start_epoch}; "
            f"best_kl={best_kl['value']:.6f}@ep{best_kl['epoch']} "
            f"best_pv={best_pv['value']:.6f}@ep{best_pv['epoch']}",
            flush=True,
        )
    run_started = time.perf_counter()

    for epoch in range(start_epoch, EPOCHS):
        policy.train(); value.train(); head.train()
        train_started = time.perf_counter()
        rows, train_metrics = train_enabled_epoch(
            train_leads, config, stack, energy_fn, policy, kernel, optimizer,
            parameters, head, plan_cfg, plan_rng, epoch, exact_sb_cache=exact_sb_cache,
        )
        train_seconds = time.perf_counter() - train_started
        write_jsonl(out_dir / "training" / f"epoch_{epoch + 1:03d}_candidates.jsonl", rows)
        joint_train_audit = train_metrics.get("joint_perm_halflife_support")
        if joint_train_audit and epoch == 0:
            write_jsonl(
                out_dir / "joint_support_audit" / "train4020_per_lead.jsonl",
                joint_train_audit["per_lead"],
            )
            write_json(
                out_dir / "joint_support_audit" / "train4020_summary.json",
                {k: v for k, v in joint_train_audit.items() if k != "per_lead"},
            )

        policy.eval(); value.eval(); head.eval()
        valid_started = time.perf_counter()
        selected, validation = validate_enabled(
            valid_leads, config, stack, energy_fn, policy, kernel, head,
            exact_sb_cache=exact_sb_cache,
        )
        valid_seconds = time.perf_counter() - valid_started
        valid_dir = out_dir / "validation" / f"epoch_{epoch + 1:03d}"
        write_jsonl(valid_dir / "selected.jsonl", selected)
        write_json(valid_dir / "summary.json", validation)
        joint_valid_audit = validation.get("joint_perm_halflife_support")
        if joint_valid_audit:
            write_jsonl(
                valid_dir / "joint_support_per_lead.jsonl",
                joint_valid_audit["per_lead"],
            )
            write_json(
                valid_dir / "joint_support_summary.json",
                {k: v for k, v in joint_valid_audit.items() if k != "per_lead"},
            )

        kl = float(validation[KL_KEY])
        pv = float(validation[PV_KEY])
        guard_cfg = dict(config.get("guardrails") or {})
        both_topologies_present = bool(
            validation.get("s5_s5_i4_count", 0) > 0
            and validation.get("r8_s5_i7_count", 0) > 0
        )
        checkpoint_eligible = bool(
            not guard_cfg.get("require_both_topologies", False)
            or both_topologies_present
        )
        improved_kl = checkpoint_eligible and kl < best_kl["value"]
        improved_pv = checkpoint_eligible and pv > best_pv["value"]
        if improved_kl:
            best_kl = {"value": kl, "epoch": epoch + 1}
            save_checkpoint(out_dir / "checkpoints" / "best_kl.pt", make_payload(epoch + 1))
        if improved_pv:
            best_pv = {"value": pv, "epoch": epoch + 1}
            save_checkpoint(out_dir / "checkpoints" / "best_pv.pt", make_payload(epoch + 1))
        # Per-epoch + latest snapshots so no epoch is lost (fresh dir; nothing
        # is overwritten across runs).
        payload = make_payload(epoch + 1)
        save_checkpoint(out_dir / "checkpoints" / f"epoch_{epoch + 1:03d}.pt", payload)
        save_checkpoint(out_dir / "checkpoints" / "latest.pt", payload)

        record = {
            "epoch": epoch + 1,
            "valid_kl": kl,
            "valid_delta_penetrance": pv,
            "best_kl_epoch": best_kl["epoch"], "best_kl_value": best_kl["value"], "kl_improved": improved_kl,
            "best_pv_epoch": best_pv["epoch"], "best_pv_value": best_pv["value"], "pv_improved": improved_pv,
            "valid_product_penetrance": validation.get("mean_product_penetrance"),
            "valid_top1_chemistry_valid_rate": validation.get("top1_chemistry_valid_rate"),
            "valid_top1_stapled_rate": validation.get("top1_stapled_rate"),
            "valid_q_star_top1_agreement": validation.get("q_star_top1_agreement"),
            "valid_q_star_spearman": validation.get("q_star_spearman"),
            "valid_edit_distance": validation.get("mean_weighted_edit_distance"),
            "both_topologies_present": both_topologies_present,
            "checkpoint_eligible": checkpoint_eligible,
            "train_loss": train_metrics.get("loss"),
            "train_plan_loss": train_metrics.get("plan_loss"),
            "train_q_star_vs_q_theta_kl": train_metrics.get("q_star_vs_q_theta_kl"),
            "train_seconds": train_seconds,
            "valid_seconds": valid_seconds,
            "epoch_seconds": train_seconds + valid_seconds,
            "train_stage_seconds": train_metrics.get("stage_seconds"),
            "exact_sb_cache": exact_sb_cache.describe() if exact_sb_cache.enabled else {"enabled": False},
            "train_joint_perm_halflife_support": (
                {k: v for k, v in (joint_train_audit or {}).items() if k != "per_lead"}
                if joint_train_audit
                else None
            ),
            "valid_joint_perm_halflife_support": (
                {
                    k: v
                    for k, v in (joint_valid_audit or {}).items()
                    if k != "per_lead"
                }
                if joint_valid_audit
                else None
            ),
        }
        history.append(record)
        write_jsonl(metrics_path, [record], mode="a")
        print(
            f"[epoch {epoch + 1}/{EPOCHS}] "
            f"valid_KL={kl:.6f} (best {best_kl['value']:.6f} @ep{best_kl['epoch']}"
            f"{' NEW' if improved_kl else ''})  "
            f"valid_deltaPV={pv:.6f} (best {best_pv['value']:.6f} @ep{best_pv['epoch']}"
            f"{' NEW' if improved_pv else ''})  "
            f"epoch_seconds={record['epoch_seconds']:.1f}",
            flush=True,
        )
        # Semantics-neutral memory hygiene: reclaim inter-epoch CUDA cache /
        # Python garbage so allocator fragmentation does not accumulate across
        # the 10 epochs. Does not touch weights, RNG, or any cached value.
        gc.collect()
        if device.type == "cuda":
            torch.cuda.empty_cache()
            print(
                f"[gpu] epoch {epoch + 1} allocated={torch.cuda.memory_allocated(device) / 2**20:.0f}MiB "
                f"reserved={torch.cuda.memory_reserved(device) / 2**20:.0f}MiB",
                flush=True,
            )

    summary = {
        "exp_name": config.get("exp_name"),
        "epochs": EPOCHS,
        "early_stopping": False,
        "train_n": FULL_TRAIN_N,
        "n_valid_leads": len(valid_leads),
        "seed": SEED,
        "best_kl": best_kl,
        "best_pv": best_pv,
        "kl_curve": [(r["epoch"], r["valid_kl"]) for r in history],
        "pv_curve": [(r["epoch"], r["valid_delta_penetrance"]) for r in history],
        "runtime_seconds": time.perf_counter() - run_started,
        "gpu_peak_mib": (torch.cuda.max_memory_allocated(device) / 2 ** 20) if device.type == "cuda" else 0,
        "checkpoints": {
            "best_kl": str(out_dir / "checkpoints" / "best_kl.pt"),
            "best_pv": str(out_dir / "checkpoints" / "best_pv.pt"),
        },
        "history": history,
    }
    write_json(out_dir / "run_summary.json", summary)
    print("\n" + json.dumps({k: v for k, v in summary.items() if k != "history"}, indent=2, default=str), flush=True)


if __name__ == "__main__":
    main()