fix(parity): resume at the next pending precision after preemption

src/gpu/modal_export.py

@@ -191,6 +191,18 @@
 app = modal.App("interscript-ml-export", image=IMAGE)
 
 
+def pending_precisions(
+    out_dir: Path, mid: str, precisions: list[str]
+) -> list[str]:
+    """Precision stages still to run: the margin report is the last
+    artifact a stage writes, so its presence means the stage (parity
+    block in the zip included) completed durably. Preemption restarts
+    resume at the next stage instead of redoing hours of decode."""
+    return [
+        p for p in precisions if not (out_dir / f"{mid}-margins-{p}.json").exists()
+    ]
+
+
 def _load_pairs(path: Path) -> list[tuple[str, str]]:
     import json
 
@@ -324,7 +336,10 @@ def stage(event: str) -> None:
     meta_path = Path("/root/interscript-ml", spec["metadata"])
     mid = re.search(r"^id:\s*(\S+)", meta_path.read_text(encoding="utf-8"), re.M).group(1)
     reports: dict[str, str] = {}
-    for precision in [q.strip() for q in precisions.split(",") if q.strip()]:
+    todo = pending_precisions(
+        out_dir, mid, [q.strip() for q in precisions.split(",") if q.strip()]
+    )
+    for precision in todo:
         zip_path = out_dir / f"{mid}-{precision}.zip"
         stage(f"onnx decode {precision}")
         report = run_parity(model, zip_path, pairs, max_len=128, reference=reference)

tests/test_parity_stage_skip.py

@@ -0,0 +1,35 @@
+"""Stage-skip resume for parity_model: preemption restarts must not
+redo precision stages whose margin reports are already durable."""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+import pytest
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))
+
+modal = pytest.importorskip("modal")
+
+from gpu.modal_export import pending_precisions  # noqa: E402
+
+
+def test_completed_stage_is_skipped(tmp_path: Path) -> None:
+    (tmp_path / "ara-diac-small-2.1-margins-fp32.json").write_text("{}")
+    got = pending_precisions(tmp_path, "ara-diac-small-2.1", ["fp32", "fp16", "int8"])
+    assert got == ["fp16", "int8"]
+
+
+def test_fresh_run_keeps_all_stages(tmp_path: Path) -> None:
+    got = pending_precisions(tmp_path, "ara-diac-small-2.1", ["fp32", "fp16", "int8"])
+    assert got == ["fp32", "fp16", "int8"]
+
+
+def test_zip_without_margin_report_still_runs(tmp_path: Path) -> None:
+    # a partial stage (zip present, margin report missing) must rerun —
+    # the margin report is the last artifact written, so its absence
+    # means the stage never completed
+    (tmp_path / "ara-diac-small-2.1-int8.zip").write_text("partial")
+    got = pending_precisions(tmp_path, "ara-diac-small-2.1", ["int8"])
+    assert got == ["int8"]