lars
ebd3e0dc71
CI / Lint (ruff check) (push) Successful in 32s
CI / Format (ruff format) (push) Successful in 30s
CI / Sync project version with tag (push) Has been skipped
CI / Type check (ty) (push) Successful in 35s
CI / Lint (ruff check) (pull_request) Successful in 33s
CI / Format (ruff format) (pull_request) Successful in 30s
CI / Type check (ty) (pull_request) Successful in 34s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Tests (push) Successful in 5m59s
CI / Bump version, tag, and update changelog on merge to master (push) Has been skipped
CI / Tests (pull_request) Successful in 4m22s
CI / Bump version, tag, and update changelog on merge to master (pull_request) Has been skipped
giant analyze compares N rollout YAMLs against one shared reference file
(all must name the same dataset, checked up front) instead of exactly one
rollout vs one reference, rendering each rollout as its own colored series
against a single reference line/panel. Series names come from a repeated
--label flag, else the YAML stem, else "rollout" for a single YAML — a
single-rollout run keeps rendering identically to before this change.
Bundle now holds a name-keyed dict of rollout sides instead of one fixed
pair, every catalog compute_partial/finalize builds a Reduced.payload
keyed the same way ("series": {name: ...}, "reference": ... as the one
distinguished non-rollout entry), and every renderer draws N series (or
N panels, for the two heatmap-shaped specs and the router/type-embedding
diagnostics, which are inherently one-matrix/one-checkpoint per rollout)
against the reference's fixed dashed-ink style.
407 lines
16 KiB
Python
407 lines
16 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_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
|