Skip to content

Make geodesic slope memory guard backend-aware for dask#2779

Merged
brendancol merged 1 commit into
mainfrom
issue-2765
Jun 1, 2026
Merged

Make geodesic slope memory guard backend-aware for dask#2779
brendancol merged 1 commit into
mainfrom
issue-2765

Conversation

@brendancol
Copy link
Copy Markdown
Contributor

Closes #2765

What changed

  • slope(method='geodesic') sized its memory guard against the full raster before dispatch. That's correct for eager NumPy/CuPy (they build a (3, H, W) float64 stack up front), but wrong for dask, which streams that stack chunk by chunk via map_overlap. The full-raster check rejected large-but-chunked rasters that would run fine.
  • Added _check_geodesic_memory_backend_aware: eager backends keep the full-raster check; dask backends check the largest spatial chunk plus its 1-cell map_overlap halo (+2 per spatial dim). A single whole-raster chunk has no memory advantage and is still rejected.

Backend coverage

The guard runs before dispatch, so the new path covers all four backends: numpy and cupy use the eager full-raster check; dask+numpy and dask+cupy use the per-chunk check. aspect(method='geodesic') has the same eager-only guard and is left for a separate change to keep this PR scoped to the reported finding.

Test plan

  • Eager raster too large for available RAM still raises MemoryError (existing test).
  • Same raster chunked small enough is allowed through (new dask test).
  • Dask array with a single whole-raster chunk is still rejected (new dask test).
  • test_geodesic_slope.py passes (23 tests).
  • test_geodesic_aspect.py still passes (shared geodesic.py untouched in behaviour).

The geodesic slope path checked working-buffer memory against the full
raster before backend dispatch. That fits eager NumPy/CuPy, which build a
(3, H, W) float64 stack up front, but the dask backends run that stack
chunk by chunk via map_overlap, so peak memory tracks the largest chunk
plus its 1-cell halo. Sizing against the full raster rejected
large-but-chunked rasters that would stream through fine.

Add _check_geodesic_memory_backend_aware: eager backends still check the
full raster; dask backends check the largest spatial chunk (+2 per dim
for the depth-1 halo). A single whole-raster chunk is still rejected.
@github-actions github-actions Bot added the performance PR touches performance-sensitive code label Jun 1, 2026
Copy link
Copy Markdown
Contributor Author

@brendancol brendancol left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

PR Review: Make geodesic slope memory guard backend-aware for dask

Blockers (must fix before merge)

None.

Suggestions (should fix, not blocking)

  • aspect.py:427 has the same eager-only guard (_check_geodesic_memory(rows, cols)) and so has the same bug. The PR scopes itself to slope per the issue and says so in the body, which is fair. File a follow-up so the aspect path doesn't quietly keep the wrong behavior.

Nits (optional improvements)

  • geodesic.py: the guard uses max(row_chunks) and max(col_chunks) independently, so the two maxima can come from different chunks and the bound can slightly overestimate the largest real chunk. That over-rejects rather than under-rejects, which is the safe direction for a memory guard, so it's fine as-is. Noting it so the conservatism is on record.

What looks good

  • The +2 per spatial dim matches the depth-1 map_overlap halo. The chunk function allocates against the haloed chunk shape, so the byte budget tracks the real peak.
  • getattr(data, "chunks", None) splits eager (no chunks attribute) from dask without importing dask in the guard.
  • Tests pin all three behaviors: eager rejected, chunked allowed, single whole-raster chunk still rejected. The chunked test first checks that the same raster is rejected eagerly, which makes the contrast explicit.

Checklist

  • Contract matches the issue: yes
  • Both backend classes covered: yes (eager full-raster, dask per-chunk)
  • NaN handling: unchanged, n/a
  • Edge cases tested: empty chunk tuple guarded with ... else 0
  • Dask chunk boundaries: yes, +2 halo accounted for
  • No premature materialization: guard reads .chunks metadata only, no compute
  • Benchmark: not needed (cheap metadata check)
  • README matrix: not needed (no API change)
  • Docstrings: present on the new helper

@brendancol
Copy link
Copy Markdown
Contributor Author

Review follow-up: the aspect.py suggestion is tracked as a separate issue (#2783) to keep this PR scoped to slope. The max() * max() nit is intended (conservative over-rejection is the safe direction for a memory guard); no change. No code changes resulted, so the existing review stands.

@brendancol brendancol merged commit 6d471a6 into main Jun 1, 2026
9 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.

Geodesic memory guard rejects chunked Dask rasters using full-raster size

1 participant