Guard basin_d8() against unbounded memory allocations (#1357)#1359
Merged
brendancol merged 1 commit intomainfrom Apr 29, 2026
Merged
Guard basin_d8() against unbounded memory allocations (#1357)#1359brendancol merged 1 commit intomainfrom
brendancol merged 1 commit intomainfrom
Conversation
basin_d8() on the numpy and cupy backends allocated H*W working arrays (~33 B/pixel CPU, ~28 B/pixel GPU) without checking available memory. A 50000x50000 input would request ~83 GB before failing. Add module-local _BYTES_PER_PIXEL / _GPU_BYTES_PER_PIXEL, _available_memory_bytes / _available_gpu_memory_bytes, _check_memory / _check_gpu_memory helpers (50% threshold) and call them from the numpy and cupy dispatch branches. The dask paths delegate to the watershed dask infrastructure and remain bounded per-tile, so they skip the guard. Tests cover the four guard outcomes plus a normal-sized success.
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.
Closes #1357.
Summary
basin_d8()on the numpy and cupy backends allocated H*W working arrays without checking available memory:fdcast (8 B),labels(8 B),state(1 B), pluspath_r/path_cint64 inside_watershed_cpu(8 + 8 B). ~33 bytes/pixel.flow_dir_f64,labels,state, and the finalcp.whereresult. ~28 bytes/pixel on the device.A 50000x50000 input on CPU would request ~83 GB before erroring out.
Fix
Add module-local
_BYTES_PER_PIXEL/_GPU_BYTES_PER_PIXEL,_available_memory_bytes/_available_gpu_memory_bytes,_check_memory/_check_gpu_memoryhelpers wired into the numpy and cupy dispatch branches before allocations. 50% threshold matches the watershed_d8 / flow_accumulation_d8 / hand_d8 series.The dask paths delegate to the watershed dask infrastructure and stay bounded per-tile by the user's chunk size, so they skip the guard.
Tests
5 new tests in
TestMemoryGuard:MemoryErrorwhen_available_memory_bytesis patched lowMemoryErrorwhen_available_gpu_memory_bytesis patched low (skipped without CUDA)Test plan
pytest xrspatial/hydro/tests/test_basin_d8.py-- 22 passedpytest xrspatial/hydro/tests/-- 644 passed