Skip to content

Fix multi_stop_search on cupy and dask+cupy backends#3637

Open
brendancol wants to merge 2 commits into
mainfrom
issue-3630
Open

Fix multi_stop_search on cupy and dask+cupy backends#3637
brendancol wants to merge 2 commits into
mainfrom
issue-3630

Conversation

@brendancol

Copy link
Copy Markdown
Contributor

Closes #3630

  • Segment results are now materialized through a _segment_to_numpy helper that computes dask arrays first and then .get()s cupy blocks, in both the stitching loop and _optimize_waypoint_order. This fixes the TypeError: Implicit conversion to a NumPy array is not allowed crash on dask+cupy input.
  • The stitched output is converted back to the input's array type (cupy for cupy input, dask with cupy blocks for dask+cupy, dask for dask+numpy), matching what a_star_search already does.
  • test_multi_stop_cupy_matches_numpy had 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(). Adds test_multi_stop_dask_cupy_matches_numpy covering the previously untested backend, including optimize_order=True.

Backend coverage: numpy unchanged, cupy and dask+numpy get output-type fixes, dask+cupy goes from crashing to working.

Test plan:

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 brendancol left a comment

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.py multi_stop_search docstring, 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:1495 the 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_numpy computes one full-size array per segment on dask inputs. Same pre-existing behavior as np.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, and np.asarray(seg.values) then hit cupy's implicit-conversion TypeError. Compute-then-get covers all four backends in one place, and _optimize_waypoint_order gets the same treatment so optimize_order=True works 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 brendancol left a comment

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

multi_stop_search crashes on dask+cupy input and returns numpy-backed output for cupy input

1 participant