Fix perlin dask OOM by dropping redundant dask.persist (#3469)#3473
Merged
Conversation
The dask+numpy and dask+cupy paths persisted the full noise array before normalizing, forcing every chunk resident at once and OOMing on large inputs. The min/ptp (min/max) reductions already share the noise subgraph within a single dask.compute call, so the persist is redundant. Remove it so peak memory stays bounded by chunk size. Adds a regression test that fails if dask.persist is reintroduced.
brendancol
commented
Jun 23, 2026
brendancol
left a comment
Contributor
Author
There was a problem hiding this comment.
PR Review: Fix perlin dask OOM by dropping redundant dask.persist (#3469)
Blockers (must fix before merge)
None.
Suggestions (should fix, not blocking)
- The fix swaps memory for compute, and that swap is worth spelling out in the PR so nobody "optimizes" the persist back in later. Without it, the noise subgraph runs twice: once for the
min/ptpreductions indask.compute, then again when the caller materializes the returned normalized array. For an OOM fix that's the right call (the noise is point-wise and deterministic, so the recompute is exact and cheap), but the code comment only covers the memory side. A half-sentence noting the recompute is intentional would save the next reader the trace.
Nits (optional improvements)
test_perlin_dask_does_not_persist_whole_arrayonly covers the dask+numpy path. The dask+cupy path got the same edit but no persist guard. Folding amonkeypatch.setattr(dask, "persist", ...)assertion intotest_perlin_dask_gpuwould pin down both. Not blocking, since the edit is symmetric and the GPU path isn't always available in CI.
What looks good
- Symmetric across both dask backends; numpy and cupy single-GPU are untouched.
- The premise checks out: the
minandptpgraphs reference the same named noise keys, so onedask.computeover both shares the subgraph and the persist really is dead weight. - Output is unchanged. The noise is deterministic, so recomputing yields identical values, and the regression test confirms the dask result still matches the numpy reference.
- The regression test fails loudly if
dask.persistcomes back, which is the right guard for this class of fix.
Checklist
- Algorithm matches reference: not affected (normalization unchanged)
- All implemented backends consistent: yes, dask+numpy and dask+cupy tests pass
- NaN handling correct: not affected
- Edge cases covered: existing coverage retained
- Dask chunk boundaries: perlin is point-wise, no overlap
- No premature materialization: improved (persist removed)
- Benchmark exists or not needed: benchmarks/benchmarks/perlin.py exists
- README feature matrix: not applicable, no API change
- Docstrings: unchanged
brendancol
commented
Jun 23, 2026
brendancol
left a comment
Contributor
Author
There was a problem hiding this comment.
Follow-up review (after 67fa8d2)
Both findings from the first pass are addressed:
- Suggestion (recompute not documented): the comment in both dask paths now states the returned lazy array recomputes the noise on materialization and that the extra pass is intentional, since the noise is point-wise and deterministic.
- Nit (dask+cupy not guarded):
test_perlin_dask_gpunow monkeypatchesdask.persistto raise, so the GPU path is covered by the same guard as the CPU path.
13 perlin tests pass locally, including the cupy and dask+cupy GPU tests on a CUDA host. No new findings.
…e-perlin-2026-06-23
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.
Closes #3469
What changed
dask.persist(data)from_perlin_dask_numpyand_perlin_dask_cupy. Both paths now reduce straight off the lazy noise array.min/ptp(andmin/maxon GPU) reductions already go out in onedask.computecall, which shares the noise subgraph between them, so each chunk is computed once and freed after both reductions read it. The persist was redundant and forced every chunk resident at once, which OOMs on large inputs.Backends
numpy and cupy single-GPU paths are untouched. dask+numpy and dask+cupy lose the persist; output is unchanged.
Test plan
xrspatial/tests/test_perlin.pypasses, including the cupy and dask+cupy GPU tests on a CUDA host (13 passed).test_perlin_dask_does_not_persist_whole_array, which monkeypatchesdask.persistto raise and confirms the dask path no longer calls it while still matching the numpy reference.