Skip to content

Fix owa() dask path and batch wpm() validation into one scheduler pass#3158

Merged
brendancol merged 5 commits into
mainfrom
deep-sweep-performance-mcda-2026-06-10-01
Jun 10, 2026
Merged

Fix owa() dask path and batch wpm() validation into one scheduler pass#3158
brendancol merged 5 commits into
mainfrom
deep-sweep-performance-mcda-2026-06-10-01

Conversation

@brendancol

Copy link
Copy Markdown
Contributor

Closes #3150

  • _sort_descending called da.sort, which does not exist in dask, so owa() raised AttributeError on any dask-backed Dataset. It now rechunks the criterion axis to a single chunk and sorts each block with np.sort via map_blocks (exact, since the axis lives in one chunk; memory stays bounded by block size). meta is forwarded so dask+cupy blocks keep their cupy meta.
  • While testing the dask+cupy path I found the order-weights array was always numpy, which cupy kernels reject. owa() now copies the order weights to the device when the data is cupy-backed (plain or dask), so both cupy paths work instead of raising TypeError.
  • _check_wpm_positive ran float(da.nanmin(arr).compute()) once per criterion, i.e. one serialized scheduler pass per layer at call time. The nanmin reductions are now collected and run through a single dask.compute() call.

Backend coverage: numpy unchanged; dask+numpy fixed (owa) and faster validation (wpm); cupy and dask+cupy owa now work and were verified on a CUDA host.

Test plan:

  • pytest xrspatial/tests/test_mcda.py (182 passed, includes 7 new tests)
  • New dask tests: owa parity vs numpy, owa uniform-order-weights == wlc, wpm parity, wpm single-scheduler-pass assertion, wpm non-positive rejection on dask
  • New GPU tests (gated on cuda_and_cupy_available): owa cupy and dask+cupy parity, run on a CUDA host
  • Graph probe: owa on 2560x2560 / 256-chunks builds ~15 tasks per chunk, no fan-in blowup

@github-actions github-actions Bot added the performance PR touches performance-sensitive code label Jun 10, 2026

@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 owa() dask path and batch wpm() validation into one scheduler pass

Blockers (must fix before merge)

None.

Suggestions (should fix, not blocking)

  • xrspatial/tests/test_mcda.py (TestOWADask): the dask tests use even 10x10 chunks only. A ragged-chunk case (e.g. chunk({"y": 7, "x": 9}) on a 20x20 grid) would catch block-alignment mistakes in the rechunk+map_blocks sort, which is exactly the code this PR adds.

Nits (optional improvements)

  • xrspatial/mcda/combine.py:387 (_sort_block_descending): a one-line note that NaN ordering matches the eager -np.sort(-data) path (NaNs sort to the same end) would save the next reader from re-deriving it.

What looks good

  • The per-block sort is exact because the criterion axis is rechunked into a single chunk first, and meta=data._meta keeps dask+cupy blocks typed correctly.
  • The scheduler-pass regression test pins the actual behavior (one dask.compute call with all three reductions) instead of timing anything.
  • The order-weights device copy fixes a second crash (plain-cupy owa) that the issue did not even cover, with parity tests on both GPU backends.

Checklist

  • Algorithm matches reference (descending sort + order weights, parity vs eager path)
  • All implemented backends produce consistent results (numpy/dask verified in tests; cupy/dask+cupy gated tests ran on a CUDA host)
  • NaN handling is correct (NaN fixture in the dask parity test)
  • Edge cases are covered by tests (non-positive rejection on dask, uniform order weights == wlc)
  • Dask chunk boundaries handled correctly (axis fully contained per block)
  • No premature materialization (validation now a single batched compute; sort path stays lazy)
  • Benchmark exists or is not needed (no mcda benchmarks exist; not added here)
  • README feature matrix: no tier changes needed (owa GPU cells already marked experimental)
  • Docstrings present and accurate (helper docstrings added)

@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.

Follow-up review after b-commit:

  • Ragged-chunk suggestion: fixed. test_owa_dask_matches_numpy is now parametrized over even (10x10) and ragged (7x9) chunkings; both pass.
  • NaN-ordering nit: fixed with a note in _sort_block_descending explaining parity with the eager -np.sort(-data) expression.

No remaining findings. 183 tests pass locally, including the GPU-gated cases on a CUDA host.

(cherry picked from commit 9585cd2ddb0ed71d05275ec1baadb0bbba9fa917)
…e-mcda-2026-06-10-01

Conflicts:
- .claude/sweep-performance-state.csv: took main's newer geotiff/rasterize/
  reproject rows (LF endings) and re-inserted this branch's mcda row.

Semantic fixups in xrspatial/tests/test_mcda.py (auto-merged textually):
- Removed strict xfail markers that now XPASS: the three owa backend
  tests fixed by this PR, plus standardize cupy piecewise, standardize
  dask+cupy categorical (#3159) and constrain attrs (#3154), whose
  fixes merged before the #3156 xfails landed.
Main's #3157 landed the same owa dask-sort and cupy order-weights
fixes this branch carried, so the merge takes main's combine.py for
those hunks and drops this branch's _sort_block_descending helper.
The wpm single-scheduler-pass validation (#3150) and its tests are
unchanged. Test-comment conflicts resolved to main's side since the
owa fix credit belongs to #3157.
@brendancol brendancol merged commit f83ee50 into main Jun 10, 2026
7 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

performance PR touches performance-sensitive code

Projects

None yet

Development

Successfully merging this pull request may close these issues.

mcda: owa() crashes on dask input; wpm() validation runs one compute per criterion

1 participant