diff --git a/core/builder/tests/test_bundle_reader.py b/core/builder/tests/test_bundle_reader.py new file mode 100644 index 0000000000..1a51e318dd --- /dev/null +++ b/core/builder/tests/test_bundle_reader.py @@ -0,0 +1,198 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import json +import struct +from pathlib import Path +from typing import Any + +import pytest + +from tensorrt_model_connect.bundle_writer import BUNDLE_MAGIC, _MAX_HEADER_SIZE + + +def _detect_duplicate_keys(pairs: list[tuple[str, Any]]) -> dict[str, Any]: + seen: dict[str, Any] = {} + for key, value in pairs: + if key in seen: + raise ValueError(f"Duplicate section name: {key!r}") + seen[key] = value + return seen + + +class BundleReader: + """Read named sections from a bundle produced by BundleWriter.""" + + def __init__(self, path: str | Path) -> None: + self._path = Path(path) + with self._path.open("rb") as f: + magic = f.read(len(BUNDLE_MAGIC)) + if magic != BUNDLE_MAGIC: + raise ValueError( + f"Invalid bundle magic signature: expected {BUNDLE_MAGIC!r}, got {magic!r}" + ) + (header_size,) = struct.unpack(" _MAX_HEADER_SIZE: + raise ValueError("bundle header exceeds the 100 MiB runtime limit") + header_bytes = f.read(header_size) + self._data_start = len(BUNDLE_MAGIC) + 8 + header_size + file_size = self._path.stat().st_size + self._data_size = file_size - self._data_start + + try: + header: dict[str, Any] = json.loads( + header_bytes, object_pairs_hook=_detect_duplicate_keys + ) + except json.JSONDecodeError as exc: + raise ValueError(f"bundle header is not valid JSON: {exc}") from exc + + if "sections" not in header: + raise ValueError("bundle header missing 'sections' key") + raw_sections = header["sections"] + if not isinstance(raw_sections, dict): + raise ValueError("bundle header 'sections' must be a JSON object") + + sections: dict[str, tuple[int, int]] = {} + for name, entry in raw_sections.items(): + offset = entry.get("offset") + length = entry.get("length") + if not isinstance(offset, int) or not isinstance(length, int): + raise ValueError( + f"section {name!r}: offset and length must be integers" + ) + if offset < 0 or length < 0: + raise ValueError( + f"section {name!r}: offset and length must be non-negative" + ) + if offset + length > self._data_size: + raise ValueError( + f"section {name!r}: goes out-of-file range " + f"(offset={offset}, length={length}, data_size={self._data_size})" + ) + sections[name] = (offset, length) + + # Overlap check: sort by offset and verify no two sections overlap. + sorted_sections = sorted(sections.items(), key=lambda kv: kv[1][0]) + for i in range(len(sorted_sections) - 1): + name_a, (off_a, len_a) = sorted_sections[i] + name_b, (off_b, _) = sorted_sections[i + 1] + if off_a + len_a > off_b: + raise ValueError( + f"Sections overlap: {name_a!r} ends at {off_a + len_a}, " + f"but {name_b!r} starts at {off_b}" + ) + + self._sections = sections + + def read_section(self, name: str) -> bytes: + """Read and return the raw bytes of a named section.""" + + if name not in self._sections: + raise KeyError(f"bundle has no section named {name!r}") + offset, length = self._sections[name] + with self._path.open("rb") as f: + f.seek(self._data_start + offset) + return f.read(length) + + + +def create_bundle(path, magic, header, sections_data): + header_str = json.dumps(header).encode("utf-8") + with open(path, "wb") as f: + f.write(magic) + f.write(struct.pack(" None: "core/builder/tests/__init__.py", "core/builder/tests/test_build.py", "core/builder/tests/test_build_cli.py", + "core/builder/tests/test_bundle_reader.py", "core/builder/tests/test_bundle_writer.py", "core/builder/tests/test_byok.py", "core/builder/tests/test_graph_transform.py",