geotiff: read_vrt(chunks=) builds a lazy dask graph#1822
Open
brendancol wants to merge 4 commits into
Open
Conversation
Before this change, read_vrt(chunks=...) materialised the full VRT mosaic on host RAM and then wrapped the resulting numpy array via .chunk(). chunks= gave no memory protection, and gpu=True + chunks= still assembled the entire mosaic on the CPU before moving to the device. The chunked path now parses the VRT XML once up front, builds one dask.delayed per destination chunk window, and assembles them via from_delayed + da.concatenate. Each task calls the existing VRT internal reader with its own window= so only the sources intersecting that window are decoded. With gpu=True each block calls cupy.asarray before returning, so the dask array is dask-on-cupy from the start. A task-count cap (50,000, matching read_geotiff_dask) refuses chunk grids that would build a scheduler-busting graph and suggests a larger chunks= value. attrs['vrt_holes'] is not populated for chunked reads; the GeoTIFFFallbackWarning still fires from each worker.
Contributor
There was a problem hiding this comment.
Pull request overview
This PR rewrites read_vrt(chunks=...) so it builds a real dask graph of per-chunk-window decode tasks instead of materialising the whole mosaic and then .chunk()-ing it (issue #1814). When chunks= is set, read_vrt now dispatches to a new _read_vrt_chunked that parses the VRT XML once, declares output dtype/coords/attrs up front, then schedules one dask.delayed(_vrt_chunk_read) per destination window (with a 50,000-task cap matching read_geotiff_dask). The eager path is unchanged.
Changes:
- Add
_read_vrt_chunkedand_vrt_chunk_readhelpers and dispatch fromread_vrtwhenchunksis set; drop the eager.chunk()post-pass. - Add a module-level
_vrt_sentinel_for_dtypemirroring the closure in the eager path, and a task-count cap (_MAX_VRT_DASK_CHUNKS = 50_000) with a helpful error message. - Add
xrspatial/geotiff/tests/test_vrt_lazy_chunks_1814.pycovering lazy construction, byte-parity with the eager path, the task cap, window+chunks, GPU+chunks (cupy meta), and multi-band+chunks.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 4 comments.
| File | Description |
|---|---|
| xrspatial/geotiff/init.py | New chunked dispatch and helpers (_read_vrt_chunked, _vrt_chunk_read, _vrt_sentinel_for_dtype), task-count cap, docstring update; eager post-decode .chunk() removed. |
| xrspatial/geotiff/tests/test_vrt_lazy_chunks_1814.py | New tests asserting lazy construction, eager/chunked parity, task cap, windowed chunks, GPU dask-cupy meta, and multi-band band-axis preservation. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+4322
to
+4348
| # float64 if any selected band has an integer dtype with a | ||
| # representable nodata sentinel (the eager path promotes that case | ||
| # on mask hits; declaring float64 up front keeps every block's | ||
| # dtype consistent with the dask array's metadata regardless of | ||
| # whether the chunk actually contains sentinel pixels). | ||
| effective_dtypes = [] | ||
| for vrt_band in selected_bands: | ||
| eff = vrt_band.dtype | ||
| for src in vrt_band.sources: | ||
| scaled = src.scale is not None and src.scale != 1.0 | ||
| offset = src.offset is not None and src.offset != 0.0 | ||
| if scaled or offset: | ||
| eff = np.dtype(np.float64) | ||
| break | ||
| effective_dtypes.append(eff) | ||
| declared_dtype = np.result_type(*effective_dtypes) | ||
|
|
||
| if declared_dtype.kind in ('u', 'i'): | ||
| promotes = False | ||
| for vrt_band in selected_bands: | ||
| if _vrt_sentinel_for_dtype(vrt_band.nodata, | ||
| declared_dtype) is not None: | ||
| promotes = True | ||
| break | ||
| if promotes: | ||
| declared_dtype = np.dtype(np.float64) | ||
|
|
Comment on lines
+4132
to
+4137
| from ._vrt import read_vrt as _read_vrt_internal | ||
|
|
||
| arr, vrt = _read_vrt_internal( | ||
| source, window=(r0, c0, r1, c1), band=band, | ||
| max_pixels=max_pixels, missing_sources=missing_sources, | ||
| ) |
Comment on lines
+4440
to
+4449
| # Surface the nodata sentinel for the selected band. The chunked | ||
| # path does not aggregate ``vrt.holes`` across tasks (per-task holes | ||
| # would need to be reduced by an extra delayed; not done here, see | ||
| # issue #1814 note in the docstring). | ||
| nodata_meta = None | ||
| if vrt.bands: | ||
| band_idx_for_nodata = band if band is not None else 0 | ||
| nodata_meta = vrt.bands[band_idx_for_nodata].nodata | ||
| if nodata_meta is not None: | ||
| attrs['nodata'] = nodata_meta |
Comment on lines
+4339
to
+4347
| if declared_dtype.kind in ('u', 'i'): | ||
| promotes = False | ||
| for vrt_band in selected_bands: | ||
| if _vrt_sentinel_for_dtype(vrt_band.nodata, | ||
| declared_dtype) is not None: | ||
| promotes = True | ||
| break | ||
| if promotes: | ||
| declared_dtype = np.dtype(np.float64) |
…th (#1822 review) Address Copilot review comments on PR #1822: (c) The eager read_vrt populates attrs['vrt_holes'] as the machine-readable partial-mosaic detection contract from #1734. The chunked path silently omitted it. Populate the attr from a parse-time os.path.exists sweep over every source referenced by the parsed VRT so callers switching from eager to chunked keep the contract. The check is a static approximation that catches the dominant missing-file case; codec-error holes still surface as per-task GeoTIFFFallbackWarning. (d) Document the static "any band declares nodata?" promotion check as an explicit approximation of the eager path's runtime mask.any(). The gate was already correct (no promotion when no band declares nodata) but the surrounding comment did not call out the parse-time-vs-runtime trade-off. Add a regression test pinning the no-nodata uint16 case to the source dtype. Defer the duplicated eager logic and N+1 VRT-XML re-parse via tracking issue #1825; TODO markers reference that issue at both call sites.
PR #1803 forwarded the caller's max_pixels to read_to_array inside read_vrt's source loop so a tiny VRT output cannot force a huge source decode (#1796). The output-window check at the source read enforces that correctly. A separate per-tile dimension check at the same call sites also consumed the caller's max_pixels, so a caller setting max_pixels as an output budget (e.g. 10_000) failed the per-tile sanity check on any normal source whose default tile size is 256x256 (= 65_536 pixels). Use MAX_PIXELS_DEFAULT for the per-tile dim check at the two call sites in _read_tiles (local) and _read_tiles_cog_http (HTTP). The output-window check at the same functions continues to enforce the user-supplied max_pixels, preserving the #1796 protection.
Contributor
Author
|
@copilot resolve the merge conflicts in this pull request |
Keep _read_vrt_chunked dispatch (handles gpu=True + chunks=) over the non-GPU-capable _read_vrt_dask added in #1807. Remove the now-dead _read_vrt_dask and _vrt_effective_dtype functions that were only reachable via the superseded dispatch branch. Auto-merged from main: _vrt.py (VRT resample window inverse #1704 + XML size cap #1815 + minimal source window #1821), test files test_read_vrt_lazy_chunks_1798.py, test_vrt_dstrect_resample_cap_1737.py, test_vrt_resample_window_inverse_1704.py, test_vrt_xml_size_cap_1815.py. Co-authored-by: brendancol <433221+brendancol@users.noreply.github.com>
Contributor
Resolved in 064baa1. The only conflict was in |
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 #1814.
Before this change, read_vrt(chunks=...) materialised the full mosaic on host RAM, then wrapped the array via .chunk(). chunks= gave no memory protection; gpu=True + chunks= still assembled the full mosaic on the CPU before moving to the device.
The chunked path now parses the VRT XML up-front, builds one dask.delayed per destination chunk window, and assembles them via from_delayed + da.concatenate. Each task calls _read_vrt_internal with its own window so only the sources intersecting that window are decoded. A task-count cap (50k, matching read_geotiff_dask) guards against chunks= values that would explode the scheduler.
attrs['vrt_holes'] is not populated for chunked reads; warnings still surface from the workers. Filing a follow-up issue if needed.
Tested: pytest xrspatial/geotiff/tests/test_vrt_lazy_chunks_1814.py