diff --git a/native/csrc/ring/drain_thread.cpp b/native/csrc/ring/drain_thread.cpp index 6fbcfa428..ecbb89ed1 100644 --- a/native/csrc/ring/drain_thread.cpp +++ b/native/csrc/ring/drain_thread.cpp @@ -13,16 +13,36 @@ #include #include #include +#include #include +#include +#include #include 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 startup; + std::future 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; diff --git a/native/csrc/ring/drain_thread.h b/native/csrc/ring/drain_thread.h index 66d03780b..55556a90c 100644 --- a/native/csrc/ring/drain_thread.h +++ b/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_; diff --git a/tests/native/ring/test_ring_engine.cu b/tests/native/ring/test_ring_engine.cu index 806a85485..3ff74f06f 100644 --- a/tests/native/ring/test_ring_engine.cu +++ b/tests/native/ring/test_ring_engine.cu @@ -7,11 +7,15 @@ #include +#include +#include #include #include #include #include +#include #include +#include #include static int g_pass = 0; @@ -239,6 +243,88 @@ static void test_zero_byte_delivery() { harness.release(task); } +struct BlockingCallbackState { + std::atomic entered{false}; + std::atomic release{false}; +}; + +static void CUDART_CB blocking_host_callback(void* data) { + auto* state = static_cast(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( + 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;