lars
a8fc019e32
Concurrent job launches in create_root_files.py can start within the same wall-clock second, and minicalosim's default seed falls back to time(NULL) in that case — so two "independent" shards could silently get identical RNG state and produce byte-identical physics. Requires the companion MINICALOSIM_SEED env-var support in the minicalosim repo. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
305 lines
10 KiB
Python
305 lines
10 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
|
|
job_seed = create_root_files.job_seed
|
|
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,
|
|
"seed": os.environ.get("MINICALOSIM_SEED"),
|
|
}}
|
|
)
|
|
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_job_seed_deterministic():
|
|
job = SimJob(detector="pbwo4", config=None, shard_index=3)
|
|
assert job_seed("steps", "gen1", job) == job_seed("steps", "gen1", job)
|
|
|
|
|
|
def test_job_seed_varies_by_shard_index():
|
|
a = SimJob(detector="pbwo4", config=None, shard_index=0)
|
|
b = SimJob(detector="pbwo4", config=None, shard_index=1)
|
|
assert job_seed("steps", "gen1", a) != job_seed("steps", "gen1", b)
|
|
|
|
|
|
def test_job_seed_varies_by_detector():
|
|
a = SimJob(detector="pbwo4", config=None, shard_index=0)
|
|
b = SimJob(detector="sampling_pb_scint", config="pb_scint", shard_index=0)
|
|
assert job_seed("steps", "gen1", a) != job_seed("steps", "gen1", b)
|
|
|
|
|
|
def test_job_seed_varies_by_gen():
|
|
job = SimJob(detector="pbwo4", config=None, shard_index=0)
|
|
assert job_seed("steps", "gen1", job) != job_seed("steps", "gen2", job)
|
|
|
|
|
|
def test_run_job_passes_deterministic_seed_env_var(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=5)
|
|
result = run_job(job, fake, 10000, tmp_path, "steps", "gen1", tmp_root)
|
|
|
|
assert result.dest is not None
|
|
payload = json.loads(result.dest.read_text())
|
|
assert payload["seed"] == str(job_seed("steps", "gen1", job))
|
|
|
|
|
|
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 not None
|
|
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
|
|
assert result.dest is not None
|
|
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)
|
|
|
|
assert result.dest is not None
|
|
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 and r.dest is not None for r in results)
|
|
dests = [r.dest for r in results if r.dest is not None]
|
|
assert {d.name for d in dests} == {f"shard-{i:03d}.root" for i in range(6)}
|
|
|
|
intervals = [json.loads(d.read_text()) for d in dests]
|
|
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
|