From eb23c70c3ab5a0baf7bb321362f5afa9baad68ba Mon Sep 17 00:00:00 2001
From: Ronald Tse <ronald.tse@gmail.com>
Date: Sun, 30 Aug 2026 22:35:37 +0200
Subject: [PATCH] fix(export): refresh the member sha table after in-place
 member replacement
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

rebuild_int8_head32 rewrites the encoder/decoder graphs inside a copy
of the shipped int8 zip but left metadata.yaml's sha256 table pointing
at the old digests — write_parity validates strictly before writing,
so every rebuild died at exactly that point (parity itself passed).
refresh_member_shas recomputes the table and rewrites the member
atomically; the rebuild calls it right after assembling the zip.
---
 src/gpu/modal_export.py  | 11 ++++++++++-
 src/imf/export.py        | 33 +++++++++++++++++++++++++++++++++
 tests/test_imf_export.py | 34 ++++++++++++++++++++++++++++++++++
 3 files changed, 77 insertions(+), 1 deletion(-)

diff --git a/src/gpu/modal_export.py b/src/gpu/modal_export.py
index 988a686..26bc8bc 100644
--- a/src/gpu/modal_export.py
+++ b/src/gpu/modal_export.py
@@ -433,7 +433,12 @@ def rebuild_int8_head32(model_id: str, limit: int = 0) -> dict:
     checkpoint = Path(spec["volume"]) / spec["checkpoint"]
     test_path = Path(spec["test_volume"]) / spec["test_data"]
 
-    from imf.export import head_matmul_names, load_byte_seq2seq, quantize_int8
+    from imf.export import (
+        head_matmul_names,
+        load_byte_seq2seq,
+        quantize_int8,
+        refresh_member_shas,
+    )
     from imf.parity import (
         reference_decode,
         run_margin_analysis,
@@ -478,6 +483,10 @@ def rebuild_int8_head32(model_id: str, limit: int = 0) -> dict:
                 else:
                     dst.writestr(name, src.read(name))
 
+    # re-quantized graphs replaced the members; the internal sha table
+    # must be refreshed or strict validation (write_parity) rejects the zip
+    refresh_member_shas(new_zip)
+
     reference = reference_decode(
         model,
         [s for s, _ in pairs],
diff --git a/src/imf/export.py b/src/imf/export.py
index 2c85130..504d738 100644
--- a/src/imf/export.py
+++ b/src/imf/export.py
@@ -263,6 +263,39 @@ def convert_fp16(model):
     return copy.deepcopy(model).half()
 
 
+def refresh_member_shas(zip_path: Path | str) -> None:
+    """Recompute the metadata.yaml sha256 table after members were
+    replaced inside a zip (the head32 re-quantization rewrites the
+    encoder/decoder graphs in place); the member set is unchanged, so
+    only the digests drift — until they are refreshed, strict zip
+    validation rejects the artifact (write_parity gates on it)."""
+
+    import hashlib
+    import os
+    import tempfile
+    import zipfile
+
+    import yaml
+
+    zip_path = Path(zip_path)
+    with zipfile.ZipFile(zip_path) as zf:
+        meta = yaml.safe_load(zf.read("metadata.yaml"))
+        for name in list(meta["sha256"]):
+            meta["sha256"][name] = hashlib.sha256(zf.read(name)).hexdigest()
+
+    with tempfile.TemporaryDirectory(dir=zip_path.parent) as tmp:
+        rewritten = Path(tmp) / "rewritten.zip"
+        with zipfile.ZipFile(zip_path) as src, zipfile.ZipFile(
+            rewritten, "w", zipfile.ZIP_DEFLATED
+        ) as dst:
+            for member in src.namelist():
+                if member == "metadata.yaml":
+                    dst.writestr(member, yaml.safe_dump(meta, sort_keys=False))
+                else:
+                    dst.writestr(member, src.read(member))
+        os.replace(rewritten, zip_path)
+
+
 def quantize_int8(
     src: Path | str, dst: Path | str, per_channel: bool = False,
     nodes_to_exclude: list[str] | None = None,
diff --git a/tests/test_imf_export.py b/tests/test_imf_export.py
index 614c60e..bc2bfc2 100644
--- a/tests/test_imf_export.py
+++ b/tests/test_imf_export.py
@@ -155,3 +155,37 @@ def test_int8_keeps_head_matmul_in_fp32(zips: dict[str, Path]) -> None:
             assert head in nodes, (graph_name, head)
             # converted heads become MatMulInteger with the same name
             assert nodes[head].op_type == "MatMul", (graph_name, head)
+
+
+def test_refresh_member_shas_after_member_replacement(zips: dict[str, Path]) -> None:
+    """The head32 rebuild replaces graphs inside an existing zip; the
+    stale sha table must make validation fail, and the refresh must
+    restore it without touching any other member."""
+    import zipfile
+
+    from imf.export import refresh_member_shas
+
+    path = zips["fixture-1.0-int8.zip"]
+    with zipfile.ZipFile(path) as src_zf:
+        members = {n: src_zf.read(n) for n in src_zf.namelist()}
+    import onnx
+
+    graph = onnx.load_from_string(members["encoder.onnx"])
+    graph.graph.name += "-tampered"
+    tampered = graph.SerializeToString()
+    members["encoder.onnx"] = tampered
+    tmp = path.with_suffix(".rewritten")
+    with zipfile.ZipFile(tmp, "w", zipfile.ZIP_DEFLATED) as dst:
+        for n, data in members.items():
+            dst.writestr(n, data)
+    tmp.replace(path)
+
+    result = validate_zip(path)
+    assert not result.ok  # stale table rejects the replaced member
+
+    refresh_member_shas(path)
+    result = validate_zip(path)
+    assert result.ok, result.errors
+    with zipfile.ZipFile(path) as after_zf:
+        assert after_zf.read("encoder.onnx") == tampered
+        assert after_zf.read("decoder.onnx") == members["decoder.onnx"]