Guard sink_d8() against unbounded memory allocations (#1356)#1361
Merged
Conversation
The numpy and cupy paths allocate H*W working buffers (labels and BFS queues on CPU, labels grid on GPU) before any sanity check on the input size. Passing a sufficiently large in-memory raster can OOM the host or device. Add per-module _BYTES_PER_PIXEL (24) and _GPU_BYTES_PER_PIXEL (8) constants and _check_memory / _check_gpu_memory helpers that raise MemoryError when the projected working set exceeds 50% of available RAM / free GPU memory. Wire the guards into the eager numpy and cupy branches of sink_d8(); dask paths skip the guard since per-tile allocations are bounded. Mirrors the pattern from #1318/#1319 and the rest of the hydro guard series.
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
1 participant
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.
Summary
Closes #1356.
sink_d8()ran the BFS connected-component kernel on the eager numpy and cupy backends with no check on the input size. Large in-memory rasters could exhaust host or device memory.This PR mirrors the pattern from #1319 and the rest of the hydro guard series.
Bytes per pixel
labelsfloat64 +queue_rint64 +queue_cint64)labelsfloat64)The caller-provided
flow_diris not double-counted.Changes
_BYTES_PER_PIXEL,_GPU_BYTES_PER_PIXEL,_available_memory_bytes,_available_gpu_memory_bytes,_check_memory,_check_gpu_memoryhelpers insink_d8.py.sink_d8()calls the matching_check_*helper on the numpy and cupy branches before dispatching. Dask branches skip the guard.TestMemoryGuard: numpy raises, normal input passes, dask skips the guard, error message points at dask, byte-per-pixel constants are pinned.Test plan
pytest xrspatial/hydro/tests/test_sink_d8.py-- 21 passedpytest xrspatial/hydro/tests/(excluding the two known dask cleanup flakes) -- 642 passed