Fix multi_stop_search on cupy and dask+cupy backends#3637
Open
brendancol wants to merge 2 commits into
Open
Conversation
Segment extraction assumed .get() meant cupy and everything else was safe for np.asarray(seg.values), which crashes on dask-of-cupy results (implicit cupy->numpy conversion). Route all backends through a _segment_to_numpy helper (compute dask, then .get() cupy) in both the stitching loop and _optimize_waypoint_order, and convert the stitched output back to the input's array type so cupy and dask inputs no longer come back numpy-backed. Fixes the no-op conversion expression in test_multi_stop_cupy_matches_numpy and adds a dask+cupy test.
brendancol
commented
Jul 3, 2026
brendancol
left a comment
Contributor
Author
There was a problem hiding this comment.
PR Review: Fix multi_stop_search on cupy and dask+cupy backends
Blockers (must fix before merge)
- none
Suggestions (should fix, not blocking)
-
xrspatial/pathfinding.pymulti_stop_searchdocstring, Returns section: now that the output matches the input's array type, saying so would make the contract explicit (a_star_search's docstring already does).
Nits (optional improvements)
-
xrspatial/pathfinding.py:1495the dask branch wraps a fully materialized numpy array back into a dask array, so the graph carries the whole stitched surface. That is inherent to the dense stitching design (each segment was already computed to numpy before this PR too), not a regression; worth remembering if multi_stop ever needs the sparse treatment a_star_search's dask path has. -
_segment_to_numpycomputes one full-size array per segment on dask inputs. Same pre-existing behavior asnp.asarray(seg.values), just centralized; noting it for a future out-of-core pass.
What looks good
- The helper fixes the actual failure mode:
hasattr(seg_data, 'get')was false for dask-of-cupy, andnp.asarray(seg.values)then hit cupy's implicit-conversion TypeError. Compute-then-get covers all four backends in one place, and_optimize_waypoint_ordergets the same treatment sooptimize_order=Trueworks too. - The no-op expression in
test_multi_stop_cupy_matches_numpy(x if not hasattr(x, 'get') else x) is gone; the test now asserts the output is cupy-backed, which is what let the downgrade slip through before. - New dask+cupy test executes on this host (CUDA available) and covers both the plain call and
optimize_order=True; 52 tests pass.
Checklist
- Algorithm matches reference/paper (no algorithmic change; extraction and output typing only)
- All implemented backends produce consistent results (numpy/cupy/dask+numpy/dask+cupy all tested)
- NaN handling is correct (unchanged)
- Edge cases are covered by tests
- Dask chunk boundaries handled correctly (output rechunked to the input's chunks)
- No premature materialization or unnecessary copies beyond the pre-existing dense stitching (see nits)
- Benchmark exists or is not needed (bug fix, no perf-critical change)
- README feature matrix updated (n/a, no new function; dask+cupy was already claimed via a_star_search)
- Docstrings present and accurate (see suggestion)
brendancol
commented
Jul 3, 2026
brendancol
left a comment
Contributor
Author
There was a problem hiding this comment.
PR Review follow-up
The docstring suggestion from the first pass is addressed in f989495: the Returns section now states the output is backed by the same array type as the input. The two nits document pre-existing dense-stitching behavior and stay as notes for a future out-of-core pass rather than changes in this PR. Tests re-run after the commit: 52 passed. No open findings.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Closes #3630
_segment_to_numpyhelper that computes dask arrays first and then.get()s cupy blocks, in both the stitching loop and_optimize_waypoint_order. This fixes theTypeError: Implicit conversion to a NumPy array is not allowedcrash on dask+cupy input.a_star_searchalready does.test_multi_stop_cupy_matches_numpyhad a no-op conversion expression (x if not hasattr(x, 'get') else x) that hid the numpy-backed output; it now asserts the output is cupy-backed and converts with.get(). Addstest_multi_stop_dask_cupy_matches_numpycovering the previously untested backend, includingoptimize_order=True.Backend coverage: numpy unchanged, cupy and dask+numpy get output-type fixes, dask+cupy goes from crashing to working.
Test plan:
pytest xrspatial/tests/test_pathfinding.py(52 passed on a CUDA host, cupy and dask+cupy tests executed)