Skip to content

test: guard the legacy pickle scan of spans above INT_MAX - #25

Draft
leehack wants to merge 1 commit into
mainfrom
fix/legacy-pickle-span
Draft

leehack wants to merge 1 commit into
mainfrom
fix/legacy-pickle-span

Conversation

@leehack

@leehack leehack commented Oct 21, 2026

Copy link
Copy Markdown
Owner

Adds a model-free regression control for the legacy checkpoint scanner defect in #23. It does not fix the defect yet: the fix belongs in upstream pickle_io, and this repo consumes upstream only through the submodule pin (AGENTS.md). Refs #23.

Root cause

skip_pickle_object scans GLOBAL, INST, PERSID and line-terminated operands with find_char(p, (int)(end - p), '\n'). read_torch_legacy_file passes it the rest of the whole checkpoint, tensor storage included. Above INT_MAX the cast goes negative and the scan finds nothing. Above UINT32_MAX it keeps only the low 32 bits, so short lines can pass while longer ones fail. Upstream master (0518b34) has the same code.

Change

tests/test_pickle_span.py compiles the pinned pickle_io.cpp with tests/native/pickle_span_test.cpp and checks:

Case Spans Pinned result
Valid GLOBAL, INST, PERSID, UNICODE, BINUNICODE pickles exact size, plain and ASan pass
Truncated or malformed pickles (missing lines, length past end, unknown opcode, no STOP) exact size, plain and ASan rejected (pass)
Malformed pickles followed by storage bytes 2^31 + 4096 and 2^32 + 16, read-only mappings rejected (pass)
The same valid pickles 2^31 + 4096 and 2^32 + 16 expected failure (#23)

When an upstream bump fixes the scanner, the last row becomes an unexpected success and fails the run, so the marker has to be removed then. The large spans are read-only anonymous mappings, so they need no memory or data file.

Upstream fix (not in this PR)

A local patch makes model_io::find_char take a size_t length, capped at INT_MAX, and drops the 13 (int)(end - p) casts in pickle_io.cpp. It applies cleanly to both the pin and upstream master. Checked locally on macOS arm64; it has not been submitted upstream.

  • With the patch applied to the submodule, this suite reports the unexpected success.
  • A 2,612,843,011-byte legacy .pth written from SD-Turbo F16 GGUF weights is rejected by the pinned runtime (unsupported PyTorch file). The patched runtime loads it through new_sd_ctx in 1.6 s.
  • The patched runtime generates identical 128x128 pixels from the GGUF, the legacy .pth, and flat and nested ZIP checkpoints (CPU, four threads, one step, seed 42).

Validation

  • python3 -m unittest discover -s tests: 53 tests OK, 6 skipped, 1 expected failure (macOS arm64, Apple clang).
  • tests.test_pickle_span with CXX=g++-16: OK, 1 expected failure.
  • Not run here: Linux. Hosted CI covers it.

Scope

Test only: no src/, export, submodule, artifact or release change. The large ZIP-checkpoint timeout in #24 is separate and untouched.

The pinned stable-diffusion.cpp narrows the remaining span to int before
scanning a pickle operand for its newline, and read_torch_legacy_file passes
the rest of the whole checkpoint, so a valid legacy checkpoint above 2 GiB is
rejected (#23). Compile the pinned pickle_io.cpp and check valid and malformed
pickles in exact-size buffers and in read-only spans above INT_MAX and
UINT32_MAX. The valid large-span case is an expected failure until an upstream
fix is pinned; a bump that fixes it fails the run until the marker goes.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant