fix: route large and renamed Flax MessagePack checkpoints#1280
Conversation
Recognize bounded, structurally valid Flax/JAX MessagePack payloads when they use misleading suffixes, while leaving generic MessagePack state maps skipped. Add routing, filtering, fail-closed, and malicious payload regression coverage plus user-facing documentation.
Performance BenchmarksCompared
|
mldangelo-oai
changed the title
|
Independent follow-up audit found and repaired a false-clean path in the large renamed MessagePack routing logic. Reproduction before Repair: bounded routing now distinguishes confirmed roots from probe exhaustion. An ambiguous renamed map remains conservatively routed, but Validation:
|
There was a problem hiding this comment.
Pull request overview
This PR improves Flax/JAX MessagePack detection so renamed checkpoints can be routed to the Flax scanner using bounded structural inspection, and ambiguous oversized renamed maps fail closed as inconclusive instead of being treated as clean.
Changes:
- Adds bounded MessagePack probing for renamed Flax/JAX checkpoint roots such as
params,opt_state, andmodel_state. - Updates Flax MessagePack scanning and registry routing to handle renamed or header-detected MessagePack checkpoints.
- Adds regression coverage and user-facing documentation for renamed/large MessagePack checkpoint behavior.
Reviewed changes
Copilot reviewed 11 out of 11 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
modelaudit/utils/file/detection.py |
Adds bounded MessagePack structure probing and renamed Flax checkpoint detection. |
modelaudit/scanners/flax_msgpack_scanner.py |
Routes can_handle through structural detection and adds fail-closed inconclusive handling. |
modelaudit/scanner_registry_metadata.py |
Registers flax_msgpack as a header-routed format. |
tests/utils/file/test_filetype.py |
Adds format-detection regressions for large renamed Flax MessagePack maps. |
tests/utils/file/test_file_filter.py |
Adds skip-filter regression coverage for disguised Flax checkpoints. |
tests/test_directory_file_filtering.py |
Verifies directory scans preserve large disguised malicious Flax MessagePack files. |
tests/test_core.py |
Adds scan, exit-code, dependency, and cache regressions for renamed/ambiguous MessagePack checkpoints. |
tests/scanners/test_flax_msgpack_scanner.py |
Adds scanner-level routing and no-unpack fail-closed regressions. |
README.md |
Documents renamed Flax/JAX MessagePack recognition. |
docs/user/compatibility-matrix.md |
Documents compatibility behavior for renamed Flax/JAX MessagePack maps. |
CHANGELOG.md |
Records the bug fix. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
|
@codex review |
|
Codex Review: Didn't find any major issues. Hooray! ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
# Conflicts: # docs/user/compatibility-matrix.md # modelaudit/utils/file/detection.py # tests/test_core.py # tests/test_directory_file_filtering.py # tests/utils/file/test_file_filter.py # tests/utils/file/test_filetype.py
# Conflicts: # modelaudit/core.py # modelaudit/scanners/archive_dispatch.py # tests/utils/file/test_filetype.py
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 4b69a72d31
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 5bfe75091b
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
|
Follow-up review remediation is now in #1379.
Validation on the follow-up branch includes independent review, focused routing/scanner/core tests ( |
Summary
flax_msgpackscanner when bounded structural inspection exposes a checkpoint-state root (params,opt_state, ormodel_state)Why
A malicious MessagePack checkpoint saved as
payload.jpgwas silently skipped. During adversarial review, an otherwise identical renamed map larger than 1 MiB with a harmlessmetadatavalue beforeparamswas also skipped, while.flaxreported the existing critical__reduce__finding.A further audit found a second failure mode in the draft: a generic 1.1 MiB MessagePack map with no checkpoint-state root could exhaust the bounded routing walk, be routed as
flax_msgpack, fully unpacked, and return a clean result with exit code0.The repaired path reaches later top-level checkpoint roots without materializing large scalar values. If a renamed structure cannot be classified within the routing budget, it is still retained for safety but now returns
inconclusivewith reasonflax_msgpack_routing_probe_limit_exceeded, is not cached, and yields exit code2withoutmsgpack.unpackb.Validation
env VIRTUAL_ENV=/Users/mdangelo/code/modelaudit/.venv UV_CACHE_DIR=/tmp/modelaudit-uv-cache uv run --active --no-sync ruff format modelaudit/ packages/modelaudit-picklescan/src packages/modelaudit-picklescan/tests tests/env VIRTUAL_ENV=/Users/mdangelo/code/modelaudit/.venv UV_CACHE_DIR=/tmp/modelaudit-uv-cache uv run --active --no-sync ruff check --fix modelaudit/ packages/modelaudit-picklescan/src packages/modelaudit-picklescan/tests tests/env VIRTUAL_ENV=/Users/mdangelo/code/modelaudit/.venv UV_CACHE_DIR=/tmp/modelaudit-uv-cache uv run --active --no-sync mypy modelaudit/ packages/modelaudit-picklescan/src packages/modelaudit-picklescan/tests tests/(Success: no issues found in 451 source files)env npm_config_cache=/private/tmp/modelaudit-npm-cache npx prettier --check CHANGELOG.md README.md docs/user/compatibility-matrix.md327 passed, 3 warningsflax_msgpack,success=False,scan_outcome=inconclusive, reasonflax_msgpack_routing_probe_limit_exceeded, exit code2env PYTHONPATH=/private/tmp/modelaudit-flax-routing VIRTUAL_ENV=/Users/mdangelo/code/modelaudit/.venv UV_CACHE_DIR=/tmp/modelaudit-uv-cache PROMPTFOO_DISABLE_TELEMETRY=1 uv run --active --no-sync pytest -n auto -m "not slow and not integration" --maxfail=1(5705 passed, 16 skipped, 21 warnings)