[Store] Skip hot-cache writeback for device-resident read destinations

mooncake-store/src/client_service.cpp

@@ -5623,14 +5623,30 @@ void Client::ProcessSlicesAsync(const std::string& key,
         return;
     }
 
+    // The writeback copies out of the read destination with a host memcpy
+    // (LocalHotCacheHandler::SubmitPutTask), so a slice that landed in device
+    // memory (e.g. batch_get_into into CUDA buffers) would fault the reader.
+    // Device-resident results are not cached through this path.
+    auto runtime_accelerator =
+        device::GetAcceleratorRegistry().RuntimeAccelerators();
+    size_t device_dest_slices = 0;
     // Identify TE transfer slices (non-local) and submit async put tasks
     for (size_t i = 0; i < slices.size(); ++i) {
+        if (runtime_accelerator.FindDeviceForPointer(slices[i].ptr) !=
+            nullptr) {
+            ++device_dest_slices;
+            continue;
+        }
         if (!hot_cache_handler_->SubmitPutTask(key, slices[i])) {
             LOG(ERROR) << "Failed to submit hot cache put task for key=" << key
                        << " slice_idx=" << i;
             return;
         }
     }
+    if (device_dest_slices > 0) {
+        LOG(INFO) << "Skipped hot-cache writeback for " << device_dest_slices
+                  << " device-resident slice(s), key=" << key;
+    }
 }
 
 bool Client::IsReplicaOnLocalMemory(const Replica::Descriptor& replica) {

mooncake-store/tests/client_local_hot_cache_test.cpp

@@ -3,6 +3,7 @@
 #include "client_buffer.h"
 #include "common/client_buffer_allocation.h"
 #include "count_min_sketch.h"
+#include "device/accelerator_registry.h"
 #include "local_hot_cache.h"
 #include "replica.h"
 #include "test_server_helpers.h"
@@ -1168,6 +1169,194 @@ TEST_F(LocalHotCacheTest, BatchGetWithHotCacheEnabled) {
     }
 }
 
+namespace {
+
+// Stand-in for a CUDA-like accelerator that claims one fixed buffer as device
+// memory, so the hot-cache writeback guard in Client::ProcessSlicesAsync can
+// be exercised on CPU-only runners. Copies fall back to memcpy: what matters
+// here is the pointer classification, not real device semantics.
+class FakeAcceleratorDevice final : public device::AcceleratorDevice {
+   public:
+    explicit FakeAcceleratorDevice(const void* device_ptr)
+        : device_ptr_(device_ptr) {}
+
+    device::AcceleratorVendor Vendor() const override {
+        return device::AcceleratorVendor::kNvidia;
+    }
+
+    bool Available(bool ensure) const override {
+        (void)ensure;
+        return true;
+    }
+
+    device::PointerInfo QueryPointer(const void* ptr) const override {
+        if (ptr == device_ptr_) {
+            return device::PointerInfo{.kind = device::MemoryKind::kDevice,
+                                       .device_id = 0};
+        }
+        return device::PointerInfo{.kind = device::MemoryKind::kHost,
+                                   .device_id = -1};
+    }
+
+    int32_t CurrentDeviceId() const override { return 0; }
+
+    void SetContext(int32_t device_id) const override { (void)device_id; }
+
+    bool Copy(void* dst, const void* src, size_t size,
+              device::CopyDirection direction) const override {
+        (void)direction;
+        std::memcpy(dst, src, size);
+        return true;
+    }
+
+    PinnedHostBuffer AllocatePinnedHost(size_t size) const override {
+        (void)size;
+        return PinnedHostBuffer();
+    }
+
+   private:
+    const void* device_ptr_;
+};
+
+}  // namespace
+
+// Regression test for the SIGSEGV reported when batch_get_into targets device
+// memory with the local hot cache enabled: ProcessSlicesAsync wrote the cache
+// block with a host memcpy out of the destination slice, which faults on a
+// CUDA pointer. A device-destined slice must be skipped while a host-destined
+// slice in the same batch is still written back.
+TEST_F(LocalHotCacheTest, BatchGetSkipsHotCacheWritebackForDeviceDestination) {
+    // Save original env vars
+    const char* original_cache_size =
+        std::getenv("MC_STORE_LOCAL_HOT_CACHE_SIZE");
+    const char* original_block_size =
+        std::getenv("MC_STORE_LOCAL_HOT_BLOCK_SIZE");
+    const char* original_threshold =
+        std::getenv("MC_STORE_LOCAL_HOT_ADMISSION_THRESHOLD");
+
+    // Enable hot cache (8MB = 2 blocks of 4MB) and admit on first access so a
+    // single BatchGet exercises the writeback path.
+    setenv("MC_STORE_LOCAL_HOT_BLOCK_SIZE", "4194304", 1);
+    setenv("MC_STORE_LOCAL_HOT_CACHE_SIZE", "8388608", 1);
+    setenv("MC_STORE_LOCAL_HOT_ADMISSION_THRESHOLD", "1", 1);
+
+    // Static storage keeps the buffer and the device alive for the rest of
+    // the process: the accelerator registry holds raw pointers to registered
+    // devices, so a stack-scoped device would dangle for later tests.
+    static char fake_device_storage[64];
+    static FakeAcceleratorDevice fake_device(fake_device_storage);
+    static device::AcceleratorDeviceRegistrar fake_registrar(fake_device);
+
+    std::string local_ip = getLocalIpAddress();
+
+    // Writer client owns the segment and the data.
+    auto writer_opt = CreateTestClient(local_ip + ":12345");
+    ASSERT_TRUE(writer_opt.has_value())
+        << "Failed to create writer client with in-proc master server";
+    auto writer = writer_opt.value();
+
+    size_t segment_size = 64 * 1024 * 1024;  // 64MB
+    void* segment_ptr = allocate_buffer_allocator_memory(segment_size);
+    ASSERT_NE(segment_ptr, nullptr) << "Failed to allocate segment memory";
+    auto mount_result = writer->MountSegment(segment_ptr, segment_size);
+    if (!mount_result.has_value()) {
+        free_memory("", segment_ptr);
+        GTEST_SKIP() << "Failed to mount segment: "
+                     << toString(mount_result.error());
+    }
+
+    const std::string dev_key = "batch_get_dev_dest_key";
+    const std::string host_key = "batch_get_host_dest_key";
+    const std::string dev_data = "payload for device destination";
+    const std::string host_data = "payload for host destination";
+
+    ReplicateConfig config;
+    config.replica_num = 1;
+    const std::pair<std::string, std::string> key_data[] = {
+        {dev_key, dev_data}, {host_key, host_data}};
+    for (const auto& [key, data] : key_data) {
+        std::vector<char> put_buffer(data.size());
+        std::memcpy(put_buffer.data(), data.data(), data.size());
+        std::vector<Slice> put_slices;
+        put_slices.emplace_back(Slice{put_buffer.data(), data.size()});
+        auto put_result = writer->Put(key, put_slices, config);
+        if (!put_result.has_value()) {
+            writer->UnmountSegment(segment_ptr, segment_size);
+            free_memory("", segment_ptr);
+            GTEST_SKIP() << "Put operation failed, skipping BatchGet test";
+        }
+    }
+
+    // Reader client owns no segment: its reads are remote transfers, which is
+    // what drives the hot-cache writeback path.
+    auto reader_opt = CreateTestClient(local_ip + ":12346");
+    ASSERT_TRUE(reader_opt.has_value())
+        << "Failed to create reader client with in-proc master server";
+    auto reader = reader_opt.value();
+    ASSERT_TRUE(reader->IsHotCacheEnabled());
+
+    ASSERT_LT(dev_data.size(), sizeof(fake_device_storage));
+    std::vector<char> host_buffer(host_data.size());
+    std::unordered_map<std::string, std::vector<Slice>> batch_slices;
+    batch_slices[dev_key].emplace_back(
+        Slice{fake_device_storage, dev_data.size()});
+    batch_slices[host_key].emplace_back(
+        Slice{host_buffer.data(), host_buffer.size()});
+
+    const std::vector<std::string> keys = {dev_key, host_key};
+    auto results = reader->BatchGet(keys, batch_slices);
+    ASSERT_EQ(results.size(), keys.size());
+    for (const auto& result : results) {
+        ASSERT_TRUE(result.has_value()) << "BatchGet should succeed";
+    }
+    EXPECT_EQ(
+        std::memcmp(fake_device_storage, dev_data.data(), dev_data.size()), 0)
+        << "Device-destined data should match";
+    EXPECT_EQ(
+        std::memcmp(host_buffer.data(), host_data.data(), host_data.size()), 0)
+        << "Host-destined data should match";
+
+    auto cache = reader->GetHotCache();
+    ASSERT_NE(cache, nullptr);
+
+    // The publish runs on the handler's worker threads; wait for the control
+    // key so the skip assertion below cannot pass vacuously.
+    bool host_cached = false;
+    for (int i = 0; i < 200 && !host_cached; ++i) {
+        host_cached = cache->HasHotKey(host_key);
+        if (!host_cached) {
+            std::this_thread::sleep_for(std::chrono::milliseconds(10));
+        }
+    }
+    EXPECT_TRUE(host_cached)
+        << "Host-destined read should be written back to the hot cache";
+    EXPECT_FALSE(cache->HasHotKey(dev_key))
+        << "Device-destined read must not be written back; the writeback "
+           "memcpy's from the destination slice on the host";
+
+    // Cleanup: unmount segment and free memory
+    reader.reset();
+    writer->UnmountSegment(segment_ptr, segment_size);
+    free_memory("", segment_ptr);
+
+    // Restore original env vars
+    if (original_cache_size) {
+        setenv("MC_STORE_LOCAL_HOT_CACHE_SIZE", original_cache_size, 1);
+    } else {
+        unsetenv("MC_STORE_LOCAL_HOT_CACHE_SIZE");
+    }
+    if (original_block_size) {
+        setenv("MC_STORE_LOCAL_HOT_BLOCK_SIZE", original_block_size, 1);
+    } else {
+        unsetenv("MC_STORE_LOCAL_HOT_BLOCK_SIZE");
+    }
+    if (original_threshold) {
+        setenv("MC_STORE_LOCAL_HOT_ADMISSION_THRESHOLD", original_threshold, 1);
+    } else {
+        unsetenv("MC_STORE_LOCAL_HOT_ADMISSION_THRESHOLD");
+    }
+}
+
 // Test deferred LRU touch ordering: accessed blocks survive eviction
 TEST_F(LocalHotCacheTest, DeferredLRUTouchOrdering) {
     // 3-block cache (48MB / 16MB = 3 blocks)