Handle multi-frame zstd in async ZIP read paths#3771
Merged
jjallaire merged 7 commits intoApr 28, 2026
Conversation
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.
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 UKGovernmentBEIS#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>
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.
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.
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).
rasmusfaber
marked this pull request as ready for review
April 27, 2026 09:35
Updated CHANGELOG with recent changes and enhancements.
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? (You can also link to an open issue here)
#3746 made the writer cap each
ZIP_ZSTANDARDzstd frame at 200 MiB of input, producing multi-frame zstd streams for large eval samples. The synchronouszipfileread path was patched (_MultiFrameZstdDecompressObjinsrc/inspect_ai/_util/zipfile.py), but the async read paths were not.Two async classes still used a single
zstandard.ZstdDecompressionObjand either crashed (ZstdError: cannot use a decompressobj multiple times) or silently dropped trailing frames on multi-frame input:compression.py::ZstdDecompressor- used byAsyncZipReaderviaCompressedToUncompressedStream. Production failure observed whenread_log_samplehit a sample whose JSON exceeded 200 MiB:compression_transcoding.py::_ZstdDecompressIterator- used byCompressedToDeflateStreamto transcode zstd → deflate for browsers (which supportContent-Encoding: deflatebut not zstd). Same bug pattern.What is the new behavior?
Both async classes now mirror the
_MultiFrameZstdDecompressObjpattern that PR #3746 already established for the sync path: track a_pendingbyte buffer forunused_datacarryover, and replace the inner decompressobj on every frame boundary.Does this PR introduce a breaking change? (What changes might users need to make in their application due to this PR?)
No. Public API and behavior on single-frame zstd entries are unchanged.