Cap zstd frames at 200 MiB to fix JS-decoder overflow on large .eval files#3746
Merged
Conversation
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
The round-trip test revealed that the decompressor path also needed fixing: the previous ZstdDecompressObjWrapper (from zipfile_zstd) used a decompressobj that stopped after the first zstd frame, leaving subsequent frames unread and causing a CRC mismatch on read-back. Add _MultiFrameZstdDecompressObj, which chains fresh inner decompressobj instances across frame boundaries, and patch _get_decompressor alongside _get_compressor so both directions handle multi-frame streams correctly.
Two issues on multi-frame writes: 1. The wrapper's compress() accumulated output with += b"...", which is O(n^2) across the 1 GiB+ compressed output of a 1.38 GiB entry — gigabytes of needless copying. Use a list + b"".join, and a memoryview over the input to avoid slice copies. 2. The factory delegated to zipfile_zstd's _get_compressor, which constructs a fresh ZstdCompressor(threads=12) per compressobj call. Share a single compressor across all frames of an entry. Before: ~9.5 s for the 1.38 GiB entry (~3x baseline). After: ~3.9 s (+25% vs single-frame baseline, matching the inherent zstd multi-frame dictionary-reset overhead).
- Resolve zipfile.ZIP_ZSTANDARD once at import time into _ZIP_ZSTANDARD, collapsing three # type: ignore suppressions into one and surfacing a missing attribute at import time instead of at call time. - Extract the threads count zipfile_zstd hardcodes (12) into _ZSTD_THREADS with a comment pointing at the source of truth, so our value and theirs can't silently drift apart. - Skip the trailing frame flush when the last compress() call landed exactly on a frame boundary; otherwise we appended an empty 9-byte zstd frame to such entries.
- Type the compressor and decompressor factories with the concrete zstandard classes instead of Any, so self._obj.flush() no longer leaks Any into a function declared to return bytes. - In the test file, resolve ZIP_ZSTANDARD via getattr like the existing test_async_zip.py does, since on Python < 3.14 stdlib zipfile does not declare the attribute and type checkers don't see zipfile_zstd's runtime monkey-patch.
rasmusfaber
marked this pull request as ready for review
April 23, 2026 12:51
Collaborator
|
@rasmusfaber, am I understanding this correctly that if the PR for |
epatey
approved these changes
Apr 23, 2026
Contributor
Author
Yes. Fixing |
5 tasks
jjallaire
added a commit
that referenced
this pull request
Apr 28, 2026
* test: extract shared zstd test helpers to test_helpers/zstd.py No behavior change. Moves moderately_compressible_payload and read_raw_compressed_entry out of test_zipfile_multiframe.py so the forthcoming async multi-frame tests can reuse them. * fix: handle multi-frame zstd in async ZIP read path A single zstandard decompressobj only spans one frame; on eof it rejects further input. The previous ZstdDecompressor crashed with 'cannot use a decompressobj multiple times' on the second chunk of any multi-frame entry -- exactly the entries PR #3746's writer produces for samples whose JSON exceeds 200 MiB. Mirror the _MultiFrameZstdDecompressObj pattern from _util/zipfile.py inside ZstdDecompressor: track a _pending buffer for unused_data carryover and replace the inner decompressobj on every frame boundary. Adds an integration test through AsyncZipReader that monkey-patches the per-frame cap to 64 KiB so the round-trip is fast (<1s). Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * fix: handle multi-frame zstd in HTTP transcoding path Same fix as the previous commit, applied to _ZstdDecompressIterator in compression_transcoding.py. This class is used by CompressedToDeflateStream when the viewer requests a sample whose on-disk encoding is zstd: we decompress zstd, recompress to deflate for the browser. Without this fix, multi-frame entries crashed at the second chunk on the same 'cannot use a decompressobj multiple times' error. Adds direct unit tests of the iterator with synthetic multi-frame input. Chunk sizes deliberately don't align to frame boundaries so the unused_data carryover is exercised. * refactor: tighten zstd decompressor invariants from post-review pass Three small post-review cleanups, no behavior change for in-spec input: 1. compression_transcoding.py::_ZstdDecompressIterator: drop the dead `or self._obj is None` half of the entry guard. `_obj` was only ever nulled immediately followed by `_exhausted = True` on the same path, so the two conditions are equivalent — `_exhausted` alone is sufficient. Tighten the `_obj` annotation to non-Optional and drop the now-redundant `self._obj = None` from the StopAsyncIteration handler. 2. compression.py::ZstdDecompressor.decompress_next: add a symmetry guard `if self._exhausted: raise StopAsyncIteration` at the top. Previously, calling decompress_next again after exhaustion (with a fresh source iterator) would silently lazy-reinit and start decoding the new bytes — the `Decompressor` Protocol advertises exhaustion as a one-way door, but the implementation didn't enforce it. Currently dormant in production because CompressedToUncompressedStream gates on `exhausted` first; this closes the latent footgun. 3. tests/test_helpers/zstd.py: annotate `name_len: int` and `extra_len: int` on the struct.unpack results. struct.unpack returns `tuple[Any, ...]`, so without the annotations the Any was leaking into the subsequent f.read() call. * docs: trim over-explanatory comments and docstrings The earlier commits in this branch added more prose than the code warrants. Pare back to one-line docstrings and drop comments that restate what the code already says, including: - The PR-reference paragraph in the integration test's docstring. - The 4-line monkey-patch explanation (kept the WHY one-liner). - The 4-line sanity-check rationale (the assertion message says it). - "Inner frame complete; carry leftover bytes..." (variable names and class context already convey it). - "Use a chunk size that does NOT align..." (chunk size is small/odd by inspection). - The "Used by:" callers list in tests/test_helpers/zstd.py. - Multi-paragraph class docstrings in ZstdDecompressor and _ZstdDecompressIterator. Also restores the pre-existing "Note: Unlike zlib, ..." comment in ZstdDecompressor to its original 3-line form (this branch had extended it). * Update CHANGELOG with latest feature additions Updated CHANGELOG with recent changes and enhancements. --------- Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com> Co-authored-by: jjallaire <jj.allaire@gmail.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR contains:
What is the current behavior?
Each entry in a
.evalZIP archive is written as a single zstd frame. For large samples (e.g. asamples/0_epoch_1.jsonof 1.38 GiB uncompressed / 1.03 GiB compressed), this produces a single multi-hundred-MiB zstd frame. The fzstd decoder used by the viewer cannot read such frames correctly: 101arrowz/fzstd#21What is the new behavior?
inspect_ai/_util/zipfile.pyadditionally monkey-patcheszipfile._get_compressorand_get_decompressorso that everyZIP_ZSTANDARDentry is written as a multi-frame zstd stream, with a new frame started every 200 MiB of input. Multi-frame zstd streams are valid per spec, decode transparently in every compliant decoder, and — because each new frame resets the decoder's state — keep every frame well under the JS decoders' thresholds.A matching
_MultiFrameZstdDecompressObjpatches the decompressor side so Python readers chain across frame boundaries within a single zip entry (zstandard.ZstdDecompressor().decompressobj()stops at the first frame by default).This only fixes future
.evalfiles. Existing files on disk still need the decoder-side fix landing separately in ts-mono's fzstd path.Performance: on the 1.38 GiB
samples/0_epoch_1.jsonfrom a real.eval, ZIP write wall-clock goes from 3.10 s baseline to 3.88 s patched (+25%) with no change in compressed size (1029.9 MB → 1029.9 MB). The regression is intrinsic to multi-frame zstd — each new frame re-initialises FSE/Huffman and match tables, contributing ~23% fixed overhead at 200 MiB frames, confirmed by isolating compression outside the ZIP path. Several non-viable mitigations were explored (copying dictionary state across frames breaks decoder independence;stream_writerwithflush(FLUSH_FRAME)performs identically to ourcompressobjrestart; the zstandard Python bindings don't exposeZSTD_CCtx_reset).Does this PR introduce a breaking change?
No. Multi-frame zstd streams are valid per spec. All
.evalconsumers that useinspect_ai's own reader, stdlibzipfile+zipfile_zstd, or any native zstd decoder are unaffected. Existing single-frame.evalfiles on disk remain readable.Other information
An alternative fix is waiting for fzstd to accept the PR and release a new version. I also investigated zstddec, but that has a similar issue with large frames.