Cut peak memory across the encode and payload paths - #309
Conversation
Audit of engine memory with peak-RSS and tracemalloc probes per workload. The core payload path is already lean (a 10M-point density build costs 13.8 MB over the data; a 10M M4 line costs 0.9 MB), so this targets the four places where transient allocation was many multiples of the result — and on this allocator a one-shot transient sets the process high-water permanently, so shrinking peaks shrinks steady-state RSS. - Color-source quantization was nominally chunk-bounded at 2^22 = 4M rows, which meant every real column still ran as one shot with three full-length f64 temporaries live at once. Chunk at 2^18 and run the scale/round stages in place; route direct-RGBA shipping (previously unchunked) through the same helper. Resolving a 2.1M-row continuous color channel drops 44.1 -> 7.3 MB, taking a colored 2.1M-point first paint from 181 -> 147 MB RSS. - The indexed-palette PNG encoder held five full frames: an intp-per-pixel np.unique inverse, a redundant astype, a list of per-row bytes, its join, and a tobytes. Stage one (h, stride+1) scanline buffer, resolve indices with searchsorted over the <=256-entry palette in row blocks, and hand the buffer straight to zlib. 1800x840 peaks at 6.1 MB instead of 49.9 and encodes in 8.6 ms instead of 15.3; at 4K it is 33 MB instead of 274 and 40 ms instead of 87. Truecolor loses one frame (-17%) at unchanged speed. - Direct-tier scatter/line/area rebuilt an all-visible row mask on every build: two isfinite passes, an &, an np.all, two N-byte temporaries -- and area also materialized an identity index vector and gathered animation keys through it. Zone maps count NaN *and* +-inf as null, so with linear axes and no nulls the mask cannot drop a row; skip it unless a log axis or a baseline column outside the x/y zone maps can. Scatter emit -35%, categorical -33%, continuous channels -16%, keyed animation payloads -19 to -35%. - SVG documents joined the mark markup twice (once inside the return f-string) and kept one Python str per point, ~50 bytes of object header each. One flat join for the document, 4096-marker blocks for the markers: a 100k-point export peaks at 26.9 MB instead of 38.7. The native rasterizer now borrows the display list instead of freezing a bytes copy of it, and facet PNG export stops copying the whole frame twice. - memory_report() reports capacity_bytes per column and canonical_capacity_bytes per store, and builds resident_array_bytes from the capacity total. A streamed column's values is a prefix view of its capacity-doubling buffer, so up to half of what it holds was invisible: a 20-append figure reported 1.92 of a real 3.20 MB. Continuous channels already reported their growth buffers this way. Figures that never appended report exactly what they did before. Every output stays byte-identical: 89 fingerprints over payloads, split payloads, density views, drill cycles, selections, append frames, and SVG/PNG/PDF/JPEG/WebP/HTML exports match origin/main, as do the animation-key paths with and without non-finite rows. Perf gate is two interleaved A/B campaigns over the 89-benchmark CodSpeed suite (arms alternated, min-of-runs, 10 full suite runs total): median +0.36%, mean -0.24%, no reproducible regression -- the two rows that crossed +5% in a single campaign (contour, weighted_ecdf) are high-variance and land in favour of this branch when measured in isolation, and weighted_ecdf is untouched native code running from the same dylib in both arms.
`_pack_entropy` built one int64 index *per output bit* twice over, so a 2.8 MB entropy stream cost over 1.5 GB of peak memory for a 6 MB image — enough to make a 4K JPEG export a plausible OOM. Pack in bounded bit passes instead, carrying each pass's sub-byte remainder into the next so the emitted bytes are exactly those of the whole-image pass (0xFF stuffing is byte-local, so it stays per-pass). Around it, the pipeline now hands memory back as it goes: the interleaved f32 source dies once the three planes exist, each plane dies when its blocks are built, the quantize chain runs in place, and the per-component token fields are released as they are gathered into MCU order. The rounding rule is now `trunc(x + copysign(0.5, x))` rather than `sign(x) * floor(|x| + 0.5)`: same round-half-away-from-zero, three in-place passes and one temporary instead of four passes and four full-size temporaries. It is exact rather than approximate — IEEE addition is sign-symmetric, so for x < 0 the sum is the exact negation of |x| + 0.5 and truncation toward zero matches that magnitude's floor; ±0 lands on ±0.0 either way and casts to 0. 1800x840: peak 108.2 -> 48.4 MB, 40.3 -> 38.3 ms. Photographic content: 1515.9 -> 399.9 MB, 292 -> 285 ms. Gradient: 127.2 -> 48.4 MB, 51.0 -> 41.5 ms. Byte-identical across 90 encodes (5 qualities x 18 images: flat/noise/gradient at five sizes down to 1x1, odd non-multiple-of-8 dimensions, RGB-vs-RGBA input, all-black, all-white), and the 89 payload/export fingerprints are unchanged. 2392 passed.
Three bookkeeping wins found by the same audit, none of which changes a byte of output. Standalone HTML export interpolated every base64 payload chunk into a wrapper string and then folded a second joined copy into the document f-string. The encoded payload *is* the bulk of the document, so that was two extra full copies of it. Append the chunk as its own part and join once: a 1M-point export peaks at 34.0 MB instead of 41.3 and is marginally faster. The newline stays *between* blocks, so a payload-free document keeps its exact byte layout too (checked against an empty figure). The rasterizer's serial point and segment passes collected `0..n` into a Vec<u32> purely to iterate it in order — 4 bytes per mark, and the serial path is the one CodSpeed measures, since raster_fanout is pinned serial under CODSPEED_ENV. The band painters now take any item iterator, so the fanned-out callers pass their bucketed indices and the serial callers pass a range. Monomorphization keeps both at the same cost per mark. `interaction._quantize_dval` clipped, scaled, and rounded through three full-window temporaries; it now reuses the clip's copy. Also drops a stale `_cpu.color`/`_cpu.size` reference from the `ship_channels` docstring — those attributes no longer exist. Validated: 89 payload/export fingerprints, 90 JPEG encodes, 30 HTML documents (including animation-progress and custom-CSS variants and an empty figure), and the animation-key paths all byte-identical; 2392 passed; 118 Rust tests pass; clippy, ruff, and ty clean and unchanged.
|
Warning Review limit reached
Next review available in: 1 minute Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (3)
📝 WalkthroughWalkthroughRendering and export paths now use bounded buffers, incremental assembly, and reduced intermediate allocations across PNG, JPEG, SVG, HTML, rasterization, payload generation, and quantization. Memory reports include streamed-column capacity, and native raster APIs accept buffer-like command inputs. ChangesPerformance and memory optimizations
Estimated code review effort: 4 (Complex) | ~60 minutes Possibly related issues
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
Merging this PR will improve performance by 16.63%
|
| Benchmark | BASE |
HEAD |
Efficiency | |
|---|---|---|---|---|
| ⚡ | test_animation_plain_payload_100k |
2.2 ms | 1.8 ms | +20.49% |
| ⚡ | test_animation_keyed_payload_100k |
3.3 ms | 2.9 ms | +12.9% |
Tip
Curious why this is faster? Comment @codspeedbot explain why this is faster on this PR, or directly use the CodSpeed MCP with your agent.
Comparing alek/mem-audit (84e0fc2) with main (62462ec)
Footnotes
-
2 benchmarks were skipped, so the baseline results were used instead. If they were deleted from the codebase, click here and archive them to remove them from the performance reports. ↩
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
python/xy/_jpeg.py (1)
357-369: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick win
partsstays alive through the return expression, undermining the documented savings.The docstring claims this costs "one joined copy plus the gathered result rather than both plus the parts," but
partsis a named local that CPython keeps alive for the full duration ofreturn np.concatenate(parts)[order](locals aren't freed until the frame returns). So during the[order]fancy-index step,parts, theconcatenateresult, and the final indexed result all coexist — effectively 3x the field size at peak, not 2x. An explicitdel partsbefore indexing would actually realize the claimed savings.♻️ Proposed fix to release `parts` before indexing
def _gathered(fields: list[list[np.ndarray]], index: int, order: np.ndarray) -> np.ndarray: parts = [f[index] for f in fields] for f in fields: f[index] = _EMPTY_I64 - return np.concatenate(parts)[order] + joined = np.concatenate(parts) + del parts + return joined[order]🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/xy/_jpeg.py` around lines 357 - 369, Update _gathered to concatenate the collected parts into a temporary result, explicitly release parts before applying order-based indexing, then return the indexed result. Preserve the existing component-release behavior and output ordering while ensuring the original parts do not remain alive during the fancy-index operation.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@python/xy/_jpeg.py`:
- Around line 357-369: Update _gathered to concatenate the collected parts into
a temporary result, explicitly release parts before applying order-based
indexing, then return the indexed result. Preserve the existing
component-release behavior and output ordering while ensuring the original parts
do not remain alive during the fancy-index operation.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 33fca85a-8dd9-4cf1-b15c-480701eb9791
📒 Files selected for processing (15)
CHANGELOG.mdpython/xy/_figure.pypython/xy/_jpeg.pypython/xy/_native.pypython/xy/_payload.pypython/xy/_png.pypython/xy/_raster.pypython/xy/_svg.pypython/xy/channels.pypython/xy/columns.pypython/xy/export.pypython/xy/facets.pypython/xy/interaction.pyspec/design-dossier.mdsrc/raster.rs
…capacity Review of #309 asked for in-repo coverage of the behaviors these commits introduced. Mutation-testing each candidate first showed the gaps were real: - The JPEG entropy packer's chunk-boundary carry was dead code in the suite. Every image in tests/test_jpeg.py has an entropy stream orders of magnitude shorter than one 2^18-bit group, so dropping the carry, padding a non-final group, or skipping its byte stuffing all kept CI green while corrupting every large export. Now checked against a bit-by-bit statement of the T.81 rule at five tiny group sizes (every byte alignment), at the production group size over a stream crossing it twice, on degenerate streams (empty, one sub-byte token, consecutive 0xFF), and end-to-end: a noisy 320x400 image encodes identically with the packer forced into a single pass. The tests pin the emitted bytes, not the grouping — moving a split point by one token is a refactor and stays green by design. - `_visible_mask_needed`'s log-axis and baseline-null conditions were invisible: removing either left all 2405 tests passing. Both are now pinned (x and y log axes, log combined with non-finite rows, symlog explicitly keeping non-positive rows, and an error_band whose baseline carries NaN/inf), so the predicate cannot be tightened past what the mask actually rejects. The null-rows condition was already covered by test_area_hostile. - The memory report's capacity accounting had assertions only for the never-appended case — the case the change deliberately leaves alone. Now the appended case is pinned end to end: capacity exceeds live length, per-column capacity matches `Column.capacity_bytes`, `resident_array_bytes` is built from the capacity total, and growth stays inside the 2x amortized bound. Each new test was verified to fail against a targeted mutation of the code it covers. Two mutations survive and should: shifting an entropy group split by one token, and treating symlog as log (which computes a mask that then rejects nothing) — both are output-preserving. Also: `_log_visible_mask` now documents the invariant it shares with `_visible_mask_needed` in both directions, and `png_truecolor` takes a real buffer union instead of `object`, with the C-contiguity requirement stated. 2412 passed; 89 payload/export fingerprints and 90 JPEG encodes unchanged; ruff and ty clean.
|
Addressed the review. Before writing anything I mutation-tested each ask, which changed what got done: The entropy-packer gap was real and worse than described. Not only does no test cross a group boundary — I confirmed that dropping the carry, padding a non-final group, or skipping its byte stuffing each left the whole suite green. Now covered against a bit-by-bit statement of the T.81 rule: five tiny group sizes (so passes end at every byte alignment), the production group size over a stream crossing it twice, degenerate streams (empty, a single sub-byte token, consecutive The tests deliberately pin the emitted bytes, not the grouping: packing is a concatenation with carries, so moving a split point by a token is output-equivalent. That mutation survives on purpose, and the test file says so, so nobody later reads it as a hole. The mask conditions: two of three were genuinely unpinned, and I checked rather than assumed. Removing the log-axis condition, or the baseline-null condition, left all 2405 tests passing. Removing the null-rows condition failed Worth noting which direction is dangerous: weakening the predicate only recomputes a mask that rejects nothing (a wasted pass, not a wrong payload), while tightening it silently ships rows that should have been dropped. The new tests pin the tightening direction. The Report assertions now cover the appended case — the case the change actually ships. Capacity exceeds live length, per-column capacity matches Both minor items taken.
One correction to the review's framing, for the record: |
CodSpeed caught a real regression from f5ea7fe: `test_html_export_line` −10.1%. That benchmark's document *is* the ~330 KB client bundle (a 100k-point line M4-decimates to a ~30 KB payload), and the new parts list interpolated `client_js` into its header part, so the join copied the bundle a second time. Locally reproducible as +1.9% wall-clock on that shape — ~16 µs, which is one memcpy of 330 KB and confirms the mechanism rather than guessing at it. The rule the assembler now follows: any string that is large relative to the document is its own part, never interpolated. That covers the client bundle, the spec JSON, the inline decoder, and each base64 chunk — so each is copied exactly once by the join, which is strictly fewer copies than main did (main duplicated the base64; this duplicated the bundle). `test_html_export_line` shape: 0.730 → 0.726 ms best-of-three-runs, medians 0.792/0.803/0.846 → 0.785/0.787/0.794. A 1M-point export is 1.248 → 1.205 ms and now peaks at 33.3 MB rather than 34.0 (main: 41.3). The empty-payload byte layout is still pinned by the 30-document check. Also takes the CodeRabbit note on `_jpeg._gathered`: a local stays alive until its frame returns, so `return np.concatenate(parts)[order]` held the pieces, the join and the gather simultaneously. Dropping `parts` and `joined` explicitly makes the docstring's claim true. Measured peak on the probe images is unchanged — their binding peak is in token construction, not the gather — but the three-way coexistence was real. 89 payload/export fingerprints, 90 JPEG encodes and 30 HTML documents unchanged; 2412 passed; ruff and ty clean.
|
Both items fixed in CodSpeed's The assembler now follows one rule: any string that is large relative to the document is its own part and is never interpolated — the client bundle, the spec JSON, the inline decoder, and each base64 chunk. Each is then copied exactly once by the join, which is strictly fewer copies than After the fix, on the benchmark's shape: 0.730 → 0.726 ms best-of-three-runs, medians 0.792/0.803/0.846 → 0.785/0.787/0.794. A 1M-point export is 1.248 → 1.205 ms and peaks at 33.3 MB versus 34.0 before and 41.3 on main — so the fix improved both axes over my own earlier version. Two notes on reading the rest of that report. The @coderabbitai's One honest qualification: the measured peak on my probe images did not move (48.4 / 399.9 / 48.4 MB). Those images' binding peak is in All other 24 checks were green before this push. |
|
Tip For best results, initiate chat on the files or code changes. 🧩 Analysis chain🏁 Script executed: #!/bin/bash
set -euo pipefail
git show --no-ext-diff --format=fuller --stat 84e0fc2
git diff --no-ext-diff 84e0fc2^ 84e0fc2 -- python/xy/export.py python/xy/_jpeg.py
ast-grep outline python/xy/export.py --items all --match 'to_html|_iter_base64_chunks|_base64_chunks'
ast-grep outline python/xy/_jpeg.py --items all --match '_gathered|encode'
rg -n -A70 -B15 'def to_html\b|def _gathered\b' python/xy/export.py python/xy/_jpeg.pyLength of output: 19318
Both findings are addressed. 🐇 ✏️ Learnings added
|
1 participant
Audit of engine memory, plus the wins from it that survived a perf gate.
Instrumented peak RSS and
tracemallocper workload in fresh processes. The finding thatshaped the PR: the payload path is already lean — a 10M-point density build costs
13.8 MB over the 160 MB of canonical data, and a 10M-point M4 line costs 0.9 MB. Nothing
there needed fixing. Every real problem was a transient, and almost all of them were in
the encoders. Since freed memory is not returned to the OS, each one-shot peak
permanently raises the process high-water, so cutting a peak cuts steady-state RSS too.
Three commits, kept unsquashed — they are independently validated and independently
revertable.
Memory: before → after
to_svg, 100k-point scatterto_html, 1M-point scatterSpeed: the same changes, measured in isolation (min-of-N, interleaved)
to_html, 1MWhat was actually wrong
Every one of these was "copying what the caller already holds":
bit — so a 2.8 MB stream peaked over 1.5 GB. It now packs in bounded bit passes, carrying
each pass's sub-byte remainder so the emitted bytes are exactly those of the whole-image
pass (0xFF stuffing is byte-local, so it stays per-pass). Around it, planes are released
as consumed, the quantize chain runs in place, and token fields are freed as they are
gathered into MCU order. The rounding rule is now
trunc(x + copysign(0.5, x)): sameround-half-away-from-zero in three in-place passes instead of four passes and four
full-size temporaries, and exact rather than approximate — IEEE addition is
sign-symmetric, so for
x < 0the sum is the exact negation of|x| + 0.5.intp-per-pixelnp.uniqueinverse (twice the image itself), a redundantastype, a list of per-rowbytes, its join, and atobytes. Now: one(h, stride+1)staging buffer,searchsortedover the ≤256-entry palette in row blocks, buffer handed straight to zlib.still ran as one shot with three simultaneous full-length f64 temporaries. Chunked at 2¹⁸
with in-place stages;
direct_rgbashipping (previously unchunked) routed through thesame helper.
isfinitepasses, an&, annp.all, two N-byte temporaries. Zone maps count NaN and±inf as null (§22/§19), so with linear axes and no nulls the mask cannot drop a row. Area
additionally built an identity index vector and gathered animation keys through it.
strper point, ~50 bytes of object header each. One flat join; markers collapse into4096-marker blocks.
their own document parts and joined once. The newline stays between blocks, so a
payload-free document keeps its exact byte layout.
0..nintoa
Vec<u32>just to iterate it in order (the band painters take an item iterator, sofanned-out callers pass bucketed indices and serial callers pass a range);
render_rasterborrows the display list instead of freezing abytescopy of it; facetPNG export drops two full-frame
tobytes(); the density quantizer keeps one temporary.Report accuracy (§27)
memory_report()now reportscapacity_bytesper column andcanonical_capacity_bytesperstore, and builds
resident_array_bytesfrom the capacity total. A streamed column'svaluesis a prefix view of its capacity-doubling growth buffer, so up to half of what theprocess holds was invisible to the report — a 20-append figure reported 1.92 MB of a real
3.20 MB. Continuous channels already reported their growth buffers this way; columns now
match. Figures that never appended report exactly what they did before. Dossier §27 updated.
Validation
Output is byte-identical everywhere. Fingerprints hashed on
origin/mainand on thisbranch and diffed:
SVG / PNG / PDF / JPEG / WebP / HTML fingerprints across 30 chart configurations,
including NaN and ±inf rows, log and symlog axes, wide categorical, and direct-RGBA
channels
non-multiple-of-8 dimensions, RGB-vs-RGBA input, all-black, all-white)
The only intentional difference is
memory_report, which gained fields.pytest -q: 2392 passed, 4 skipped.cargo test: 118 passed. clippy,ruff check,ruff format --checkclean;ty checkreports the same 13 pre-existing diagnostics as thebase commit.
Perf gate: three interleaved A/B campaigns over the 89-benchmark CodSpeed suite in
walltime mode — arms alternated (
ref mem mem ref ref mem), full suite per arm,min-of-runs per arm, 16 suite runs total, and the native dylib swapped per arm alongside the
Python tree. Final campaign: median +0.24%, mean −0.49%, 0 of 89 rows slower by >5%.
Note for the CodSpeed comment
Expect noise on rows this branch cannot touch. Every single-campaign outlier above +5%
failed to reproduce:
contour_core_2dread +9.3% in one campaign and +2.1% in another; it has ~17% intra-armspread, and in an isolated in-process interleaved benchmark this branch produced the
fastest run of all six (3.43 ms vs the baseline's best 4.65 ms).
weighted_ecdf_native(+6.4% in one campaign) andsample_mask(+4.7%) are untouchednative kernels running from the same dylib in both arms — noise by construction.
native_png_export_scatterread +0.79% in-suite while an in-process interleavedmicro-benchmark showed equal-or-better medians and the segments path reproducibly −0.8%.
Reporting it rather than hiding it.
select_lasso_message_1mis the repo's known bimodal row.Deliberately not in this PR
Documented rather than coded, because each is a trade or a bigger change than a memory fix
should smuggle in:
MeanColorCellis 40 B/cell (kernels.rs), andtiles::build_colorholds it alive whilematerializing the count and color planes: a 168 MB transient at the 2048² default for a
2.1M-point trace, and
MEAN_COLOR_ACCUM_BUDGET_BYTES(1 GB) permits 4 workers — 672 MB ona row-heavy trace. The exact u64 sums are a documented bitwise-determinism guarantee, so
the cell cannot shrink without changing that contract. Both available levers (lower the
worker budget; pick
base_dimfrom row count) regress real-world build wall-clock that nobenchmark covers, so they are trades, not wins.
range_indices/bin_2d_indicesallocate 4N regardless of result size — a 400 MBtransient at 100M rows to return a few KB for a deep-zoom window. Chunking adds two extra
copies of the result exactly when the result is the whole column; count-then-fill doubles
a memory-bound scan. The sound fix is structural:
interaction.pyscans the view windowand
_padded_drill_windowthen scans up to three superset windows, discarding all butthe winner — one scan of the widest rung plus O(rows) filters could replace k+1 full O(N)
scans, a CPU win as well as a memory win. It needs a rows-restricted range kernel (only
polygon_selecttakesrowstoday), i.e. a new C-ABI entry point and anABI_VERSIONbump, and it lands on the two already-red drill rows. Worth its own PR with its own A/B.
to_pdfinflates the SVG into an ElementTree: 29.4 MB peak for a 1.07 MB PDF at only20k points, versus 5.8 MB for the same figure's SVG. Needs a streaming converter or
conversion from the payload instead of from serialized SVG.
to_html(path=...)still materializes the whole document before writing it; streamingparts to the file would remove one full copy but changes the return contract.
196k cells × 40 B); capping the fan-out at 2 halves it and costs wall-clock.
serial, partly because the fan-out estimate needs a full pass first.
memory_reportshould surface the transient high-water. After a colored 2.1M drillsession RSS is 382 MB while
resident_array_bytesreports 120 MB; most of the gap isone-shot transients the allocator keeps rather than anything retained. §27's promise
argues for an observed-RSS / peak-transient line.
Summary by CodeRabbit