diff --git a/src/plugins/ucx/ucx_backend.cpp b/src/plugins/ucx/ucx_backend.cpp index d031d9bd66..5dde2ec172 100644 --- a/src/plugins/ucx/ucx_backend.cpp +++ b/src/plugins/ucx/ucx_backend.cpp @@ -26,6 +26,7 @@ #include #include #include +#include #include #include #include "absl/strings/numbers.h" @@ -133,13 +134,19 @@ class nixlUcxBackendReqH : public nixlBackendReqH { virtual void release() { - // TODO: Error log: uncompleted requests found! Cancelling ... for (nixlUcxReq req : requests_) { const nixl_status_t ret = nixl::ucx::ucsToNixlStatus(ucp_request_check_status(req)); if (ret == NIXL_IN_PROG) { - // TODO: Need process this properly. - // it may not be enough to cancel UCX request + // A no-op for RMA - ucp_request_cancel() only acts on tag receives - but + // kept so tag-based operations still get cancelled before the drain. worker_->reqCancel(req); + if (!drainRequest(req)) { + NIXL_ERROR << "Request " << req + << " is still in flight after the drain timeout. It is not " + "safe to deregister the memory backing this transfer yet: " + "UCX dereferences the ucp_mem_h when the operation " + "eventually completes."; + } } worker_->reqRelease(req); } @@ -147,6 +154,53 @@ class nixlUcxBackendReqH : public nixlBackendReqH { conn_.reset(); } + // Progress the worker until req reaches a terminal state, so that UCX is done with the + // ucp_mem_h before the caller is free to deregister it. Returns false if the request is + // still in flight when the deadline expires. + // + // RMA operations are posted with UCP_OP_ATTR_FIELD_MEMH, and UCX keeps that ucp_mem_h in + // the request as a plain pointer (ucp_datatype_iter contig.memh): a user memh is not + // reference counted. On completion ucp_datatype_iter_cleanup() calls ucp_memh_put() on + // it, which dereferences memh->context and memh->parent. If the caller deregistered in + // the meantime, ucp_mem_unmap() has already ucs_free()d the memh and the completion + // faults inside ucp_memh_put(). + // + // Releasing the request object itself is safe either way: ucp_request_free() on an + // uncompleted request only marks it released, and UCX returns it to the request pool + // when it completes. What the drain protects is the memory handle, not the request. + // + // The deadline keeps release() from blocking forever when a transfer is wedged on an + // endpoint that is alive but no longer making progress. + [[nodiscard]] bool + drainRequest(nixlUcxReq req) const { + using namespace std::chrono_literals; + + const auto timeout = + nixl::config::getValueDefaulted("NIXL_UCX_REQUEST_DRAIN_TIMEOUT", 10'000ms); + const auto warning_interval = + nixl::config::getValueDefaulted("NIXL_UCX_WARNING_TIMEOUT", 5'000ms); + auto next_warning = warning_interval; + + const auto start = std::chrono::steady_clock::now(); + while (nixl::ucx::ucsToNixlStatus(ucp_request_check_status(req)) == NIXL_IN_PROG) { + const auto elapsed = std::chrono::steady_clock::now() - start; + if (elapsed > timeout) { + return false; + } + + if (elapsed > next_warning) { + NIXL_WARN << "Still draining in-flight request " << req << " after " + << next_warning.count() << " ms"; + next_warning += warning_interval; + } + + if (worker_->progress() == 0) { + std::this_thread::sleep_for(1ms); + } + } + return true; + } + [[nodiscard]] virtual nixl_status_t status() { if (requests_.empty()) { @@ -547,12 +601,59 @@ class nixlUcxCompositeBackendReqH : public nixlUcxBackendReqH { if (sharedState_) { // Set failed status to stop progress chunks sharedState_->status.store(NIXL_ERR_NOT_FOUND); + + if (!waitForPendingChunks()) { + NIXL_ERROR << *this << " still has " << sharedState_->pendingReqs.load() + << " chunk request(s) in flight after the drain timeout; " + "deregistering the memory backing this transfer is not safe."; + } + // Reset shared state - it will be effectively released when the last chunk // resets the shared state pointer sharedState_.reset(); } } + // Wait for chunks already in flight to finish. Setting the failed status stops new + // chunks from starting, but it does not complete requests that are already posted, + // and resetting the shared state only drops our reference to it. Returns false if + // chunks are still in flight when the deadline expires. + // + // Returning while a chunk request is still in flight lets the caller deregister its + // memory: ucp_mem_unmap() then frees the memory handle while a zcopy completion is + // still pending, and that completion later dereferences the freed handle on a + // progress thread (use-after-free inside ucp_memh_put). + // + // The deadline keeps release() from blocking forever when a chunk is wedged on an + // endpoint that is alive but no longer making progress. + [[nodiscard]] bool + waitForPendingChunks() const { + using namespace std::chrono_literals; + + const auto timeout = + nixl::config::getValueDefaulted("NIXL_UCX_REQUEST_DRAIN_TIMEOUT", 10'000ms); + const auto warning_interval = + nixl::config::getValueDefaulted("NIXL_UCX_WARNING_TIMEOUT", 5'000ms); + auto next_warning = warning_interval; + + const auto start = std::chrono::steady_clock::now(); + while (sharedState_->pendingReqs.load() > 0) { + const auto elapsed = std::chrono::steady_clock::now() - start; + if (elapsed > timeout) { + return false; + } + + if (elapsed > next_warning) { + NIXL_WARN << *this << " still waiting for " << sharedState_->pendingReqs.load() + << " in-flight chunk(s) after " << next_warning.count() << " ms"; + next_warning += warning_interval; + } + + std::this_thread::yield(); + } + return true; + } + [[nodiscard]] nixl_status_t status() override { getWorker()->progressLoop(); diff --git a/test/gtest/test_transfer.cpp b/test/gtest/test_transfer.cpp index 9bbeb859d6..156de57b47 100644 --- a/test/gtest/test_transfer.cpp +++ b/test/gtest/test_transfer.cpp @@ -26,7 +26,10 @@ #include #include #include +#include +#include #include +#include #include #include #include @@ -664,6 +667,117 @@ TEST_P(TestTransferTelemetry, GetXferTelemetryDisabled) { EXPECT_LE(lig.getIgnoredCount(), 1); } +// Releasing a transfer handle while its requests are still in flight must not return until +// UCX is done with the memory handles the transfer was posted with. NIXL posts RMA with +// UCP_OP_ATTR_FIELD_MEMH, and UCX keeps that ucp_mem_h in the request as a plain pointer, +// dereferencing it from ucp_memh_put() when the operation completes. A caller that +// deregisters as soon as release() returns therefore frees the memh from under UCX. +// +// The data path is pinned to TCP so that a large transfer is genuinely asynchronous: over +// shm/self UCX copies inline, the post completes immediately and there is nothing in flight +// for release() to drain. +class TestTransferRelease : public TestTransfer { +protected: + void + SetUp() override { + env.addVar("NIXL_TELEMETRY_ENABLE", "n"); + env.addVar("UCX_TLS", "tcp"); + addAgent(0); + addAgent(1); + } + + static void + fillBuffers(const std::vector &buffers, uint8_t value) { + for (const auto &buffer : buffers) { + std::memset( + reinterpret_cast(static_cast(buffer)), value, buffer.getSize()); + } + } + + static size_t + countBytes(const std::vector &buffers, uint8_t value) { + size_t found = 0; + for (const auto &buffer : buffers) { + const auto *data = reinterpret_cast(static_cast(buffer)); + found += static_cast(std::count(data, data + buffer.getSize(), value)); + } + return found; + } +}; + +TEST_P(TestTransferRelease, InFlightXferIsDrainedBeforeReleaseReturns) { + // The descriptor count is above the fixture's split_batch_size so that the threadpool + // parameterisations exercise the composite handle rather than the plain one. + constexpr size_t size = 1024 * 1024; + constexpr size_t count = 64; + constexpr uint8_t src_pattern = 0xab; + constexpr uint8_t dst_pattern = 0x00; + + std::vector local_buffers, remote_buffers; + createRegisteredMem(getAgent(0), size, count, DRAM_SEG, local_buffers); + createRegisteredMem(getAgent(1), size, count, DRAM_SEG, remote_buffers); + fillBuffers(remote_buffers, src_pattern); + fillBuffers(local_buffers, dst_pattern); + + exchangeMD(0, 1); + + nixlXferReqH *xfer_req = nullptr; + ASSERT_EQ(getAgent(0).createXferReq(NIXL_READ, + makeDescList(local_buffers, DRAM_SEG), + makeDescList(remote_buffers, DRAM_SEG), + getAgentName(1), + xfer_req), + NIXL_SUCCESS); + ASSERT_NE(xfer_req, nullptr); + + const nixl_status_t post_status = getAgent(0).postXferReq(xfer_req); + ASSERT_TRUE((post_status == NIXL_SUCCESS) || (post_status == NIXL_IN_PROG)); + + // release() only progresses the local worker. Over TCP the peer has to progress too, so + // stand in for a live remote process while the drain runs; without a progress thread + // nothing else would advance agent 1. + std::atomic pump_peer{true}; + std::thread peer_thread([&]() { + nixl_notifs_t notifs; + while (pump_peer.load(std::memory_order_relaxed)) { + getAgent(1).getNotifs(notifs); + } + }); + + // A timed-out drain is a failure of this test, not of the run, so keep it out of the + // global problem counter and assert on it here instead. + const LogIgnoreGuard lig_drain("Still draining in-flight request"); + const LogIgnoreGuard lig_timeout("is still in flight after the drain timeout"); + + const nixl_status_t release_status = getAgent(0).releaseXferReq(xfer_req); + pump_peer = false; + peer_thread.join(); + + EXPECT_EQ(release_status, NIXL_SUCCESS); + EXPECT_EQ(lig_timeout.getIgnoredCount(), 0u) << "release() gave up before the read finished"; + + // A read completes only once the data has landed locally, so a release() that drained + // implies the whole destination is written by the time it returns. Without the drain the + // read is still outstanding here and the destination is only partly filled. + if (post_status == NIXL_IN_PROG) { + EXPECT_EQ(countBytes(local_buffers, src_pattern), size * count) + << "releaseXferReq() returned while the read was still in flight"; + } else { + Logger() << "transfer completed inline, nothing was in flight to drain"; + } + + // Deregistering is what would free the ucp_mem_h from under an in-flight request. + invalidateMD(0, 1); + deregisterMem(getAgent(0), local_buffers, DRAM_SEG); + deregisterMem(getAgent(1), remote_buffers, DRAM_SEG); +} + +NIXL_INSTANTIATE_TEST(ucx, TestTransferRelease, "UCX", true, 2, 0, ""); +NIXL_INSTANTIATE_TEST(ucx_no_pt, TestTransferRelease, "UCX", false, 2, 0, ""); +// The threadpool parameterisations take the composite handle path. +NIXL_INSTANTIATE_TEST(ucx_threadpool, TestTransferRelease, "UCX", true, 6, 4, ""); +NIXL_INSTANTIATE_TEST(ucx_threadpool_no_pt, TestTransferRelease, "UCX", false, 6, 4, ""); + NIXL_INSTANTIATE_TEST(ucx, TestTransfer, "UCX", true, 2, 0, ""); NIXL_INSTANTIATE_TEST(ucx_no_pt, TestTransfer, "UCX", false, 2, 0, ""); NIXL_INSTANTIATE_TEST(ucx_threadpool, TestTransfer, "UCX", true, 6, 4, "");