lars
2358a75ee1
CI / Sync project version with tag (pull_request) Has been skipped
CI / Format (ruff format) (pull_request) Successful in 43s
CI / Lint (ruff check) (pull_request) Successful in 54s
CI / Type check (ty) (pull_request) Successful in 53s
CI / Tests (pull_request) Successful in 2m33s
CI / Bump version, tag, and update changelog on merge to master (pull_request) Has been skipped
CI / Publish package to Gitea package registry (pull_request) Has been skipped
Closes the roadmap's long-standing "no eval-latency number exists for any configuration" gap. Instruments `giant rollout` to record per-physical-step wall-clock cost in its YAML sidecar, adds a measured Geant4/miniCaloSim per-step reference (giant/analysis/geant4_reference.py, from a 3-energy, 4-event-count-per-energy local benchmark), and wires both into a new eval_cost_per_step PlotSpec in the giant analyze gallery. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
423 lines
17 KiB
Python
423 lines
17 KiB
Python
"""Tests for the rollout-YAML(s) → run-directory flow, compute, and submit."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import pyarrow.parquet as pq
|
|
import pytest
|
|
import yaml
|
|
|
|
from giant.analysis import (
|
|
RunMeta,
|
|
SubmitConfig,
|
|
catalog_ids,
|
|
compute_one,
|
|
compute_reduced,
|
|
derive_run_dir,
|
|
load_rollout_yaml,
|
|
load_rollout_yamls,
|
|
merge_one,
|
|
prep,
|
|
write_submit,
|
|
)
|
|
from giant.analysis.catalog import get_spec
|
|
from giant.analysis.condor import Context
|
|
from giant.analysis.reduced import Partial, Reduced
|
|
from giant.constants import PREDICT_COORD_METADATA_KEY, ROLLOUT_COORD_VALUE
|
|
from tests.test_analysis_reduce import _reference_frame, _rollout_frame
|
|
|
|
|
|
def _write_rollout(path: Path) -> None:
|
|
tbl = _rollout_frame().collect().to_arrow()
|
|
tbl = tbl.replace_schema_metadata({PREDICT_COORD_METADATA_KEY: ROLLOUT_COORD_VALUE})
|
|
pq.write_table(tbl, path)
|
|
|
|
|
|
def _write_inputs(tmp_path: Path) -> Path:
|
|
"""Materialize rollout+reference parquet and a rollout YAML; return the YAML path."""
|
|
rollout = tmp_path / "rollout.parquet"
|
|
reference = tmp_path / "reference.parquet"
|
|
_write_rollout(rollout)
|
|
_reference_frame().collect().write_parquet(reference)
|
|
|
|
yaml_path = tmp_path / "run.yaml"
|
|
yaml_path.write_text(
|
|
yaml.safe_dump(
|
|
{
|
|
"prediction_id": "abcd1234ef",
|
|
"output": str(rollout),
|
|
"dataset": str(reference),
|
|
"checkpoint": "/ckpt/best.pt",
|
|
"kind": "rollout",
|
|
"energy_cutoff": 0.1,
|
|
"steps": 10,
|
|
}
|
|
)
|
|
)
|
|
return yaml_path
|
|
|
|
|
|
def _write_two_inputs(tmp_path: Path) -> tuple[Path, Path]:
|
|
"""Two rollout YAMLs (distinct output files) sharing one reference file."""
|
|
reference = tmp_path / "reference.parquet"
|
|
_reference_frame().collect().write_parquet(reference)
|
|
|
|
paths = []
|
|
for tag, pred_id in (("a", "aaaa1111ef"), ("b", "bbbb2222ef")):
|
|
rollout = tmp_path / f"rollout_{tag}.parquet"
|
|
_write_rollout(rollout)
|
|
yaml_path = tmp_path / f"run_{tag}.yaml"
|
|
yaml_path.write_text(
|
|
yaml.safe_dump(
|
|
{
|
|
"prediction_id": pred_id,
|
|
"output": str(rollout),
|
|
"dataset": str(reference),
|
|
"checkpoint": f"/ckpt/{tag}.pt",
|
|
"kind": "rollout",
|
|
"energy_cutoff": 0.1,
|
|
"steps": 10,
|
|
}
|
|
)
|
|
)
|
|
paths.append(yaml_path)
|
|
return paths[0], paths[1]
|
|
|
|
|
|
def _fake_venv(repo_dir: Path) -> None:
|
|
"""Stand in for a `uv sync`'d venv: write_submit checks `.venv/bin/giant` exists."""
|
|
giant = repo_dir / ".venv" / "bin" / "giant"
|
|
giant.parent.mkdir(parents=True, exist_ok=True)
|
|
giant.write_text("#!/bin/bash\n")
|
|
giant.chmod(0o755)
|
|
|
|
|
|
def _prep(rollout_yamls, run_dir: str | Path | None = None, chunks: int = 1, labels=None) -> Path:
|
|
"""``prep`` with small test-sized context bins/sampling."""
|
|
return prep(
|
|
rollout_yamls,
|
|
run_dir,
|
|
n_chunks=chunks,
|
|
labels=labels,
|
|
n_energy_bins=2,
|
|
n_marginal_bins=8,
|
|
top_k_pdg=3,
|
|
sample_rows=1000,
|
|
)
|
|
|
|
|
|
def test_load_rollout_yaml_requires_paths(tmp_path: Path):
|
|
bad = tmp_path / "bad.yaml"
|
|
bad.write_text(yaml.safe_dump({"output": "x.parquet"})) # no dataset
|
|
with pytest.raises(ValueError):
|
|
load_rollout_yaml(bad)
|
|
|
|
|
|
def test_load_rollout_yamls_single_defaults_to_rollout_name(tmp_path: Path):
|
|
yaml_path = _write_inputs(tmp_path)
|
|
loaded, reference = load_rollout_yamls([yaml_path])
|
|
assert [lr.name for lr in loaded] == ["rollout"]
|
|
assert reference.endswith("reference.parquet")
|
|
|
|
|
|
def test_load_rollout_yamls_multi_defaults_to_stem(tmp_path: Path):
|
|
a, b = _write_two_inputs(tmp_path)
|
|
loaded, _ = load_rollout_yamls([a, b])
|
|
assert [lr.name for lr in loaded] == ["run_a", "run_b"]
|
|
|
|
|
|
def test_load_rollout_yamls_explicit_labels(tmp_path: Path):
|
|
a, b = _write_two_inputs(tmp_path)
|
|
loaded, _ = load_rollout_yamls([a, b], labels=["flow", "wgan"])
|
|
assert [lr.name for lr in loaded] == ["flow", "wgan"]
|
|
|
|
|
|
def test_load_rollout_yamls_label_count_mismatch(tmp_path: Path):
|
|
a, b = _write_two_inputs(tmp_path)
|
|
with pytest.raises(ValueError, match="--label"):
|
|
load_rollout_yamls([a, b], labels=["only-one"])
|
|
|
|
|
|
def test_load_rollout_yamls_rejects_duplicate_names(tmp_path: Path):
|
|
a, b = _write_two_inputs(tmp_path)
|
|
with pytest.raises(ValueError, match="collide"):
|
|
load_rollout_yamls([a, b], labels=["same", "same"])
|
|
|
|
|
|
def test_load_rollout_yamls_rejects_mismatched_reference(tmp_path: Path):
|
|
a, _ = _write_two_inputs(tmp_path)
|
|
other_ref = tmp_path / "other_reference.parquet"
|
|
_reference_frame().collect().write_parquet(other_ref)
|
|
c = tmp_path / "run_c.yaml"
|
|
c.write_text(
|
|
yaml.safe_dump(
|
|
{"prediction_id": "cccc3333ef", "output": str(tmp_path / "rollout_c.parquet"), "dataset": str(other_ref)}
|
|
)
|
|
)
|
|
_write_rollout(tmp_path / "rollout_c.parquet")
|
|
with pytest.raises(ValueError, match="same reference"):
|
|
load_rollout_yamls([a, c])
|
|
|
|
|
|
def test_derive_run_dir_next_to_rollout():
|
|
y = {"output": "/data/roll.parquet", "prediction_id": "abcd1234ef", "dataset": "d"}
|
|
assert derive_run_dir([y]) == Path("/data/analysis_abcd1234")
|
|
assert derive_run_dir([y], "/somewhere") == Path("/somewhere")
|
|
|
|
|
|
def test_derive_run_dir_default_base():
|
|
y = {"output": "/data/roll.parquet", "prediction_id": "abcd1234ef", "dataset": "d"}
|
|
assert derive_run_dir([y], default_base="/work/lbogner/giant2/analysis_runs") == Path(
|
|
"/work/lbogner/giant2/analysis_runs/analysis_abcd1234"
|
|
)
|
|
# an explicit run_dir still wins over default_base
|
|
assert derive_run_dir([y], "/somewhere", default_base="/other") == Path("/somewhere")
|
|
|
|
|
|
def test_derive_run_dir_multi_rollout_joins_tags():
|
|
ys = [{"output": f"/data/roll_{i}.parquet", "prediction_id": f"tag{i}xxxx", "dataset": "d"} for i in range(2)]
|
|
assert derive_run_dir(ys, default_base="/base") == Path("/base/analysis_tag0xxxx-tag1xxxx")
|
|
|
|
|
|
def test_derive_run_dir_many_rollouts_truncates_with_plus_count():
|
|
ys = [{"output": f"/data/roll_{i}.parquet", "prediction_id": f"tag{i}xxxx", "dataset": "d"} for i in range(5)]
|
|
run_dir = derive_run_dir(ys, default_base="/base")
|
|
assert run_dir == Path("/base/analysis_tag0xxxx-tag1xxxx-tag2xxxx-plus2")
|
|
|
|
|
|
def test_prep_lays_out_run_dir(tmp_path: Path):
|
|
yaml_path = _write_inputs(tmp_path)
|
|
run_dir = _prep([yaml_path])
|
|
assert run_dir == tmp_path / "analysis_abcd1234"
|
|
assert (run_dir / "shared.json").exists()
|
|
ctx = Context.load(run_dir / "shared.json")
|
|
assert set(ctx.var_ranges) == {"step_length", "edep", "delta_e", "post_E"}
|
|
meta = RunMeta.load(run_dir / "run_meta.json")
|
|
assert meta.reference.endswith("reference.parquet")
|
|
assert [ro["name"] for ro in meta.rollouts] == ["rollout"]
|
|
assert meta.rollouts[0]["plot_meta"]["checkpoint"] == "/ckpt/best.pt"
|
|
assert "best.pt" in meta.title
|
|
assert meta.n_chunks == 1
|
|
assert meta.rows_per_chunk == [meta.total_rows] # single chunk holds everything
|
|
assert meta.total_rows == 8 # 5 rollout rows + 3 reference rows
|
|
|
|
|
|
def test_prep_multi_rollout_lays_out_run_dir(tmp_path: Path):
|
|
a, b = _write_two_inputs(tmp_path)
|
|
run_dir = _prep([a, b], labels=["flow", "wgan"])
|
|
meta = RunMeta.load(run_dir / "run_meta.json")
|
|
assert [ro["name"] for ro in meta.rollouts] == ["flow", "wgan"]
|
|
assert meta.rollouts[0]["plot_meta"]["checkpoint"] == "/ckpt/a.pt"
|
|
assert meta.rollouts[1]["plot_meta"]["checkpoint"] == "/ckpt/b.pt"
|
|
# 5 rows from each rollout + 3 from the shared reference
|
|
assert meta.total_rows == 13
|
|
|
|
|
|
def test_prep_splits_rows_per_chunk(tmp_path: Path):
|
|
run_dir = _prep([_write_inputs(tmp_path)], chunks=2)
|
|
meta = RunMeta.load(run_dir / "run_meta.json")
|
|
assert len(meta.rows_per_chunk) == 2
|
|
assert sum(meta.rows_per_chunk) == meta.total_rows == 8
|
|
|
|
|
|
def test_reprep_clears_stale_partials_from_a_different_chunk_count(tmp_path: Path):
|
|
"""Re-prepping with a different n_chunks must not leave old chunk
|
|
partials on disk for merge_one to silently merge against the new
|
|
context (they'd be keyed/sized for the old n_chunks)."""
|
|
yaml_path = _write_inputs(tmp_path)
|
|
run_dir = _prep([yaml_path], chunks=2)
|
|
compute_one("marginal_edep", run_dir, chunk_index=0)
|
|
compute_one("marginal_edep", run_dir, chunk_index=1)
|
|
stale = run_dir / "reduced_partial" / "marginal_edep__0.json"
|
|
assert stale.exists()
|
|
(run_dir / "reduced").mkdir(exist_ok=True)
|
|
(run_dir / "reduced" / "marginal_edep.json").write_text("{}")
|
|
|
|
_prep([yaml_path], run_dir, chunks=1)
|
|
|
|
assert not stale.exists()
|
|
assert not (run_dir / "reduced" / "marginal_edep.json").exists()
|
|
assert (run_dir / "shared.json").exists() # prep's own fresh output untouched
|
|
|
|
|
|
def test_compute_one_from_run_dir(tmp_path: Path):
|
|
run_dir = _prep([_write_inputs(tmp_path)])
|
|
out = compute_one("marginal_edep", run_dir)
|
|
assert out == run_dir / "reduced_partial" / "marginal_edep__0.json"
|
|
partial = Partial.load(out)
|
|
assert partial.id == "marginal_edep" and partial.chunk == 0
|
|
assert "r" in partial.data and "t" in partial.data
|
|
assert list(partial.data["r"]) == ["rollout"]
|
|
|
|
|
|
def test_timing_survives_plot_meta_to_compute_one(tmp_path: Path):
|
|
yaml_path = _write_inputs(tmp_path)
|
|
d = yaml.safe_load(yaml_path.read_text())
|
|
d["timing"] = {"us_per_step": 7.0, "write_us_per_step": 1.0}
|
|
yaml_path.write_text(yaml.safe_dump(d))
|
|
|
|
run_dir = _prep([yaml_path])
|
|
meta = RunMeta.load(run_dir / "run_meta.json")
|
|
assert meta.rollouts[0]["plot_meta"]["timing"] == {"us_per_step": 7.0, "write_us_per_step": 1.0}
|
|
|
|
out = compute_one("eval_cost_per_step", run_dir)
|
|
reduced = Reduced(**Partial.load(out).data["reduced"])
|
|
assert reduced.kind == "bar"
|
|
assert reduced.payload["series"]["rollout"] == [7.0, 1.0, 8.0]
|
|
|
|
|
|
def test_compute_reduced_explicit_paths(tmp_path: Path):
|
|
run_dir = _prep([_write_inputs(tmp_path)])
|
|
meta = RunMeta.load(run_dir / "run_meta.json")
|
|
rollouts = [{"name": ro["name"], "path": ro["path"]} for ro in meta.rollouts]
|
|
out = compute_reduced(
|
|
"marginal_step_length",
|
|
rollouts,
|
|
meta.reference,
|
|
run_dir / "shared.json",
|
|
tmp_path / "r.json",
|
|
)
|
|
assert Partial.load(out).id == "marginal_step_length"
|
|
|
|
|
|
def test_merge_one_produces_reduced(tmp_path: Path):
|
|
run_dir = _prep([_write_inputs(tmp_path)])
|
|
compute_one("marginal_edep", run_dir)
|
|
out = merge_one("marginal_edep", run_dir)
|
|
assert out == run_dir / "reduced" / "marginal_edep.json"
|
|
reduced = Reduced.load(out)
|
|
assert reduced.id == "marginal_edep"
|
|
assert len(reduced.payload["series"]["rollout"]) == len(reduced.payload["edges"]) - 1
|
|
|
|
|
|
def test_merge_one_fails_loudly_on_missing_chunk(tmp_path: Path):
|
|
run_dir = _prep([_write_inputs(tmp_path)], chunks=2)
|
|
compute_one("marginal_edep", run_dir, chunk_index=0) # chunk 1 never computed
|
|
with pytest.raises(FileNotFoundError, match="missing chunk"):
|
|
merge_one("marginal_edep", run_dir)
|
|
|
|
|
|
def test_chunked_compute_and_merge_matches_unchunked(tmp_path: Path):
|
|
(tmp_path / "a").mkdir()
|
|
(tmp_path / "b").mkdir()
|
|
unchunked_dir = _prep([_write_inputs(tmp_path / "a")])
|
|
compute_one("marginal_step_length", unchunked_dir)
|
|
unchunked = Reduced.load(merge_one("marginal_step_length", unchunked_dir))
|
|
|
|
chunked_dir = _prep([_write_inputs(tmp_path / "b")], chunks=2)
|
|
for k in range(2):
|
|
compute_one("marginal_step_length", chunked_dir, chunk_index=k)
|
|
chunked = Reduced.load(merge_one("marginal_step_length", chunked_dir))
|
|
|
|
assert chunked.payload == unchunked.payload
|
|
|
|
|
|
def test_two_rollout_compute_and_merge_produces_both_series(tmp_path: Path):
|
|
a, b = _write_two_inputs(tmp_path)
|
|
run_dir = _prep([a, b], labels=["flow", "wgan"])
|
|
compute_one("marginal_edep", run_dir)
|
|
reduced = Reduced.load(merge_one("marginal_edep", run_dir))
|
|
assert list(reduced.payload["series"]) == ["flow", "wgan"]
|
|
assert "reference" in reduced.payload
|
|
|
|
|
|
def test_compute_reduced_rejects_out_of_range_chunk(tmp_path: Path):
|
|
run_dir = _prep([_write_inputs(tmp_path)]) # n_chunks=1 (default)
|
|
with pytest.raises(ValueError, match="out of range"):
|
|
compute_one("marginal_edep", run_dir, chunk_index=1)
|
|
|
|
|
|
def test_write_submit_description(tmp_path: Path):
|
|
run_dir = _prep([_write_inputs(tmp_path)])
|
|
_fake_venv(tmp_path)
|
|
cfg = SubmitConfig(run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path)
|
|
txt = write_submit(cfg).read_text()
|
|
assert "universe = docker" in txt
|
|
assert "docker_image = cverstege/alma9-gridjob" in txt
|
|
assert "requirements = TARGET.ProvidesETPResources" in txt
|
|
assert "accounting_group = cms" in txt
|
|
assert "+RequestWalltime = $(walltime)" in txt
|
|
assert "queue plotid,chunk,walltime from" in txt
|
|
jobs = [line.split(",") for line in (run_dir / "jobs.txt").read_text().split()]
|
|
assert [i for i, _, _ in jobs] == catalog_ids()
|
|
assert all(k == "0" for _, k, _ in jobs) # n_chunks=1 default
|
|
assert all(int(w) > 0 for _, _, w in jobs)
|
|
wrapper = run_dir / "run_compute.sh"
|
|
assert wrapper.exists() and (wrapper.stat().st_mode & 0o111)
|
|
body = wrapper.read_text()
|
|
assert "giant analyze compute-one --id" in body
|
|
assert "--chunk" in body and "--run-dir" in body
|
|
|
|
|
|
def test_write_submit_requires_synced_venv(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
|
|
run_dir = _prep([_write_inputs(tmp_path)])
|
|
cfg = SubmitConfig(run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path)
|
|
# No `giant` next to the (fake) active interpreter, so this falls through
|
|
# to repo_dir/.venv/bin/giant, which _write_inputs/_prep also didn't create.
|
|
monkeypatch.setattr(sys, "executable", str(tmp_path / "not-a-venv" / "bin" / "python"))
|
|
with pytest.raises(FileNotFoundError, match="uv sync"):
|
|
write_submit(cfg)
|
|
|
|
|
|
def test_write_submit_remote_flag(tmp_path: Path):
|
|
run_dir = _prep([_write_inputs(tmp_path)])
|
|
_fake_venv(tmp_path)
|
|
cfg = SubmitConfig(run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path, remote=True)
|
|
txt = write_submit(cfg).read_text()
|
|
assert "+RemoteJob = True" in txt
|
|
assert "ProvidesETPResources" not in txt
|
|
|
|
|
|
def test_write_submit_chunks_respect_chunkable(tmp_path: Path):
|
|
assert get_spec("router_gating").chunkable is False
|
|
run_dir = _prep([_write_inputs(tmp_path)], chunks=4)
|
|
_fake_venv(tmp_path)
|
|
cfg = SubmitConfig(run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path, n_chunks=4)
|
|
write_submit(cfg)
|
|
jobs = [line.split(",") for line in (run_dir / "jobs.txt").read_text().split()]
|
|
counts: dict[str, int] = {}
|
|
for spec_id, _, _ in jobs:
|
|
counts[spec_id] = counts.get(spec_id, 0) + 1
|
|
assert counts["marginal_edep"] == 4
|
|
assert counts["router_gating"] == 1 # chunkable=False, ignores n_chunks
|
|
|
|
|
|
def test_write_submit_rejects_n_chunks_mismatch_with_run_meta(tmp_path: Path):
|
|
"""cfg.n_chunks must match the n_chunks the run_dir was actually prepped
|
|
with — RunMeta.rows_per_chunk is sized to the prepped value, so a
|
|
mismatch would otherwise surface as a confusing IndexError deep inside
|
|
_job_walltimes instead of a clear error here."""
|
|
run_dir = _prep([_write_inputs(tmp_path)], chunks=2)
|
|
_fake_venv(tmp_path)
|
|
cfg = SubmitConfig(run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path, n_chunks=4)
|
|
with pytest.raises(ValueError, match="n_chunks"):
|
|
write_submit(cfg)
|
|
|
|
|
|
def test_estimate_runtime_s_scales_with_rows_and_margin():
|
|
from giant.analysis import RUNTIME_SAFETY_MARGIN, estimate_runtime_s
|
|
from giant.analysis.runtime_estimate import _FIXED_OVERHEAD_S
|
|
|
|
assert RUNTIME_SAFETY_MARGIN > 0
|
|
small = estimate_runtime_s("marginal_edep", 1_000)
|
|
large = estimate_runtime_s("marginal_edep", 100_000_000)
|
|
assert small >= (1 + RUNTIME_SAFETY_MARGIN) * _FIXED_OVERHEAD_S
|
|
assert large > small # bigger chunk -> longer estimate
|
|
|
|
|
|
def test_write_submit_walltime_grows_with_chunk_rows(tmp_path: Path):
|
|
"""A chunked run's later job walltimes track that chunk's row count."""
|
|
from giant.analysis.runtime_estimate import estimate_runtime_s
|
|
|
|
run_dir = _prep([_write_inputs(tmp_path)], chunks=2)
|
|
meta = RunMeta.load(run_dir / "run_meta.json")
|
|
_fake_venv(tmp_path)
|
|
cfg = SubmitConfig(run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path, n_chunks=2)
|
|
write_submit(cfg)
|
|
jobs = {(i, int(k)): int(w) for i, k, w in (line.split(",") for line in (run_dir / "jobs.txt").read_text().split())}
|
|
for chunk in range(2):
|
|
expected = estimate_runtime_s("marginal_edep", meta.rows_per_chunk[chunk])
|
|
assert jobs[("marginal_edep", chunk)] == expected
|