diff --git a/src/inspect_ai/_util/zipfile.py b/src/inspect_ai/_util/zipfile.py index 6215059d72..fe3dc51a21 100644 --- a/src/inspect_ai/_util/zipfile.py +++ b/src/inspect_ai/_util/zipfile.py @@ -1,7 +1,26 @@ +"""ZIP file helpers and monkey-patches for zstd support. + +On Python < 3.14 we import ``zipfile_zstd`` to monkey-patch zstandard +compression into the stdlib ``zipfile`` module. Python 3.14+ handles zstd +natively via stdlib. + +Additionally, we install a second monkey-patch on ``zipfile._get_compressor`` +that caps each emitted zstd frame at ``_MAX_INPUT_PER_FRAME`` bytes of input. +Large single zstd frames (>= 256 MiB compressed) trigger an overflow bug in +the pure-JS ``fzstd`` decoder used by our viewers; multi-framing keeps every +frame well under that threshold. +""" + +from __future__ import annotations + import logging import sys import zipfile -from typing import Any +from collections.abc import Callable +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + import zstandard logger = logging.getLogger(__name__) @@ -9,9 +28,156 @@ if sys.version_info < (3, 14): import zipfile_zstd # type: ignore[import-not-found, import-untyped] # noqa: F401 +# Resolve once so the rest of the module can use a plain int; also fails loudly +# at import time rather than at first call if the attribute is ever missing. +_ZIP_ZSTANDARD: int = zipfile.ZIP_ZSTANDARD # type: ignore[attr-defined] + zipfile_compress_kwargs: dict[str, Any] = { - "compression": zipfile.ZIP_ZSTANDARD, # type: ignore[attr-defined] + "compression": _ZIP_ZSTANDARD, "compresslevel": None, } + +# 200 MiB. Well under fzstd's 256 MiB (2^28) compressed-frame overflow +# threshold, applied to *input* bytes which compress smaller. +_MAX_INPUT_PER_FRAME = 200 * 1024 * 1024 + +# Matches zipfile_zstd's hardcoded default so multi-frame and single-frame zip +# writes use the same thread count. If zipfile_zstd ever changes its default, +# update here too. +_ZSTD_THREADS = 12 + + +class _MultiFrameZstdCompressObj: + """A zstd compressobj that chunks its output into multiple frames. + + Wraps a ``zstandard`` compressobj and flushes it (finalizing the current + frame) and replaces it (starting a new frame) every + ``_MAX_INPUT_PER_FRAME`` bytes of input. Multi-frame zstd streams are + valid per spec -- any compliant decoder reads them transparently. + """ + + def __init__(self, factory: Callable[[], zstandard.ZstdCompressionObj]) -> None: + self._factory = factory + self._obj = factory() + self._input_bytes = 0 + + def compress(self, data: bytes) -> bytes: + view = memoryview(data) + pieces: list[bytes] = [] + offset = 0 + n = len(view) + while offset < n: + remaining_cap = _MAX_INPUT_PER_FRAME - self._input_bytes + end = min(offset + remaining_cap, n) + chunk = view[offset:end] + pieces.append(self._obj.compress(chunk)) + self._input_bytes += end - offset + offset = end + if self._input_bytes >= _MAX_INPUT_PER_FRAME: + pieces.append(self._obj.flush()) + self._obj = self._factory() + self._input_bytes = 0 + return b"".join(pieces) + + def flush(self) -> bytes: + # If the last ``compress()`` call landed exactly on the frame boundary, + # ``self._obj`` was replaced with a fresh compressobj that has received + # no bytes. Flushing it would append an empty 9-byte trailing frame. + if self._input_bytes == 0: + return b"" + return self._obj.flush() + + +class _MultiFrameZstdDecompressObj: + """A zstd decompressobj that transparently spans multiple frames. + + ``zstandard.ZstdDecompressor().decompressobj()`` stops at the first frame + boundary and marks ``eof=True``. When the compressor splits large entries + into multiple frames (see ``_MultiFrameZstdCompressObj``), the reader must + recognise the frame boundary, start a fresh inner decompressobj, and + continue until all compressed bytes have been consumed. + + The stdlib ``zipfile._read1`` path for non-deflate compression does: + + data = self._decompressor.decompress(data) + self._eof = self._decompressor.eof or self._compress_left <= 0 + + We return ``eof=False`` always so that ``self._eof`` is driven purely by + ``self._compress_left <= 0`` (all compressed bytes fed). Meanwhile, + ``decompress`` buffers leftover bytes from a completed frame and feeds them + into the next inner decompressobj. + """ + + def __init__(self) -> None: + import zstandard as zstd # local import -- already a hard dep + + self._dctx: zstandard.ZstdDecompressor = zstd.ZstdDecompressor() + self._obj: zstandard.ZstdDecompressionObj = self._dctx.decompressobj() + self._pending: bytes = b"" + + def decompress(self, data: bytes) -> bytes: + self._pending += data + out = b"" + while self._pending: + result = self._obj.decompress(self._pending) + out += result + if self._obj.eof: + # Inner frame complete; unused_data holds the start of the next. + self._pending = self._obj.unused_data + self._obj = self._dctx.decompressobj() + else: + # All pending input consumed; wait for more. + self._pending = b"" + break + return out + + def flush(self) -> bytes: + return b"" + + @property + def eof(self) -> bool: + # Always False: let compress_left drive the outer EOF check. + return False + + +def _install_multiframe_patches() -> None: + """Install multi-frame zstd compressor and decompressor patches. + + Idempotent. Wraps whatever ``_get_compressor`` / ``_get_decompressor`` were + installed before us (stdlib on Py >= 3.14; ``zipfile_zstd``'s versions on + Py < 3.14), so compression level and thread count are preserved. + """ + if getattr(zipfile, "_inspect_ai_multiframe_installed", False): + return + + original_compressor = zipfile._get_compressor # type: ignore[attr-defined] + original_decompressor = zipfile._get_decompressor # type: ignore[attr-defined] + + def patched_compressor(compress_type: int, compresslevel: int | None = None) -> Any: + if compress_type == _ZIP_ZSTANDARD: + # Share one ``ZstdCompressor`` across all frames of the entry. + # Delegating to ``original_compressor`` would instead create a + # fresh ``ZstdCompressor(threads=N)`` per frame, re-initialising + # its thread pool every 200 MiB. + import zstandard + + level = 3 if compresslevel is None else compresslevel + compressor = zstandard.ZstdCompressor(level=level, threads=_ZSTD_THREADS) + return _MultiFrameZstdCompressObj(compressor.compressobj) + return original_compressor(compress_type, compresslevel) + + def patched_decompressor(compress_type: int) -> Any: + if compress_type == _ZIP_ZSTANDARD: + return _MultiFrameZstdDecompressObj() + return original_decompressor(compress_type) + + zipfile._get_compressor = patched_compressor # type: ignore[attr-defined] + zipfile._get_decompressor = patched_decompressor # type: ignore[attr-defined] + zipfile._inspect_ai_multiframe_installed = True # type: ignore[attr-defined] + + +_install_multiframe_patches() + + __all__ = ["zipfile_compress_kwargs"] diff --git a/tests/util/test_zipfile_multiframe.py b/tests/util/test_zipfile_multiframe.py new file mode 100644 index 0000000000..70cbcfe157 --- /dev/null +++ b/tests/util/test_zipfile_multiframe.py @@ -0,0 +1,135 @@ +"""Tests for multi-frame zstd compression of ZIP entries. + +Producer-side fix: every ZIP_ZSTANDARD entry should be a multi-frame zstd +stream, capped at 200 MiB of input per frame, so JS decoders with per-frame +size limits (fzstd @ 256 MiB) can decode large inspect_ai .eval files. +""" + +from __future__ import annotations + +import zipfile +from pathlib import Path + +import pytest +import zstandard + +# Importing this module installs the zstd compression patches (both the +# zipfile_zstd delegation on Python < 3.14 and our multi-frame wrapper). +import inspect_ai._util.zipfile # noqa: F401 + +MAX_INPUT_PER_FRAME = 200 * 1024 * 1024 # must match the value in _util/zipfile.py +ZSTD_MAGIC = b"\x28\xb5\x2f\xfd" + +# zipfile_zstd adds ZIP_ZSTANDARD (93) to zipfile at import time; use getattr +# so type checkers on Python < 3.14 (where stdlib zipfile doesn't declare it) +# stay happy. +ZIP_ZSTANDARD: int = getattr(zipfile, "ZIP_ZSTANDARD", 93) + + +def _moderately_compressible_payload(total_bytes: int) -> bytes: + """Build a deterministic payload that compresses at a realistic ratio. + + Pure random data bypasses zstd's match-finder (compressor emits raw + blocks); pure-constant data collapses to tiny output. We want something + in between so the compressed output is a non-trivial fraction of input — + close to the JSON-like eval workload we care about. + """ + fragments = [ + b'{"role": "assistant", "content": "The quick brown fox jumps over the lazy dog."}\n', + b'{"role": "user", "content": "What is the capital of France? Please answer in one sentence."}\n', + b'{"tool_calls": [{"name": "search", "arguments": {"query": "weather in paris"}}]}\n', + b'{"metadata": {"timestamp": "2026-04-23T10:00:00Z", "model": "claude-opus-4-7"}}\n', + ] + chunk = b"".join(fragments) + n = (total_bytes // len(chunk)) + 1 + return (chunk * n)[:total_bytes] + + +@pytest.fixture(scope="module") +def large_payload() -> bytes: + """~450 MiB of moderately compressible JSON-like bytes.""" + return _moderately_compressible_payload(450 * 1024 * 1024) + + +def _read_raw_compressed_entry(zip_path: Path, entry_name: str) -> bytes: + """Return the raw compressed bytes stored for a zip entry (no decompression).""" + with zipfile.ZipFile(zip_path) as zf: + info = zf.getinfo(entry_name) + with open(zip_path, "rb") as f: + f.seek(info.header_offset) + # Skip local file header to reach the data. Using ZipFile's private + # helper is fragile; instead, decode the local header manually. + import struct + + sig_version_etc = f.read(30) + assert sig_version_etc[:4] == b"PK\x03\x04", "not a local file header" + name_len = struct.unpack(" None: + """A ~450 MiB payload should produce ≥3 zstd frames, each ≤ 200 MiB uncompressed.""" + zip_path = tmp_path / "large.zip" + with zipfile.ZipFile(zip_path, "w", compression=ZIP_ZSTANDARD) as zf: + zf.writestr("big.json", large_payload) + + raw = _read_raw_compressed_entry(zip_path, "big.json") + + magic_count = raw.count(ZSTD_MAGIC) + assert magic_count >= 3, ( + f"expected ≥3 zstd frames in a 450 MiB entry, found {magic_count}" + ) + + frame_sizes = [len(out) for _, out in _iter_frames(raw)] + assert len(frame_sizes) == magic_count, ( + f"magic count {magic_count} != frames walked {len(frame_sizes)}" + ) + for i, size in enumerate(frame_sizes): + assert size <= MAX_INPUT_PER_FRAME, ( + f"frame {i} decompressed to {size} bytes, exceeds cap {MAX_INPUT_PER_FRAME}" + ) + + +def test_small_entry_emits_single_frame(tmp_path: Path) -> None: + """A 1 MiB entry should still be exactly one zstd frame.""" + zip_path = tmp_path / "small.zip" + payload = _moderately_compressible_payload(1 * 1024 * 1024) + with zipfile.ZipFile(zip_path, "w", compression=ZIP_ZSTANDARD) as zf: + zf.writestr("small.json", payload) + + raw = _read_raw_compressed_entry(zip_path, "small.json") + assert raw.count(ZSTD_MAGIC) == 1, ( + f"expected exactly 1 frame for small entry, got {raw.count(ZSTD_MAGIC)}" + ) + + +def test_large_entry_round_trip(tmp_path: Path, large_payload: bytes) -> None: + """Writing and reading back a multi-frame entry must preserve bytes exactly.""" + zip_path = tmp_path / "rt.zip" + with zipfile.ZipFile(zip_path, "w", compression=ZIP_ZSTANDARD) as zf: + zf.writestr("big.json", large_payload) + + with zipfile.ZipFile(zip_path) as zf: + got = zf.read("big.json") + + assert got == large_payload, ( + f"round-trip mismatch: input {len(large_payload)} bytes, got {len(got)} bytes" + )