diff --git a/CMake/resolve_dependency_modules/cudf.cmake b/CMake/resolve_dependency_modules/cudf.cmake index 5874ce383e9..447e23d9a08 100644 --- a/CMake/resolve_dependency_modules/cudf.cmake +++ b/CMake/resolve_dependency_modules/cudf.cmake @@ -40,12 +40,12 @@ set( set(VELOX_rmm_SOURCE_URL "https://github.com/rapidsai/rmm/archive/${VELOX_rmm_COMMIT}.tar.gz") velox_resolve_dependency_url(rmm) -# kvikio commit 257d08b from 2026-10-01 (main branch) +# kvikio commit 4114a21 from 2026-10-07 (main branch) set(VELOX_kvikio_VERSION 26.12) -set(VELOX_kvikio_COMMIT 257d08b41ef4e566cba6f5f64cb365e681961a3b) +set(VELOX_kvikio_COMMIT 4114a218184956f4bb652f2028980269d709a7cd) set( VELOX_kvikio_BUILD_SHA256_CHECKSUM - 15117616e0f7228e2324e0e64c21f3fa9305ffbe6b92d842f967e9b2ceb47bfc + 161e986cc0e253c32cdc196c18bc2cdeeae51347b672aee7616a572691b6998a ) set( VELOX_kvikio_SOURCE_URL diff --git a/velox/common/caching/AsyncDataCache.cpp b/velox/common/caching/AsyncDataCache.cpp index c02b6f95928..6b9a63229d8 100644 --- a/velox/common/caching/AsyncDataCache.cpp +++ b/velox/common/caching/AsyncDataCache.cpp @@ -251,7 +251,8 @@ std::unique_ptr CacheShard::getFreeEntryLocked() { std::optional CacheShard::lookupLocked( RawFileCacheKey key, uint64_t size, - folly::SemiFuture* wait) { + folly::SemiFuture* wait, + CacheEntrySizePolicy sizePolicy) { ++eventCounter_; auto it = entryMap_.find(key); if (it == entryMap_.end()) { @@ -266,7 +267,8 @@ std::optional CacheShard::lookupLocked( return CachePin{}; } // size=0 (from find()) always passes since entry size is non-negative. - if (foundEntry->size() < size) { + if (sizePolicy == CacheEntrySizePolicy::kRequireAtLeast && + foundEntry->size() < size) { // This can happen if different load quanta apply to access via different // connectors. This is not an error but still worth logging. VELOX_CACHE_LOG_EVERY_MS(WARNING, 1'000) @@ -296,11 +298,12 @@ CachePin CacheShard::findOrCreate( RawFileCacheKey key, uint64_t size, bool contiguous, - folly::SemiFuture* wait) { + folly::SemiFuture* wait, + CacheEntrySizePolicy sizePolicy) { AsyncDataCacheEntry* entryToInit = nullptr; { std::lock_guard l(mutex_); - auto result = lookupLocked(key, size, wait); + auto result = lookupLocked(key, size, wait, sizePolicy); if (result.has_value()) { return std::move(result.value()); } @@ -874,9 +877,10 @@ CachePin AsyncDataCache::findOrCreate( RawFileCacheKey key, uint64_t size, bool contiguous, - folly::SemiFuture* wait) { + folly::SemiFuture* wait, + CacheEntrySizePolicy sizePolicy) { const int shard = std::hash()(key) & shardMask_; - return shards_[shard]->findOrCreate(key, size, contiguous, wait); + return shards_[shard]->findOrCreate(key, size, contiguous, wait, sizePolicy); } std::optional AsyncDataCache::find( diff --git a/velox/common/caching/AsyncDataCache.h b/velox/common/caching/AsyncDataCache.h index 40d9a4b4ed1..9bd8c59da94 100644 --- a/velox/common/caching/AsyncDataCache.h +++ b/velox/common/caching/AsyncDataCache.h @@ -120,6 +120,13 @@ struct RawFileCacheKey { return offset == other.offset && fileNum == other.fileNum; } }; + +/// Controls whether an entry that begins at the requested key but is shorter +/// than the requested size is stale or is a reusable prefix. +enum class CacheEntrySizePolicy : uint8_t { + kRequireAtLeast, + kAllowSmaller, +}; } // namespace facebook::velox::cache namespace std { @@ -625,12 +632,15 @@ class CacheShard { : cache_(cache), maxWriteRatio_(maxWriteRatio) {} /// See AsyncDataCache::findOrCreate. If 'contiguous' is true, the - /// entry's data is allocated as a single contiguous region. + /// entry's data is allocated as a single contiguous region. 'sizePolicy' + /// decides how an existing entry at 'key' that is shorter than 'size' is + /// treated. See AsyncDataCache::findOrCreate for details. CachePin findOrCreate( RawFileCacheKey key, uint64_t size, bool contiguous = false, - folly::SemiFuture* readyFuture = nullptr); + folly::SemiFuture* readyFuture = nullptr, + CacheEntrySizePolicy sizePolicy = CacheEntrySizePolicy::kRequireAtLeast); /// Finds a cache entry for 'key'. Returns a shared-mode pin if the entry /// exists and is not exclusive. Returns an empty pin (inside optional) if @@ -731,7 +741,8 @@ class CacheShard { std::optional lookupLocked( RawFileCacheKey key, uint64_t size, - folly::SemiFuture* waitFuture); + folly::SemiFuture* waitFuture, + CacheEntrySizePolicy sizePolicy = CacheEntrySizePolicy::kRequireAtLeast); void tryAddFreeEntry(std::unique_ptr&& entry); @@ -895,11 +906,22 @@ class AsyncDataCache : public memory::Cache { /// When the future is realized, the caller may retry findOrCreate(). /// runtime error with code kNoCacheSpace if there is no space to /// create the new entry after evicting any unpinned content. + /// + /// 'sizePolicy' decides how an existing shared entry at 'key' that is + /// shorter than 'size' is treated. With kRequireAtLeast (the default) the + /// entry is stale: it is evicted and a new exclusive entry of 'size' is + /// created for the caller to fill. With kAllowSmaller the shorter entry is + /// returned as a valid prefix and counted as a hit. The cache does not + /// track the uncovered suffix. The caller must read the returned entry's + /// size(), and request [key.offset + size(), key.offset + 'size') under a + /// separate key, which may in turn hit, miss, or return another shorter + /// prefix. CachePin findOrCreate( RawFileCacheKey key, uint64_t size, bool contiguous = false, - folly::SemiFuture* waitFuture = nullptr); + folly::SemiFuture* waitFuture = nullptr, + CacheEntrySizePolicy sizePolicy = CacheEntrySizePolicy::kRequireAtLeast); /// Finds a cache entry for 'key'. Returns a shared-mode pin if the entry /// exists and is not exclusive. Returns an empty pin (inside optional) if @@ -978,9 +1000,11 @@ class AsyncDataCache : public memory::Cache { const std::vector& keys, const SizeFunc& sizeFunc, const ProcessPin& processPin, - bool contiguous = false) { + bool contiguous = false, + CacheEntrySizePolicy sizePolicy = CacheEntrySizePolicy::kRequireAtLeast) { for (size_t i = 0; i < keys.size(); ++i) { - auto pin = findOrCreate(keys[i], sizeFunc(i), contiguous); + auto pin = + findOrCreate(keys[i], sizeFunc(i), contiguous, nullptr, sizePolicy); if (pin.empty() || pin.checkedEntry()->isShared()) { continue; } diff --git a/velox/common/caching/tests/AsyncDataCacheTest.cpp b/velox/common/caching/tests/AsyncDataCacheTest.cpp index 7df81482038..b6f29458b6a 100644 --- a/velox/common/caching/tests/AsyncDataCacheTest.cpp +++ b/velox/common/caching/tests/AsyncDataCacheTest.cpp @@ -1291,7 +1291,29 @@ TEST_P(AsyncDataCacheTest, staleEntry) { ASSERT_EQ(stats.numEntries, 1); ASSERT_EQ(stats.numHit, 1); + // A caller that can assemble one read from several consecutive cache entries + // may consume this entry as a valid prefix instead of invalidating and + // reloading it. + const auto beforePrefix = cache_->refreshStats(); + auto prefixPin = cache_->findOrCreate( + key, + 2 * size, + /*contiguous=*/false, + &wait, + CacheEntrySizePolicy::kAllowSmaller); + ASSERT_FALSE(prefixPin.empty()); + ASSERT_TRUE(wait.isReady()); + ASSERT_FALSE(prefixPin.entry()->isExclusive()); + ASSERT_EQ(prefixPin.entry()->size(), size); + prefixPin.clear(); + stats = cache_->refreshStats(); + ASSERT_EQ(stats.numEntries, 1); + auto delta = stats - beforePrefix; + ASSERT_EQ(delta.numStales, 0); + ASSERT_EQ(delta.numHit, 1); + // Stale cache access with large cache size. + const auto beforeStale = cache_->refreshStats(); auto stalePin = cache_->findOrCreate(key, 2 * size, /*contiguous=*/false, &wait); ASSERT_FALSE(stalePin.empty()); @@ -1299,9 +1321,10 @@ TEST_P(AsyncDataCacheTest, staleEntry) { ASSERT_TRUE(stalePin.entry()->isExclusive()); stalePin.entry()->setExclusiveToShared(); stats = cache_->refreshStats(); - ASSERT_EQ(stats.numStales, 1); ASSERT_EQ(stats.numEntries, 1); - ASSERT_EQ(stats.numHit, 1); + delta = stats - beforeStale; + ASSERT_EQ(delta.numStales, 1); + ASSERT_EQ(delta.numHit, 0); } TEST_P(AsyncDataCacheTest, shrinkCache) { diff --git a/velox/experimental/cudf/CudfConfig.h b/velox/experimental/cudf/CudfConfig.h index 999796bd1af..4ddbf7e2486 100644 --- a/velox/experimental/cudf/CudfConfig.h +++ b/velox/experimental/cudf/CudfConfig.h @@ -65,6 +65,8 @@ struct CudfConfig { static constexpr const char* kUcxExchangeLogLevel{"cudf.exchange_log_level"}; static constexpr const char* kUcxPartitionedOutputBatchRows{ "cudf.partitioned_output_batch_rows"}; + static constexpr const char* kCudfKvikioCacheCopyThreads{ + "cudf.kvikio_cache_copy_threads"}; /// Query session configs for the cuDF Operators. static constexpr const char* kCudfTopNBatchSize{"cudf.topk_batch_size"}; @@ -121,6 +123,10 @@ struct CudfConfig { /// disable accumulation. int64_t partitionedOutputBatchRows{10'000}; + /// Threads that copy KvikIO data from the async data cache to the GPU. One + /// pool for the whole process, created by the first cached GPU read. + int32_t kvikioCacheCopyThreads{16}; + /// Memory resource for cuDF. /// Possible values are (cuda, pool, async, arena, managed, managed_pool). std::string memoryResource{"async"}; diff --git a/velox/experimental/cudf/connectors/hive/CMakeLists.txt b/velox/experimental/cudf/connectors/hive/CMakeLists.txt index 60dd4f859c4..6c7fc0c342d 100644 --- a/velox/experimental/cudf/connectors/hive/CMakeLists.txt +++ b/velox/experimental/cudf/connectors/hive/CMakeLists.txt @@ -15,6 +15,7 @@ add_library( velox_cudf_hive_connector BufferedInputDataSource.cpp + CachingDataSource.cpp CudfHiveConfig.cpp CudfHiveConnector.cpp CudfHiveConnectorSplit.cpp diff --git a/velox/experimental/cudf/connectors/hive/CachingDataSource.cpp b/velox/experimental/cudf/connectors/hive/CachingDataSource.cpp new file mode 100644 index 00000000000..c2e5975abf1 --- /dev/null +++ b/velox/experimental/cudf/connectors/hive/CachingDataSource.cpp @@ -0,0 +1,606 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "velox/experimental/cudf/CudfConfig.h" +#include "velox/experimental/cudf/connectors/hive/CachingDataSource.h" + +#include "velox/common/caching/FileIds.h" + +#include +#include +#include + +#include +#include + +#include + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +// A cuDF datasource for a file on S3 that serves every read (host_read, +// host_read_async, device_read and device_read_async) from Velox's async data +// cache, and does remote reads through KvikIO for the bytes it lacks. +// +// The requested byte range is served piece by piece, starting at its offset. +// Each piece comes from one async data cache entry, looked up by (file, offset +// of the piece): +// - Entry found: copy its bytes. The entry may be shorter than the rest of +// the range (CacheEntrySizePolicy::kAllowSmaller). The next piece then +// starts where this entry ends. +// - No entry: do a remote read of the rest of the range into a new entry, +// then copy its bytes. +// - Entry being filled by another read: wait for it, then look it up again. +// - The async data cache cannot hold the piece (not enough free memory for +// it, or 2 GiB or more, the most an entry can hold): do a remote read of +// the rest of the range without the async data cache. +// +// Device reads copy each piece from host memory into a pinned buffer, then to +// the GPU. device_read_async() uses two thread pools: a fill thread brings the +// range into the async data cache, waiting for its remote reads, then a copy +// thread copies it to the GPU. The returned future is ready once the bytes are +// in GPU memory. + +namespace facebook::velox::cudf_velox::connector::hive { +namespace { + +// AsyncDataCache stores entry sizes in an int32_t. Larger pieces are read +// without the cache. +constexpr size_t kMaxCacheEntryBytes = std::numeric_limits::max(); + +// Threads that put ranges in the cache. A fill thread waits while KvikIO reads +// a miss. The thread count equals KVIKIO_REMOTE_IO_MAX_CONCURRENT_REQUESTS +// (default 256), or 64 when that is 0 (concurrency unlimited). +folly::Executor* fillExecutor() { + static auto* executor = [] { + size_t threads = 256; + if (const char* value = + std::getenv("KVIKIO_REMOTE_IO_MAX_CONCURRENT_REQUESTS")) { + threads = folly::to(value); + } + if (threads == 0) { + threads = 64; + } + return new folly::CPUThreadPoolExecutor( + threads, std::make_shared("CudfCacheFill")); + }(); + return executor; +} + +// Threads that copy a range to the GPU, each through its own pinned buffer. +// The thread count is cudf.kvikio_cache_copy_threads, kept low to limit pinned +// memory. A copy thread waits for a remote read only for the part of a range +// that the async data cache cannot hold. +folly::Executor* copyExecutor() { + static auto* executor = [] { + const auto threads = CudfConfig::getInstance().kvikioCacheCopyThreads; + VELOX_USER_CHECK_GT( + threads, + 0, + "{} must be positive", + CudfConfig::kCudfKvikioCacheCopyThreads); + return new folly::CPUThreadPoolExecutor( + threads, std::make_shared("CudfCacheCopy")); + }(); + return executor; +} + +int deviceOf(cuda::stream_ref stream) { + int device{}; +#if defined(CUDART_VERSION) && CUDART_VERSION >= 12080 + CUDF_CUDA_TRY(cudaStreamGetDevice(stream.get(), &device)); +#else + CUDF_CUDA_TRY(cudaGetDevice(&device)); +#endif + return device; +} + +// Blocks until every copy queued on 'stream' has finished. The caller may then +// use or free the destination. +void waitForStream(cuda::stream_ref stream, int device) { + const rmm::cuda_set_device_raii onDevice{rmm::cuda_device_id{device}}; + try { + CUDF_CUDA_TRY(cudaStreamSynchronize(stream.get())); + } catch (...) { + // Make sure no copy still writes to the destination. + if (cudaDeviceSynchronize() != cudaSuccess) { + std::terminate(); + } + throw; + } +} + +// Pinned host memory used as a bounce buffer for copies from the async data +// cache to the GPU. Bytes read without the async data cache use it too. +class PinnedBuffer { + public: + // Waits until the last copy out of the buffer has finished, grows the buffer + // to at least 'bytes' if needed, and returns it. + uint8_t* reserve(size_t bytes) { + if (lastCopy_ != nullptr) { + CUDF_CUDA_TRY(cudaEventSynchronize(lastCopy_)); + } + if (size_ < bytes) { + if (data_ != nullptr) { + allocator_.deallocate_sync(data_, size_); + data_ = nullptr; + size_ = 0; + } + data_ = static_cast(allocator_.allocate_sync(bytes)); + size_ = bytes; + } + return data_; + } + + // Queues a copy of the first 'bytes' bytes of the buffer to 'dst' on the GPU. + void copyToGpu(size_t bytes, uint8_t* dst, cuda::stream_ref stream) { + CUDF_CUDA_TRY(cudaMemcpyAsync( + dst, data_, bytes, cudaMemcpyHostToDevice, stream.get())); + if (lastCopy_ == nullptr) { + CUDF_CUDA_TRY( + cudaEventCreateWithFlags(&lastCopy_, cudaEventDisableTiming)); + } + CUDF_CUDA_TRY(cudaEventRecord(lastCopy_, stream.get())); + } + + private: + rmm::host_device_async_resource_ref allocator_ = + cudf::get_pinned_memory_resource(); + uint8_t* data_{nullptr}; + size_t size_{0}; + // Recorded after the last copy out of 'data_'. + cudaEvent_t lastCopy_{nullptr}; +}; + +// One pinned buffer per thread and GPU, kept until the process exits. +PinnedBuffer& pinnedBuffer(int device) { + static thread_local std::unordered_map buffers; + return buffers[device]; +} + +// One device read: 'bytes' bytes at file 'offset' go to 'dst' on 'device', +// copied on 'stream'. +struct DeviceRead { + size_t offset; + size_t bytes; + uint8_t* dst; + cuda::stream_ref stream; + int device; +}; + +// Holds the KvikIO source and does the reads. Async work keeps it alive, since +// cuDF may destroy the CachingDataSource before that work finishes. +class Reader { + public: + Reader( + std::unique_ptr source, + std::string_view path, + cache::AsyncDataCache* dataCache, + std::shared_ptr stats) + : kvikioSource_(std::move(source)), + cache_(dataCache), + fileId_(fileIds(), std::string(path)), + ioStats_(std::move(stats)) { + VELOX_CHECK_NOT_NULL(kvikioSource_); + VELOX_CHECK_NOT_NULL(cache_); + } + + size_t fileSize() const { + return kvikioSource_->size(); + } + + bool isDeviceReadPreferred(size_t bytes) const { + return kvikioSource_->is_device_read_preferred(bytes); + } + + // Returns how many of the 'bytes' bytes at 'offset' are inside the file. + size_t trimToFile(size_t offset, size_t bytes) const { + const auto fileSize = kvikioSource_->size(); + return offset < fileSize ? std::min(bytes, fileSize - offset) : 0; + } + + // Puts the range from 'offset' in the cache, piece by piece, and returns a + // pin per piece. Stops at the first piece that cannot be cached. The rest of + // the range is then read without the cache. + std::vector cacheRange(size_t offset, size_t bytes) { + std::vector pins; + size_t cached = 0; + while (cached < bytes) { + auto pin = cachePiece(offset + cached, bytes - cached); + if (pin.empty()) { + break; + } + cached += pieceSize(pin, bytes - cached); + pins.push_back(std::move(pin)); + } + return pins; + } + + size_t readToHost(size_t offset, size_t size, uint8_t* dst) { + const auto bytes = trimToFile(offset, size); + if (bytes == 0) { + return 0; + } + VELOX_CHECK_NOT_NULL(dst); + size_t copied = 0; + while (copied < bytes) { + const auto wanted = bytes - copied; + auto pin = cachePiece(offset + copied, wanted); + if (pin.empty()) { + // Keep the pieces already copied. Read only the rest. + const auto bytesRead = + kvikioSource_->host_read(offset + copied, wanted, dst + copied); + countBytes("cudfKvikioCacheBypassBytes", bytesRead); + return copied + bytesRead; + } + const auto pieceBytes = pieceSize(pin, wanted); + copyFromEntry(pin, pieceBytes, dst + copied); + copied += pieceBytes; + } + return copied; + } + + std::unique_ptr readToHost( + size_t offset, + size_t size) { + std::vector bytes(trimToFile(offset, size)); + bytes.resize(readToHost(offset, bytes.size(), bytes.data())); + return cudf::io::datasource::buffer::create(std::move(bytes)); + } + + // Queues the GPU copies of the cached pieces in 'pins', from cacheRange(), + // then reads the rest of the range without the cache. Returns the number of + // bytes queued. + size_t copyRangeToGpu( + std::vector pins, + const DeviceRead& read) { + const rmm::cuda_set_device_raii onDevice{rmm::cuda_device_id{read.device}}; + size_t queued = 0; + try { + for (auto& pin : pins) { + const auto pieceBytes = pieceSize(pin, read.bytes - queued); + copyCachedPieceToGpu( + std::move(pin), pieceBytes, read.dst + queued, read); + queued += pieceBytes; + } + if (queued < read.bytes) { + queued += copyUncachedToGpu( + read.offset + queued, read.bytes - queued, read.dst + queued, read); + } + return queued; + } catch (...) { + // Wait for the copies already queued before rethrowing. + waitForStream(read.stream, read.device); + throw; + } + } + + private: + void countBytes(const std::string& name, size_t bytes) { + if (ioStats_) { + ioStats_->addCounter( + name, RuntimeCounter(bytes, RuntimeCounter::Unit::kBytes)); + } + } + + void count(const std::string& name) { + if (ioStats_) { + ioStats_->addCounter(name, RuntimeCounter(1)); + } + } + + // How many of the 'wanted' bytes the cache entry in 'pin' holds. An entry + // may be shorter than the request (CacheEntrySizePolicy::kAllowSmaller). + static size_t pieceSize(const cache::CachePin& pin, size_t wanted) { + const auto entrySize = pin.checkedEntry()->size(); + VELOX_CHECK_GT(entrySize, 0, "A cache entry must hold at least one byte"); + return std::min(wanted, entrySize); + } + + // Finds the cache entry at 'offset', or creates an empty one. Returns: + // - a shared pin: the entry is filled and ready to read; + // - an exclusive pin: a new, empty entry that the caller must fill; + // - an empty pin: the cache cannot be used for this piece now. + // When the empty pin is because another reader is filling the entry, and + // 'otherReaderFill' is not null, it is set to wait for that reader. + cache::CachePin findOrCreateEntry( + size_t offset, + size_t size, + folly::SemiFuture* otherReaderFill) { + cache::CachePin pin; + try { + pin = cache_->findOrCreate( + {fileId_.id(), offset}, + size, + /*contiguous=*/true, + otherReaderFill, + cache::CacheEntrySizePolicy::kAllowSmaller); + } catch (const VeloxRuntimeError& error) { + if (error.errorCode() != error_code::kNoCacheSpace) { + throw; + } + // The cache has no room. + return {}; + } + if (pin.empty()) { + // Another reader is filling the entry. + return pin; + } + auto* entry = pin.checkedEntry(); + entry->getAndClearFirstUseFlag(); + if (entry->isExclusive() && !entry->hasContiguousData()) { + // KvikIO can only read into one block of memory, and the new entry + // is not one block. + return {}; + } + return pin; + } + + // Returns a pin on a filled cache entry that starts at 'offset', reading the + // bytes from KvikIO on a miss. The entry may hold fewer than 'size' bytes. + // Returns an empty pin when the piece must be read without the cache. + cache::CachePin cachePiece(size_t offset, size_t size) { + // Too big for one async data cache entry. + if (size > kMaxCacheEntryBytes) { + return {}; + } + auto otherReaderFill = folly::SemiFuture::makeEmpty(); + auto pin = findOrCreateEntry(offset, size, &otherReaderFill); + // Another read is filling this entry. Wait for it, then look it up again. + if (pin.empty() && otherReaderFill.valid()) { + std::move(otherReaderFill) + .via(&folly::QueuedImmediateExecutor::instance()) + .wait(); + pin = findOrCreateEntry(offset, size, nullptr); + } + // The async data cache cannot be used for this piece: no room, or the + // entry is still being filled after the wait. + if (pin.empty()) { + return pin; + } + auto* entry = pin.checkedEntry(); + // Hit: the entry is already filled. It may hold fewer than 'size' bytes. + if (!entry->isExclusive()) { + const auto bytes = pieceSize(pin, size); + countBytes("cudfKvikioCacheHitBytes", bytes); + if (bytes < size) { + count("cudfKvikioCacheReusedPrefixes"); + countBytes("cudfKvikioCacheReusedPrefixBytes", bytes); + } + return pin; + } + // Miss: this read created the entry (exclusive) and fills it from S3. + const auto bytesRead = kvikioSource_->host_read( + offset, size, reinterpret_cast(entry->contiguousData())); + // Never make a partly filled entry readable. + VELOX_CHECK_EQ(bytesRead, size, "Short KvikIO read while filling cache"); + // Filled: make the entry readable by other reads (shared). + entry->setExclusiveToShared(); + countBytes("cudfKvikioCacheMissBytes", size); + return pin; + } + + static void + copyFromEntry(const cache::CachePin& pin, size_t bytes, uint8_t* dst) { + auto* entry = pin.checkedEntry(); + if (entry->hasContiguousData()) { + std::memcpy(dst, entry->contiguousData(), bytes); + return; + } + // Entries made by other readers may be split into several blocks. + for (const auto& block : entry->dataRanges(bytes)) { + std::memcpy(dst, block.data(), block.size()); + dst += block.size(); + } + } + + // Queues the GPU copy of 'bytes' bytes of a cached piece, through this + // thread's pinned buffer. + static void copyCachedPieceToGpu( + cache::CachePin pin, + size_t bytes, + uint8_t* dst, + const DeviceRead& read) { + auto& buffer = pinnedBuffer(read.device); + copyFromEntry(pin, bytes, buffer.reserve(bytes)); + // The bytes are in the pinned buffer. The cache may evict the entry. + pin.clear(); + buffer.copyToGpu(bytes, dst, read.stream); + } + + // Reads 'bytes' bytes at 'offset' from KvikIO, without the cache, and queues + // their GPU copy through this thread's pinned buffer. Returns the number of + // bytes read. + size_t copyUncachedToGpu( + size_t offset, + size_t bytes, + uint8_t* dst, + const DeviceRead& read) { + auto& buffer = pinnedBuffer(read.device); + const auto bytesRead = + kvikioSource_->host_read(offset, bytes, buffer.reserve(bytes)); + countBytes("cudfKvikioCacheBypassBytes", bytesRead); + buffer.copyToGpu(bytesRead, dst, read.stream); + return bytesRead; + } + + std::unique_ptr kvikioSource_; + cache::AsyncDataCache* cache_; + StringIdLease fileId_; + std::shared_ptr ioStats_; +}; + +// The datasource cuDF reads from. The reads themselves are done by Reader. +class CachingDataSource final : public cudf::io::datasource { + public: + CachingDataSource( + std::unique_ptr kvikioSource, + std::string_view path, + cache::AsyncDataCache* cache, + std::shared_ptr ioStats) + : reader_( + std::make_shared( + std::move(kvikioSource), + path, + cache, + std::move(ioStats))) {} + + size_t size() const override { + return reader_->fileSize(); + } + + bool supports_device_read() const override { + return true; + } + + bool is_device_read_preferred(size_t bytes) const override { + return reader_->isDeviceReadPreferred(bytes); + } + + std::unique_ptr host_read(size_t offset, size_t bytes) override { + return reader_->readToHost(offset, bytes); + } + + size_t host_read(size_t offset, size_t bytes, uint8_t* dst) override { + return reader_->readToHost(offset, bytes, dst); + } + + // The host reads run when cuDF calls get(), on cuDF's thread. + std::future> host_read_async( + size_t offset, + size_t bytes) override { + return std::async(std::launch::deferred, [reader = reader_, offset, bytes] { + return reader->readToHost(offset, bytes); + }); + } + + std::future host_read_async(size_t offset, size_t bytes, uint8_t* dst) + override { + return std::async( + std::launch::deferred, [reader = reader_, offset, bytes, dst] { + return reader->readToHost(offset, bytes, dst); + }); + } + + std::future device_read_async( + size_t offset, + size_t bytes, + uint8_t* dst, + cuda::stream_ref stream) override { + bytes = reader_->trimToFile(offset, bytes); + if (bytes == 0) { + std::promise nothingToRead; + nothingToRead.set_value(0); + return nothingToRead.get_future(); + } + VELOX_CHECK_NOT_NULL(dst); + const DeviceRead read{offset, bytes, dst, stream, deviceOf(stream)}; + // Ready once the copy thread has queued every GPU copy of the read. + auto copiesQueuedPromise = std::make_shared>(); + auto copiesQueued = copiesQueuedPromise->get_future(); + + // Put the range in the cache on a fill thread, then copy it to the GPU on + // a copy thread. + fillExecutor()->add([reader = reader_, read, copiesQueuedPromise] { + try { + auto pins = reader->cacheRange(read.offset, read.bytes); + copyExecutor()->add([reader, + read, + copiesQueuedPromise, + pins = std::move(pins)]() mutable { + try { + copiesQueuedPromise->set_value( + reader->copyRangeToGpu(std::move(pins), read)); + } catch (...) { + copiesQueuedPromise->set_exception(std::current_exception()); + } + }); + } catch (...) { + copiesQueuedPromise->set_exception(std::current_exception()); + } + }); + + // Return a future whose get() waits until the bytes are in GPU memory. + return std::async( + std::launch::deferred, + [copiesQueued = std::move(copiesQueued), read]() mutable { + const auto bytes = copiesQueued.get(); + waitForStream(read.stream, read.device); + return bytes; + }); + } + + size_t device_read( + size_t offset, + size_t bytes, + uint8_t* dst, + cuda::stream_ref stream) override { + bytes = reader_->trimToFile(offset, bytes); + if (bytes == 0) { + return 0; + } + VELOX_CHECK_NOT_NULL(dst); + const DeviceRead read{offset, bytes, dst, stream, deviceOf(stream)}; + // Runs on the caller's thread, not on the two pools. + const auto queued = + reader_->copyRangeToGpu(reader_->cacheRange(offset, bytes), read); + waitForStream(read.stream, read.device); + return queued; + } + + std::unique_ptr + device_read(size_t offset, size_t bytes, cuda::stream_ref stream) override { + const auto readSize = reader_->trimToFile(offset, bytes); + rmm::device_buffer result( + readSize, stream, cudf::get_current_device_resource_ref()); + const auto bytesRead = device_read( + offset, readSize, static_cast(result.data()), stream); + result.resize(bytesRead, stream); + return buffer::create(std::move(result)); + } + + private: + std::shared_ptr reader_; +}; + +} // namespace + +std::unique_ptr maybeCacheKvikioDataSource( + std::unique_ptr kvikioSource, + std::string_view path, + cache::AsyncDataCache* cache, + bool cacheable, + std::shared_ptr ioStats) { + if (cache == nullptr || !cacheable) { + return kvikioSource; + } + return std::make_unique( + std::move(kvikioSource), path, cache, std::move(ioStats)); +} + +} // namespace facebook::velox::cudf_velox::connector::hive diff --git a/velox/experimental/cudf/connectors/hive/CachingDataSource.h b/velox/experimental/cudf/connectors/hive/CachingDataSource.h new file mode 100644 index 00000000000..22b782519cd --- /dev/null +++ b/velox/experimental/cudf/connectors/hive/CachingDataSource.h @@ -0,0 +1,42 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "velox/common/caching/AsyncDataCache.h" +#include "velox/common/file/File.h" + +#include + +#include +#include + +namespace facebook::velox::cudf_velox::connector::hive { + +/// Returns 'kvikioSource' wrapped so that its reads go through the async data +/// cache, or 'kvikioSource' itself when the async data cache is null or the +/// split is not cacheable. Presto marks splits cacheable only under soft or +/// hard affinity scheduling (Hive config hive.node-selection-strategy, session +/// property node_selection_strategy). With Presto's default NO_PREFERENCE, the +/// async data cache is therefore never used. +std::unique_ptr maybeCacheKvikioDataSource( + std::unique_ptr kvikioSource, + std::string_view path, + cache::AsyncDataCache* cache, + bool cacheable, + std::shared_ptr ioStats = nullptr); + +} // namespace facebook::velox::cudf_velox::connector::hive diff --git a/velox/experimental/cudf/connectors/hive/CudfHiveConnectorSplit.cpp b/velox/experimental/cudf/connectors/hive/CudfHiveConnectorSplit.cpp index a7a3f5f5266..f924718ee44 100644 --- a/velox/experimental/cudf/connectors/hive/CudfHiveConnectorSplit.cpp +++ b/velox/experimental/cudf/connectors/hive/CudfHiveConnectorSplit.cpp @@ -19,6 +19,9 @@ #include +#include +#include +#include #include #include @@ -34,6 +37,29 @@ std::string stripFilePrefix(const std::string& targetPath) { } } // namespace +std::optional knownKvikioFileSize(const CudfHiveConnectorSplit& split) { + // The whole file's size, not this split's range. + const auto it = split.infoColumns.find("$file_size"); + if (it == split.infoColumns.end() || it->second.empty()) { + return std::nullopt; + } + const auto& text = it->second; + int64_t parsedSize; + const auto [end, error] = + std::from_chars(text.data(), text.data() + text.size(), parsedSize); + if (error != std::errc{} || end != text.data() + text.size() || + parsedSize < 0) { + return std::nullopt; + } + const auto fileSize = static_cast(parsedSize); + if (fileSize > std::numeric_limits::max() || split.start > fileSize || + (split.length != std::numeric_limits::max() && + split.length > fileSize - split.start)) { + return std::nullopt; + } + return static_cast(fileSize); +} + std::string CudfHiveConnectorSplit::toString() const { return fmt::format("CudfHive: {}", filePath); } @@ -53,8 +79,12 @@ CudfHiveConnectorSplit::CudfHiveConnectorSplit( uint64_t _start, uint64_t _length, int64_t _splitWeight, - const std::unordered_map& _infoColumns) - : facebook::velox::connector::ConnectorSplit(connectorId, _splitWeight), + const std::unordered_map& _infoColumns, + bool _cacheable) + : facebook::velox::connector::ConnectorSplit( + connectorId, + _splitWeight, + _cacheable), filePath(stripFilePrefix(_filePath)), start(_start), length(_length), @@ -69,6 +99,7 @@ std::shared_ptr CudfHiveConnectorSplit::create( const auto start = static_cast(obj["start"].asInt()); const auto length = static_cast(obj["length"].asInt()); const auto splitWeight = obj["splitWeight"].asInt(); + const auto cacheable = obj.getDefault("cacheable", true).asBool(); std::unordered_map infoColumns; for (const auto& [key, value] : obj["infoColumns"].items()) { @@ -76,7 +107,13 @@ std::shared_ptr CudfHiveConnectorSplit::create( } return std::make_shared( - connectorId, filePath, start, length, splitWeight, infoColumns); + connectorId, + filePath, + start, + length, + splitWeight, + infoColumns, + cacheable); } folly::dynamic CudfHiveConnectorSplit::serialize() const { @@ -86,6 +123,7 @@ folly::dynamic CudfHiveConnectorSplit::serialize() const { obj["start"] = start; obj["length"] = length; obj["splitWeight"] = splitWeight; + obj["cacheable"] = cacheable; folly::dynamic infoColumnsObj = folly::dynamic::object; for (const auto& [key, value] : infoColumns) { diff --git a/velox/experimental/cudf/connectors/hive/CudfHiveConnectorSplit.h b/velox/experimental/cudf/connectors/hive/CudfHiveConnectorSplit.h index 7bd186ab048..0338913c688 100644 --- a/velox/experimental/cudf/connectors/hive/CudfHiveConnectorSplit.h +++ b/velox/experimental/cudf/connectors/hive/CudfHiveConnectorSplit.h @@ -26,6 +26,7 @@ struct source_info; } // namespace cudf #include +#include #include #include @@ -50,7 +51,8 @@ struct CudfHiveConnectorSplit uint64_t _start = 0, uint64_t _length = std::numeric_limits::max(), int64_t _splitWeight = 0, - const std::unordered_map& _infoColumns = {}); + const std::unordered_map& _infoColumns = {}, + bool _cacheable = true); std::string toString() const override; std::string getFileName() const; @@ -67,6 +69,11 @@ struct CudfHiveConnectorSplit const folly::dynamic& obj); }; +/// The size of the whole file, which Presto sends with the split as +/// '$file_size'. Returns nullopt if it is missing, malformed, or smaller than +/// the split's range. +std::optional knownKvikioFileSize(const CudfHiveConnectorSplit& split); + class CudfHiveConnectorSplitBuilder { public: explicit CudfHiveConnectorSplitBuilder(std::string filePath) @@ -94,6 +101,11 @@ class CudfHiveConnectorSplitBuilder { return *this; } + CudfHiveConnectorSplitBuilder& cacheable(bool cacheable) { + cacheable_ = cacheable; + return *this; + } + CudfHiveConnectorSplitBuilder& connectorId(const std::string& connectorId) { connectorId_ = connectorId; return *this; @@ -101,7 +113,13 @@ class CudfHiveConnectorSplitBuilder { std::shared_ptr build() const { return std::make_shared( - connectorId_, filePath_, start_, length_, splitWeight_, infoColumns_); + connectorId_, + filePath_, + start_, + length_, + splitWeight_, + infoColumns_, + cacheable_); } private: @@ -110,6 +128,7 @@ class CudfHiveConnectorSplitBuilder { uint64_t length_{std::numeric_limits::max()}; std::string connectorId_; int64_t splitWeight_{0}; + bool cacheable_{true}; std::unordered_map infoColumns_ = {}; }; diff --git a/velox/experimental/cudf/connectors/hive/CudfHiveDataSource.cpp b/velox/experimental/cudf/connectors/hive/CudfHiveDataSource.cpp index 2d1f679978f..ae690bc576b 100644 --- a/velox/experimental/cudf/connectors/hive/CudfHiveDataSource.cpp +++ b/velox/experimental/cudf/connectors/hive/CudfHiveDataSource.cpp @@ -234,7 +234,8 @@ void CudfHiveDataSource::convertSplit(std::shared_ptr split) { .start(hiveSplit->start) .length(hiveSplit->length) .connectorId(hiveSplit->connectorId) - .splitWeight(hiveSplit->splitWeight); + .splitWeight(hiveSplit->splitWeight) + .cacheable(hiveSplit->cacheable); for (auto const& infoColumn : hiveSplit->infoColumns) { cudfHiveSplitBuilder.infoColumn(infoColumn.first, infoColumn.second); } diff --git a/velox/experimental/cudf/connectors/hive/CudfSplitReader.cpp b/velox/experimental/cudf/connectors/hive/CudfSplitReader.cpp index 010fa6f5a8d..f7ea0607531 100644 --- a/velox/experimental/cudf/connectors/hive/CudfSplitReader.cpp +++ b/velox/experimental/cudf/connectors/hive/CudfSplitReader.cpp @@ -16,6 +16,7 @@ #include "velox/experimental/cudf/CudfNoDefaults.h" #include "velox/experimental/cudf/connectors/hive/BufferedInputDataSource.h" +#include "velox/experimental/cudf/connectors/hive/CachingDataSource.h" #include "velox/experimental/cudf/connectors/hive/CudfSplitReader.h" #include "velox/experimental/cudf/exec/GpuResources.h" #include "velox/experimental/cudf/exec/VeloxCudfInterop.h" @@ -530,6 +531,33 @@ bool CudfSplitReader::hasSplitSpecificPushdownFilter() const { return hasSplitSpecificPushdownFilter_; } +namespace { + +std::unique_ptr makeKvikioDataSource( + const CudfHiveConnectorSplit& split, + const std::shared_ptr& ioStats) { + const auto knownSize = knownKvikioFileSize(split); + // Open the whole file (offset 0, no size estimate), not this split's range, + // because the reader also reads the footer at the end of the file. With the + // file size known, opening sends no HEAD request to S3. + auto source = cudf::io::datasource::create( + split.filePath, /*offset=*/0, /*max_size_estimate=*/0, knownSize); + if (ioStats) { + ioStats->addCounter( + knownSize.has_value() ? "cudfKvikioKnownSizeOpens" + : "cudfKvikioUnknownSizeOpens", + RuntimeCounter(1)); + } + return maybeCacheKvikioDataSource( + std::move(source), + split.filePath, + velox::cache::AsyncDataCache::getInstance(), + split.cacheable, + ioStats); +} + +} // namespace + void CudfSplitReader::setupCudfDataSource() { if (dataSource_) { return; @@ -551,9 +579,7 @@ void CudfSplitReader::setupCudfDataSource() { if (not useBufferedInput) { VLOG(1) << fmt::format( "Using KvikIO data source for file: {}", split_->filePath); - dataSource_ = std::move( - cudf::io::make_datasources(cudf::io::source_info{split_->filePath}) - .front()); + dataSource_ = makeKvikioDataSource(*split_, ioStats_); return; } @@ -581,9 +607,7 @@ void CudfSplitReader::setupCudfDataSource() { LOG(WARNING) << fmt::format( "Failed to generate file handle cache for file. Falling back to KvikIO. Path: {}", split_->filePath); - dataSource_ = std::move( - cudf::io::make_datasources(cudf::io::source_info{split_->filePath}) - .front()); + dataSource_ = makeKvikioDataSource(*split_, ioStats_); return; } @@ -615,9 +639,7 @@ void CudfSplitReader::setupCudfDataSource() { LOG(WARNING) << fmt::format( "Failed to create buffered input data source for file. Falling back to the KvikIO. Path: {}", split_->filePath); - dataSource_ = std::move( - cudf::io::make_datasources(cudf::io::source_info{split_->filePath}) - .front()); + dataSource_ = makeKvikioDataSource(*split_, ioStats_); return; } dataSource_ = diff --git a/velox/experimental/cudf/exec/ToCudf.cpp b/velox/experimental/cudf/exec/ToCudf.cpp index a3f7c2c6bfb..e39662b0b50 100644 --- a/velox/experimental/cudf/exec/ToCudf.cpp +++ b/velox/experimental/cudf/exec/ToCudf.cpp @@ -504,6 +504,10 @@ void CudfConfig::initialize( partitionedOutputBatchRows = folly::to(config[kUcxPartitionedOutputBatchRows]); } + if (config.find(kCudfKvikioCacheCopyThreads) != config.end()) { + kvikioCacheCopyThreads = + folly::to(config[kCudfKvikioCacheCopyThreads]); + } if (config.find(kCudfLogFallback) != config.end()) { logFallback = folly::to(config[kCudfLogFallback]); } diff --git a/velox/experimental/cudf/tests/CMakeLists.txt b/velox/experimental/cudf/tests/CMakeLists.txt index 3b5a12a6de5..295449206f0 100644 --- a/velox/experimental/cudf/tests/CMakeLists.txt +++ b/velox/experimental/cudf/tests/CMakeLists.txt @@ -119,6 +119,12 @@ velox_add_cudf_test( LIBS velox_cudf_exec_test_lib velox_cudf_hive_connector ) +velox_add_cudf_test( + NAME velox_cudf_caching_datasource_test + SOURCES Main.cpp CachingDataSourceTest.cpp + LIBS ${CUDF_TEST_DEFAULT_LIBS} velox_cudf_hive_connector +) + velox_add_cudf_test( NAME velox_cudf_decimal_aggregation_test SOURCES Main.cpp DecimalAggregationTest.cpp diff --git a/velox/experimental/cudf/tests/CachingDataSourceTest.cpp b/velox/experimental/cudf/tests/CachingDataSourceTest.cpp new file mode 100644 index 00000000000..1ff53f1ac9b --- /dev/null +++ b/velox/experimental/cudf/tests/CachingDataSourceTest.cpp @@ -0,0 +1,841 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "velox/experimental/cudf/CudfConfig.h" +#include "velox/experimental/cudf/connectors/hive/CachingDataSource.h" + +#include "velox/common/caching/FileIds.h" +#include "velox/common/caching/SsdCache.h" +#include "velox/common/memory/MallocAllocator.h" + +#include +#include + +#include +#include + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace facebook::velox::cudf_velox::connector::hive { +namespace { +using namespace std::chrono_literals; + +struct ReadState { + using Requests = std::vector>; + + void recordRead(size_t offset, size_t bytes) { + std::lock_guard lock(mutex); + requests.emplace_back(offset, bytes); + ++reads; + } + + Requests recordedReads() { + std::lock_guard lock(mutex); + return requests; + } + + std::string data = std::string(8192, 'x'); + std::atomic reads{0}; + std::atomic shortRead{false}; + std::atomic failRead{false}; + std::atomic destroyed{false}; + std::shared_future gate; + std::mutex mutex; + Requests requests; +}; + +class MemorySource : public cudf::io::datasource { + public: + explicit MemorySource(std::shared_ptr state) + : state_(std::move(state)) {} + ~MemorySource() override { + state_->destroyed = true; + } + size_t size() const override { + return state_->data.size(); + } + bool supports_device_read() const override { + return true; + } + bool is_device_read_preferred(size_t) const override { + return true; + } + std::unique_ptr host_read(size_t offset, size_t bytes) + override { + std::vector result( + offset < size() ? std::min(bytes, size() - offset) : 0); + result.resize(host_read(offset, result.size(), result.data())); + return datasource::buffer::create(std::move(result)); + } + size_t host_read(size_t offset, size_t bytes, uint8_t* dst) override { + state_->recordRead(offset, bytes); + if (state_->gate.valid()) { + state_->gate.wait(); + } + VELOX_CHECK(!state_->failRead.exchange(false), "Injected remote failure"); + bytes = offset < size() ? std::min(bytes, size() - offset) : 0; + if (state_->shortRead.exchange(false) && bytes > 0) { + --bytes; + } + if (bytes > 0) { + std::memcpy(dst, state_->data.data() + offset, bytes); + } + return bytes; + } + + private: + std::shared_ptr state_; +}; + +// Holds the reads of an AsyncMemorySource until the test completes them. +struct AsyncReadControl { + std::mutex mutex; + std::vector> pending; + bool open{false}; + std::atomic submitted{0}; + std::atomic failSubmission{false}; + + void add(std::function fill) { + { + std::lock_guard lock(mutex); + if (!open) { + pending.push_back(std::move(fill)); + ++submitted; + return; + } + } + ++submitted; + fill(); + } + + void completeNewest() { + std::function fill; + { + std::lock_guard lock(mutex); + if (pending.empty()) { + return; + } + fill = std::move(pending.back()); + pending.pop_back(); + } + fill(); + } + + void releaseAll() { + std::vector> fills; + { + std::lock_guard lock(mutex); + open = true; + fills.swap(pending); + } + for (auto& fill : fills) { + fill(); + } + } +}; + +class AsyncMemorySource : public MemorySource { + public: + AsyncMemorySource( + std::shared_ptr state, + std::shared_ptr control) + : MemorySource(state), state_(state), control_(std::move(control)) {} + + // Like cuDF's KvikIO datasource. + size_t host_read(size_t offset, size_t bytes, uint8_t* dst) override { + return host_read_async(offset, bytes, dst).get(); + } + + std::future host_read_async(size_t offset, size_t bytes, uint8_t* dst) + override { + VELOX_CHECK( + !control_->failSubmission.exchange(false), + "Injected submission failure"); + auto promise = std::make_shared>(); + auto future = promise->get_future(); + state_->recordRead(offset, bytes); + control_->add([state = state_, offset, bytes, dst, promise]() mutable { + try { + VELOX_CHECK( + !state->failRead.exchange(false), "Injected remote failure"); + if (state->shortRead.exchange(false) && bytes != 0) { + --bytes; + } + std::memcpy(dst, state->data.data() + offset, bytes); + promise->set_value(bytes); + } catch (...) { + promise->set_exception(std::current_exception()); + } + }); + return future; + } + + private: + std::shared_ptr state_; + std::shared_ptr control_; +}; + +bool waitUntil(const std::function& condition) { + const auto deadline = std::chrono::steady_clock::now() + 5s; + while (!condition() && std::chrono::steady_clock::now() < deadline) { + std::this_thread::sleep_for(1ms); + } + return condition(); +} + +struct StreamGate { + std::promise entered; + std::promise release; + std::shared_future released = release.get_future().share(); +}; + +void CUDART_CB waitForStreamGate(void* argument) { + auto* gate = static_cast(argument); + gate->entered.set_value(); + gate->released.wait(); +} + +class CachingDataSourceTest : public testing::Test { + protected: + static void SetUpTestSuite() { + // The executors are created once per process: 2 copy threads and 8 fill + // threads for all tests. + CudfConfig::getInstance().kvikioCacheCopyThreads = 2; + setenv("KVIKIO_REMOTE_IO_MAX_CONCURRENT_REQUESTS", "8", 1); + } + + std::unique_ptr source(bool cacheable = true) { + return maybeCacheKvikioDataSource( + std::make_unique(read_), + path_, + cache_.get(), + cacheable, + stats_); + } + std::string fromDevice(const uint8_t* device, size_t bytes) { + std::string result(bytes, '\0'); + CUDF_CUDA_TRY( + cudaMemcpy(result.data(), device, bytes, cudaMemcpyDeviceToHost)); + return result; + } + + std::unique_ptr asyncSource( + const std::shared_ptr& control) { + return maybeCacheKvikioDataSource( + std::make_unique(read_, control), + path_, + cache_.get(), + true, + stats_); + } + + void fillPattern(size_t bytes) { + read_->data.resize(bytes); + for (size_t i = 0; i < bytes; ++i) { + read_->data[i] = static_cast(i % 251); + } + } + + std::shared_ptr allocator_ = + std::make_shared( + memory::MemoryAllocator::Options{ + .capacity = 16 << 20, + .reservationByteLimit = 0}); + std::shared_ptr cache_ = + cache::AsyncDataCache::create(allocator_.get()); + std::shared_ptr read_ = std::make_shared(); + std::shared_ptr stats_ = std::make_shared(); + const std::string path_ = "s3://test/caching-datasource"; +}; + +TEST_F(CachingDataSourceTest, asyncFillsExceedCopyThreadsAndStayBounded) { + constexpr size_t kRequests = 24; + constexpr size_t kBytes = 4096; + fillPattern(kRequests * kBytes); + auto control = std::make_shared(); + auto input = asyncSource(control); + rmm::cuda_stream stream; + rmm::device_buffer destination( + read_->data.size(), stream, cudf::get_current_device_resource_ref()); + auto* dst = static_cast(destination.data()); + std::vector> reads; + // Complete the held reads on exit, even after an ASSERT failure. + auto release = folly::makeGuard([&] { control->releaseAll(); }); + for (size_t i = 0; i < kRequests; ++i) { + reads.push_back( + input->device_read_async(i * kBytes, kBytes, dst + i * kBytes, stream)); + } + ASSERT_TRUE(waitUntil([&] { return control->submitted.load() >= 8; })); + std::this_thread::sleep_for(20ms); + EXPECT_EQ(control->submitted, 8) + << "Eight fill threads, more than the two copy threads, run one fill each"; + control->releaseAll(); + for (auto& read : reads) { + EXPECT_EQ(read.get(), kBytes); + } + EXPECT_EQ(fromDevice(dst, read_->data.size()), read_->data); + EXPECT_EQ(read_->reads, kRequests); + EXPECT_EQ(control->submitted, kRequests); +} + +class CachingDataSourceFragmentTest : public CachingDataSourceTest, + public testing::WithParamInterface { +}; + +TEST_P(CachingDataSourceFragmentTest, deviceReadsReuseCachedFragments) { + const bool nonContiguous = GetParam(); + // Many fragments, with nonuniform data to catch incorrect source offsets or + // overwriting an earlier destination. + constexpr size_t kFragments = 40; + constexpr size_t kFragmentBytes = 4096; + constexpr size_t kTailBytes = 4096 + 731; + constexpr size_t kPrefixBytes = kFragments * kFragmentBytes; + constexpr size_t kBytes = kPrefixBytes + kTailBytes; + fillPattern(kBytes); + StringIdLease file(fileIds(), path_); + for (size_t offset = 0; offset < kPrefixBytes; offset += kFragmentBytes) { + auto pin = cache_->findOrCreate( + {file.id(), offset}, kFragmentBytes, !nonContiguous); + auto* entry = pin.checkedEntry(); + ASSERT_TRUE(entry->isExclusive()); + ASSERT_EQ(entry->hasContiguousData(), !nonContiguous); + size_t copied = 0; + for (const auto& range : entry->dataRanges(kFragmentBytes)) { + std::memcpy( + range.data(), read_->data.data() + offset + copied, range.size()); + copied += range.size(); + } + entry->setExclusiveToShared(); + } + auto control = std::make_shared(); + control->releaseAll(); + auto input = asyncSource(control); + rmm::cuda_stream stream; + rmm::device_buffer destination( + kBytes, stream, cudf::get_current_device_resource_ref()); + auto* dst = static_cast(destination.data()); + EXPECT_EQ(input->device_read_async(0, kBytes, dst, stream).get(), kBytes); + EXPECT_EQ(fromDevice(dst, kBytes), read_->data); + EXPECT_EQ( + read_->recordedReads(), + (ReadState::Requests{{kPrefixBytes, kTailBytes}})); + auto metrics = stats_->stats(); + EXPECT_EQ(metrics.at("cudfKvikioCacheHitBytes").sum, kPrefixBytes); + EXPECT_EQ(metrics.at("cudfKvikioCacheMissBytes").sum, kTailBytes); + EXPECT_EQ(metrics.at("cudfKvikioCacheReusedPrefixes").sum, kFragments); + EXPECT_EQ(metrics.at("cudfKvikioCacheReusedPrefixBytes").sum, kPrefixBytes); + + // Synchronous overloads and a smaller final fragment use the same chain. + EXPECT_EQ(input->device_read(0, kBytes - 17, dst, stream), kBytes - 17); + auto owned = input->device_read(0, kBytes, stream); + EXPECT_EQ(fromDevice(owned->data(), owned->size()), read_->data); + EXPECT_EQ(fromDevice(dst, kBytes - 17), read_->data.substr(0, kBytes - 17)); + EXPECT_EQ(read_->reads, 1); +} + +INSTANTIATE_TEST_SUITE_P( + CachingDataSourceFragmentTest, + CachingDataSourceFragmentTest, + testing::Values(false, true), + [](const testing::TestParamInfo& info) { + return info.param ? "nonContiguous" : "contiguous"; + }); + +TEST_F( + CachingDataSourceTest, + longerAsyncReadWaitsForPrefixThenFillsOnlySuffix) { + auto control = std::make_shared(); + auto input = asyncSource(control); + rmm::cuda_stream stream; + rmm::device_buffer destination( + 5120, stream, cudf::get_current_device_resource_ref()); + auto* dst = static_cast(destination.data()); + std::vector> reads; + auto release = folly::makeGuard([&] { control->releaseAll(); }); + reads.push_back(input->device_read_async(0, 1024, dst, stream)); + ASSERT_TRUE(waitUntil([&] { return control->submitted.load() == 1; })); + reads.push_back(input->device_read_async(0, 4096, dst + 1024, stream)); + ASSERT_TRUE( + waitUntil([&] { return cache_->refreshStats().numWaitExclusive > 0; })); + control->completeNewest(); + ASSERT_TRUE(waitUntil([&] { return control->submitted.load() == 2; })); + EXPECT_EQ(reads.front().get(), 1024); + EXPECT_EQ( + read_->recordedReads(), (ReadState::Requests{{0, 1024}, {1024, 3072}})); + control->releaseAll(); + EXPECT_EQ(reads.back().get(), 4096); + EXPECT_EQ(fromDevice(dst, 5120), std::string(5120, 'x')); + EXPECT_EQ(stats_->stats().at("cudfKvikioCacheHitBytes").sum, 1024); + EXPECT_EQ(stats_->stats().at("cudfKvikioCacheMissBytes").sum, 4096); +} + +TEST_F(CachingDataSourceTest, asyncSharedFillAndOutOfOrderCompletion) { + read_->data.assign(16384, 'a'); + auto control = std::make_shared(); + auto input = asyncSource(control); + rmm::cuda_stream stream; + rmm::device_buffer destination( + 16384, stream, cudf::get_current_device_resource_ref()); + auto* dst = static_cast(destination.data()); + std::vector> reads; + auto release = folly::makeGuard([&] { control->releaseAll(); }); + reads.push_back(input->device_read_async(0, 4096, dst, stream)); + ASSERT_TRUE(waitUntil([&] { return control->submitted.load() == 1; })); + // Two followers wait for the pending fill on fill threads. They must not + // block a later, independent fill. + reads.push_back(input->device_read_async(0, 4096, dst + 4096, stream)); + reads.push_back(input->device_read_async(0, 4096, dst + 8192, stream)); + reads.push_back(input->device_read_async(4096, 4096, dst + 12288, stream)); + ASSERT_TRUE(waitUntil([&] { return control->submitted.load() == 2; })); + control->completeNewest(); + ASSERT_TRUE(waitUntil([&] { + return stats_->stats().contains("cudfKvikioCacheMissBytes"); + })) << "A ready range must complete while the oldest range is still pending"; + EXPECT_EQ(reads.back().get(), 4096); + EXPECT_EQ(fromDevice(dst + 12288, 4096), std::string(4096, 'a')); + control->releaseAll(); + for (size_t i = 0; i < reads.size() - 1; ++i) { + EXPECT_EQ(reads[i].get(), 4096); + } + EXPECT_EQ(control->submitted, 2); + EXPECT_EQ(stats_->stats().at("cudfKvikioCacheHitBytes").sum, 8192); + EXPECT_EQ(fromDevice(dst, 16384), std::string(16384, 'a')); +} + +TEST_F(CachingDataSourceTest, asyncFailedAndShortFillsAreNotCached) { + auto control = std::make_shared(); + control->releaseAll(); + auto input = asyncSource(control); + rmm::cuda_stream stream; + rmm::device_buffer destination( + 4096, stream, cudf::get_current_device_resource_ref()); + auto* dst = static_cast(destination.data()); + read_->failRead = true; + EXPECT_THROW( + input->device_read_async(0, 4096, dst, stream).get(), VeloxRuntimeError); + read_->shortRead = true; + EXPECT_THROW( + input->device_read_async(0, 4096, dst, stream).get(), VeloxRuntimeError); + control->failSubmission = true; + EXPECT_THROW( + input->device_read_async(0, 4096, dst, stream).get(), VeloxRuntimeError); + EXPECT_EQ(input->device_read_async(0, 4096, dst, stream).get(), 4096); + EXPECT_EQ(input->device_read_async(0, 4096, dst, stream).get(), 4096); + EXPECT_EQ(control->submitted, 3); + EXPECT_EQ(fromDevice(dst, 4096), read_->data.substr(0, 4096)); +} + +TEST_F(CachingDataSourceTest, asyncCapacityFailureUsesUncachedFallback) { + read_->data.assign(32 << 20, 'b'); + auto control = std::make_shared(); + control->releaseAll(); + auto input = asyncSource(control); + rmm::cuda_stream stream; + rmm::device_buffer destination( + read_->data.size(), stream, cudf::get_current_device_resource_ref()); + auto* dst = static_cast(destination.data()); + EXPECT_EQ( + input->device_read_async(0, read_->data.size(), dst, stream).get(), + read_->data.size()); + EXPECT_EQ(control->submitted, 1); + EXPECT_EQ( + stats_->stats().at("cudfKvikioCacheBypassBytes").sum, read_->data.size()); + EXPECT_EQ(fromDevice(dst, read_->data.size()), read_->data); + EXPECT_EQ(cache_->refreshStats().numEntries, 0); +} + +TEST_F(CachingDataSourceTest, missesHitsAndSynchronousOverloads) { + fillPattern(128 << 10); + auto input = source(); + rmm::cuda_stream stream; + constexpr size_t kSize = 17 * 4096 + 37; + rmm::device_buffer destination( + kSize, stream, cudf::get_current_device_resource_ref()); + auto* dst = static_cast(destination.data()); + EXPECT_EQ(input->device_read_async(19, kSize, dst, stream).get(), kSize); + EXPECT_EQ(input->device_read(19, kSize, dst, stream), kSize); + auto owned = input->device_read(19, kSize, stream); + EXPECT_EQ(fromDevice(dst, kSize), read_->data.substr(19, kSize)); + EXPECT_EQ( + fromDevice(owned->data(), owned->size()), read_->data.substr(19, kSize)); + EXPECT_EQ(read_->reads, 1); + const auto metrics = stats_->stats(); + EXPECT_EQ(metrics.at("cudfKvikioCacheMissBytes").sum, kSize); + EXPECT_EQ(metrics.at("cudfKvikioCacheHitBytes").sum, 2 * kSize); + EXPECT_EQ(input->host_read(19, kSize)->size(), kSize); + EXPECT_EQ(read_->reads, 1); +} + +TEST_F(CachingDataSourceTest, cacheOffAndNonCacheableReturnSourceUnchanged) { + for (const bool cacheable : {false, true}) { + auto source = std::make_unique(read_); + auto* original = source.get(); + auto actual = maybeCacheKvikioDataSource( + std::move(source), + path_, + cacheable ? nullptr : cache_.get(), + cacheable); + EXPECT_EQ(actual.get(), original); + } +} + +TEST_F(CachingDataSourceTest, hostMissHitAndFileIdentity) { + auto input = source(); + auto first = input->host_read(17, 512); + auto second = input->host_read(17, 512); + EXPECT_EQ(read_->reads, 1); + EXPECT_EQ(std::memcmp(first->data(), second->data(), 512), 0); + const auto metrics = stats_->stats(); + EXPECT_EQ(metrics.at("cudfKvikioCacheMissBytes").sum, 512); + EXPECT_EQ(metrics.at("cudfKvikioCacheHitBytes").sum, 512); + auto anotherFile = maybeCacheKvikioDataSource( + std::make_unique(read_), + "s3://test/another-file", + cache_.get(), + true); + anotherFile->host_read(17, 512); + EXPECT_EQ(read_->reads, 2); +} + +TEST_F(CachingDataSourceTest, failedAndShortReadsAreNotCached) { + auto input = source(); + read_->shortRead = true; + EXPECT_THROW(input->host_read(31, 100), VeloxRuntimeError); + read_->failRead = true; + EXPECT_THROW(input->host_read(31, 100), VeloxRuntimeError); + EXPECT_EQ(input->host_read(31, 100)->size(), 100); + EXPECT_EQ(read_->reads, 3); + input->host_read(31, 100); + EXPECT_EQ(read_->reads, 3); +} + +TEST_F(CachingDataSourceTest, nonContiguousCacheHit) { + StringIdLease file(fileIds(), path_); + auto pin = cache_->findOrCreate({file.id(), 0}, read_->data.size(), false); + auto* entry = pin.checkedEntry(); + ASSERT_FALSE(entry->hasContiguousData()); + for (const auto& range : entry->dataRanges(read_->data.size())) { + std::memset(range.data(), 'p', range.size()); + } + entry->setExclusiveToShared(); + auto input = source(); + auto bytes = input->host_read(0, read_->data.size()); + EXPECT_EQ( + std::string(reinterpret_cast(bytes->data()), bytes->size()), + std::string(read_->data.size(), 'p')); + EXPECT_EQ(read_->reads, 0); +} + +TEST_F(CachingDataSourceTest, emptyEofAndOverflowingEnd) { + auto input = source(); + EXPECT_EQ(input->host_read(0, 0)->size(), 0); + EXPECT_EQ(input->host_read(input->size(), 1)->size(), 0); + EXPECT_EQ( + input->host_read(std::numeric_limits::max(), 1, nullptr), 0); + EXPECT_EQ(read_->reads, 0); + EXPECT_EQ( + input->host_read(input->size() - 3, std::numeric_limits::max()) + ->size(), + 3); + EXPECT_EQ(read_->reads, 1); +} + +TEST_F(CachingDataSourceTest, longerRangeReusesShortPrefix) { + auto input = source(); + input->host_read(1, 100); + EXPECT_EQ(input->host_read(1, 200)->size(), 200); + EXPECT_EQ(read_->reads, 2); + input->host_read(1, 50); + EXPECT_EQ(read_->reads, 2); + EXPECT_EQ( + read_->recordedReads(), (ReadState::Requests{{1, 100}, {101, 100}})); + StringIdLease file(fileIds(), path_); + auto pin = cache_->findOrCreate({file.id(), 1}, 100, true); + EXPECT_EQ(pin.checkedEntry()->size(), 100); +} + +TEST_F( + CachingDataSourceTest, + hostReadsReuseFragmentChainAndClampFinalFragment) { + fillPattern(read_->data.size()); + auto input = source(); + input->host_read(0, 1024); + input->host_read(1024, 512); + input->host_read(1536, 1024); + std::vector dst(4096); + EXPECT_EQ(input->host_read(0, dst.size(), dst.data()), dst.size()); + EXPECT_EQ(std::memcmp(dst.data(), read_->data.data(), dst.size()), 0); + EXPECT_EQ( + read_->recordedReads(), + (ReadState::Requests{ + {0, 1024}, {1024, 512}, {1536, 1024}, {2560, 1536}})); + auto metrics = stats_->stats(); + EXPECT_EQ(metrics.at("cudfKvikioCacheMissBytes").sum, 4096); + EXPECT_EQ(metrics.at("cudfKvikioCacheHitBytes").sum, 2560); + EXPECT_EQ(metrics.at("cudfKvikioCacheReusedPrefixes").sum, 3); + EXPECT_EQ(metrics.at("cudfKvikioCacheReusedPrefixBytes").sum, 2560); + auto owned = input->host_read_async(0, 1800).get(); + EXPECT_EQ(owned->size(), 1800); + EXPECT_EQ(std::memcmp(owned->data(), read_->data.data(), 1800), 0); + EXPECT_EQ( + input->host_read_async(0, dst.size(), dst.data()).get(), dst.size()); + EXPECT_EQ(std::memcmp(dst.data(), read_->data.data(), dst.size()), 0); + EXPECT_EQ(read_->reads, 4); +} + +TEST_F(CachingDataSourceTest, failedHostSuffixPreservesCachedPrefix) { + auto input = source(); + input->host_read(0, 1024); + read_->shortRead = true; + EXPECT_THROW(input->host_read(0, 4096), VeloxRuntimeError); + read_->failRead = true; + EXPECT_THROW(input->host_read(0, 4096), VeloxRuntimeError); + EXPECT_EQ(input->host_read(0, 4096)->size(), 4096); + EXPECT_EQ( + read_->recordedReads(), + (ReadState::Requests{ + {0, 1024}, {1024, 3072}, {1024, 3072}, {1024, 3072}})); +} + +TEST_F(CachingDataSourceTest, uncachedFallbackReadsOnlyMissingSuffix) { + read_->data.assign(32 << 20, 'b'); + auto control = std::make_shared(); + control->releaseAll(); + auto input = asyncSource(control); + input->host_read(0, 4096); + // Hold the prefix so an unsuccessful large cache allocation cannot evict it. + StringIdLease file(fileIds(), path_); + auto pin = cache_->findOrCreate({file.id(), 0}, 4096, true); + auto host = input->host_read(0, read_->data.size()); + EXPECT_EQ(host->size(), read_->data.size()); + EXPECT_EQ(std::memcmp(host->data(), read_->data.data(), host->size()), 0); + rmm::cuda_stream stream; + rmm::device_buffer destination( + read_->data.size(), stream, cudf::get_current_device_resource_ref()); + auto* dst = static_cast(destination.data()); + EXPECT_EQ( + input->device_read_async(0, read_->data.size(), dst, stream).get(), + read_->data.size()); + EXPECT_EQ(fromDevice(dst, read_->data.size()), read_->data); + // Short uncached reads return the prefix plus the actual suffix length. + read_->shortRead = true; + EXPECT_EQ( + input->device_read(0, read_->data.size(), dst, stream), + read_->data.size() - 1); + EXPECT_EQ( + read_->recordedReads(), + (ReadState::Requests{ + {0, 4096}, + {4096, read_->data.size() - 4096}, + {4096, read_->data.size() - 4096}, + {4096, read_->data.size() - 4096}})); +} + +TEST_F(CachingDataSourceTest, cacheCapacityFailureFallsBackWithoutCaching) { + read_->data.assign(32 << 20, 'b'); + auto input = source(); + auto result = input->host_read(0, read_->data.size()); + EXPECT_EQ(result->size(), read_->data.size()); + EXPECT_EQ(result->data()[result->size() - 1], 'b'); + EXPECT_EQ(read_->reads, 1); + EXPECT_EQ( + stats_->stats().at("cudfKvikioCacheBypassBytes").sum, result->size()); +} + +TEST_F(CachingDataSourceTest, concurrentReadersShareOneFill) { + auto input = source(); + std::promise start; + auto started = start.get_future().share(); + std::vector> readers; + for (int i = 0; i < 8; ++i) { + readers.push_back(std::async(std::launch::async, [&input, started] { + started.wait(); + auto result = input->host_read(0, 4096); + EXPECT_EQ(result->size(), 4096); + EXPECT_EQ(result->data()[4095], 'x'); + })); + } + start.set_value(); + for (auto& reader : readers) { + reader.get(); + } + EXPECT_EQ(read_->reads, 1); +} + +TEST_F(CachingDataSourceTest, hostFutureKeepsSourceAlive) { + auto input = source(); + auto future = input->host_read_async(0, 100); + input.reset(); + EXPECT_FALSE(read_->destroyed); + EXPECT_EQ(future.get()->size(), 100); + EXPECT_TRUE(read_->destroyed); +} + +TEST_F(CachingDataSourceTest, synchronousReadsUseCacheAndClampRanges) { + auto input = source(); + rmm::cuda_stream stream; + rmm::device_buffer destination( + 200, stream, cudf::get_current_device_resource_ref()); + auto* dst = static_cast(destination.data()); + // Exercise a miss and both synchronous cache-hit overloads. + EXPECT_EQ(input->device_read(9, 200, dst, stream), 200); + EXPECT_EQ(input->device_read(9, 200, dst, stream), 200); + auto owned = input->device_read(9, 200, stream); + EXPECT_EQ(read_->reads, 1); + EXPECT_EQ(fromDevice(dst, 200), read_->data.substr(9, 200)); + EXPECT_EQ( + fromDevice(owned->data(), owned->size()), read_->data.substr(9, 200)); + + EXPECT_EQ(input->device_read(0, 0, nullptr, stream), 0); + EXPECT_EQ(input->device_read(input->size(), 1, nullptr, stream), 0); + EXPECT_EQ( + input->device_read( + std::numeric_limits::max(), 1, nullptr, stream), + 0); + EXPECT_EQ( + input->device_read( + input->size() - 3, std::numeric_limits::max(), dst, stream), + 3); + EXPECT_EQ(fromDevice(dst, 3), read_->data.substr(input->size() - 3)); + EXPECT_THROW(input->device_read(0, 1, nullptr, stream), VeloxRuntimeError); +} + +TEST_F(CachingDataSourceTest, synchronousReadsFromSingleThreadCaller) { + for (const bool owning : {false, true}) { + SCOPED_TRACE(owning); + rmm::cuda_stream stream; + rmm::device_buffer destination( + 100, stream, cudf::get_current_device_resource_ref()); + auto* dst = static_cast(destination.data()); + folly::CPUThreadPoolExecutor pool(1); + auto input = maybeCacheKvikioDataSource( + std::make_unique(read_), + path_, + cache_.get(), + true, + stats_); + // Exercise each overload on a cache hit from a single-threaded caller. + ASSERT_EQ(input->host_read(0, 100)->size(), 100); + std::unique_ptr owned; + std::promise completed; + auto result = completed.get_future(); + pool.add([&] { + try { + if (owning) { + owned = input->device_read(0, 100, stream); + completed.set_value(owned->size()); + } else { + completed.set_value(input->device_read(0, 100, dst, stream)); + } + } catch (...) { + completed.set_exception(std::current_exception()); + } + }); + const auto initial = result.wait_for(5s); + EXPECT_EQ(initial, std::future_status::ready) + << "Synchronous read did not complete from the caller executor"; + EXPECT_EQ(result.get(), 100); + EXPECT_EQ( + fromDevice(owning ? owned->data() : dst, 100), std::string(100, 'x')); + pool.join(); + } +} + +TEST_F(CachingDataSourceTest, synchronousReadWaitsForH2D) { + auto input = source(); + rmm::cuda_stream stream; + rmm::device_buffer destination( + 100, stream, cudf::get_current_device_resource_ref()); + stream.synchronize(); + auto* dst = static_cast(destination.data()); + StreamGate gate; + CUDF_CUDA_TRY(cudaLaunchHostFunc(stream.value(), waitForStreamGate, &gate)); + gate.entered.get_future().wait(); + auto waiter = std::async(std::launch::async, [&] { + return input->device_read(10, 100, dst, stream); + }); + EXPECT_EQ(waiter.wait_for(50ms), std::future_status::timeout); + gate.release.set_value(); + EXPECT_EQ(waiter.get(), 100); + EXPECT_EQ(fromDevice(dst, 100), read_->data.substr(10, 100)); +} + +TEST_F(CachingDataSourceTest, asyncReadGetWaitsForH2D) { + auto input = source(); + rmm::cuda_stream stream; + rmm::device_buffer destination( + 100, stream, cudf::get_current_device_resource_ref()); + stream.synchronize(); + auto* dst = static_cast(destination.data()); + StreamGate gate; + CUDF_CUDA_TRY(cudaLaunchHostFunc(stream.value(), waitForStreamGate, &gate)); + gate.entered.get_future().wait(); + auto read = input->device_read_async(10, 100, dst, stream); + auto waiter = std::async(std::launch::async, [&] { return read.get(); }); + EXPECT_EQ(waiter.wait_for(50ms), std::future_status::timeout); + gate.release.set_value(); + EXPECT_EQ(waiter.get(), 100); + EXPECT_EQ(fromDevice(dst, 100), read_->data.substr(10, 100)); +} + +TEST_F(CachingDataSourceTest, deviceFutureKeepsSourceAlive) { + std::promise release; + read_->gate = release.get_future().share(); + auto input = source(); + rmm::cuda_stream stream; + rmm::device_buffer destination( + 100, stream, cudf::get_current_device_resource_ref()); + auto* dst = static_cast(destination.data()); + auto future = input->device_read_async(10, 100, dst, stream); + input.reset(); + EXPECT_FALSE(read_->destroyed); + release.set_value(); + EXPECT_EQ(future.get(), 100); + EXPECT_EQ(fromDevice(dst, 100), read_->data.substr(10, 100)); +} + +TEST_F(CachingDataSourceTest, usesTheDestinationStreamsDevice) { + int count = 0; + CUDF_CUDA_TRY(cudaGetDeviceCount(&count)); + if (count < 2) { + GTEST_SKIP() << "Requires two CUDA devices"; + } + auto input = source(); + for (int device : {0, 1, 0}) { + const rmm::cuda_set_device_raii scope{rmm::cuda_device_id{device}}; + rmm::cuda_stream stream; + rmm::device_buffer destination( + 100, stream, cudf::get_current_device_resource_ref()); + auto* dst = static_cast(destination.data()); + EXPECT_EQ(input->device_read_async(0, 100, dst, stream).get(), 100); + EXPECT_EQ(fromDevice(dst, 100), std::string(100, 'x')); + } + EXPECT_EQ(read_->reads, 1); +} + +} // namespace +} // namespace facebook::velox::cudf_velox::connector::hive