fix(parity): stage completeness = zip parity block + margins json

src/gpu/modal_export.py

@@ -200,17 +200,39 @@ def normalize_precisions(precisions: "str | list[str]") -> list[str]:
     return [q.strip() for q in precisions if q.strip()]
 
 
+def _zip_has_parity(zip_path: Path) -> bool:
+    """A stage writes its parity block into the zip BEFORE the margins
+    report; both must be present for the stage to count as complete
+    (a standalone margin_model run produces the json without the zip
+    block, which is not shippable)."""
+    if not zip_path.exists():
+        return False
+    import zipfile
+
+    try:
+        with zipfile.ZipFile(zip_path) as zf:
+            import yaml
+
+            meta = yaml.safe_load(zf.read("metadata.yaml"))
+    except (OSError, KeyError, ValueError, zipfile.BadZipFile):
+        return False
+    return isinstance(meta, dict) and "parity" in meta
+
+
 def pending_precisions(
     out_dir: Path, mid: str, precisions: "str | 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."""
+    """Precision stages still to run. A stage is durably complete only
+    when BOTH its margin report exists and its zip carries the parity
+    block. Preemption restarts resume at the next stage instead of
+    redoing hours of decode."""
     return [
         p
         for p in normalize_precisions(precisions)
-        if not (out_dir / f"{mid}-margins-{p}.json").exists()
+        if not (
+            (out_dir / f"{mid}-margins-{p}.json").exists()
+            and _zip_has_parity(out_dir / f"{mid}-{p}.zip")
+        )
     ]
 
 

tests/test_parity_stage_skip.py

@@ -15,7 +15,19 @@
 from gpu.modal_export import pending_precisions  # noqa: E402
 
 
+def _write_zip(path: Path, parity: bool) -> None:
+    import zipfile
+    import yaml
+
+    meta = {"id": "m", "precision": "int4"}
+    if parity:
+        meta["parity"] = {"cer_delta": 0.05}
+    with zipfile.ZipFile(path, "w") as zf:
+        zf.writestr("metadata.yaml", yaml.safe_dump(meta))
+
+
 def test_completed_stage_is_skipped(tmp_path: Path) -> None:
+    _write_zip(tmp_path / "ara-diac-small-2.1-fp32.zip", parity=True)
     (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"]
@@ -39,6 +51,7 @@ def test_list_input_from_entrypoint(tmp_path: Path) -> None:
     # the parity/margins entrypoints pass precisions.split(",") — a
     # list — into the remote functions; the functions must accept both
     # forms (direct ::parity_model CLI invocation passes a string)
+    _write_zip(tmp_path / "ara-diac-small-2.1-fp32.zip", parity=True)
     (tmp_path / "ara-diac-small-2.1-margins-fp32.json").write_text("{}")
     got = pending_precisions(tmp_path, "ara-diac-small-2.1", ["fp32", "int8"])
     assert got == ["int8"]
@@ -47,3 +60,23 @@ def test_list_input_from_entrypoint(tmp_path: Path) -> None:
 def test_string_input_strips_whitespace(tmp_path: Path) -> None:
     got = pending_precisions(tmp_path, "m", "fp32, int8")
     assert got == ["fp32", "int8"]
+
+
+def test_margins_without_zip_parity_reruns(tmp_path: Path) -> None:
+    # a standalone margin_model run writes the margins json but never
+    # the zip's parity block; such a stage must rerun, not skip
+    (tmp_path / "m-margins-int4.json").write_text("{}")
+    _write_zip(tmp_path / "m-int4.zip", parity=False)
+    assert pending_precisions(tmp_path, "m", ["int4"]) == ["int4"]
+
+
+def test_margins_and_zip_parity_skip(tmp_path: Path) -> None:
+    (tmp_path / "m-margins-int4.json").write_text("{}")
+    _write_zip(tmp_path / "m-int4.zip", parity=True)
+    assert pending_precisions(tmp_path, "m", ["int4"]) == []
+
+
+def test_corrupt_zip_is_pending(tmp_path: Path) -> None:
+    (tmp_path / "m-margins-int4.json").write_text("{}")
+    (tmp_path / "m-int4.zip").write_text("not a zip")
+    assert pending_precisions(tmp_path, "m", ["int4"]) == ["int4"]