From 17a53143c35955e6631bad22b76a9b63d1bc51e2 Mon Sep 17 00:00:00 2001 From: Ronald Tse Date: Fri, 4 Sep 2026 08:27:31 +0200 Subject: [PATCH] fix(parity): a stage is complete only with zip parity block AND margins json MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The layerdrop chain skipped int4 because yesterday's standalone margin_model run had written margins-int4.json — but margin_model never writes the zip's parity block, so the int4 zip shipped parity-less and would fail the strict validator. pending_precisions now requires both markers; corrupt or missing zips are pending. --- src/gpu/modal_export.py | 32 +++++++++++++++++++++++++++----- tests/test_parity_stage_skip.py | 33 +++++++++++++++++++++++++++++++++ 2 files changed, 60 insertions(+), 5 deletions(-) diff --git a/src/gpu/modal_export.py b/src/gpu/modal_export.py index ab19052..34f3af1 100644 --- a/src/gpu/modal_export.py +++ b/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") + ) ] diff --git a/tests/test_parity_stage_skip.py b/tests/test_parity_stage_skip.py index 29f7c49..3cd3789 100644 --- a/tests/test_parity_stage_skip.py +++ b/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"]