giant/giant/analysis/condor.py

"""HTCondor orchestration driven by one or more ``giant rollout`` YAML sidecars.

A rollout writes a YAML sidecar (``giant/cli.py:_write_prediction_ref`` +
rollout extras) that already names both files we need and carries the run's
provenance:

* ``output``  — the rollout steps parquet (the *generated* side)
* ``dataset`` — the file the rollout was seeded from, i.e. the held-out real
  steps (the *reference* side)
* ``checkpoint``, ``geometry_oracle``, ``energy_cutoff``, ``steps``, ... —
  metadata that flows straight into every plot's gallery ``metadata.yaml``.

The analysis takes N such YAMLs — one series per rollout, all required to
share the same ``dataset`` (the premise is "N candidates vs one ground
truth") — resolves each one's series name (``load_rollout_yamls``), derives
its own **run directory** next to the first rollout's parquet, and lays
everything out under it:

    <run_dir>/shared.json     fixed bin edges / group sets (prep)
    <run_dir>/run_meta.json   resolved rollout/reference paths + plot metadata
    <run_dir>/reduced_partial/<id>__<chunk>.json   one per (plot, chunk) job
    <run_dir>/reduced/<id>.json   merged, per plot
    <run_dir>/plots/<family>/<id>.pdf   rendered locally

Job model (one condor job per (plot, chunk), compute/merge/render split):

1. ``prep`` runs once on the submit node — reads the YAML, resolves the shared
   context from a subsample, writes ``shared.json`` + ``run_meta.json``
   (including the run's configured ``n_chunks``).
2. one job per catalog id x chunk index runs ``giant analyze compute-one
   --run-dir`` on a worker — a single streaming pass over that
   ``event_id``-disjoint chunk, writing ``reduced_partial/<id>__<chunk>.json``
   (polars/numpy only, no LaTeX). Specs marked ``chunkable=False``
   (``PlotSpec``, ``catalog.py``) always run as a single chunk.
3. a *local* ``giant analyze render`` first merges every plot's chunk partials
   (``merge_all`` — sums/concatenates them and re-derives any data-dependent
   histogram edges or mean/std, per ``PlotSpec.finalize``) into
   ``reduced/<id>.json``, then renders those into the styled PDF + gallery tree
   (that step imports plotstyle/LaTeX).

Files on ``/ceph`` or ``/work`` are reached via ``ProvidesETPResources``; no
HTCondor file transfer of the multi-GB inputs.
"""

from __future__ import annotations

import json
import shutil
import sys
from collections.abc import Sequence
from dataclasses import dataclass, field
from pathlib import Path

import polars as pl
import yaml

from giant.analysis.catalog import Bundle, catalog_ids, get_spec
from giant.analysis.context import Context, build_context
from giant.analysis.reduced import Partial
from giant.analysis.runtime_estimate import estimate_runtime_s
from giant.analysis.sources import RolloutSpec, Side, open_side

# Keys copied verbatim from a rollout YAML into each plot's gallery metadata.
_PLOT_META_KEYS = (
    "prediction_id",
    "checkpoint",
    "output",
    "dataset",
    "geometry_oracle",
    "energy_cutoff",
    "max_steps",
    "steps",
    "max_tracks_per_event",
    "escape_threshold",
    "n_events",
    "n_seed_events",
    "timestamp",
    "comment",
    "weights",
    "batch_size",
    "device",
    "rollout_seed",
    "n_rows",
    "termination_reason_counts",
    "model_config",
    "training_epoch",
    "best_val_loss",
    "training_config",
    "training_meta",
    # Diagnostic — only present when giant rollout ran under
    # stage2_model.particle_type.target="embedding" (see giant/cli.py's
    # rollout command and giant.rollout.L1DistCollector); absent otherwise,
    # which the type_embedding_l1_distance PlotSpec (catalog.py) reads as
    # "not applicable to this checkpoint".
    "type_embedding_l1_dist",
)


# ---------------------------------------------------------------------------
# rollout-YAML → run directory
# ---------------------------------------------------------------------------


def load_rollout_yaml(path: str | Path) -> dict:
    """Load a ``giant rollout`` YAML sidecar, requiring the two file paths."""
    d = yaml.safe_load(Path(path).read_text())
    for key in ("output", "dataset"):
        if key not in d:
            raise ValueError(
                f"{path} is not a rollout YAML (missing {key!r}); expected the "
                "sidecar `giant rollout` writes next to the checkpoint"
            )
    if d.get("kind") not in (None, "rollout"):
        raise ValueError(f"{path} has kind={d.get('kind')!r}, not a rollout YAML")
    return d


@dataclass
class LoadedRollout:
    """One rollout YAML plus its resolved series ``name`` (see ``load_rollout_yamls``)."""

    name: str
    yaml: dict


def load_rollout_yamls(
    paths: Sequence[str | Path], labels: Sequence[str] | None = None
) -> tuple[list[LoadedRollout], str]:
    """Load every rollout YAML, resolve each one's series name, and verify they
    all share one reference (``dataset``) file — the premise is "N candidates
    vs one ground truth", not N independent comparisons.

    Names: an explicit ``labels[i]`` if given (``labels`` must be empty or
    exactly ``len(paths)`` long); otherwise the YAML's stem for N>1, or
    ``"rollout"`` for the single-YAML case — matching today's one-series
    legend/payload key, so a single-rollout run renders identically to
    before this feature existed. Raises ``ValueError`` if two rollouts
    resolve to the same name, or if the YAMLs don't all name the same
    ``dataset``.
    """
    if labels and len(labels) != len(paths):
        raise ValueError(f"--label given {len(labels)} time(s) but {len(paths)} rollout YAML(s) were passed")
    yamls = [load_rollout_yaml(p) for p in paths]
    if labels:
        names = list(labels)
    elif len(paths) == 1:
        names = ["rollout"]
    else:
        names = [Path(p).stem for p in paths]
    if len(set(names)) != len(names):
        dupes = sorted({n for n in names if names.count(n) > 1})
        raise ValueError(f"rollout series names collide: {dupes} — pass --label to disambiguate")

    references = {str(y["dataset"]) for y in yamls}
    if len(references) > 1:
        detail = "\n".join(f"  {p}: dataset={y['dataset']!r}" for p, y in zip(paths, yamls))
        raise ValueError(
            "all rollout YAMLs must be seeded from the same reference (dataset) "
            f"file — got {len(references)} distinct ones:\n{detail}"
        )

    return [LoadedRollout(name=n, yaml=y) for n, y in zip(names, yamls)], yamls[0]["dataset"]


def _run_tag(y: dict) -> str:
    rollout = Path(y["output"])
    return str(y.get("prediction_id") or rollout.stem)[:8]


def derive_run_dir(
    rollout_yamls: list[dict],
    run_dir: str | Path | None = None,
    default_base: str | Path | None = None,
) -> Path:
    """Analysis output directory.

    Precedence: an explicit ``run_dir`` always wins. Otherwise
    ``default_base / analysis_<tag>`` if ``default_base`` is given (the CLI
    passes the repo's gitignored ``analysis_runs/``, so run directories don't
    pile up on ``/ceph`` next to the rollout parquet). Falls back to next to
    the *first* rollout's parquet — the original convention — for callers
    that don't care where the run directory lives.

    ``tag`` is a single rollout's ``prediction_id``/output stem (matching
    today's single-rollout convention exactly) when there's only one; for
    N>1 it joins up to three tags with ``-``, then ``-plus<K>`` for any
    beyond that, so a many-rollout run still gets a short, stable directory
    name.
    """
    if run_dir is not None:
        return Path(run_dir)
    tags = [_run_tag(y) for y in rollout_yamls]
    if len(tags) == 1:
        tag = tags[0]
    else:
        shown, rest = tags[:3], tags[3:]
        tag = "-".join(shown) + (f"-plus{len(rest)}" if rest else "")
    base = Path(default_base) if default_base is not None else Path(rollout_yamls[0]["output"]).parent
    return base / f"analysis_{tag}"


def _plot_meta(rollout_yaml: dict) -> dict:
    return {k: rollout_yaml[k] for k in _PLOT_META_KEYS if k in rollout_yaml}


@dataclass
class RunMeta:
    """Resolved paths + plot metadata for one analysis run (``run_meta.json``).

    ``rollouts`` is ``[{"name", "path", "plot_meta"}, ...]``, insertion order
    = the order rollouts were given on the CLI (and so the order every
    ``Reduced.payload["series"]`` dict is built in — see ``catalog.py``).
    """

    rollouts: list[dict]
    reference: str
    run_dir: str
    title: str
    n_chunks: int = 1
    # combined rollout+reference row count of each event_id-disjoint chunk,
    # and the dataset total — inputs to `runtime_estimate.estimate_runtime_s`.
    # Empty/0 on run directories written before this field existed.
    rows_per_chunk: list[int] = field(default_factory=list)
    total_rows: int = 0

    def save(self, path: str | Path) -> None:
        Path(path).write_text(json.dumps(self.__dict__, indent=2))

    @classmethod
    def load(cls, path: str | Path) -> "RunMeta":
        return cls(**json.loads(Path(path).read_text()))


def _rows_per_chunk(rollouts: list[str | Path], reference: str | Path, n_chunks: int) -> list[int]:
    """Combined rollout+reference row count of each ``event_id % n_chunks`` chunk.

    One cheap streaming ``group_by`` per side (just the ``event_id`` column) —
    the sizing input every job's estimated walltime
    (``runtime_estimate.estimate_runtime_s``) is computed from.
    """

    def counts(lf: pl.LazyFrame) -> pl.DataFrame:
        return (
            lf.select((pl.col("event_id") % n_chunks).alias("_c"))
            .group_by("_c")
            .agg(pl.len().alias("n"))
            .collect(engine="streaming")
        )

    out = [0] * n_chunks
    sides = [open_side(reference, Side.reference)] + [open_side(r, Side.rollout) for r in rollouts]
    for lf in sides:
        df = counts(lf)
        for c, n in zip(df["_c"].to_list(), df["n"].to_list()):
            out[c] += n
    return out


def prep(
    rollout_yamls: Sequence[str | Path],
    run_dir: str | Path | None = None,
    n_chunks: int = 1,
    default_base: str | Path | None = None,
    labels: Sequence[str] | None = None,
    **ctx_kwargs,
) -> Path:
    """Read the rollout YAML(s), build the shared context, and lay out the run dir.

    Writes ``shared.json`` + ``run_meta.json`` and returns the run directory.
    ``n_chunks`` is the run-level chunk count every ``compute-one``/``merge-one``
    job reads back out of ``run_meta.json`` (via ``RunMeta.n_chunks``), so it is
    resolved once here rather than re-passed (and risking disagreement) at every
    later step. See ``derive_run_dir`` for how ``run_dir``/``default_base``
    resolve the actual directory, and ``load_rollout_yamls`` for how
    ``labels``/YAML stems resolve each rollout's series name.

    Clears any existing ``reduced_partial/``/``reduced/`` from a prior prep of
    this same ``run_dir``: partial files carry no record of what context
    (``n_chunks``, bin edges, group sets) they were computed under, so
    re-prepping with a different ``n_chunks``/``**ctx_kwargs`` (or after the
    rollout/reference files changed) would otherwise let ``merge_one`` silently
    merge stale partials against the new ``shared.json``.
    """
    loaded, reference = load_rollout_yamls(list(rollout_yamls), labels)
    run_path = derive_run_dir([lr.yaml for lr in loaded], run_dir, default_base=default_base)
    run_path.mkdir(parents=True, exist_ok=True)

    for stale in ("reduced_partial", "reduced"):
        stale_dir = run_path / stale
        if stale_dir.exists():
            shutil.rmtree(stale_dir)

    rollout_specs = [RolloutSpec(name=lr.name, source=lr.yaml["output"]) for lr in loaded]
    ctx = build_context(rollout_specs, reference, **ctx_kwargs)
    ctx.save(run_path / "shared.json")

    rows_per_chunk = _rows_per_chunk([lr.yaml["output"] for lr in loaded], reference, n_chunks)

    rollouts_meta = [
        {"name": lr.name, "path": str(lr.yaml["output"]), "plot_meta": _plot_meta(lr.yaml)} for lr in loaded
    ]
    ckpts = ", ".join(Path(lr.yaml.get("checkpoint", "")).name or "rollout" for lr in loaded)

    RunMeta(
        rollouts=rollouts_meta,
        reference=str(reference),
        run_dir=str(run_path),
        title=f"GIANT rollout analysis — {ckpts}",
        n_chunks=n_chunks,
        rows_per_chunk=rows_per_chunk,
        total_rows=sum(rows_per_chunk),
    ).save(run_path / "run_meta.json")
    return run_path


# ---------------------------------------------------------------------------
# per-(plot, chunk) compute (what each condor job runs)
# ---------------------------------------------------------------------------


def compute_reduced(
    spec_id: str,
    rollouts: list[dict],
    reference: str | Path,
    shared: str | Path,
    out: str | Path,
    chunk_index: int = 0,
    n_chunks: int = 1,
) -> Path:
    """Core: run one (plot, chunk)'s partial reduction against explicit paths.

    ``rollouts``: ``[{"name", "path", "checkpoint"?, "type_embedding_l1_dist"?},
    ...]``, one per rollout series (insertion order preserved through to every
    plot's ``Reduced.payload["series"]``).

    Writes a ``Partial`` JSON — the raw, not-yet-merged output of
    ``PlotSpec.compute_partial`` — never a finished ``Reduced``; ``merge_one``
    is what combines every chunk's ``Partial`` for a plot into the final
    ``Reduced``. Specs with ``chunkable=False`` always run as a single chunk
    regardless of ``n_chunks``.
    """
    ctx = Context.load(shared)
    spec = get_spec(spec_id)
    effective_n = n_chunks if spec.chunkable else 1
    if not (0 <= chunk_index < effective_n):
        raise ValueError(
            f"{spec_id}: chunk_index={chunk_index} out of range for n_chunks={effective_n} (chunkable={spec.chunkable})"
        )
    rollout_specs = [
        RolloutSpec(
            name=r["name"],
            source=r["path"],
            checkpoint=r.get("checkpoint"),
            type_embedding_l1_dist=r.get("type_embedding_l1_dist"),
        )
        for r in rollouts
    ]
    bundle = Bundle.open(rollout_specs, reference, ctx, chunk=(chunk_index, effective_n))
    partial = Partial(
        id=spec_id,
        family=spec.family,
        chunk=chunk_index,
        data=spec.compute_partial(bundle),
    )
    out = Path(out)
    partial.save(out)
    return out


def compute_one(spec_id: str, run_dir: str | Path, chunk_index: int = 0) -> Path:
    """Run one (plot, chunk)'s partial reduction from a prepped run directory."""
    run_path = Path(run_dir)
    meta = RunMeta.load(run_path / "run_meta.json")
    rollouts = [
        {
            "name": ro["name"],
            "path": ro["path"],
            "checkpoint": ro["plot_meta"].get("checkpoint"),
            "type_embedding_l1_dist": ro["plot_meta"].get("type_embedding_l1_dist"),
        }
        for ro in meta.rollouts
    ]
    return compute_reduced(
        spec_id,
        rollouts,
        meta.reference,
        run_path / "shared.json",
        run_path / "reduced_partial" / f"{spec_id}__{chunk_index}.json",
        chunk_index=chunk_index,
        n_chunks=meta.n_chunks,
    )


# ---------------------------------------------------------------------------
# per-plot merge (the join step ``render`` runs before rendering)
# ---------------------------------------------------------------------------


def merge_one(spec_id: str, run_dir: str | Path) -> Path:
    """Merge every chunk's partial for one plot into the final ``Reduced`` JSON.

    Fails loudly if fewer partials exist than the run's configured chunk count
    for this plot — that is what catches an incomplete/failed condor job
    instead of silently rendering a plot from partial data. Idempotent: safe
    to call again (e.g. from ``render_run``) once all chunks are in.
    """
    run_path = Path(run_dir)
    meta = RunMeta.load(run_path / "run_meta.json")
    ctx = Context.load(run_path / "shared.json")
    spec = get_spec(spec_id)
    effective_n = meta.n_chunks if spec.chunkable else 1

    partial_dir = run_path / "reduced_partial"
    found = {p.chunk: p for p in (Partial.load(jf) for jf in partial_dir.glob(f"{spec_id}__*.json"))}
    missing = sorted(set(range(effective_n)) - set(found))
    if missing:
        raise FileNotFoundError(
            f"{spec_id}: missing chunk partial(s) {missing} of {effective_n} "
            f"under {partial_dir} — did every compute-one job finish?"
        )

    parts = [found[k].data for k in range(effective_n)]
    reduced = spec.finalize(parts, ctx)
    out = run_path / "reduced" / f"{spec_id}.json"
    reduced.save(out)
    return out


def merge_all(run_dir: str | Path) -> list[Path]:
    """Merge every catalog plot's chunk partials into ``reduced/<id>.json``."""
    return [merge_one(spec_id, run_dir) for spec_id in catalog_ids()]


# ---------------------------------------------------------------------------
# submit description
# ---------------------------------------------------------------------------


@dataclass
class SubmitConfig:
    run_dir: Path
    accounting_group: str
    repo_dir: Path
    docker_image: str = "cverstege/alma9-gridjob"
    request_memory_mb: int = 8192
    request_cpus: int = 1
    remote: bool = False  # +RemoteJob (grid I/O) vs ProvidesETPResources (local files)
    n_chunks: int = 1  # per-plot data chunks; ignored for chunkable=False specs


_WRAPPER = """#!/bin/bash
set -euo pipefail
cd {repo_dir}
exec {giant_exe} analyze compute-one --id "$1" --chunk "$2" --run-dir {run_dir}
"""


def _submit_description(cfg: SubmitConfig, wrapper: Path, jobs_file: Path) -> str:
    reqs_attrs = "+RemoteJob = True\n" if cfg.remote else "requirements = TARGET.ProvidesETPResources\n"
    return (
        "universe = docker\n"
        f"docker_image = {cfg.docker_image}\n"
        f"executable = {wrapper}\n"
        "arguments = $(plotid) $(chunk)\n"
        "should_transfer_files = YES\n"
        "when_to_transfer_output = ON_EXIT\n"
        f"request_memory = {cfg.request_memory_mb}\n"
        f"request_cpus = {cfg.request_cpus}\n"
        "+RequestWalltime = $(walltime)\n"
        f"accounting_group = {cfg.accounting_group}\n"
        f"{reqs_attrs}"
        f"output = {cfg.run_dir}/logs/$(plotid)__$(chunk).out\n"
        f"error = {cfg.run_dir}/logs/$(plotid)__$(chunk).err\n"
        f"log = {cfg.run_dir}/logs/condor.log\n"
        f"queue plotid,chunk,walltime from {jobs_file}\n"
    )


def _job_walltimes(run_dir: Path, ids: list[str], n_chunks: int) -> list[tuple[str, int, int]]:
    """``(spec_id, chunk, walltime_s)`` for every job, sized from ``run_meta.json``.

    Row counts come from ``prep``'s ``RunMeta.rows_per_chunk``/``total_rows``;
    ``chunkable=False`` specs (router diagnostics) always use the dataset
    total since they run as a single job regardless of ``n_chunks``.
    """
    meta = RunMeta.load(run_dir / "run_meta.json")
    jobs: list[tuple[str, int, int]] = []
    for spec_id in ids:
        chunkable = get_spec(spec_id).chunkable
        chunks = range(n_chunks) if chunkable else [0]
        for chunk in chunks:
            n_rows = meta.rows_per_chunk[chunk] if chunkable else meta.total_rows
            jobs.append((spec_id, chunk, estimate_runtime_s(spec_id, n_rows)))
    return jobs


def _resolve_giant_executable(repo_dir: Path) -> Path:
    """Path to the ``giant`` entry point to bake into the condor wrapper script.

    Prefers the venv currently running this process (``sys.executable``'s
    sibling ``giant``) so a submit from a non-default venv (e.g. ``--extra
    cuda`` on a dev box) doesn't silently pick up a different one; falls back
    to ``repo_dir/.venv/bin/giant`` for the case this is invoked from outside
    any venv (e.g. a system Python).
    """
    active = Path(sys.executable).parent / "giant"
    if active.exists():
        return active
    venv_giant = repo_dir / ".venv" / "bin" / "giant"
    if not venv_giant.exists():
        raise FileNotFoundError(
            f"no `giant` executable found next to {sys.executable} or at "
            f"{venv_giant} — condor jobs run it directly (no `uv` on the "
            f"worker image), so run `uv sync --extra cpu` in {repo_dir} "
            "before submitting."
        )
    return venv_giant


def write_submit(cfg: SubmitConfig, ids: list[str] | None = None) -> Path:
    """Write the wrapper script, (plot, chunk) job list, and HTCondor submit
    description.

    Each catalog id gets ``cfg.n_chunks`` jobs, except ``chunkable=False``
    specs (the router diagnostics), which always get exactly one regardless of
    ``cfg.n_chunks``. Every job's ``+RequestWalltime`` is estimated from its
    chunk's row count (``runtime_estimate.estimate_runtime_s``, requires
    ``run_meta.json`` from ``prep`` to already carry ``rows_per_chunk``).
    Returns the submit description path (``<run_dir>/analyze.sub``). Does not
    submit — call ``condor_submit`` on the returned file.

    ``cfg.n_chunks`` and the run directory's own ``RunMeta.n_chunks`` (fixed by
    ``prep``, and what ``RunMeta.rows_per_chunk`` was sized against) are two
    independent values — checked equal up front so a mismatch is a clear error
    here rather than an ``IndexError`` out of ``_job_walltimes``.
    """
    giant_exe = _resolve_giant_executable(cfg.repo_dir)

    ids = ids or catalog_ids()
    run_dir = cfg.run_dir
    meta = RunMeta.load(run_dir / "run_meta.json")
    if cfg.n_chunks != meta.n_chunks:
        raise ValueError(
            f"SubmitConfig.n_chunks={cfg.n_chunks} does not match the "
            f"n_chunks this run directory was prepped with "
            f"(RunMeta.n_chunks={meta.n_chunks} in {run_dir}/run_meta.json) — "
            "re-run `prep` with the desired n_chunks, or fix cfg.n_chunks to "
            "match it."
        )
    (run_dir / "logs").mkdir(parents=True, exist_ok=True)
    (run_dir / "reduced").mkdir(parents=True, exist_ok=True)
    (run_dir / "reduced_partial").mkdir(parents=True, exist_ok=True)

    wrapper = run_dir / "run_compute.sh"
    wrapper.write_text(_WRAPPER.format(repo_dir=cfg.repo_dir, giant_exe=giant_exe, run_dir=run_dir))
    wrapper.chmod(0o755)

    jobs = _job_walltimes(run_dir, ids, cfg.n_chunks)
    jobs_file = run_dir / "jobs.txt"
    jobs_file.write_text("\n".join(f"{i},{k},{w}" for i, k, w in jobs) + "\n")

    sub = run_dir / "analyze.sub"
    sub.write_text(_submit_description(cfg, wrapper, jobs_file))
    return sub