giant/giant/analysis/condor.py

"""HTCondor orchestration driven by a ``giant rollout`` YAML sidecar.

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``.

So the analysis takes that one YAML as input, derives its own **run directory**
next to the rollout 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/<id>.json   one per compute job
    <run_dir>/plots/<family>/<id>.pdf   rendered locally

Job model (one condor job per plot, compute/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``.
2. one job per catalog id runs ``giant analyze compute-one --run-dir`` on a
   worker — a single streaming pass writing ``reduced/<id>.json`` (polars/numpy
   only, no LaTeX).
3. a final *local* ``giant analyze render`` turns 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
from dataclasses import dataclass
from pathlib import Path

import yaml

from giant.analysis.catalog import Bundle, catalog_ids, get_spec
from giant.analysis.context import Context, build_context

# 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",
    "n_seed_events",
    "timestamp",
    "comment",
    "model_config",
    "training_epoch",
    "best_val_loss",
)


# ---------------------------------------------------------------------------
# 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


def derive_run_dir(rollout_yaml: dict, run_dir: str | Path | None = None) -> Path:
    """Analysis output directory, next to the rollout parquet unless overridden."""
    if run_dir is not None:
        return Path(run_dir)
    rollout = Path(rollout_yaml["output"])
    tag = str(rollout_yaml.get("prediction_id") or rollout.stem)[:8]
    return rollout.parent / 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``)."""

    rollout: str
    reference: str
    run_dir: str
    title: str
    plot_meta: dict

    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 prep(
    rollout_yaml: str | Path,
    run_dir: str | Path | None = None,
    **ctx_kwargs,
) -> Path:
    """Read the rollout YAML, build the shared context, and lay out the run dir.

    Writes ``shared.json`` + ``run_meta.json`` and returns the run directory.
    """
    y = load_rollout_yaml(rollout_yaml)
    run_path = derive_run_dir(y, run_dir)
    run_path.mkdir(parents=True, exist_ok=True)

    rollout, reference = y["output"], y["dataset"]
    ctx = build_context(rollout, reference, **ctx_kwargs)
    ctx.save(run_path / "shared.json")

    ckpt = Path(y.get("checkpoint", "")).name or "rollout"
    RunMeta(
        rollout=str(rollout),
        reference=str(reference),
        run_dir=str(run_path),
        title=f"GIANT rollout analysis — {ckpt}",
        plot_meta=_plot_meta(y),
    ).save(run_path / "run_meta.json")
    return run_path


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


def compute_reduced(
    spec_id: str,
    rollout: str | Path,
    reference: str | Path,
    shared: str | Path,
    out: str | Path,
) -> Path:
    """Core: run one plot's reduction against explicit paths → ``Reduced`` JSON."""
    ctx = Context.load(shared)
    bundle = Bundle.open(rollout, reference, ctx)
    reduced = get_spec(spec_id).compute(bundle)
    out = Path(out)
    reduced.save(out)
    return out


def compute_one(spec_id: str, run_dir: str | Path) -> Path:
    """Run one plot's reduction from a prepped run directory."""
    run_path = Path(run_dir)
    meta = RunMeta.load(run_path / "run_meta.json")
    return compute_reduced(
        spec_id,
        meta.rollout,
        meta.reference,
        run_path / "shared.json",
        run_path / "reduced" / f"{spec_id}.json",
    )


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


@dataclass
class SubmitConfig:
    run_dir: Path
    accounting_group: str
    repo_dir: Path
    docker_image: str = "mschnepf/slc7-condocker"
    request_memory_mb: int = 4096
    request_cpus: int = 1
    request_walltime_s: int = 3600
    remote: bool = False  # +RemoteJob (grid I/O) vs ProvidesETPResources (local files)


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


def _submit_description(cfg: SubmitConfig, wrapper: Path, ids_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)\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"
        f"+RequestWalltime = {cfg.request_walltime_s}\n"
        f"accounting_group = {cfg.accounting_group}\n"
        f"{reqs_attrs}"
        f"output = {cfg.run_dir}/logs/$(plotid).out\n"
        f"error = {cfg.run_dir}/logs/$(plotid).err\n"
        f"log = {cfg.run_dir}/logs/condor.log\n"
        f"queue plotid from {ids_file}\n"
    )


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

    Returns the submit description path (``<run_dir>/analyze.sub``). Does not
    submit — call ``condor_submit`` on the returned file.
    """
    ids = ids or catalog_ids()
    run_dir = cfg.run_dir
    (run_dir / "logs").mkdir(parents=True, exist_ok=True)
    (run_dir / "reduced").mkdir(parents=True, exist_ok=True)

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

    ids_file = run_dir / "plotids.txt"
    ids_file.write_text("\n".join(ids) + "\n")

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