Lazy assembly for hand_mfd dask path (#1416)#1417
Open
brendancol wants to merge 1 commit intomainfrom
Open
Conversation
The dask assembly phase in `hand_mfd._hand_mfd_dask` previously computed every tile via `blocks[iy, ix].compute()` into a list-of-lists of numpy arrays, then called `da.block(rows)`. Driver memory held the full grid at once, defeating the per-tile streaming used elsewhere in the hydro subpackage. Replace the eager block-list with `da.map_blocks` over `(flow_accum_da, elev_da)`, slicing `fractions_da` per tile inside the closure. The converged `boundaries` and `frac_bdry` snapshots are captured by the closure so each tile is recomputed lazily on demand. Matches the pattern already used by `hand_dinf._hand_dinf_dask`. The boundary-propagation phase already streamed one tile at a time, so it is unchanged. Add two regression tests: - `test_dask_result_is_lazy` checks the result preserves chunking. - `test_dask_assembly_does_not_materialize_all_tiles` checks the dask graph contains one task per output chunk rather than collapsing to a single materialized array.
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
1 participant
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.
Summary
hand_mfd._hand_mfd_dasknow assembles viada.map_blocksinstead ofcomputing every tile into a list and stitching with
da.block. Drivermemory scales with chunk size again, matching
hand_dinfand the restof the hydro subpackage.
time; only the assembly phase is changed.
chunk rather than collapsing to a single materialized array.
Closes #1416.
Test plan
pytest xrspatial/hydro/tests/test_hand_mfd.py— 20 passedpytest xrspatial/hydro/tests/— 774 passedon this branch.