Add tooling for a versioned geant_steps dataset layout
Introduces raw/<kind>/<gen>/<detector>/shard-NNN.root and processed/<kind>/<gen>/<schema>/<detector>/shard-NNN.parquet as the dataset convention, plus scripts to operate on it: migrate_geant_steps.py for the one-time move into this layout, bump_dataset_version.py to cut new gen/schema versions with a logged reason, steps_to_parquet_parallel.py to convert ROOT shards to parquet in parallel and place them correctly, and create_root_files.py to generate new ROOT shards via a minicalosim executable. The loader gains .manifest file support so pools/ (train/dev/ holdout shard lists) can be passed straight to `giant train`. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,95 @@
|
||||
import importlib.util
|
||||
from pathlib import Path
|
||||
|
||||
# scripts/ is not an installed package — load the module straight from its path.
|
||||
_SPEC = importlib.util.spec_from_file_location(
|
||||
"bump_dataset_version",
|
||||
Path(__file__).resolve().parents[1] / "scripts" / "bump_dataset_version.py",
|
||||
)
|
||||
bump_dataset_version = importlib.util.module_from_spec(_SPEC)
|
||||
_SPEC.loader.exec_module(bump_dataset_version)
|
||||
|
||||
plan_bump_gen = bump_dataset_version.plan_bump_gen
|
||||
plan_bump_schema = bump_dataset_version.plan_bump_schema
|
||||
apply_bump = bump_dataset_version.apply_bump
|
||||
|
||||
|
||||
def test_bump_gen_starts_at_gen1_when_none_exist(tmp_path):
|
||||
dirs, log_line = plan_bump_gen(tmp_path, "steps", "first generation", None, "2026-01-01")
|
||||
assert dirs == [
|
||||
tmp_path / "raw" / "steps" / "gen1",
|
||||
tmp_path / "processed" / "steps" / "gen1" / "schema1",
|
||||
]
|
||||
assert "`gen1`" in log_line
|
||||
assert "first generation" in log_line
|
||||
|
||||
|
||||
def test_bump_gen_increments_past_existing(tmp_path):
|
||||
(tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
|
||||
(tmp_path / "raw" / "steps" / "gen2").mkdir(parents=True)
|
||||
dirs, _ = plan_bump_gen(tmp_path, "steps", "next gen", None, "2026-01-01")
|
||||
assert dirs[0] == tmp_path / "raw" / "steps" / "gen3"
|
||||
|
||||
|
||||
def test_bump_gen_checks_both_raw_and_processed_trees(tmp_path):
|
||||
# processed/ is ahead of raw/ — next gen must still be past the max of both.
|
||||
(tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
|
||||
(tmp_path / "processed" / "steps" / "gen4" / "schema1").mkdir(parents=True)
|
||||
dirs, _ = plan_bump_gen(tmp_path, "steps", "next gen", None, "2026-01-01")
|
||||
assert dirs[0] == tmp_path / "raw" / "steps" / "gen5"
|
||||
|
||||
|
||||
def test_bump_gen_kinds_are_independent(tmp_path):
|
||||
(tmp_path / "raw" / "steps" / "gen5").mkdir(parents=True)
|
||||
dirs, _ = plan_bump_gen(tmp_path, "hits", "first hits gen", None, "2026-01-01")
|
||||
assert dirs[0] == tmp_path / "raw" / "hits" / "gen1"
|
||||
|
||||
|
||||
def test_bump_schema_starts_at_schema1_for_a_fresh_gen(tmp_path):
|
||||
(tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
|
||||
dirs, log_line = plan_bump_schema(
|
||||
tmp_path, "steps", "gen1", "added e_sec column", None, "2026-01-01"
|
||||
)
|
||||
assert dirs == [tmp_path / "processed" / "steps" / "gen1" / "schema1"]
|
||||
assert "`gen1`/`schema1`" in log_line
|
||||
|
||||
|
||||
def test_bump_schema_increments_within_its_gen(tmp_path):
|
||||
(tmp_path / "processed" / "steps" / "gen1" / "schema1").mkdir(parents=True)
|
||||
(tmp_path / "processed" / "steps" / "gen1" / "schema2").mkdir(parents=True)
|
||||
dirs, _ = plan_bump_schema(tmp_path, "steps", "gen1", "next schema", None, "2026-01-01")
|
||||
assert dirs == [tmp_path / "processed" / "steps" / "gen1" / "schema3"]
|
||||
|
||||
|
||||
def test_bump_schema_does_not_see_other_gens_schemas(tmp_path):
|
||||
(tmp_path / "processed" / "steps" / "gen1" / "schema5").mkdir(parents=True)
|
||||
(tmp_path / "raw" / "steps" / "gen2").mkdir(parents=True)
|
||||
dirs, _ = plan_bump_schema(tmp_path, "steps", "gen2", "fresh schema for gen2", None, "2026-01-01")
|
||||
assert dirs == [tmp_path / "processed" / "steps" / "gen2" / "schema1"]
|
||||
|
||||
|
||||
def test_bump_schema_rejects_nonexistent_gen(tmp_path):
|
||||
try:
|
||||
plan_bump_schema(tmp_path, "steps", "gen9", "oops", None, "2026-01-01")
|
||||
assert False, "expected SystemExit"
|
||||
except SystemExit:
|
||||
pass
|
||||
|
||||
|
||||
def test_apply_bump_creates_dirs_and_appends_log(tmp_path):
|
||||
dirs, log_line = plan_bump_gen(tmp_path, "steps", "reason A", "alice", "2026-01-01")
|
||||
apply_bump(tmp_path, dirs, log_line)
|
||||
for d in dirs:
|
||||
assert d.is_dir()
|
||||
text = (tmp_path / "VERSIONS.md").read_text()
|
||||
assert "reason A" in text
|
||||
assert "alice" in text
|
||||
|
||||
|
||||
def test_apply_bump_appends_without_clobbering_existing_log(tmp_path):
|
||||
(tmp_path / "VERSIONS.md").write_text("# Dataset versions\n\n- existing entry\n")
|
||||
dirs, log_line = plan_bump_gen(tmp_path, "steps", "reason B", None, "2026-01-02")
|
||||
apply_bump(tmp_path, dirs, log_line)
|
||||
text = (tmp_path / "VERSIONS.md").read_text()
|
||||
assert "existing entry" in text
|
||||
assert "reason B" in text
|
||||
@@ -0,0 +1,274 @@
|
||||
import importlib.util
|
||||
import json
|
||||
import stat
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
# scripts/ is not an installed package — load the module straight from its path.
|
||||
_SPEC = importlib.util.spec_from_file_location(
|
||||
"create_root_files",
|
||||
Path(__file__).resolve().parents[1] / "scripts" / "create_root_files.py",
|
||||
)
|
||||
create_root_files = importlib.util.module_from_spec(_SPEC)
|
||||
_SPEC.loader.exec_module(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
|
||||
|
||||
|
||||
def test_build_parser_defaults():
|
||||
args = create_root_files.build_parser().parse_args(
|
||||
[
|
||||
"--executable",
|
||||
"fake",
|
||||
"--detector",
|
||||
"pbwo4",
|
||||
"--num-files",
|
||||
"2",
|
||||
"--events-per-file",
|
||||
"10000",
|
||||
"--gen",
|
||||
"gen1",
|
||||
]
|
||||
)
|
||||
assert args.jobs == 4
|
||||
assert args.kind == "steps"
|
||||
assert args.dataset_root == "/ceph/lbogner/geant_steps"
|
||||
assert args.execute is False
|
||||
@@ -0,0 +1,60 @@
|
||||
import pytest
|
||||
|
||||
from giant.data.loader import find_parquet_files
|
||||
|
||||
|
||||
def _touch(path):
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.touch()
|
||||
return path
|
||||
|
||||
|
||||
def test_find_parquet_files_single_file(tmp_path):
|
||||
f = _touch(tmp_path / "shard-000.parquet")
|
||||
assert find_parquet_files(f) == [f]
|
||||
|
||||
|
||||
def test_find_parquet_files_directory_glob(tmp_path):
|
||||
a = _touch(tmp_path / "shard-000.parquet")
|
||||
b = _touch(tmp_path / "shard-001.parquet")
|
||||
_touch(tmp_path / "not_a_parquet.root")
|
||||
assert find_parquet_files(tmp_path) == sorted([a, b])
|
||||
|
||||
|
||||
def test_find_parquet_files_empty_directory_raises(tmp_path):
|
||||
with pytest.raises(FileNotFoundError):
|
||||
find_parquet_files(tmp_path)
|
||||
|
||||
|
||||
def test_manifest_resolves_relative_to_its_own_directory(tmp_path):
|
||||
target = _touch(tmp_path / "processed" / "pbwo4" / "shard-000.parquet")
|
||||
manifest_dir = tmp_path / "pools" / "pbwo4"
|
||||
manifest_dir.mkdir(parents=True)
|
||||
manifest = manifest_dir / "full.manifest"
|
||||
manifest.write_text("../../processed/pbwo4/shard-000.parquet\n")
|
||||
|
||||
assert find_parquet_files(manifest) == [target.resolve()]
|
||||
|
||||
|
||||
def test_manifest_skips_blank_lines_and_comments(tmp_path):
|
||||
target = _touch(tmp_path / "shard-000.parquet")
|
||||
manifest = tmp_path / "full.manifest"
|
||||
manifest.write_text("\n# a comment\nshard-000.parquet\n\n")
|
||||
|
||||
assert find_parquet_files(manifest) == [target.resolve()]
|
||||
|
||||
|
||||
def test_manifest_missing_file_raises(tmp_path):
|
||||
manifest = tmp_path / "full.manifest"
|
||||
manifest.write_text("does_not_exist.parquet\n")
|
||||
|
||||
with pytest.raises(FileNotFoundError):
|
||||
find_parquet_files(manifest)
|
||||
|
||||
|
||||
def test_manifest_with_no_entries_raises(tmp_path):
|
||||
manifest = tmp_path / "full.manifest"
|
||||
manifest.write_text("# only comments\n")
|
||||
|
||||
with pytest.raises(FileNotFoundError):
|
||||
find_parquet_files(manifest)
|
||||
@@ -0,0 +1,186 @@
|
||||
import importlib.util
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
# scripts/ is not an installed package — load the module straight from its path.
|
||||
_SPEC = importlib.util.spec_from_file_location(
|
||||
"steps_to_parquet_parallel",
|
||||
Path(__file__).resolve().parents[1] / "scripts" / "steps_to_parquet_parallel.py",
|
||||
)
|
||||
steps_to_parquet_parallel = importlib.util.module_from_spec(_SPEC)
|
||||
_SPEC.loader.exec_module(steps_to_parquet_parallel)
|
||||
|
||||
run_parallel = steps_to_parquet_parallel.run_parallel
|
||||
resolve_destination = steps_to_parquet_parallel.resolve_destination
|
||||
latest_schema_tag = steps_to_parquet_parallel.latest_schema_tag
|
||||
DestinationError = steps_to_parquet_parallel.DestinationError
|
||||
|
||||
|
||||
def _write_fake_executable(tmp_path: Path, marker_dir: Path) -> Path:
|
||||
"""A stand-in for steps_to_parquet.py: records its own start/end time per
|
||||
input file (so tests can check overlap), then exits 1 if "fail" is in the
|
||||
filename, else 0. Accepts (and ignores) the same flags as the real script.
|
||||
"""
|
||||
fake = tmp_path / "fake_steps_to_parquet.py"
|
||||
fake.write_text(
|
||||
f"""
|
||||
import json, sys, time
|
||||
from pathlib import Path
|
||||
|
||||
marker_dir = Path({str(marker_dir)!r})
|
||||
root_file = sys.argv[1]
|
||||
name = Path(root_file).stem
|
||||
start = time.time()
|
||||
time.sleep(0.2)
|
||||
end = time.time()
|
||||
(marker_dir / f"{{name}}.json").write_text(json.dumps({{"start": start, "end": end}}))
|
||||
print(f"converted {{root_file}}")
|
||||
sys.exit(1 if "fail" in name else 0)
|
||||
"""
|
||||
)
|
||||
return fake
|
||||
|
||||
|
||||
def test_runs_one_job_per_file_and_reports_success(tmp_path):
|
||||
marker_dir = tmp_path / "markers"
|
||||
marker_dir.mkdir()
|
||||
fake = _write_fake_executable(tmp_path, marker_dir)
|
||||
|
||||
files = [str(tmp_path / f"shard-{i:03d}.root") for i in range(3)]
|
||||
results = run_parallel(files, jobs=4, steps_to_parquet_path=fake)
|
||||
|
||||
assert {r[0] for r in results} == set(files)
|
||||
assert all(code == 0 for _, code, _, _ in results)
|
||||
assert all((marker_dir / f"{Path(f).stem}.json").exists() for f in files)
|
||||
|
||||
|
||||
def test_failures_are_reported_with_nonzero_exit_code(tmp_path):
|
||||
marker_dir = tmp_path / "markers"
|
||||
marker_dir.mkdir()
|
||||
fake = _write_fake_executable(tmp_path, marker_dir)
|
||||
|
||||
files = [str(tmp_path / "shard-000.root"), str(tmp_path / "shard-fail.root")]
|
||||
results = run_parallel(files, jobs=4, steps_to_parquet_path=fake)
|
||||
|
||||
codes = {Path(f).stem: code for f, code, _, _ in results}
|
||||
assert codes["shard-000"] == 0
|
||||
assert codes["shard-fail"] == 1
|
||||
|
||||
|
||||
def test_jobs_caps_concurrency(tmp_path):
|
||||
marker_dir = tmp_path / "markers"
|
||||
marker_dir.mkdir()
|
||||
fake = _write_fake_executable(tmp_path, marker_dir)
|
||||
|
||||
files = [str(tmp_path / f"shard-{i:03d}.root") for i in range(6)]
|
||||
run_parallel(files, jobs=2, steps_to_parquet_path=fake)
|
||||
|
||||
intervals = []
|
||||
for f in files:
|
||||
data = json.loads((marker_dir / f"{Path(f).stem}.json").read_text())
|
||||
intervals.append((data["start"], data["end"]))
|
||||
|
||||
# max number of intervals overlapping any given instant must not exceed jobs
|
||||
events = sorted([(s, 1) for s, _ in intervals] + [(e, -1) for _, e in intervals])
|
||||
concurrent = 0
|
||||
peak = 0
|
||||
for _, delta in events:
|
||||
concurrent += delta
|
||||
peak = max(peak, concurrent)
|
||||
assert peak <= 2
|
||||
|
||||
|
||||
def test_default_jobs_is_four():
|
||||
args = steps_to_parquet_parallel.build_parser().parse_args(["dummy.root"])
|
||||
assert args.jobs == 4
|
||||
|
||||
|
||||
def test_output_for_is_passed_through_as_output_flag(tmp_path):
|
||||
marker_dir = tmp_path / "markers"
|
||||
marker_dir.mkdir()
|
||||
fake = tmp_path / "fake_steps_to_parquet.py"
|
||||
fake.write_text(
|
||||
f"""
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
marker_dir = Path({str(marker_dir)!r})
|
||||
root_file = sys.argv[1]
|
||||
name = Path(root_file).stem
|
||||
output = sys.argv[sys.argv.index("--output") + 1] if "--output" in sys.argv else ""
|
||||
(marker_dir / f"{{name}}.txt").write_text(output)
|
||||
sys.exit(0)
|
||||
"""
|
||||
)
|
||||
root_file = str(tmp_path / "shard-000.root")
|
||||
dest = tmp_path / "processed" / "shard-000.parquet"
|
||||
run_parallel(
|
||||
[root_file], jobs=1, steps_to_parquet_path=fake, output_for={root_file: dest}
|
||||
)
|
||||
assert (marker_dir / "shard-000.txt").read_text() == str(dest)
|
||||
|
||||
|
||||
def _make_dataset(tmp_path: Path, schemas: list[str] | None = None) -> Path:
|
||||
root_file = tmp_path / "raw" / "steps" / "gen1" / "pbwo4" / "shard-000.root"
|
||||
root_file.parent.mkdir(parents=True)
|
||||
root_file.touch()
|
||||
for schema in schemas or []:
|
||||
(tmp_path / "processed" / "steps" / "gen1" / schema).mkdir(parents=True)
|
||||
return root_file
|
||||
|
||||
|
||||
def test_resolve_destination_uses_latest_schema(tmp_path):
|
||||
root_file = _make_dataset(tmp_path, schemas=["schema1", "schema3", "schema2"])
|
||||
dest = resolve_destination(root_file, tmp_path, schema_override=None)
|
||||
assert dest == tmp_path / "processed" / "steps" / "gen1" / "schema3" / "pbwo4" / "shard-000.parquet"
|
||||
|
||||
|
||||
def test_resolve_destination_schema_override_wins(tmp_path):
|
||||
root_file = _make_dataset(tmp_path, schemas=["schema1", "schema3"])
|
||||
dest = resolve_destination(root_file, tmp_path, schema_override="schema9")
|
||||
assert dest == tmp_path / "processed" / "steps" / "gen1" / "schema9" / "pbwo4" / "shard-000.parquet"
|
||||
|
||||
|
||||
def test_resolve_destination_errors_without_any_schema(tmp_path):
|
||||
root_file = _make_dataset(tmp_path, schemas=[])
|
||||
try:
|
||||
resolve_destination(root_file, tmp_path, schema_override=None)
|
||||
assert False, "expected DestinationError"
|
||||
except DestinationError:
|
||||
pass
|
||||
|
||||
|
||||
def test_resolve_destination_errors_outside_dataset_root(tmp_path):
|
||||
other_root = tmp_path / "other"
|
||||
other_root.mkdir()
|
||||
stray_file = other_root / "shard-000.root"
|
||||
stray_file.touch()
|
||||
dataset_root = tmp_path / "dataset"
|
||||
dataset_root.mkdir()
|
||||
try:
|
||||
resolve_destination(stray_file, dataset_root, schema_override=None)
|
||||
assert False, "expected DestinationError"
|
||||
except DestinationError:
|
||||
pass
|
||||
|
||||
|
||||
def test_resolve_destination_errors_on_wrong_shape(tmp_path):
|
||||
# missing the <gen> segment
|
||||
bad_file = tmp_path / "raw" / "steps" / "pbwo4" / "shard-000.root"
|
||||
bad_file.parent.mkdir(parents=True)
|
||||
bad_file.touch()
|
||||
try:
|
||||
resolve_destination(bad_file, tmp_path, schema_override=None)
|
||||
assert False, "expected DestinationError"
|
||||
except DestinationError:
|
||||
pass
|
||||
|
||||
|
||||
def test_latest_schema_tag_returns_none_when_missing(tmp_path):
|
||||
assert latest_schema_tag(tmp_path / "does" / "not" / "exist") is None
|
||||
|
||||
|
||||
def test_dataset_root_and_schema_flags_default(tmp_path):
|
||||
args = steps_to_parquet_parallel.build_parser().parse_args(["dummy.root"])
|
||||
assert args.dataset_root == "/ceph/lbogner/geant_steps"
|
||||
assert args.schema is None
|
||||
Reference in New Issue
Block a user