Fix stale-partial reuse and n_chunks mismatch in analysis condor pipeline

- prep() now clears reduced_partial/ and reduced/ on every (re-)run.
  Partial files carry no record of what context (n_chunks, bin edges,
  group sets) they were computed under, so re-prepping the same run_dir
  with a different --chunks/--bins/--top-pdg (or after the rollout was
  regenerated) previously left old partials on disk that merge_one
  would silently merge against the new shared.json — producing a
  wrong-but-plausible reduced/*.json with no error.
- write_submit() now checks SubmitConfig.n_chunks against the run
  directory's own RunMeta.n_chunks (fixed at prep time, and what
  rows_per_chunk is sized against) and raises a clear error on
  mismatch, instead of an uncaught IndexError deep in _job_walltimes.

Each fix has a regression test.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

giant/analysis/condor.py

@@ -42,6 +42,7 @@ HTCondor file transfer of the multi-GB inputs.
 from __future__ import annotations
 
 import json
+import shutil
 from dataclasses import dataclass, field
 from pathlib import Path
 
@@ -194,11 +195,23 @@ def prep(
     resolved once here rather than re-passed (and risking disagreement) at every
     later step. See ``derive_run_dir`` for how ``run_dir``/``default_base``
     resolve the actual directory.
+
+    Clears any existing ``reduced_partial/``/``reduced/`` from a prior prep of
+    this same ``run_dir``: partial files carry no record of what context
+    (``n_chunks``, bin edges, group sets) they were computed under, so
+    re-prepping with a different ``n_chunks``/``**ctx_kwargs`` (or after the
+    rollout/reference files changed) would otherwise let ``merge_one`` silently
+    merge stale partials against the new ``shared.json``.
     """
     y = load_rollout_yaml(rollout_yaml)
     run_path = derive_run_dir(y, run_dir, default_base=default_base)
     run_path.mkdir(parents=True, exist_ok=True)
 
+    for stale in ("reduced_partial", "reduced"):
+        stale_dir = run_path / stale
+        if stale_dir.exists():
+            shutil.rmtree(stale_dir)
+
     rollout, reference = y["output"], y["dataset"]
     ctx = build_context(rollout, reference, **ctx_kwargs)
     ctx.save(run_path / "shared.json")
@@ -403,6 +416,11 @@ def write_submit(cfg: SubmitConfig, ids: list[str] | None = None) -> Path:
     ``run_meta.json`` from ``prep`` to already carry ``rows_per_chunk``).
     Returns the submit description path (``<run_dir>/analyze.sub``). Does not
     submit — call ``condor_submit`` on the returned file.
+
+    ``cfg.n_chunks`` and the run directory's own ``RunMeta.n_chunks`` (fixed by
+    ``prep``, and what ``RunMeta.rows_per_chunk`` was sized against) are two
+    independent values — checked equal up front so a mismatch is a clear error
+    here rather than an ``IndexError`` out of ``_job_walltimes``.
     """
     venv_giant = cfg.repo_dir / ".venv" / "bin" / "giant"
     if not venv_giant.exists():
@@ -414,6 +432,15 @@ def write_submit(cfg: SubmitConfig, ids: list[str] | None = None) -> Path:
 
     ids = ids or catalog_ids()
     run_dir = cfg.run_dir
+    meta = RunMeta.load(run_dir / "run_meta.json")
+    if cfg.n_chunks != meta.n_chunks:
+        raise ValueError(
+            f"SubmitConfig.n_chunks={cfg.n_chunks} does not match the "
+            f"n_chunks this run directory was prepped with "
+            f"(RunMeta.n_chunks={meta.n_chunks} in {run_dir}/run_meta.json) — "
+            "re-run `prep` with the desired n_chunks, or fix cfg.n_chunks to "
+            "match it."
+        )
     (run_dir / "logs").mkdir(parents=True, exist_ok=True)
     (run_dir / "reduced").mkdir(parents=True, exist_ok=True)
     (run_dir / "reduced_partial").mkdir(parents=True, exist_ok=True)

tests/test_condor.py

@@ -121,6 +121,26 @@ def test_prep_splits_rows_per_chunk(tmp_path: Path):
     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)
@@ -237,6 +257,20 @@ def test_write_submit_chunks_respect_chunkable(tmp_path: Path):
     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