ASV better memory benchmarks - #1609
Open
cmdupuis3 wants to merge 5 commits into
Open
Conversation
Three benchmarks reported a near-identical "improvement" on every PR regardless
of what the PR touched:
578M -> 391M 0.68 face_bounds.FaceBounds.peakmem_face_bounds(geoflow-small)
708M -> 390M 0.55 face_bounds.FaceBounds.peakmem_face_bounds(quad-hexagon)
498M -> 384M 0.77 mpas_ocean.Gradient.peakmem_gradient('480km')
They were not measuring the operation under test. asv's peakmem_* records the
max RSS of the whole process, and per asv's docs it "also counts memory usage
during the setup routine". Profiling the face_bounds params:
grid import +open_grid +.bounds attributable to op
quad-hexagon ( 24K) 236MB 265MB 301MB 36MB (12%)
geoflow-small (1.1M) 236MB 263MB 487MB 224MB (46%)
outCSne8 ( 48K) 235MB 281MB 317MB 35MB (11%)
oQU480 (4.6M) 236MB 270MB 306MB 36MB (12%)
`import uxarray` alone is ~226 MB and constant to within 1 MB. oQU480 is 190x
larger than quad-hexagon yet both attributed ~36 MB to Grid.bounds -- the
benchmark was nearly insensitive to its own workload. The part that did vary was
numba: uxarray has 81 @njit(cache=True) kernels and both affected paths go
through them (uxarray/grid/bounds.py, uxarray/core/gradient.py). With a cold JIT
cache the quad-hexagon case peaked at 599 MB, with a warm one 298 MB -- a ratio
of 0.50, matching the ratios seen in CI. Which side of an `asv continuous`
comparison paid the compile cost depended on run order, not on the code under
review.
peakmem_* could not be repaired in place, because there was nothing to measure.
Across every operation the five peakmem benchmarks covered, against every mesh
in the suite including the 98 MB / 28,571-face oQU120:
Grid.bounds ~1 MB open_grid 0-10 MB (lazy; mostly allocator noise)
gradient 0-0.1 MB integrate 0.0 MB
open_dataset 0.0 MB
Two better instruments were evaluated and rejected:
- Sampled peak RSS. Validates cleanly (a known 200 MB allocation measures as
200.0 MB) and is immune to process history. But the warm-up call needed to
keep JIT out of the measurement leaves freed pages in the allocator, so RSS
*growth* undercounts -- every operation measured 0.0-0.1 MB this way.
- tracemalloc. Measures allocation volume rather than RSS growth, so it would
sidestep page reuse, but it does not observe numba NRT allocations, which is
where uxarray's array memory is allocated.
What remains is deterministic size accounting via track_* benchmarks, which also
sidesteps the setup confound entirely -- track_* does not count setup. Verified
byte-identical across cold JIT, warm JIT, and an independent second cold JIT for
all 14 parameter combinations, and it scales correctly with the mesh
(quad-hexagon reports 128 bytes = 4 faces x 32). It does not capture transient
peaks, but the measurements above show those are ~1 MB, far below anything worth
gating on, and no available instrument captures them reliably here.
The rationale is recorded in benchmarks/_memsize.py so the next person does not
reintroduce peakmem_*. The commented-out mem_* stubs in quad_hexagon.py, an
earlier attempt at the same thing, are removed.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
The previous commit removed peakmem_* because the per-operation memory it
claimed to measure was ~1 MB against a reported 400-700 MB. But the larger
question those benchmarks were reaching for -- "how much memory does it take to
open this grid and compute on it, from a cold start" -- is real, and asv can
answer it with its own peakmem_* once two things are controlled:
1. Process history. ru_maxrss is a high-water mark that never falls, so
imports, setup() and earlier work in the same process set a floor under the
result. asv already spawns a separate process per benchmark *and per
parameter* (runner._run_benchmark_single_param), so this is handled as long
as the benchmark classes declare no setup() -- asv counts setup memory
towards peakmem_*, which is exactly the confound its own docs warn about.
2. numba JIT cache warmth. Compiling uxarray's 81 @njit(cache=True) kernels
costs a few hundred MB of transient RSS, so an otherwise identical run
peaked at 599 MB cold and 298 MB warm. Whichever side of an `asv continuous`
comparison happened to compile paid that cost.
setup_cache handles (2). asv runs it in its own process
(runner.Spawner.create_setup_cache), so LLVM's footprint never enters the
high-water mark of the processes that do the measuring; they load the compiled
kernels from numba's on-disk cache instead. It runs once per commit, so both
sides of a comparison are warmed symmetrically. Every parameter is warmed, not
just one -- the grids do not all reach the same njit signatures.
Measured via `asv run --python=same --quick -b peakmem`, twice warm and once
after deleting every .nbi/.nbc in the installed uxarray:
warm warm COLD spread
import uxarray 280M 282M 280M 0.7%
open+bounds quad-hexagon 349M 349M 346M 0.9%
open+bounds geoflow-small 348M 348M 348M 0.0%
open+bounds outCSne8 365M 364M 370M 1.6%
open+bounds oQU480 354M 356M 351M 1.4%
gradient 480km 358M 355M 354M 1.1%
gradient 120km 371M 370M 381M 2.9%
The cold column is the condition that used to halve the result; the cache
repopulated from 0 to 33 entries during that run, confirming setup_cache did the
compiling. Everything sits inside 2.9%, against asv's default 10% factor.
peakmem_import_uxarray is included deliberately. ~280 MB of every row above is
just importing uxarray, so tracking it on its own means a heavy new top-level
import shows up as itself instead of silently inflating everything else.
Also moves _memsize into benchmarks/helpers/ and gives it an __init__.py, so the
package structure is explicit rather than relying on namespace packages.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
cmdupuis3
added
bug
ASV BenchmarkingBenchmark Comparison ResultsBenchmarks that have stayed the same:
Benchmarks that have got worse:
|
Collaborator
|
Thank you for working on this and opening up this PR! I have a few quick thoughts before reviewing:
|
Collaborator
Author
|
@Sevans711 as per convo, 1 is really a question for @erogluorhan or @rajeeja. On 2, these are included because the peakmem is not necessarily the same as the total memory usage in a chunked situation. This particular batch of benchmarks is a little verbose, but it basically serves as a regression test against the kinds of spurious improvements in other PRs, so that we have a better idea of which things are bloating the peakmem metrics. |
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
3 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.
Closes #1605
Overview
The current behavior of a few peakmem benchmarks entrains module imports and caching, which dwarf the actual memory usage supposedly being tested. This PR offers benchmarks of low-level memory behavior as well as warm setups to prevent the persistent "every PR improves memory by the same amount" issue in #1605.
Expected Usage
Most of these benchmarks are useful as-is, but there are also a couple of benchmarks of pure uxarray imports to compare with if the illusory memory improvement ever returns.
PR Checklist
General
Testing
Documentation