lars
ce54a44dac
CI / Sync project version with tag (pull_request) Has been skipped
CI / Format (ruff format) (pull_request) Successful in 1m21s
CI / Type check (ty) (pull_request) Successful in 1m21s
CI / Lint (ruff check) (pull_request) Successful in 1m26s
CI / Tests (pull_request) Successful in 2m41s
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
# Conflicts: # .bumpversion.toml # giant/cli.py # pyproject.toml # uv.lock
328 lines
12 KiB
Python
328 lines
12 KiB
Python
"""Tests for the rollout-YAML(s) → run-directory flow, compute, and merge."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
|
|
import pyarrow.parquet as pq
|
|
import pytest
|
|
import yaml
|
|
|
|
from giant.analysis import (
|
|
RunMeta,
|
|
compute_one,
|
|
compute_reduced,
|
|
derive_run_dir,
|
|
load_rollout_yaml,
|
|
load_rollout_yamls,
|
|
merge_one,
|
|
prep,
|
|
)
|
|
from giant.analysis.run 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 _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_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
|