Stitch multi_stop_search segments lazily on dask backends

.claude/sweep-performance-state.csv

@@ -31,7 +31,7 @@ mcda,2026-06-10,SAFE,memory-bound,2,3150,"2 HIGH fixed in PR #3158: owa() dask p
 morphology,2026-06-20,SAFE,compute-bound,1,3401,memory guard fired on full lazy-dask shape (false MemoryError); skip guard for dask-backed inputs; eager numpy/cupy guard preserved
 multispectral,2026-05-02,SAFE,compute-bound,0,,"Re-audit 2026-05-02 after PRs 1292 (true_color memory guard) and 1301 (validate_arrays in true_color). Verified SAFE. No HIGH. MEDIUM: da.stack in _true_color_dask/_true_color_dask_cupy at L1702/L1731 creates (1,1,1,1) chunks along band axis (4 bands so impact is minor, scheduling overhead not OOM). LOW: np.zeros((h,w,4)) at L1681 then full overwrite -- np.empty would suffice. All 17 indices use plain map_blocks with no halo; 8192x8192 ndvi graph is 80 tasks, evi/arvi/ebbi 112 tasks."
 normalize,2026-03-31T18:00:00Z,SAFE,compute-bound,0,1124,Boolean indexing replaced with lazy nanmin/nanmax/nanmean/nanstd.
-pathfinding,2026-04-15T12:00:00Z,SAFE,compute-bound,0,false-positive,Downgraded. CuPy .get() is required -- A* has no GPU kernel. Per-pixel .compute() is only 2 calls for start/goal validation. seg.values in multi_stop_search collects already-computed results for stitching.
+pathfinding,2026-07-08,RISKY,compute-bound,1,3660,"HIGH (#3660/PR #3666, fixed in-tree): multi_stop_search materialized the full grid on dask (eager np.full + per-segment .compute + full-array reads in _optimize_waypoint_order); 133.8MB->11.0MB peak on 2000x2000/250-chunks repro; now lazy da.where stitching + single-block _cost_at_pixel; peakmem asv benchmark added. MEDIUM (#3661/PR #3667, fixed in-tree): optimize_order built its symmetric cost matrix with N(N-1) A* runs instead of N(N-1)/2; mirrored now, snap=True keeps both directions. Verdict RISKY post-fix: raster access is chunk-bound (LRU 128-chunk cache) but sparse-A* frontier dicts (g_cost/parent/visited) scale with the explored region, so adversarial start/goal pairs on huge grids can still exhaust RAM; realistic corridor routing is fine. LOW (documented, not fixed): _nearest_neighbor_2opt recomputes full _tour_cost per 2-opt candidate instead of O(1) delta (matters only for hundreds of waypoints); _a_star_dask re-fetches invariant f_u friction inside the 8-neighbor loop; parent_ys/parent_xs use np.ones*NONE instead of np.full in the numba kernel; dask start/goal checks issue two separate scalar .compute()s; f_min recomputed per a_star_search call for dask friction (N-1 full nanmin scans per multi-stop route, noted in #3660 as follow-up). cupy backend is a documented CPU fallback (single get()/asarray round trip, executed OK on this GPU host); dask+cupy parity verified locally. cuda-available."
 perlin,2026-07-02,SAFE,compute-bound,0,3469,"Re-audit 2026-07-02 (CUDA host). Prior HIGH #3469 (dask.persist -> full array resident, WILL OOM) CONFIRMED FIXED in-tree: no persist calls in perlin.py; both _perlin_dask_numpy (~L136) and _perlin_dask_cupy (~L293) use dask.compute(min, ptp/max) sharing the named noise subgraph, so noise computed once per reduction and freed; two monkeypatch regression tests guard against persist returning. Graph probe (no compute) on 20000x20000/2000-chunks: 530 tasks/100 chunks = 5.3 tasks/chunk, linear fan-in, reductions are tree reductions bounded by chunk size -> SAFE. GPU register check: _perlin_gpu 36 regs/thread, _perlin_gpu_xy 40 regs/thread (no pressure); ~12 float locals < 20 so 22x22 block fine; cupy/dask+cupy paths stay on GPU, no host round trip. No .values/np.asarray on dask/cupy arrays; map_blocks meta is empty np.array(()) idiom (not a materialization). No new CRITICAL/HIGH/MEDIUM. LOW (not fixed): _gradient (L53) np.zeros then overwrites every element -> np.empty would suffice; _gradient uses nb.prange without parallel=True (degrades to range). 31 perlin tests pass incl GPU+dask+GPU. No PR opened (prior HIGH already fixed)."
 polygon_clip,2026-06-10,SAFE,graph-bound,0,3191,"crop=True picked tiny leading edge chunk as rasterize mask size -> 13169-task graph; fixed to max(rc),max(cc) -> 1045 tasks. crop=False/numpy/cupy clean. Cat1-5 clean. GPU+dask+cupy run-validated."
 polygonize,2026-06-12,RISKY,compute-bound,0,3303,"Pass 3 (2026-06-12): re-audit after #2817/#2913/#3041. 0 HIGH. 1 MEDIUM fixed (#3303): _compute_region_value_ranges ran a pure-Python per-pixel loop (95% of float chunk time; 0.283s of 0.299s on 1024x1024, float chunks ~30x int) and re-ran _calculate_regions on an already-labelled block; moved to jitted _region_ranges_scan + _polygonize_numpy_regions label reuse (0.299s -> 0.015s/chunk). Side fix: w_match/s_match flags were always-truthy (_is_close numba overload generator called from pure Python returns impl function); output-neutral by chunk geometry, now computed correctly in jit. Cat1/2 clean (dask.compute batching is the documented #2673 design). Cat3 validated on GPU: cupy int/float + dask+cupy run end-to-end, single documented transfer, no round-trip. Cat4/5 LOW unchanged: _calculate_regions_cupy per-unique-value labeling (low impact); per-polygon Python classify loop in _polygonize_chunk dominates only on pathological many-polygon chunks (788K polys -> 7.8s). Cat6 RISKY unchanged: driver accumulates O(total polygons); 32-chunk batches bound transient peak. 527 polygonize tests + 40 new pass."

benchmarks/benchmarks/pathfinding.py

@@ -1,8 +1,40 @@
-from xrspatial.pathfinding import a_star_search
+import numpy as np
+import xarray as xr
+
+from xrspatial.pathfinding import a_star_search, multi_stop_search
 
 from .common import get_xr_dataarray
 
 
+class MultiStopSearchDaskMemory:
+    """Memory contract of the dask multi-stop path (issue #3660).
+
+    Peak memory must scale with the chunk size and the explored
+    corridor, not the full grid.  The grid is large (128 MB) but the
+    waypoints are close together, so a regression back to eager
+    full-grid stitching shows up as a step change in peak memory while
+    the benchmark itself stays fast.
+    """
+
+    def setup(self):
+        n = 4000  # 4000 x 4000 float64 = 128 MB, chunks 500 x 500 = 2 MB
+        import dask.array as da
+        data = da.ones((n, n), chunks=(500, 500), dtype='float64')
+        self.agg = xr.DataArray(
+            data, dims=['y', 'x'], attrs={'res': (1.0, 1.0)})
+        self.agg['y'] = np.linspace(n - 1, 0, n)
+        self.agg['x'] = np.linspace(0, n - 1, n)
+        # 3 waypoints inside a 200-pixel corner neighbourhood
+        self.waypoints = [
+            (float(n - 1), 0.0),
+            (float(n - 101), 100.0),
+            (float(n - 201), 200.0),
+        ]
+
+    def peakmem_multi_stop_search(self):
+        multi_stop_search(self.agg, self.waypoints)
+
+
 class AStarSearch:
     params = ([10, 100, 300], [4, 8], ["numpy"])
     param_names = ("nx", "connectivity", "type")

xrspatial/pathfinding.py

@@ -1332,6 +1332,21 @@ def _segment_to_numpy(seg_data):
     return np.asarray(seg_data)
 
 
+def _cost_at_pixel(seg_data, py, px):
+    """Read the cost of one pixel of an a_star_search result as a float.
+
+    On dask backends this computes only the block containing the pixel
+    (a cheap delayed block from ``_path_to_dask_array``), so the full
+    segment is never materialized.
+    """
+    value = seg_data[py, px]
+    if da is not None and isinstance(value, da.Array):
+        value = value.compute()
+    if hasattr(value, 'get'):  # cupy scalar -> numpy
+        value = value.get()
+    return float(value)
+
+
 def _optimize_waypoint_order(surface, waypoints, barriers, x, y,
                              connectivity, snap, friction, search_radius):
     """Build pairwise cost matrix and solve TSP with fixed endpoints.
@@ -1350,10 +1365,11 @@ def _pair_cost(a, b):
             snap_start=snap, snap_goal=snap,
             friction=friction, search_radius=search_radius,
         )
-        seg_vals = _segment_to_numpy(seg.data)
         goal_py, goal_px = _get_pixel_id(waypoints[b], surface, x, y)
-        goal_cost = seg_vals[goal_py, goal_px]
-        return float(goal_cost) if np.isfinite(goal_cost) else INF
+        # Single-pixel read: on dask backends this computes only the
+        # block containing the goal instead of the whole segment.
+        goal_cost = _cost_at_pixel(seg.data, goal_py, goal_px)
+        return goal_cost if np.isfinite(goal_cost) else INF
 
     for i in range(n):
         dist[i][i] = 0.0
@@ -1399,6 +1415,10 @@ def multi_stop_search(surface: xr.DataArray,
     minimize total travel cost (TSP), keeping the first and last
     waypoints fixed.
 
+    Dask-backed surfaces are routed sparsely and the segments are
+    stitched lazily, so the full grid is never materialized in memory;
+    peak memory scales with the chunk size and the explored corridor.
+
     Parameters
     ----------
     surface : xr.DataArray or xr.Dataset
@@ -1497,7 +1517,11 @@ def multi_stop_search(surface: xr.DataArray,
             )
 
     # --- Segment-by-segment routing ---
-    path_data = np.full(surface.shape, np.nan, dtype=np.float64)
+    if _is_dask:
+        # Built lazily below; never allocate the full grid in memory.
+        path_data = None
+    else:
+        path_data = np.full(surface.shape, np.nan, dtype=np.float64)
     cumulative_cost = 0.0
     segment_costs = []
 
@@ -1512,44 +1536,68 @@ def multi_stop_search(surface: xr.DataArray,
             snap_start=snap, snap_goal=snap,
             friction=friction, search_radius=search_radius,
         )
-        seg_vals = _segment_to_numpy(seg.data)
 
         goal_py, goal_px = waypoint_pixels[i + 1]
 
-        # If snap is on, the actual goal pixel may differ from the
-        # requested one.  Find the pixel with maximum finite cost
-        # (the true goal of this segment).
-        if snap and not np.isfinite(seg_vals[goal_py, goal_px]):
-            finite = np.isfinite(seg_vals)
-            if finite.any():
-                max_idx = np.nanargmax(seg_vals)
-                goal_py, goal_px = np.unravel_index(max_idx, seg_vals.shape)
-                waypoint_pixels[i + 1] = (goal_py, goal_px)
-
-        seg_goal_cost = seg_vals[goal_py, goal_px]
-
-        if not np.isfinite(seg_goal_cost):
-            raise ValueError(
-                f"no path between waypoints {i} and {i + 1}")
-
-        mask = np.isfinite(seg_vals)
-        if i > 0:
-            # Don't overwrite the junction pixel (set by previous segment)
-            sp_y, sp_x = waypoint_pixels[i]
-            mask[sp_y, sp_x] = False
+        if _is_dask:
+            # Lazy stitching: reading the goal cost computes only the
+            # block containing it, and the cumulative-cost overlay stays
+            # chunked.  snap is rejected for dask inputs above, so the
+            # goal pixel is always the requested one.  Junction pixels
+            # need no masking: the overwriting value (segment-start cost
+            # 0 plus the cumulative offset) equals what the previous
+            # segment wrote there.
+            seg_goal_cost = _cost_at_pixel(seg.data, goal_py, goal_px)
+            if not np.isfinite(seg_goal_cost):
+                raise ValueError(
+                    f"no path between waypoints {i} and {i + 1}")
+
+            # Each segment adds one da.where layer over every chunk, so
+            # the task graph grows as n_waypoints * n_chunks elementwise
+            # tasks.  _MAX_WAYPOINTS bounds this; the eager alternative
+            # materializes the full grid.
+            shifted = seg.data + cumulative_cost  # NaN stays NaN
+            if path_data is None:
+                path_data = shifted
+            else:
+                path_data = da.where(
+                    da.isfinite(shifted), shifted, path_data)
+        else:
+            seg_vals = _segment_to_numpy(seg.data)
+
+            # If snap is on, the actual goal pixel may differ from the
+            # requested one.  Find the pixel with maximum finite cost
+            # (the true goal of this segment).
+            if snap and not np.isfinite(seg_vals[goal_py, goal_px]):
+                finite = np.isfinite(seg_vals)
+                if finite.any():
+                    max_idx = np.nanargmax(seg_vals)
+                    goal_py, goal_px = np.unravel_index(
+                        max_idx, seg_vals.shape)
+                    waypoint_pixels[i + 1] = (goal_py, goal_px)
+
+            seg_goal_cost = seg_vals[goal_py, goal_px]
+
+            if not np.isfinite(seg_goal_cost):
+                raise ValueError(
+                    f"no path between waypoints {i} and {i + 1}")
+
+            mask = np.isfinite(seg_vals)
+            if i > 0:
+                # Don't overwrite the junction pixel
+                # (set by previous segment)
+                sp_y, sp_x = waypoint_pixels[i]
+                mask[sp_y, sp_x] = False
+
+            path_data[mask] = seg_vals[mask] + cumulative_cost
 
-        path_data[mask] = seg_vals[mask] + cumulative_cost
         segment_costs.append(float(seg_goal_cost))
         cumulative_cost += seg_goal_cost
 
-    # Match the input's array type, like a_star_search does
-    if _is_dask:
-        chunks = surface_data.chunks
-        path_data = da.from_array(path_data, chunks=chunks)
-        if has_cuda_and_cupy() and is_dask_cupy(surface):
-            import cupy
-            path_data = path_data.map_blocks(cupy.asarray)
-    elif has_cuda_and_cupy() and is_cupy_array(surface_data):
+    # Match the input's array type, like a_star_search does.  On dask
+    # backends path_data is already a lazy array with the surface's
+    # chunking (and cupy blocks for dask+cupy).
+    if not _is_dask and has_cuda_and_cupy() and is_cupy_array(surface_data):
         import cupy
         path_data = cupy.asarray(path_data)
 

xrspatial/tests/test_pathfinding.py

@@ -266,6 +266,30 @@ def _make_raster(data, dims=None, res=None, backend='numpy', chunks=(3, 3)):
     return raster
 
 
+def _tracking_np_full(threshold_cells, large_allocs):
+    """Build an np.full replacement that records allocations of
+    ``threshold_cells`` cells or more into *large_allocs*.
+
+    Used to assert that dask code paths never materialise full-size
+    arrays.  Patch with ``patch('numpy.full', side_effect=...)``.
+    """
+    original_full = np.full
+
+    def tracking_full(shape, *args, **kwargs):
+        result = original_full(shape, *args, **kwargs)
+        if hasattr(shape, '__len__'):
+            total = 1
+            for s in shape:
+                total *= s
+        else:
+            total = shape
+        if total >= threshold_cells:
+            large_allocs.append(('full', shape))
+        return result
+
+    return tracking_full
+
+
 # -----------------------------------------------------------------------
 # Weighted A* tests — parametrized for numpy and dask
 # -----------------------------------------------------------------------
@@ -493,23 +517,11 @@ def test_dask_no_large_numpy_arrays():
     goal = (0.0, float(width - 1))
 
     # Track large numpy allocations
-    original_full = np.full
     large_allocs = []
 
-    def tracking_full(shape, *args, **kwargs):
-        result = original_full(shape, *args, **kwargs)
-        if hasattr(shape, '__len__'):
-            total = 1
-            for s in shape:
-                total *= s
-        else:
-            total = shape
-        if total >= height * width:
-            large_allocs.append(('full', shape))
-        return result
-
     from unittest.mock import patch
-    with patch('numpy.full', side_effect=tracking_full):
+    with patch('numpy.full',
+               side_effect=_tracking_np_full(height * width, large_allocs)):
         result = a_star_search(agg, start, goal, friction=friction_agg)
 
     # Result should be dask-backed
@@ -1165,6 +1177,91 @@ def test_multi_stop_dask_matches_numpy():
     )
 
 
+@pytest.mark.skipif(not has_dask_array(), reason="Requires dask.Array")
+def test_multi_stop_dask_no_large_numpy_arrays():
+    """Dask multi-stop should not materialise full-size numpy arrays.
+
+    Same guard as test_dask_no_large_numpy_arrays, but for
+    multi_stop_search (issue #3660): segments must be stitched lazily
+    and per-segment goal costs read from a single block.
+    """
+    height, width = 50, 60
+    data = np.ones((height, width))
+
+    agg = _make_raster(data, backend='dask+numpy', chunks=(25, 30))
+
+    wp0 = (float(height - 1), 0.0)
+    wp1 = (float(height // 2), float(width // 2))
+    wp2 = (0.0, float(width - 1))
+
+    large_allocs = []
+
+    from unittest.mock import patch
+    with patch('numpy.full',
+               side_effect=_tracking_np_full(height * width, large_allocs)):
+        result = multi_stop_search(agg, [wp0, wp1, wp2])
+
+    # Result should be dask-backed
+    assert isinstance(result.data, da.Array)
+
+    # No full-size arrays should have been allocated while routing
+    assert len(large_allocs) == 0, (
+        f"Unexpected large allocations: {large_allocs}")
+
+    # And it should still match the numpy backend
+    agg_np = _make_raster(data, backend='numpy')
+    expected = multi_stop_search(agg_np, [wp0, wp1, wp2])
+    np.testing.assert_allclose(
+        np.asarray(result.values),
+        expected.values,
+        equal_nan=True, atol=1e-10,
+    )
+    np.testing.assert_allclose(
+        result.attrs['total_cost'], expected.attrs['total_cost'],
+        atol=1e-10,
+    )
+    np.testing.assert_allclose(
+        result.attrs['segment_costs'], expected.attrs['segment_costs'],
+        atol=1e-10,
+    )
+
+
+@pytest.mark.skipif(not has_dask_array(), reason="Requires dask.Array")
+def test_multi_stop_optimize_order_dask_no_large_numpy_arrays():
+    """optimize_order on dask must not materialise full-size arrays either."""
+    height, width = 50, 60
+    data = np.ones((height, width))
+
+    agg = _make_raster(data, backend='dask+numpy', chunks=(25, 30))
+
+    wp0 = (float(height - 1), 0.0)
+    wp1 = (0.0, float(width - 1))
+    wp2 = (float(height - 1), float(width // 2))
+    wp3 = (0.0, float(width // 2))
+
+    large_allocs = []
+
+    from unittest.mock import patch
+    with patch('numpy.full',
+               side_effect=_tracking_np_full(height * width, large_allocs)):
+        result = multi_stop_search(
+            agg, [wp0, wp1, wp2, wp3], optimize_order=True)
+
+    assert isinstance(result.data, da.Array)
+    assert len(large_allocs) == 0, (
+        f"Unexpected large allocations: {large_allocs}")
+
+    # Optimization behavior itself must match the numpy backend
+    agg_np = _make_raster(data, backend='numpy')
+    expected = multi_stop_search(
+        agg_np, [wp0, wp1, wp2, wp3], optimize_order=True)
+    assert result.attrs['waypoint_order'] == expected.attrs['waypoint_order']
+    np.testing.assert_allclose(
+        result.attrs['total_cost'], expected.attrs['total_cost'],
+        atol=1e-10,
+    )
+
+
 @cuda_and_cupy_available
 def test_multi_stop_cupy_matches_numpy():
     """CuPy multi-stop should match numpy results."""