[Fix] Bind ring drain workers to their CUDA device

native/csrc/ring/drain_thread.cpp

@@ -13,16 +13,36 @@
 #include <ATen/ATen.h>
 #include <cassert>
 #include <chrono>
+#include <cstdio>
 #include <cstring>
+#include <future>
+#include <string>
 #include <stdexcept>
 
 namespace ring {
 
+namespace {
+
+void report_cuda_failure(const char* operation, cudaError_t error) {
+    if (error == cudaSuccess) return;
+    std::fprintf(stderr, "[drain] ERROR: %s failed: %s\n",
+                 operation, cudaGetErrorString(error));
+    std::fflush(stderr);
+}
+
+}  // namespace
+
 // ---------------------------------------------------------------------------
 DrainThread::DrainThread(RingState& rs, PinnedStaging& staging,
                          const RingConfig& cfg)
     : ring_(rs), staging_(staging), cfg_(cfg)
 {
+    cudaError_t error = cudaGetDevice(&owner_device_);
+    if (error != cudaSuccess) {
+        throw std::runtime_error(
+            std::string("DrainThread: cudaGetDevice failed: ") +
+            cudaGetErrorString(error));
+    }
     if (cudaStreamCreateWithFlags(&stream_, cudaStreamNonBlocking) != cudaSuccess)
         throw std::runtime_error("DrainThread: cudaStreamCreate failed");
 }
@@ -33,8 +53,32 @@ DrainThread::~DrainThread() noexcept {
 }
 
 void DrainThread::start() {
+    std::promise<cudaError_t> startup;
+    std::future<cudaError_t> startup_result = startup.get_future();
     running_.store(true, std::memory_order_relaxed);
-    thread_ = std::thread([this] { loop(); });
+    try {
+        thread_ = std::thread(
+            [this, startup = std::move(startup)]() mutable {
+                cudaError_t error = cudaSetDevice(owner_device_);
+                startup.set_value(error);
+                if (error != cudaSuccess) {
+                    running_.store(false, std::memory_order_relaxed);
+                    return;
+                }
+                loop();
+            });
+    } catch (...) {
+        running_.store(false, std::memory_order_relaxed);
+        throw;
+    }
+
+    cudaError_t error = startup_result.get();
+    if (error != cudaSuccess) {
+        if (thread_.joinable()) thread_.join();
+        throw std::runtime_error(
+            std::string("DrainThread: cudaSetDevice failed: ") +
+            cudaGetErrorString(error));
+    }
 }
 
 void DrainThread::stop() {
@@ -289,7 +333,7 @@ void DrainThread::loop() {
     }
 
     // Final flush
-    cudaDeviceSynchronize();
+    report_cuda_failure("cudaDeviceSynchronize", cudaDeviceSynchronize());
     do_full_flush();
 }
 
@@ -357,7 +401,8 @@ void DrainThread::flush_state_update(uint64_t flush_count, uint64_t flush_bytes)
 }
 
 void DrainThread::sync_stream() {
-    cudaStreamSynchronize(stream_);
+    report_cuda_failure("cudaStreamSynchronize",
+                        cudaStreamSynchronize(stream_));
 }
 
 // ---------------------------------------------------------------------------
@@ -384,10 +429,10 @@ void DrainThread::enqueue_d2h(uint64_t flush_bytes) {
                         ring_.payload_buf + gpu_cursor,
                         chunk, cudaMemcpyDeviceToHost, stream_);
         if (err != cudaSuccess) {
-            RING_DBG("[enqueue_d2h] cudaMemcpyAsync FAILED: %s\n",
-                    cudaGetErrorString(err));
+            report_cuda_failure("cudaMemcpyAsync", err);
+        } else {
+            RING_DBG("[enqueue_d2h] chunk=%d enqueued OK\n", chunk_idx);
         }
-        RING_DBG("[enqueue_d2h] chunk=%d enqueued OK\n", chunk_idx);
 
         remaining  -= chunk;
         gpu_cursor  = (gpu_cursor + chunk) % gpu_cap;

native/csrc/ring/drain_thread.h

@@ -74,6 +74,7 @@ class DrainThread {
     RingState&      ring_;
     PinnedStaging&  staging_;
     RingConfig      cfg_;
+    int             owner_device_{-1};
     cudaStream_t    stream_{};
 
     std::thread             thread_;

tests/native/ring/test_ring_engine.cu

@@ -7,11 +7,15 @@
 
 #include <cuda_runtime.h>
 
+#include <atomic>
+#include <chrono>
 #include <cstdint>
 #include <cstdio>
 #include <cstdlib>
 #include <cstring>
+#include <future>
 #include <memory>
+#include <thread>
 #include <vector>
 
 static int g_pass = 0;
@@ -239,6 +243,88 @@ static void test_zero_byte_delivery() {
     harness.release(task);
 }
 
+struct BlockingCallbackState {
+    std::atomic<bool> entered{false};
+    std::atomic<bool> release{false};
+};
+
+static void CUDART_CB blocking_host_callback(void* data) {
+    auto* state = static_cast<BlockingCallbackState*>(data);
+    state->entered.store(true, std::memory_order_release);
+    while (!state->release.load(std::memory_order_acquire)) {
+        std::this_thread::sleep_for(std::chrono::milliseconds(1));
+    }
+}
+
+static void test_drain_worker_binds_owner_device() {
+    banner("drain worker binds the ring owner device");
+
+    int device_count = 0;
+    CUDA_CHECK(cudaGetDeviceCount(&device_count));
+    if (device_count < 2) {
+        std::printf("[ SKIP ] requires two CUDA devices\n");
+        return;
+    }
+
+    int original_device = 0;
+    CUDA_CHECK(cudaGetDevice(&original_device));
+
+    BlockingCallbackState callback;
+    cudaStream_t blocked_stream{};
+    CUDA_CHECK(cudaSetDevice(0));
+    CUDA_CHECK(cudaStreamCreateWithFlags(&blocked_stream,
+                                         cudaStreamNonBlocking));
+    CUDA_CHECK(cudaLaunchHostFunc(blocked_stream, blocking_host_callback,
+                                  &callback));
+
+    const auto callback_deadline =
+        std::chrono::steady_clock::now() + std::chrono::seconds(5);
+    while (!callback.entered.load(std::memory_order_acquire) &&
+           std::chrono::steady_clock::now() < callback_deadline) {
+        std::this_thread::sleep_for(std::chrono::milliseconds(1));
+    }
+    const bool callback_entered =
+        callback.entered.load(std::memory_order_acquire);
+    EXPECT(callback_entered);
+
+    bool stopped_without_waiting_for_device_zero = false;
+    if (callback_entered) {
+        CUDA_CHECK(cudaSetDevice(1));
+        ring::RingConfig cfg = make_config();
+        ring::AllocatedRing allocated(cfg);
+        allocated.init();
+        ring::PinnedStaging staging;
+        staging.init(cfg.effective_staging_bytes());
+        auto drain = std::make_unique<ring::DrainThread>(
+            allocated.state(), staging, cfg);
+
+        CUDA_CHECK(cudaSetDevice(0));
+        drain->start();
+        auto stopped = std::async(std::launch::async, [&drain] {
+            drain->stop();
+        });
+        stopped_without_waiting_for_device_zero =
+            stopped.wait_for(std::chrono::seconds(2)) ==
+            std::future_status::ready;
+
+        callback.release.store(true, std::memory_order_release);
+        CUDA_CHECK(cudaSetDevice(0));
+        CUDA_CHECK(cudaStreamSynchronize(blocked_stream));
+        stopped.get();
+
+        CUDA_CHECK(cudaSetDevice(1));
+        drain.reset();
+    } else {
+        callback.release.store(true, std::memory_order_release);
+        CUDA_CHECK(cudaStreamSynchronize(blocked_stream));
+    }
+
+    CUDA_CHECK(cudaSetDevice(0));
+    CUDA_CHECK(cudaStreamDestroy(blocked_stream));
+    CUDA_CHECK(cudaSetDevice(original_device));
+    EXPECT(stopped_without_waiting_for_device_zero);
+}
+
 int main() {
     setbuf(stdout, nullptr);
     ring::set_ring_null_mode(false);
@@ -249,6 +335,7 @@ int main() {
     test_prefix_force_flush();
     test_repeated_wrap_delivery();
     test_zero_byte_delivery();
+    test_drain_worker_binds_owner_device();
 
     std::printf("Results: %d passed, %d failed\n", g_pass, g_fail);
     return g_fail == 0 ? 0 : 1;