Skip morphology memory guard for lazy dask inputs (#3401)#3410
Open
brendancol wants to merge 4 commits into
Open
Skip morphology memory guard for lazy dask inputs (#3401)#3410brendancol wants to merge 4 commits into
brendancol wants to merge 4 commits into
Conversation
_check_kernel_memory budgets a full padded float64 copy of the input, which only the eager numpy/cupy backends allocate. Dask processes the raster chunk-by-chunk via map_overlap, so peak memory scales with chunk size, not the full shape. Running the full-shape guard on a lazy dask array raised a false MemoryError that rejected large workloads which would have run fine. Skip the guard when the input is dask-backed; keep it for numpy/cupy. Add a regression test that a large lazy dask raster builds the graph under a tiny memory budget without raising.
brendancol
commented
Jun 20, 2026
brendancol
left a comment
Contributor
Author
There was a problem hiding this comment.
PR Review: Skip morphology memory guard for lazy dask inputs (#3401)
Blockers
None.
Suggestions
None.
Nits
xrspatial/morphology.py:451-- The skip is keyed on the input being anydask.Array, regardless of chunking. A user who passes a single-chunk dask array (chunks=(-1, -1)) now bypasses the guard even though that chunk function materializes a full padded copy per block. Unusual chunking choice, and the eager backends still guard, so not worth blocking on. A one-line note in the comment that the skip assumes reasonable chunking would help the next reader.
What looks good
- The fix matches the cause in the issue: the guard budgets a full padded copy that only the eager backends allocate, so skipping it for dask is correct.
- The condition reuses
has_dask_array()andisinstance(agg.data, da.Array), the same patternArrayTypeFunctionMappinguses to detect the dask backends, so it stays consistent with dispatch. - The eager numpy and cupy paths are untouched, so the existing guard tests still apply.
- The regression test builds the graph only and never calls
.compute(), and uses a tiny monkeypatched memory budget so it actually exercises the skip rather than passing by luck.
Checklist
- Algorithm unchanged; only the guard gate changed: ok
- Backends: numpy/cupy guard preserved, dask+numpy/dask+cupy guard skipped: ok
- NaN handling: unchanged
- Edge cases: new dask test added; eager guard tests intact
- Dask chunk boundaries: unchanged (map_overlap untouched)
- No premature materialization: confirmed, graph-only test
- Benchmark: not needed (gate change, not a hot path)
- Docstrings: comment added explaining the skip rationale
brendancol
commented
Jun 20, 2026
brendancol
left a comment
Contributor
Author
There was a problem hiding this comment.
Follow-up review (after b1bafaa)
The one nit from the prior review is resolved: the guard-skip comment now notes that the skip assumes reasonable chunking and that a single giant chunk would still materialize a full padded copy per block. No code-path change, so the test results stand (16 memory/guard/dask tests pass; full file 48 pass).
No remaining blockers, suggestions, or nits.
…e-morphology-2026-06-20
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 #3401
What changed
_dispatch()inxrspatial/morphology.pyno longer runs the full-shape memory guard when the input is dask-backed. The guard budgets a full padded float64 copy of the input, which only the eager numpy/cupy backends allocate. Dask processes the raster chunk-by-chunk viamap_overlap, so peak memory scales with chunk size, not the full shape.MemoryErrorand rejected work that would have run fine.Backend coverage
Test plan
.compute()).test_morphology.pysuite passes (48 tests).