Skip to content

Guard stream_order_mfd() against unbounded memory allocations (#1349)#1352

Merged
brendancol merged 1 commit intomainfrom
issue-1349
Apr 29, 2026
Merged

Guard stream_order_mfd() against unbounded memory allocations (#1349)#1352
brendancol merged 1 commit intomainfrom
issue-1349

Conversation

@brendancol
Copy link
Copy Markdown
Contributor

Fixes #1349.

stream_order_mfd allocated several H*W working arrays in its eager numpy and cupy backends and a (8, H, W) float64 fractions copy on top, with no upfront check that the result will fit in memory. A 50000x50000 raster wants ~262 GB on the CPU path, so the failure mode was an OOM kill rather than a clean Python exception.

This adds the per-module memory guard helpers (_BYTES_PER_PIXEL, _GPU_BYTES_PER_PIXEL, _available_memory_bytes, _available_gpu_memory_bytes, _check_memory, _check_gpu_memory) used across the rest of the hydro series and wires them into the public dispatch before any allocations.

Bytes per pixel

CPU peak working set, Strahler worst case:

array dtype B/px
frac input copy (8,H,W) float64 64
stream_mask int8 1
order float64 8
in_degree int32 4
max_in float64 8
cnt_max int32 4
queue_r int64 8
queue_c int64 8

Total ~105 B/px.

GPU peak working set ~101 B/px. Use 105 B/px for both as a conservative budget.

Behaviour

  • numpy and cupy paths now raise MemoryError with a clear message when projected use exceeds 50% of available memory.
  • dask paths skip the guard. Per-tile allocations are bounded by the chunk size and the dask scheduler handles spilling.
  • The error message names the grid dimensions and points users at dask for out-of-core processing.

Tests

Five new tests in TestMemoryGuard:

  • numpy path raises with patched memory
  • normal-size raster passes the guard
  • dask path bypasses the guard
  • error message mentions dimensions and dask
  • cupy path raises with patched GPU memory

Full file: 14 tests pass. Full hydro suite: 643/644 pass; the one failure is the known test_stream_link_dask_temp_cleanup flake unrelated to this change.

Companion issues

Same pattern as #1318 / #1319, #1321 / #1324, #1322, #1328 / #1331, #1333, #1337 / #1341, #1343.

The eager numpy and cupy backends allocate several H*W working arrays
plus an (8, H, W) float64 fractions copy with no upfront check that
the result fits in memory.  A 50000x50000 raster needs ~262 GB on the
CPU path, so the failure mode is an OOM kill rather than a clean
exception.

Add the per-module memory guard helpers used across the rest of the
hydro series and wire them into the public dispatch before any
allocations.  Threshold is 50% of available memory; dask paths are
unaffected because per-tile allocations are bounded.

CPU budget 105 B/px (Strahler worst case), GPU budget 105 B/px.
@github-actions github-actions Bot added the performance PR touches performance-sensitive code label Apr 29, 2026
@brendancol brendancol merged commit 7f00051 into main Apr 29, 2026
11 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

performance PR touches performance-sensitive code

Projects

None yet

Development

Successfully merging this pull request may close these issues.

stream_order_mfd: no memory guard on H*W working arrays

1 participant