Skip to content

multi_stop_search materializes the full grid on dask backends, defeating the sparse A* design #3660

Description

@brendancol

a_star_search has a memory-safe dask path: sparse Python A* with an LRU chunk cache, and a lazy output assembled from per-chunk delayed blocks. There is even a regression test (test_dask_no_large_numpy_arrays) pinning that no full-size numpy arrays get allocated.

multi_stop_search undoes that. On a dask-backed surface it:

  • allocates path_data = np.full(surface.shape, np.nan, dtype=np.float64) eagerly (xrspatial/pathfinding.py line 1473), whatever the backend
  • calls .compute() on each segment's full-size lazy result inside the segment loop, via _segment_to_numpy (line 1488)
  • wraps the finished numpy array back into dask with da.from_array (line 1521), so the "dask" output is a wrapper around an array that already lived in RAM

_optimize_waypoint_order has the same problem: it materializes a full h x w numpy array per waypoint pair (N(N-1) of them) just to read one pixel.

Peak memory therefore scales with the full grid, not the chunk size. The point of the dask backend is out-of-core pathfinding on grids that don't fit in RAM; routing 3 waypoints across such a grid OOMs even though each individual a_star_search call is fine.

Repro, executed on this host (2000x2000 float64 = 32 MB, chunks 250x250 = 0.5 MB):

import tracemalloc
from unittest.mock import patch
import numpy as np, dask.array as da, xarray as xr
from xrspatial import a_star_search, multi_stop_search

H = W = 2000
raster = xr.DataArray(da.from_array(np.ones((H, W)), chunks=(250, 250)),
                      dims=['y', 'x'], attrs={'res': (1.0, 1.0)})
raster['y'] = np.linspace(H - 1, 0, H)
raster['x'] = np.linspace(0, W - 1, W)
wps = [(float(H - 1), 0.0), (float(H // 2), float(W // 2)), (0.0, float(W - 1))]

# np.full tracking patch (same approach as test_dask_no_large_numpy_arrays)
# omitted here for brevity; full script attached to the PR.
result = multi_stop_search(raster, wps)

Output:

full array size          : 32.0 MB
chunk size               : 0.50 MB
a_star_search  peak      : 30.1 MB, full-size np.full allocs: []
multi_stop_search peak   : 133.8 MB, full-size np.full allocs: [('np.full', (2000, 2000))]
multi peak / full array  : 4.18x
multi peak / chunk       : 268x

(The 30 MB a_star_search peak is frontier-dict overhead, which scales with the explored corridor, not the grid.)

Suggested fix: stitch segments lazily. Each segment's lazy array already carries the surface's chunking, so the cumulative-cost overlay can be da.where(da.isfinite(seg + offset), seg + offset, acc), and the per-segment goal cost only needs the block containing the goal pixel. Junction pixels need no special casing on this path: the overwriting value is identical (segment-start cost 0 plus the cumulative offset), and snap is already rejected for dask inputs. The same single-pixel read fixes _optimize_waypoint_order.

Related, not in scope here: with a dask friction surface, every a_star_search call recomputes f_min with a full nanmin pass over friction, so an N-waypoint route does N-1 full scans of the friction array (N(N-1) with optimize_order=True).

Backends affected: dask+numpy and dask+cupy. The numpy and cupy paths are in-memory by design and unaffected.

Found by the performance sweep (deep-sweep, 2026-07-08). OOM verdict for this path at the 30 TB target workload: WILL OOM.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    area:pathfindingArea: pathfindingbugSomething isn't workingdaskDask backend / chunked arraysoomOut-of-memory risk with large datasetsperformancePR touches performance-sensitive codeseverity:highSweep finding: HIGHsweep-performanceFound by /sweep-performance

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions