Repository navigation
feat(cudf): Serve KvikIO reads through AsyncDataCache - #18941
Open
kingcrimsontianyu wants to merge 29 commits into
Open
kingcrimsontianyu wants to merge 29 commits into
kingcrimsontianyu wants to merge 29 commits into
Conversation
✅ Deploy Preview for meta-velox canceled.
|
meta-cla
Bot
added
the
CLA Signed
kjmph
added a commit
to rapidsai/velox
that referenced
this pull request
Sep 11, 2026
Adapt Tianyu Liu's KvikIO/AsyncDataCache implementation from Velox PR facebookincubator#18941 to the split-reader and asynchronous lifetime interfaces. Reuse AsyncDataCache for KvikIO-backed Hive reads so repeated object ranges can be served from host memory. Route asynchronous device reads with absent cache keys to a dedicated executor, and reuse pinned staging buffers for cache-to-device copies. Keep BufferedInput on its existing reader path. When caching is disabled or a split is not cacheable, return the original KvikIO datasource so its native asynchronous and batched reads remain available. Retain read ownership through asynchronous completion, fence pending device copies before releasing their destinations, and reject short cache fills. Handle page-backed hits and cache-capacity fallback, and add cache traffic counters and focused host/device tests. Cache entries use the requested file offset rather than aligned cache pages. Pinned staging retains its per-thread, per-device high-water capacity until process exit; this does not add a bounded staging pool or host-cache-only preloading. Upstream-PR: facebookincubator#18941 Upstream-commit: a0347ef Co-authored-by: Tianyu Liu <tialiu@nvidia.com>
kjmph
added a commit
to kjmph/velox-testing
that referenced
this pull request
Sep 11, 2026
Support the cache reader adapted from Tianyu Liu's Velox PR #18941. This change provides launcher integration for that reader. Add an explicit kvikio-cache reader mode for workers that support the KvikIO cache datasource. Enable AsyncDataCache while retaining KvikIO reads, using the same host-memory budget and soft-affinity placement as buffered-cache. Keep kvikio as the cache-disabled mode. Apply the direct KvikIO defaults to both modes, reject AWS-only receive and TCP MSS options, and reconcile coordinator affinity when switching between cached and uncached readers. Test reader selection, generated catalog properties, worker memory budgets, affinity transitions, and incompatible launch options. Related-PR: facebookincubator/velox#18941
kjmph
added a commit
to rapidsai/velox
that referenced
this pull request
Sep 11, 2026
Our adaptation of PR facebookincubator#18941 routes synchronous device reads through the asynchronous executor and waits for completion. A cache-hit caller already occupying the connector executor can exhaust its threads and deadlock waiting for its own queued work. Perform synchronous reads inline while preserving destination-device selection and the CUDA completion fence. Leave asynchronous scheduling and cache selection unchanged. Cover both synchronous overloads on a one-thread executor, verify that synchronous reads do not enqueue work, and retain H2D completion checks. Test empty, EOF, overflowing and null-buffer requests. Fixes: 7c2cea7 Upstream-PR: facebookincubator#18941
kjmph
added a commit
to rapidsai/velox
that referenced
this pull request
Sep 14, 2026
Cache-backed GPU reads copy received bytes from pageable cache storage into pinned staging before transferring them to the device. Separate registration lifetime from cache-entry lifetime so cache fills and H2D can use the same backing storage. Add an optional allocator-owned backing interface to AsyncDataCache and a registered slab pool shared by the BufferedInput and KvikIO GPU readers. Suballocate stable, page-aligned slices without moving live entries. Grow regular slabs from 64 MiB through 128 MiB to 256 MiB, and use dedicated page-rounded blocks for larger contiguous requests. Inherit the allocating worker's memory policy and restrict slab reuse to its CPU NUMA node. Charge complete slabs to the existing root allocator and a process-wide registration budget, including unused space and pending registrations. Keep the feature disabled by default with a 32-GiB budget when enabled. Tiny or pre-existing pageable entries, budget pressure, and recoverable allocation or registration failures retain ordinary cache backing and the staging fallback without refetching received bytes. Other CUDA and UCX registrations are outside this budget. Return slices on eviction, invalidation, and failed fills without retaining idle cache pins. Cache clear preserves reusable prepared slabs; allocator pressure and shrink reclaim empty slabs and report only backing capacity actually freed. Keep CUDA calls and backing frees outside cache shard and pool locks. Retain cache pins and DMA leases through copy completion, and unregister before freeing slab storage. Tear down backing during quiescent cache shutdown while the allocator and CUDA remain available. Prepare new slabs synchronously before publishing their slices. Leave ordinary CPU caches, cache-off reads, and storage scheduling unchanged; this does not pre-register the entire cache or move live data between pinned and pageable tiers. Add counters for registration, fallback, pool capacity, and H2D. Cover shared slabs, geometric growth, fragmented reuse, clear/refill, memory pressure accounting, invalidation, registration failures, shutdown, concurrent leases, and outstanding DMA in both GPU readers. Related-PR: facebookincubator#18941
kjmph
added a commit
to rapidsai/velox
that referenced
this pull request
Sep 14, 2026
The KvikIO cache reader adapted from Velox PR facebookincubator#18941, with the shared executor introduced by #116, still performs synchronous host reads for device-read cache misses. Pending remote reads occupy executor threads, limiting fill concurrency even when the transport can service more requests. Submit misses through the delegate's caller-buffer host_read_async API. Keep pending host futures and exclusive-cache-entry waits off the executor, and dispatch ready reads independently for cache publication and H2D submission. Bridge the std::future interface with one readiness thread that sleeps when idle and polls pending work at 100-microsecond intervals. Deferred-only delegates remain supported on the executor. Bound admitted logical reads before allocating cache entries. Reuse a positive KVIKIO_REMOTE_IO_MAX_CONCURRENT_REQUESTS, with a limit of 64 when unset or zero. This window is separate from the transport's HTTP request slots and does not add a byte budget or increase GPU driver count, preload depth, or destination allocations. Publish exclusive entries only after complete, successful fills. Drain outstanding writes before releasing failed allocations, preserve the destination completion fence when a future is consumed or discarded, and retain registered sources through H2D. Keep synchronous calls inline and preserve capacity fallback, BufferedInput dispatch, and the original cache-off KvikIO datasource. Add counters for admission delay, active fills, exclusive waits, deferred delegates, and fallback. Cover concurrency beyond executor size, bounded admission, shared fills, out-of-order completion, short and failed reads, capacity fallback, and datasource/destination lifetime on future discard. Upstream-PR: facebookincubator#18941 Related-PR: #116
kjmph
added a commit
to rapidsai/velox
that referenced
this pull request
Sep 14, 2026
The KvikIO read-through cache requires an entry at the requested offset to cover the entire read. A larger request can invalidate a useful smaller entry and fetch its bytes again, even though BufferedInput already supports consuming smaller cached prefixes. Use AsyncDataCache's kAllowSmaller policy for host and device reads. Consume each cached prefix, continue lookup at the next offset, and fill only the remaining suffix on a miss. Support contiguous and page-run backing, and count the actual bytes served by each fragment rather than charging the full request repeatedly. Preserve inline synchronous reads and asynchronous device-read fills. Retain every registered fragment through the logical read's final H2D fence, including while a later suffix is pending. Drain submitted copies before releasing sources after a suffix failure or discarded future. If the remainder cannot be cached, bypass only that remainder without refetching an already consumed prefix. This is exact-offset fragment reuse, not general interval lookup or compaction of cached storage. Reads beginning inside an existing entry, or encountering a cached island beyond an uncached gap, can still fetch overlapping bytes. Keep oversized-request bypass, pool budgets, BufferedInput, and cache-off KvikIO behavior unchanged. Add reused-prefix counters and tests for exact delegate offsets and lengths, long fragment chains, mixed backing, concurrent unequal-size readers, suffix-only fallback, failed fills, and H2D source lifetime. Related-PR: facebookincubator#18941
kingcrimsontianyu
force-pushed
the
kvikio_async_v2
branch
from
September 21, 2026 21:32
a0347ef to
214d9f2
Compare
Co-authored-by: Kyle Hubert <kjmph@users.noreply.github.com>
Co-authored-by: Kyle Hubert <kjmph@users.noreply.github.com>
kingcrimsontianyu
force-pushed
the
kvikio_async_v2
branch
from
September 22, 2026 02:18
b180da2 to
2166ba2
Compare
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Collaborator
|
@kingcrimsontianyu is this ready for review? |
8 of 27 tasks
# Conflicts: # CMake/resolve_dependency_modules/cudf.cmake # velox/experimental/cudf/connectors/hive/CudfSplitReader.cpp
kingcrimsontianyu
added a commit
to kingcrimsontianyu/velox
that referenced
this pull request
Oct 13, 2026
kingcrimsontianyu
marked this pull request as ready for review
October 16, 2026 16:03
kingcrimsontianyu
requested review from
bdice,
devavret,
jinchengchenghh,
karthikeyann,
majetideepak,
mhaseeb123 and
shrshi
as code owners
October 16, 2026 16:03
This branch has not been deployed
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
4 participants
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.
PR description
The native KvikIO datasource bypasses Velox's I/O stack and has no cache. Every query therefore re-reads the same bytes from remote storage.
This PR wraps it in a
CachingDataSourcebacked byAsyncDataCache. Reads are served from cache entries (with(file, offset)as the key). On a miss, the bytes are read from remote storage into a new entry. A shorter entry at the same offset is reused as a prefix, and only the rest is read from remote storage (#19126). Device reads copy through a per-thread pinned buffer to the GPU.device_read_async()fills the cache on one thread pool and copies to the GPU on a second, smaller one (cudf.kvikio_cache_copy_threads, default 16).The wrapper is used only when the async data cache is enabled and the split is cacheable. Otherwise the native KvikIO source is returned unchanged. Also this PR forwards the split's
cacheableflag and the file size Presto sends, which avoids a HEAD request per file open.This PR also bumps KvikIO version, so that we can use per-NIC binding and non-temporal copy introduced in rapidsai/kvikio#1052.
Depends on #19126
This PR is extracted from the
WXD-launchintegration branch (https://github.com/rapidsai/velox/tree/WXD-launch) made by @kjmph.Performance result
wxd-upstreamis velox upstream780d58a0c(2026-10-11) merged with this PR and feat(ucx-exchange): Register the UCX transport and enable its build #18695.wxd-upstreamis the branch from feat(scheduler): Add deterministic bounded cache-affinity placement prestodb/presto#28622, which is presto upstream2f2aa624a6(2026-10-13) merged with feat(native): Select the cuDF UCX exchange transport per edge and start its Communicator prestodb/presto#28398.