lars
d5853d5a75
Replace the five separately-hyphenated uv entry points (steps-to-parquet, steps-to-parquet-parallel, migrate-geant-steps, bump-dataset-version, create-root-files) plus the unregistered hparam_scan.py with one `dwarf` command exposing convert/migrate/bump-gen/bump-schema/status/ update-manifest/create-manifest/make-root/hparam-scan as subcommands. Each scripts/*.py module now only holds argparse-free business logic; scripts/dwarf.py wires it up with Typer, matching giant/cli.py's style. `dwarf convert` merges the old serial/parallel conversion scripts behind a --jobs flag (default 1: sequential with plain -o; >1: dataset-layout fan-out via subprocess). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
247 lines
8.5 KiB
Python
247 lines
8.5 KiB
Python
import json
|
|
import stat
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from scripts import create_root_files
|
|
|
|
parse_detector_spec = create_root_files.parse_detector_spec
|
|
next_shard_index = create_root_files.next_shard_index
|
|
plan_jobs = create_root_files.plan_jobs
|
|
run_job = create_root_files.run_job
|
|
run_all = create_root_files.run_all
|
|
SimJob = create_root_files.SimJob
|
|
PlanError = create_root_files.PlanError
|
|
|
|
|
|
def _write_fake_executable(
|
|
path: Path, *, output_count: int = 1, exit_code: int = 0, sleep: float = 0.0
|
|
) -> Path:
|
|
"""Stand-in for run_pbwo4/run_sampling: writes *output_count* .root files
|
|
into its own cwd (so callers can verify each job gets an isolated workdir
|
|
and that the workdir ends up holding *only* the .root output, matching
|
|
real executables that may also drop other side-effect files). Each .root
|
|
file's content is a JSON record of argv/cwd/timing, so a test can recover
|
|
that info from the moved file after the workdir is gone. Also writes a
|
|
non-.root side file unconditionally, to catch any cleanup code that
|
|
assumes the workdir is empty after the .root is moved out.
|
|
"""
|
|
path.write_text(
|
|
f"""#!/usr/bin/env python3
|
|
import json, os, sys, time
|
|
|
|
start = time.time()
|
|
time.sleep({sleep})
|
|
end = time.time()
|
|
payload = json.dumps(
|
|
{{"argv": sys.argv[1:], "cwd": os.getcwd(), "start": start, "end": end}}
|
|
)
|
|
for i in range({output_count}):
|
|
with open(f"out_{{i}}.root", "w") as f:
|
|
f.write(payload)
|
|
with open("side_effect.log", "w") as f:
|
|
f.write("not a root file")
|
|
sys.exit({exit_code})
|
|
"""
|
|
)
|
|
path.chmod(path.stat().st_mode | stat.S_IEXEC | stat.S_IXGRP | stat.S_IXOTH)
|
|
return path
|
|
|
|
|
|
def test_parse_detector_spec_with_config():
|
|
assert parse_detector_spec("sampling_pb_scint:pb_scint") == ("sampling_pb_scint", "pb_scint")
|
|
|
|
|
|
def test_parse_detector_spec_without_config():
|
|
assert parse_detector_spec("pbwo4") == ("pbwo4", None)
|
|
|
|
|
|
def test_parse_detector_spec_rejects_empty_parts():
|
|
with pytest.raises(PlanError):
|
|
parse_detector_spec(":pb_scint")
|
|
with pytest.raises(PlanError):
|
|
parse_detector_spec("sampling_pb_scint:")
|
|
|
|
|
|
def test_next_shard_index_missing_dir(tmp_path):
|
|
assert next_shard_index(tmp_path / "nope") == 0
|
|
|
|
|
|
def test_next_shard_index_continues_past_existing(tmp_path):
|
|
d = tmp_path / "pbwo4"
|
|
d.mkdir()
|
|
(d / "shard-000.root").touch()
|
|
(d / "shard-005.root").touch()
|
|
(d / "not-a-shard.root").touch()
|
|
assert next_shard_index(d) == 6
|
|
|
|
|
|
def test_plan_jobs_rejects_missing_gen(tmp_path):
|
|
with pytest.raises(PlanError):
|
|
plan_jobs(["pbwo4"], num_files=2, dataset_root=tmp_path, kind="steps", gen="gen1")
|
|
|
|
|
|
def test_plan_jobs_rejects_malformed_gen(tmp_path):
|
|
(tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
|
|
with pytest.raises(PlanError):
|
|
plan_jobs(["pbwo4"], num_files=2, dataset_root=tmp_path, kind="steps", gen="notagen")
|
|
|
|
|
|
def test_plan_jobs_continues_from_existing_shards(tmp_path):
|
|
gen_dir = tmp_path / "raw" / "steps" / "gen1"
|
|
(gen_dir / "pbwo4").mkdir(parents=True)
|
|
(gen_dir / "pbwo4" / "shard-000.root").touch()
|
|
(gen_dir / "pbwo4" / "shard-001.root").touch()
|
|
|
|
jobs = plan_jobs(["pbwo4"], num_files=3, dataset_root=tmp_path, kind="steps", gen="gen1")
|
|
|
|
assert [j.shard_index for j in jobs] == [2, 3, 4]
|
|
assert all(j.detector == "pbwo4" and j.config is None for j in jobs)
|
|
|
|
|
|
def test_plan_jobs_multiple_detectors_each_start_independently(tmp_path):
|
|
gen_dir = tmp_path / "raw" / "steps" / "gen1"
|
|
(gen_dir / "sampling_pb_scint").mkdir(parents=True)
|
|
(gen_dir / "sampling_pb_scint" / "shard-003.root").touch()
|
|
(gen_dir / "sampling_fe_scint").mkdir(parents=True)
|
|
|
|
jobs = plan_jobs(
|
|
["sampling_pb_scint:pb_scint", "sampling_fe_scint:fe_scint"],
|
|
num_files=2,
|
|
dataset_root=tmp_path,
|
|
kind="steps",
|
|
gen="gen1",
|
|
)
|
|
|
|
by_detector = {}
|
|
for j in jobs:
|
|
by_detector.setdefault(j.detector, []).append(j.shard_index)
|
|
assert by_detector["sampling_pb_scint"] == [4, 5]
|
|
assert by_detector["sampling_fe_scint"] == [0, 1]
|
|
|
|
|
|
def test_run_job_moves_output_to_correct_shard_path(tmp_path):
|
|
fake = _write_fake_executable(tmp_path / "fake_exe.py")
|
|
gen_dir = tmp_path / "raw" / "steps" / "gen1"
|
|
gen_dir.mkdir(parents=True)
|
|
tmp_root = tmp_path / ".sim-tmp"
|
|
tmp_root.mkdir()
|
|
|
|
job = SimJob(detector="pbwo4", config=None, shard_index=7)
|
|
result = run_job(job, fake, 10000, tmp_path, "steps", "gen1", tmp_root)
|
|
|
|
assert result.ok
|
|
assert result.dest == gen_dir / "pbwo4" / "shard-007.root"
|
|
assert result.dest.is_file()
|
|
assert not any(tmp_root.iterdir()) # workdir cleaned up
|
|
|
|
|
|
def test_run_job_passes_config_arg_and_isolates_cwd(tmp_path):
|
|
fake = _write_fake_executable(tmp_path / "fake_exe.py")
|
|
(tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
|
|
tmp_root = tmp_path / ".sim-tmp"
|
|
tmp_root.mkdir()
|
|
|
|
job = SimJob(detector="sampling_pb_scint", config="pb_scint", shard_index=0)
|
|
result = run_job(job, fake, 10000, tmp_path, "steps", "gen1", tmp_root)
|
|
|
|
assert result.ok
|
|
payload = json.loads(result.dest.read_text())
|
|
assert payload["argv"] == ["pb_scint", "10000"]
|
|
# ran in its own scratch workdir under .sim-tmp, not directly in dataset_root
|
|
assert payload["cwd"] != str(tmp_path)
|
|
assert str(tmp_root) in payload["cwd"]
|
|
|
|
|
|
def test_run_job_omits_config_arg_when_none(tmp_path):
|
|
fake = _write_fake_executable(tmp_path / "fake_exe.py")
|
|
(tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
|
|
tmp_root = tmp_path / ".sim-tmp"
|
|
tmp_root.mkdir()
|
|
|
|
job = SimJob(detector="pbwo4", config=None, shard_index=0)
|
|
result = run_job(job, fake, 10000, tmp_path, "steps", "gen1", tmp_root)
|
|
|
|
payload = json.loads(result.dest.read_text())
|
|
assert payload["argv"] == ["10000"]
|
|
|
|
|
|
def test_run_job_fails_when_executable_errors(tmp_path):
|
|
fake = _write_fake_executable(tmp_path / "fake_exe.py", exit_code=1)
|
|
(tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
|
|
tmp_root = tmp_path / ".sim-tmp"
|
|
tmp_root.mkdir()
|
|
|
|
job = SimJob(detector="pbwo4", config=None, shard_index=0)
|
|
result = run_job(job, fake, 10000, tmp_path, "steps", "gen1", tmp_root)
|
|
|
|
assert not result.ok
|
|
assert "exited 1" in result.message
|
|
|
|
|
|
def test_run_job_fails_when_no_root_file_produced(tmp_path):
|
|
fake = _write_fake_executable(tmp_path / "fake_exe.py", output_count=0)
|
|
(tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
|
|
tmp_root = tmp_path / ".sim-tmp"
|
|
tmp_root.mkdir()
|
|
|
|
job = SimJob(detector="pbwo4", config=None, shard_index=0)
|
|
result = run_job(job, fake, 10000, tmp_path, "steps", "gen1", tmp_root)
|
|
|
|
assert not result.ok
|
|
assert "found 0" in result.message
|
|
|
|
|
|
def test_run_job_fails_when_multiple_root_files_produced(tmp_path):
|
|
fake = _write_fake_executable(tmp_path / "fake_exe.py", output_count=2)
|
|
(tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
|
|
tmp_root = tmp_path / ".sim-tmp"
|
|
tmp_root.mkdir()
|
|
|
|
job = SimJob(detector="pbwo4", config=None, shard_index=0)
|
|
result = run_job(job, fake, 10000, tmp_path, "steps", "gen1", tmp_root)
|
|
|
|
assert not result.ok
|
|
assert "found 2" in result.message
|
|
|
|
|
|
def test_run_job_refuses_to_overwrite_existing_shard(tmp_path):
|
|
fake = _write_fake_executable(tmp_path / "fake_exe.py")
|
|
detector_dir = tmp_path / "raw" / "steps" / "gen1" / "pbwo4"
|
|
detector_dir.mkdir(parents=True)
|
|
(detector_dir / "shard-000.root").write_text("already here")
|
|
tmp_root = tmp_path / ".sim-tmp"
|
|
tmp_root.mkdir()
|
|
|
|
job = SimJob(detector="pbwo4", config=None, shard_index=0)
|
|
result = run_job(job, fake, 10000, tmp_path, "steps", "gen1", tmp_root)
|
|
|
|
assert not result.ok
|
|
assert "overwrite" in result.message
|
|
assert (detector_dir / "shard-000.root").read_text() == "already here"
|
|
|
|
|
|
def test_run_all_caps_concurrency(tmp_path):
|
|
fake = _write_fake_executable(tmp_path / "fake_exe.py", sleep=0.2)
|
|
(tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
|
|
tmp_root = tmp_path / ".sim-tmp"
|
|
tmp_root.mkdir()
|
|
|
|
jobs = [SimJob(detector="pbwo4", config=None, shard_index=i) for i in range(6)]
|
|
results = run_all(jobs, fake, 10000, tmp_path, "steps", "gen1", max_workers=2, tmp_root=tmp_root)
|
|
|
|
assert all(r.ok for r in results)
|
|
assert {r.dest.name for r in results} == {f"shard-{i:03d}.root" for i in range(6)}
|
|
|
|
intervals = [json.loads(r.dest.read_text()) for r in results]
|
|
events = sorted(
|
|
[(p["start"], 1) for p in intervals] + [(p["end"], -1) for p in intervals]
|
|
)
|
|
concurrent = 0
|
|
peak = 0
|
|
for _, delta in events:
|
|
concurrent += delta
|
|
peak = max(peak, concurrent)
|
|
assert peak <= 2
|