WIP prefetch - #23901
Conversation
Adds an async prefetch pipeline for hybrid scan SplitScan tasks: per-split prefetch coroutines run on the same event loop as scan_node and read_chunk, offloading pruning and byte-range computation to the existing thread pool rather than a dedicated one. Pinned host memory is reserved via reserve_or_fail by default (HybridScanPrefetchMemoryMode. FAIL_FAST), never waiting on memory something else may need, with reserve_memory (WAIT) available as an opt-in alternative. Byte ranges are coalesced before issuing reads, and a per-producer ordering mechanism (HybridScanPrefetchOrderingMode.ORDERED, default) keeps reservation attempts in consumption order so a split due for consumption soon can't lose its reservation to one that isn't. PrefetchedByteRanges covers both HybridScanPassMode variants (a single combined batch for SINGLE_PASS, separate filter/payload batches for TWO_PASS). Demotion of pinned memory under pressure is deliberately left as a TODO for both memory modes rather than built speculatively.
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
github-actions
Bot
added
Python
4257c6f to
4b3a609
Compare
| # A handle already opened while sizing this file during footer prefetch | ||
| # (see `_prefetch_parquet_footers_for_paths`), reused here instead of | ||
| # paying for a second HTTP HEAD request via `remote_handle()`. Only an | ||
| # __init__-time argument, not a stored field. |
There was a problem hiding this comment.
Can we double check this? My understanding is that kvikio.RemoteFile.open_* does not issue an HTTP request so long as size is provided. All else equal, I'd much prefer to pass around metadata things like (url, size) rather than stateful, open file handles.
My hope / thought here is that sharing / reusing an open file handle shouldn't really have a performance advantage: the actual, physical HTTP requests induced by a .read(nbytes, offset) are potentially going to be split into smaller HTTP requests, so there shouldn't really be sharing of resources at the RemoteFile level.
There was a problem hiding this comment.
Yeah will do. This still left over from #23317 IIRC
| if not isinstance(first, SplitScan): | ||
| return False |
There was a problem hiding this comment.
This is a temporary limitation right? Is it just bindings for hybrid scan's multi-file API that's blocking us?
There was a problem hiding this comment.
Yes, temporary.
Is it just bindings for hybrid scan's multi-file API that's blocking us?
Yes Ill open up for review today and ping for reviews
| if first.total_splits > total_row_groups: | ||
| return False |
There was a problem hiding this comment.
Why do we require this condition for prefetching?
| while i < n: | ||
| group_start = offset | ||
| group_file_start = ranges[i].offset | ||
| group_file_end = group_file_start + ranges[i].size | ||
| offset += ranges[i].size | ||
| j = i + 1 | ||
| while j < n and ranges[j].offset == group_file_end: | ||
| group_file_end += ranges[j].size | ||
| offset += ranges[j].size | ||
| j += 1 | ||
| futures.append( | ||
| handle.pread( | ||
| host[group_start:offset], | ||
| size=group_file_end - group_file_start, | ||
| file_offset=group_file_start, | ||
| ) | ||
| ) | ||
| i = j |
There was a problem hiding this comment.
It seems like this loop is mixing some non-trivial endpoint adjustment logic and IO. Could we maybe split those? Then whatever function is actually driving the I/O would call the endpoint adjustment, followed by the handle.preads.
| from cudf_polars.streaming.io import SplitScan | ||
|
|
||
|
|
||
| @nvtx_annotate_cudf_polars(message="issue_pread_calls") |
There was a problem hiding this comment.
This is fine, but would we expect this to be very fast? IIUC, .pread() returns a Future, so this shouldn't really block & take time, right?
There was a problem hiding this comment.
In general it doesnt take much time. But I noticed a few slow calls, and decided to offload it
2 participants
Description
Checklist