ASV better memory benchmarks - #1609
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:
|
|
Thank you for working on this and opening up this PR! I have a few quick thoughts before reviewing:
|
|
@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. |
|
Just noticed, the |
In regards to peak-mem vs static object size (nbytes) though, I think each of them has its own value (the former gives the peak memory used throughout an operation while the latter gives you the static memory size needed for an object, and they should not be considered alternative to each other. |
erogluorhan
left a comment
erogluorhan
left a comment
There was a problem hiding this comment.
This is looking great to me; just a few comments below.
There was a problem hiding this comment.
Can you please clarify why this is not converted to the new way? No setup_cache needed for it (since no njit)?
| def time_integrate(self, resolution): | ||
| self.uxds[data_var].integrate() | ||
|
|
||
| def peakmem_integrate(self, resolution): |
There was a problem hiding this comment.
Why no IntegratePeakMem class defined like others?
|
@erogluorhan Thank you for clarifying these points, that helps, though I do still have a few follow-up thoughts below: For my question (1), that makes sense, it also sounds good to me to have it go into a different PR. Flagging #1647 as the relevant feature request (created after I left my comment above) for things like
Indeed, it makes sense to me that peak memory is a very different measurement from static memory size needed for an object. My understanding is that no other benchmarks in the benchmarking suite were measuring static memory size, and that the original issue reported here doesn't really require any static memory size checks in order to solve it. That's the main reason I was surprised these were added here, and wanted to know more about their purpose.
I believe I understand what you mean. Is this rephrasing correct?:
If that is correct, my suspicion here is still that adding to the pytest Assuming I understand correctly, those are my last reservations about it. If the overall preference here is still to include these as benchmarks instead of putting them in the test suite, I will get on board with that decision and review with that in mind! |
|
@Sevans711 I think I'm seeing what you're getting at. These benchmarks don't really track how a chunked peakmem metric would actually behave, they basically just measure the total memory consumption. For now though, I think it's probably okay, because we're kind of accepting that the whole grid is going to be materialized elsewhere anyway, and until we have a strong incentive to get an end-to-end chunked grid implementation, the stakes here are low. That said, I think we might be able to get something like a true peakmem metric with something like this: import tracemalloc
__all__ = ["peak_allocated"]
def peak_allocated(build):
tracemalloc.start()
try:
tracemalloc.reset_peak()
build()
_, peak = tracemalloc.get_traced_memory()
finally:
tracemalloc.stop()
return peak
...
class FaceAreas(GridBenchmark):
def track_peakmem_face_areas(self, resolution):
return peak_allocated(lambda: self.uxgrid.face_areas) |
Let @cmdupuis3 answer all your questions about if static memory measurements needed as part of this PR's work or whether now desired in benchmarking, or if they should into pytest, etc. but I'll just share my thoughts on a few things below:
Yeah, Also, we had ASV always run with PRs until some time when we decided not every PR needed it (and also the above speed cost IIRC). For the time being, drawing a line between performance and unit testing, and leaving the former to ASV and mostly @cmdupuis3 's efforts while leaving the latter to We can always decide to get ASV to always run again if that's needed and if the speed cost is worth it. |
Wait, actually this confused me a bit. I didn't intend to make any claims about the peakmem benchmarks, only the track_nbytes benchmarks. I believe my previous comment is focused entirely on the track_nbytes benchmarks, how it is a completely separate measurement from peakmem, and asking whether track_nbytes should be included in the benchmarking suite. Can you clarify:
If the answer to (3) is "yes" then I would agree it's probably okay for now, for the purposes of this PR, because the benchmarks aren't using chunked grids (right?).
Thank you for considering my points above and for clarifying this! This explanation helped me understand clearly the motivation for wanting track_nbytes as a benchmark rather than a pytest test. Rephrasing to make sure I understand: the design decisions are (A) to keep performance-related checks in benchmarking, and (B) that "static object size" is performance-related, not an issue of correctness. These seem reasonable enough, and moving forwards I will keep them in mind.
Oh! I didn't realize there was an easy way to see a graph of how the ASV numbers change across commits, I thought you might need to click one ASV job at a time. A quick google seems to indicate While looking into that yml file briefly I also noticed it has instructions to post to https://github.com/UXARRAY/uxarray-asv, however that repo seems to have not been updated for over a year. (Something like |
|
@Sevans711 I am referring to the What I was saying in my last comment is that what we're measuring with In practice, I think the issue isn't too important because we are currently expecting to materialize the whole grid anyway, so the static object size is actually a reasonable proxy for what the memory usage behavior currently does. To answer 3, the peakmem benchmarks being replaced here were never measuring memory usage correctly. My goal here is to have them be "less wrong," but there could still be room for improvement. |
Ah, thank you for clarifying, and this is a great point.
This might be a naive question, but why can't you just run the functions once during setup? I think that's the pattern used by other benchmarks to avoid the transient "cold start" overhead leaking into the tests. For example, something like this: class FaceBounds:
def setup(self, grid_path):
tmp = ux.open_grid(grid_path)
tmp.bounds # cold start memory costs of imports, numba caching, etc. will happen here
self.uxgrid = ux.open_grid(grid_path) # make a fresh object for the actual tests below
def peakmem_face_bounds(self, grid_path):
"""Peak memory usage obtain ``Grid.face_bounds."""
face_bounds = self.uxgrid.boundsWould this not give a proper peakmem measurement of just |
3 participants
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