lars
c3e5956718
write_submit baked in cfg.repo_dir/.venv/bin/giant unconditionally, which breaks when submitting from a differently-named or non-default venv (e.g. --extra cuda). Prefer the giant executable next to sys.executable (the venv actually running the submit), falling back to repo_dir/.venv/bin/giant.
479 lines
18 KiB
Python
479 lines
18 KiB
Python
"""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_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 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 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",
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 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,
|
|
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 rollout parquet — the original convention — for callers that don't
|
|
care where the run directory lives.
|
|
"""
|
|
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]
|
|
base = Path(default_base) if default_base is not None else rollout.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``)."""
|
|
|
|
rollout: str
|
|
reference: str
|
|
run_dir: str
|
|
title: str
|
|
plot_meta: dict
|
|
n_chunks: int = 1
|
|
# 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(
|
|
rollout: str | Path, reference: str | Path, n_chunks: int
|
|
) -> list[int]:
|
|
"""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
|
|
for lf in (open_side(rollout, Side.rollout), open_side(reference, Side.reference)):
|
|
df = counts(lf)
|
|
for c, n in zip(df["_c"].to_list(), df["n"].to_list()):
|
|
out[c] += n
|
|
return out
|
|
|
|
|
|
def prep(
|
|
rollout_yaml: str | Path,
|
|
run_dir: str | Path | None = None,
|
|
n_chunks: int = 1,
|
|
default_base: 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.
|
|
``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.
|
|
|
|
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``.
|
|
"""
|
|
y = load_rollout_yaml(rollout_yaml)
|
|
run_path = derive_run_dir(y, 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, reference = y["output"], y["dataset"]
|
|
ctx = build_context(rollout, reference, **ctx_kwargs)
|
|
ctx.save(run_path / "shared.json")
|
|
|
|
rows_per_chunk = _rows_per_chunk(rollout, reference, n_chunks)
|
|
|
|
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),
|
|
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,
|
|
rollout: str | Path,
|
|
reference: str | Path,
|
|
shared: str | Path,
|
|
out: str | Path,
|
|
checkpoint: str | None = None,
|
|
chunk_index: int = 0,
|
|
n_chunks: int = 1,
|
|
) -> Path:
|
|
"""Core: run one (plot, chunk)'s partial reduction against explicit paths.
|
|
|
|
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 "
|
|
f"n_chunks={effective_n} (chunkable={spec.chunkable})"
|
|
)
|
|
bundle = Bundle.open(
|
|
rollout, reference, ctx, checkpoint=checkpoint, 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")
|
|
return compute_reduced(
|
|
spec_id,
|
|
meta.rollout,
|
|
meta.reference,
|
|
run_path / "shared.json",
|
|
run_path / "reduced_partial" / f"{spec_id}__{chunk_index}.json",
|
|
checkpoint=meta.plot_meta.get("checkpoint"),
|
|
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
|