giant/tests/test_condor.py

"""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_prediction_yaml,
    load_prediction_yamls,
    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, PREDICT_TRUTH_METADATA_KEY, ROLLOUT_COORD_VALUE
from tests.test_analysis_prediction import _global_prediction_frame
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 _write_prediction(path: Path, coord: str = "global") -> None:
    tbl = _global_prediction_frame().collect().to_arrow()
    tbl = tbl.replace_schema_metadata({PREDICT_COORD_METADATA_KEY: coord, PREDICT_TRUTH_METADATA_KEY: "1"})
    pq.write_table(tbl, path)


def _write_prediction_yaml(tmp_path: Path, reference: Path, tag: str = "p", coord: str = "global") -> Path:
    pred = tmp_path / f"pred_{tag}.parquet"
    _write_prediction(pred, coord=coord)
    yaml_path = tmp_path / f"pred_{tag}.yaml"
    yaml_path.write_text(
        yaml.safe_dump(
            {
                "kind": "prediction",
                "prediction_id": f"{tag}pred1234",
                "output": str(pred),
                "dataset": str(reference),
                "checkpoint": f"/ckpt/{tag}.pt",
            }
        )
    )
    return yaml_path


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, prediction_yamls=()) -> Path:
    """``prep`` with small test-sized context bins/sampling."""
    return prep(
        rollout_yamls,
        run_dir,
        n_chunks=chunks,
        labels=labels,
        prediction_yamls=prediction_yamls,
        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_load_prediction_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_prediction_yaml(bad)


def test_load_prediction_yaml_rejects_rollout_kind(tmp_path: Path):
    y = tmp_path / "r.yaml"
    y.write_text(yaml.safe_dump({"output": "x.parquet", "dataset": "d.parquet", "kind": "rollout"}))
    with pytest.raises(ValueError, match="kind"):
        load_prediction_yaml(y)


def test_load_prediction_yamls_single_defaults_to_prediction_name(tmp_path: Path):
    reference = tmp_path / "reference.parquet"
    _reference_frame().collect().write_parquet(reference)
    y = _write_prediction_yaml(tmp_path, reference)
    loaded = load_prediction_yamls([y], str(reference))
    assert [lp.name for lp in loaded] == ["prediction"]
    assert loaded[0].coord == "global"


def test_load_prediction_yamls_multi_defaults_to_stem_and_labels(tmp_path: Path):
    reference = tmp_path / "reference.parquet"
    _reference_frame().collect().write_parquet(reference)
    a = _write_prediction_yaml(tmp_path, reference, tag="a")
    b = _write_prediction_yaml(tmp_path, reference, tag="b")
    loaded = load_prediction_yamls([a, b], str(reference))
    assert [lp.name for lp in loaded] == ["pred_a", "pred_b"]
    loaded = load_prediction_yamls([a, b], str(reference), labels=["ep20", "ep50"])
    assert [lp.name for lp in loaded] == ["ep20", "ep50"]


def test_load_prediction_yamls_rejects_mismatched_reference(tmp_path: Path):
    reference = tmp_path / "reference.parquet"
    _reference_frame().collect().write_parquet(reference)
    other_ref = tmp_path / "other_reference.parquet"
    _reference_frame().collect().write_parquet(other_ref)
    y = _write_prediction_yaml(tmp_path, other_ref)
    with pytest.raises(ValueError, match="same reference"):
        load_prediction_yamls([y], str(reference))


def test_load_prediction_yamls_rejects_mixed_coord(tmp_path: Path):
    reference = tmp_path / "reference.parquet"
    _reference_frame().collect().write_parquet(reference)
    a = _write_prediction_yaml(tmp_path, reference, tag="a", coord="global")
    b = _write_prediction_yaml(tmp_path, reference, tag="b", coord="local")
    with pytest.raises(ValueError, match="coord"):
        load_prediction_yamls([a, b], str(reference))


def test_prep_with_prediction_writes_run_meta(tmp_path: Path):
    rollout_yaml = _write_inputs(tmp_path)
    reference = load_rollout_yaml(rollout_yaml)["dataset"]
    pred_yaml = _write_prediction_yaml(tmp_path, Path(reference))
    run_dir = _prep([rollout_yaml], prediction_yamls=[pred_yaml])
    meta = RunMeta.load(run_dir / "run_meta.json")
    assert [p["name"] for p in meta.predictions] == ["prediction"]
    assert meta.predictions[0]["plot_meta"]["checkpoint"] == "/ckpt/p.pt"

    computed = compute_one("pred_marginal_edep", run_dir, chunk_index=0)
    partial = Partial.load(computed)
    assert partial.data["available"]


def test_prep_forwards_predict_only_metadata_keys(tmp_path: Path):
    """A rich `giant predict` sidecar's provenance/timing keys reach
    run_meta.json's plot_meta, same as a rollout's do — a thin legacy
    sidecar (no such keys) still loads fine (see _write_prediction_yaml)."""
    rollout_yaml = _write_inputs(tmp_path)
    reference = load_rollout_yaml(rollout_yaml)["dataset"]
    pred = tmp_path / "pred_rich.parquet"
    _write_prediction(pred, coord="global")
    yaml_path = tmp_path / "pred_rich.yaml"
    yaml_path.write_text(
        yaml.safe_dump(
            {
                "kind": "prediction",
                "prediction_id": "richpred12",
                "output": str(pred),
                "dataset": str(reference),
                "checkpoint": "/ckpt/rich.pt",
                "coord": "global",
                "has_truth": True,
                "schema_version": "3",
                "n_input_rows": 1000,
                "n_files": 1,
                "n_skipped_rows": 3,
                "unknown_pdg_counts": {"999999": 3},
                "batch_size_auto": False,
                "timing": {"us_per_step": 12.5},
            }
        )
    )
    run_dir = _prep([rollout_yaml], prediction_yamls=[yaml_path])
    meta = RunMeta.load(run_dir / "run_meta.json")
    plot_meta = meta.predictions[0]["plot_meta"]
    assert plot_meta["coord"] == "global"
    assert plot_meta["has_truth"] is True
    assert plot_meta["n_input_rows"] == 1000
    assert plot_meta["n_skipped_rows"] == 3
    assert plot_meta["unknown_pdg_counts"] == {"999999": 3}
    assert plot_meta["timing"] == {"us_per_step": 12.5}


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_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