From 44418d6020435547ee8224c36a1ad1a7ace8065c Mon Sep 17 00:00:00 2001 From: Brendan Collins Date: Mon, 1 Jun 2026 10:33:33 -0700 Subject: [PATCH] Make geodesic memory guard backend-aware for dask (#2765) 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. --- xrspatial/geodesic.py | 29 ++++++++++++++++++ xrspatial/slope.py | 5 ++-- xrspatial/tests/test_geodesic_slope.py | 41 ++++++++++++++++++++++++++ 3 files changed, 72 insertions(+), 3 deletions(-) diff --git a/xrspatial/geodesic.py b/xrspatial/geodesic.py index 4b92665d7..546330ff0 100644 --- a/xrspatial/geodesic.py +++ b/xrspatial/geodesic.py @@ -82,6 +82,35 @@ def _check_geodesic_memory(rows, cols, func_name): ) +def _check_geodesic_memory_backend_aware(agg, func_name): + """Backend-aware front end for :func:`_check_geodesic_memory`. + + Eager NumPy / CuPy backends build a ``(3, H, W)`` float64 stack for the + whole raster up front, so the guard sizes against the full ``(rows, cols)``. + + Dask backends run the same stack chunk by chunk via ``map_overlap``, so + peak working memory is bounded by the largest chunk plus its 1-cell halo, + not the full raster. Sizing the guard against the full raster would reject + large-but-chunked rasters that would otherwise stream through fine, so for + Dask we check the largest chunk instead. + """ + data = getattr(agg, "data", agg) + + chunks = getattr(data, "chunks", None) + if chunks is None: + # Eager (NumPy / CuPy): size against the full raster. + rows, cols = agg.shape[-2], agg.shape[-1] + _check_geodesic_memory(rows, cols, func_name=func_name) + return + + # Dask: size against the largest spatial chunk, plus the 1-cell halo that + # ``map_overlap`` adds on each side (depth (0, 1, 1) -> +2 per spatial dim). + row_chunks, col_chunks = chunks[-2], chunks[-1] + rows = (max(row_chunks) if row_chunks else 0) + 2 + cols = (max(col_chunks) if col_chunks else 0) + 2 + _check_geodesic_memory(rows, cols, func_name=func_name) + + # ===================================================================== # CPU (Numba ngjit) primitives # ===================================================================== diff --git a/xrspatial/slope.py b/xrspatial/slope.py index 56b41df3f..58a422a14 100644 --- a/xrspatial/slope.py +++ b/xrspatial/slope.py @@ -23,7 +23,7 @@ class cupy(object): # local modules from xrspatial.dataset_support import supports_dataset -from xrspatial.geodesic import (INV_2R, WGS84_A2, WGS84_B2, _check_geodesic_memory, +from xrspatial.geodesic import (INV_2R, WGS84_A2, WGS84_B2, _check_geodesic_memory_backend_aware, _cpu_geodesic_slope, _run_gpu_geodesic_slope) from xrspatial.utils import (Z_UNITS, ArrayTypeFunctionMapping, _boundary_to_dask, _extract_latlon_coords, _pad_array, _validate_boundary, @@ -393,8 +393,7 @@ def slope(agg: xr.DataArray, ) z_factor = Z_UNITS[z_unit] - rows, cols = agg.shape[-2], agg.shape[-1] - _check_geodesic_memory(rows, cols, func_name='slope') + _check_geodesic_memory_backend_aware(agg, func_name='slope') lat_2d, lon_2d = _extract_latlon_coords(agg) diff --git a/xrspatial/tests/test_geodesic_slope.py b/xrspatial/tests/test_geodesic_slope.py index d22a379dd..21776fa61 100644 --- a/xrspatial/tests/test_geodesic_slope.py +++ b/xrspatial/tests/test_geodesic_slope.py @@ -273,6 +273,47 @@ def test_planar_method_skips_guard(self, monkeypatch): assert result.shape == (8, 8) +@dask_array_available +class TestGeodesicSlopeMemoryGuardDask: + """The dask geodesic backend streams the raster chunk by chunk, so the + memory guard must size against the largest chunk, not the full raster. + A raster that would be rejected eagerly should be allowed once it is + chunked small enough to fit.""" + + def test_chunked_raster_allowed_when_eager_would_reject(self, monkeypatch): + # 1 MB available. The 200x200 raster needs ~2.2 MB eagerly (56 B/cell) + # and trips the eager guard, but 20x20 chunks need only ~25 KB each. + monkeypatch.setattr( + 'xrspatial.geodesic._available_memory_bytes', lambda: 1024 * 1024 + ) + elev = _flat_surface(H=200, W=200) + # eager numpy of the same size is rejected (sanity check the budget). + r_np = _make_geo_raster(elev, 40.0, 41.0, 10.0, 11.0) + with pytest.raises(MemoryError, match="slope"): + slope(r_np, method='geodesic') + # same raster, chunked small — guard must let it through. + r_da = _make_geo_raster( + elev, 40.0, 41.0, 10.0, 11.0, + backend='dask+numpy', chunks=(20, 20), + ) + result = slope(r_da, method='geodesic') + assert result.compute().shape == (200, 200) + + def test_single_huge_chunk_still_rejected(self, monkeypatch): + # A dask array whose only chunk spans the whole raster has no memory + # advantage over eager, so the guard must still reject it. + monkeypatch.setattr( + 'xrspatial.geodesic._available_memory_bytes', lambda: 1024 * 1024 + ) + elev = _flat_surface(H=200, W=200) + r_da = _make_geo_raster( + elev, 40.0, 41.0, 10.0, 11.0, + backend='dask+numpy', chunks=(200, 200), + ) + with pytest.raises(MemoryError, match="slope"): + slope(r_da, method='geodesic') + + # --------------------------------------------------------------------------- # Tests — cross-backend consistency # ---------------------------------------------------------------------------