From 3d5665e97afd21edf54339b63a60a1ac85cb58e0 Mon Sep 17 00:00:00 2001
From: Ronald Tse <ronald.tse@gmail.com>
Date: Thu, 3 Sep 2026 17:22:37 +0200
Subject: [PATCH] fix(parity): resume at the next pending precision after
 preemption

Each restart redid every precision stage (~4h of fp32+fp16 decode)
because the loop had no completion check. The margin report is the
last artifact a stage writes, so its presence on the volume marks the
stage durably complete; restarts now skip straight to the pending one.
---
 src/gpu/modal_export.py         | 17 +++++++++++++++-
 tests/test_parity_stage_skip.py | 35 +++++++++++++++++++++++++++++++++
 2 files changed, 51 insertions(+), 1 deletion(-)
 create mode 100644 tests/test_parity_stage_skip.py

diff --git a/src/gpu/modal_export.py b/src/gpu/modal_export.py
index cbe368a..bec6117 100644
--- a/src/gpu/modal_export.py
+++ b/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)
diff --git a/tests/test_parity_stage_skip.py b/tests/test_parity_stage_skip.py
new file mode 100644
index 0000000..13fb229
--- /dev/null
+++ b/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"]