From 12cf2b3d5ea9590650ccd2a02d5081fd0956c956 Mon Sep 17 00:00:00 2001 From: Brendan Collins Date: Mon, 15 Jun 2026 06:45:23 -0700 Subject: [PATCH] Guard cost_distance numpy dask chunk path with _check_memory (#3343) The numpy map_overlap chunk function called _cost_distance_kernel directly, allocating several height*width arrays per chunk without the _check_memory guard that _cost_distance_numpy and the cupy/iterative paths apply. An oversized chunk could exhaust a worker with an opaque allocator error instead of a MemoryError pointing at max_cost= or smaller chunks. Add _check_memory(h, w) inside the chunk closure so every numpy allocation path is guarded consistently. --- xrspatial/cost_distance.py | 1 + xrspatial/tests/test_cost_distance.py | 25 +++++++++++++++++++++++++ 2 files changed, 26 insertions(+) diff --git a/xrspatial/cost_distance.py b/xrspatial/cost_distance.py index a080b06cd..5f7325b23 100644 --- a/xrspatial/cost_distance.py +++ b/xrspatial/cost_distance.py @@ -1125,6 +1125,7 @@ def _make_chunk_func(cellsize_x, cellsize_y, max_cost, target_values, def _chunk(source_block, friction_block): h, w = source_block.shape + _check_memory(h, w) return _cost_distance_kernel( source_block, friction_block, h, w, cellsize_x, cellsize_y, max_cost, diff --git a/xrspatial/tests/test_cost_distance.py b/xrspatial/tests/test_cost_distance.py index 309a515e2..555420272 100644 --- a/xrspatial/tests/test_cost_distance.py +++ b/xrspatial/tests/test_cost_distance.py @@ -759,6 +759,31 @@ def test_numpy_memory_guard_passes_for_small_raster(): np.testing.assert_allclose(out[0, 1], 1.0, atol=1e-5) +@pytest.mark.skipif(da is None, reason="dask not installed") +def test_dask_map_overlap_chunk_memory_guard_raises(): + """The numpy map_overlap chunk path also honours _check_memory (#3343).""" + from unittest.mock import patch + + source = np.zeros((16, 16)) + source[0, 0] = 1.0 + friction = np.ones((16, 16)) + + # finite max_cost with friction=1, cellsize=1 -> pad=3, which is < the + # 8x8 chunk size, so the per-chunk map_overlap path is taken (not the + # iterative tile Dijkstra). + raster = _make_raster(source, backend='dask+numpy', chunks=(8, 8)) + fric = _make_raster(friction, backend='dask+numpy', chunks=(8, 8)) + + result = cost_distance(raster, fric, max_cost=2.0) + + # The guard runs inside the dask task, so it fires at compute time. + with patch( + 'xrspatial.cost_distance._available_memory_bytes', return_value=1000 + ): + with pytest.raises(MemoryError, match="max_cost"): + result.compute() + + # ----------------------------------------------------------------------- # Memory guard on CuPy GPU path (Issue #1262) # -----------------------------------------------------------------------