train.py

"""Drive a single architecture end-to-end and record the whole run.

Every invocation produces one directory under runs/<id>/ containing the
manifest, the raw event stream, the generated dataset, derived metrics
and curves, per-question predictions, captured stdout, and a notes.md.
Nothing is ever overwritten — new runs get new timestamped ids.
"""

from __future__ import annotations
import hashlib
import json
import sys
import time
import traceback
from dataclasses import asdict
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, TextIO

import registry
import runs as runs_lib
from data import generate_dataset
from evaluation import EvalResult, evaluate
from manifest import build_manifest, update_status, write_manifest
from recorder import RunRecorder

REPO_ROOT = Path(__file__).resolve().parent
RUNS_DIR = REPO_ROOT / "runs"


# ---------- stdout capture -----------------------------------------------

class _Tee:
    """Write to multiple streams at once. Used so verbose runs print to
    the terminal AND land in stdout.log under the run dir."""

    def __init__(self, *streams: TextIO) -> None:
        self._streams = streams

    def write(self, s: str) -> int:
        n = 0
        for st in self._streams:
            n = st.write(s)
            st.flush()
        return n

    def flush(self) -> None:
        for st in self._streams:
            st.flush()


# ---------- run id --------------------------------------------------------

def _make_run_id(arch_name: str, seed: int, arch_config: dict[str, Any]) -> str:
    now = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H-%M-%S")
    h = hashlib.sha256(
        json.dumps(
            {"arch": arch_name, "seed": seed, "cfg": arch_config, "now": now},
            sort_keys=True,
        ).encode()
    ).hexdigest()[:4]
    return f"{now}_{arch_name}_seed{seed}_{h}"


# ---------- dataset serialization ----------------------------------------

def _serialize_dataset(dataset: Any) -> dict[str, Any]:
    return {
        "vocab": dataset.vocab,
        "entity_ids": dataset.entity_ids,
        "relation_ids": dataset.relation_ids,
        "lessons": [
            {"idx": L.idx, "facts": [[f.subject, f.relation, f.object] for f in L.facts]}
            for L in dataset.lessons
        ],
        "questions": [
            {
                "qid": q.qid,
                "kind": q.kind,
                "subject": q.subject,
                "relation": q.relation,
                "relation2": q.relation2,
                "relation3": q.relation3,
                "candidate": q.candidate,
                "source_lessons": list(q.source_lessons),
                "gold_object": q.gold_object,
                "gold_is_true": q.gold_is_true,
            }
            for q in dataset.questions
        ],
    }


# ---------- derived artifacts --------------------------------------------

def _derive_metrics_json(result: EvalResult) -> dict[str, Any]:
    return {
        "arch_name": result.arch_name,
        "seed": result.seed,
        "n_lessons": result.n_lessons,
        "recall_accuracy": result.recall_accuracy,
        "composition_accuracy": result.composition_accuracy,
        "composition_3hop_accuracy": result.composition_3hop_accuracy,
        "negative_accuracy": result.negative_accuracy,
        "overall_accuracy": result.overall_accuracy,
    }


def _derive_predictions_jsonl(run_dir: Path) -> None:
    """Extract question_answered events from events.jsonl into a flat
    predictions.jsonl view suitable for pandas/jq analysis."""
    events_path = run_dir / "events.jsonl"
    out_path = run_dir / "predictions.jsonl"
    if not events_path.exists():
        return
    with events_path.open() as src, out_path.open("w") as dst:
        for line in src:
            try:
                e = json.loads(line)
            except json.JSONDecodeError:
                continue
            if e.get("event") != "question_answered":
                continue
            dst.write(json.dumps({
                "checkpoint_lesson": e["checkpoint_lesson"],
                "qid": e["qid"],
                "kind": e["kind"],
                "prediction": e["prediction"],
                "correct": e["correct"],
            }) + "\n")


# ---------- notes template ------------------------------------------------

def _write_notes(run_dir: Path, run_id: str, note: str | None) -> None:
    p = run_dir / "notes.md"
    if p.exists():
        return
    if note:
        p.write_text(
            f"# Run {run_id}\n\n"
            f"## Hypothesis\n{note}\n\n"
            f"## Result\n<!-- fill in after inspecting metrics -->\n\n"
            f"## Surprises\n<!-- -->\n\n"
            f"## Next\n<!-- -->\n"
        )
    else:
        p.write_text(runs_lib.NOTES_TEMPLATE.format(run_id=run_id))


# ---------- index row -----------------------------------------------------

def _index_row(run_id: str, result: EvalResult, status: str) -> dict[str, Any]:
    return {
        "run_id": run_id,
        "arch_name": result.arch_name,
        "seed": result.seed,
        "recall_accuracy": result.recall_accuracy,
        "composition_accuracy": result.composition_accuracy,
        "negative_accuracy": result.negative_accuracy,
        "overall_accuracy": result.overall_accuracy,
        "status": status,
        "recorded_at": datetime.now(timezone.utc).isoformat(timespec="seconds"),
    }


# ---------- main entry ----------------------------------------------------

def train_one(
    arch_name: str,
    seed: int = 42,
    note: str | None = None,
    verbose: bool = False,
    dataset_kwargs: dict[str, Any] | None = None,
    arch_kwargs: dict[str, Any] | None = None,
) -> tuple[str, EvalResult]:
    """Run one architecture end-to-end, producing a full run directory.

    Args:
      arch_kwargs: overrides for the architecture constructor (from
        `--cfg key=value` on the CLI). Passed as keyword arguments to
        `cls(...)`, and their effect is recorded in the manifest via
        the architecture's own `config()`.

    Returns (run_id, result).
    """
    registry.discover()
    cls = registry.get(arch_name)
    dataset_config = dict(dataset_kwargs or {"seed": seed})
    dataset = generate_dataset(**dataset_config)
    arch_kwargs = dict(arch_kwargs or {})

    # Peek at the arch config with user overrides applied, so the run_id
    # hash and the manifest both reflect the overrides. Instantiate a
    # throwaway probe with the same kwargs the real system will get.
    probe = cls(dataset=dataset, seed=seed, **arch_kwargs)
    arch_config = probe.config()
    del probe

    run_id = _make_run_id(arch_name, seed, arch_config)
    run_dir = RUNS_DIR / run_id
    run_dir.mkdir(parents=True, exist_ok=True)

    manifest = build_manifest(
        run_id=run_id,
        arch_name=arch_name,
        seed=seed,
        dataset_config=dataset_config,
        arch_config=arch_config,
        argv=sys.argv,
        repo_root=REPO_ROOT,
    )
    write_manifest(run_dir, manifest)

    (run_dir / "dataset.json").write_text(
        json.dumps(_serialize_dataset(dataset), indent=2) + "\n"
    )

    stdout_path = run_dir / "stdout.log"
    stdout_f = stdout_path.open("w", encoding="utf-8")
    real_stdout = sys.stdout
    sys.stdout = _Tee(real_stdout, stdout_f) if verbose else _Tee(stdout_f)

    t0 = time.time()
    result: EvalResult | None = None
    try:
        with RunRecorder(run_dir) as recorder:
            recorder.log_event("run_started", run_id=run_id, manifest=manifest)
            recorder.log_event(
                "dataset_generated",
                n_lessons=len(dataset.lessons),
                n_questions=len(dataset.questions),
                vocab_size=len(dataset.vocab),
            )
            system = cls(dataset=dataset, seed=seed, recorder=recorder, **arch_kwargs)
            recorder.log_event(
                "system_instantiated",
                arch_name=arch_name,
                arch_config=system.config(),
            )
            result = evaluate(system, dataset, recorder=recorder, verbose=verbose)
            recorder.log_event(
                "run_finished",
                recall_accuracy=result.recall_accuracy,
                composition_accuracy=result.composition_accuracy,
                negative_accuracy=result.negative_accuracy,
                overall_accuracy=result.overall_accuracy,
                duration_s=time.time() - t0,
            )
        duration = time.time() - t0
        (run_dir / "metrics.json").write_text(
            json.dumps(_derive_metrics_json(result), indent=2) + "\n"
        )
        (run_dir / "curves.json").write_text(
            json.dumps(result.curves, indent=2) + "\n"
        )
        _derive_predictions_jsonl(run_dir)
        update_status(run_dir, "completed", duration_s=duration)
        _write_notes(run_dir, run_id, note)
        runs_lib.append_index(_index_row(run_id, result, "completed"))
    except Exception as e:
        duration = time.time() - t0
        tb = traceback.format_exc()
        try:
            with (run_dir / "events.jsonl").open("a") as f:
                f.write(json.dumps({
                    "ts": datetime.now(timezone.utc).isoformat(timespec="milliseconds"),
                    "event": "run_failed",
                    "error": str(e),
                    "traceback": tb,
                }) + "\n")
        except Exception:
            pass
        update_status(run_dir, "failed", duration_s=duration, error=str(e))
        _write_notes(run_dir, run_id, note)
        raise
    finally:
        sys.stdout = real_stdout
        stdout_f.close()

    assert result is not None
    return run_id, result