File size: 5,691 Bytes
5081377 47382f6 5081377 47382f6 5081377 47382f6 5081377 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 | """Validate the canonical Nemotron-3-Diarization debug sideload archive."""
from __future__ import annotations
import argparse
import hashlib
import json
import posixpath
import zipfile
from pathlib import PurePosixPath
from typing import Any
ADAPTER_ID = "nemotron3-diarization-low-latency-v1"
FUSED_ADAPTER_ID = "nemotron3-diarization-low-latency-v2"
SUPPORTED_ADAPTER_IDS = {ADAPTER_ID, FUSED_ADAPTER_ID}
MANIFEST_PATH = "model-manifest.json"
MAX_ENTRIES = 64
MAX_ENTRY_BYTES = 128 * 1024 * 1024
MAX_TOTAL_BYTES = 256 * 1024 * 1024
def parse_args() -> argparse.Namespace:
"""Parse archive validation arguments."""
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("archive", type=str)
return parser.parse_args()
def _validate_path(path: str) -> None:
"""Reject non-canonical, absolute, or traversing ZIP paths."""
pure = PurePosixPath(path)
if (
not path
or path.startswith("/")
or "\\" in path
or any(part in ("", ".", "..") for part in pure.parts)
or posixpath.normpath(path) != path
):
raise ValueError(f"Unsafe ZIP path: {path!r}")
def validate_archive(path: str) -> dict[str, Any]:
"""Verify structure, limits, manifest identity, and every payload hash."""
with zipfile.ZipFile(path) as archive:
entries = archive.infolist()
names = [entry.filename for entry in entries]
if len(entries) > MAX_ENTRIES:
raise ValueError(f"Too many ZIP entries: {len(entries)}")
if len(names) != len(set(names)):
raise ValueError("Duplicate ZIP entry")
if MANIFEST_PATH not in names:
raise ValueError(f"Missing {MANIFEST_PATH}")
total = 0
for entry in entries:
_validate_path(entry.filename)
if entry.is_dir():
raise ValueError(f"Directory entries are not allowed: {entry.filename}")
if entry.file_size > MAX_ENTRY_BYTES:
raise ValueError(f"ZIP entry is too large: {entry.filename}")
total += entry.file_size
if total > MAX_TOTAL_BYTES:
raise ValueError("ZIP uncompressed size exceeds limit")
unix_type = (entry.external_attr >> 16) & 0o170000
if unix_type not in (0, 0o100000):
raise ValueError(f"Non-regular ZIP entry: {entry.filename}")
manifest = json.loads(archive.read(MANIFEST_PATH))
if (
manifest.get("schema_version") != 1
or manifest.get("adapter_id") not in SUPPORTED_ADAPTER_IDS
):
raise ValueError("Unsupported sideload manifest")
inventory = manifest.get("files")
if not isinstance(inventory, list):
raise ValueError("Manifest files must be a list")
expected = {MANIFEST_PATH}
inventory_paths: set[str] = set()
for item in inventory:
item_path = item.get("path")
if not isinstance(item_path, str):
raise ValueError("Manifest file path must be a string")
_validate_path(item_path)
if item_path in inventory_paths:
raise ValueError(f"Duplicate manifest path: {item_path}")
inventory_paths.add(item_path)
expected.add(item_path)
try:
entry = archive.getinfo(item_path)
except KeyError as error:
raise ValueError(f"Missing inventoried file: {item_path}") from error
content = archive.read(entry)
if entry.file_size != item.get("bytes"):
raise ValueError(f"Size mismatch: {item_path}")
if hashlib.sha256(content).hexdigest() != item.get("sha256"):
raise ValueError(f"SHA-256 mismatch: {item_path}")
if item_path.endswith(".tflite") and content[4:8] != b"TFL3":
raise ValueError(f"Invalid LiteRT flatbuffer: {item_path}")
if set(names) != expected:
raise ValueError(f"Unexpected ZIP entries: {sorted(set(names) ^ expected)}")
graphs = manifest.get("ordered_graphs", [])
graph_paths = [graph.get("path") for graph in graphs]
if manifest["adapter_id"] == FUSED_ADAPTER_ID:
expected_graphs = [
"graphs/feature_stacker.tflite",
"graphs/encoder_head.tflite",
]
else:
expected_graphs = ["graphs/feature_stacker.tflite"]
expected_graphs.extend(
f"graphs/encoder_layer_{index:02d}.tflite" for index in range(31)
)
expected_graphs.append("graphs/classification_head.tflite")
if graph_paths != expected_graphs:
raise ValueError("Unexpected graph count or invocation order")
if manifest.get("total_uncompressed_bytes") != sum(
int(item["bytes"]) for item in inventory
):
raise ValueError("Manifest total_uncompressed_bytes mismatch")
return manifest
def main() -> int:
"""Validate one archive and print its release identity."""
args = parse_args()
manifest = validate_archive(args.archive)
print(
json.dumps(
{
"adapter_id": manifest["adapter_id"],
"evaluation_only": manifest["evaluation_only"],
"graph_count": len(manifest["ordered_graphs"]),
"source_revision": manifest["source_revision"],
"total_uncompressed_bytes": manifest["total_uncompressed_bytes"],
},
indent=2,
sort_keys=True,
)
)
return 0
if __name__ == "__main__":
raise SystemExit(main())
|