Repository navigation
[QDP] Refuse .npy loads that would exceed host memory (#787) - #1554
Draft
andrewmusselman wants to merge 3 commits into
Draft
andrewmusselman wants to merge 3 commits into
andrewmusselman wants to merge 3 commits into
Conversation
NumpyReader read the whole array before looking at its size, so a large (or corrupt) .npy file could OOM-kill the process. It now parses only the header first, estimates prod(shape) * itemsize with checked arithmetic, and applies a limit before any data is read: - Auto (default): refuse when estimate * 2 > 0.8 * MemAvailable, as proposed in the PR #777 review. If available memory is unknown the load is allowed and a warning is logged (previous behaviour). - QDP_NPY_MAX_BYTES=<bytes|16G|512MiB>: explicit cap on the estimate. - QDP_NPY_MAX_BYTES=off: disable. Invalid values are errors. - NumpyReader::with_memory_limit() sets the policy without touching the environment. The decision is a pure function, check_npy_memory(header, limit, available), so it is tested without allocating. The error names the estimate, the limit, how to override it, and suggests Parquet streaming. A second check compares the file length to the header so truncated files fail before allocation, even with the guard off. MemAvailable comes from a new platform::available_host_memory_bytes() (/proc/meminfo on Linux, None elsewhere). Inside a container that is the host's figure, not the cgroup limit; reviewers may want a cgroup-aware follow-up. The header parser is ours rather than ndarray-npy's because that crate keeps its header module private; it uses py_literal, which ndarray-npy already depends on, so no new crate enters the tree. Part of #787 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Integration tests write tiny .npy files whose header claims a huge shape, e.g. (10^12, 10) float64 (~72.76 TiB), with no body. They check that the reader fails fast on the size check (explicit and auto limits), on the truncation check (guard disabled), and on shape overflow, and that small C- and Fortran-order files still load under a tight limit. QDP_NPY_MAX_BYTES tests mutate the process environment, so they live in their own test binary (tests/numpy_env.rs) with a single #[test] function; no other thread can observe a half-set variable. Part of #787 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Describe the header-based size check on the Readers page: the auto formula, accepted QDP_NPY_MAX_BYTES values, behaviour when MemAvailable is unknown or containerised, the truncation check, and the Parquet streaming alternative. Add a troubleshooting row to Getting Started for the new error message. Closes #787 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.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
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.
Purpose of PR
NumpyReaderread the whole array before looking at its size, so a large or corrupt.npyfile could OOM-kill the process. Following @rich7420's design comment on #787, the reader now parses only the header, estimates the in-memory size, and refuses the load before reading any data when it would not fit. The limit is automatic by default and configurable.Related Issues or PRs
Closes #787 (follow-up from the #777 review; claimed by @ryankert01 in the thread)
Changes Made
qdp-core/src/readers/numpy.rs:py_literal, already in the tree via ndarray-npy);prod(shape) * itemsizewith checked arithmetic;check_npy_memory(header, limit, available);estimate * 2 > 0.8 * MemAvailable; unknown memory allows the load with a warning;QDP_NPY_MAX_BYTES=<bytes|16G|512MiB>oroff; invalid values are errors;NumpyReader::with_memory_limit();qdp-core/src/platform/host_memory.rs:available_host_memory_bytes()(/proc/meminfoon Linux,Noneelsewhere).(10^12, 10)float64) intests/numpy.rs; env-var tests in their own binarytests/numpy_env.rs.docs/qdp/readers.mdsection and a troubleshooting row indocs/qdp/getting-started.md.Testing
cargo test -p qdp-core --test numpy --test numpy_env: 11 + 1 passed (main: 4 innumpy.rs).cargo test -p qdp-core --lib: 112 passed, 3 failed (main 97/3; same 3 GPU tests; 15 new unit tests).cargo test -p qdp-core --no-fail-fast: 215 passed, 145 failed (main 192 passed; the failures are the same no-GPU set).cargo fmt, clippy-D warnings,cargo check -p qdp-python --all-targets: clean.clippy --all-features(needs libtorch). In containers,MemAvailableis the host's figure, not the cgroup limit.Merging
No textual conflicts, but if the #723 PR merges first, two of its new tests in
qdp-core/tests/numpy_edge_cases.rs(truncated and non-NPY files) expect the old parser's wording. This PR's size pre-check rejects both files earlier with clearer messages. Theafter-723/version adds one commit updating those two expectations; both files are still rejected.🤖 Generated with Claude Code
https://claude.ai/code/session_01GuXmXRXdNF43h2FoQcNAZb