diff --git a/cpp/include/tensorrt_llm/runtime/utils/mpiTags.h b/cpp/include/tensorrt_llm/runtime/utils/mpiTags.h index 32c086c84ee9..caa0651d2a2f 100644 --- a/cpp/include/tensorrt_llm/runtime/utils/mpiTags.h +++ b/cpp/include/tensorrt_llm/runtime/utils/mpiTags.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2021-2023, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2021-2026, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -71,7 +71,23 @@ enum class MpiTag : int // KvCacheEventManager kKvCacheEventSize = 1026, - kKvCacheEvent = 1027 + kKvCacheEvent = 1027, + + // Fault-tolerant NCCL communicator rendezvous. Each ownership path and + // protocol phase has a distinct tag so stale or simultaneous messages + // cannot be decoded as another phase's payload. + kNcclCommReady = 1028, + kNcclCommUniqueId = 1029, + kNcclCommAck = 1030, + kNcclPpReady = 1031, + kNcclPpUniqueId = 1032, + kNcclPpAck = 1033, + + // Seeds used with the canonical group to derive MPI_Comm_create_group + // tags for dedicated, pre-failure NCCL rendezvous channels. These are + // separate from message tags and distinguish raw-op and PP ownership. + kNcclCommControl = 1034, + kNcclPpControl = 1035 }; } // namespace tensorrt_llm::mpi diff --git a/cpp/include/tensorrt_llm/runtime/utils/ncclHostApi.h b/cpp/include/tensorrt_llm/runtime/utils/ncclHostApi.h new file mode 100644 index 000000000000..a477547902b0 --- /dev/null +++ b/cpp/include/tensorrt_llm/runtime/utils/ncclHostApi.h @@ -0,0 +1,52 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * 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 + +namespace tensorrt_llm::runtime +{ + +// NCCL host APIs must not run concurrently from different threads for +// communicators associated with the same CUDA device. TRT-LLM can own both a +// PP communicator and several cached raw-op communicators, each with its own +// watchdog, so a per-communicator mutex is insufficient. A process-wide gate +// is deliberately stricter than a per-device gate and avoids relying on CUDA +// thread-local device state in watchdog threads. +// +// Recursive locking lets low-level error paths abort a communicator while the +// operation wrapper already owns the gate. A grouped operation keeps one lock +// alive from ncclGroupStart through ncclGroupEnd. +inline std::recursive_mutex& getNcclHostApiMutex() +{ + // Communicator registries and watchdogs are also process-lifetime + // singletons, with cross-translation-unit destruction order unspecified. + // Keep the gate alive until process exit so late communicator teardown can + // never touch an already-destroyed mutex. + static auto* mutex = new std::recursive_mutex; + return *mutex; +} + +using NcclHostApiLock = std::unique_lock; + +inline NcclHostApiLock acquireNcclHostApiLock() +{ + return NcclHostApiLock{getNcclHostApiMutex()}; +} + +} // namespace tensorrt_llm::runtime diff --git a/cpp/include/tensorrt_llm/runtime/utils/ncclUniqueIdRendezvous.h b/cpp/include/tensorrt_llm/runtime/utils/ncclUniqueIdRendezvous.h new file mode 100644 index 000000000000..67b76aa9df3c --- /dev/null +++ b/cpp/include/tensorrt_llm/runtime/utils/ncclUniqueIdRendezvous.h @@ -0,0 +1,102 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * 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 "tensorrt_llm/common/config.h" +#include "tensorrt_llm/runtime/utils/mpiUtils.h" + +#include +#include +#include +#include + +#if ENABLE_MULTI_DEVICE +#include +#endif + +namespace tensorrt_llm::runtime +{ + +#if ENABLE_MULTI_DEVICE + +struct NcclUniqueIdRendezvousTags +{ + mpi::MpiTag ready; + mpi::MpiTag id; + mpi::MpiTag ack; +}; + +//! An MPI control channel dedicated to NCCL unique-ID rendezvous. +//! +//! The channel is created collectively by the initial NCCL group while every +//! member is healthy. It retains a process-lifetime MPI communicator whose +//! error handler is MPI_ERRORS_RETURN, leaving the parent/session communicator unchanged. The +//! original parent-communicator ranks are retained so a later survivor subset +//! can keep using stable world-rank IDs even though MPI ranks in this channel +//! are compact. +//! +//! This isolates rendezvous traffic and error handling, but it is not a ULFM +//! communicator repair. Post-failure point-to-point progress still requires an +//! MPI implementation and launcher configured to let survivors continue. +class NcclUniqueIdRendezvousComm +{ +public: + NcclUniqueIdRendezvousComm(mpi::MpiComm comm, std::vector worldRanks, int worldRank); + + NcclUniqueIdRendezvousComm(NcclUniqueIdRendezvousComm const&) = delete; + NcclUniqueIdRendezvousComm& operator=(NcclUniqueIdRendezvousComm const&) = delete; + + [[nodiscard]] mpi::MpiComm const& mpiComm() const noexcept; + [[nodiscard]] std::vector const& worldRanks() const noexcept; + [[nodiscard]] int worldRank() const noexcept; + [[nodiscard]] int commRank(int worldRank) const; + [[nodiscard]] int worldRank(int commRank) const; + +private: + mpi::MpiComm mComm; + std::vector mWorldRanks; + int mWorldRank; +}; + +//! Create a dedicated control channel over initialRanks. +//! +//! Every member of initialRanks must call this before any process failure. +//! Non-members must not call. creationTagSeed and the canonical group derive a +//! bounded, deterministic MPI creation tag so concurrently-created overlapping +//! groups do not normally cross-pair. +std::shared_ptr createNcclUniqueIdRendezvousComm( + std::vector const& initialRanks, int worldRank, mpi::MpiComm const& parentComm, int creationTagSeed); + +// Exchange a fresh NCCL unique ID among only activeRanks. READY/ID/ACK tokens +// prevent delayed eager MPI messages from an earlier timed-out attempt from +// pairing different IDs across ranks. rendezvousId is a coordinator-provided +// logical attempt identity: all ranks in one attempt must pass the same value, +// and retries for the same logical communicator must use strictly increasing +// values. The protocol discards older same-communicator messages after seeing +// a later value while retaining messages for future values and other groups. +// All ranks use original/global rank IDs. +// Callers must keep one logical communicator per (control channel, tag set, +// active-rank set), or construct multiple instances in the same order on every +// rank. Recovery reuses the pre-failure control channel and performs no MPI +// collective. +ncclUniqueId exchangeNcclUniqueId(std::vector const& activeRanks, NcclUniqueIdRendezvousComm const& controlComm, + NcclUniqueIdRendezvousTags tags, std::uint64_t rendezvousId, std::chrono::steady_clock::time_point deadline); + +#endif // ENABLE_MULTI_DEVICE + +} // namespace tensorrt_llm::runtime diff --git a/cpp/tensorrt_llm/common/attentionOp.cpp b/cpp/tensorrt_llm/common/attentionOp.cpp index 0ea70da0d73b..c9c571bdd0b4 100644 --- a/cpp/tensorrt_llm/common/attentionOp.cpp +++ b/cpp/tensorrt_llm/common/attentionOp.cpp @@ -33,6 +33,7 @@ #include "tensorrt_llm/runtime/utils/debugUtils.h" #include "tensorrt_llm/runtime/utils/mpiUtils.h" #include +#include #include using namespace tensorrt_llm::kernels; @@ -363,20 +364,26 @@ int AttentionOp::ulyssesContextPreprocess(T const* input, T* output, T* buffer, // Do all to all #if ENABLE_MULTI_DEVICE - ncclGroupStart(); - for (int cpIdx = 0; cpIdx < mCpSize; cpIdx++) { - if (cpIdx != mCpRank) + auto commLease = acquireComm(mCpNcclComm); + auto const comm = commLease.get(); + auto const watchdogToken = commLease.begin(stream, "NCCL send/recv(ulysses context preprocess)"); + commLease.groupStart("ncclGroupStart(ulysses context preprocess)"); + auto checkNcclEnqueue = [&](ncclResult_t result) + { commLease.checkEnqueue(result, "NCCL send/recv(ulysses context preprocess)"); }; + for (int cpIdx = 0; cpIdx < mCpSize; cpIdx++) { - NCCLCHECK(ncclSend(output + cpIdx * (partialTokenNum * getHeadSize() * partialHeads), - (partialTokenNum * getHeadSize() * partialHeads), (*getDtypeMap())[mType], cpIdx, *mCpNcclComm, - stream)); - NCCLCHECK(ncclRecv(buffer + cpIdx * (partialTokenNum * getHeadSize() * partialHeads), - (partialTokenNum * getHeadSize() * partialHeads), (*getDtypeMap())[mType], cpIdx, *mCpNcclComm, - stream)); + if (cpIdx != mCpRank) + { + checkNcclEnqueue(ncclSend(output + cpIdx * (partialTokenNum * getHeadSize() * partialHeads), + (partialTokenNum * getHeadSize() * partialHeads), (*getDtypeMap())[mType], cpIdx, comm, stream)); + checkNcclEnqueue(ncclRecv(buffer + cpIdx * (partialTokenNum * getHeadSize() * partialHeads), + (partialTokenNum * getHeadSize() * partialHeads), (*getDtypeMap())[mType], cpIdx, comm, stream)); + } } + commLease.groupEnd("ncclGroupEnd(ulysses context preprocess)"); + commLease.track(watchdogToken, stream); } - ncclGroupEnd(); sync_check_cuda_error(stream); #endif // ENABLE_MULTI_DEVICE @@ -422,28 +429,36 @@ int AttentionOp::ulyssesContextPostprocess(T* input, T* output, T* buffer, Enque // all-to-all #if ENABLE_MULTI_DEVICE size_t const elementNum = partialTokenNum * getHeadSize() * mNumAttnHeads; - ncclGroupStart(); - for (int cpIdx = 0; cpIdx < mCpSize; cpIdx++) { - if (cpIdx != mCpRank) + auto commLease = acquireComm(mCpNcclComm); + auto const comm = commLease.get(); + auto const watchdogToken = commLease.begin(stream, "NCCL send/recv(ulysses context postprocess)"); + commLease.groupStart("ncclGroupStart(ulysses context postprocess)"); + auto checkNcclEnqueue = [&](ncclResult_t result) + { commLease.checkEnqueue(result, "NCCL send/recv(ulysses context postprocess)"); }; + for (int cpIdx = 0; cpIdx < mCpSize; cpIdx++) { - if (mFP8AttenOutput) - { - NCCLCHECK(ncclSend(reinterpret_cast<__nv_fp8_e4m3*>(buffer) + cpIdx * elementNum, elementNum, ncclInt8, - cpIdx, *mCpNcclComm, stream)); - NCCLCHECK(ncclRecv(reinterpret_cast<__nv_fp8_e4m3*>(input) + cpIdx * elementNum, elementNum, ncclInt8, - cpIdx, *mCpNcclComm, stream)); - } - else + if (cpIdx != mCpRank) { - NCCLCHECK(ncclSend( - buffer + cpIdx * elementNum, elementNum, (*getDtypeMap())[mType], cpIdx, *mCpNcclComm, stream)); - NCCLCHECK(ncclRecv( - input + cpIdx * elementNum, elementNum, (*getDtypeMap())[mType], cpIdx, *mCpNcclComm, stream)); + if (mFP8AttenOutput) + { + checkNcclEnqueue(ncclSend(reinterpret_cast<__nv_fp8_e4m3*>(buffer) + cpIdx * elementNum, elementNum, + ncclInt8, cpIdx, comm, stream)); + checkNcclEnqueue(ncclRecv(reinterpret_cast<__nv_fp8_e4m3*>(input) + cpIdx * elementNum, elementNum, + ncclInt8, cpIdx, comm, stream)); + } + else + { + checkNcclEnqueue(ncclSend( + buffer + cpIdx * elementNum, elementNum, (*getDtypeMap())[mType], cpIdx, comm, stream)); + checkNcclEnqueue( + ncclRecv(input + cpIdx * elementNum, elementNum, (*getDtypeMap())[mType], cpIdx, comm, stream)); + } } } + commLease.groupEnd("ncclGroupEnd(ulysses context postprocess)"); + commLease.track(watchdogToken, stream); } - ncclGroupEnd(); #endif // ENABLE_MULTI_DEVICE // transpose_1_reverse + view @@ -489,20 +504,26 @@ int AttentionOp::ulyssesGenerationPreprocess( #if ENABLE_MULTI_DEVICE auto const partialHeads = mNumAttnHeads + 2 * mNumAttnKVHeads; - ncclGroupStart(); - for (int cpIdx = 0; cpIdx < mCpSize; cpIdx++) { - if (cpIdx != mCpRank) + auto commLease = acquireComm(mCpNcclComm); + auto const comm = commLease.get(); + auto const watchdogToken = commLease.begin(stream, "NCCL send/recv(ulysses generation preprocess)"); + commLease.groupStart("ncclGroupStart(ulysses generation preprocess)"); + auto checkNcclEnqueue = [&](ncclResult_t result) + { commLease.checkEnqueue(result, "NCCL send/recv(ulysses generation preprocess)"); }; + for (int cpIdx = 0; cpIdx < mCpSize; cpIdx++) { - NCCLCHECK(ncclSend(buffer + cpIdx * (partialTokenNum * getHeadSize() * partialHeads), - (partialTokenNum * getHeadSize() * partialHeads), (*getDtypeMap())[mType], cpIdx, *mCpNcclComm, - stream)); - NCCLCHECK(ncclRecv(output + cpIdx * (partialTokenNum * getHeadSize() * partialHeads), - (partialTokenNum * getHeadSize() * partialHeads), (*getDtypeMap())[mType], cpIdx, *mCpNcclComm, - stream)); + if (cpIdx != mCpRank) + { + checkNcclEnqueue(ncclSend(buffer + cpIdx * (partialTokenNum * getHeadSize() * partialHeads), + (partialTokenNum * getHeadSize() * partialHeads), (*getDtypeMap())[mType], cpIdx, comm, stream)); + checkNcclEnqueue(ncclRecv(output + cpIdx * (partialTokenNum * getHeadSize() * partialHeads), + (partialTokenNum * getHeadSize() * partialHeads), (*getDtypeMap())[mType], cpIdx, comm, stream)); + } } + commLease.groupEnd("ncclGroupEnd(ulysses generation preprocess)"); + commLease.track(watchdogToken, stream); } - ncclGroupEnd(); sync_check_cuda_error(stream); #endif // ENABLE_MULTI_DEVICE return 0; @@ -525,28 +546,36 @@ int AttentionOp::ulyssesGenerationPostprocess(T* input, T* output, T* buffer, in // do all-to-all #if ENABLE_MULTI_DEVICE size_t const elementNum = partialTokenNum * getHeadSize() * mNumAttnHeads; - ncclGroupStart(); - for (int cpIdx = 0; cpIdx < mCpSize; cpIdx++) { - if (cpIdx != mCpRank) + auto commLease = acquireComm(mCpNcclComm); + auto const comm = commLease.get(); + auto const watchdogToken = commLease.begin(stream, "NCCL send/recv(ulysses generation postprocess)"); + commLease.groupStart("ncclGroupStart(ulysses generation postprocess)"); + auto checkNcclEnqueue = [&](ncclResult_t result) + { commLease.checkEnqueue(result, "NCCL send/recv(ulysses generation postprocess)"); }; + for (int cpIdx = 0; cpIdx < mCpSize; cpIdx++) { - if (mFP8AttenOutput) + if (cpIdx != mCpRank) { - NCCLCHECK(ncclSend(reinterpret_cast<__nv_fp8_e4m3*>(input) + cpIdx * elementNum, elementNum, ncclInt8, - cpIdx, *mCpNcclComm, stream)); - NCCLCHECK(ncclRecv(reinterpret_cast<__nv_fp8_e4m3*>(buffer) + cpIdx * elementNum, elementNum, ncclInt8, - cpIdx, *mCpNcclComm, stream)); - } - else - { - NCCLCHECK(ncclSend( - input + cpIdx * elementNum, elementNum, (*getDtypeMap())[mType], cpIdx, *mCpNcclComm, stream)); - NCCLCHECK(ncclRecv( - buffer + cpIdx * elementNum, elementNum, (*getDtypeMap())[mType], cpIdx, *mCpNcclComm, stream)); + if (mFP8AttenOutput) + { + checkNcclEnqueue(ncclSend(reinterpret_cast<__nv_fp8_e4m3*>(input) + cpIdx * elementNum, elementNum, + ncclInt8, cpIdx, comm, stream)); + checkNcclEnqueue(ncclRecv(reinterpret_cast<__nv_fp8_e4m3*>(buffer) + cpIdx * elementNum, elementNum, + ncclInt8, cpIdx, comm, stream)); + } + else + { + checkNcclEnqueue( + ncclSend(input + cpIdx * elementNum, elementNum, (*getDtypeMap())[mType], cpIdx, comm, stream)); + checkNcclEnqueue(ncclRecv( + buffer + cpIdx * elementNum, elementNum, (*getDtypeMap())[mType], cpIdx, comm, stream)); + } } } + commLease.groupEnd("ncclGroupEnd(ulysses generation postprocess)"); + commLease.track(watchdogToken, stream); } - ncclGroupEnd(); #endif // ENABLE_MULTI_DEVICE // do transpose_1_reverse @@ -3198,11 +3227,13 @@ int AttentionOp::initialize() noexcept return 0; } #if ENABLE_MULTI_DEVICE - if (mCpSize > 1 && COMM_SESSION.getSize() > 1) + if (mCpSize > 1) { - TLLM_LOG_TRACE("%s start for rank %d", __PRETTY_FUNCTION__, COMM_SESSION.getRank()); + TLLM_CHECK_WITH_INFO(mCpGroup.size() == static_cast(mCpSize), + "Context-parallel group size %zu does not match configured CP size %d", mCpGroup.size(), mCpSize); + TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); mCpNcclComm = getComm(mCpGroup); - TLLM_LOG_TRACE("%s stop for rank %d", __PRETTY_FUNCTION__, COMM_SESSION.getRank()); + TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); } #endif // ENABLE_MULTI_DEVICE return 0; diff --git a/cpp/tensorrt_llm/common/ncclUtils.cpp b/cpp/tensorrt_llm/common/ncclUtils.cpp index e36cfddcd404..980fc502fd92 100644 --- a/cpp/tensorrt_llm/common/ncclUtils.cpp +++ b/cpp/tensorrt_llm/common/ncclUtils.cpp @@ -21,8 +21,13 @@ #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/cudaUtils.h" #include "tensorrt_llm/common/logger.h" +#include "tensorrt_llm/common/opUtils.h" +#include "tensorrt_llm/runtime/utils/ncclHostApi.h" +#include +#include #include #include +#include namespace { @@ -70,6 +75,7 @@ struct NcclMemGuard { if (ptr) { + auto const hostApiLock = tensorrt_llm::runtime::acquireNcclHostApiLock(); TLLM_NCCL_CHECK_WARN(ncclMemFree(ptr)); } } @@ -97,6 +103,62 @@ namespace constexpr int kNcclWindowMinRuntimeVersion = NCCL_VERSION(2, 28, 0); constexpr int kNcclGb10WindowFixedVersion = NCCL_VERSION(2, 30, 4); constexpr int kGb10RealSmVersion = 121; +constexpr auto kRawNcclPollInterval = std::chrono::milliseconds{1}; + +struct RawNcclCallCompletion +{ + ncclResult_t result; + bool completed; +}; + +std::chrono::milliseconds getRawNcclCallTimeout() +{ + constexpr int64_t defaultTimeoutMs = 5000; + auto const* value = std::getenv("TRTLLM_NCCL_NONBLOCKING_TIMEOUT_MS"); + if (value != nullptr) + { + try + { + auto const parsed = std::stoll(value); + if (parsed > 0) + { + return std::chrono::milliseconds{parsed}; + } + } + catch (...) + { + } + } + return std::chrono::milliseconds{defaultTimeoutMs}; +} + +// The caller must hold the process-wide NCCL host-API gate from the initial +// call through this poll. NCCL permits only ncclCommGetAsyncError after a +// nonblocking API returns ncclInProgress. +RawNcclCallCompletion completeRawNcclCall( + ncclComm_t comm, ncclResult_t initialResult, std::chrono::steady_clock::time_point deadline) noexcept +{ + if (initialResult != ncclInProgress) + { + return {initialResult, true}; + } + + while (std::chrono::steady_clock::now() < deadline) + { + ncclResult_t asyncResult = ncclInProgress; + auto const queryResult = ncclCommGetAsyncError(comm, &asyncResult); + if (queryResult == ncclSuccess && asyncResult != ncclInProgress) + { + return {asyncResult, true}; + } + if (queryResult != ncclSuccess && queryResult != ncclInProgress) + { + return {queryResult, true}; + } + std::this_thread::sleep_for(kRawNcclPollInterval); + } + return {ncclInProgress, false}; +} bool isGb10Platform(int realSmVersion, bool isIntegrated) { @@ -107,6 +169,7 @@ bool isGb10Platform(int realSmVersion, bool isIntegrated) bool queryNcclWindowSupported() { #if NCCL_VERSION_CODE >= NCCL_VERSION(2, 28, 0) + auto const hostApiLock = tensorrt_llm::runtime::acquireNcclHostApiLock(); int version = 0; if (ncclGetVersion(&version) != ncclSuccess) { @@ -211,8 +274,11 @@ bool isNcclWindowSupported() NcclCommResourceManager& NcclCommResourceManager::getInstance() noexcept { - static NcclCommResourceManager instance; - return instance; + // Communicator registries and watchdogs are process-lifetime. Keep their + // resource owner alive until process exit instead of relying on undefined + // cross-translation-unit singleton destruction order. + static auto* instance = new NcclCommResourceManager; + return *instance; } NcclCommResourceManager::~NcclCommResourceManager() @@ -271,18 +337,44 @@ void NcclCommResourceManager::registerResource(ncclComm_t comm, ResourceCleanupF debugName ? debugName : "unnamed", static_cast(comm), resources.size()); } -void NcclCommResourceManager::cleanupResources(ncclComm_t comm) noexcept +void NcclCommResourceManager::beginAbortCleanup(ncclComm_t comm) noexcept { - if (!comm) + if (comm == nullptr) { return; } + std::lock_guard lock(mMutex); + mAbortCleanups.insert(comm); +} + +void NcclCommResourceManager::waitForAbortCleanup(ncclComm_t comm) noexcept +{ + if (comm == nullptr) + { + return; + } + std::unique_lock lock(mMutex); + mAbortCleanupComplete.wait(lock, [this, comm]() { return mAbortCleanups.find(comm) == mAbortCleanups.end(); }); +} + +bool NcclCommResourceManager::cleanupResources(ncclComm_t comm, bool communicatorAborted) noexcept +{ + if (!comm) + { + return true; + } // Check if we're in the process of being destroyed // If so, skip cleanup - the destructor will handle it proactively if (mIsDestroying.load(std::memory_order_acquire)) { - return; + if (communicatorAborted) + { + std::lock_guard lock(mMutex); + mAbortCleanups.erase(comm); + mAbortCleanupComplete.notify_all(); + } + return false; } std::vector resourcesToClean; @@ -297,20 +389,42 @@ void NcclCommResourceManager::cleanupResources(ncclComm_t comm) noexcept // Double-check after acquiring lock (destruction may have started) if (mIsDestroying.load(std::memory_order_acquire)) { - return; + if (communicatorAborted) + { + mAbortCleanups.erase(comm); + mAbortCleanupComplete.notify_all(); + } + return false; + } + + if (communicatorAborted) + { + mAbortCleanups.insert(comm); + } + if (mCleanupsInProgress.find(comm) != mCleanupsInProgress.end()) + { + // Another thread already owns the callbacks. Its completion + // path rechecks the retirement marker, including one that was + // installed after that owner started, before waking handle + // reusers. + return true; } auto it = mCommResources.find(comm); if (it == mCommResources.end()) { // Nothing registered for this comm, nothing to clean up - return; + if (mAbortCleanups.erase(comm) != 0) + { + mAbortCleanupComplete.notify_all(); + } + return true; } // Move resources out (preserves order) and remove from map resourcesToClean = std::move(it->second); mCommResources.erase(it); - + mCleanupsInProgress.insert(comm); // Logging may fail during static destruction, so wrap in try-catch try { @@ -326,12 +440,13 @@ void NcclCommResourceManager::cleanupResources(ncclComm_t comm) noexcept { // If mutex access fails during static destruction, just return. // This prevents segfaults when the singleton is being destroyed. - return; + return false; } } // Clean up outside the lock to avoid deadlocks if cleanup functions try to access the manager // Order is preserved: resources are cleaned up in registration order + bool cleanupSucceeded = true; for (auto& [cleanup, name] : resourcesToClean) { try @@ -346,10 +461,11 @@ void NcclCommResourceManager::cleanupResources(ncclComm_t comm) noexcept { // Ignore logging failures during static destruction } - cleanup(); + cleanupSucceeded = cleanup() && cleanupSucceeded; } catch (std::exception const& e) { + cleanupSucceeded = false; try { TLLM_LOG_ERROR("[NCCLUtil] Exception during cleanup of resource '%s' for NCCL comm %p: %s", @@ -362,6 +478,7 @@ void NcclCommResourceManager::cleanupResources(ncclComm_t comm) noexcept } catch (...) { + cleanupSucceeded = false; try { TLLM_LOG_ERROR("[NCCLUtil] Unknown exception during cleanup of resource '%s' for NCCL comm %p", @@ -373,6 +490,22 @@ void NcclCommResourceManager::cleanupResources(ncclComm_t comm) noexcept } } } + + { + std::lock_guard lock(mMutex); + mCleanupsInProgress.erase(comm); + if (mAbortCleanups.erase(comm) != 0) + { + mAbortCleanupComplete.notify_all(); + } + } + return cleanupSucceeded; +} + +bool NcclCommResourceManager::isAbortCleanup(ncclComm_t comm) const noexcept +{ + std::lock_guard lock(mMutex); + return mAbortCleanups.find(comm) != mAbortCleanups.end(); } bool NcclCommResourceManager::hasResources(ncclComm_t comm) const noexcept @@ -396,93 +529,171 @@ size_t NcclCommResourceManager::getResourceCount(ncclComm_t comm) const noexcept NCCLWindowAllocator& NCCLWindowAllocator::getInstance() { - static NCCLWindowAllocator instance; - return instance; + static auto* instance = new NCCLWindowAllocator; + return *instance; } NCCLWindowBuffer NCCLWindowAllocator::requestBuffer(ncclComm_t comm, size_t size) +{ + return requestBufferImpl(comm, size, nullptr, nullptr); +} + +NCCLWindowBuffer NCCLWindowAllocator::requestBuffer( + std::shared_ptr const& comm, size_t size, ncclComm_t* associatedComm) +{ + TLLM_CHECK_WITH_INFO(comm != nullptr, "NCCL communicator cannot be null"); + if (associatedComm != nullptr) + { + *associatedComm = nullptr; + } + return requestBufferImpl(nullptr, size, &comm, associatedComm); +} + +NCCLWindowBuffer NCCLWindowAllocator::requestBufferImpl( + ncclComm_t comm, size_t size, std::shared_ptr const* managedComm, ncclComm_t* associatedComm) { if (!isNcclWindowSupported()) { return NCCLWindowBuffer(); } - TLLM_CHECK_WITH_INFO(comm != nullptr, "NCCL communicator cannot be null"); TLLM_CHECK_WITH_INFO(size > 0, "Buffer size must be greater than 0"); - std::lock_guard lock(mMutex); + bool const managedFaultTolerance = managedComm != nullptr && isNcclFaultToleranceEnabled(); + if (managedComm == nullptr) + { + TLLM_CHECK_WITH_INFO(comm != nullptr, "NCCL communicator cannot be null"); + NcclCommResourceManager::getInstance().waitForAbortCleanup(comm); + } + // Pool hits need only mMutex. On a miss, release all communicator/pool + // locks, serialize allocation attempts, then recheck: another thread may + // have populated the pool while this thread waited. This preserves + // collective ordering without adding a second lock to the steady-state + // pooled-buffer path or inverting the host-gate -> pool-lock cleanup order. + std::unique_lock allocationLock(mAllocationMutex, std::defer_lock); + int handle = 0; + while (true) + { + std::unique_ptr validationLease; + if (managedComm != nullptr) + { + if (managedFaultTolerance) + { + // Keep lock ordering consistent with abort cleanup: + // communicator state / NCCL host gate precede the pool mutex. + validationLease = std::make_unique(acquireComm(*managedComm)); + comm = validationLease->get(); + } + else + { + TLLM_CHECK_WITH_INFO(*managedComm != nullptr, "NCCL communicator cannot be null"); + comm = **managedComm; + } + } + TLLM_CHECK_WITH_INFO(comm != nullptr, "NCCL communicator cannot be null"); + if (associatedComm != nullptr) + { + *associatedComm = comm; + } + std::unique_lock poolLock(mMutex); - // Register cleanup callback for this communicator if not already registered - // This is cheap even if no buffers exist yet - cleanup will just return early - registerBufferCleanup(comm); + // Register cleanup callback for this communicator if not already registered. + registerBufferCleanup(comm); + // Keep validationLease until poolLock releases. Lease destruction can + // drain abort cleanup, which re-enters this non-recursive pool mutex. - // Check if we have an available buffer of at least the requested size for this communicator - // Use best-fit: find the smallest buffer that's >= requested size - auto& commBuffers = mBufferPool[comm]; - auto bestFit = commBuffers.end(); - size_t bestFitSize = std::numeric_limits::max(); + auto& commBuffers = mBufferPool[comm]; + auto bestFit = commBuffers.end(); + size_t bestFitSize = std::numeric_limits::max(); - for (auto it = commBuffers.begin(); it != commBuffers.end(); ++it) - { - if (!it->inUse && it->buffer.size >= size && it->buffer.size < bestFitSize) + for (auto it = commBuffers.begin(); it != commBuffers.end(); ++it) { - bestFit = it; - bestFitSize = it->buffer.size; + if (!it->inUse && it->buffer.size >= size && it->buffer.size < bestFitSize) + { + bestFit = it; + bestFitSize = it->buffer.size; + } } - } - if (bestFit != commBuffers.end()) - { - bestFit->inUse = true; - TLLM_LOG_TRACE( - "[NCCLUtil] Reusing NCCL window buffer for comm %p: handle=%d, ptr=%p, size=%zu (requested: %zu)", - static_cast(comm), bestFit->buffer.handle, bestFit->buffer.ptr, bestFit->buffer.size, size); - return bestFit->buffer; - } + if (bestFit != commBuffers.end()) + { + bestFit->inUse = true; + TLLM_LOG_TRACE( + "[NCCLUtil] Reusing NCCL window buffer for comm %p: handle=%d, ptr=%p, size=%zu (requested: %zu)", + static_cast(comm), bestFit->buffer.handle, bestFit->buffer.ptr, bestFit->buffer.size, size); + return bestFit->buffer; + } - // If a previous allocateAndRegisterBuffer call collectively failed for this comm at a size - // no larger than this request, do not retry the known-failing new allocation path. Smaller - // requests and already-pooled buffers can still use NCCL windows. - auto const failureIt = mMinSymmetricFailureSize.find(comm); - if (failureIt != mMinSymmetricFailureSize.end() && size >= failureIt->second) - { - TLLM_LOG_DEBUG("[NCCLUtil] Skipping NCCL window allocation for comm %p, size=%zu; known failure threshold=%zu", - static_cast(comm), size, failureIt->second); - return NCCLWindowBuffer(); - } + auto const failureIt = mMinSymmetricFailureSize.find(comm); + if (failureIt != mMinSymmetricFailureSize.end() && size >= failureIt->second) + { + TLLM_LOG_DEBUG( + "[NCCLUtil] Skipping NCCL window allocation for comm %p, size=%zu; known failure threshold=%zu", + static_cast(comm), size, failureIt->second); + return NCCLWindowBuffer(); + } - // No available buffer found, avoid registration during CUDA graph capture - auto stream = at::cuda::getCurrentCUDAStream(); - cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone; - auto capture_err = cudaStreamIsCapturing(stream, &capture_status); - if (capture_err != cudaSuccess) - { - TLLM_LOG_DEBUG("[NCCLUtil] cudaStreamIsCapturing failed: %s", cudaGetErrorString(capture_err)); - } - if (capture_err == cudaSuccess && capture_status != cudaStreamCaptureStatusNone) - { - TLLM_LOG_DEBUG("[NCCLUtil] Skipping NCCL window allocation during capture for comm %p (requested: %zu)", - static_cast(comm), size); - return NCCLWindowBuffer(); + auto stream = at::cuda::getCurrentCUDAStream(); + cudaStreamCaptureStatus captureStatus = cudaStreamCaptureStatusNone; + auto const captureResult = cudaStreamIsCapturing(stream, &captureStatus); + if (captureResult != cudaSuccess) + { + TLLM_LOG_DEBUG("[NCCLUtil] cudaStreamIsCapturing failed: %s", cudaGetErrorString(captureResult)); + } + if (captureResult == cudaSuccess && captureStatus != cudaStreamCaptureStatusNone) + { + TLLM_LOG_DEBUG("[NCCLUtil] Skipping NCCL window allocation during capture for comm %p (requested: %zu)", + static_cast(comm), size); + return NCCLWindowBuffer(); + } + + if (!allocationLock.owns_lock()) + { + poolLock.unlock(); + validationLease.reset(); + allocationLock.lock(); + continue; + } + + TLLM_LOG_TRACE( + "[NCCLUtil] Allocating new NCCL window buffer for comm %p, size=%zu", static_cast(comm), size); + handle = static_cast(commBuffers.size()); + poolLock.unlock(); + validationLease.reset(); + break; } - // No available buffer found, allocate a new one - TLLM_LOG_TRACE( - "[NCCLUtil] Allocating new NCCL window buffer for comm %p, size=%zu", static_cast(comm), size); - int handle = static_cast(commBuffers.size()); - NCCLWindowBuffer buffer = allocateAndRegisterBuffer(comm, size, handle); + // CUDA allocation must not hold the communicator mutex or process-wide + // NCCL gate: a watchdog needs both in order to abort stalled work. + std::unique_ptr completionLease; + NCCLWindowBuffer buffer = allocateAndRegisterBuffer(comm, size, handle, managedComm, &completionLease); // Only cache valid buffers. allocateAndRegisterBuffer returns an empty buffer when any rank // failed ncclMemAlloc (collective fallback to plain allreduce); caching it would leak a // permanently "in use" empty entry per request because releaseBuffer is a no-op for nullptr. if (buffer.isValid()) { - commBuffers.push_back({buffer, true}); + // The managed registration lease remains live until the buffer is in + // the pool. That closes the gap where the watchdog could abort and run + // cleanup after registration but before the resource became visible. + { + std::lock_guard poolLock(mMutex); + mBufferPool[comm].push_back({buffer, true}); + } + completionLease.reset(); } else { + std::unique_ptr healthyLease; + if (managedFaultTolerance) + { + healthyLease = std::make_unique(acquireComm(*managedComm)); + TLLM_CHECK_WITH_INFO( + healthyLease->get() == comm, "NCCL communicator changed while recording a window-allocation fallback"); + } // The collective allreduce inside allocateAndRegisterBuffer agreed that this request // cannot use symmetric memory on at least one rank. Remember the smallest failing // request size so repeated too-large autotuner probes do not keep stressing this path. + std::lock_guard poolLock(mMutex); recordSymmetricFailureLocked(comm, size); } @@ -607,7 +818,8 @@ bool NCCLWindowAllocator::isCommValid(ncclComm_t comm) const noexcept return comm != nullptr; } -NCCLWindowBuffer NCCLWindowAllocator::allocateAndRegisterBuffer(ncclComm_t comm, size_t size, int handle) +NCCLWindowBuffer NCCLWindowAllocator::allocateAndRegisterBuffer(ncclComm_t comm, size_t size, int handle, + std::shared_ptr const* managedComm, std::unique_ptr* completionLease) { // Step 1: Pre-allocate the rank-sync flag before ncclMemAlloc. ncclMemAlloc can fail // asymmetrically with ncclUnhandledCudaError on configurations where the symmetric/VMM path @@ -626,7 +838,10 @@ NCCLWindowBuffer NCCLWindowAllocator::allocateAndRegisterBuffer(ncclComm_t comm, // the stream-ordered flag copy and collective fallback so the failing rank still reaches // ncclAllReduce with the other ranks. void* ncclPtr = nullptr; - TLLM_NCCL_CHECK_WARN(ncclMemAlloc(&ncclPtr, size)); + { + auto const hostApiLock = tensorrt_llm::runtime::acquireNcclHostApiLock(); + TLLM_NCCL_CHECK_WARN(ncclMemAlloc(&ncclPtr, size)); + } int const localAllocOk = (ncclPtr != nullptr) ? 1 : 0; NcclMemGuard ncclGuard{ncclPtr}; // frees ncclPtr on any early return or exception clearCudaErrorIfSymmetricAllocationFailed(localAllocOk); @@ -640,8 +855,66 @@ NCCLWindowBuffer NCCLWindowAllocator::allocateAndRegisterBuffer(ncclComm_t comm, TLLM_CUDA_CHECK_WARN( cudaMemcpyAsync(rankSyncFlag, &localAllocOk, sizeof(localAllocOk), cudaMemcpyHostToDevice, stream)); } - TLLM_NCCL_CHECK(ncclAllReduce(rankSyncFlag, rankSyncFlag, 1, ncclInt32, ncclMin, comm, stream)); - TLLM_CUDA_CHECK_WARN(cudaStreamSynchronize(stream)); + if (managedComm != nullptr) + { + try + { + std::uint64_t watchdogToken = 0; + { + auto commLease = acquireComm(*managedComm); + auto const currentComm = commLease.get(); + TLLM_CHECK_WITH_INFO(currentComm == comm, + "NCCL communicator changed while allocating a window buffer; retry with the active communicator"); + watchdogToken = commLease.begin(stream, "NCCL window allocation agreement"); + auto const allocAgreementResult + = ncclAllReduce(rankSyncFlag, rankSyncFlag, 1, ncclInt32, ncclMin, currentComm, stream); + commLease.check(allocAgreementResult, "ncclAllReduce(window allocation agreement)"); + commLease.track(watchdogToken, stream); + } + if (isNcclFaultToleranceEnabled()) + { + waitCommOperation(*managedComm, watchdogToken, "NCCL window allocation agreement"); + } + else + { + TLLM_CUDA_CHECK_WARN(cudaStreamSynchronize(stream)); + } + } + catch (...) + { + // A failed NCCL kernel may still hold the temporary flag even + // after communicator abort. Quarantine four bytes instead of + // letting cudaFree introduce a device-wide wait on recovery. + static_cast(flagGuard.release()); + throw; + } + } + else + { + RawNcclCallCompletion completion{ncclSuccess, true}; + { + auto const hostApiLock = tensorrt_llm::runtime::acquireNcclHostApiLock(); + auto const allocAgreementResult + = ncclAllReduce(rankSyncFlag, rankSyncFlag, 1, ncclInt32, ncclMin, comm, stream); + completion = completeRawNcclCall( + comm, allocAgreementResult, std::chrono::steady_clock::now() + getRawNcclCallTimeout()); + } + if (!completion.completed) + { + // The background enqueue may still acquire the flag after the + // local deadline. Keep its storage alive rather than racing it. + static_cast(flagGuard.release()); + TLLM_THROW("NCCL error: window allocation agreement timed out while the communicator was in progress"); + } + if (completion.result != ncclSuccess) + { + // NCCL documents device work as indeterminate after an async + // communicator error. Quarantine the flag before surfacing it. + static_cast(flagGuard.release()); + TLLM_NCCL_CHECK(completion.result); + } + TLLM_CUDA_CHECK_WARN(cudaStreamSynchronize(stream)); + } int allAllocOk = 0; TLLM_CUDA_CHECK(cudaMemcpy(&allAllocOk, rankSyncFlag, sizeof(int), cudaMemcpyDeviceToHost)); @@ -663,10 +936,60 @@ NCCLWindowBuffer NCCLWindowAllocator::allocateAndRegisterBuffer(ncclComm_t comm, // Failure here is non-fatal: warn and fall back to regular allreduce. // ncclGuard frees ncclPtr on return. ncclWindow_t window = nullptr; - ncclResult_t const regResult = ncclCommWindowRegister(comm, ncclPtr, size, &window, NCCL_WIN_COLL_SYMMETRIC); - TLLM_NCCL_CHECK_WARN(regResult); + ncclResult_t regResult = ncclSuccess; + if (managedComm != nullptr) + { + auto commLease = acquireComm(*managedComm); + auto const currentComm = commLease.get(); + TLLM_CHECK_WITH_INFO(currentComm == comm, + "NCCL communicator changed while registering a window buffer; retry with the active communicator"); + regResult = ncclCommWindowRegister(currentComm, ncclPtr, size, &window, NCCL_WIN_COLL_SYMMETRIC); + try + { + regResult = commLease.checkOptional(regResult, "ncclCommWindowRegister"); + } + catch (...) + { + // Registration may still reference the allocation after an + // asynchronous communicator failure. Quarantine it instead + // of calling ncclMemFree on the recovery path. + static_cast(ncclGuard.release()); + throw; + } + if (regResult == ncclSuccess && completionLease != nullptr && isNcclFaultToleranceEnabled()) + { + *completionLease = std::make_unique(std::move(commLease)); + } + } + else + { + ncclResult_t initialResult = ncclSuccess; + RawNcclCallCompletion completion{ncclSuccess, true}; + auto const hostApiLock = tensorrt_llm::runtime::acquireNcclHostApiLock(); + initialResult = ncclCommWindowRegister(comm, ncclPtr, size, &window, NCCL_WIN_COLL_SYMMETRIC); + completion + = completeRawNcclCall(comm, initialResult, std::chrono::steady_clock::now() + getRawNcclCallTimeout()); + if (!completion.completed) + { + // Registration can still complete in the background and retain + // ncclPtr. The unmanaged caller owns communicator recovery, so + // quarantine the allocation locally and surface the timeout. + static_cast(ncclGuard.release()); + TLLM_THROW("NCCL error: window registration timed out while the communicator was in progress"); + } + regResult = completion.result; + if (initialResult == ncclInProgress && regResult != ncclSuccess) + { + // Once an asynchronous registration fails, NCCL does not promise + // that the allocation is no longer referenced. + static_cast(ncclGuard.release()); + TLLM_THROW("NCCL error: window registration failed asynchronously: %s", ncclGetErrorString(regResult)); + } + } if (regResult != ncclSuccess) { + TLLM_LOG_WARNING("NCCL window registration failed for comm %p; falling back to regular allreduce: %s", + static_cast(comm), ncclGetErrorString(regResult)); return NCCLWindowBuffer{}; } @@ -709,24 +1032,38 @@ void NCCLWindowAllocator::registerBufferCleanup(ncclComm_t comm) // Register cleanup with the resource manager NcclCommResourceManager::getInstance().registerResource( - comm, [this, comm]() { this->cleanupBuffersForComm(comm); }, "NCCLWindowAllocator"); + comm, [this, comm]() { return this->cleanupBuffersForComm(comm); }, "NCCLWindowAllocator"); } -void NCCLWindowAllocator::cleanupBuffersForComm(ncclComm_t comm) noexcept +bool NCCLWindowAllocator::cleanupBuffersForComm(ncclComm_t comm) noexcept { if (!comm) { - return; + return true; } - // Synchronize CUDA to ensure all operations using these buffers are complete - // before we deregister windows and free memory - cudaError_t cudaErr = cudaDeviceSynchronize(); - if (cudaErr != cudaSuccess) + bool const communicatorAborted = NcclCommResourceManager::getInstance().isAbortCleanup(comm); + bool cudaSynchronized = true; + + // Never hold the process-wide NCCL host gate across a CUDA synchronization: + // work from another failed communicator may be what the synchronization is + // waiting for, and that communicator's watchdog needs the gate to abort it. + tensorrt_llm::runtime::NcclHostApiLock hostApiLock; + if (!communicatorAborted) { - TLLM_LOG_WARNING("[NCCLUtil] cudaDeviceSynchronize failed with error: %d before cleanup for comm %p", cudaErr, - static_cast(comm)); - // Continue anyway - the sync failure might be from a previous error + cudaError_t cudaErr = cudaDeviceSynchronize(); + if (cudaErr != cudaSuccess) + { + cudaSynchronized = false; + TLLM_LOG_WARNING("[NCCLUtil] cudaDeviceSynchronize failed with error: %d before cleanup for comm %p", + cudaErr, static_cast(comm)); + } + if (cudaSynchronized) + { + // Healthy cleanup uses NCCL deregistration/free below. Acquire the gate + // before mMutex to preserve the communicator -> allocator lock order. + hostApiLock = tensorrt_llm::runtime::acquireNcclHostApiLock(); + } } std::lock_guard lock(mMutex); @@ -735,7 +1072,7 @@ void NCCLWindowAllocator::cleanupBuffersForComm(ncclComm_t comm) noexcept if (mRegisteredComms.find(comm) == mRegisteredComms.end()) { // Already cleaned up or never registered - return; + return true; } auto commIt = mBufferPool.find(comm); @@ -744,7 +1081,7 @@ void NCCLWindowAllocator::cleanupBuffersForComm(ncclComm_t comm) noexcept // No buffers to clean up, but mark as cleaned mRegisteredComms.erase(comm); mMinSymmetricFailureSize.erase(comm); - return; + return true; } TLLM_LOG_TRACE( @@ -772,25 +1109,50 @@ void NCCLWindowAllocator::cleanupBuffersForComm(ncclComm_t comm) noexcept TLLM_LOG_DEBUG("[NCCLUtil] NCCL window allocator teardown for comm %p: %zu buffers, %zu bytes total", static_cast(comm), commIt->second.size(), totalBytes); + if (communicatorAborted || !cudaSynchronized) + { + // Tensor storages or non-NCCL CUDA work may still reference these + // allocations. There is no safe stream fence once the communicator + // has failed, so quarantine them for the remaining process lifetime + // instead of risking a UAF or another blocking CUDA call. Rank failure + // is rare and process restart ultimately reclaims the memory. + TLLM_LOG_WARNING("[NCCLUtil] Quarantining %zu NCCL window buffers (%zu bytes) after %s", commIt->second.size(), + totalBytes, communicatorAborted ? "communicator abort" : "failed CUDA synchronization"); + mBufferPool.erase(commIt); + mRegisteredComms.erase(comm); + mMinSymmetricFailureSize.erase(comm); + return communicatorAborted; + } + + bool quarantineRemainingBuffers = false; for (auto& entry : commIt->second) { if (entry.buffer.isValid()) { + if (quarantineRemainingBuffers) + { + continue; + } // Deregister the window - the communicator is still valid at this point // (cleanup happens before ncclCommDestroy), but we need to be careful // if buffers are still in use by active operations - if (entry.buffer.window && comm) + if (entry.buffer.window && comm && !communicatorAborted) { // Note: Even if buffer is marked inUse, we must deregister since // the communicator is being destroyed. The communicator is valid, // but we should handle potential errors gracefully. - ncclResult_t result = ncclCommWindowDeregister(comm, entry.buffer.window); - if (result != ncclSuccess) + auto const initialResult = ncclCommWindowDeregister(comm, entry.buffer.window); + auto const completion = completeRawNcclCall( + comm, initialResult, std::chrono::steady_clock::now() + getRawNcclCallTimeout()); + if (!completion.completed || completion.result != ncclSuccess) { TLLM_LOG_WARNING( - "[NCCLUtil] ncclCommWindowDeregister failed with error: %d for comm %p, " - "window %p (buffer inUse: %d)", - result, static_cast(comm), static_cast(entry.buffer.window), entry.inUse); + "[NCCLUtil] ncclCommWindowDeregister %s for comm %p, window %p " + "(result: %s, buffer inUse: %d); quarantining this and remaining buffers", + completion.completed ? "failed" : "timed out", static_cast(comm), + static_cast(entry.buffer.window), ncclGetErrorString(completion.result), entry.inUse); + quarantineRemainingBuffers = true; + continue; } } @@ -820,6 +1182,7 @@ void NCCLWindowAllocator::cleanupBuffersForComm(ncclComm_t comm) noexcept mBufferPool.erase(commIt); mRegisteredComms.erase(comm); mMinSymmetricFailureSize.erase(comm); + return !quarantineRemainingBuffers; } #endif // NCCL_VERSION_CODE >= NCCL_VERSION(2, 28, 0) diff --git a/cpp/tensorrt_llm/common/ncclUtils.h b/cpp/tensorrt_llm/common/ncclUtils.h index fdc88f1aaaba..2b629cf72212 100644 --- a/cpp/tensorrt_llm/common/ncclUtils.h +++ b/cpp/tensorrt_llm/common/ncclUtils.h @@ -30,11 +30,13 @@ #include #include +#include #include #include #include #include #include +#include #include #include #include @@ -61,6 +63,10 @@ TRTLLM_NAMESPACE_BEGIN +class NcclCommLease; +bool isNcclFaultToleranceEnabled(); +std::optional getCommAsyncError(std::shared_ptr const& comm); + namespace common::nccl_util { @@ -68,8 +74,10 @@ namespace common::nccl_util // NCCL Resource Management //============================================================================== -// Resource cleanup function type. Called before the NCCL communicator is destroyed. -using ResourceCleanupFunc = std::function; +// Resource cleanup function type. Called before the NCCL communicator is +// destroyed. False tells the owner that graceful communicator destruction is +// unsafe and it must take the abort/quarantine path instead. +using ResourceCleanupFunc = std::function; // Manages resources associated with NCCL communicators. Thread-safe singleton that maintains // a pool of resources per NCCL comm. Resources are automatically cleaned up when the @@ -88,7 +96,18 @@ class NcclCommResourceManager // the shared_ptr deleter before ncclCommDestroy. // Thread-safe: Uses global mutex to serialize cleanup operations. // Order-preserving: Resources are cleaned up in registration order. - void cleanupResources(ncclComm_t comm) noexcept; + [[nodiscard]] bool cleanupResources(ncclComm_t comm, bool communicatorAborted = false) noexcept; + + // Mark a successfully-aborted native handle as retired before releasing + // the NCCL host gate. A later communicator whose native pointer is reused + // waits for the old pointer-keyed cleanup before it can be published. + void beginAbortCleanup(ncclComm_t comm) noexcept; + void waitForAbortCleanup(ncclComm_t comm) noexcept; + + // True only while abort-specific callbacks for comm are running. Resource + // owners use this to skip operations that require a live communicator or + // a device-wide synchronization (both can hang on a failed communicator). + bool isAbortCleanup(ncclComm_t comm) const noexcept; // Check if a communicator has registered resources. bool hasResources(ncclComm_t comm) const noexcept; @@ -108,7 +127,13 @@ class NcclCommResourceManager using ResourceEntry = std::pair; mutable std::mutex mMutex; + std::condition_variable mAbortCleanupComplete; std::unordered_map> mCommResources; + std::unordered_set mAbortCleanups; + // Track every callback owner, not just callers explicitly marked as + // aborted. A normal cleanup may already own the callbacks when another + // thread retires the same native handle. + std::unordered_set mCleanupsInProgress; std::atomic mIsDestroying{false}; }; @@ -134,6 +159,7 @@ class NcclCommResource { mCleanup(mResource); } + return true; }, debugName); } @@ -223,8 +249,18 @@ class NCCLWindowAllocator // If an unused buffer of at least the requested size exists for this communicator, it will be reused. // Uses best-fit strategy: selects the smallest available buffer that meets the size requirement. // Otherwise, a new buffer is allocated and registered. + // Caller-owned/raw communicator overload. It supports a nonblocking call + // result by bounded polling, but it cannot coordinate with the TRT-LLM + // registry's abort/replacement lock. Registry-managed communicators must + // use the shared_ptr overload below. NCCLWindowBuffer requestBuffer(ncclComm_t comm, size_t size); + // Managed-communicator overload. CUDA allocation and tensor construction + // happen without an operation lease; the allocator acquires short leases + // only around NCCL calls so a watchdog can abort stalled work. + NCCLWindowBuffer requestBuffer( + std::shared_ptr const& comm, size_t size, ncclComm_t* associatedComm = nullptr); + // Search for a buffer by pointer. Returns an invalid buffer if not found. // This matches the UBManager.search_buffer() interface. NCCLWindowBuffer searchBuffer(ncclComm_t comm, void* ptr) const; @@ -264,7 +300,10 @@ class NCCLWindowAllocator ~NCCLWindowAllocator() = default; // Allocate a new buffer and register it with NCCL as a window - NCCLWindowBuffer allocateAndRegisterBuffer(ncclComm_t comm, size_t size, int handle); + NCCLWindowBuffer requestBufferImpl( + ncclComm_t comm, size_t size, std::shared_ptr const* managedComm, ncclComm_t* associatedComm); + NCCLWindowBuffer allocateAndRegisterBuffer(ncclComm_t comm, size_t size, int handle, + std::shared_ptr const* managedComm, std::unique_ptr* completionLease); // Record a failed new symmetric allocation (assumes mMutex is already locked). void recordSymmetricFailureLocked(ncclComm_t comm, size_t size); @@ -282,7 +321,7 @@ class NCCLWindowAllocator void registerBufferCleanup(ncclComm_t comm); // Cleanup all buffers for a specific communicator - void cleanupBuffersForComm(ncclComm_t comm) noexcept; + [[nodiscard]] bool cleanupBuffersForComm(ncclComm_t comm) noexcept; struct BufferEntry { @@ -291,6 +330,7 @@ class NCCLWindowAllocator }; mutable std::mutex mMutex; + std::mutex mAllocationMutex; std::unordered_map> mBufferPool; std::unordered_set mRegisteredComms; // Smallest request size that is known to fail collectively for each communicator. @@ -305,11 +345,12 @@ class ScopedNCCLWindowBuffer public: ScopedNCCLWindowBuffer(std::shared_ptr comm, size_t size) : mComm(std::move(comm)) + , mRawComm(nullptr) , mBuffer{} { - if (mComm && *mComm) + if (mComm) { - mBuffer = NCCLWindowAllocator::getInstance().requestBuffer(*mComm, size); + mBuffer = NCCLWindowAllocator::getInstance().requestBuffer(mComm, size, &mRawComm); } } @@ -317,7 +358,7 @@ class ScopedNCCLWindowBuffer { if (mBuffer.isValid()) { - NCCLWindowAllocator::getInstance().releaseBuffer(*mComm, mBuffer.ptr); + NCCLWindowAllocator::getInstance().releaseBuffer(mRawComm, mBuffer.ptr); } } @@ -348,6 +389,7 @@ class ScopedNCCLWindowBuffer private: std::shared_ptr mComm; + ncclComm_t mRawComm; NCCLWindowBuffer mBuffer; }; @@ -357,6 +399,7 @@ class ScopedNCCLWindowBuffer inline std::pair createNCCLWindowTensor( std::shared_ptr comm, at::IntArrayRef shape, torch::ScalarType dtype) { + ncclComm_t rawComm = nullptr; // Calculate buffer size int64_t buffer_size = std::accumulate(shape.begin(), shape.end(), 1LL, std::multiplies()) * torch::elementSize(dtype); @@ -376,7 +419,7 @@ inline std::pair createNCCLWindowTensor( auto& allocator = NCCLWindowAllocator::getInstance(); NCCLWindowBuffer buffer; - if (!comm || !*comm) + if (!comm || (!isNcclFaultToleranceEnabled() && !*comm)) { TLLM_LOG_DEBUG("[createNCCLWindowTensor] null comm; returning invalid buffer"); return std::make_pair(torch::Tensor(), NCCLWindowBuffer()); @@ -384,10 +427,17 @@ inline std::pair createNCCLWindowTensor( try { - buffer = allocator.requestBuffer(*comm, buffer_size); + buffer = allocator.requestBuffer(comm, buffer_size, &rawComm); } catch (std::exception const& e) { + // Local allocation failures may fall back to a regular CUDA tensor. + // Never hide a watchdog error: continuing would reuse an aborted + // communicator on the fallback collective. + if (isNcclFaultToleranceEnabled() && getCommAsyncError(comm).has_value()) + { + throw; + } TLLM_LOG_DEBUG("[createNCCLWindowTensor] requestBuffer failed; returning invalid buffer: %s", e.what()); return std::make_pair(torch::Tensor(), NCCLWindowBuffer()); } @@ -398,9 +448,16 @@ inline std::pair createNCCLWindowTensor( TLLM_LOG_DEBUG("[createNCCLWindowTensor] invalid buffer returned from requestBuffer; returning invalid buffer"); return std::make_pair(torch::Tensor(), NCCLWindowBuffer()); } + TLLM_CHECK_WITH_INFO(rawComm != nullptr, "NCCL window allocator returned a buffer without its communicator"); // Create custom deleter that releases the buffer - auto deleter = [comm, ptr = buffer.ptr](void*) { NCCLWindowAllocator::getInstance().releaseBuffer(*comm, ptr); }; + auto deleter = [comm = std::move(comm), rawComm, ptr = buffer.ptr](void*) + { + // Keep a generic caller-owned communicator alive until its window + // tensor releases the buffer; rawComm remains the stable pool key. + static_cast(comm); + NCCLWindowAllocator::getInstance().releaseBuffer(rawComm, ptr); + }; // Create tensor from the buffer auto tensor = torch::from_blob(buffer.ptr, shape, strides_vec, deleter, torch::dtype(dtype).device(torch::kCUDA)); diff --git a/cpp/tensorrt_llm/common/opUtils.cpp b/cpp/tensorrt_llm/common/opUtils.cpp index 35236bd1c8b4..900c2abffe67 100644 --- a/cpp/tensorrt_llm/common/opUtils.cpp +++ b/cpp/tensorrt_llm/common/opUtils.cpp @@ -20,15 +20,26 @@ #include "tensorrt_llm/runtime/ipcNvlsMemory.h" #include "tensorrt_llm/runtime/utils/mpiTags.h" #include "tensorrt_llm/runtime/utils/mpiUtils.h" +#include "tensorrt_llm/runtime/utils/ncclHostApi.h" +#include "tensorrt_llm/runtime/utils/ncclUniqueIdRendezvous.h" #include "cuda.h" #include #include #include +#include +#include +#include +#include #include +#include +#include #include +#include #include +#include +#include TRTLLM_NAMESPACE_BEGIN #if ENABLE_MULTI_DEVICE @@ -52,16 +63,55 @@ std::unordered_map* getDtypeMap() namespace { -// Get NCCL unique ID for a group of ranks. -ncclUniqueId getUniqueId(std::set const& group) +using Clock = std::chrono::steady_clock; + +std::chrono::milliseconds getPositiveDurationFromEnv(char const* name, int64_t defaultMs) +{ + auto const* value = std::getenv(name); + if (value == nullptr) + { + return std::chrono::milliseconds{defaultMs}; + } + try + { + auto const parsed = std::stoll(value); + if (parsed > 0) + { + return std::chrono::milliseconds{parsed}; + } + } + catch (...) + { + } + TLLM_LOG_WARNING("Ignoring invalid %s=%s; expected a positive integer in milliseconds", name, value); + return std::chrono::milliseconds{defaultMs}; +} + +std::string groupToString(std::set const& group) +{ + std::ostringstream oss; + for (auto it = group.begin(); it != group.end(); ++it) + { + if (it != group.begin()) + { + oss << ','; + } + oss << *it; + } + return oss.str(); +} + +ncclUniqueId getLegacyUniqueId(std::set const& group) { auto const rank = COMM_SESSION.getRank(); - TLLM_LOG_TRACE("%s start for rank %d", __PRETTY_FUNCTION__, rank); - ncclUniqueId id; + ncclUniqueId id{}; if (rank == *group.begin()) { - NCCLCHECK_THROW(ncclGetUniqueId(&id)); - for (auto it = std::next(std::begin(group), 1); it != group.end(); ++it) + { + auto hostApiLock = tensorrt_llm::runtime::acquireNcclHostApiLock(); + NCCLCHECK_THROW(ncclGetUniqueId(&id)); + } + for (auto it = std::next(group.begin()); it != group.end(); ++it) { COMM_SESSION.sendValue(id, *it, tensorrt_llm::mpi::MpiTag::kDefault); } @@ -70,115 +120,1357 @@ ncclUniqueId getUniqueId(std::set const& group) { COMM_SESSION.recvValue(id, *group.begin(), tensorrt_llm::mpi::MpiTag::kDefault); } - TLLM_LOG_TRACE("%s stop for rank %d", __PRETTY_FUNCTION__, rank); return id; } -} // namespace -std::shared_ptr getComm(std::set const& group) +void configureNcclEnvironment() { - auto const rank = COMM_SESSION.getRank(); - TLLM_LOG_TRACE("%s start for rank %d", __PRETTY_FUNCTION__, rank); - static std::map, std::shared_ptr> commMap; - static std::mutex mutex; - std::lock_guard lock(mutex); - std::ostringstream oss; - int index = 0; - for (auto const& rank : group) +#if defined(_WIN32) + // Need static connection initialization for accurate KV cache size estimation + if (getenv("NCCL_RUNTIME_CONNECT") == nullptr) + _putenv_s("NCCL_RUNTIME_CONNECT", "0"); + // Disable graph register to avoid startup hangs + if (getenv("NCCL_GRAPH_REGISTER") == nullptr) + _putenv_s("NCCL_GRAPH_REGISTER", "0"); +#else + setenv("NCCL_RUNTIME_CONNECT", "0", 0); + setenv("NCCL_GRAPH_REGISTER", "0", 0); + // NCCL aborts during init if it tries NVLS multicast but the fabric/IMEX + // plane can't bind it. Disable NVLS when it isn't actually usable so NCCL + // falls back to NVLink P2P. No-overwrite preserves an explicit user setting. + if (!tensorrt_llm::runtime::ipcNvlsSupported()) + { + setenv("NCCL_NVLS_ENABLE", "0", 0); + } +#endif // _WIN32 +} + +class NcclCommState +{ +public: + explicit NcclCommState(std::set group, + std::shared_ptr controlComm = nullptr, bool recovery = false, + std::uint64_t rendezvousId = 0) + : mGroup(std::move(group)) + , mControlComm(std::move(controlComm)) + , mWorldRank(mControlComm != nullptr ? mControlComm->worldRank() : COMM_SESSION.getRank()) + , mFaultToleranceEnabled(tensorrt_llm::isNcclFaultToleranceEnabled()) + , mPollInterval(mFaultToleranceEnabled + ? getPositiveDurationFromEnv("TRTLLM_NCCL_WATCHDOG_POLL_INTERVAL_MS", 100) + : std::chrono::milliseconds{100}) + , mOperationTimeout(mFaultToleranceEnabled + ? getPositiveDurationFromEnv("TRTLLM_NCCL_NONBLOCKING_TIMEOUT_MS", 5000) + : std::chrono::milliseconds{5000}) + , mInitTimeout(mFaultToleranceEnabled + ? (recovery ? getPositiveDurationFromEnv("TRTLLM_NCCL_RECOVERY_TIMEOUT_MS", 5000) + : getPositiveDurationFromEnv("TRTLLM_NCCL_INIT_TIMEOUT_MS", 120000)) + : std::chrono::milliseconds{120000}) { - if (index != 0) + TLLM_CHECK_WITH_INFO(!recovery || mFaultToleranceEnabled, + "NCCL error: recovery communicator construction requires TLLM_FAULT_TOLERANCE_MODE=1"); + if (mFaultToleranceEnabled && mControlComm == nullptr) { - oss << ","; + std::vector const initialRanks(mGroup.begin(), mGroup.end()); + mControlComm = tensorrt_llm::runtime::createNcclUniqueIdRendezvousComm( + initialRanks, mWorldRank, COMM_SESSION, static_cast(tensorrt_llm::mpi::MpiTag::kNcclCommControl)); + } + initialize(rendezvousId); + if (!mFaultToleranceEnabled) + { + return; + } + try + { + mWatchdog = std::thread([this]() { watchdogLoop(); }); + } + catch (...) + { + ncclComm_t abortedComm = nullptr; + if (mComm != nullptr) + { + auto const comm = mComm; + ncclResult_t abortResult; + { + auto const hostApiLock = tensorrt_llm::runtime::acquireNcclHostApiLock(); + abortResult = ncclCommAbort(comm); + if (abortResult == ncclSuccess) + { + common::nccl_util::NcclCommResourceManager::getInstance().beginAbortCleanup(comm); + } + } + if (abortResult == ncclSuccess) + { + mComm = nullptr; + abortedComm = comm; + } + else + { + TLLM_LOG_ERROR("NCCL communicator leaked after watchdog thread creation and abort failed: %s", + ncclGetErrorString(abortResult)); + } + } + if (abortedComm != nullptr) + { + static_cast( + common::nccl_util::NcclCommResourceManager::getInstance().cleanupResources(abortedComm, true)); + } + throw; } - oss << rank; - index++; } - auto groupStr = oss.str(); - auto it = commMap.find(group); - if (it != commMap.end()) + + ~NcclCommState() { - auto ncclComm = it->second; - TLLM_LOG_TRACE("NCCL comm for group(%s) is cached for rank %d", groupStr.c_str(), rank); - return ncclComm; + { + std::lock_guard lock(mMutex); + mStop = true; + } + mWakeup.notify_all(); + if (mWatchdog.joinable()) + { + mWatchdog.join(); + } + + cleanupAbortedResources(); + + std::lock_guard lock(mMutex); + if (mComm != nullptr) + { + auto const comm = mComm; + mComm = nullptr; + if (mError.has_value() || mFaultToleranceEnabled) + { + ncclResult_t abortResult; + { + auto const hostApiLock = tensorrt_llm::runtime::acquireNcclHostApiLock(); + abortResult = ncclCommAbort(comm); + if (abortResult == ncclSuccess) + { + common::nccl_util::NcclCommResourceManager::getInstance().beginAbortCleanup(comm); + } + } + if (abortResult == ncclSuccess) + { + static_cast( + common::nccl_util::NcclCommResourceManager::getInstance().cleanupResources(comm, true)); + quarantinePendingOperationsLocked(); + } + else + { + // Do not release registered buffers while a failed abort + // may have left kernels referencing them. + try + { + TLLM_LOG_WARNING("Leaking failed NCCL communicator for group (%s) after abort error: %s", + groupToString(mGroup).c_str(), ncclGetErrorString(abortResult)); + } + catch (...) + { + } + } + if (abortResult != ncclSuccess) + { + quarantinePendingOperationsLocked(); + } + destroyPooledEventsLocked(); + return; + } + bool const resourcesCleaned + = common::nccl_util::NcclCommResourceManager::getInstance().cleanupResources(comm); + if (!resourcesCleaned) + { + ncclResult_t abortResult; + { + auto const hostApiLock = tensorrt_llm::runtime::acquireNcclHostApiLock(); + abortResult = ncclCommAbort(comm); + } + if (abortResult != ncclSuccess) + { + try + { + TLLM_LOG_WARNING( + "NCCL resource cleanup did not complete and ncclCommAbort failed for group (%s): %s", + groupToString(mGroup).c_str(), ncclGetErrorString(abortResult)); + } + catch (...) + { + } + } + quarantinePendingOperationsLocked(); + destroyPooledEventsLocked(); + return; + } + ncclResult_t result; + { + auto const hostApiLock = tensorrt_llm::runtime::acquireNcclHostApiLock(); + ncclResult_t asyncResult = ncclInProgress; + auto const queryResult = ncclCommGetAsyncError(comm, &asyncResult); + if (queryResult != ncclSuccess || asyncResult != ncclSuccess) + { + // The watcher is already stopped. Re-sample immediately + // before destroy so a late asynchronous error cannot turn + // ncclCommDestroy into an unbounded blocking call. + result = ncclCommAbort(comm); + if (result != ncclSuccess) + { + try + { + TLLM_LOG_WARNING("Final NCCL health check failed and abort failed for group (%s): %s", + groupToString(mGroup).c_str(), ncclGetErrorString(result)); + } + catch (...) + { + } + } + quarantinePendingOperationsLocked(); + destroyPooledEventsLocked(); + return; + } + result = ncclCommDestroy(comm); + } + if (result != ncclSuccess) + { + try + { + TLLM_LOG_WARNING("ncclCommDestroy failed for group (%s): %s", groupToString(mGroup).c_str(), + ncclGetErrorString(result)); + } + catch (...) + { + } + } + if (result == ncclSuccess) + { + destroyCompletedWatchdogEventsLocked(); + } + else + { + quarantinePendingOperationsLocked(); + destroyPooledEventsLocked(); + } + return; + } + destroyPooledEventsLocked(); } - TLLM_LOG_TRACE("Init NCCL comm for group(%s) for rank %d", groupStr.c_str(), rank); - ncclUniqueId id = getUniqueId(group); - int groupRank = 0; - for (auto const& currentRank : group) + NcclCommState(NcclCommState const&) = delete; + NcclCommState& operator=(NcclCommState const&) = delete; + + ncclComm_t requireCommLocked() const { - if (rank == currentRank) - break; - ++groupRank; + if (mError.has_value()) + { + TLLM_THROW("NCCL error: communicator was aborted: %s", mError->c_str()); + } + TLLM_CHECK_WITH_INFO(mComm != nullptr, "NCCL error: communicator was aborted"); + return mComm; } - TLLM_CHECK(static_cast(groupRank) < group.size()); - std::shared_ptr ncclComm(new ncclComm_t, - [](ncclComm_t* comm) + + void checkResultLocked(ncclResult_t result, char const* operation) + { + if (result == ncclSuccess) { - if (!comm) + return; + } + + auto const deadline = Clock::now() + mOperationTimeout; + while (result == ncclInProgress && Clock::now() < deadline) + { + ncclResult_t asyncResult = ncclInProgress; + auto const queryResult = ncclCommGetAsyncError(requireCommLocked(), &asyncResult); + if (queryResult != ncclSuccess && queryResult != ncclInProgress) + { + result = queryResult; + break; + } + if (queryResult == ncclSuccess && asyncResult == ncclSuccess) { return; } + if (queryResult == ncclSuccess && asyncResult != ncclInProgress) + { + result = asyncResult; + break; + } + std::this_thread::sleep_for(std::chrono::milliseconds{1}); + } + + std::string message; + if (result == ncclInProgress) + { + message = std::string(operation) + " timed out while the NCCL communicator was in progress"; + } + else + { + message = std::string(operation) + " failed with NCCL error: " + ncclGetErrorString(result); + } + setErrorAndAbortLocked(message); + TLLM_THROW("NCCL error: communicator was aborted: %s", message.c_str()); + } + + ncclResult_t checkOptionalResultLocked(ncclResult_t result, char const* operation) + { + if (result == ncclSuccess || result == ncclInvalidArgument) + { + return result; + } + if (result != ncclInProgress) + { + checkResultLocked(result, operation); + } - // STEP 1: Clean up resources and destroy NCCL communicator if it's valid - if (*comm) + auto const deadline = Clock::now() + mOperationTimeout; + while (Clock::now() < deadline) + { + ncclResult_t asyncResult = ncclInProgress; + auto const queryResult = ncclCommGetAsyncError(requireCommLocked(), &asyncResult); + if (queryResult != ncclSuccess && queryResult != ncclInProgress) + { + checkResultLocked(queryResult, "ncclCommGetAsyncError(optional NCCL operation)"); + } + if (queryResult == ncclSuccess) { - // Clean up all registered resources FIRST - // The cleanupResources function uses a destruction guard to safely handle - // static destruction order issues - it will return early if the singleton - // is being destroyed (in which case the destructor handles cleanup proactively) - tensorrt_llm::common::nccl_util::NcclCommResourceManager::getInstance().cleanupResources(*comm); + if (asyncResult == ncclSuccess) + { + return asyncResult; + } + if (asyncResult != ncclInProgress) + { + // ncclInvalidArgument is a benign capability fallback only + // when the optional API returns it synchronously. Once the + // call reports ncclInProgress, every eventual non-success + // result is an asynchronous communicator failure. + checkResultLocked(asyncResult, operation); + } + } + std::this_thread::sleep_for(std::chrono::milliseconds{1}); + } + + auto const message = std::string(operation) + " timed out while the optional NCCL operation was in progress"; + setErrorAndAbortLocked(message); + TLLM_THROW("NCCL error: communicator was aborted: %s", message.c_str()); + } - // Now destroy the NCCL communicator - ncclResult_t result = ncclCommDestroy(*comm); - if (result != ncclSuccess) + void checkEnqueueResultLocked(ncclResult_t result, char const* operation) + { + if (result == ncclSuccess || result == ncclInProgress) + { + return; + } + auto const message = std::string(operation) + " failed with NCCL error: " + ncclGetErrorString(result); + setErrorAndAbortLocked(message); + TLLM_THROW("NCCL error: communicator was aborted: %s", message.c_str()); + } + + void waitOperation(std::uint64_t token, char const* operation) + { + TLLM_CHECK_WITH_INFO(token != 0, "NCCL error: watchdog operation token must not be zero in FT mode"); + + while (true) + { + std::string pendingError; + { + std::lock_guard lock(mMutex); + if (mError.has_value()) { - // Logging may fail during static destruction, so wrap in try-catch - try + pendingError = *mError; + } + else + { + TLLM_CHECK_WITH_INFO(token < mNextOperationToken, + "NCCL error: unknown watchdog operation token %llu", static_cast(token)); + auto const pending = std::find_if(mPendingOperations.begin(), mPendingOperations.end(), + [token](PendingOperation const& candidate) { return candidate.token == token; }); + // The watchdog may have observed and retired a fast + // operation between track() releasing its lease and this + // waiter acquiring the state mutex. + if (pending == mPendingOperations.end()) { - TLLM_LOG_WARNING("ncclCommDestroy failed with error: %d", result); + return; } - catch (...) + TLLM_CHECK_WITH_INFO(pending->completion != nullptr, + "NCCL error: watchdog operation token %llu has not been tracked", + static_cast(token)); + + auto const now = Clock::now(); + if (pending->start != nullptr) + { + auto const startResult = cudaEventQuery(pending->start); + if (startResult == cudaSuccess) + { + recycleEventLocked(pending->start); + pending->start = nullptr; + pending->deadline = now + mOperationTimeout; + pending->armed = true; + } + else if (startResult != cudaErrorNotReady) + { + pendingError = pending->name + " start marker failed: " + cudaGetErrorString(startResult); + } + } + + bool completionReady = false; + if (pendingError.empty() && pending->start == nullptr) + { + auto const completionResult = cudaEventQuery(pending->completion); + if (completionResult == cudaSuccess) + { + completionReady = true; + } + else if (completionResult != cudaErrorNotReady) + { + pendingError + = pending->name + " failed with CUDA error: " + cudaGetErrorString(completionResult); + } + else if (pending->armed && now >= pending->deadline) + { + pendingError = pending->name + " timed out before its CUDA completion marker"; + } + } + + bool communicatorReady = false; + if (pendingError.empty()) + { + auto const hostApiLock = tensorrt_llm::runtime::acquireNcclHostApiLock(); + ncclResult_t asyncResult = ncclInProgress; + auto const queryResult = ncclCommGetAsyncError(requireCommLocked(), &asyncResult); + if (queryResult != ncclSuccess && queryResult != ncclInProgress) + { + pendingError = std::string(operation) + + " failed while querying NCCL: " + ncclGetErrorString(queryResult); + } + else if (queryResult == ncclSuccess && asyncResult != ncclSuccess + && asyncResult != ncclInProgress) + { + pendingError = std::string(operation) + + " failed with NCCL error: " + ncclGetErrorString(asyncResult); + } + else + { + communicatorReady = queryResult == ncclSuccess && asyncResult == ncclSuccess; + } + } + + if (pendingError.empty() && completionReady && communicatorReady) { - // Ignore logging failures during static destruction + recycleEventLocked(pending->completion); + mPendingOperations.erase(pending); + return; + } + if (!pendingError.empty()) + { + setErrorAndAbortLocked(pendingError); + pendingError = mError.value_or(pendingError); } } + } - // Clear the communicator value before freeing the pointer - *comm = nullptr; + if (!pendingError.empty()) + { + cleanupAbortedResources(); + TLLM_THROW("NCCL error: communicator was aborted: %s", pendingError.c_str()); } + std::this_thread::sleep_for(std::chrono::milliseconds{1}); + } + } - // STEP 2: Always free the pointer memory (regardless of whether *comm was valid) - delete comm; - }); -#if defined(_WIN32) - // Need static connection initialization for accurate KV cache size estimation - if (getenv("NCCL_RUNTIME_CONNECT") == nullptr) - _putenv_s("NCCL_RUNTIME_CONNECT", "0"); - // Disable graph register to avoid startup hangs - if (getenv("NCCL_GRAPH_REGISTER") == nullptr) - _putenv_s("NCCL_GRAPH_REGISTER", "0"); -#else - setenv("NCCL_RUNTIME_CONNECT", "0", 0); - setenv("NCCL_GRAPH_REGISTER", "0", 0); - // NCCL aborts during init if it tries NVLS multicast but the fabric/IMEX - // plane can't bind it. Disable NVLS when it isn't actually usable so NCCL - // falls back to NVLink P2P. No-overwrite preserves an explicit user setting. - if (!tensorrt_llm::runtime::ipcNvlsSupported()) + uint64_t beginOperationLocked(cudaStream_t stream, char const* operation) { - setenv("NCCL_NVLS_ENABLE", "0", 0); + requireCommLocked(); + if (!mFaultToleranceEnabled) + { + return 0; + } + + cudaStreamCaptureStatus captureStatus = cudaStreamCaptureStatusNone; + auto const captureResult = cudaStreamIsCapturing(stream, &captureStatus); + if (captureResult != cudaSuccess) + { + auto const message = std::string(operation) + + " failed to query CUDA stream capture state: " + cudaGetErrorString(captureResult); + setErrorAndAbortLocked(message); + TLLM_THROW("NCCL error: communicator was aborted: %s", message.c_str()); + } + if (captureStatus != cudaStreamCaptureStatusNone) + { + TLLM_CHECK_WITH_INFO(!mFaultToleranceEnabled, + "NCCL error: CUDA graph capture of managed NCCL operations is disabled while " + "TLLM_FAULT_TOLERANCE_MODE=1; captured graph nodes retain the old communicator across recovery"); + // An event recorded while capturing belongs to the graph and may + // not execute until long after capture ends. Arming a wall-clock + // deadline here would therefore report a false timeout before the + // graph is launched. Non-FT graph launches retain legacy async + // error monitoring; FT mode above requires eager launches so every + // operation receives the completion deadline below. + TLLM_LOG_DEBUG("Skipping eager NCCL completion deadline for %s during CUDA graph capture", operation); + return 0; + } + + if (mPendingOperations.size() >= kMaxPendingOperations) + { + auto const message = std::string(operation) + " exceeded the bounded NCCL watchdog operation queue"; + setErrorAndAbortLocked(message); + TLLM_THROW("NCCL error: communicator was aborted: %s", message.c_str()); + } + + auto const token = mNextOperationToken++; + cudaEvent_t start = acquireEventLocked(); + mPendingOperations.push_back({token, start, nullptr, {}, false, operation}); + auto const result = cudaEventRecord(start, stream); + if (result != cudaSuccess) + { + auto const message + = std::string(operation) + " failed to record its start marker: " + cudaGetErrorString(result); + setErrorAndAbortLocked(message); + TLLM_THROW("NCCL error: communicator was aborted: %s", message.c_str()); + } + + mWakeup.notify_all(); + return token; } -#endif // _WIN32 -#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 29, 0) - ncclConfig_t config = NCCL_CONFIG_INITIALIZER; - config.graphUsageMode = 1; - NCCLCHECK_THROW(ncclCommInitRankConfig(ncclComm.get(), group.size(), id, groupRank, &config)); -#else - NCCLCHECK_THROW(ncclCommInitRank(ncclComm.get(), group.size(), id, groupRank)); -#endif // NCCL_VERSION_CODE >= NCCL_VERSION(2, 29, 0) - commMap[group] = ncclComm; - TLLM_LOG_TRACE("%s stop for rank %d", __PRETTY_FUNCTION__, rank); - return ncclComm; + + void finishOperationLocked(uint64_t token, cudaStream_t stream) + { + requireCommLocked(); + if (token == 0) + { + return; + } + + auto const operation = std::find_if(mPendingOperations.begin(), mPendingOperations.end(), + [token](PendingOperation const& pending) { return pending.token == token; }); + TLLM_CHECK_WITH_INFO(operation != mPendingOperations.end(), "NCCL error: unknown watchdog operation token %llu", + static_cast(token)); + TLLM_CHECK_WITH_INFO(operation->completion == nullptr, + "NCCL error: watchdog operation token %llu was already completed", static_cast(token)); + + operation->completion = acquireEventLocked(); + auto const result = cudaEventRecord(operation->completion, stream); + if (result != cudaSuccess) + { + auto const message + = operation->name + " failed to record its completion marker: " + cudaGetErrorString(result); + setErrorAndAbortLocked(message); + TLLM_THROW("NCCL error: communicator was aborted: %s", message.c_str()); + } + mWakeup.notify_all(); + } + + void abort(std::string const& reason) + { + { + std::lock_guard lock(mMutex); + auto hostApiLock = tensorrt_llm::runtime::acquireNcclHostApiLock(); + setErrorAndAbortLocked(reason); + TLLM_CHECK_WITH_INFO(mComm == nullptr, + "NCCL error: communicator abort failed; refusing to initialize a replacement while the old handle " + "is live"); + } + cleanupAbortedResources(); + } + + std::optional getError() const + { + std::lock_guard lock(mMutex); + return mError; + } + + std::shared_ptr const& controlComm() const noexcept + { + return mControlComm; + } + + int worldRank() const noexcept + { + return mWorldRank; + } + + ncclComm_t* slot() noexcept + { + return &mComm; + } + + std::mutex& mutex() noexcept + { + return mMutex; + } + + void cleanupAbortedResources() noexcept + { + std::vector abortedComms; + { + std::lock_guard lock(mMutex); + abortedComms.swap(mPendingAbortCleanup); + } + for (auto const comm : abortedComms) + { + static_cast(common::nccl_util::NcclCommResourceManager::getInstance().cleanupResources(comm, true)); + } + } + +private: + void initialize(std::uint64_t rendezvousId) + { + configureNcclEnvironment(); + auto const rank = mWorldRank; + TLLM_CHECK_WITH_INFO(!mGroup.empty(), "NCCL communicator group must not be empty"); + auto const rankIt = mGroup.find(rank); + TLLM_CHECK_WITH_INFO(rankIt != mGroup.end(), "Global rank %d is not in NCCL communicator group (%s)", rank, + groupToString(mGroup).c_str()); + + auto const groupRank = static_cast(std::distance(mGroup.begin(), rankIt)); + if (!mFaultToleranceEnabled) + { + auto const id = getLegacyUniqueId(mGroup); + auto hostApiLock = tensorrt_llm::runtime::acquireNcclHostApiLock(); + ncclComm_t comm = nullptr; +#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 29, 0) + ncclConfig_t config = NCCL_CONFIG_INITIALIZER; + config.graphUsageMode = 1; + auto const result = ncclCommInitRankConfig(&comm, static_cast(mGroup.size()), id, groupRank, &config); +#else + auto const result = ncclCommInitRank(&comm, static_cast(mGroup.size()), id, groupRank); +#endif // NCCL_VERSION_CODE >= NCCL_VERSION(2, 29, 0) + if (result != ncclSuccess) + { + if (comm != nullptr) + { + static_cast(ncclCommAbort(comm)); + } + TLLM_THROW("NCCL communicator initialization failed for group (%s) with NCCL error: %s", + groupToString(mGroup).c_str(), ncclGetErrorString(result)); + } + hostApiLock.unlock(); + common::nccl_util::NcclCommResourceManager::getInstance().waitForAbortCleanup(comm); + mComm = comm; + return; + } + + TLLM_CHECK_WITH_INFO(mControlComm != nullptr, "NCCL rendezvous control channel is not initialized"); + TLLM_CHECK_WITH_INFO(rank == mControlComm->worldRank(), + "NCCL error: MPI rank %d does not match rendezvous world rank %d", rank, mControlComm->worldRank()); + auto const deadline = Clock::now() + mInitTimeout; + std::vector const group(mGroup.begin(), mGroup.end()); + auto const id = tensorrt_llm::runtime::exchangeNcclUniqueId(group, *mControlComm, + {tensorrt_llm::mpi::MpiTag::kNcclCommReady, tensorrt_llm::mpi::MpiTag::kNcclCommUniqueId, + tensorrt_llm::mpi::MpiTag::kNcclCommAck}, + rendezvousId, deadline); + ncclConfig_t config = NCCL_CONFIG_INITIALIZER; + config.blocking = 0; +#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 29, 0) + config.graphUsageMode = 1; +#endif // NCCL_VERSION_CODE >= NCCL_VERSION(2, 29, 0) + + auto hostApiLock = tensorrt_llm::runtime::acquireNcclHostApiLock(); + ncclComm_t comm = nullptr; + auto result = ncclCommInitRankConfig(&comm, static_cast(mGroup.size()), id, groupRank, &config); + while (result == ncclInProgress && Clock::now() < deadline) + { + ncclResult_t asyncResult = ncclInProgress; + auto const queryResult = ncclCommGetAsyncError(comm, &asyncResult); + if (queryResult != ncclSuccess && queryResult != ncclInProgress) + { + result = queryResult; + break; + } + if (queryResult == ncclSuccess && asyncResult == ncclSuccess) + { + result = ncclSuccess; + break; + } + if (queryResult == ncclSuccess && asyncResult != ncclInProgress) + { + result = asyncResult; + break; + } + std::this_thread::sleep_for(std::chrono::milliseconds{1}); + } + + if (result != ncclSuccess) + { + if (comm != nullptr) + { + ncclCommAbort(comm); + } + if (result == ncclInProgress) + { + TLLM_THROW( + "NCCL error: communicator initialization timed out for group (%s)", groupToString(mGroup).c_str()); + } + TLLM_THROW("NCCL communicator initialization failed for group (%s) with NCCL error: %s", + groupToString(mGroup).c_str(), ncclGetErrorString(result)); + } + hostApiLock.unlock(); + common::nccl_util::NcclCommResourceManager::getInstance().waitForAbortCleanup(comm); + mComm = comm; + } + + void setErrorAndAbortLocked(std::string const& reason) + { + auto const hostApiLock = tensorrt_llm::runtime::acquireNcclHostApiLock(); + if (!mError.has_value()) + { + mError = reason; + } + if (mComm == nullptr) + { + return; + } + + auto const comm = mComm; + // Abort first: graceful window cleanup synchronizes the entire device + // and can hang forever while a failed NCCL kernel is still resident. + auto const abortResult = ncclCommAbort(comm); + if (abortResult == ncclSuccess) + { + common::nccl_util::NcclCommResourceManager::getInstance().beginAbortCleanup(comm); + mComm = nullptr; + quarantinePendingOperationsLocked(); + // Cleanup is deferred until the operation lease and any allocator + // lock have unwound. Running callbacks here can recursively enter + // NCCLWindowAllocator and deadlock its non-recursive mutex. + mPendingAbortCleanup.push_back(comm); + } + else + { + // Keep both the handle and registered buffers alive. Releasing + // them while NCCL kernels may still reference them is unsafe; a + // later control-path call can retry the idempotent abort. + TLLM_LOG_WARNING("ncclCommAbort failed for group (%s): %s", groupToString(mGroup).c_str(), + ncclGetErrorString(abortResult)); + } + } + + void quarantinePendingOperationsLocked() noexcept + { + if (!mPendingOperations.empty()) + { + // cudaEventDestroy is asynchronous, but the event may still be + // referenced by work from the failed communicator. Keep the + // handles alive for the rest of the process instead of making + // assumptions about CUDA progress after abort. + TLLM_LOG_WARNING( + "Quarantining %zu NCCL completion events after communicator abort", mPendingOperations.size()); + } + mPendingOperations.clear(); + } + + void destroyCompletedWatchdogEventsLocked() noexcept + { + size_t quarantined = 0; + for (auto const& operation : mPendingOperations) + { + for (auto const event : {operation.start, operation.completion}) + { + if (event == nullptr) + { + continue; + } + if (cudaEventQuery(event) == cudaSuccess) + { + static_cast(cudaEventDestroy(event)); + } + else + { + ++quarantined; + } + } + } + mPendingOperations.clear(); + if (quarantined != 0) + { + TLLM_LOG_WARNING("Quarantining %zu incomplete NCCL watchdog events during healthy teardown", quarantined); + } + for (auto const event : mEventPool) + { + static_cast(cudaEventDestroy(event)); + } + mEventPool.clear(); + } + + void destroyPooledEventsLocked() noexcept + { + for (auto const event : mEventPool) + { + static_cast(cudaEventDestroy(event)); + } + mEventPool.clear(); + } + + cudaEvent_t acquireEventLocked() + { + if (!mEventPool.empty()) + { + auto const event = mEventPool.back(); + mEventPool.pop_back(); + return event; + } + + cudaEvent_t event = nullptr; + auto const result = cudaEventCreateWithFlags(&event, cudaEventDisableTiming); + if (result != cudaSuccess) + { + auto const message + = std::string("failed to allocate an NCCL watchdog event: ") + cudaGetErrorString(result); + setErrorAndAbortLocked(message); + TLLM_THROW("NCCL error: communicator was aborted: %s", message.c_str()); + } + return event; + } + + void recycleEventLocked(cudaEvent_t event) noexcept + { + if (event == nullptr) + { + return; + } + if (mEventPool.size() < kMaxPooledEvents) + { + mEventPool.push_back(event); + } + else + { + // The event has completed, so releasing pool overflow cannot race + // device work. Keeping the pool bounded avoids unbounded steady- + // state driver resources under bursty workloads. + static_cast(cudaEventDestroy(event)); + } + } + + void watchdogLoop() noexcept + { + while (true) + { + std::unique_lock lock(mMutex); + if (mWakeup.wait_for(lock, mPollInterval, [this]() { return mStop; })) + { + return; + } + if (mComm == nullptr) + { + if (mError.has_value()) + { + lock.unlock(); + cleanupAbortedResources(); + return; + } + continue; + } + if (mError.has_value()) + { + { + auto const hostApiLock = tensorrt_llm::runtime::acquireNcclHostApiLock(); + setErrorAndAbortLocked(*mError); + } + lock.unlock(); + cleanupAbortedResources(); + continue; + } + + std::string pendingError; + auto const now = Clock::now(); + for (auto operation = mPendingOperations.begin(); operation != mPendingOperations.end();) + { + if (operation->start != nullptr) + { + auto const startResult = cudaEventQuery(operation->start); + if (startResult == cudaSuccess) + { + recycleEventLocked(operation->start); + operation->start = nullptr; + operation->deadline = now + mOperationTimeout; + operation->armed = true; + } + else if (startResult != cudaErrorNotReady) + { + pendingError = operation->name + " start marker failed: " + cudaGetErrorString(startResult); + break; + } + else + { + ++operation; + continue; + } + } + + if (operation->completion == nullptr) + { + if (operation->armed && now >= operation->deadline) + { + pendingError = operation->name + " timed out before its completion marker was recorded"; + break; + } + ++operation; + continue; + } + auto const eventResult = cudaEventQuery(operation->completion); + if (eventResult == cudaSuccess) + { + recycleEventLocked(operation->completion); + operation = mPendingOperations.erase(operation); + continue; + } + if (eventResult != cudaErrorNotReady) + { + pendingError = operation->name + " failed with CUDA error: " + cudaGetErrorString(eventResult); + break; + } + if (operation->armed && now >= operation->deadline) + { + pendingError = operation->name + " timed out before its CUDA completion marker"; + break; + } + ++operation; + } + + { + auto const hostApiLock = tensorrt_llm::runtime::acquireNcclHostApiLock(); + if (!pendingError.empty()) + { + setErrorAndAbortLocked(pendingError); + } + else + { + ncclResult_t asyncResult = ncclSuccess; + auto const queryResult = ncclCommGetAsyncError(mComm, &asyncResult); + if (queryResult != ncclSuccess && queryResult != ncclInProgress) + { + setErrorAndAbortLocked(std::string("ncclCommGetAsyncError failed with NCCL error: ") + + ncclGetErrorString(queryResult)); + } + else if (asyncResult != ncclSuccess && asyncResult != ncclInProgress) + { + setErrorAndAbortLocked(std::string("NCCL watchdog observed asynchronous NCCL error: ") + + ncclGetErrorString(asyncResult)); + } + } + } + lock.unlock(); + cleanupAbortedResources(); + } + } + + std::set mGroup; + std::shared_ptr mControlComm; + int const mWorldRank; + bool mFaultToleranceEnabled; + ncclComm_t mComm{nullptr}; + mutable std::mutex mMutex; + std::condition_variable mWakeup; + bool mStop{false}; + std::optional mError; + std::vector mPendingAbortCleanup; + + struct PendingOperation + { + uint64_t token; + cudaEvent_t start; + cudaEvent_t completion; + Clock::time_point deadline; + bool armed; + std::string name; + }; + + std::vector mPendingOperations; + std::vector mEventPool; + uint64_t mNextOperationToken{1}; + static constexpr size_t kMaxPendingOperations = 4096; + static constexpr size_t kMaxPooledEvents = 256; + std::chrono::milliseconds mPollInterval; + std::chrono::milliseconds mOperationTimeout; + std::chrono::milliseconds mInitTimeout; + std::thread mWatchdog; +}; + +class NcclCommRegistry +{ +public: + static NcclCommRegistry& instance() + { + // Raw communicators, watchdogs, and their CUDA/MPI dependencies are + // process-lifetime resources. Avoid cross-translation-unit static + // destruction where the resource manager or allocator could disappear + // before a live watchdog/state attempts final cleanup. + static auto* registry = new NcclCommRegistry; + return *registry; + } + + std::shared_ptr get(std::set const& group) + { + std::lock_guard lock(mMutex); + auto it = mComms.find(group); + if (it == mComms.end()) + { + auto state = std::make_shared(group); + mSlots[state->slot()] = state; + it = mComms.emplace(group, std::move(state)).first; + } + return std::shared_ptr{it->second, it->second->slot()}; + } + + std::shared_ptr find(std::shared_ptr const& comm) + { + TLLM_CHECK_WITH_INFO(comm != nullptr, "NCCL communicator holder is null"); + std::lock_guard lock(mMutex); + auto const it = mSlots.find(comm.get()); + TLLM_CHECK_WITH_INFO(it != mSlots.end(), "NCCL communicator is not managed by the TRT-LLM registry"); + auto state = it->second.lock(); + TLLM_CHECK_WITH_INFO(state != nullptr, "NCCL communicator registry entry has expired"); + return state; + } + + std::shared_ptr find(std::set const& group) + { + std::lock_guard lock(mMutex); + auto const it = mComms.find(group); + TLLM_CHECK_WITH_INFO( + it != mComms.end(), "No cached NCCL communicator exists for group (%s)", groupToString(group).c_str()); + return it->second; + } + + void abortAndReinit(std::set const& oldGroup, std::set const& activeGroup, std::uint64_t rendezvousId) + { + std::lock_guard const recoveryLock(mRecoveryMutex); + TLLM_CHECK_WITH_INFO(rendezvousId != 0, + "NCCL error: recovery rendezvous ID must be nonzero; zero is reserved for initial bootstrap"); + TLLM_CHECK_WITH_INFO(!activeGroup.empty(), "active NCCL rank group must not be empty"); + TLLM_CHECK_WITH_INFO(std::includes(oldGroup.begin(), oldGroup.end(), activeGroup.begin(), activeGroup.end()), + "active NCCL rank group (%s) must be a subset of old group (%s)", groupToString(activeGroup).c_str(), + groupToString(oldGroup).c_str()); + TLLM_CHECK_WITH_INFO(tensorrt_llm::isNcclFaultToleranceEnabled(), + "NCCL error: communicator reinitialization requires TLLM_FAULT_TOLERANCE_MODE=1; otherwise captured " + "CUDA graphs may retain the old communicator"); + std::shared_ptr oldState; + std::shared_ptr existingTargetState; + { + std::lock_guard lock(mMutex); + auto const old = mComms.find(oldGroup); + if (old != mComms.end()) + { + oldState = old->second; + } + auto const existingTarget = mComms.find(activeGroup); + if (existingTarget != mComms.end()) + { + existingTargetState = existingTarget->second; + } + } + TLLM_CHECK_WITH_INFO( + oldState != nullptr, "No cached NCCL communicator exists for group (%s)", groupToString(oldGroup).c_str()); + auto const& controlComm = oldState->controlComm(); + TLLM_CHECK_WITH_INFO(controlComm != nullptr, "NCCL rendezvous control channel is not initialized"); + auto const worldRank = controlComm->worldRank(); + TLLM_CHECK_WITH_INFO( + activeGroup.count(worldRank) != 0, "Failed rank %d must not call NCCL abort_and_reinit", worldRank); + oldState->abort("NCCL error: communicator was aborted for survivor reinitialization"); + if (existingTargetState != nullptr && existingTargetState != oldState) + { + // Explicit recovery is a distributed generation change. Never let + // one rank reuse a locally healthy cached target communicator + // while another rank enters the fresh-ID rendezvous. + existingTargetState->abort("NCCL error: cached target communicator was retired for reinitialization"); + } + + // Do not mutate oldState: existing operations retain it and must fail + // immediately instead of silently switching rank numbering underneath + // an in-flight call. Subsequent operations request activeGroup. + std::lock_guard lock(mMutex); + auto state = std::make_shared(activeGroup, controlComm, true, rendezvousId); + mSlots[state->slot()] = state; + mComms[activeGroup] = std::move(state); + } + + std::optional getError(std::set const& group) + { + std::shared_ptr state; + { + std::lock_guard lock(mMutex); + auto const it = mComms.find(group); + if (it == mComms.end()) + { + return std::nullopt; + } + state = it->second; + } + return state->getError(); + } + +private: + std::mutex mRecoveryMutex; + std::mutex mMutex; + std::map, std::shared_ptr> mComms; + std::unordered_map> mSlots; +}; + +} // namespace + +struct NcclCommLease::Impl +{ + explicit Impl(std::shared_ptr state_) + : state(std::move(state_)) + , lock(state->mutex()) + { + // Complete the call_once-protected capability probe before taking the + // host gate. The probe itself briefly takes that gate for + // ncclGetVersion; reversing the order can deadlock with a concurrent + // first-time capability query. + static_cast(common::nccl_util::isNcclWindowSupported()); + hostApiLock = tensorrt_llm::runtime::acquireNcclHostApiLock(); + state->requireCommLocked(); + } + + ~Impl() + { + if (groupActive) + { + auto const result = ncclGroupEnd(); + groupActive = false; + try + { + state->checkResultLocked(result, "ncclGroupEnd(exception cleanup)"); + } + catch (...) + { + // Destructors cannot propagate. checkResultLocked has already + // latched the classifier-friendly error and aborted the comm. + } + } + if (hostApiLock.owns_lock()) + { + hostApiLock.unlock(); + } + if (lock.owns_lock()) + { + lock.unlock(); + } + state->cleanupAbortedResources(); + } + + std::shared_ptr state; + std::unique_lock lock; + tensorrt_llm::runtime::NcclHostApiLock hostApiLock; + bool groupActive{false}; +}; + +NcclCommLease::NcclCommLease(std::unique_ptr impl) + : mImpl(std::move(impl)) +{ +} + +NcclCommLease::NcclCommLease(ncclComm_t legacyComm) + : mLegacyComm(legacyComm) +{ +} + +NcclCommLease::~NcclCommLease() +{ + if (mLegacyGroupActive) + { + auto const result = ncclGroupEnd(); + mLegacyGroupActive = false; + if (result != ncclSuccess) + { + TLLM_LOG_WARNING( + "Failed to close legacy NCCL group during exception cleanup: %s", ncclGetErrorString(result)); + } + } +} + +NcclCommLease::NcclCommLease(NcclCommLease&& other) noexcept + : mImpl(std::move(other.mImpl)) + , mLegacyComm(std::exchange(other.mLegacyComm, nullptr)) + , mLegacyGroupActive(std::exchange(other.mLegacyGroupActive, false)) +{ +} + +NcclCommLease& NcclCommLease::operator=(NcclCommLease&& other) noexcept +{ + if (this != &other) + { + if (mLegacyGroupActive) + { + static_cast(ncclGroupEnd()); + } + mImpl = std::move(other.mImpl); + mLegacyComm = std::exchange(other.mLegacyComm, nullptr); + mLegacyGroupActive = std::exchange(other.mLegacyGroupActive, false); + } + return *this; +} + +ncclComm_t NcclCommLease::get() const +{ + if (mImpl == nullptr) + { + TLLM_CHECK_WITH_INFO(mLegacyComm != nullptr, "NCCL communicator lease was moved from"); + return mLegacyComm; + } + return mImpl->state->requireCommLocked(); +} + +void NcclCommLease::check(ncclResult_t result, char const* operation) const +{ + if (mImpl == nullptr) + { + (void) operation; + TLLM_CHECK_WITH_INFO(mLegacyComm != nullptr, "NCCL communicator lease was moved from"); + NCCLCHECK_THROW(result); + return; + } + mImpl->state->checkResultLocked(result, operation); +} + +ncclResult_t NcclCommLease::checkOptional(ncclResult_t result, char const* operation) const +{ + if (mImpl == nullptr) + { + (void) operation; + TLLM_CHECK_WITH_INFO(mLegacyComm != nullptr, "NCCL communicator lease was moved from"); + if (result != ncclInvalidArgument) + { + NCCLCHECK_THROW(result); + } + return result; + } + return mImpl->state->checkOptionalResultLocked(result, operation); +} + +void NcclCommLease::checkEnqueue(ncclResult_t result, char const* operation) const +{ + if (mImpl == nullptr) + { + check(result, operation); + return; + } + if (result != ncclSuccess && result != ncclInProgress && mImpl->groupActive) + { + // Close the thread-local group before aborting. Otherwise the first + // collective on a replacement communicator can be nested inside the + // failed group. + static_cast(ncclGroupEnd()); + mImpl->groupActive = false; + } + mImpl->state->checkEnqueueResultLocked(result, operation); +} + +uint64_t NcclCommLease::begin(cudaStream_t stream, char const* operation) const +{ + if (mImpl == nullptr) + { + (void) stream; + (void) operation; + TLLM_CHECK_WITH_INFO(mLegacyComm != nullptr, "NCCL communicator lease was moved from"); + return 0; + } + return mImpl->state->beginOperationLocked(stream, operation); +} + +void NcclCommLease::track(uint64_t token, cudaStream_t stream) const +{ + if (mImpl == nullptr) + { + (void) token; + (void) stream; + TLLM_CHECK_WITH_INFO(mLegacyComm != nullptr, "NCCL communicator lease was moved from"); + return; + } + mImpl->state->finishOperationLocked(token, stream); +} + +void NcclCommLease::groupStart(char const* operation) const +{ + if (mImpl == nullptr) + { + TLLM_CHECK_WITH_INFO(mLegacyComm != nullptr, "NCCL communicator lease was moved from"); + TLLM_CHECK_WITH_INFO(!mLegacyGroupActive, "NCCL error: nested NCCL groups are not supported by this lease"); + (void) operation; + NCCLCHECK_THROW(ncclGroupStart()); + mLegacyGroupActive = true; + return; + } + TLLM_CHECK_WITH_INFO(!mImpl->groupActive, "NCCL error: nested NCCL groups are not supported by this lease"); + mImpl->state->checkResultLocked(ncclGroupStart(), operation); + mImpl->groupActive = true; +} + +void NcclCommLease::groupEnd(char const* operation) const +{ + if (mImpl == nullptr) + { + TLLM_CHECK_WITH_INFO(mLegacyComm != nullptr, "NCCL communicator lease was moved from"); + TLLM_CHECK_WITH_INFO(mLegacyGroupActive, "NCCL error: no NCCL group is active on this lease"); + auto const result = ncclGroupEnd(); + mLegacyGroupActive = false; + (void) operation; + NCCLCHECK_THROW(result); + return; + } + TLLM_CHECK_WITH_INFO(mImpl->groupActive, "NCCL error: no NCCL group is active on this lease"); + auto const result = ncclGroupEnd(); + mImpl->groupActive = false; + mImpl->state->checkResultLocked(result, operation); +} + +std::shared_ptr getComm(std::set const& group) +{ + auto const groupString = groupToString(group); + // Do not query COMM_SESSION here: after survivor reinitialization that + // communicator may include a dead process and retain its fatal handler. + TLLM_LOG_TRACE("Get NCCL comm for group(%s)", groupString.c_str()); + return NcclCommRegistry::instance().get(group); +} + +int getCommWorldRank(std::shared_ptr const& comm) +{ + return NcclCommRegistry::instance().find(comm)->worldRank(); +} + +bool isNcclFaultToleranceEnabled() +{ + // Communicator mode cannot change after the first managed communicator is + // used: its blocking configuration and watchdog ownership are fixed. + static bool const enabled = tensorrt_llm::mpi::isFaultToleranceModeEnabled(); + return enabled; +} + +NcclCommLease acquireComm(std::shared_ptr const& comm) +{ + // FT mode is a startup-only process setting. Keep the default-off path + // equivalent to the legacy blocking communicator: no registry lookup, + // heap allocation, communicator mutex, or process-wide NCCL gate per op. + if (!isNcclFaultToleranceEnabled()) + { + TLLM_CHECK_WITH_INFO(comm != nullptr && *comm != nullptr, "NCCL communicator is null"); + return NcclCommLease{*comm}; + } + auto state = NcclCommRegistry::instance().find(comm); + return NcclCommLease{std::make_unique(std::move(state))}; +} + +void waitCommOperation(std::shared_ptr const& comm, std::uint64_t token, char const* operation) +{ + if (!isNcclFaultToleranceEnabled()) + { + TLLM_CHECK_WITH_INFO(token == 0, "NCCL error: legacy communicator received a watchdog operation token"); + return; + } + auto state = NcclCommRegistry::instance().find(comm); + state->waitOperation(token, operation); +} + +void abortAndReinitComm(std::set const& oldGroup, std::set const& activeGroup, std::uint64_t rendezvousId) +{ + NcclCommRegistry::instance().abortAndReinit(oldGroup, activeGroup, rendezvousId); +} + +std::optional getCommAsyncError(std::set const& group) +{ + return NcclCommRegistry::instance().getError(group); +} + +std::optional getCommAsyncError(std::shared_ptr const& comm) +{ + auto state = NcclCommRegistry::instance().find(comm); + return state->getError(); } #endif // ENABLE_MULTI_DEVICE diff --git a/cpp/tensorrt_llm/common/opUtils.h b/cpp/tensorrt_llm/common/opUtils.h index 72e5a5ea3e09..306c9cffd2e1 100644 --- a/cpp/tensorrt_llm/common/opUtils.h +++ b/cpp/tensorrt_llm/common/opUtils.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -29,6 +29,8 @@ #include #endif // ENABLE_MULTI_DEVICE +#include +#include #include #include #include @@ -37,6 +39,7 @@ #include #include #include +#include #include "tensorrt_llm/common/nvmlWrapper.h" @@ -216,8 +219,100 @@ inline bool isBuilding() std::unordered_map* getDtypeMap(); +/// Access to a cached NCCL communicator. +/// +/// NCCL communicators used by the raw TRT-LLM collectives are configured as +/// nonblocking so that the fault-tolerance watchdog can abort them. NCCL does +/// not allow host API calls and ncclCommAbort to race on the same communicator, +/// so every FT raw collective must hold this serialized lease from its first +/// communicator call through ncclGroupEnd (when a group is used). When FT is +/// disabled the lease is an inline, non-locking wrapper around the legacy +/// blocking communicator. +class NcclCommLease +{ +public: + ~NcclCommLease(); + + NcclCommLease(NcclCommLease const&) = delete; + NcclCommLease& operator=(NcclCommLease const&) = delete; + NcclCommLease(NcclCommLease&&) noexcept; + NcclCommLease& operator=(NcclCommLease&&) noexcept; + + /// Return the native communicator, or throw a classifier-friendly error + /// when its watchdog has aborted it. + [[nodiscard]] ncclComm_t get() const; + + /// Complete a nonblocking NCCL host call. ncclInProgress is polled via + /// ncclCommGetAsyncError before another host API is issued. + void check(ncclResult_t result, char const* operation) const; + + /// Complete an optional NCCL capability call. An immediate + /// ncclInvalidArgument is returned as a documented, nonfatal unsupported + /// result. Once a call returns ncclInProgress, every eventual non-success + /// result is latched and aborted as an asynchronous communicator failure. + [[nodiscard]] ncclResult_t checkOptional(ncclResult_t result, char const* operation) const; + + /// Validate an operation queued inside ncclGroupStart/ncclGroupEnd. + /// ncclInProgress is completed by the enclosing ncclGroupEnd call. + void checkEnqueue(ncclResult_t result, char const* operation) const; + + /// Record a stream marker immediately before an asynchronous NCCL + /// operation. The timeout is armed only when this marker completes, so + /// unrelated work already queued on the stream cannot cause a false + /// timeout. + [[nodiscard]] uint64_t begin(cudaStream_t stream, char const* operation) const; + + /// Record the matching completion marker after the NCCL operation has + /// been enqueued. The watchdog aborts if it does not complete by the + /// deadline armed by begin(). A zero token denotes the default-off legacy + /// path. FT mode rejects capture so a recovered communicator can never be + /// bypassed by stale graph nodes. + void track(uint64_t token, cudaStream_t stream) const; + + /// Start/end an NCCL group tracked by this lease. If any intervening C++ + /// operation throws, the lease destructor closes the group before + /// releasing the shared NCCL host-API gate. + void groupStart(char const* operation) const; + void groupEnd(char const* operation) const; + +private: + struct Impl; + explicit NcclCommLease(std::unique_ptr impl); + explicit NcclCommLease(ncclComm_t legacyComm); + std::unique_ptr mImpl; + ncclComm_t mLegacyComm{nullptr}; + mutable bool mLegacyGroupActive{false}; + + friend NcclCommLease acquireComm(std::shared_ptr const& comm); +}; + std::shared_ptr getComm(std::set const& group); +/// Return the immutable original/global MPI rank associated with a managed communicator. +int getCommWorldRank(std::shared_ptr const& comm); + +/// Return the startup-cached raw-NCCL fault-tolerance mode. +bool isNcclFaultToleranceEnabled(); + +/// Acquire exclusive host-API access to a cached NCCL communicator. +NcclCommLease acquireComm(std::shared_ptr const& comm); + +/// Abort the communicator for oldGroup and initialize a fresh communicator +/// containing only activeGroup. Both groups contain global MPI ranks. Only +/// ranks in activeGroup may call this function. +void abortAndReinitComm(std::set const& oldGroup, std::set const& activeGroup, std::uint64_t rendezvousId); + +/// Wait for a begin()/track() operation without retaining its lease. The +/// operation timeout is armed only after its start marker reaches the head of +/// the stream, and other communicators remain free to enqueue or abort. +void waitCommOperation(std::shared_ptr const& comm, std::uint64_t token, char const* operation); + +/// Return the stable watchdog error for a cached communicator, if any. +std::optional getCommAsyncError(std::set const& group); + +/// Return the stable watchdog error for the communicator referenced by comm. +std::optional getCommAsyncError(std::shared_ptr const& comm); + #endif // ENABLE_MULTI_DEVICE //! To save GPU memory, all the plugins share the same cublas and cublasLt handle globally. diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu b/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu index 0c4431b8eaef..0f2d453c363a 100644 --- a/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu +++ b/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu @@ -741,7 +741,7 @@ __device__ void vectorized_combine_impl(T* dst_typed_base, int size_per_token, i { int target_rank = ptrs.topk_target_ranks[local_token_idx * TOP_K + k]; int dst_idx = ptrs.topk_send_indices[local_token_idx * TOP_K + k]; - if (dst_idx < 0) + if (dst_idx < 0 || !is_rank_active(ptrs.active_rank_mask, target_rank)) { acc[k].fill(0.0f); continue; @@ -766,8 +766,12 @@ __device__ void vectorized_combine_impl(T* dst_typed_base, int size_per_token, i #pragma unroll for (int k = 0; k < TOP_K; ++k) { - if (ptrs.topk_send_indices[local_token_idx * TOP_K + k] < 0) + int target_rank = ptrs.topk_target_ranks[local_token_idx * TOP_K + k]; + int dst_idx = ptrs.topk_send_indices[local_token_idx * TOP_K + k]; + if (dst_idx < 0 || !is_rank_active(ptrs.active_rank_mask, target_rank)) + { continue; // acc[k] already holds 0.0f from fill() above + } #pragma unroll for (int j = elems_per_vec - 1; j >= 0; --j) acc[k][j] = static_cast(reinterpret_cast(&acc[k])[j]); diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h b/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h index 9a6f3904c501..138ca92e71a8 100644 --- a/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h +++ b/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h @@ -93,8 +93,8 @@ struct CombineKernelPointers int const* topk_send_indices; // dst index per k, -1 for duplicates // Active-rank bitmask: see DispatchKernelPointers::active_rank_mask. Combine skips flag - // writes/waits to/from masked peers; per-token accumulation uses topk_send_indices[k] < 0 - // (set by dispatch) to skip dead-targeted slots, so no explicit mask check is needed there. + // writes/waits to/from masked peers and also skips per-token accumulation for ranks that + // become inactive between dispatch and combine. uint64_t active_rank_mask[kRankMaskWords]; }; diff --git a/cpp/tensorrt_llm/runtime/CMakeLists.txt b/cpp/tensorrt_llm/runtime/CMakeLists.txt index ca81fbb0f6cd..9c0c0f388b05 100644 --- a/cpp/tensorrt_llm/runtime/CMakeLists.txt +++ b/cpp/tensorrt_llm/runtime/CMakeLists.txt @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & # AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); you may not @@ -16,6 +16,7 @@ include(FetchContent) set(SRCS utils/mpiUtils.cpp + utils/ncclUniqueIdRendezvous.cpp utils/numpyUtils.cpp utils/runtimeUtils.cpp utils/debugUtils.cu diff --git a/cpp/tensorrt_llm/runtime/ncclCommunicator.cpp b/cpp/tensorrt_llm/runtime/ncclCommunicator.cpp index dc79fd5f48e2..51014d56a0a5 100644 --- a/cpp/tensorrt_llm/runtime/ncclCommunicator.cpp +++ b/cpp/tensorrt_llm/runtime/ncclCommunicator.cpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -18,6 +18,18 @@ #include "tensorrt_llm/common/logger.h" #include "tensorrt_llm/runtime/utils/multiDeviceUtils.h" +#include "tensorrt_llm/runtime/utils/ncclHostApi.h" +#include "tensorrt_llm/runtime/utils/ncclUniqueIdRendezvous.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include #if ENABLE_MULTI_DEVICE #include @@ -27,6 +39,15 @@ using namespace tensorrt_llm::runtime; namespace { +using namespace std::chrono_literals; + +constexpr auto kInitialReadyTimeout = 120s; +#if ENABLE_MULTI_DEVICE +constexpr auto kReadyPollInterval = 1ms; +constexpr auto kRecoveryReadyTimeout = 5s; +constexpr auto kWatcherPollInterval = 100ms; +#endif + #if ENABLE_MULTI_DEVICE ncclDataType_t toNcclType(nvinfer1::DataType dataType) @@ -43,19 +64,154 @@ ncclDataType_t toNcclType(nvinfer1::DataType dataType) #if ENABLE_BF16 case nvinfer1::DataType::kBF16: return ncclBfloat16; #endif // ENABLE_BF16 - default: TLLM_THROW("Unsupported data type: %d", static_cast(dataType)); + default: TLLM_THROW("NCCL error: unsupported data type: %d", static_cast(dataType)); } } + +std::string ncclErrorMessage(char const* operation, ncclResult_t result) +{ + std::ostringstream message; + message << "NCCL error during " << operation << ": " << ncclGetErrorString(result); + return message.str(); +} + +std::chrono::milliseconds getOperationTimeout() +{ + static auto const timeout = [] + { + constexpr int64_t defaultTimeoutMs = 5000; + auto const* value = std::getenv("TRTLLM_NCCL_NONBLOCKING_TIMEOUT_MS"); + if (value != nullptr) + { + try + { + auto const parsed = std::stoll(value); + if (parsed > 0) + { + return std::chrono::milliseconds{parsed}; + } + } + catch (...) + { + } + TLLM_LOG_WARNING( + "Ignoring invalid TRTLLM_NCCL_NONBLOCKING_TIMEOUT_MS=%s; expected positive milliseconds", value); + } + return std::chrono::milliseconds{defaultTimeoutMs}; + }(); + return timeout; +} + #endif // ENABLE_MULTI_DEVICE } // namespace +NcclCommunicator::NcclCommunicator(ncclComm_t comm) + : mFaultToleranceEnabled{mpi::isFaultToleranceModeEnabled()} + , mComm{comm} +{ +#if ENABLE_MULTI_DEVICE + TLLM_CHECK_WITH_INFO(mComm != nullptr, "NCCL error: cannot wrap a null communicator"); + + auto const hostApiLock = acquireNcclHostApiLock(); + try + { + if (mFaultToleranceEnabled) + { + waitUntilReady(mComm, "wrapping communicator", std::chrono::steady_clock::now() + kInitialReadyTimeout); + } + TLLM_NCCL_CHECK(ncclCommCount(mComm, &mInitialWorldSize)); + TLLM_NCCL_CHECK(ncclCommUserRank(mComm, &mWorldRank)); + mActiveRanks.resize(mInitialWorldSize); + std::iota(mActiveRanks.begin(), mActiveRanks.end(), 0); + } + catch (...) + { + ncclCommAbort(mComm); + mComm = nullptr; + throw; + } + startWatcher(); +#else + (void) comm; +#endif // ENABLE_MULTI_DEVICE +} + +NcclCommunicator::NcclCommunicator(int worldSize, int rank, mpi::MpiComm const& mpiComm) + : mInitialWorldSize{worldSize} + , mWorldRank{rank} + , mFaultToleranceEnabled{mpi::isFaultToleranceModeEnabled()} +{ + TLLM_CHECK_WITH_INFO(worldSize > 0, "NCCL error: communicator world size must be positive, got %d", worldSize); + TLLM_CHECK_WITH_INFO( + rank >= 0 && rank < worldSize, "NCCL error: communicator rank must be in [0, %d), got %d", worldSize, rank); + + mActiveRanks.resize(worldSize); + std::iota(mActiveRanks.begin(), mActiveRanks.end(), 0); +#if ENABLE_MULTI_DEVICE + TLLM_CHECK_WITH_INFO(mpiComm.getRank() == rank, "NCCL error: world rank %d does not match MPI bootstrap rank %d", + rank, mpiComm.getRank()); + TLLM_CHECK_WITH_INFO(mpiComm.getSize() >= worldSize, "NCCL error: world size %d exceeds MPI bootstrap size %d", + worldSize, mpiComm.getSize()); + if (mFaultToleranceEnabled) + { + mControlComm = createNcclUniqueIdRendezvousComm( + mActiveRanks, rank, mpiComm, static_cast(mpi::MpiTag::kNcclPpControl)); + mComm = createComm(mActiveRanks, rank, *mControlComm, /*rendezvousId=*/0, kInitialReadyTimeout); + } + else + { + // Preserve the legacy blocking MPI-broadcast bootstrap when FT is + // disabled. It creates no control communicator or watcher thread. + mComm = createLegacyComm(worldSize, rank, mpiComm); + } +#else + mComm = nullptr; +#endif // ENABLE_MULTI_DEVICE + startWatcher(); +} + void NcclCommunicator::send( void const* sendbuff, size_t count, nvinfer1::DataType dataType, int peer, CudaStream const& stream) const { #if ENABLE_MULTI_DEVICE - TLLM_NCCL_CHECK(ncclSend(sendbuff, count, toNcclType(dataType), peer, mComm, stream.get())); + if (!mFaultToleranceEnabled) + { + TLLM_NCCL_CHECK(ncclSend(sendbuff, count, toNcclType(dataType), peer, mComm, stream.get())); + return; + } + std::lock_guard lock(mCommMutex); + auto const hostApiLock = acquireNcclHostApiLock(); + checkUsableLocked(); + int const commPeer = getCommPeerRankLocked(peer); + auto const watchdogToken = beginOperationLocked(stream.get(), "ncclSend"); + ncclResult_t const result = ncclSend(sendbuff, count, toNcclType(dataType), commPeer, mComm, stream.get()); + if (result == ncclInProgress) + { + try + { + waitUntilReady(mComm, "ncclSend enqueue", std::chrono::steady_clock::now() + kRecoveryReadyTimeout); + } + catch (std::exception const& error) + { + static_cast(abortLocked( + std::string{"NCCL error: communicator was aborted after ncclSend failed: "} + error.what())); + TLLM_THROW("%s", mAsyncError.c_str()); + } + catch (...) + { + static_cast(abortLocked("NCCL error: communicator was aborted after ncclSend failed")); + TLLM_THROW("%s", mAsyncError.c_str()); + } + } + else if (result != ncclSuccess) + { + static_cast( + abortLocked("NCCL error: communicator was aborted after " + ncclErrorMessage("ncclSend", result))); + TLLM_THROW("%s", mAsyncError.c_str()); + } + finishOperationLocked(watchdogToken, stream.get()); #else - TLLM_THROW("Multi device support is disabled."); + TLLM_THROW("NCCL error: multi device support is disabled."); #endif // ENABLE_MULTI_DEVICE } @@ -63,23 +219,126 @@ void NcclCommunicator::receive( void* sendbuff, size_t count, nvinfer1::DataType dataType, int peer, CudaStream const& stream) const { #if ENABLE_MULTI_DEVICE - TLLM_NCCL_CHECK(ncclRecv(sendbuff, count, toNcclType(dataType), peer, mComm, stream.get())); + if (!mFaultToleranceEnabled) + { + TLLM_NCCL_CHECK(ncclRecv(sendbuff, count, toNcclType(dataType), peer, mComm, stream.get())); + return; + } + std::lock_guard lock(mCommMutex); + auto const hostApiLock = acquireNcclHostApiLock(); + checkUsableLocked(); + int const commPeer = getCommPeerRankLocked(peer); + auto const watchdogToken = beginOperationLocked(stream.get(), "ncclRecv"); + ncclResult_t const result = ncclRecv(sendbuff, count, toNcclType(dataType), commPeer, mComm, stream.get()); + if (result == ncclInProgress) + { + try + { + waitUntilReady(mComm, "ncclRecv enqueue", std::chrono::steady_clock::now() + kRecoveryReadyTimeout); + } + catch (std::exception const& error) + { + static_cast(abortLocked( + std::string{"NCCL error: communicator was aborted after ncclRecv failed: "} + error.what())); + TLLM_THROW("%s", mAsyncError.c_str()); + } + catch (...) + { + static_cast(abortLocked("NCCL error: communicator was aborted after ncclRecv failed")); + TLLM_THROW("%s", mAsyncError.c_str()); + } + } + else if (result != ncclSuccess) + { + static_cast( + abortLocked("NCCL error: communicator was aborted after " + ncclErrorMessage("ncclRecv", result))); + TLLM_THROW("%s", mAsyncError.c_str()); + } + finishOperationLocked(watchdogToken, stream.get()); #else - TLLM_THROW("Multi device support is disabled."); + TLLM_THROW("NCCL error: multi device support is disabled."); #endif // ENABLE_MULTI_DEVICE } -ncclComm_t NcclCommunicator::createComm(int worldSize, int rank, mpi::MpiComm const& mpiComm) +ncclComm_t NcclCommunicator::createComm(std::vector const& activeRanks, int worldRank, + NcclUniqueIdRendezvousComm const& controlComm, std::uint64_t rendezvousId, std::chrono::milliseconds readyTimeout) { #if ENABLE_MULTI_DEVICE + TLLM_CHECK_WITH_INFO(!activeRanks.empty(), "NCCL error: communicator active-rank set must not be empty"); + auto const rankIt = std::find(activeRanks.begin(), activeRanks.end(), worldRank); + TLLM_CHECK_WITH_INFO( + rankIt != activeRanks.end(), "NCCL error: world rank %d is not present in the active-rank set", worldRank); + // Do not use an MPI collective here. During recovery, failed ranks remain + // members of the pre-failure control communicator. Only the survivors + // exchange the fresh NCCL ID point-to-point. + auto const deadline = std::chrono::steady_clock::now() + readyTimeout; + TLLM_CHECK_WITH_INFO(controlComm.worldRank() == worldRank, + "NCCL error: world rank %d does not match rendezvous control world rank %d", worldRank, + controlComm.worldRank()); + ncclUniqueId const id = exchangeNcclUniqueId(activeRanks, controlComm, + {mpi::MpiTag::kNcclPpReady, mpi::MpiTag::kNcclPpUniqueId, mpi::MpiTag::kNcclPpAck}, rendezvousId, deadline); + + ncclComm_t comm{nullptr}; +// Need static connection initialization for accurate KV cache size estimation +#if defined(_WIN32) + if (getenv("NCCL_RUNTIME_CONNECT") == nullptr) + _putenv_s("NCCL_RUNTIME_CONNECT", "0"); +#else + setenv("NCCL_RUNTIME_CONNECT", "0", 0); +#endif // _WIN32 + + ncclConfig_t config = NCCL_CONFIG_INITIALIZER; + config.blocking = 0; + int const commRank = static_cast(std::distance(activeRanks.begin(), rankIt)); + int const commSize = static_cast(activeRanks.size()); + auto const hostApiLock = acquireNcclHostApiLock(); + ncclResult_t const result = ncclCommInitRankConfig(&comm, commSize, id, commRank, &config); + if (result != ncclSuccess && result != ncclInProgress) + { + if (comm != nullptr) + { + ncclCommAbort(comm); + } + TLLM_THROW("%s", ncclErrorMessage("ncclCommInitRankConfig", result).c_str()); + } + + try + { + waitUntilReady(comm, "ncclCommInitRankConfig", deadline); + } + catch (...) + { + if (comm != nullptr) + { + ncclCommAbort(comm); + } + throw; + } + return comm; +#else + (void) activeRanks; + (void) worldRank; + (void) controlComm; + (void) rendezvousId; + (void) readyTimeout; + // Python runtime requires instantiation of a communicator even though it may never be used to enable + // pipeline parallel code-path. To enable this, have an empty communicator with uninitialized state. + return nullptr; +#endif // ENABLE_MULTI_DEVICE +} + +ncclComm_t NcclCommunicator::createLegacyComm(int worldSize, int rank, mpi::MpiComm const& mpiComm) +{ +#if ENABLE_MULTI_DEVICE ncclUniqueId id; if (rank == 0) { - ncclGetUniqueId(&id); + auto const hostApiLock = acquireNcclHostApiLock(); + TLLM_NCCL_CHECK(ncclGetUniqueId(&id)); } mpiComm.bcastValue(id, 0); - ncclComm_t comm; + // Need static connection initialization for accurate KV cache size estimation #if defined(_WIN32) if (getenv("NCCL_RUNTIME_CONNECT") == nullptr) @@ -87,21 +346,577 @@ ncclComm_t NcclCommunicator::createComm(int worldSize, int rank, mpi::MpiComm co #else setenv("NCCL_RUNTIME_CONNECT", "0", 0); #endif // _WIN32 + + ncclComm_t comm{nullptr}; + auto const hostApiLock = acquireNcclHostApiLock(); TLLM_NCCL_CHECK(ncclCommInitRank(&comm, worldSize, id, rank)); return comm; #else - // Python runtime requires instantiation of a communicator even though it may never be used to enable - // pipeline parallel code-path. To enable this, have an empty communicator with uninitialized state. + (void) worldSize; + (void) rank; + (void) mpiComm; return nullptr; #endif // ENABLE_MULTI_DEVICE } +void NcclCommunicator::waitUntilReady( + ncclComm_t comm, char const* operation, std::chrono::steady_clock::time_point deadline) +{ +#if ENABLE_MULTI_DEVICE + auto const hostApiLock = acquireNcclHostApiLock(); + TLLM_CHECK_WITH_INFO(comm != nullptr, "NCCL error: communicator was aborted while waiting for %s", operation); + while (true) + { + ncclResult_t asyncError = ncclInProgress; + ncclResult_t const queryResult = ncclCommGetAsyncError(comm, &asyncError); + if (queryResult != ncclSuccess && queryResult != ncclInProgress) + { + TLLM_THROW("%s", ncclErrorMessage("ncclCommGetAsyncError", queryResult).c_str()); + } + if (queryResult == ncclSuccess && asyncError == ncclSuccess) + { + return; + } + if (queryResult == ncclSuccess && asyncError != ncclInProgress) + { + TLLM_THROW("%s", ncclErrorMessage(operation, asyncError).c_str()); + } + if (std::chrono::steady_clock::now() >= deadline) + { + TLLM_THROW("NCCL error: timed out waiting for %s", operation); + } + std::this_thread::sleep_for(kReadyPollInterval); + } +#else + (void) comm; + (void) operation; + (void) deadline; +#endif // ENABLE_MULTI_DEVICE +} + +void NcclCommunicator::checkUsableLocked() const +{ + TLLM_CHECK_WITH_INFO(mAsyncError.empty(), "%s", mAsyncError.c_str()); + TLLM_CHECK_WITH_INFO(mComm != nullptr, "NCCL error: communicator was aborted"); +} + +bool NcclCommunicator::abortLocked(std::string reason) const noexcept +{ +#if ENABLE_MULTI_DEVICE + auto const hostApiLock = acquireNcclHostApiLock(); + if (mComm == nullptr) + { + if (mAsyncError.empty()) + { + mAsyncError = std::move(reason); + } + return true; + } + + ncclComm_t const comm = mComm; + ncclResult_t const result = ncclCommAbort(comm); + if (result == ncclSuccess) + { + mComm = nullptr; + mAsyncError = std::move(reason); + quarantinePendingOperationsLocked(); + return true; + } + + bool const firstAbortFailure = reason.find("ncclCommAbort failed") == std::string::npos; + mAsyncError = std::move(reason); + if (firstAbortFailure) + { + try + { + mAsyncError += "; ncclCommAbort failed: "; + mAsyncError += ncclGetErrorString(result); + } + catch (...) + { + } + } + if (firstAbortFailure) + { + TLLM_LOG_ERROR("Failed to abort NCCL communicator: %s", ncclGetErrorString(result)); + } + return false; +#else + (void) reason; + return true; +#endif // ENABLE_MULTI_DEVICE +} + +int NcclCommunicator::getCommPeerRankLocked(int worldRank) const +{ + auto const peer = std::find(mActiveRanks.begin(), mActiveRanks.end(), worldRank); + TLLM_CHECK_WITH_INFO(peer != mActiveRanks.end(), + "NCCL error: peer world rank %d is not active in the current communicator", worldRank); + return static_cast(std::distance(mActiveRanks.begin(), peer)); +} + +uint64_t NcclCommunicator::beginOperationLocked(cudaStream_t stream, char const* operation) const +{ +#if ENABLE_MULTI_DEVICE + checkUsableLocked(); + if (!mFaultToleranceEnabled) + { + return 0; + } + cudaStreamCaptureStatus captureStatus = cudaStreamCaptureStatusNone; + auto const captureResult = cudaStreamIsCapturing(stream, &captureStatus); + if (captureResult != cudaSuccess) + { + static_cast(abortLocked(std::string{"NCCL error: "} + operation + + " failed to query CUDA stream capture state: " + cudaGetErrorString(captureResult))); + TLLM_THROW("%s", mAsyncError.c_str()); + } + if (captureStatus != cudaStreamCaptureStatusNone) + { + TLLM_CHECK_WITH_INFO(!mFaultToleranceEnabled, + "NCCL error: CUDA graph capture of PP NCCL operations is disabled while " + "TLLM_FAULT_TOLERANCE_MODE=1; captured graph nodes retain the old communicator across recovery"); + TLLM_LOG_DEBUG("Skipping eager NCCL completion deadline for %s during CUDA graph capture", operation); + return 0; + } + if (mPendingOperations.size() >= kMaxPendingOperations) + { + static_cast(abortLocked( + std::string{"NCCL error: "} + operation + " exceeded the bounded NCCL watchdog operation queue")); + TLLM_THROW("%s", mAsyncError.c_str()); + } + + auto const token = mNextOperationToken++; + auto const start = acquireEventLocked(); + mPendingOperations.push_back({token, start, nullptr, {}, false, operation}); + auto const result = cudaEventRecord(start, stream); + if (result != cudaSuccess) + { + static_cast(abortLocked(std::string{"NCCL error: "} + operation + + " failed to record its start marker: " + cudaGetErrorString(result))); + TLLM_THROW("%s", mAsyncError.c_str()); + } + mWatcherWakeup.notify_all(); + return token; +#else + (void) stream; + (void) operation; + return 0; +#endif // ENABLE_MULTI_DEVICE +} + +void NcclCommunicator::finishOperationLocked(uint64_t token, cudaStream_t stream) const +{ +#if ENABLE_MULTI_DEVICE + checkUsableLocked(); + if (token == 0) + { + return; + } + auto const operation = std::find_if(mPendingOperations.begin(), mPendingOperations.end(), + [token](PendingOperation const& pending) { return pending.token == token; }); + TLLM_CHECK_WITH_INFO(operation != mPendingOperations.end(), "NCCL error: unknown watchdog operation token %llu", + static_cast(token)); + TLLM_CHECK_WITH_INFO(operation->completion == nullptr, + "NCCL error: watchdog operation token %llu was already completed", static_cast(token)); + + operation->completion = acquireEventLocked(); + auto const result = cudaEventRecord(operation->completion, stream); + if (result != cudaSuccess) + { + static_cast(abortLocked(std::string{"NCCL error: "} + operation->name + + " failed to record its completion marker: " + cudaGetErrorString(result))); + TLLM_THROW("%s", mAsyncError.c_str()); + } + mWatcherWakeup.notify_all(); +#else + (void) token; + (void) stream; +#endif // ENABLE_MULTI_DEVICE +} + +cudaEvent_t NcclCommunicator::acquireEventLocked() const +{ +#if ENABLE_MULTI_DEVICE + if (!mEventPool.empty()) + { + auto const event = mEventPool.back(); + mEventPool.pop_back(); + return event; + } + cudaEvent_t event = nullptr; + auto const result = cudaEventCreateWithFlags(&event, cudaEventDisableTiming); + if (result != cudaSuccess) + { + static_cast( + abortLocked(std::string{"NCCL error: failed to allocate a watchdog event: "} + cudaGetErrorString(result))); + TLLM_THROW("%s", mAsyncError.c_str()); + } + return event; +#else + return nullptr; +#endif // ENABLE_MULTI_DEVICE +} + +void NcclCommunicator::recycleEventLocked(cudaEvent_t event) const noexcept +{ +#if ENABLE_MULTI_DEVICE + if (event == nullptr) + { + return; + } + if (mEventPool.size() < kMaxPooledEvents) + { + mEventPool.push_back(event); + } + else + { + static_cast(cudaEventDestroy(event)); + } +#else + (void) event; +#endif // ENABLE_MULTI_DEVICE +} + +void NcclCommunicator::quarantinePendingOperationsLocked() const noexcept +{ + if (!mPendingOperations.empty()) + { + // The failed communicator may still have device work referencing + // these events. Drop host ownership without destroying their handles; + // process restart reclaims the deliberately quarantined resources. + TLLM_LOG_WARNING( + "Quarantining %zu PP NCCL completion markers after communicator abort", mPendingOperations.size()); + } + mPendingOperations.clear(); +} + +void NcclCommunicator::destroyCompletedWatchdogEventsLocked() const noexcept +{ +#if ENABLE_MULTI_DEVICE + size_t quarantined = 0; + for (auto const& operation : mPendingOperations) + { + for (auto const event : {operation.start, operation.completion}) + { + if (event == nullptr) + { + continue; + } + if (cudaEventQuery(event) == cudaSuccess) + { + static_cast(cudaEventDestroy(event)); + } + else + { + ++quarantined; + } + } + } + mPendingOperations.clear(); + if (quarantined != 0) + { + TLLM_LOG_WARNING("Quarantining %zu incomplete PP NCCL watchdog events during healthy teardown", quarantined); + } + destroyPooledEventsLocked(); +#endif // ENABLE_MULTI_DEVICE +} + +void NcclCommunicator::destroyPooledEventsLocked() const noexcept +{ +#if ENABLE_MULTI_DEVICE + for (auto const event : mEventPool) + { + static_cast(cudaEventDestroy(event)); + } + mEventPool.clear(); +#endif // ENABLE_MULTI_DEVICE +} + +void NcclCommunicator::abort() +{ +#if ENABLE_MULTI_DEVICE + std::lock_guard lock(mCommMutex); + bool const abortSucceeded = abortLocked("NCCL error: communicator was aborted by the fault-tolerance control path"); + TLLM_CHECK_WITH_INFO(abortSucceeded, "%s", mAsyncError.c_str()); +#endif // ENABLE_MULTI_DEVICE +} + +void NcclCommunicator::abortAndReinit(std::vector const& activeRanks, std::uint64_t rendezvousId) +{ +#if ENABLE_MULTI_DEVICE + std::lock_guard lock(mCommMutex); + TLLM_CHECK_WITH_INFO(mFaultToleranceEnabled, + "NCCL error: communicator reinitialization requires TLLM_FAULT_TOLERANCE_MODE=1; otherwise captured CUDA " + "graphs may retain the old communicator"); + TLLM_CHECK_WITH_INFO( + rendezvousId != 0, "NCCL error: recovery rendezvous ID must be positive; zero is reserved for bootstrap"); + TLLM_CHECK_WITH_INFO(mControlComm != nullptr, + "NCCL error: communicator created from a raw handle has no MPI bootstrap for reinitialization"); + TLLM_CHECK_WITH_INFO(!activeRanks.empty(), "NCCL error: active-rank set must not be empty"); + std::set const requested(activeRanks.begin(), activeRanks.end()); + TLLM_CHECK_WITH_INFO(requested.size() == activeRanks.size(), "NCCL error: active-rank set contains duplicates"); + TLLM_CHECK_WITH_INFO(requested.find(mWorldRank) != requested.end(), + "NCCL error: survivor rank %d must be present in active_ranks", mWorldRank); + + for (int const rank : requested) + { + TLLM_CHECK_WITH_INFO(rank >= 0 && rank < mInitialWorldSize, + "NCCL error: active world rank %d is outside [0, %d)", rank, mInitialWorldSize); + TLLM_CHECK_WITH_INFO(std::find(mActiveRanks.begin(), mActiveRanks.end(), rank) != mActiveRanks.end(), + "NCCL error: cannot re-add world rank %d after it was removed", rank); + } + + std::vector nextActiveRanks; + nextActiveRanks.reserve(requested.size()); + for (int const worldRank : mActiveRanks) + { + if (requested.find(worldRank) != requested.end()) + { + nextActiveRanks.push_back(worldRank); + } + } + + TLLM_CHECK_WITH_INFO(mControlComm->worldRank() == mWorldRank, + "NCCL error: world rank %d does not match rendezvous control world rank %d", mWorldRank, + mControlComm->worldRank()); + // This is deliberately a fresh-ID rebuild rather than ncclCommShrink. The + // watcher may already have destroyed the poisoned communicator, and the + // failed rank must not be required by any recovery collective. + bool const abortSucceeded + = abortLocked("NCCL error: communicator was aborted before survivor-only reinitialization"); + TLLM_CHECK_WITH_INFO(abortSucceeded, + "NCCL error: communicator reinitialization refused because ncclCommAbort did not succeed: %s", + mAsyncError.c_str()); + + ncclComm_t nextComm{nullptr}; + try + { + nextComm = createComm(nextActiveRanks, mWorldRank, *mControlComm, rendezvousId, kRecoveryReadyTimeout); + } + catch (std::exception const& error) + { + mAsyncError = std::string{"NCCL error: communicator reinitialization failed after abort: "} + error.what(); + TLLM_THROW("%s", mAsyncError.c_str()); + } + catch (...) + { + mAsyncError = "NCCL error: communicator reinitialization failed after abort"; + TLLM_THROW("%s", mAsyncError.c_str()); + } + + mComm = nextComm; + mActiveRanks = std::move(nextActiveRanks); + mAsyncError.clear(); +#else + (void) activeRanks; + (void) rendezvousId; + TLLM_THROW("NCCL error: multi device support is disabled."); +#endif // ENABLE_MULTI_DEVICE +} + +std::string NcclCommunicator::getAsyncError() const +{ + std::lock_guard lock(mCommMutex); + return mAsyncError; +} + +std::vector NcclCommunicator::getActiveRanks() const +{ + std::lock_guard lock(mCommMutex); + return mActiveRanks; +} + +void NcclCommunicator::startWatcher() +{ +#if ENABLE_MULTI_DEVICE + if (mFaultToleranceEnabled && mComm != nullptr) + { + try + { + mWatcherThread = std::thread(&NcclCommunicator::watchAsyncErrors, this); + } + catch (...) + { + auto const hostApiLock = acquireNcclHostApiLock(); + ncclCommAbort(mComm); + mComm = nullptr; + throw; + } + } +#endif // ENABLE_MULTI_DEVICE +} + +void NcclCommunicator::stopWatcher() noexcept +{ + mStopWatcher.store(true, std::memory_order_release); + mWatcherWakeup.notify_all(); + if (mWatcherThread.joinable()) + { + mWatcherThread.join(); + } +} + +void NcclCommunicator::watchAsyncErrors() noexcept +{ +#if ENABLE_MULTI_DEVICE + try + { + while (!mStopWatcher.load(std::memory_order_acquire)) + { + { + // NCCL communicators are not generally thread-safe. Serialize + // the monitor with send/receive, abort, and reinit host calls. + std::lock_guard lock(mCommMutex); + if (mComm != nullptr && !mAsyncError.empty()) + { + // A previous abort attempt failed. Keep the handle so the + // watcher can retry; a rebuild is forbidden until abort + // succeeds. + static_cast(abortLocked(mAsyncError)); + } + else if (mComm != nullptr) + { + std::string error; + auto const now = std::chrono::steady_clock::now(); + for (auto operation = mPendingOperations.begin(); operation != mPendingOperations.end();) + { + if (operation->start != nullptr) + { + auto const startResult = cudaEventQuery(operation->start); + if (startResult == cudaSuccess) + { + recycleEventLocked(operation->start); + operation->start = nullptr; + operation->deadline = now + getOperationTimeout(); + operation->armed = true; + } + else if (startResult != cudaErrorNotReady) + { + error = operation->name + " start marker failed: " + cudaGetErrorString(startResult); + break; + } + else + { + ++operation; + continue; + } + } + if (operation->completion == nullptr) + { + if (operation->armed && now >= operation->deadline) + { + error = operation->name + " timed out before its completion marker was recorded"; + break; + } + ++operation; + continue; + } + + auto const eventResult = cudaEventQuery(operation->completion); + if (eventResult == cudaSuccess) + { + recycleEventLocked(operation->completion); + operation = mPendingOperations.erase(operation); + continue; + } + if (eventResult != cudaErrorNotReady) + { + error = operation->name + " completion marker failed: " + cudaGetErrorString(eventResult); + break; + } + if (operation->armed && now >= operation->deadline) + { + error = operation->name + " timed out before its CUDA completion marker"; + break; + } + ++operation; + } + + if (error.empty()) + { + auto const hostApiLock = acquireNcclHostApiLock(); + ncclResult_t asyncError = ncclSuccess; + ncclResult_t const queryResult = ncclCommGetAsyncError(mComm, &asyncError); + if (queryResult != ncclSuccess && queryResult != ncclInProgress) + { + error = ncclErrorMessage("ncclCommGetAsyncError watcher", queryResult); + } + else if (queryResult == ncclSuccess && asyncError != ncclSuccess + && asyncError != ncclInProgress) + { + error = ncclErrorMessage("asynchronous communicator operation", asyncError); + } + } + + if (!error.empty()) + { + static_cast(abortLocked( + "NCCL error: communicator was aborted by the async-error watcher after " + error)); + TLLM_LOG_ERROR("%s", mAsyncError.c_str()); + } + } + } + + std::unique_lock waitLock(mWatcherWaitMutex); + mWatcherWakeup.wait_for( + waitLock, kWatcherPollInterval, [this]() { return mStopWatcher.load(std::memory_order_acquire); }); + } + } + catch (...) + { + // Never allow a diagnostic side thread to terminate the process. + } +#endif // ENABLE_MULTI_DEVICE +} + NcclCommunicator::~NcclCommunicator() { #if ENABLE_MULTI_DEVICE - if (mComm && ncclCommDestroy(mComm) != ncclSuccess) + if (!mFaultToleranceEnabled) + { + if (mComm != nullptr && ncclCommDestroy(mComm) != ncclSuccess) + { + TLLM_LOG_WARNING("Failed to destroy NCCL communicator."); + } + return; + } +#endif // ENABLE_MULTI_DEVICE + stopWatcher(); +#if ENABLE_MULTI_DEVICE + std::lock_guard lock(mCommMutex); + auto const hostApiLock = acquireNcclHostApiLock(); + if (mComm != nullptr) + { + ncclResult_t asyncError = ncclSuccess; + ncclResult_t const queryResult = ncclCommGetAsyncError(mComm, &asyncError); + bool const healthy = mAsyncError.empty() && queryResult == ncclSuccess && asyncError == ncclSuccess; + // ncclCommDestroy is an intra-node collective. Python model-engine + // reference counts are process-local, so their final release is not a + // cross-rank teardown barrier; after a peer failure, the one healthy + // sample above is also inherently racy. In FT mode use the local abort + // path and never stop the watcher only to enter an unbounded graceful + // destroy while holding the process-wide NCCL gate. + bool const useAbort = !healthy || mFaultToleranceEnabled; + ncclResult_t const result = useAbort ? ncclCommAbort(mComm) : ncclCommDestroy(mComm); + if (result != ncclSuccess) + { + TLLM_LOG_WARNING("Failed to release NCCL communicator: %s", ncclGetErrorString(result)); + } + mComm = nullptr; + if (!useAbort && result == ncclSuccess) + { + destroyCompletedWatchdogEventsLocked(); + } + else + { + quarantinePendingOperationsLocked(); + destroyPooledEventsLocked(); + } + } + else { - TLLM_LOG_WARNING("Failed to destroy NCCL communicator."); + destroyPooledEventsLocked(); } #endif // ENABLE_MULTI_DEVICE } diff --git a/cpp/tensorrt_llm/runtime/ncclCommunicator.h b/cpp/tensorrt_llm/runtime/ncclCommunicator.h index 76cce4beab8a..388030e234b5 100644 --- a/cpp/tensorrt_llm/runtime/ncclCommunicator.h +++ b/cpp/tensorrt_llm/runtime/ncclCommunicator.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -21,20 +21,43 @@ #include "tensorrt_llm/runtime/utils/mpiUtils.h" #include "tensorrt_llm/runtime/worldConfig.h" +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + struct ncclComm; typedef struct ncclComm* ncclComm_t; namespace tensorrt_llm::runtime { +class NcclUniqueIdRendezvousComm; + class NcclCommunicator { public: - explicit NcclCommunicator(ncclComm_t comm) - : mComm{comm} {}; - - explicit NcclCommunicator(int worldSize, int rank, mpi::MpiComm const& mpiComm = COMM_SESSION) - : mComm{createComm(worldSize, rank, mpiComm)} {}; + //! Wrap an existing communicator. + //! + //! In fault-tolerance mode the caller is responsible for creating `comm` + //! in non-blocking mode. The default-off path preserves the legacy + //! blocking communicator contract. + explicit NcclCommunicator(ncclComm_t comm); + + //! Create a communicator bootstrapped by `mpiComm`. + //! + //! Construction establishes and retains a dedicated MPI control channel while + //! the initial group is healthy. Recovery reuses that channel without an + //! MPI collective; `mpiComm` therefore need not outlive this object. + explicit NcclCommunicator(int worldSize, int rank, mpi::MpiComm const& mpiComm = COMM_SESSION); explicit NcclCommunicator(WorldConfig const& worldConfig, mpi::MpiComm const& mpiComm = COMM_SESSION) : NcclCommunicator{worldConfig.getSize(), worldConfig.getRank(), mpiComm} {}; @@ -55,15 +78,86 @@ class NcclCommunicator receive(buf.data(), buf.getSize(), buf.getDataType(), peer, stream); } + //! Abort the current communicator. Idempotent. + void abort(); + + //! Abort and replace the current communicator with a fresh communicator + //! containing only `activeRanks`. + //! + //! Ranks are the original world-rank IDs supplied to the constructor. The + //! local rank must be present, and previously removed ranks cannot be added + //! again. Every survivor must call this method with the same rank set. The + //! positive `rendezvousId` must identify the same logical recovery attempt + //! on every survivor; zero is reserved for the initial bootstrap. The NCCL + //! unique ID is exchanged point-to-point over the dedicated control + //! channel, so failed ranks do not participate in recovery. + //! + //! Communicators created by the raw-handle constructor cannot be rebuilt + //! because that constructor has no MPI bootstrap communicator. + void abortAndReinit(std::vector const& activeRanks, std::uint64_t rendezvousId); + + //! Return the latched NCCL abort/error reason, or an empty string. + [[nodiscard]] std::string getAsyncError() const; + + //! Return the original world-rank IDs in the current communicator. + [[nodiscard]] std::vector getActiveRanks() const; + private: void send( void const* sendbuff, size_t count, nvinfer1::DataType dataType, int peer, CudaStream const& stream) const; void receive(void* sendbuff, size_t count, nvinfer1::DataType dataType, int peer, CudaStream const& stream) const; - static ncclComm_t createComm(int worldSize, int rank, mpi::MpiComm const& mpiComm); - - ncclComm_t mComm; + static ncclComm_t createComm(std::vector const& activeRanks, int worldRank, + NcclUniqueIdRendezvousComm const& controlComm, std::uint64_t rendezvousId, + std::chrono::milliseconds readyTimeout); + static ncclComm_t createLegacyComm(int worldSize, int rank, mpi::MpiComm const& mpiComm); + static void waitUntilReady(ncclComm_t comm, char const* operation, std::chrono::steady_clock::time_point deadline); + + void startWatcher(); + void stopWatcher() noexcept; + void watchAsyncErrors() noexcept; + [[nodiscard]] bool abortLocked(std::string reason) const noexcept; + void checkUsableLocked() const; + [[nodiscard]] int getCommPeerRankLocked(int worldRank) const; + [[nodiscard]] uint64_t beginOperationLocked(cudaStream_t stream, char const* operation) const; + void finishOperationLocked(uint64_t token, cudaStream_t stream) const; + [[nodiscard]] cudaEvent_t acquireEventLocked() const; + void recycleEventLocked(cudaEvent_t event) const noexcept; + void quarantinePendingOperationsLocked() const noexcept; + void destroyCompletedWatchdogEventsLocked() const noexcept; + void destroyPooledEventsLocked() const noexcept; + + int mInitialWorldSize{0}; + int mWorldRank{0}; + bool mFaultToleranceEnabled{false}; + std::vector mActiveRanks; + std::shared_ptr mControlComm; + + mutable std::mutex mCommMutex; + mutable ncclComm_t mComm{nullptr}; + mutable std::string mAsyncError; + + struct PendingOperation + { + uint64_t token; + cudaEvent_t start; + cudaEvent_t completion; + std::chrono::steady_clock::time_point deadline; + bool armed; + std::string name; + }; + + mutable std::vector mPendingOperations; + mutable std::vector mEventPool; + mutable uint64_t mNextOperationToken{1}; + static constexpr size_t kMaxPendingOperations = 4096; + static constexpr size_t kMaxPooledEvents = 256; + + std::atomic mStopWatcher{false}; + std::mutex mWatcherWaitMutex; + mutable std::condition_variable mWatcherWakeup; + std::thread mWatcherThread; }; } // namespace tensorrt_llm::runtime diff --git a/cpp/tensorrt_llm/runtime/utils/ncclUniqueIdRendezvous.cpp b/cpp/tensorrt_llm/runtime/utils/ncclUniqueIdRendezvous.cpp new file mode 100644 index 000000000000..7d13162d5dc9 --- /dev/null +++ b/cpp/tensorrt_llm/runtime/utils/ncclUniqueIdRendezvous.cpp @@ -0,0 +1,787 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * 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 "tensorrt_llm/runtime/utils/ncclUniqueIdRendezvous.h" + +#if ENABLE_MULTI_DEVICE + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/runtime/utils/ncclHostApi.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace tensorrt_llm::runtime +{ +namespace +{ + +using Clock = std::chrono::steady_clock; + +constexpr std::uint32_t kProtocolVersion = 2; +constexpr auto kPollInterval = std::chrono::milliseconds{1}; +constexpr auto kReadyRefreshInterval = std::chrono::milliseconds{100}; +constexpr std::size_t kMaxTokensPerPeer = 1024; +constexpr std::size_t kMaxPendingMessages = 4096; + +struct AttemptToken +{ + std::uint64_t incarnation; + std::uint64_t sequence; +}; + +bool operator==(AttemptToken const& lhs, AttemptToken const& rhs) noexcept +{ + return lhs.incarnation == rhs.incarnation && lhs.sequence == rhs.sequence; +} + +struct ReadyMessage +{ + std::uint32_t version; + std::uint32_t reserved; + std::uint64_t communicatorKey; + std::uint64_t rendezvousId; + AttemptToken clientToken; +}; + +struct IdMessage +{ + std::uint32_t version; + std::uint32_t reserved; + std::uint64_t communicatorKey; + std::uint64_t rendezvousId; + AttemptToken clientToken; + AttemptToken serverToken; + std::uint64_t idDigest; + ncclUniqueId id; +}; + +struct AckMessage +{ + std::uint32_t version; + std::uint32_t reserved; + std::uint64_t communicatorKey; + std::uint64_t rendezvousId; + AttemptToken clientToken; + AttemptToken serverToken; + std::uint64_t idDigest; +}; + +template +struct ReceivedMessage +{ + int source; + Message message; +}; + +struct RendezvousContext +{ + std::mutex exchangeMutex; + std::deque> pendingReady; + std::deque> pendingIds; + std::deque> pendingAcks; +}; + +struct ControlCommCacheEntry +{ + ControlCommCacheEntry(MPI_Group parentGroup_, std::vector initialRanks_, int worldRank_, int creationTagSeed_) + : parentGroup(parentGroup_) + , initialRanks(std::move(initialRanks_)) + , worldRank(worldRank_) + , creationTagSeed(creationTagSeed_) + { + } + + std::mutex mutex; + MPI_Group parentGroup{MPI_GROUP_NULL}; + std::vector initialRanks; + int worldRank; + int creationTagSeed; + std::shared_ptr comm; + std::exception_ptr failure; +}; + +std::uint64_t groupIdentity(std::vector const& ranks) noexcept +{ + std::uint64_t hash = 14695981039346656037ULL; + for (int const rank : ranks) + { + auto const* bytes = reinterpret_cast(&rank); + for (std::size_t index = 0; index < sizeof(rank); ++index) + { + hash ^= bytes[index]; + hash *= 1099511628211ULL; + } + } + return hash; +} + +using ContextKey = std::tuple; + +ContextKey getContextKey(NcclUniqueIdRendezvousTags const& tags, NcclUniqueIdRendezvousComm const& controlComm) noexcept +{ + return {MPI_Comm_c2f(static_cast(controlComm.mpiComm())), groupIdentity(controlComm.worldRanks()), + static_cast(tags.ready), static_cast(tags.id), static_cast(tags.ack)}; +} + +RendezvousContext& getRendezvousContext( + NcclUniqueIdRendezvousTags const& tags, NcclUniqueIdRendezvousComm const& controlComm) +{ + // Keep contexts immortal for the same reason as the NCCL host gate: a + // communicator may be torn down from another process-lifetime singleton. + static auto* registryMutex = new std::mutex; + static auto* contexts = new std::map>; + std::lock_guard lock(*registryMutex); + auto& context = (*contexts)[getContextKey(tags, controlComm)]; + if (context == nullptr) + { + context = std::make_unique(); + } + return *context; +} + +std::shared_ptr getControlCommCacheEntry( + std::vector const& initialRanks, int worldRank, mpi::MpiComm const& parentComm, int creationTagSeed) +{ + // A control communicator cannot be freed safely after a member dies. Keep + // one channel per topology for the process lifetime and reuse it across PP + // engine teardown/recreation instead of leaking a new MPI context each + // time. Retain the parent's MPI group as a stable process-membership + // identity: MPI communicator handles may be recycled after a custom parent + // is freed. Per-entry locking permits unrelated groups to create + // concurrently. + static auto* registryMutex = new std::mutex; + static auto* entries = new std::vector>; + + MPI_Group parentGroup = MPI_GROUP_NULL; + int result = MPI_Comm_group(static_cast(parentComm), &parentGroup); + TLLM_CHECK_WITH_INFO( + result == MPI_SUCCESS, "NCCL error: failed to retain MPI parent group identity (code %d)", result); + try + { + std::lock_guard lock(*registryMutex); + for (auto const& entry : *entries) + { + if (entry->initialRanks != initialRanks || entry->worldRank != worldRank + || entry->creationTagSeed != creationTagSeed) + { + continue; + } + int comparison = MPI_UNEQUAL; + result = MPI_Group_compare(entry->parentGroup, parentGroup, &comparison); + TLLM_CHECK_WITH_INFO( + result == MPI_SUCCESS, "NCCL error: failed to compare MPI parent group identity (code %d)", result); + if (comparison == MPI_IDENT) + { + MPI_Group_free(&parentGroup); + return entry; + } + } + + auto entry = std::make_shared(parentGroup, initialRanks, worldRank, creationTagSeed); + entries->push_back(entry); + parentGroup = MPI_GROUP_NULL; + return entry; + } + catch (...) + { + if (parentGroup != MPI_GROUP_NULL) + { + MPI_Group_free(&parentGroup); + } + throw; + } +} + +template +bool takePending(std::deque>& pending, int source, std::uint64_t key, + std::uint64_t rendezvousId, ReceivedMessage& result) +{ + auto found = pending.begin(); + for (; found != pending.end(); ++found) + { + if ((source == MPI_ANY_SOURCE || found->source == source) && found->message.communicatorKey == key + && found->message.rendezvousId == rendezvousId) + { + break; + } + } + if (found == pending.end()) + { + return false; + } + result = *found; + pending.erase(found); + return true; +} + +template +void discardObsoletePending( + std::deque>& pending, std::uint64_t key, std::uint64_t rendezvousId) +{ + pending.erase(std::remove_if(pending.begin(), pending.end(), + [key, rendezvousId](auto const& received) + { + return received.message.version == kProtocolVersion && received.message.communicatorKey == key + && received.message.rendezvousId < rendezvousId; + }), + pending.end()); +} + +template +bool isCurrentAttempt(Message const& message, std::uint64_t key, std::uint64_t rendezvousId) noexcept +{ + return message.communicatorKey == key && message.rendezvousId == rendezvousId; +} + +template +bool shouldStashForLater(Message const& message, std::uint64_t key, std::uint64_t rendezvousId) noexcept +{ + if (message.version != kProtocolVersion) + { + return false; + } + // Rendezvous IDs must increase for retries of one base communicator. + // Messages from an older retry can never be consumed again, while a peer + // may legitimately have entered a future retry or a different group first. + return message.communicatorKey != key || message.rendezvousId > rendezvousId; +} + +template +void stashPending(std::deque>& pending, ReceivedMessage received) +{ + TLLM_CHECK_WITH_INFO( + pending.size() < kMaxPendingMessages, "NCCL error: communicator rendezvous pending-message limit exceeded"); + pending.push_back(std::move(received)); +} + +std::uint64_t fnv1a(void const* data, std::size_t size, std::uint64_t hash = 14695981039346656037ULL) noexcept +{ + auto const* bytes = static_cast(data); + for (std::size_t index = 0; index < size; ++index) + { + hash ^= bytes[index]; + hash *= 1099511628211ULL; + } + return hash; +} + +std::uint64_t communicatorKey(std::vector const& activeRanks, NcclUniqueIdRendezvousTags const& tags) noexcept +{ + std::uint64_t hash = fnv1a(activeRanks.data(), activeRanks.size() * sizeof(activeRanks.front())); + std::array const tagValues{ + static_cast(tags.ready), static_cast(tags.id), static_cast(tags.ack)}; + return fnv1a(tagValues.data(), tagValues.size() * sizeof(tagValues.front()), hash); +} + +std::uint64_t processIncarnation() +{ + static std::uint64_t const incarnation = []() + { + std::random_device random; + std::array entropy{random(), random(), random(), random()}; + auto const wallClock = std::chrono::system_clock::now().time_since_epoch().count(); + std::uint64_t hash = fnv1a(entropy.data(), entropy.size() * sizeof(entropy.front())); + hash = fnv1a(&wallClock, sizeof(wallClock), hash); + return hash == 0 ? std::uint64_t{1} : hash; + }(); + return incarnation; +} + +AttemptToken makeAttemptToken() +{ + static std::atomic sequence{1}; + auto const value = sequence.fetch_add(1, std::memory_order_relaxed); + TLLM_CHECK_WITH_INFO(value != 0, "NCCL error: communicator rendezvous token sequence overflow"); + return AttemptToken{processIncarnation(), value}; +} + +int abandonRequest(MPI_Request& request) noexcept +{ + int firstError = MPI_SUCCESS; + if (request != MPI_REQUEST_NULL) + { + firstError = MPI_Cancel(&request); + // MPI_Wait can block forever after a peer failure, defeating the + // recovery deadline. Free the local request handle and let MPI retire + // it asynchronously; callers deliberately leak the tiny backing + // payload so the implementation can no longer access freed storage. + int const freeResult = MPI_Request_free(&request); + if (firstError == MPI_SUCCESS) + { + firstError = freeResult; + } + } + return firstError; +} + +void waitForSend(MPI_Request& request, Clock::time_point deadline, char const* operation) +{ + while (true) + { + int complete = 0; + int const result = MPI_Test(&request, &complete, MPI_STATUS_IGNORE); + if (result != MPI_SUCCESS) + { + int const cleanupResult = abandonRequest(request); + TLLM_THROW( + "NCCL error: MPI_Test failed while %s (code %d, cleanup code %d)", operation, result, cleanupResult); + } + if (complete != 0) + { + return; + } + if (Clock::now() >= deadline) + { + int const cleanupResult = abandonRequest(request); + TLLM_THROW("NCCL error: timed out while %s (MPI cleanup code %d)", operation, cleanupResult); + } + std::this_thread::sleep_for(kPollInterval); + } +} + +template +void sendMessage(Message const& message, int destination, mpi::MpiTag tag, mpi::MpiComm const& mpiComm, + Clock::time_point deadline, char const* operation) +{ + auto payload = std::make_unique(message); + MPI_Request request = MPI_REQUEST_NULL; + int const result = MPI_Isend(payload.get(), static_cast(sizeof(message)), MPI_BYTE, destination, + static_cast(tag), static_cast(mpiComm), &request); + if (result != MPI_SUCCESS) + { + int const cleanupResult = abandonRequest(request); + static_cast(payload.release()); + TLLM_THROW("NCCL error: MPI_Isend failed while %s to rank %d (code %d, cleanup code %d)", operation, + destination, result, cleanupResult); + } + try + { + waitForSend(request, deadline, operation); + } + catch (...) + { + static_cast(payload.release()); + throw; + } +} + +bool probeMessage(int source, mpi::MpiTag tag, mpi::MpiComm const& mpiComm, MPI_Status& status) +{ + int available = 0; + int const result = MPI_Iprobe(source, static_cast(tag), static_cast(mpiComm), &available, &status); + TLLM_CHECK_WITH_INFO(result == MPI_SUCCESS, "NCCL error: MPI_Iprobe failed with code %d", result); + return available != 0; +} + +template +Message receiveMessage(MPI_Status const& probedStatus, mpi::MpiTag tag, mpi::MpiComm const& mpiComm, + Clock::time_point deadline, char const* operation) +{ + auto message = std::make_unique(); + MPI_Request request = MPI_REQUEST_NULL; + int const result = MPI_Irecv(message.get(), static_cast(sizeof(Message)), MPI_BYTE, probedStatus.MPI_SOURCE, + static_cast(tag), static_cast(mpiComm), &request); + if (result != MPI_SUCCESS) + { + int const cleanupResult = abandonRequest(request); + static_cast(message.release()); + TLLM_THROW("NCCL error: MPI_Irecv failed while %s from rank %d (code %d, cleanup code %d)", operation, + probedStatus.MPI_SOURCE, result, cleanupResult); + } + + while (true) + { + int complete = 0; + MPI_Status status{}; + int const testResult = MPI_Test(&request, &complete, &status); + if (testResult != MPI_SUCCESS) + { + int const cleanupResult = abandonRequest(request); + static_cast(message.release()); + TLLM_THROW("NCCL error: MPI_Test failed while %s (code %d, cleanup code %d)", operation, testResult, + cleanupResult); + } + if (complete != 0) + { + int count = 0; + int const countResult = MPI_Get_count(&status, MPI_BYTE, &count); + TLLM_CHECK_WITH_INFO(countResult == MPI_SUCCESS && count == static_cast(sizeof(Message)), + "NCCL error: invalid MPI rendezvous payload while %s", operation); + return *message; + } + if (Clock::now() >= deadline) + { + int const cleanupResult = abandonRequest(request); + static_cast(message.release()); + TLLM_THROW("NCCL error: timed out while %s (MPI cleanup code %d)", operation, cleanupResult); + } + std::this_thread::sleep_for(kPollInterval); + } +} + +bool isActivePeer(std::vector const& activeRanks, int rank, int root) +{ + return rank != root && std::binary_search(activeRanks.begin(), activeRanks.end(), rank); +} + +} // namespace + +NcclUniqueIdRendezvousComm::NcclUniqueIdRendezvousComm(mpi::MpiComm comm, std::vector worldRanks, int worldRank) + : mComm(std::move(comm)) + , mWorldRanks(std::move(worldRanks)) + , mWorldRank(worldRank) +{ + TLLM_CHECK_WITH_INFO(!mWorldRanks.empty(), "NCCL error: rendezvous control group must not be empty"); + TLLM_CHECK_WITH_INFO(std::is_sorted(mWorldRanks.begin(), mWorldRanks.end()), + "NCCL error: rendezvous control ranks must be in canonical order"); + TLLM_CHECK_WITH_INFO(std::adjacent_find(mWorldRanks.begin(), mWorldRanks.end()) == mWorldRanks.end(), + "NCCL error: rendezvous control ranks contain duplicates"); + TLLM_CHECK_WITH_INFO(mComm.getSize() == static_cast(mWorldRanks.size()), + "NCCL error: rendezvous control communicator size %d does not match rank map size %zu", mComm.getSize(), + mWorldRanks.size()); + TLLM_CHECK_WITH_INFO(worldRank == this->worldRank(mComm.getRank()), + "NCCL error: world rank %d does not match rendezvous control rank %d", worldRank, mComm.getRank()); +} + +mpi::MpiComm const& NcclUniqueIdRendezvousComm::mpiComm() const noexcept +{ + return mComm; +} + +std::vector const& NcclUniqueIdRendezvousComm::worldRanks() const noexcept +{ + return mWorldRanks; +} + +int NcclUniqueIdRendezvousComm::worldRank() const noexcept +{ + return mWorldRank; +} + +int NcclUniqueIdRendezvousComm::commRank(int worldRank) const +{ + auto const found = std::lower_bound(mWorldRanks.begin(), mWorldRanks.end(), worldRank); + TLLM_CHECK_WITH_INFO(found != mWorldRanks.end() && *found == worldRank, + "NCCL error: world rank %d is not present in the rendezvous control channel", worldRank); + return static_cast(std::distance(mWorldRanks.begin(), found)); +} + +int NcclUniqueIdRendezvousComm::worldRank(int commRank) const +{ + TLLM_CHECK_WITH_INFO(commRank >= 0 && commRank < static_cast(mWorldRanks.size()), + "NCCL error: rendezvous control rank %d is outside [0, %zu)", commRank, mWorldRanks.size()); + return mWorldRanks[static_cast(commRank)]; +} + +std::shared_ptr createNcclUniqueIdRendezvousComm( + std::vector const& initialRanks, int worldRank, mpi::MpiComm const& parentComm, int creationTagSeed) +{ + TLLM_CHECK_WITH_INFO(!initialRanks.empty(), "NCCL error: rendezvous control group must not be empty"); + TLLM_CHECK_WITH_INFO(std::is_sorted(initialRanks.begin(), initialRanks.end()), + "NCCL error: rendezvous control ranks must be in canonical order"); + TLLM_CHECK_WITH_INFO(std::adjacent_find(initialRanks.begin(), initialRanks.end()) == initialRanks.end(), + "NCCL error: rendezvous control ranks contain duplicates"); + TLLM_CHECK_WITH_INFO(std::binary_search(initialRanks.begin(), initialRanks.end(), worldRank), + "NCCL error: world rank %d is not a member of the rendezvous control group", worldRank); + TLLM_CHECK_WITH_INFO(parentComm.getRank() == worldRank, + "NCCL error: world rank %d does not match MPI parent rank %d", worldRank, parentComm.getRank()); + TLLM_CHECK_WITH_INFO(initialRanks.front() >= 0 && initialRanks.back() < parentComm.getSize(), + "NCCL error: rendezvous control ranks must be within MPI parent size %d", parentComm.getSize()); + + auto cacheEntry = getControlCommCacheEntry(initialRanks, worldRank, parentComm, creationTagSeed); + std::lock_guard const cacheLock(cacheEntry->mutex); + if (cacheEntry->comm != nullptr) + { + return cacheEntry->comm; + } + if (cacheEntry->failure != nullptr) + { + // Creation is collective over the initial group. Never let one rank + // retry that collective after a prior asymmetric local failure while + // peers may already have cached a successfully-created channel. + std::rethrow_exception(cacheEntry->failure); + } + + void* tagUpperBoundValue = nullptr; + int hasTagUpperBound = 0; + int result + = MPI_Comm_get_attr(static_cast(parentComm), MPI_TAG_UB, &tagUpperBoundValue, &hasTagUpperBound); + TLLM_CHECK_WITH_INFO( + result == MPI_SUCCESS, "NCCL error: failed to query MPI_TAG_UB for the control channel (code %d)", result); + TLLM_CHECK_WITH_INFO(hasTagUpperBound != 0 && tagUpperBoundValue != nullptr, + "NCCL error: MPI parent communicator does not expose MPI_TAG_UB"); + int const tagUpperBound = *static_cast(tagUpperBoundValue); + constexpr int kFirstControlCreationTag = static_cast(mpi::MpiTag::kNcclPpControl) + 1; + TLLM_CHECK_WITH_INFO(creationTagSeed >= 0, "NCCL error: rendezvous control creation-tag seed must be nonnegative"); + int const tagParity = creationTagSeed & 1; + TLLM_CHECK_WITH_INFO(tagUpperBound >= kFirstControlCreationTag + tagParity, + "NCCL error: MPI_TAG_UB %d is too small for NCCL control-channel creation", tagUpperBound); + std::uint64_t tagHash = groupIdentity(initialRanks); + tagHash ^= static_cast(creationTagSeed) + 0x9e3779b97f4a7c15ULL + (tagHash << 6) + (tagHash >> 2); + // Even/odd slots separate raw-op and PP seeds; the group hash makes lazy + // creation order-independent within each ownership path. + auto const tagSlots = static_cast(tagUpperBound - kFirstControlCreationTag - tagParity) / 2 + 1; + int const creationTag = kFirstControlCreationTag + tagParity + static_cast(2 * (tagHash % tagSlots)); + + MPI_Group parentGroup = MPI_GROUP_NULL; + MPI_Group controlGroup = MPI_GROUP_NULL; + MPI_Comm controlComm = MPI_COMM_NULL; + auto cleanup = [&]() noexcept + { + if (controlGroup != MPI_GROUP_NULL) + { + MPI_Group_free(&controlGroup); + } + if (parentGroup != MPI_GROUP_NULL) + { + MPI_Group_free(&parentGroup); + } + // Once MPI_Comm_create_group has returned a non-null communicator, + // never call collective MPI_Comm_free from an exception path. A + // local allocation or error-handler failure can be asymmetric, so + // peers may already have published their process-lifetime channel. + // Leaking that rare failed construction is safer than deadlocking. + }; + + try + { + result = MPI_Comm_group(static_cast(parentComm), &parentGroup); + TLLM_CHECK_WITH_INFO(result == MPI_SUCCESS, "NCCL error: failed to inspect MPI parent group (code %d)", result); + result = MPI_Group_incl(parentGroup, static_cast(initialRanks.size()), initialRanks.data(), &controlGroup); + TLLM_CHECK_WITH_INFO(result == MPI_SUCCESS, "NCCL error: failed to create MPI control group (code %d)", result); + // MPI_Comm_create_group is collective only over initialRanks. This is + // the sole control-channel collective and must happen before failure. + result = MPI_Comm_create_group(static_cast(parentComm), controlGroup, creationTag, &controlComm); + TLLM_CHECK_WITH_INFO( + result == MPI_SUCCESS, "NCCL error: failed to create MPI control communicator (code %d)", result); + TLLM_CHECK_WITH_INFO(controlComm != MPI_COMM_NULL, + "NCCL error: MPI returned a null rendezvous control communicator for member rank %d", worldRank); + result = MPI_Comm_set_errhandler(controlComm, MPI_ERRORS_RETURN); + TLLM_CHECK_WITH_INFO(result == MPI_SUCCESS, + "NCCL error: failed to configure MPI_ERRORS_RETURN on the control channel (code %d)", result); + + MPI_Group_free(&controlGroup); + controlGroup = MPI_GROUP_NULL; + MPI_Group_free(&parentGroup); + parentGroup = MPI_GROUP_NULL; + // Never call MPI_Comm_free on a successfully-published FT control + // channel: freeing after a member dies is not guaranteed to make + // progress. The communicator intentionally lives until process exit. + auto ownedComm = mpi::MpiComm{controlComm, false}; + auto resultComm = std::make_shared(std::move(ownedComm), initialRanks, worldRank); + controlComm = MPI_COMM_NULL; + cacheEntry->comm = std::move(resultComm); + return cacheEntry->comm; + } + catch (...) + { + cacheEntry->failure = std::current_exception(); + cleanup(); + throw; + } +} + +ncclUniqueId exchangeNcclUniqueId(std::vector const& activeRanks, NcclUniqueIdRendezvousComm const& controlComm, + NcclUniqueIdRendezvousTags tags, std::uint64_t rendezvousId, Clock::time_point deadline) +{ + auto const worldRank = controlComm.worldRank(); + auto const& mpiComm = controlComm.mpiComm(); + TLLM_CHECK_WITH_INFO(!activeRanks.empty(), "NCCL error: communicator rendezvous group must not be empty"); + TLLM_CHECK_WITH_INFO(std::is_sorted(activeRanks.begin(), activeRanks.end()), + "NCCL error: communicator rendezvous ranks must be in canonical order"); + TLLM_CHECK_WITH_INFO(std::adjacent_find(activeRanks.begin(), activeRanks.end()) == activeRanks.end(), + "NCCL error: communicator rendezvous ranks contain duplicates"); + TLLM_CHECK_WITH_INFO(std::binary_search(activeRanks.begin(), activeRanks.end(), worldRank), + "NCCL error: world rank %d is not active in communicator rendezvous", worldRank); + TLLM_CHECK_WITH_INFO(std::includes(controlComm.worldRanks().begin(), controlComm.worldRanks().end(), + activeRanks.begin(), activeRanks.end()), + "NCCL error: communicator rendezvous group is not a subset of the pre-failure control group"); + + auto& context = getRendezvousContext(tags, controlComm); + std::unique_lock const exchangeLock(context.exchangeMutex); + int const root = activeRanks.front(); + std::uint64_t const key = communicatorKey(activeRanks, tags); + discardObsoletePending(context.pendingReady, key, rendezvousId); + discardObsoletePending(context.pendingIds, key, rendezvousId); + discardObsoletePending(context.pendingAcks, key, rendezvousId); + if (worldRank == root) + { + ncclUniqueId id{}; + { + auto const hostApiLock = acquireNcclHostApiLock(); + auto const result = ncclGetUniqueId(&id); + TLLM_CHECK_WITH_INFO( + result == ncclSuccess, "NCCL error: ncclGetUniqueId failed: %s", ncclGetErrorString(result)); + } + if (activeRanks.size() == 1) + { + return id; + } + + AttemptToken const serverToken = makeAttemptToken(); + std::uint64_t const digest = fnv1a(&id, sizeof(id)); + std::unordered_map> issuedTokens; + std::unordered_set acknowledged; + while (acknowledged.size() + 1 < activeRanks.size()) + { + TLLM_CHECK_WITH_INFO( + Clock::now() < deadline, "NCCL error: timed out waiting for survivor communicator rendezvous"); + bool progressed = false; + ReceivedMessage receivedReady{}; + bool haveReady = takePending(context.pendingReady, MPI_ANY_SOURCE, key, rendezvousId, receivedReady); + MPI_Status status{}; + if (!haveReady && probeMessage(MPI_ANY_SOURCE, tags.ready, mpiComm, status)) + { + progressed = true; + receivedReady = {controlComm.worldRank(status.MPI_SOURCE), + receiveMessage(status, tags.ready, mpiComm, deadline, "receiving NCCL READY")}; + haveReady = isCurrentAttempt(receivedReady.message, key, rendezvousId); + if (!haveReady && shouldStashForLater(receivedReady.message, key, rendezvousId)) + { + stashPending(context.pendingReady, receivedReady); + } + } + if (haveReady) + { + progressed = true; + ReadyMessage const& ready = receivedReady.message; + int const source = receivedReady.source; + if (ready.version == kProtocolVersion && isCurrentAttempt(ready, key, rendezvousId) + && isActivePeer(activeRanks, source, root)) + { + auto& tokens = issuedTokens[source]; + if (std::find(tokens.begin(), tokens.end(), ready.clientToken) == tokens.end()) + { + TLLM_CHECK_WITH_INFO(tokens.size() < kMaxTokensPerPeer, + "NCCL error: too many stale communicator rendezvous attempts from rank %d", source); + tokens.push_back(ready.clientToken); + } + IdMessage const proposal{ + kProtocolVersion, 0, key, rendezvousId, ready.clientToken, serverToken, digest, id}; + sendMessage( + proposal, controlComm.commRank(source), tags.id, mpiComm, deadline, "sending NCCL unique ID"); + } + } + + ReceivedMessage receivedAck{}; + bool haveAck = takePending(context.pendingAcks, MPI_ANY_SOURCE, key, rendezvousId, receivedAck); + if (!haveAck && probeMessage(MPI_ANY_SOURCE, tags.ack, mpiComm, status)) + { + progressed = true; + receivedAck = {controlComm.worldRank(status.MPI_SOURCE), + receiveMessage(status, tags.ack, mpiComm, deadline, "receiving NCCL ACK")}; + haveAck = isCurrentAttempt(receivedAck.message, key, rendezvousId); + if (!haveAck && shouldStashForLater(receivedAck.message, key, rendezvousId)) + { + stashPending(context.pendingAcks, receivedAck); + } + } + if (haveAck) + { + progressed = true; + AckMessage const& ack = receivedAck.message; + int const source = receivedAck.source; + auto const issued = issuedTokens.find(source); + if (ack.version == kProtocolVersion && isCurrentAttempt(ack, key, rendezvousId) + && ack.serverToken == serverToken && ack.idDigest == digest && issued != issuedTokens.end() + && std::find(issued->second.begin(), issued->second.end(), ack.clientToken) != issued->second.end()) + { + acknowledged.insert(source); + } + } + + if (!progressed) + { + TLLM_CHECK_WITH_INFO( + Clock::now() < deadline, "NCCL error: timed out waiting for survivor communicator rendezvous"); + std::this_thread::sleep_for(kPollInterval); + } + } + return id; + } + + AttemptToken const clientToken = makeAttemptToken(); + ReadyMessage const ready{kProtocolVersion, 0, key, rendezvousId, clientToken}; + int const rootCommRank = controlComm.commRank(root); + sendMessage(ready, rootCommRank, tags.ready, mpiComm, deadline, "sending NCCL READY"); + auto nextReadyRefresh = Clock::now() + kReadyRefreshInterval; + while (true) + { + TLLM_CHECK_WITH_INFO(Clock::now() < deadline, "NCCL error: timed out waiting for survivor communicator ID"); + if (Clock::now() >= nextReadyRefresh) + { + // A root attempt can consume READY and fail before accepting our + // ACK. Refresh the same idempotent client token so a subsequent + // root attempt can converge without waiting for this call to end. + sendMessage(ready, rootCommRank, tags.ready, mpiComm, deadline, "refreshing NCCL READY"); + nextReadyRefresh = Clock::now() + kReadyRefreshInterval; + } + ReceivedMessage receivedProposal{}; + bool haveProposal = takePending(context.pendingIds, root, key, rendezvousId, receivedProposal); + MPI_Status status{}; + if (!haveProposal && probeMessage(rootCommRank, tags.id, mpiComm, status)) + { + receivedProposal = {controlComm.worldRank(status.MPI_SOURCE), + receiveMessage(status, tags.id, mpiComm, deadline, "receiving NCCL unique ID")}; + haveProposal = isCurrentAttempt(receivedProposal.message, key, rendezvousId); + if (!haveProposal && shouldStashForLater(receivedProposal.message, key, rendezvousId)) + { + stashPending(context.pendingIds, receivedProposal); + } + } + if (!haveProposal) + { + TLLM_CHECK_WITH_INFO(Clock::now() < deadline, "NCCL error: timed out waiting for survivor communicator ID"); + std::this_thread::sleep_for(kPollInterval); + continue; + } + + IdMessage const& proposal = receivedProposal.message; + if (proposal.version != kProtocolVersion || !isCurrentAttempt(proposal, key, rendezvousId) + || !(proposal.clientToken == clientToken) || proposal.idDigest != fnv1a(&proposal.id, sizeof(proposal.id))) + { + sendMessage(ready, rootCommRank, tags.ready, mpiComm, deadline, "refreshing NCCL READY"); + nextReadyRefresh = Clock::now() + kReadyRefreshInterval; + continue; + } + + AckMessage const ack{ + kProtocolVersion, 0, key, rendezvousId, clientToken, proposal.serverToken, proposal.idDigest}; + sendMessage(ack, rootCommRank, tags.ack, mpiComm, deadline, "sending NCCL ACK"); + return proposal.id; + } +} + +} // namespace tensorrt_llm::runtime + +#endif // ENABLE_MULTI_DEVICE diff --git a/cpp/tensorrt_llm/thop/allgatherOp.cpp b/cpp/tensorrt_llm/thop/allgatherOp.cpp index 0d92aa966901..961aa8f7befa 100644 --- a/cpp/tensorrt_llm/thop/allgatherOp.cpp +++ b/cpp/tensorrt_llm/thop/allgatherOp.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -56,9 +56,9 @@ class AllgatherOp int initialize() { - TLLM_LOG_TRACE("%s start for rank %d", __PRETTY_FUNCTION__, COMM_SESSION.getRank()); + TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); mNcclComm = getComm(mGroup); - TLLM_LOG_TRACE("%s stop for rank %d", __PRETTY_FUNCTION__, COMM_SESSION.getRank()); + TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); return 0; } @@ -71,12 +71,11 @@ class AllgatherOp int64_t sum_sizes = sizes.has_value() ? std::accumulate(sizes.value().begin(), sizes.value().end(), 0, std::plus<>{}) : 0; std::vector output_list; + std::vector streams; + bool const trackOperations = isNcclFaultToleranceEnabled(); output_list.reserve(input_list.size()); - ncclGroupStart(); for (auto const& input : input_list) { - auto stream = at::cuda::getCurrentCUDAStream(input.get_device()); - auto type = tensorrt_llm::runtime::TorchUtils::dataType(input.scalar_type()); std::vector outputShape = input.sizes().vec(); if (sizes.has_value()) { @@ -86,29 +85,68 @@ class AllgatherOp { outputShape[0] *= mGroup.size(); } - auto output = torch::empty(outputShape, input.options()); + output_list.push_back(torch::empty(outputShape, input.options())); + if (trackOperations) + { + auto const stream = at::cuda::getCurrentCUDAStream(input.get_device()); + if (std::find(streams.begin(), streams.end(), stream) == streams.end()) + { + streams.push_back(stream); + } + } + } + + auto commLease = acquireComm(mNcclComm); + auto const comm = commLease.get(); + std::vector watchdogTokens; + if (trackOperations) + { + watchdogTokens.reserve(streams.size()); + for (auto const stream : streams) + { + watchdogTokens.push_back(commLease.begin(stream, "ncclAllGather")); + } + } + commLease.groupStart("ncclGroupStart(allgather)"); + for (size_t index = 0; index < input_list.size(); ++index) + { + auto const& input = input_list[index]; + auto& output = output_list[index]; + auto stream = at::cuda::getCurrentCUDAStream(input.get_device()); + auto type = tensorrt_llm::runtime::TorchUtils::dataType(input.scalar_type()); if (use_nccl_allgather) { - ncclAllGather(input.data_ptr(), output.mutable_data_ptr(), input.numel(), (*getDtypeMap())[type], - *mNcclComm, stream); + commLease.checkEnqueue(ncclAllGather(input.data_ptr(), output.mutable_data_ptr(), input.numel(), + (*getDtypeMap())[type], comm, stream), + "ncclAllGather"); } else { + auto const& outputShape = output.sizes(); size_t numel_base = std::accumulate(outputShape.cbegin() + 1, outputShape.cend(), 1, std::multiplies<>{}); int64_t split_offset = 0; for (int root = 0; root < static_cast(mGroup.size()); ++root) { auto split_size = sizes.value()[root]; - ncclBroadcast(input.data_ptr(), - output.index({torch::indexing::Slice(split_offset, torch::indexing::None)}).mutable_data_ptr(), - numel_base * split_size, (*getDtypeMap())[type], root, *mNcclComm, stream); + commLease.checkEnqueue( + ncclBroadcast(input.data_ptr(), + output.index({torch::indexing::Slice(split_offset, torch::indexing::None)}) + .mutable_data_ptr(), + numel_base * split_size, (*getDtypeMap())[type], root, comm, stream), + "ncclBroadcast(allgather)"); split_offset += split_size; } } - output_list.push_back(output); } - NCCLCHECK_THROW(ncclGroupEnd()); + commLease.groupEnd("ncclGroupEnd(allgather)"); + if (trackOperations) + { + for (size_t index = 0; index < streams.size(); ++index) + { + commLease.track(watchdogTokens[index], streams[index]); + } + } return output_list; } diff --git a/cpp/tensorrt_llm/thop/allreduceOp.cpp b/cpp/tensorrt_llm/thop/allreduceOp.cpp index 9f4de156ea55..3ab92117e6cf 100644 --- a/cpp/tensorrt_llm/thop/allreduceOp.cpp +++ b/cpp/tensorrt_llm/thop/allreduceOp.cpp @@ -265,9 +265,13 @@ class AllreduceOp int getRank() const { - return std::visit( - overloaded{[&](std::shared_ptr const&) { return COMM_SESSION.getRank(); }, - [&](c10::intrusive_ptr const& torchPg) { return get_world_pg()->getRank(); }}, + return std::visit(overloaded{[&](std::shared_ptr const&) + { + TLLM_CHECK_WITH_INFO( + mRank >= 0, "Raw AllreduceOp must be initialized before querying its rank"); + return mRank; + }, + [&](c10::intrusive_ptr const&) { return get_world_pg()->getRank(); }}, mNcclComm); } @@ -284,6 +288,12 @@ class AllreduceOp AllReduceStrategyType runtime_strategy = selectImplementation(seq_len, hidden_size); + if (mNcclComm.index() == 0 && isNcclFaultToleranceEnabled()) + { + TLLM_CHECK_WITH_INFO(runtime_strategy == AllReduceStrategyType::NCCL, + "NCCL error: TLLM_FAULT_TOLERANCE_MODE=1 requires a watchdog-managed NCCL allreduce strategy"); + } + // Log runtime strategy auto const rank = getRank(); TLLM_LOG_DEBUG( @@ -308,10 +318,14 @@ class AllreduceOp int initialize() { - TLLM_LOG_TRACE("%s start for rank %d", __PRETTY_FUNCTION__, getRank()); + int const initialRank = mNcclComm.index() == 0 ? -1 : getRank(); + TLLM_LOG_TRACE("%s start for rank %d", __PRETTY_FUNCTION__, initialRank); if (mNcclComm.index() == 0) { + TLLM_CHECK_WITH_INFO(!isNcclFaultToleranceEnabled() || mStrategy == AllReduceStrategyType::NCCL, + "NCCL error: TLLM_FAULT_TOLERANCE_MODE=1 requires a watchdog-managed NCCL allreduce strategy"); mNcclComm = getComm(mGroup); + mRank = getCommWorldRank(std::get<0>(mNcclComm)); } if (mStrategy != AllReduceStrategyType::NCCL && mStrategy != AllReduceStrategyType::UB) { @@ -415,8 +429,12 @@ class AllreduceOp auto stream = at::cuda::getCurrentCUDAStream(input.get_device()); int size = input.numel(); reduce_output = torch::empty_like(input); - NCCLCHECK_THROW(ncclAllReduce(input.data_ptr(), reduce_output.mutable_data_ptr(), size, - (*getDtypeMap())[mType], ncclSum, *rawComm, stream)); + auto commLease = acquireComm(rawComm); + auto const watchdogToken = commLease.begin(stream, "ncclAllReduce"); + commLease.check(ncclAllReduce(input.data_ptr(), reduce_output.mutable_data_ptr(), size, + (*getDtypeMap())[mType], ncclSum, commLease.get(), stream), + "ncclAllReduce"); + commLease.track(watchdogToken, stream); }, [&](c10::intrusive_ptr& torchPg) { @@ -462,8 +480,6 @@ class AllreduceOp // From here on, we have a raw NCCL comm - can proceed with window registration auto rawComm = std::get<0>(mNcclComm); - ncclComm_t comm = *rawComm; - TLLM_CHECK_WITH_INFO(comm != nullptr, "NCCL communicator is null"); TLLM_LOG_DEBUG("[runNCCLAllReduceSymmetric] Using raw NCCL comm path (not ProcessGroup)"); using tensorrt_llm::common::nccl_util::NCCLWindowAllocator; @@ -486,7 +502,13 @@ class AllreduceOp double const a = -4986.43478503; double const b = 156716.52177552; int nRanks; - NCCLCHECK_THROW(ncclCommCount(comm, &nRanks)); + ncclComm_t comm = nullptr; + { + auto commLease = acquireComm(rawComm); + comm = commLease.get(); + TLLM_CHECK_WITH_INFO(comm != nullptr, "NCCL communicator is null"); + commLease.check(ncclCommCount(comm, &nRanks), "ncclCommCount"); + } size_t minRegistrationThreshold = static_cast(std::max(0.0, a * nRanks + b)) * input.element_size(); // Disable window registration if neither NVLink nor MNNVL is supported // TODO replace in NCCL 2.29 with comm query @@ -566,7 +588,16 @@ class AllreduceOp } // Perform allreduce - NCCLCHECK_THROW(ncclAllReduce(inputPtr, outputPtr, size, (*getDtypeMap())[mType], ncclSum, comm, stream)); + { + auto commLease = acquireComm(rawComm); + auto const activeComm = commLease.get(); + TLLM_CHECK_WITH_INFO(activeComm == comm, "NCCL communicator changed while preparing a symmetric allreduce"); + auto const watchdogToken = commLease.begin(stream, "ncclAllReduce(symmetric)"); + commLease.check( + ncclAllReduce(inputPtr, outputPtr, size, (*getDtypeMap())[mType], ncclSum, activeComm, stream), + "ncclAllReduce(symmetric)"); + commLease.track(watchdogToken, stream); + } if (mOp == AllReduceFusionOp::NONE) { @@ -1466,6 +1497,7 @@ class AllreduceOp AllReduceFusionOp mOp; float mEps; std::variant, c10::intrusive_ptr> mNcclComm; + int mRank{-1}; }; } // namespace @@ -1488,7 +1520,7 @@ void preallocateNCCLWindowBuffer( } auto const commPtr = getComm(groupSet); - if (!commPtr || *commPtr == nullptr) + if (!commPtr) { TLLM_LOG_DEBUG("[preallocateNCCLWindowBuffers] NCCL comm is null; skipping preallocation"); return; @@ -1496,7 +1528,11 @@ void preallocateNCCLWindowBuffer( using tensorrt_llm::common::nccl_util::NCCLWindowAllocator; auto& allocator = NCCLWindowAllocator::getInstance(); - ncclComm_t const comm = *commPtr; + ncclComm_t comm = nullptr; + { + auto commLease = acquireComm(commPtr); + comm = commLease.get(); + } int64_t const numTokens = input.size(0); int64_t const elementsPerToken = input.numel() / numTokens; @@ -1518,7 +1554,7 @@ void preallocateNCCLWindowBuffer( { for (int64_t i = 0; i < buffersPerSize; ++i) { - auto buffer = allocator.requestBuffer(comm, bufferSize); + auto buffer = allocator.requestBuffer(commPtr, bufferSize); if (!buffer.isValid()) { break; @@ -1528,6 +1564,14 @@ void preallocateNCCLWindowBuffer( } catch (std::exception const& e) { + // A local allocation failure is a best-effort preallocation miss. A + // communicator failure is not: preserve the classifier-friendly NCCL + // exception so the executor enters recovery instead of continuing on + // the aborted native handle. + if (getCommAsyncError(commPtr).has_value()) + { + throw; + } TLLM_LOG_DEBUG("[preallocateNCCLWindowBuffer] requestBuffer failed for %zu bytes: %s", bufferSize, e.what()); } diff --git a/cpp/tensorrt_llm/thop/alltoallOp.cpp b/cpp/tensorrt_llm/thop/alltoallOp.cpp index 8a775b1cf6b8..8ae38fcf1ba4 100644 --- a/cpp/tensorrt_llm/thop/alltoallOp.cpp +++ b/cpp/tensorrt_llm/thop/alltoallOp.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -44,10 +44,10 @@ class AllToAllHelixOp int initialize() { - TLLM_LOG_TRACE("%s start for rank %d", __PRETTY_FUNCTION__, COMM_SESSION.getRank()); + TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); mNcclComm = getComm(mGroup); TLLM_CHECK_WITH_INFO(mGroup.size() > 0, "group size should be greater than 0"); - TLLM_LOG_TRACE("%s stop for rank %d", __PRETTY_FUNCTION__, COMM_SESSION.getRank()); + TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); return 0; } @@ -64,25 +64,37 @@ class AllToAllHelixOp TORCH_CHECK(input.is_contiguous(), "input must be contiguous"); } std::vector output_list(static_cast(num_lists_)); - auto stream = at::cuda::getCurrentCUDAStream(input_list[0].get_device()); - ncclGroupStart(); for (int il = 0; il < num_lists_; ++il) { - auto off = il * num_ranks; + auto const off = il * num_ranks; auto output_shape = input_list[off].sizes().vec(); output_shape.insert(output_shape.begin(), num_ranks); - auto output = torch::empty(output_shape, input_list[off].options()); - output_list[il] = output; + output_list[il] = torch::empty(output_shape, input_list[off].options()); + } + + auto stream = at::cuda::getCurrentCUDAStream(input_list[0].get_device()); + auto commLease = acquireComm(mNcclComm); + auto const comm = commLease.get(); + auto const watchdogToken = commLease.begin(stream, "ncclSend/ncclRecv(alltoall)"); + commLease.groupStart("ncclGroupStart(alltoall)"); + for (int il = 0; il < num_lists_; ++il) + { + auto off = il * num_ranks; + auto& output = output_list[il]; auto type = tensorrt_llm::runtime::TorchUtils::dataType(input_list[off].scalar_type()); auto nccl_type = (*getDtypeMap())[type]; for (int r = 0; r < num_ranks; ++r) { auto const& input = input_list[off + r]; - ncclSend(input.data_ptr(), input.numel(), nccl_type, r, *mNcclComm, stream); - ncclRecv(output[r].mutable_data_ptr(), output[r].numel(), nccl_type, r, *mNcclComm, stream); + commLease.checkEnqueue( + ncclSend(input.data_ptr(), input.numel(), nccl_type, r, comm, stream), "ncclSend(alltoall)"); + commLease.checkEnqueue( + ncclRecv(output[r].mutable_data_ptr(), output[r].numel(), nccl_type, r, comm, stream), + "ncclRecv(alltoall)"); } } - NCCLCHECK_THROW(ncclGroupEnd()); + commLease.groupEnd("ncclGroupEnd(alltoall)"); + commLease.track(watchdogToken, stream); return output_list; } diff --git a/cpp/tensorrt_llm/thop/ncclCommunicatorOp.cpp b/cpp/tensorrt_llm/thop/ncclCommunicatorOp.cpp index a45fa955a00b..d588f61e557a 100644 --- a/cpp/tensorrt_llm/thop/ncclCommunicatorOp.cpp +++ b/cpp/tensorrt_llm/thop/ncclCommunicatorOp.cpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -16,19 +16,85 @@ #include "tensorrt_llm/thop/ncclCommunicatorOp.h" +#include "tensorrt_llm/common/opUtils.h" #include "tensorrt_llm/runtime/iBuffer.h" +#include +#include +#include +#include +#include + namespace tr = tensorrt_llm::runtime; TRTLLM_NAMESPACE_BEGIN namespace torch_ext { +namespace +{ + +std::mutex& pipelineCommunicatorMutex() +{ + static auto* mutex = new std::mutex; + return *mutex; +} + +struct PipelineCommunicatorCache +{ + int64_t worldSize{-1}; + int64_t rank{-1}; + std::shared_ptr communicator; +}; + +PipelineCommunicatorCache& pipelineCommunicatorCache() +{ + // Python model-engine reference counts are process-local and therefore + // cannot define a cross-rank NCCL teardown boundary. Keep the FT PP + // communicator alive for the process lifetime so rank-skewed wrapper churn + // never makes one rank reuse an old communicator while a peer enters a + // fresh-ID rendezvous. Coordinated teardown belongs to the Phase-2 owner. + static auto* cache = new PipelineCommunicatorCache; + return *cache; +} + +} // namespace NcclCommunicatorOp::NcclCommunicatorOp(int64_t worldSize, int64_t rank) : mRank(static_cast(rank)) { - mPipelineComm = std::make_shared(worldSize, rank); + TLLM_CHECK_WITH_INFO(worldSize > 0 && worldSize <= std::numeric_limits::max(), + "NCCL error: world size is out of range: %lld", static_cast(worldSize)); + TLLM_CHECK_WITH_INFO(rank >= 0 && rank < worldSize && rank <= std::numeric_limits::max(), + "NCCL error: rank is out of range: %lld", static_cast(rank)); + + if (!mpi::isFaultToleranceModeEnabled()) + { + // The legacy MPI-broadcast bootstrap assigns independent NCCL IDs, so + // multiple PP communicators remain valid when fault tolerance is off. + mPipelineComm = std::make_shared( + static_cast(worldSize), static_cast(rank)); + return; + } + + std::lock_guard const lock(pipelineCommunicatorMutex()); + auto& cache = pipelineCommunicatorCache(); + if (cache.communicator == nullptr) + { + cache.communicator = std::make_shared( + static_cast(worldSize), static_cast(rank)); + cache.worldSize = worldSize; + cache.rank = rank; + } + else + { + TLLM_CHECK_WITH_INFO(cache.worldSize == worldSize && cache.rank == rank, + "NCCL error: the process-lifetime PP communicator was initialized for world size %lld rank %lld, " + "not world size %lld rank %lld", + static_cast(cache.worldSize), static_cast(cache.rank), + static_cast(worldSize), static_cast(rank)); + } + mPipelineComm = cache.communicator; } void NcclCommunicatorOp::send(th::Tensor tensor, int64_t toRank) const @@ -49,6 +115,91 @@ void NcclCommunicatorOp::recv(th::Tensor& tensor, int64_t fromRank) const mPipelineComm->receive(*tr::IBuffer::wrap(ptr, size), static_cast(fromRank), cudaStream); } +void NcclCommunicatorOp::abort() +{ + mPipelineComm->abort(); +} + +void NcclCommunicatorOp::abortAndReinit(std::vector const& activeRanks, int64_t rendezvousId) +{ + TLLM_CHECK_WITH_INFO(rendezvousId > 0, "NCCL error: recovery rendezvous ID must be positive, got %lld", + static_cast(rendezvousId)); + std::vector ranks; + ranks.reserve(activeRanks.size()); + for (int64_t const rank : activeRanks) + { + TLLM_CHECK_WITH_INFO(rank >= 0 && rank <= std::numeric_limits::max(), + "NCCL error: active rank is out of range: %lld", static_cast(rank)); + ranks.push_back(static_cast(rank)); + } + mPipelineComm->abortAndReinit(ranks, static_cast(rendezvousId)); +} + +std::string NcclCommunicatorOp::getAsyncError() const +{ + return mPipelineComm->getAsyncError(); +} + +std::vector NcclCommunicatorOp::getActiveRanks() const +{ + auto const ranks = mPipelineComm->getActiveRanks(); + return std::vector(ranks.begin(), ranks.end()); +} + +void ncclCommAbortAndReinit( + std::vector const& oldGroup, std::vector const& activeGroup, int64_t rendezvousId) +{ +#if ENABLE_MULTI_DEVICE + TLLM_CHECK_WITH_INFO(rendezvousId > 0, "NCCL error: recovery rendezvous ID must be positive, got %lld", + static_cast(rendezvousId)); + std::set oldRanks; + std::set activeRanks; + for (auto const rank : oldGroup) + { + TLLM_CHECK_WITH_INFO(rank >= 0 && rank <= std::numeric_limits::max(), + "NCCL error: old-group rank is out of range: %lld", static_cast(rank)); + oldRanks.insert(static_cast(rank)); + } + for (auto const rank : activeGroup) + { + TLLM_CHECK_WITH_INFO(rank >= 0 && rank <= std::numeric_limits::max(), + "NCCL error: active-group rank is out of range: %lld", static_cast(rank)); + activeRanks.insert(static_cast(rank)); + } + TLLM_CHECK_WITH_INFO(oldRanks.size() == oldGroup.size(), "NCCL error: old group contains duplicate ranks"); + TLLM_CHECK_WITH_INFO(activeRanks.size() == activeGroup.size(), "NCCL error: active group contains duplicate ranks"); + abortAndReinitComm(oldRanks, activeRanks, static_cast(rendezvousId)); +#else + (void) oldGroup; + (void) activeGroup; + (void) rendezvousId; + TLLM_THROW("NCCL error: multi device support is disabled."); +#endif +} + +c10::optional ncclCommGetAsyncError(std::vector const& group) +{ +#if ENABLE_MULTI_DEVICE + std::set ranks; + for (auto const rank : group) + { + TLLM_CHECK_WITH_INFO(rank >= 0 && rank <= std::numeric_limits::max(), + "NCCL error: group rank is out of range: %lld", static_cast(rank)); + ranks.insert(static_cast(rank)); + } + TLLM_CHECK_WITH_INFO(ranks.size() == group.size(), "NCCL error: group contains duplicate ranks"); + auto error = getCommAsyncError(ranks); + if (error.has_value()) + { + return *error; + } + return c10::nullopt; +#else + (void) group; + return c10::nullopt; +#endif +} + } // namespace torch_ext TRTLLM_NAMESPACE_END @@ -57,4 +208,20 @@ static auto trtllmNcclCommunicator = torch::jit::class_("trtllm", "NcclCommunicatorOp") .def(torch::jit::init()) .def("send", &tensorrt_llm::torch_ext::NcclCommunicatorOp::send) - .def("recv", &tensorrt_llm::torch_ext::NcclCommunicatorOp::recv); + .def("recv", &tensorrt_llm::torch_ext::NcclCommunicatorOp::recv) + .def("abort", &tensorrt_llm::torch_ext::NcclCommunicatorOp::abort) + .def("abort_and_reinit", &tensorrt_llm::torch_ext::NcclCommunicatorOp::abortAndReinit) + .def("get_async_error", &tensorrt_llm::torch_ext::NcclCommunicatorOp::getAsyncError) + .def("get_active_ranks", &tensorrt_llm::torch_ext::NcclCommunicatorOp::getActiveRanks); + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def("nccl_comm_abort_and_reinit(int[] old_group, int[] active_group, int rendezvous_id) -> ()"); + m.def("nccl_comm_get_async_error(int[] group) -> str?"); +} + +TORCH_LIBRARY_IMPL(trtllm, CompositeExplicitAutograd, m) +{ + m.impl("nccl_comm_abort_and_reinit", &tensorrt_llm::torch_ext::ncclCommAbortAndReinit); + m.impl("nccl_comm_get_async_error", &tensorrt_llm::torch_ext::ncclCommGetAsyncError); +} diff --git a/cpp/tensorrt_llm/thop/ncclCommunicatorOp.h b/cpp/tensorrt_llm/thop/ncclCommunicatorOp.h index 38f4d215ac1a..a19378c1482c 100644 --- a/cpp/tensorrt_llm/thop/ncclCommunicatorOp.h +++ b/cpp/tensorrt_llm/thop/ncclCommunicatorOp.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -18,7 +18,10 @@ #include "tensorrt_llm/common/config.h" #include "tensorrt_llm/runtime/ncclCommunicator.h" #include "tensorrt_llm/thop/thUtils.h" +#include #include +#include +#include namespace th = torch; @@ -34,6 +37,10 @@ class NcclCommunicatorOp : public th::jit::CustomClassHolder void send(th::Tensor tensor, int64_t toRank) const; void recv(th::Tensor& tensor, int64_t fromRank) const; + void abort(); + void abortAndReinit(std::vector const& activeRanks, int64_t rendezvousId); + [[nodiscard]] std::string getAsyncError() const; + [[nodiscard]] std::vector getActiveRanks() const; private: int32_t mRank; diff --git a/cpp/tensorrt_llm/thop/outputTensor.h b/cpp/tensorrt_llm/thop/outputTensor.h index d9cde554852f..eac594652251 100644 --- a/cpp/tensorrt_llm/thop/outputTensor.h +++ b/cpp/tensorrt_llm/thop/outputTensor.h @@ -70,7 +70,7 @@ inline std::pair allocate_output(std::vector co { TLLM_LOG_DEBUG("[allocate_output] getComm threw (MPI disabled?): %s; fallback to default", e.what()); } - if (commPtr && *commPtr != nullptr) + if (commPtr) { auto [tensor, buffer] = tensorrt_llm::common::nccl_util::createNCCLWindowTensor(commPtr, output_size, dtype); diff --git a/cpp/tensorrt_llm/thop/reducescatterOp.cpp b/cpp/tensorrt_llm/thop/reducescatterOp.cpp index 40f89e40ff75..762a23efd1af 100644 --- a/cpp/tensorrt_llm/thop/reducescatterOp.cpp +++ b/cpp/tensorrt_llm/thop/reducescatterOp.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -28,7 +28,11 @@ #include #endif // ENABLE_MULTI_DEVICE +#include #include +#include +#include +#include #include #include @@ -57,7 +61,11 @@ class ReducescatterOp { TLLM_LOG_TRACE("%s start for rank %d", __PRETTY_FUNCTION__, -1); mNcclComm = getComm(mGroup); - TLLM_LOG_TRACE("%s stop for rank %d", __PRETTY_FUNCTION__, -1); + auto const rank = getCommWorldRank(mNcclComm); + auto const rankIt = mGroup.find(rank); + TLLM_CHECK_WITH_INFO(rankIt != mGroup.end(), "Global rank %d is not in the reduce-scatter group", rank); + mGroupRank = static_cast(std::distance(mGroup.begin(), rankIt)); + TLLM_LOG_TRACE("%s stop for rank %d", __PRETTY_FUNCTION__, rank); return 0; } @@ -67,25 +75,14 @@ class ReducescatterOp bool use_nccl_reducescatter = !sizes.has_value() || std::all_of(sizes.value().begin(), sizes.value().end(), [&sizes](int64_t size) { return size == sizes.value()[0]; }); - int groupRank = 0; - if (sizes.has_value()) - { - auto rank = COMM_SESSION.getRank(); - for (auto const& currentRank : mGroup) - { - if (rank == currentRank) - break; - ++groupRank; - } - TLLM_CHECK(static_cast(groupRank) < mGroup.size()); - } + int const groupRank = sizes.has_value() ? mGroupRank : 0; + TLLM_CHECK_WITH_INFO(!sizes.has_value() || groupRank >= 0, "ReducescatterOp must be initialized before use"); std::vector output_list; + std::vector streams; + bool const trackOperations = isNcclFaultToleranceEnabled(); output_list.reserve(input_list.size()); - ncclGroupStart(); for (auto const& input : input_list) { - auto stream = at::cuda::getCurrentCUDAStream(input.get_device()); - auto type = tensorrt_llm::runtime::TorchUtils::dataType(input.scalar_type()); std::vector outputShape = input.sizes().vec(); if (sizes.has_value()) { @@ -95,29 +92,68 @@ class ReducescatterOp { outputShape[0] = outputShape[0] / mGroup.size(); } - auto output = torch::empty(outputShape, input.options()); + output_list.push_back(torch::empty(outputShape, input.options())); + if (trackOperations) + { + auto const stream = at::cuda::getCurrentCUDAStream(input.get_device()); + if (std::find(streams.begin(), streams.end(), stream) == streams.end()) + { + streams.push_back(stream); + } + } + } + + auto commLease = acquireComm(mNcclComm); + auto const comm = commLease.get(); + std::vector watchdogTokens; + if (trackOperations) + { + watchdogTokens.reserve(streams.size()); + for (auto const stream : streams) + { + watchdogTokens.push_back(commLease.begin(stream, "ncclReduceScatter")); + } + } + commLease.groupStart("ncclGroupStart(reducescatter)"); + for (size_t index = 0; index < input_list.size(); ++index) + { + auto const& input = input_list[index]; + auto& output = output_list[index]; + auto stream = at::cuda::getCurrentCUDAStream(input.get_device()); + auto type = tensorrt_llm::runtime::TorchUtils::dataType(input.scalar_type()); if (use_nccl_reducescatter) { - ncclReduceScatter(input.data_ptr(), output.mutable_data_ptr(), output.numel(), (*getDtypeMap())[type], - ncclSum, *mNcclComm, stream); + commLease.checkEnqueue(ncclReduceScatter(input.data_ptr(), output.mutable_data_ptr(), output.numel(), + (*getDtypeMap())[type], ncclSum, comm, stream), + "ncclReduceScatter"); } else { + auto const& outputShape = output.sizes(); size_t numel_base = std::accumulate(outputShape.cbegin() + 1, outputShape.cend(), 1, std::multiplies<>{}); int64_t split_offset = 0; for (int root = 0; root < static_cast(mGroup.size()); ++root) { auto split_size = sizes.value()[root]; - ncclReduce(input.index({torch::indexing::Slice(split_offset, torch::indexing::None)}).data_ptr(), - output.mutable_data_ptr(), numel_base * split_size, (*getDtypeMap())[type], ncclSum, root, - *mNcclComm, stream); + commLease.checkEnqueue( + ncclReduce( + input.index({torch::indexing::Slice(split_offset, torch::indexing::None)}).data_ptr(), + output.mutable_data_ptr(), numel_base * split_size, (*getDtypeMap())[type], ncclSum, root, + comm, stream), + "ncclReduce(reducescatter)"); split_offset += split_size; } } - output_list.push_back(output); } - NCCLCHECK_THROW(ncclGroupEnd()); + commLease.groupEnd("ncclGroupEnd(reducescatter)"); + if (trackOperations) + { + for (size_t index = 0; index < streams.size(); ++index) + { + commLease.track(watchdogTokens[index], streams[index]); + } + } return output_list; } @@ -129,6 +165,7 @@ class ReducescatterOp private: std::set mGroup; std::shared_ptr mNcclComm; + int mGroupRank{-1}; }; class ReducescatterPgOp diff --git a/cpp/tests/unit_tests/multi_gpu/CMakeLists.txt b/cpp/tests/unit_tests/multi_gpu/CMakeLists.txt index 669456551ed2..9787b650e066 100644 --- a/cpp/tests/unit_tests/multi_gpu/CMakeLists.txt +++ b/cpp/tests/unit_tests/multi_gpu/CMakeLists.txt @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & +# SPDX-FileCopyrightText: Copyright (c) 2023-2026 NVIDIA CORPORATION & # AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); you may not @@ -19,6 +19,133 @@ add_gtest(cacheTransceiverTest cacheTransceiverTest.cpp) target_link_libraries(cacheTransceiverTest PRIVATE ${Python3_LIBRARIES}) add_gtest(mpiUtilsTest mpiUtilsTest.cpp) +add_gtest(ncclUniqueIdRendezvousTest ncclUniqueIdRendezvousTest.cpp) +if(ENABLE_MULTI_DEVICE) + set(NCCL_RENDEZVOUS_RAW_RECOVERY_TEST + NcclUniqueIdRendezvousTest.RawCommunicatorRebuildRunsPostRecoveryCollective + ) + set(NCCL_RENDEZVOUS_PP_RECOVERY_TEST + NcclUniqueIdRendezvousTest.PipelineCommunicatorRebuildRunsPostRecoverySendRecv + ) + set(NCCL_RENDEZVOUS_PROCESS_DEATH_TEST + NcclUniqueIdRendezvousTest.ProcessDeathRendezvousRequiresOptInFaultTolerantMpiLauncher + ) + set(NCCL_RENDEZVOUS_ISOLATED_TESTS + "${NCCL_RENDEZVOUS_RAW_RECOVERY_TEST}:${NCCL_RENDEZVOUS_PP_RECOVERY_TEST}:${NCCL_RENDEZVOUS_PROCESS_DEATH_TEST}" + ) + add_test( + NAME ncclUniqueIdRendezvousTest.mpi3 + COMMAND + ${MPIEXEC_EXECUTABLE} ${MPIEXEC_NUMPROC_FLAG} 3 ${MPIEXEC_PREFLAGS} + $ + --gtest_filter=NcclUniqueIdRendezvousTest.*-${NCCL_RENDEZVOUS_ISOLATED_TESTS} + ${MPIEXEC_POSTFLAGS}) + set_tests_properties( + ncclUniqueIdRendezvousTest.mpi3 + PROPERTIES ENVIRONMENT + "CUDA_MODULE_LOADING=LAZY;TLLM_FAULT_TOLERANCE_MODE=1" + PROCESSORS 3 TIMEOUT 180) + + # Raw communicator replacement mutates process-global NCCL state and needs one + # distinct physical GPU per MPI rank. Keep it in its own CTest process. + add_test( + NAME ncclUniqueIdRendezvousTest.rawRecoveryMpi3 + COMMAND + ${MPIEXEC_EXECUTABLE} ${MPIEXEC_NUMPROC_FLAG} 3 ${MPIEXEC_PREFLAGS} + $ + --gtest_filter=${NCCL_RENDEZVOUS_RAW_RECOVERY_TEST} ${MPIEXEC_POSTFLAGS}) + set_tests_properties( + ncclUniqueIdRendezvousTest.rawRecoveryMpi3 + PROPERTIES ENVIRONMENT + "CUDA_MODULE_LOADING=LAZY;TLLM_FAULT_TOLERANCE_MODE=1" + PROCESSORS 3 TIMEOUT 180) + + # The runtime PP wrapper owns a distinct communicator implementation and + # rendezvous tag set. Exercise its real post-recovery send/receive path in a + # separate process so it cannot share native singleton state with raw ops. + add_test( + NAME ncclUniqueIdRendezvousTest.ppRecoveryMpi3 + COMMAND + ${MPIEXEC_EXECUTABLE} ${MPIEXEC_NUMPROC_FLAG} 3 ${MPIEXEC_PREFLAGS} + $ + --gtest_filter=${NCCL_RENDEZVOUS_PP_RECOVERY_TEST} ${MPIEXEC_POSTFLAGS}) + set_tests_properties( + ncclUniqueIdRendezvousTest.ppRecoveryMpi3 + PROPERTIES ENVIRONMENT + "CUDA_MODULE_LOADING=LAZY;TLLM_FAULT_TOLERANCE_MODE=1" + PROCESSORS 3 TIMEOUT 180) + + option(TRTLLM_ENABLE_NCCL_PROCESS_DEATH_TEST + "Enable the destructive NCCL rendezvous test that exits one MPI rank" + OFF) + set(TRTLLM_NCCL_PROCESS_DEATH_MPIEXEC_PREFLAGS + "" + CACHE + STRING + "Semicolon-separated launcher recovery flags for the NCCL process-death test" + ) + if(TRTLLM_ENABLE_NCCL_PROCESS_DEATH_TEST) + if("${TRTLLM_NCCL_PROCESS_DEATH_MPIEXEC_PREFLAGS}" STREQUAL "") + message( + FATAL_ERROR + "TRTLLM_ENABLE_NCCL_PROCESS_DEATH_TEST requires launcher-specific " + "TRTLLM_NCCL_PROCESS_DEATH_MPIEXEC_PREFLAGS (for example OpenMPI recovery flags)" + ) + endif() + # Keep a launcher-level signal from bypassing the sentinel properties. The + # CMake wrapper translates it to a normal nonzero process result. + set(NCCL_RENDEZVOUS_PROCESS_DEATH_ENV + CUDA_MODULE_LOADING=LAZY + TLLM_FAULT_TOLERANCE_MODE=1 + TRTLLM_TEST_NCCL_PROCESS_DEATH=1 + TRTLLM_NCCL_NONBLOCKING_TIMEOUT_MS=20000 + TRTLLM_NCCL_RECOVERY_TIMEOUT_MS=20000) + add_test( + NAME ncclUniqueIdRendezvousTest.processDeathMpi3 + COMMAND + ${CMAKE_COMMAND} -E env ${MPIEXEC_EXECUTABLE} ${MPIEXEC_NUMPROC_FLAG} 3 + ${MPIEXEC_PREFLAGS} ${TRTLLM_NCCL_PROCESS_DEATH_MPIEXEC_PREFLAGS} + $ + --gtest_filter=${NCCL_RENDEZVOUS_PROCESS_DEATH_TEST} + ${MPIEXEC_POSTFLAGS}) + set_tests_properties( + ncclUniqueIdRendezvousTest.processDeathMpi3 + PROPERTIES ENVIRONMENT + "${NCCL_RENDEZVOUS_PROCESS_DEATH_ENV}" + PROCESSORS + 3 + TIMEOUT + 60 + RUN_SERIAL + TRUE + # A launcher may report failure solely because the injected + # victim intentionally skips MPI_Finalize. The root emits this + # sentinel only after survivor-only communicator rebuild and a + # successful allreduce, so it is the authoritative pass + # condition. + PASS_REGULAR_EXPRESSION + "TRTLLM_PROCESS_DEATH_RENDEZVOUS_OK" + FAIL_REGULAR_EXPRESSION + "TRTLLM_PROCESS_DEATH_RENDEZVOUS_FAIL") + endif() + + add_test( + NAME ncclUniqueIdRendezvousTest.defaultOff + COMMAND + ${MPIEXEC_EXECUTABLE} ${MPIEXEC_NUMPROC_FLAG} 1 ${MPIEXEC_PREFLAGS} + $ + --gtest_filter=NcclUniqueIdRendezvousTest.FaultToleranceDisabledPreservesLegacyOperationPath + ${MPIEXEC_POSTFLAGS}) + set_tests_properties( + ncclUniqueIdRendezvousTest.defaultOff + PROPERTIES + ENVIRONMENT + "CUDA_MODULE_LOADING=LAZY;TLLM_FAULT_TOLERANCE_MODE=0;TRTLLM_TEST_FT_DEFAULT_OFF=1" + PROCESSORS + 1 + TIMEOUT + 60) +endif() add_gtest(userBufferTest userBufferTest.cpp) add_gtest(ncclUtilsTest ncclUtilsTest.cpp) target_link_libraries(ncclUtilsTest PRIVATE ${Python3_LIBRARIES}) diff --git a/cpp/tests/unit_tests/multi_gpu/ncclUniqueIdRendezvousTest.cpp b/cpp/tests/unit_tests/multi_gpu/ncclUniqueIdRendezvousTest.cpp new file mode 100644 index 000000000000..93c2b6d0a6c5 --- /dev/null +++ b/cpp/tests/unit_tests/multi_gpu/ncclUniqueIdRendezvousTest.cpp @@ -0,0 +1,971 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * 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 "tensorrt_llm/runtime/utils/ncclUniqueIdRendezvous.h" +#include "tensorrt_llm/common/opUtils.h" +#include "tensorrt_llm/runtime/ncclCommunicator.h" +#include "tensorrt_llm/runtime/utils/ncclHostApi.h" + +#include + +#if ENABLE_MULTI_DEVICE + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace +{ + +using tensorrt_llm::runtime::NcclUniqueIdRendezvousComm; +using tensorrt_llm::runtime::createNcclUniqueIdRendezvousComm; +using tensorrt_llm::runtime::exchangeNcclUniqueId; + +constexpr int kTestControlTagSeed = 0x4e43; +constexpr std::uint64_t kSubsetRendezvousId = 1; +constexpr std::uint64_t kStaleAttemptRendezvousId = 2; +constexpr std::uint64_t kRetryRendezvousId = 3; +constexpr std::uint64_t kStaleFloodFirstRendezvousId = 100; +constexpr std::uint64_t kStaleFloodCurrentRendezvousId = 10'000; +constexpr int kStaleFloodAttemptCount = 4'097; +constexpr std::uint64_t kRawRecoveryRendezvousId = 4; +constexpr std::uint64_t kRawCleanupRendezvousId = 5; +constexpr std::uint64_t kPpRecoveryRendezvousId = 6; +constexpr std::uint64_t kPpExcludedCleanupRendezvousId = 7; +constexpr std::uint64_t kProcessDeathRendezvousId = 8; + +tensorrt_llm::runtime::NcclUniqueIdRendezvousTags rawRendezvousTags() +{ + return {tensorrt_llm::mpi::MpiTag::kNcclCommReady, tensorrt_llm::mpi::MpiTag::kNcclCommUniqueId, + tensorrt_llm::mpi::MpiTag::kNcclCommAck}; +} + +std::vector allSessionRanks() +{ + std::vector ranks(COMM_SESSION.getSize()); + std::iota(ranks.begin(), ranks.end(), 0); + return ranks; +} + +std::shared_ptr createTestControlComm() +{ + return createNcclUniqueIdRendezvousComm( + allSessionRanks(), COMM_SESSION.getRank(), COMM_SESSION, kTestControlTagSeed); +} + +class EnvironmentVariableGuard +{ +public: + EnvironmentVariableGuard(char const* name, char const* value) + : mName(name) + { + if (auto const* previous = std::getenv(mName.c_str()); previous != nullptr) + { + mPrevious = previous; + } +#if defined(_WIN32) + _putenv_s(mName.c_str(), value); +#else + setenv(mName.c_str(), value, 1); +#endif + } + + ~EnvironmentVariableGuard() + { +#if defined(_WIN32) + _putenv_s(mName.c_str(), mPrevious.value_or("").c_str()); +#else + if (mPrevious.has_value()) + { + setenv(mName.c_str(), mPrevious->c_str(), 1); + } + else + { + unsetenv(mName.c_str()); + } +#endif + } + + EnvironmentVariableGuard(EnvironmentVariableGuard const&) = delete; + EnvironmentVariableGuard& operator=(EnvironmentVariableGuard const&) = delete; + +private: + std::string mName; + std::optional mPrevious; +}; + +bool isEnvironmentFlagEnabled(char const* name) +{ + auto const* value = std::getenv(name); + return value != nullptr && std::string{value} == "1"; +} + +bool configureDistinctCudaDeviceForEveryRank(std::string& failure) +{ + int const rank = COMM_SESSION.getRank(); + int const worldSize = COMM_SESSION.getSize(); + int deviceCount = 0; + cudaUUID_t deviceUuid{}; + int localReady = 0; + auto const countResult = cudaGetDeviceCount(&deviceCount); + if (countResult == cudaSuccess && deviceCount > 0) + { + int const device = rank % deviceCount; + auto const setResult = cudaSetDevice(device); + auto const uuidResult = setResult == cudaSuccess ? cudaDeviceGetUuid(&deviceUuid, device) : setResult; + localReady = setResult == cudaSuccess && uuidResult == cudaSuccess ? 1 : 0; + } + + std::vector gatheredReady(static_cast(worldSize)); + COMM_SESSION.allgather(&localReady, gatheredReady.data(), 1, tensorrt_llm::mpi::MpiType::kINT32); + std::vector gatheredUuids( + static_cast(worldSize) * static_cast(sizeof(deviceUuid.bytes))); + COMM_SESSION.allgather(deviceUuid.bytes, gatheredUuids.data(), static_cast(sizeof(deviceUuid.bytes)), + tensorrt_llm::mpi::MpiType::kBYTE); + + if (!std::all_of(gatheredReady.begin(), gatheredReady.end(), [](int ready) { return ready != 0; })) + { + failure = "every MPI rank must have a usable CUDA device"; + return false; + } + + std::unordered_set uniqueUuids; + for (int peer = 0; peer < worldSize; ++peer) + { + auto const* bytes = gatheredUuids.data() + + static_cast(peer) * static_cast(sizeof(deviceUuid.bytes)); + uniqueUuids.emplace(bytes, bytes + sizeof(deviceUuid.bytes)); + } + if (uniqueUuids.size() != static_cast(worldSize)) + { + failure = "every MPI rank must use a distinct physical GPU"; + return false; + } + return true; +} + +[[noreturn]] void exitProcessDeathTestFailure(int rank, std::string const& failure) noexcept +{ + std::fprintf(stderr, "TRTLLM_PROCESS_DEATH_RENDEZVOUS_FAIL rank=%d: %s\n", rank, failure.c_str()); + std::fflush(stderr); + std::_Exit(EXIT_FAILURE); +} + +void CUDART_CB waitForRelease(void* data) +{ + auto* release = static_cast*>(data); + while (!release->load(std::memory_order_acquire)) + { + std::this_thread::sleep_for(std::chrono::milliseconds{1}); + } +} + +class HostCallbackReleaseGuard +{ +public: + explicit HostCallbackReleaseGuard(std::atomic& release) + : mRelease(release) + { + } + + ~HostCallbackReleaseGuard() + { + mRelease.store(true, std::memory_order_release); + } + +private: + std::atomic& mRelease; +}; + +TEST(NcclUniqueIdRendezvousTest, DedicatedControlCommPreservesParentErrorHandlerAndRankMap) +{ + MPI_Errhandler parentBefore = MPI_ERRHANDLER_NULL; + ASSERT_EQ(MPI_Comm_get_errhandler(static_cast(COMM_SESSION), &parentBefore), MPI_SUCCESS); + + auto control = createTestControlComm(); + + MPI_Errhandler parentAfter = MPI_ERRHANDLER_NULL; + MPI_Errhandler controlHandler = MPI_ERRHANDLER_NULL; + ASSERT_EQ(MPI_Comm_get_errhandler(static_cast(COMM_SESSION), &parentAfter), MPI_SUCCESS); + ASSERT_EQ(MPI_Comm_get_errhandler(static_cast(control->mpiComm()), &controlHandler), MPI_SUCCESS); + EXPECT_EQ(parentAfter, parentBefore); + EXPECT_EQ(controlHandler, MPI_ERRORS_RETURN); + EXPECT_EQ(control->worldRank(), COMM_SESSION.getRank()); + EXPECT_EQ(control->worldRank(control->commRank(COMM_SESSION.getRank())), COMM_SESSION.getRank()); + + ASSERT_EQ(MPI_Errhandler_free(&parentBefore), MPI_SUCCESS); + ASSERT_EQ(MPI_Errhandler_free(&parentAfter), MPI_SUCCESS); + ASSERT_EQ(MPI_Errhandler_free(&controlHandler), MPI_SUCCESS); +} + +TEST(NcclUniqueIdRendezvousTest, DedicatedControlCommCacheReusesSameTopology) +{ + auto const firstControl = createTestControlComm(); + auto const secondControl = createTestControlComm(); + + ASSERT_NE(firstControl, nullptr); + ASSERT_NE(secondControl, nullptr); + EXPECT_EQ(firstControl.get(), secondControl.get()); + EXPECT_EQ(MPI_Comm_c2f(static_cast(firstControl->mpiComm())), + MPI_Comm_c2f(static_cast(secondControl->mpiComm()))); +} + +TEST(NcclUniqueIdRendezvousTest, DedicatedControlCommCacheSurvivesCustomParentHandleLifetime) +{ + auto const sessionControl = createTestControlComm(); + std::shared_ptr firstDuplicateControl; + { + MPI_Comm duplicate = MPI_COMM_NULL; + ASSERT_EQ(MPI_Comm_dup(static_cast(COMM_SESSION), &duplicate), MPI_SUCCESS); + auto duplicateParent = tensorrt_llm::mpi::MpiComm{duplicate, true}; + firstDuplicateControl = createNcclUniqueIdRendezvousComm( + allSessionRanks(), COMM_SESSION.getRank(), duplicateParent, kTestControlTagSeed); + } + + std::shared_ptr secondDuplicateControl; + { + MPI_Comm duplicate = MPI_COMM_NULL; + ASSERT_EQ(MPI_Comm_dup(static_cast(COMM_SESSION), &duplicate), MPI_SUCCESS); + auto duplicateParent = tensorrt_llm::mpi::MpiComm{duplicate, true}; + secondDuplicateControl = createNcclUniqueIdRendezvousComm( + allSessionRanks(), COMM_SESSION.getRank(), duplicateParent, kTestControlTagSeed); + } + + EXPECT_EQ(firstDuplicateControl.get(), sessionControl.get()); + EXPECT_EQ(secondDuplicateControl.get(), sessionControl.get()); +} + +TEST(NcclUniqueIdRendezvousTest, NonContiguousSurvivorSubsetExchangesWithoutExcludedRank) +{ + auto control = createTestControlComm(); + int const rank = COMM_SESSION.getRank(); + int const worldSize = COMM_SESSION.getSize(); + if (worldSize < 2) + { + GTEST_SKIP() << "Test requires at least two MPI processes"; + } + + // With three or more processes rank 1 is deliberately excluded, proving + // that recovery rendezvous does not contain a hidden world collective. + std::vector activeRanks{0, worldSize - 1}; + std::array localId{}; + if (std::binary_search(activeRanks.begin(), activeRanks.end(), rank)) + { + auto const id = exchangeNcclUniqueId(activeRanks, *control, rawRendezvousTags(), kSubsetRendezvousId, + std::chrono::steady_clock::now() + std::chrono::seconds{10}); + std::memcpy(localId.data(), &id, sizeof(id)); + } + + // This synchronization is test-only and happens after the survivor + // exchange. The production recovery path performs no MPI collective. + COMM_SESSION.barrier(); + std::vector gathered(static_cast(worldSize) * localId.size()); + COMM_SESSION.allgather( + localId.data(), gathered.data(), static_cast(localId.size()), tensorrt_llm::mpi::MpiType::kBYTE); + + auto const* first = gathered.data(); + auto const* last = gathered.data() + static_cast(worldSize - 1) * localId.size(); + EXPECT_TRUE(std::any_of(first, first + localId.size(), [](unsigned char value) { return value != 0; })); + EXPECT_TRUE(std::equal(first, first + localId.size(), last)); + if (worldSize >= 3) + { + auto const* excluded = gathered.data() + localId.size(); + EXPECT_TRUE(std::all_of(excluded, excluded + localId.size(), [](unsigned char value) { return value == 0; })); + } +} + +TEST(NcclUniqueIdRendezvousTest, RetryGenerationCannotPairWithStaleAttemptMessages) +{ + auto const control = createTestControlComm(); + int const rank = COMM_SESSION.getRank(); + int const worldSize = COMM_SESSION.getSize(); + if (worldSize < 2) + { + GTEST_SKIP() << "Test requires at least two MPI processes"; + } + + std::vector const activeRanks{0, worldSize - 1}; + bool const isRoot = rank == activeRanks.front(); + bool const isClient = rank == activeRanks.back(); + bool staleAttemptTimedOut = false; + bool retrySucceeded = false; + std::array retryId{}; + + // The root times out generation A before the client starts it, then enters + // generation B. The client deliberately sends a late generation-A READY + // while the root is already waiting in B. That message must be stashed, + // not accepted as B's READY; both ranks converge only after the client also + // advances to B. + control->mpiComm().barrier(); + if (isRoot) + { + try + { + static_cast(exchangeNcclUniqueId(activeRanks, *control, rawRendezvousTags(), + kStaleAttemptRendezvousId, std::chrono::steady_clock::now() + std::chrono::milliseconds{100})); + } + catch (std::exception const& error) + { + staleAttemptTimedOut = std::string{error.what()}.find("timed out") != std::string::npos; + } + + try + { + auto const id = exchangeNcclUniqueId(activeRanks, *control, rawRendezvousTags(), kRetryRendezvousId, + std::chrono::steady_clock::now() + std::chrono::seconds{5}); + std::memcpy(retryId.data(), &id, sizeof(id)); + retrySucceeded = true; + } + catch (std::exception const&) + { + // Assert after every rank reaches the test-only barrier. If the + // generation key regresses, the late A exchange can consume B and + // leave this retry waiting until its bounded deadline. + } + } + else if (isClient) + { + std::this_thread::sleep_for(std::chrono::milliseconds{500}); + try + { + static_cast(exchangeNcclUniqueId(activeRanks, *control, rawRendezvousTags(), + kStaleAttemptRendezvousId, std::chrono::steady_clock::now() + std::chrono::milliseconds{150})); + } + catch (std::exception const& error) + { + staleAttemptTimedOut = std::string{error.what()}.find("timed out") != std::string::npos; + } + + try + { + auto const id = exchangeNcclUniqueId(activeRanks, *control, rawRendezvousTags(), kRetryRendezvousId, + std::chrono::steady_clock::now() + std::chrono::seconds{5}); + std::memcpy(retryId.data(), &id, sizeof(id)); + retrySucceeded = true; + } + catch (std::exception const&) + { + // Keep the MPI test convergent on the regression path; the + // assertions below retain the failure details. + } + } + + control->mpiComm().barrier(); + std::vector gathered(static_cast(worldSize) * retryId.size()); + control->mpiComm().allgather( + retryId.data(), gathered.data(), static_cast(retryId.size()), tensorrt_llm::mpi::MpiType::kBYTE); + + if (isRoot || isClient) + { + EXPECT_TRUE(staleAttemptTimedOut); + EXPECT_TRUE(retrySucceeded); + } + auto const* rootId = gathered.data(); + auto const* clientId = gathered.data() + static_cast(worldSize - 1) * retryId.size(); + EXPECT_TRUE(std::any_of(rootId, rootId + retryId.size(), [](unsigned char value) { return value != 0; })); + EXPECT_TRUE(std::equal(rootId, rootId + retryId.size(), clientId)); +} + +TEST(NcclUniqueIdRendezvousTest, RepeatedStaleGenerationsDoNotExhaustPendingQueue) +{ + auto const control = createTestControlComm(); + int const rank = COMM_SESSION.getRank(); + int const worldSize = COMM_SESSION.getSize(); + if (worldSize < 2) + { + GTEST_SKIP() << "Test requires at least two MPI processes"; + } + + std::vector const activeRanks{0, worldSize - 1}; + bool const isRoot = rank == activeRanks.front(); + bool const isClient = rank == activeRanks.back(); + bool currentAttemptSucceeded = false; + int staleAttemptsTimedOut = 0; + std::array currentId{}; + + // The old opaque-key protocol stashed every lower-generation READY while + // the root waited on the current generation, then failed at its 4,096 + // pending-message cap. Distinct increasing stale IDs honor the caller + // contract while proving that same-communicator generations older than the + // root's current ID are discarded instead of retained indefinitely. + control->mpiComm().barrier(); + if (isRoot) + { + try + { + auto const id = exchangeNcclUniqueId(activeRanks, *control, rawRendezvousTags(), + kStaleFloodCurrentRendezvousId, std::chrono::steady_clock::now() + std::chrono::seconds{120}); + std::memcpy(currentId.data(), &id, sizeof(id)); + currentAttemptSucceeded = true; + } + catch (std::exception const&) + { + // Converge at the final test-only barrier even when the old queue + // limit is hit; the assertions below report the regression. + } + } + else if (isClient) + { + std::this_thread::sleep_for(std::chrono::milliseconds{100}); + for (int attempt = 0; attempt < kStaleFloodAttemptCount; ++attempt) + { + auto const staleId = kStaleFloodFirstRendezvousId + static_cast(attempt); + try + { + static_cast(exchangeNcclUniqueId(activeRanks, *control, rawRendezvousTags(), staleId, + std::chrono::steady_clock::now() + std::chrono::milliseconds{2})); + } + catch (std::exception const& error) + { + if (std::string{error.what()}.find("timed out waiting for survivor communicator ID") + != std::string::npos) + { + ++staleAttemptsTimedOut; + } + } + } + + try + { + auto const id = exchangeNcclUniqueId(activeRanks, *control, rawRendezvousTags(), + kStaleFloodCurrentRendezvousId, std::chrono::steady_clock::now() + std::chrono::seconds{60}); + std::memcpy(currentId.data(), &id, sizeof(id)); + currentAttemptSucceeded = true; + } + catch (std::exception const&) + { + // The root may already have failed on the old pending-queue cap. + } + } + + control->mpiComm().barrier(); + std::vector gathered(static_cast(worldSize) * currentId.size()); + control->mpiComm().allgather( + currentId.data(), gathered.data(), static_cast(currentId.size()), tensorrt_llm::mpi::MpiType::kBYTE); + + if (isRoot || isClient) + { + EXPECT_TRUE(currentAttemptSucceeded); + } + if (isClient) + { + EXPECT_EQ(staleAttemptsTimedOut, kStaleFloodAttemptCount); + } + auto const* rootId = gathered.data(); + auto const* clientId = gathered.data() + static_cast(worldSize - 1) * currentId.size(); + EXPECT_TRUE(std::any_of(rootId, rootId + currentId.size(), [](unsigned char value) { return value != 0; })); + EXPECT_TRUE(std::equal(rootId, rootId + currentId.size(), clientId)); +} + +TEST(NcclUniqueIdRendezvousTest, WaitOperationDoesNotHoldGlobalGateForUnrelatedStreamPrefix) +{ + EnvironmentVariableGuard const faultToleranceMode{"TLLM_FAULT_TOLERANCE_MODE", "1"}; + EnvironmentVariableGuard const operationTimeout{"TRTLLM_NCCL_NONBLOCKING_TIMEOUT_MS", "200"}; + EnvironmentVariableGuard const watchdogPoll{"TRTLLM_NCCL_WATCHDOG_POLL_INTERVAL_MS", "10"}; + + int deviceCount = 0; + ASSERT_EQ(cudaGetDeviceCount(&deviceCount), cudaSuccess); + if (deviceCount == 0) + { + GTEST_SKIP() << "Test requires a visible CUDA device"; + } + int const rank = COMM_SESSION.getRank(); + ASSERT_EQ(cudaSetDevice(rank % deviceCount), cudaSuccess); + + auto comm = tensorrt_llm::getComm({rank}); + cudaStream_t stream = nullptr; + ASSERT_EQ(cudaStreamCreate(&stream), cudaSuccess); + std::atomic releasePrefix{false}; + HostCallbackReleaseGuard const releaseGuard{releasePrefix}; + ASSERT_EQ(cudaLaunchHostFunc(stream, waitForRelease, &releasePrefix), cudaSuccess); + + std::uint64_t token = 0; + { + auto lease = tensorrt_llm::acquireComm(comm); + token = lease.begin(stream, "operation behind unrelated stream prefix"); + lease.track(token, stream); + } + + std::atomic waiterStarted{false}; + std::exception_ptr waiterError; + std::thread waiter( + [&]() + { + try + { + waiterStarted.store(true, std::memory_order_release); + tensorrt_llm::waitCommOperation(comm, token, "operation behind unrelated stream prefix"); + } + catch (...) + { + waiterError = std::current_exception(); + } + }); + while (!waiterStarted.load(std::memory_order_acquire)) + { + std::this_thread::yield(); + } + // Give the waiter time to enter its polling path before contending for the + // process-wide NCCL host gate from an unrelated control thread. + std::this_thread::sleep_for(std::chrono::milliseconds{100}); + + std::atomic gateAcquired{false}; + std::thread gateContender( + [&]() + { + auto gate = tensorrt_llm::runtime::acquireNcclHostApiLock(); + gateAcquired.store(true, std::memory_order_release); + }); + auto const gateDeadline = std::chrono::steady_clock::now() + std::chrono::seconds{2}; + while (!gateAcquired.load(std::memory_order_acquire) && std::chrono::steady_clock::now() < gateDeadline) + { + std::this_thread::sleep_for(std::chrono::milliseconds{1}); + } + bool const acquiredBeforePrefixRelease = gateAcquired.load(std::memory_order_acquire); + + releasePrefix.store(true, std::memory_order_release); + gateContender.join(); + waiter.join(); + + EXPECT_TRUE(acquiredBeforePrefixRelease) + << "waitCommOperation retained the process-wide NCCL host gate while waiting for an unrelated stream prefix"; + EXPECT_FALSE(waiterError); + EXPECT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + EXPECT_EQ(cudaStreamDestroy(stream), cudaSuccess); +} + +TEST(NcclUniqueIdRendezvousTest, FaultToleranceRejectsGraphCaptureAndTimesOutStalledOperation) +{ + EnvironmentVariableGuard const faultToleranceMode{"TLLM_FAULT_TOLERANCE_MODE", "1"}; + EnvironmentVariableGuard const operationTimeout{"TRTLLM_NCCL_NONBLOCKING_TIMEOUT_MS", "100"}; + EnvironmentVariableGuard const watchdogPoll{"TRTLLM_NCCL_WATCHDOG_POLL_INTERVAL_MS", "10"}; + + int deviceCount = 0; + ASSERT_EQ(cudaGetDeviceCount(&deviceCount), cudaSuccess); + if (deviceCount == 0) + { + GTEST_SKIP() << "Test requires a visible CUDA device"; + } + int const rank = COMM_SESSION.getRank(); + ASSERT_EQ(cudaSetDevice(rank % deviceCount), cudaSuccess); + + std::set const singletonGroup{rank}; + auto comm = tensorrt_llm::getComm(singletonGroup); + cudaStream_t stream = nullptr; + int* captureScratch = nullptr; + ASSERT_EQ(cudaStreamCreate(&stream), cudaSuccess); + ASSERT_EQ(cudaMalloc(reinterpret_cast(&captureScratch), sizeof(*captureScratch)), cudaSuccess); + + ASSERT_EQ(cudaStreamBeginCapture(stream, cudaStreamCaptureModeThreadLocal), cudaSuccess); + ASSERT_EQ(cudaMemsetAsync(captureScratch, 0, sizeof(*captureScratch), stream), cudaSuccess); + { + auto lease = tensorrt_llm::acquireComm(comm); + EXPECT_THROW(static_cast(lease.begin(stream, "captured NCCL test operation")), std::exception); + } + cudaGraph_t graph = nullptr; + ASSERT_EQ(cudaStreamEndCapture(stream, &graph), cudaSuccess); + if (graph != nullptr) + { + ASSERT_EQ(cudaGraphDestroy(graph), cudaSuccess); + } + ASSERT_EQ(cudaFree(captureScratch), cudaSuccess); + + std::atomic releaseCallback{false}; + HostCallbackReleaseGuard const releaseGuard{releaseCallback}; + { + auto lease = tensorrt_llm::acquireComm(comm); + auto const token = lease.begin(stream, "injected stalled NCCL operation"); + ASSERT_EQ(cudaLaunchHostFunc(stream, waitForRelease, &releaseCallback), cudaSuccess); + lease.track(token, stream); + } + + auto const deadline = std::chrono::steady_clock::now() + std::chrono::seconds{5}; + std::optional error; + while (std::chrono::steady_clock::now() < deadline) + { + error = tensorrt_llm::getCommAsyncError(singletonGroup); + if (error.has_value()) + { + break; + } + std::this_thread::sleep_for(std::chrono::milliseconds{10}); + } + EXPECT_TRUE(error.has_value()); + if (error.has_value()) + { + EXPECT_NE(error->find("timed out"), std::string::npos); + } + + releaseCallback.store(true, std::memory_order_release); + EXPECT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + EXPECT_EQ(cudaStreamDestroy(stream), cudaSuccess); +} + +TEST(NcclUniqueIdRendezvousTest, FaultToleranceDisabledPreservesLegacyOperationPath) +{ + if (!isEnvironmentFlagEnabled("TRTLLM_TEST_FT_DEFAULT_OFF")) + { + GTEST_SKIP() << "Run in the isolated default-off CTest process"; + } + ASSERT_FALSE(tensorrt_llm::mpi::isFaultToleranceModeEnabled()); + + int deviceCount = 0; + ASSERT_EQ(cudaGetDeviceCount(&deviceCount), cudaSuccess); + if (deviceCount == 0) + { + GTEST_SKIP() << "Test requires a visible CUDA device"; + } + int const rank = COMM_SESSION.getRank(); + ASSERT_EQ(cudaSetDevice(rank % deviceCount), cudaSuccess); + + std::set const singletonGroup{rank}; + auto comm = tensorrt_llm::getComm(singletonGroup); + auto lease = tensorrt_llm::acquireComm(comm); + EXPECT_EQ(lease.begin(nullptr, "default-off legacy operation"), 0); + EXPECT_EQ(lease.get(), *comm); + EXPECT_THROW( + tensorrt_llm::abortAndReinitComm(singletonGroup, singletonGroup, kRawRecoveryRendezvousId), std::exception); +} + +TEST(NcclUniqueIdRendezvousTest, ProcessDeathRendezvousRequiresOptInFaultTolerantMpiLauncher) +{ + if (!isEnvironmentFlagEnabled("TRTLLM_TEST_NCCL_PROCESS_DEATH")) + { + GTEST_SKIP() << "This destructive test requires a dedicated fault-tolerant MPI launcher"; + } + int const rank = COMM_SESSION.getRank(); + int const worldSize = COMM_SESSION.getSize(); + if (worldSize != 3) + { + exitProcessDeathTestFailure(rank, "the opt-in process-death test requires exactly three MPI processes"); + } + + try + { + std::string gpuFailure; + if (!configureDistinctCudaDeviceForEveryRank(gpuFailure)) + { + exitProcessDeathTestFailure(rank, gpuFailure); + } + + // Establish both the production raw-NCCL communicator and its retained + // MPI control channel while every rank is healthy. The separate test + // control channel supplies only the pre-failure synchronization below. + auto const control = createTestControlComm(); + auto const fullRanks = allSessionRanks(); + std::set const fullGroup(fullRanks.begin(), fullRanks.end()); + auto const fullComm = tensorrt_llm::getComm(fullGroup); + if (fullComm == nullptr || *fullComm == nullptr) + { + throw std::runtime_error("healthy full-group NCCL communicator is null"); + } + control->mpiComm().barrier(); + if (rank == 1) + { + std::_Exit(EXIT_SUCCESS); + } + + // This probes the stronger Mode-A contract: after one process exits, + // surviving ranks must still exchange a fresh NCCL ID over the retained + // MPI channel, abort the poisoned full communicator, and execute real + // data-plane work on its replacement. + std::this_thread::sleep_for(std::chrono::milliseconds{500}); + std::set const survivorGroup{0, 2}; + tensorrt_llm::abortAndReinitComm(fullGroup, survivorGroup, kProcessDeathRendezvousId); + auto const rebuiltComm = tensorrt_llm::getComm(survivorGroup); + if (rebuiltComm == nullptr || *rebuiltComm == nullptr) + { + throw std::runtime_error("rebuilt survivor NCCL communicator is null"); + } + + auto const checkCuda = [](cudaError_t result, char const* operation) + { + if (result != cudaSuccess) + { + throw std::runtime_error(std::string{operation} + " failed: " + cudaGetErrorString(result)); + } + }; + int const hostInput = rank + 1; + int hostOutput = 0; + int* deviceInput = nullptr; + int* deviceOutput = nullptr; + cudaStream_t stream = nullptr; + checkCuda(cudaMalloc(reinterpret_cast(&deviceInput), sizeof(hostInput)), "cudaMalloc(input)"); + checkCuda(cudaMalloc(reinterpret_cast(&deviceOutput), sizeof(hostOutput)), "cudaMalloc(output)"); + checkCuda(cudaStreamCreate(&stream), "cudaStreamCreate"); + checkCuda(cudaMemcpyAsync(deviceInput, &hostInput, sizeof(hostInput), cudaMemcpyHostToDevice, stream), + "cudaMemcpyAsync(input)"); + + std::uint64_t operationToken = 0; + { + auto lease = tensorrt_llm::acquireComm(rebuiltComm); + operationToken = lease.begin(stream, "ncclAllReduce(process-death recovery test)"); + lease.check(ncclAllReduce(deviceInput, deviceOutput, 1, ncclInt32, ncclSum, lease.get(), stream), + "ncclAllReduce(process-death recovery test)"); + lease.track(operationToken, stream); + } + tensorrt_llm::waitCommOperation(rebuiltComm, operationToken, "ncclAllReduce(process-death recovery test)"); + checkCuda( + cudaMemcpy(&hostOutput, deviceOutput, sizeof(hostOutput), cudaMemcpyDeviceToHost), "cudaMemcpy(output)"); + constexpr int kExpectedSurvivorSum = 4; + if (hostOutput != kExpectedSurvivorSum) + { + throw std::runtime_error("post-recovery allreduce returned " + std::to_string(hostOutput) + ", expected " + + std::to_string(kExpectedSurvivorSum)); + } + + if (rank == *survivorGroup.begin()) + { + // Some fault-tolerant launchers still return nonzero because the + // victim intentionally skipped MPI_Finalize. CTest keys off this + // flushed sentinel, reached only after survivor communicator + // rebuild and the expected allreduce, instead of that exit code. + std::fprintf(stdout, "TRTLLM_PROCESS_DEATH_RENDEZVOUS_OK\n"); + std::fflush(stdout); + } + std::_Exit(EXIT_SUCCESS); + } + catch (std::exception const& error) + { + exitProcessDeathTestFailure(rank, error.what()); + } + catch (...) + { + exitProcessDeathTestFailure(rank, "unknown exception"); + } +} + +TEST(NcclUniqueIdRendezvousTest, RawCommunicatorRebuildRunsPostRecoveryCollective) +{ + int const rank = COMM_SESSION.getRank(); + int const worldSize = COMM_SESSION.getSize(); + if (worldSize < 3) + { + GTEST_SKIP() << "Test requires at least three MPI processes"; + } + EnvironmentVariableGuard const faultToleranceMode{"TLLM_FAULT_TOLERANCE_MODE", "1"}; + + int deviceCount = 0; + ASSERT_EQ(cudaGetDeviceCount(&deviceCount), cudaSuccess); + if (deviceCount == 0) + { + GTEST_SKIP() << "Test requires a visible CUDA device"; + } + int const device = rank % deviceCount; + ASSERT_EQ(cudaSetDevice(device), cudaSuccess); + + // A launcher may expose all GPUs to every rank or bind each rank to one + // different physical GPU. Compare UUIDs rather than local ordinals so both + // layouts work, and never run multiple NCCL ranks on the same device. + cudaUUID_t deviceUuid{}; + ASSERT_EQ(cudaDeviceGetUuid(&deviceUuid, device), cudaSuccess); + std::vector gatheredUuids( + static_cast(worldSize) * static_cast(sizeof(deviceUuid.bytes))); + COMM_SESSION.allgather(deviceUuid.bytes, gatheredUuids.data(), static_cast(sizeof(deviceUuid.bytes)), + tensorrt_llm::mpi::MpiType::kBYTE); + std::unordered_set uniqueUuids; + for (int peer = 0; peer < worldSize; ++peer) + { + auto const* bytes = gatheredUuids.data() + + static_cast(peer) * static_cast(sizeof(deviceUuid.bytes)); + uniqueUuids.emplace(bytes, bytes + sizeof(deviceUuid.bytes)); + } + if (uniqueUuids.size() != static_cast(worldSize)) + { + GTEST_SKIP() << "Test requires one distinct physical GPU per MPI rank"; + } + + auto const fullRanks = allSessionRanks(); + std::set const fullGroup(fullRanks.begin(), fullRanks.end()); + std::set const survivorGroup{0, worldSize - 1}; + auto const testSync = createTestControlComm(); + auto oldComm = tensorrt_llm::getComm(fullGroup); + ASSERT_NE(oldComm, nullptr); + EXPECT_EQ(tensorrt_llm::getCommWorldRank(oldComm), rank); + + // Keep the old registry state alive so its storage cannot be recycled; + // pointer inequality then proves recovery installed a fresh state. + std::shared_ptr cachedSurvivorComm; + ncclComm_t* cachedSurvivorState = nullptr; + if (survivorGroup.count(rank) != 0) + { + cachedSurvivorComm = tensorrt_llm::getComm(survivorGroup); + ASSERT_NE(cachedSurvivorComm, nullptr); + cachedSurvivorState = cachedSurvivorComm.get(); + } + + // Ensure the survivor-target cache entry exists before any rank begins + // recovery. Use a dedicated test communicator so the excluded rank does + // not enter a parent-communicator collective while survivors create their + // subgroup control channel. + testSync->mpiComm().barrier(); + + if (survivorGroup.count(rank) != 0) + { + tensorrt_llm::abortAndReinitComm(fullGroup, survivorGroup, kRawRecoveryRendezvousId); + auto rebuiltComm = tensorrt_llm::getComm(survivorGroup); + ASSERT_NE(rebuiltComm, nullptr); + EXPECT_EQ(tensorrt_llm::getCommWorldRank(rebuiltComm), rank); + EXPECT_NE(rebuiltComm.get(), cachedSurvivorState); + EXPECT_EQ(tensorrt_llm::getComm(survivorGroup).get(), rebuiltComm.get()); + + int const hostInput = rank + 1; + int hostOutput = 0; + int* deviceInput = nullptr; + int* deviceOutput = nullptr; + cudaStream_t stream = nullptr; + ASSERT_EQ(cudaMalloc(reinterpret_cast(&deviceInput), sizeof(hostInput)), cudaSuccess); + ASSERT_EQ(cudaMalloc(reinterpret_cast(&deviceOutput), sizeof(hostOutput)), cudaSuccess); + ASSERT_EQ(cudaStreamCreate(&stream), cudaSuccess); + ASSERT_EQ( + cudaMemcpyAsync(deviceInput, &hostInput, sizeof(hostInput), cudaMemcpyHostToDevice, stream), cudaSuccess); + + std::uint64_t operationToken = 0; + { + auto lease = tensorrt_llm::acquireComm(rebuiltComm); + operationToken = lease.begin(stream, "ncclAllReduce(recovery test)"); + lease.check(ncclAllReduce(deviceInput, deviceOutput, 1, ncclInt32, ncclSum, lease.get(), stream), + "ncclAllReduce(recovery test)"); + lease.track(operationToken, stream); + } + tensorrt_llm::waitCommOperation(rebuiltComm, operationToken, "ncclAllReduce(recovery test)"); + + ASSERT_EQ(cudaMemcpy(&hostOutput, deviceOutput, sizeof(hostOutput), cudaMemcpyDeviceToHost), cudaSuccess); + EXPECT_EQ(hostOutput, worldSize + 1); + EXPECT_EQ(cudaStreamDestroy(stream), cudaSuccess); + EXPECT_EQ(cudaFree(deviceOutput), cudaSuccess); + EXPECT_EQ(cudaFree(deviceInput), cudaSuccess); + + // Leave only a singleton communicator for process teardown. This is + // test cleanup, not another simulated survivor-membership decision. + tensorrt_llm::abortAndReinitComm(survivorGroup, {rank}, kRawCleanupRendezvousId); + } + else + { + // The excluded rank deliberately does not participate in the survivor + // rendezvous. Once survivors have started rebuilding, replace its old + // local communicator with a singleton so test-process teardown never + // calls ncclCommDestroy on a half-aborted full-group communicator. + tensorrt_llm::abortAndReinitComm(fullGroup, {rank}, kRawCleanupRendezvousId); + } + + // Test-only synchronization after every local recovery path has finished. + testSync->mpiComm().barrier(); +} + +TEST(NcclUniqueIdRendezvousTest, PipelineCommunicatorRebuildRunsPostRecoverySendRecv) +{ + int const rank = COMM_SESSION.getRank(); + int const worldSize = COMM_SESSION.getSize(); + if (worldSize < 3) + { + GTEST_SKIP() << "Test requires at least three MPI processes"; + } + EnvironmentVariableGuard const faultToleranceMode{"TLLM_FAULT_TOLERANCE_MODE", "1"}; + + int deviceCount = 0; + ASSERT_EQ(cudaGetDeviceCount(&deviceCount), cudaSuccess); + if (deviceCount == 0) + { + GTEST_SKIP() << "Test requires a visible CUDA device"; + } + int const device = rank % deviceCount; + ASSERT_EQ(cudaSetDevice(device), cudaSuccess); + + cudaUUID_t deviceUuid{}; + ASSERT_EQ(cudaDeviceGetUuid(&deviceUuid, device), cudaSuccess); + std::vector gatheredUuids( + static_cast(worldSize) * static_cast(sizeof(deviceUuid.bytes))); + COMM_SESSION.allgather(deviceUuid.bytes, gatheredUuids.data(), static_cast(sizeof(deviceUuid.bytes)), + tensorrt_llm::mpi::MpiType::kBYTE); + std::unordered_set uniqueUuids; + for (int peer = 0; peer < worldSize; ++peer) + { + auto const* bytes = gatheredUuids.data() + + static_cast(peer) * static_cast(sizeof(deviceUuid.bytes)); + uniqueUuids.emplace(bytes, bytes + sizeof(deviceUuid.bytes)); + } + if (uniqueUuids.size() != static_cast(worldSize)) + { + GTEST_SKIP() << "Test requires one distinct physical GPU per MPI rank"; + } + + auto const testSync = createTestControlComm(); + auto communicator = std::make_unique(worldSize, rank); + std::vector const survivorRanks{0, worldSize - 1}; + bool const isSurvivor = std::binary_search(survivorRanks.begin(), survivorRanks.end(), rank); + + testSync->mpiComm().barrier(); + if (isSurvivor) + { + communicator->abortAndReinit(survivorRanks, kPpRecoveryRendezvousId); + EXPECT_EQ(communicator->getActiveRanks(), survivorRanks); + } + else + { + communicator->abortAndReinit({rank}, kPpExcludedCleanupRendezvousId); + EXPECT_EQ(communicator->getActiveRanks(), std::vector{rank}); + } + EXPECT_TRUE(communicator->getAsyncError().empty()); + + int* deviceValue = nullptr; + ASSERT_EQ(cudaMalloc(reinterpret_cast(&deviceValue), sizeof(*deviceValue)), cudaSuccess); + int hostValue = rank == survivorRanks.front() ? 0x1a07 : 0; + ASSERT_EQ(cudaMemcpy(deviceValue, &hostValue, sizeof(hostValue), cudaMemcpyHostToDevice), cudaSuccess); + auto buffer = tensorrt_llm::runtime::IBuffer::wrap(deviceValue, nvinfer1::DataType::kINT32, 1); + tensorrt_llm::runtime::CudaStream stream; + + if (rank == survivorRanks.front()) + { + communicator->send(*buffer, survivorRanks.back(), stream); + stream.synchronize(); + } + else if (rank == survivorRanks.back()) + { + communicator->receive(*buffer, survivorRanks.front(), stream); + stream.synchronize(); + ASSERT_EQ(cudaMemcpy(&hostValue, deviceValue, sizeof(hostValue), cudaMemcpyDeviceToHost), cudaSuccess); + EXPECT_EQ(hostValue, 0x1a07); + } + + buffer.reset(); + ASSERT_EQ(cudaFree(deviceValue), cudaSuccess); + testSync->mpiComm().barrier(); + communicator.reset(); + testSync->mpiComm().barrier(); +} + +} // namespace + +#endif // ENABLE_MULTI_DEVICE diff --git a/cpp/tests/unit_tests/multi_gpu/ncclUtilsTest.cpp b/cpp/tests/unit_tests/multi_gpu/ncclUtilsTest.cpp index 4f3309431a2f..5c8aeab428f8 100644 --- a/cpp/tests/unit_tests/multi_gpu/ncclUtilsTest.cpp +++ b/cpp/tests/unit_tests/multi_gpu/ncclUtilsTest.cpp @@ -20,7 +20,9 @@ #include "tensorrt_llm/common/logger.h" #include "tensorrt_llm/common/opUtils.h" #include "tensorrt_llm/runtime/utils/mpiUtils.h" +#include "tensorrt_llm/runtime/utils/ncclHostApi.h" +#include #include #include #include @@ -87,34 +89,46 @@ TEST(NCCLWindowSupportTest, RuntimeVersionAndGB10Gate) // Helper function to create a split communicator for testing // This allows us to test cleanup behavior explicitly by controlling the lifetime -std::shared_ptr createSplitComm(ncclComm_t parentComm, int color, int key) +std::shared_ptr createSplitComm(std::shared_ptr const& parentComm, int color, int key) { - ncclComm_t newComm; - ncclResult_t result = ncclCommSplit(parentComm, color, key, &newComm, nullptr); - if (result != ncclSuccess) + ncclComm_t newComm = nullptr; + ncclConfig_t config = NCCL_CONFIG_INITIALIZER; + // The cached parent is nonblocking. Make the test-owned child blocking so + // the legacy raw allocator overload and its simple deleter have an + // explicit blocking-communicator contract. + config.blocking = 1; { - TLLM_THROW("ncclCommSplit failed with error: %d", result); + auto lease = tensorrt_llm::acquireComm(parentComm); + auto const result = ncclCommSplit(lease.get(), color, key, &newComm, &config); + lease.check(result, "ncclCommSplit(ncclUtilsTest)"); + } + if (newComm == nullptr) + { + TLLM_THROW("ncclCommSplit returned a null child communicator"); } // Create a shared_ptr with custom deleter that cleans up resources first return std::shared_ptr(new ncclComm_t(newComm), [](ncclComm_t* comm) { - if (comm && *comm) + if (comm != nullptr && *comm != nullptr) { // STEP 1: Clean up all registered resources FIRST - tensorrt_llm::common::nccl_util::NcclCommResourceManager::getInstance().cleanupResources(*comm); + bool const resourcesCleaned + = tensorrt_llm::common::nccl_util::NcclCommResourceManager::getInstance().cleanupResources(*comm); - // STEP 2: Now destroy the NCCL communicator - ncclResult_t result = ncclCommDestroy(*comm); + // STEP 2: Serialize with the cached parent's watchdog and + // destroy only after resource cleanup fully completed. + auto const hostApiLock = tensorrt_llm::runtime::acquireNcclHostApiLock(); + ncclResult_t const result = resourcesCleaned ? ncclCommDestroy(*comm) : ncclCommAbort(*comm); if (result != ncclSuccess) { - TLLM_LOG_WARNING("ncclCommDestroy failed with error: %d", result); + TLLM_LOG_WARNING("NCCL test communicator release failed with error: %d", result); } - - // STEP 3: Free the memory - delete comm; + *comm = nullptr; } + // STEP 3: Free the holder even if communicator creation failed. + delete comm; }); } @@ -122,6 +136,106 @@ std::shared_ptr createSplitComm(ncclComm_t parentComm, int color, in // NcclCommResourceManager Tests //============================================================================== +TEST(NcclCommResourceManagerStandaloneTest, ReusedHandleWaitsForAbortCleanup) +{ + auto& manager = nccl_util::NcclCommResourceManager::getInstance(); + int handleStorage = 0; + auto const comm = reinterpret_cast(&handleStorage); + std::atomic waitStarted{false}; + std::atomic waitCompleted{false}; + + manager.beginAbortCleanup(comm); + std::thread waiter( + [&]() + { + waitStarted.store(true, std::memory_order_release); + manager.waitForAbortCleanup(comm); + waitCompleted.store(true, std::memory_order_release); + }); + while (!waitStarted.load(std::memory_order_acquire)) + { + std::this_thread::yield(); + } + EXPECT_FALSE(waitCompleted.load(std::memory_order_acquire)); + + EXPECT_TRUE(manager.cleanupResources(comm, true)); + waiter.join(); + EXPECT_TRUE(waitCompleted.load(std::memory_order_acquire)); +} + +TEST(NcclCommResourceManagerStandaloneTest, ConcurrentCleanupCannotReleaseRetiredHandleEarly) +{ + auto& manager = nccl_util::NcclCommResourceManager::getInstance(); + int handleStorage = 0; + auto const comm = reinterpret_cast(&handleStorage); + std::atomic cleanupEntered{false}; + std::atomic releaseCleanup{false}; + std::atomic cleanupSucceeded{false}; + manager.registerResource( + comm, + [&]() + { + cleanupEntered.store(true, std::memory_order_release); + while (!releaseCleanup.load(std::memory_order_acquire)) + { + std::this_thread::yield(); + } + return true; + }, + "BlockingAbortCleanup"); + manager.beginAbortCleanup(comm); + + std::thread cleanupOwner( + [&]() { cleanupSucceeded.store(manager.cleanupResources(comm, true), std::memory_order_release); }); + while (!cleanupEntered.load(std::memory_order_acquire)) + { + std::this_thread::yield(); + } + + EXPECT_TRUE(manager.cleanupResources(comm, true)); + EXPECT_TRUE(manager.isAbortCleanup(comm)); + releaseCleanup.store(true, std::memory_order_release); + cleanupOwner.join(); + EXPECT_TRUE(cleanupSucceeded.load(std::memory_order_acquire)); + EXPECT_FALSE(manager.isAbortCleanup(comm)); +} + +TEST(NcclCommResourceManagerStandaloneTest, NormalCleanupOwnsRetirementUntilCallbacksComplete) +{ + auto& manager = nccl_util::NcclCommResourceManager::getInstance(); + int handleStorage = 0; + auto const comm = reinterpret_cast(&handleStorage); + std::atomic cleanupEntered{false}; + std::atomic releaseCleanup{false}; + std::atomic cleanupSucceeded{false}; + manager.registerResource( + comm, + [&]() + { + cleanupEntered.store(true, std::memory_order_release); + while (!releaseCleanup.load(std::memory_order_acquire)) + { + std::this_thread::yield(); + } + return true; + }, + "BlockingNormalCleanup"); + std::thread cleanupOwner( + [&]() { cleanupSucceeded.store(manager.cleanupResources(comm, false), std::memory_order_release); }); + while (!cleanupEntered.load(std::memory_order_acquire)) + { + std::this_thread::yield(); + } + + manager.beginAbortCleanup(comm); + EXPECT_TRUE(manager.cleanupResources(comm, true)); + EXPECT_TRUE(manager.isAbortCleanup(comm)); + releaseCleanup.store(true, std::memory_order_release); + cleanupOwner.join(); + EXPECT_TRUE(cleanupSucceeded.load(std::memory_order_acquire)); + EXPECT_FALSE(manager.isAbortCleanup(comm)); +} + class NcclCommResourceManagerTest : public ::testing::Test { protected: @@ -170,12 +284,18 @@ TEST_F(NcclCommResourceManagerTest, ResourceRegistration) auto& manager = nccl_util::NcclCommResourceManager::getInstance(); // Create a separate comm using split for this test - auto testComm = createSplitComm(*mComm, 0, mRank); + auto testComm = createSplitComm(mComm, 0, mRank); // Register a resource bool cleanupCalled = false; manager.registerResource( - *testComm, [&cleanupCalled]() { cleanupCalled = true; }, "TestResource"); + *testComm, + [&cleanupCalled]() + { + cleanupCalled = true; + return true; + }, + "TestResource"); EXPECT_TRUE(manager.hasResources(*testComm)); EXPECT_EQ(manager.getResourceCount(*testComm), 1); @@ -202,15 +322,33 @@ TEST_F(NcclCommResourceManagerTest, MultipleResources) auto& manager = nccl_util::NcclCommResourceManager::getInstance(); // Create a separate comm using split for this test - auto testComm = createSplitComm(*mComm, 0, mRank); + auto testComm = createSplitComm(mComm, 0, mRank); std::vector cleanupOrder; manager.registerResource( - *testComm, [&cleanupOrder]() { cleanupOrder.push_back(1); }, "Resource1"); + *testComm, + [&cleanupOrder]() + { + cleanupOrder.push_back(1); + return true; + }, + "Resource1"); manager.registerResource( - *testComm, [&cleanupOrder]() { cleanupOrder.push_back(2); }, "Resource2"); + *testComm, + [&cleanupOrder]() + { + cleanupOrder.push_back(2); + return true; + }, + "Resource2"); manager.registerResource( - *testComm, [&cleanupOrder]() { cleanupOrder.push_back(3); }, "Resource3"); + *testComm, + [&cleanupOrder]() + { + cleanupOrder.push_back(3); + return true; + }, + "Resource3"); EXPECT_EQ(manager.getResourceCount(*testComm), 3); @@ -229,22 +367,34 @@ TEST_F(NcclCommResourceManagerTest, ResourceCount) auto& manager = nccl_util::NcclCommResourceManager::getInstance(); // Create a separate comm using split for this test - auto testComm = createSplitComm(*mComm, 0, mRank); + auto testComm = createSplitComm(mComm, 0, mRank); EXPECT_FALSE(manager.hasResources(*testComm)); EXPECT_EQ(manager.getResourceCount(*testComm), 0); manager.registerResource( - *testComm, []() {}, "Test1"); + *testComm, []() { return true; }, "Test1"); EXPECT_EQ(manager.getResourceCount(*testComm), 1); manager.registerResource( - *testComm, []() {}, "Test2"); + *testComm, []() { return true; }, "Test2"); EXPECT_EQ(manager.getResourceCount(*testComm), 2); testComm.reset(); } +TEST_F(NcclCommResourceManagerTest, CleanupFailureIsReportedToCommunicatorOwner) +{ + auto& manager = nccl_util::NcclCommResourceManager::getInstance(); + auto testComm = createSplitComm(mComm, 0, mRank); + + manager.registerResource( + *testComm, []() { return false; }, "InjectedCleanupFailure"); + + EXPECT_FALSE(manager.cleanupResources(*testComm)); + EXPECT_FALSE(manager.hasResources(*testComm)); +} + //============================================================================== // NCCLWindowAllocator Tests //============================================================================== @@ -302,7 +452,7 @@ TEST_F(NCCLWindowAllocatorTest, BasicAllocation) auto& allocator = nccl_util::NCCLWindowAllocator::getInstance(); const size_t bufferSize = 1024 * 1024; // 1MB - auto buffer = allocator.requestBuffer(*mComm, bufferSize); + auto buffer = allocator.requestBuffer(mComm, bufferSize); EXPECT_TRUE(buffer.isValid()); EXPECT_NE(buffer.ptr, nullptr); @@ -326,7 +476,7 @@ TEST_F(NCCLWindowAllocatorTest, BufferReuse) const size_t bufferSize = 512 * 1024; // 512KB // Allocate first buffer - auto buffer1 = allocator.requestBuffer(*mComm, bufferSize); + auto buffer1 = allocator.requestBuffer(mComm, bufferSize); EXPECT_TRUE(buffer1.isValid()); void* ptr1 = buffer1.ptr; @@ -334,7 +484,7 @@ TEST_F(NCCLWindowAllocatorTest, BufferReuse) allocator.releaseBuffer(*mComm, ptr1); // Request another buffer of the same size - should reuse - auto buffer2 = allocator.requestBuffer(*mComm, bufferSize); + auto buffer2 = allocator.requestBuffer(mComm, bufferSize); EXPECT_TRUE(buffer2.isValid()); EXPECT_EQ(buffer2.ptr, ptr1); // Should be the same buffer @@ -346,9 +496,9 @@ TEST_F(NCCLWindowAllocatorTest, BestFitReuse) auto& allocator = nccl_util::NCCLWindowAllocator::getInstance(); // Allocate buffers of different sizes - auto buffer1MB = allocator.requestBuffer(*mComm, 1024 * 1024); - auto buffer2MB = allocator.requestBuffer(*mComm, 2 * 1024 * 1024); - auto buffer512KB = allocator.requestBuffer(*mComm, 512 * 1024); + auto buffer1MB = allocator.requestBuffer(mComm, 1024 * 1024); + auto buffer2MB = allocator.requestBuffer(mComm, 2 * 1024 * 1024); + auto buffer512KB = allocator.requestBuffer(mComm, 512 * 1024); void* ptr1MB = buffer1MB.ptr; void* ptr2MB = buffer2MB.ptr; @@ -360,7 +510,7 @@ TEST_F(NCCLWindowAllocatorTest, BestFitReuse) allocator.releaseBuffer(*mComm, ptr512KB); // Request 768KB - should reuse 1MB (best fit, smallest that fits) - auto buffer768KB = allocator.requestBuffer(*mComm, 768 * 1024); + auto buffer768KB = allocator.requestBuffer(mComm, 768 * 1024); EXPECT_TRUE(buffer768KB.isValid()); EXPECT_EQ(buffer768KB.ptr, ptr1MB); // Should reuse 1MB buffer EXPECT_EQ(buffer768KB.size, 1024 * 1024); // Original size @@ -371,7 +521,7 @@ TEST_F(NCCLWindowAllocatorTest, BestFitReuse) TEST_F(NCCLWindowAllocatorTest, FailureCacheIsSizeAwareForNewAllocations) { auto& allocator = nccl_util::NCCLWindowAllocator::getInstance(); - auto testComm = createSplitComm(*mComm, 0, mRank); + auto testComm = createSplitComm(mComm, 0, mRank); constexpr size_t failureSize = 1024 * 1024; nccl_util::NCCLWindowAllocatorTestAccess::recordSymmetricFailure(allocator, *testComm, failureSize); @@ -391,7 +541,7 @@ TEST_F(NCCLWindowAllocatorTest, FailureCacheIsSizeAwareForNewAllocations) TEST_F(NCCLWindowAllocatorTest, FailureCacheDoesNotDisableReusableBuffers) { auto& allocator = nccl_util::NCCLWindowAllocator::getInstance(); - auto testComm = createSplitComm(*mComm, 0, mRank); + auto testComm = createSplitComm(mComm, 0, mRank); auto buffer1MB = allocator.requestBuffer(*testComm, 1024 * 1024); ASSERT_TRUE(buffer1MB.isValid()); @@ -416,7 +566,7 @@ TEST_F(NCCLWindowAllocatorTest, FailureCacheDoesNotDisableReusableBuffers) TEST_F(NCCLWindowAllocatorTest, FailureCacheKeepsSmallestFailureSize) { auto& allocator = nccl_util::NCCLWindowAllocator::getInstance(); - auto testComm = createSplitComm(*mComm, 0, mRank); + auto testComm = createSplitComm(mComm, 0, mRank); nccl_util::NCCLWindowAllocatorTestAccess::recordSymmetricFailure(allocator, *testComm, 2 * 1024 * 1024); nccl_util::NCCLWindowAllocatorTestAccess::recordSymmetricFailure(allocator, *testComm, 1024 * 1024); @@ -459,7 +609,7 @@ TEST_F(NCCLWindowAllocatorTest, MultipleBuffers) // Allocate multiple buffers for (int i = 0; i < 5; ++i) { - auto buffer = allocator.requestBuffer(*mComm, bufferSize); + auto buffer = allocator.requestBuffer(mComm, bufferSize); EXPECT_TRUE(buffer.isValid()); ptrs.push_back(buffer.ptr); } @@ -482,7 +632,7 @@ TEST_F(NCCLWindowAllocatorTest, SearchBuffer) auto& allocator = nccl_util::NCCLWindowAllocator::getInstance(); const size_t bufferSize = 128 * 1024; - auto buffer = allocator.requestBuffer(*mComm, bufferSize); + auto buffer = allocator.requestBuffer(mComm, bufferSize); // Test searchBuffer auto found = allocator.searchBuffer(*mComm, buffer.ptr); @@ -505,7 +655,7 @@ TEST_F(NCCLWindowAllocatorTest, GetWindowAndSize) auto& allocator = nccl_util::NCCLWindowAllocator::getInstance(); const size_t bufferSize = 64 * 1024; - auto buffer = allocator.requestBuffer(*mComm, bufferSize); + auto buffer = allocator.requestBuffer(mComm, bufferSize); // Test getWindow auto window = allocator.getWindow(*mComm, buffer.ptr); @@ -530,7 +680,7 @@ TEST_F(NCCLWindowAllocatorTest, GetBufferInfo) auto& allocator = nccl_util::NCCLWindowAllocator::getInstance(); const size_t bufferSize = 32 * 1024; - auto buffer = allocator.requestBuffer(*mComm, bufferSize); + auto buffer = allocator.requestBuffer(mComm, bufferSize); auto info = allocator.getBufferInfo(*mComm, buffer.ptr); EXPECT_TRUE(info.isValid()); @@ -570,7 +720,7 @@ TEST_F(NCCLWindowAllocatorTest, CleanupOnCommDestroy) auto& allocator = nccl_util::NCCLWindowAllocator::getInstance(); // Create a separate comm using split for this test - auto testComm = createSplitComm(*mComm, 0, mRank); + auto testComm = createSplitComm(mComm, 0, mRank); // Store the raw comm value before destruction ncclComm_t rawComm = *testComm; @@ -627,8 +777,8 @@ TEST_F(NCCLWindowAllocatorTest, MultipleComms) auto& allocator = nccl_util::NCCLWindowAllocator::getInstance(); // Create two different communicators using split (different colors) - auto comm1 = createSplitComm(*mComm, 0, mRank); - auto comm2 = createSplitComm(*mComm, 1, mRank); + auto comm1 = createSplitComm(mComm, 0, mRank); + auto comm2 = createSplitComm(mComm, 1, mRank); const size_t bufferSize = 4 * 1024; @@ -831,6 +981,25 @@ TEST_F(CreateNCCLWindowTensorTest, TensorDeleterReleasesBuffer) EXPECT_GE(allocator.getBufferCount(*mComm), 1); } +TEST_F(CreateNCCLWindowTensorTest, TensorKeepsGenericCommunicatorOwnerAlive) +{ + if (tensorrt_llm::isNcclFaultToleranceEnabled()) + { + GTEST_SKIP() << "Generic communicator holders use the legacy default-off allocator contract"; + } + auto comm = createSplitComm(mComm, 0, mRank); + std::weak_ptr weakComm = comm; + std::vector shape = {16, 16}; + auto [tensor, buffer] = nccl_util::createNCCLWindowTensor(comm, shape, torch::kFloat32); + ASSERT_TRUE(tensor.defined()); + ASSERT_TRUE(buffer.isValid()); + + comm.reset(); + EXPECT_FALSE(weakComm.expired()); + tensor = torch::Tensor(); + EXPECT_TRUE(weakComm.expired()); +} + TEST_F(CreateNCCLWindowTensorTest, MultipleTensors) { using nccl_util::createNCCLWindowTensor; diff --git a/tensorrt_llm/_mnnvl_utils.py b/tensorrt_llm/_mnnvl_utils.py index af4157786c06..334438484fec 100644 --- a/tensorrt_llm/_mnnvl_utils.py +++ b/tensorrt_llm/_mnnvl_utils.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -18,7 +18,7 @@ import platform import sys from dataclasses import dataclass -from typing import List, Optional, Union +from typing import List, Optional, Tuple, Union import pynvml import torch @@ -400,19 +400,50 @@ class HelixCpMnnvlMemory(MnnvlMemory): its own isolated state separate from MnnvlMemory. """ + comm_topology = None + + @staticmethod + def get_comm_topology( + mapping: Mapping, + ) -> Tuple[int, int, int, Tuple[int, ...]]: + """Return the exact local CP topology used by ``MPI_Comm_split``.""" + return ( + int(mapping.world_size), + int(mapping.rank), + int(mapping.cp_rank), + tuple(int(rank) for rank in mapping.cp_group), + ) + @classmethod def get_comm(cls, mapping: Mapping): """Get CP-based communicator (ranks grouped by PP+TP+MOE_TP, ordered by CP rank).""" + topology = cls.get_comm_topology(mapping) if cls.comm is not None: + if cls.comm_topology != topology: + raise RuntimeError( + "Helix MNNVL CP communicator is already initialized for " + f"topology {cls.comm_topology}; cannot reuse it for " + f"incompatible topology {topology}" + ) return cls.comm comm = mpi_comm().Split( mapping.pp_rank * mapping.tp_size + mapping.tp_rank, mapping.cp_rank, ) cls.comm = comm + cls.comm_topology = topology return comm +def get_helix_cp_mnnvl_topology( + mapping: Mapping, +) -> Optional[Tuple[int, int, int, Tuple[int, ...]]]: + """Return the Helix MNNVL topology relevant to communicator sharing.""" + if mapping.has_cp_helix() and not mapping.cp_config.get("use_nccl_for_alltoall", True): + return HelixCpMnnvlMemory.get_comm_topology(mapping) + return None + + def init_helix_cp_comm(mapping: Mapping) -> None: """Pre-initialize the Helix CP communicator. diff --git a/tensorrt_llm/_torch/alltoall_watchdog.py b/tensorrt_llm/_torch/alltoall_watchdog.py new file mode 100644 index 000000000000..2b71aaefa257 --- /dev/null +++ b/tensorrt_llm/_torch/alltoall_watchdog.py @@ -0,0 +1,783 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. +"""Host-side watchdog for MoE AlltoAll completion flags. + +The NVLinkOneSided kernels signal each collective by writing the current +``flag_val`` into the rank-local completion flag table. A dead peer in the +silent-spin failure mode never writes its slot, so this watchdog polls the same +table from a CPU thread and reports peers whose flags do not reach the expected +generation before a bounded timeout. + +Timeouts are detection events, not committed membership changes. The optional +EP health object supplies the already-committed active peers to watch, and is +read-only here. Higher-layer recovery coordination consumes ``on_timeout`` +events and commits a new membership only after the data plane and expert +placement are ready for the same generation. +""" + +from __future__ import annotations + +import threading +import time +from collections import deque +from collections.abc import Callable, Mapping, MutableMapping, Sequence +from dataclasses import dataclass +from typing import Protocol + +import torch + +from tensorrt_llm._utils import prefer_pinned +from tensorrt_llm.logger import logger as tllm_logger + +DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S = 5.0 +DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S = 0.1 +UNKNOWN_COMPLETION_FLAG = -(2**63) +_COMPLETION_FLAG_MASK = (1 << 32) - 1 +_COMPLETION_FLAG_HALF_RANGE = 1 << 31 +_ACTIVE_RANK_MASK_WORD_BITS = 64 +_ACTIVE_RANK_MASK_WORDS = 2 +_ACTIVE_RANK_MASK_WORD_MASK = (1 << _ACTIVE_RANK_MASK_WORD_BITS) - 1 +_WORKSPACE_WATCHDOG_STATE_KEY = "alltoall_watchdog_shared_state" +_WORKSPACE_WATCHDOG_STATE_INIT_LOCK = threading.Lock() + + +def _normalize_completion_flag(value: int) -> int: + return int(value) & _COMPLETION_FLAG_MASK + + +def _completion_flag_reached(observed: int, expected: int) -> bool: + if observed == UNKNOWN_COMPLETION_FLAG: + return False + # Counter values are ordered modulo uint32. A queued watch cannot lag by + # half the counter space within the bounded watchdog timeout. + distance = (_normalize_completion_flag(observed) - expected) & _COMPLETION_FLAG_MASK + return distance < _COMPLETION_FLAG_HALF_RANGE + + +class CompletionFlagReader(Protocol): + """Reads one phase's rank-local completion flag row.""" + + def read_completion_flags(self, phase: str) -> Sequence[int]: + """Return ``ep_size`` flag values for ``phase``.""" + + +class EPGroupHealthSnapshotLike(Protocol): + """Atomic fields required from a committed-membership snapshot.""" + + @property + def mask(self) -> int: + """Return the committed active-rank bitmask.""" + + @property + def generation(self) -> int: + """Return the matching committed-membership generation.""" + + +class EPGroupHealthLike(Protocol): + """Read-only committed EP membership used by AlltoAll frontends.""" + + def get_mask(self) -> int: + """Return the active-rank bitmask.""" + + def get_mask_words(self) -> tuple[int, ...]: + """Return the active-rank bitmask split into uint64 words.""" + + def snapshot(self) -> EPGroupHealthSnapshotLike: + """Return one atomic committed mask and generation snapshot.""" + + +@dataclass(frozen=True) +class ActiveRankMaskSnapshot: + """One dispatch epoch's rank mask and optional committed generation. + + ``committed_generation`` is populated only when ``active_rank_mask`` was + read from the coordinator-owned EP health object. Explicit caller masks do + not inherit a generation from unrelated health state. + """ + + active_rank_mask: torch.Tensor | None + committed_generation: int | None + + +class CompletionFlagReadTimeout(TimeoutError): + """Raised when the host watchdog cannot read completion flags in time.""" + + +@dataclass(frozen=True) +class AlltoAllWatchdogTimeout: + """Detection details emitted when an AlltoAll phase times out. + + ``missing_ranks`` contains suspects relative to the committed active mask + captured for this collective. Receiving this event does not publish a new + active mask. + """ + + phase: str + expected_flag: int + observed_flags: tuple[int, ...] + missing_ranks: tuple[int, ...] + elapsed_s: float + poll_timed_out: bool = False + + +@dataclass(frozen=True) +class _CollectiveWatch: + phase: str + expected_flag: int + active_mask: int + start_s: float + + +@dataclass(frozen=True) +class _SharedWatchdogConfig: + ep_size: int + ep_rank: int + timeout_s: float + poll_interval_s: float + health: EPGroupHealthLike | None + on_timeout: Callable[[AlltoAllWatchdogTimeout], None] | None + + +@dataclass +class _WorkspaceWatchdogState: + lock: threading.Lock + generation: int = 0 + watchdog: AlltoAllWatchdog | None = None + config: _SharedWatchdogConfig | None = None + ref_count: int = 0 + + +class _TorchCompletionFlagReader: + """Completion-flag reader backed by the MoE AlltoAll workspace tensor.""" + + def __init__( + self, + workspace: torch.Tensor, + ep_rank: int, + ep_size: int, + dispatch_completion_flags_offset: int, + combine_completion_flags_offset: int, + device_copy_timeout_s: float = DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, + ) -> None: + if workspace.dim() != 2: + raise ValueError("workspace must be a 2D tensor [ep_size, size_per_rank]") + if not 0 <= ep_rank < ep_size: + raise ValueError(f"ep_rank must be in [0, {ep_size}), got {ep_rank}") + if workspace.size(0) != ep_size: + raise ValueError( + f"workspace first dimension must equal ep_size={ep_size}, got {workspace.size(0)}" + ) + self._workspace = workspace + self._ep_rank = ep_rank + self._ep_size = ep_size + self._offsets = { + "dispatch": int(dispatch_completion_flags_offset), + "combine": int(combine_completion_flags_offset), + } + self._device_copy_timeout_s = float(device_copy_timeout_s) + self._copy_stream: torch.cuda.Stream | None = None + self._host_flags: torch.Tensor | None = None + self._copy_event: torch.cuda.Event | None = None + self._retired_copies: list[tuple[torch.Tensor, torch.cuda.Event]] = [] + if workspace.device.type == "cuda": + self._copy_stream = torch.cuda.Stream(device=workspace.device) + + def _prune_retired_copies(self) -> None: + self._retired_copies = [ + (host_flags, event) for host_flags, event in self._retired_copies if not event.query() + ] + + def _read_cuda_flags(self, flags: torch.Tensor) -> tuple[int, ...]: + assert self._copy_stream is not None + self._prune_retired_copies() + + if self._host_flags is None: + self._host_flags = torch.empty( + (self._ep_size,), + dtype=torch.int32, + device="cpu", + pin_memory=prefer_pinned(), + ) + if self._copy_event is None: + self._copy_event = torch.cuda.Event(blocking=False) + host_flags = self._host_flags + event = self._copy_event + with torch.cuda.device(flags.device), torch.cuda.stream(self._copy_stream): + host_flags.copy_(flags.detach(), non_blocking=True) + event.record(self._copy_stream) + + deadline_s = time.monotonic() + self._device_copy_timeout_s + while not event.query(): + remaining_s = deadline_s - time.monotonic() + if remaining_s <= 0: + self._retired_copies.append((host_flags, event)) + self._host_flags = None + self._copy_event = None + raise CompletionFlagReadTimeout( + "timed out copying AlltoAll completion flags to host" + ) + time.sleep(min(remaining_s, 0.001)) + + return tuple(int(v) for v in host_flags.tolist()) + + def read_completion_flags(self, phase: str) -> tuple[int, ...]: + offset = self._offsets[phase] + end = offset + self._ep_size * 4 + flags = self._workspace[self._ep_rank, offset:end].view(torch.int32) + if flags.device.type == "cuda": + return self._read_cuda_flags(flags) + if flags.device.type != "cpu": + flags = flags.detach().cpu() + return tuple(int(v) for v in flags.tolist()) + + +class AlltoAllWatchdogCoordinator: + """Shared watchdog plumbing for MoE AlltoAll frontends. + + ``health`` is the committed membership snapshot source. This class reads + it when a dispatch starts but never publishes detected failures into it. + """ + + def __init__( + self, + *, + workspace_state: MutableMapping[str, object], + workspace: torch.Tensor, + metainfo: torch.Tensor, + metainfo_index: Mapping[str, int], + ep_rank: int, + health: EPGroupHealthLike | None = None, + ) -> None: + self._workspace_state = workspace_state + self._workspace = workspace + self._metainfo = metainfo + self._metainfo_index = metainfo_index + self._ep_rank = ep_rank + self._health = health + + def _find_shared_state(self) -> _WorkspaceWatchdogState | None: + state = self._workspace_state.get(_WORKSPACE_WATCHDOG_STATE_KEY) + if state is None: + return None + if not isinstance(state, _WorkspaceWatchdogState): + raise TypeError("invalid shared AlltoAll watchdog workspace state") + return state + + def _get_shared_state(self) -> _WorkspaceWatchdogState: + state = self._find_shared_state() + if state is not None: + return state + with _WORKSPACE_WATCHDOG_STATE_INIT_LOCK: + state = self._workspace_state.get(_WORKSPACE_WATCHDOG_STATE_KEY) + if state is None: + state = _WorkspaceWatchdogState(lock=threading.Lock()) + self._workspace_state[_WORKSPACE_WATCHDOG_STATE_KEY] = state + if not isinstance(state, _WorkspaceWatchdogState): + raise TypeError("invalid shared AlltoAll watchdog workspace state") + return state + + def read_current_flag_val(self) -> int: + flag_val_offset = int(self._metainfo[self._metainfo_index["FLAG_VAL_OFFSET_INDEX"]].item()) + flag_val = self._workspace[self._ep_rank, flag_val_offset : flag_val_offset + 4].view( + torch.int32 + ) + if flag_val.device.type != "cpu": + flag_val = flag_val.detach().cpu() + return _normalize_completion_flag(int(flag_val.item())) + + def capture_active_rank_mask( + self, active_rank_mask: torch.Tensor | None + ) -> ActiveRankMaskSnapshot: + """Capture a caller override or one coherent committed-mask epoch. + + The returned tensor does not alias a caller-owned override, so it can + safely represent the mask used by one dispatch/combine pair. When the + mask comes from ``health``, derive both fixed-width ABI words from one + atomic mask/generation snapshot. + """ + if active_rank_mask is not None: + return ActiveRankMaskSnapshot( + active_rank_mask=active_rank_mask.detach().clone(), + committed_generation=None, + ) + if self._health is None: + return ActiveRankMaskSnapshot( + active_rank_mask=None, + committed_generation=None, + ) + + health_snapshot = self._health.snapshot() + mask_words = tuple( + (health_snapshot.mask >> (_ACTIVE_RANK_MASK_WORD_BITS * index)) + & _ACTIVE_RANK_MASK_WORD_MASK + for index in range(_ACTIVE_RANK_MASK_WORDS) + ) + return ActiveRankMaskSnapshot( + active_rank_mask=torch.tensor( + mask_words, + dtype=torch.uint64, + device="cpu", + ), + committed_generation=health_snapshot.generation, + ) + + def active_rank_mask_tensor(self, active_rank_mask: torch.Tensor | None) -> torch.Tensor | None: + """Return only the tensor from :meth:`capture_active_rank_mask`. + + Dispatch frontends must retain the full snapshot so combine can verify + the committed generation before launching its kernel. + """ + return self.capture_active_rank_mask(active_rank_mask).active_rank_mask + + def active_rank_mask_for_combine( + self, + dispatch_snapshot: ActiveRankMaskSnapshot, + requested_active_rank_mask: torch.Tensor | None, + ) -> torch.Tensor | None: + """Validate the dispatch epoch and return its captured rank mask. + + A coordinator-owned generation change between dispatch and combine is + an invalid asynchronous commit, so fail closed before launching the + combine kernel. Explicit caller masks are independent of health and + therefore carry no committed generation. + """ + dispatch_active_rank_mask = dispatch_snapshot.active_rank_mask + committed_generation = dispatch_snapshot.committed_generation + if committed_generation is not None: + if self._health is None: + raise RuntimeError("committed rank-mask snapshot has no EP health source") + current_generation = self._health.snapshot().generation + if current_generation != committed_generation: + raise RuntimeError( + "committed EP membership changed between dispatch and combine " + f"(generation {committed_generation} -> {current_generation}); " + "aborting collective epoch" + ) + + if requested_active_rank_mask is None: + return dispatch_active_rank_mask + requested_mask = self.active_rank_mask_tensor(requested_active_rank_mask) + assert requested_mask is not None + if dispatch_active_rank_mask is None or not torch.equal( + dispatch_active_rank_mask, requested_mask + ): + raise ValueError("active_rank_mask must match the mask captured at dispatch") + return dispatch_active_rank_mask + + def active_mask_int(self, active_rank_mask: torch.Tensor | None) -> int | None: + if active_rank_mask is not None: + mask_cpu = active_rank_mask.detach().cpu() + return sum(int(word) << (64 * idx) for idx, word in enumerate(mask_cpu.tolist())) + if self._health is not None: + return self._health.get_mask() + return None + + def acquire_watchdog( + self, + *, + ep_size: int, + timeout_s: float, + poll_interval_s: float, + on_timeout: Callable[[AlltoAllWatchdogTimeout], None] | None = None, + ) -> AlltoAllWatchdog: + config = _SharedWatchdogConfig( + ep_size=int(ep_size), + ep_rank=self._ep_rank, + timeout_s=float(timeout_s), + poll_interval_s=float(poll_interval_s), + health=self._health, + on_timeout=on_timeout, + ) + state = self._get_shared_state() + with state.lock: + if state.watchdog is None: + state.generation = self.read_current_flag_val() + state.watchdog = AlltoAllWatchdog.from_workspace( + workspace=self._workspace, + metainfo=self._metainfo, + metainfo_index=self._metainfo_index, + ep_rank=self._ep_rank, + ep_size=ep_size, + timeout_s=timeout_s, + poll_interval_s=poll_interval_s, + health=self._health, + on_timeout=on_timeout, + ) + state.config = config + elif not self._matching_config(state.config, config): + raise ValueError( + "AlltoAll wrappers sharing a workspace must use the same " + "watchdog configuration and EP health object" + ) + state.ref_count += 1 + watchdog = state.watchdog + assert watchdog is not None + return watchdog + + @staticmethod + def _matching_config( + existing: _SharedWatchdogConfig | None, + requested: _SharedWatchdogConfig, + ) -> bool: + return ( + existing is not None + and existing.ep_size == requested.ep_size + and existing.ep_rank == requested.ep_rank + and existing.timeout_s == requested.timeout_s + and existing.poll_interval_s == requested.poll_interval_s + and existing.health is requested.health + and existing.on_timeout is requested.on_timeout + ) + + def release_watchdog(self, watchdog: AlltoAllWatchdog) -> None: + state = self._get_shared_state() + watchdog_to_stop: AlltoAllWatchdog | None = None + with state.lock: + if state.watchdog is not watchdog or state.ref_count <= 0: + raise RuntimeError("attempted to release an unregistered AlltoAll watchdog") + state.ref_count -= 1 + if state.ref_count == 0: + watchdog_to_stop = state.watchdog + state.watchdog = None + state.config = None + if watchdog_to_stop is not None: + watchdog_to_stop.stop() + + def watch_collective( + self, + watchdog: AlltoAllWatchdog | None, + phase: str, + active_rank_mask: torch.Tensor | None, + ) -> None: + if watchdog is None: + state = self._find_shared_state() + if state is not None: + with state.lock: + if state.watchdog is not None: + state.generation = (state.generation + 1) & _COMPLETION_FLAG_MASK + return + active_mask = self.active_mask_int(active_rank_mask) + state = self._get_shared_state() + with state.lock: + if state.watchdog is not watchdog: + raise RuntimeError("AlltoAll watchdog is not registered for this workspace") + state.generation = (state.generation + 1) & _COMPLETION_FLAG_MASK + watchdog.watch( + phase=phase, + expected_flag=state.generation, + active_mask=active_mask, + ) + + +class AlltoAllWatchdog: + """Background host thread that watches AlltoAll completion flags. + + The watchdog is intentionally opt-in. Callers queue phases with + :meth:`watch`; the thread polls them in FIFO order so a queued combine cannot + hide a still-spinning dispatch. A timeout is reported through + ``on_timeout`` without mutating the committed EP membership. + """ + + VALID_PHASES = frozenset({"dispatch", "combine"}) + + def __init__( + self, + *, + ep_size: int, + ep_rank: int, + completion_reader: CompletionFlagReader, + timeout_s: float = DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S, + poll_interval_s: float = DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, + health: EPGroupHealthLike | None = None, + on_timeout: Callable[[AlltoAllWatchdogTimeout], None] | None = None, + ) -> None: + if ep_size <= 0: + raise ValueError(f"ep_size must be > 0, got {ep_size}") + if not 0 <= ep_rank < ep_size: + raise ValueError(f"ep_rank must be in [0, {ep_size}), got {ep_rank}") + if timeout_s <= 0: + raise ValueError(f"timeout_s must be > 0, got {timeout_s}") + if poll_interval_s <= 0: + raise ValueError(f"poll_interval_s must be > 0, got {poll_interval_s}") + + self._ep_size = int(ep_size) + self._ep_rank = int(ep_rank) + self._completion_reader = completion_reader + self._timeout_s = float(timeout_s) + self._poll_interval_s = float(poll_interval_s) + self._health = health + self._on_timeout = on_timeout + + self._cv = threading.Condition() + self._queue: deque[_CollectiveWatch] = deque() + self._closed = False + self._stopping = False + self._thread: threading.Thread | None = None + self._last_error: BaseException | None = None + + @classmethod + def from_workspace( + cls, + *, + workspace: torch.Tensor, + metainfo: torch.Tensor, + metainfo_index: Mapping[str, int], + ep_rank: int, + ep_size: int, + timeout_s: float = DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S, + poll_interval_s: float = DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, + health: EPGroupHealthLike | None = None, + on_timeout: Callable[[AlltoAllWatchdogTimeout], None] | None = None, + ) -> "AlltoAllWatchdog": + """Build a watchdog from the MoE AlltoAll workspace and metainfo.""" + dispatch_offset = int( + metainfo[metainfo_index["DISPATCH_COMPLETION_FLAGS_OFFSET_INDEX"]].item() + ) + combine_offset = int( + metainfo[metainfo_index["COMBINE_COMPLETION_FLAGS_OFFSET_INDEX"]].item() + ) + reader = _TorchCompletionFlagReader( + workspace=workspace, + ep_rank=ep_rank, + ep_size=ep_size, + dispatch_completion_flags_offset=dispatch_offset, + combine_completion_flags_offset=combine_offset, + device_copy_timeout_s=poll_interval_s, + ) + return cls( + ep_size=ep_size, + ep_rank=ep_rank, + completion_reader=reader, + timeout_s=timeout_s, + poll_interval_s=poll_interval_s, + health=health, + on_timeout=on_timeout, + ) + + @property + def last_error(self) -> BaseException | None: + """Return the last polling-thread error, if any.""" + with self._cv: + return self._last_error + + def start(self) -> None: + """Start the background polling thread. Idempotent.""" + with self._cv: + if self._closed: + raise RuntimeError("cannot start a stopped AlltoAllWatchdog") + if self._thread is not None and self._thread.is_alive(): + return + self._stopping = False + self._thread = threading.Thread( + target=self._run, + name=f"AlltoAllWatchdog-rank{self._ep_rank}", + daemon=True, + ) + self._thread.start() + + def stop(self, timeout_s: float | None = None) -> None: + """Stop the polling thread and wait for it to exit.""" + with self._cv: + self._closed = True + self._stopping = True + self._queue.clear() + self._cv.notify_all() + thread = self._thread + if thread is not None and thread is not threading.current_thread(): + thread.join(timeout=timeout_s) + + def watch( + self, + *, + phase: str, + expected_flag: int, + active_mask: int | None = None, + ) -> None: + """Queue a just-launched AlltoAll phase for watchdog polling.""" + if phase not in self.VALID_PHASES: + raise ValueError(f"phase must be one of {sorted(self.VALID_PHASES)}, got {phase!r}") + if not 0 <= expected_flag <= _COMPLETION_FLAG_MASK: + raise ValueError( + f"expected_flag must be in [0, {_COMPLETION_FLAG_MASK}], got {expected_flag}" + ) + if active_mask is None: + if self._health is not None: + active_mask = self._health.get_mask() + else: + active_mask = (1 << self._ep_size) - 1 + if not (active_mask >> self._ep_rank) & 1: + raise ValueError("active_mask must include the local ep_rank") + + self.start() + with self._cv: + if self._closed: + raise RuntimeError("cannot queue a stopped AlltoAllWatchdog") + self._queue.append( + _CollectiveWatch( + phase=phase, + expected_flag=int(expected_flag), + active_mask=int(active_mask), + start_s=time.monotonic(), + ) + ) + self._cv.notify_all() + + def wait_until_idle(self, timeout_s: float) -> bool: + """Wait until all queued phases complete or timeout handling clears them.""" + deadline = time.monotonic() + timeout_s + with self._cv: + while self._queue: + remaining = deadline - time.monotonic() + if remaining <= 0: + return False + self._cv.wait(timeout=remaining) + return True + + def __enter__(self) -> "AlltoAllWatchdog": + self.start() + return self + + def __exit__(self, exc_type, exc, tb) -> None: + self.stop(timeout_s=1.0) + + def _active_ranks(self, active_mask: int) -> tuple[int, ...]: + return tuple(rank for rank in range(self._ep_size) if (active_mask >> rank) & 1) + + def _phase_complete(self, watch: _CollectiveWatch, observed_flags: tuple[int, ...]) -> bool: + return all( + _completion_flag_reached(observed_flags[rank], watch.expected_flag) + for rank in self._active_ranks(watch.active_mask) + ) + + def _missing_ranks( + self, watch: _CollectiveWatch, observed_flags: tuple[int, ...] + ) -> tuple[int, ...]: + return tuple( + rank + for rank in self._active_ranks(watch.active_mask) + if not _completion_flag_reached(observed_flags[rank], watch.expected_flag) + ) + + def _handle_timeout( + self, + watch: _CollectiveWatch, + observed_flags: tuple[int, ...], + *, + poll_timed_out: bool = False, + ) -> None: + elapsed_s = time.monotonic() - watch.start_s + missing_ranks = self._missing_ranks(watch, observed_flags) + event = AlltoAllWatchdogTimeout( + phase=watch.phase, + expected_flag=watch.expected_flag, + observed_flags=observed_flags, + missing_ranks=missing_ranks, + elapsed_s=elapsed_s, + poll_timed_out=poll_timed_out, + ) + if poll_timed_out: + tllm_logger.error( + "AlltoAll watchdog could not read completion flags on rank %d " + "during %s before timeout %.3fs; expected flag %d, active " + "ranks %s, observed flags %s; reporting detection event only", + self._ep_rank, + watch.phase, + elapsed_s, + watch.expected_flag, + list(self._active_ranks(watch.active_mask)), + list(observed_flags), + ) + else: + tllm_logger.warning( + "AlltoAll watchdog timeout on rank %d during %s: expected flag %d, " + "missing ranks %s, observed flags %s", + self._ep_rank, + watch.phase, + watch.expected_flag, + list(missing_ranks), + list(observed_flags), + ) + if self._on_timeout is not None: + self._on_timeout(event) + + def _stop_after_error(self, exc: Exception) -> None: + with self._cv: + self._last_error = exc + self._closed = True + self._stopping = True + self._queue.clear() + self._cv.notify_all() + + def _run(self) -> None: + last_observed_flags = tuple(UNKNOWN_COMPLETION_FLAG for _ in range(self._ep_size)) + poll_timed_out = False + while True: + with self._cv: + while not self._queue and not self._stopping: + self._cv.wait() + if self._stopping: + return + watch = self._queue[0] + + try: + observed_flags = tuple( + _normalize_completion_flag(int(v)) + for v in self._completion_reader.read_completion_flags(watch.phase) + ) + if len(observed_flags) != self._ep_size: + raise RuntimeError( + f"completion reader returned {len(observed_flags)} flags; " + f"expected ep_size={self._ep_size}" + ) + last_observed_flags = observed_flags + poll_timed_out = False + except CompletionFlagReadTimeout: + observed_flags = last_observed_flags + poll_timed_out = True + except Exception as exc: # noqa: BLE001 - keep watchdog failures visible. + self._stop_after_error(exc) + tllm_logger.error("AlltoAll watchdog stopped after polling error: %s", exc) + return + + if self._phase_complete(watch, observed_flags): + with self._cv: + if self._queue and self._queue[0] is watch: + self._queue.popleft() + self._cv.notify_all() + last_observed_flags = tuple(UNKNOWN_COMPLETION_FLAG for _ in range(self._ep_size)) + poll_timed_out = False + continue + + if time.monotonic() - watch.start_s >= self._timeout_s: + try: + self._handle_timeout(watch, observed_flags, poll_timed_out=poll_timed_out) + except Exception as exc: # noqa: BLE001 - keep watchdog failures visible. + self._stop_after_error(exc) + tllm_logger.error( + "AlltoAll watchdog stopped after timeout handling error: %s", exc + ) + return + with self._cv: + # The GPU stream is no longer trustworthy once a collective + # times out. Drop queued follow-on phases so they do not + # produce duplicate or misleading reports. + self._queue.clear() + self._cv.notify_all() + last_observed_flags = tuple(UNKNOWN_COMPLETION_FLAG for _ in range(self._ep_size)) + poll_timed_out = False + continue + + with self._cv: + self._cv.wait(timeout=self._poll_interval_s) diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index 7b5d199624fa..9d5e7a6cb013 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -56,6 +56,12 @@ # BufferKind is bound from C++; see cpp/tensorrt_llm/thop/outputTensor.h (torch_ext::BufferKind). from tensorrt_llm.bindings.internal.thop import BufferKind +# Communicator mode is startup-only. Keep this local cached flag so importing +# the custom-op registry does not recursively initialize the distributed +# package merely to filter allreduce tactics. +NCCL_FAULT_TOLERANCE_ENABLED = os.environ.get( + "TLLM_FAULT_TOLERANCE_MODE") == "1" + def _init_deep_gemm_pdl() -> None: try: @@ -74,6 +80,16 @@ def _init_deep_gemm_pdl() -> None: logger.warning(f"Failed to initialize DeepGEMM PDL: {err}") +def _use_cuda_graph_for_allreduce_tuning() -> bool: + """Whether raw-NCCL tactics may be captured by the autotuner. + + The 1a.7 communicator can be replaced after failure, so captured NCCL + nodes are intentionally rejected in FT mode until graph invalidation and + recapture land in 1a.11. Eager profiling preserves valid tactic timings. + """ + return os.environ.get("TLLM_FAULT_TOLERANCE_MODE") != "1" + + _init_deep_gemm_pdl() @@ -2200,6 +2216,8 @@ def get_valid_tactics( AllReduceStrategy.NCCL_SYMMETRIC.value, AllReduceStrategy.NCCL.value, ] + if NCCL_FAULT_TOLERANCE_ENABLED: + return [AllReduceStrategy.NCCL.value] # Fallback in allreduceOp is set to NCCL_SYMMETRIC as default # So we need to check if the workspace size is too large to avoid hanging. workspace_size = inputs[0].numel() * inputs[0].element_size() @@ -2242,10 +2260,17 @@ def forward( return input if tactic == -1: # tactic == -1 means the autotuner cache missed for this shape. - # Fall back to NCCL_SYMMETRIC; asymmetric ncclMemAlloc failures are handled - # by a cross-rank barrier in NCCLWindowAllocator which falls back - # to plain NCCL. - tactic = AllReduceStrategy.NCCL_SYMMETRIC.value + # FT mode must avoid NCCL_SYMMETRIC's MPI-local topology discovery; + # otherwise retain the optimized default-off fallback. + tactic = (AllReduceStrategy.NCCL.value + if NCCL_FAULT_TOLERANCE_ENABLED else + AllReduceStrategy.NCCL_SYMMETRIC.value) + + if (NCCL_FAULT_TOLERANCE_ENABLED + and tactic != AllReduceStrategy.NCCL.value): + raise RuntimeError( + "NCCL error: fault-tolerant allreduce tuning selected a " + "custom tactic without a peer-failure abort path") return torch.ops.trtllm.allreduce( input, @@ -2315,7 +2340,8 @@ def _inputs_pre_hook_register_nccl_symmetric_memory_window( tuning_config = replace( AllReduceRunner.tuning_config, - inputs_pre_hook=_inputs_pre_hook_register_nccl_symmetric_memory_window) + inputs_pre_hook=_inputs_pre_hook_register_nccl_symmetric_memory_window, + use_cuda_graph=_use_cuda_graph_for_allreduce_tuning()) _, best_tactic = tuner.choose_one( "trtllm::tunable_allreduce::allreduce", diff --git a/tensorrt_llm/_torch/distributed/communicator.py b/tensorrt_llm/_torch/distributed/communicator.py index 6fbbf16a19df..4c7dbc63b7fd 100644 --- a/tensorrt_llm/_torch/distributed/communicator.py +++ b/tensorrt_llm/_torch/distributed/communicator.py @@ -1,6 +1,23 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + import math import pickle # nosec B403 +import threading from abc import ABC, abstractmethod +from dataclasses import dataclass from enum import IntEnum from functools import lru_cache, wraps from typing import Any, Callable, List, Optional, Tuple @@ -16,7 +33,11 @@ except Exception: MPI = None # deferred; functions will error if used when ENABLE_MULTI_DEVICE is True -from tensorrt_llm._mnnvl_utils import init_helix_cp_comm +from tensorrt_llm._mnnvl_utils import (get_helix_cp_mnnvl_topology, + init_helix_cp_comm) +from tensorrt_llm._torch.distributed.nccl_fault_tolerance import ( + NCCL_FAULT_TOLERANCE_ENABLED, _canonical_recovery_generation, + _recovery_rendezvous_id) from tensorrt_llm._utils import (mpi_allgather, mpi_barrier, mpi_comm, mpi_disabled, mpi_isend, mpi_isend_object, mpi_recv, mpi_recv_object, mpi_send, @@ -73,6 +94,448 @@ def reduce_op_to_mpi(op: ReduceOp) -> MPI.Op: return _reduce_op_to_mpi_dict[op] +@dataclass(frozen=True) +class MpiFtSubcommSetup: + """Collectively validated WideEP FT communicator startup state.""" + + comm: Any + local_rank: int + ep_size: int + ulfm_available: bool + + +_MPI_FT_PROCESS_LIFETIME_REFS: list[Any] = [] +_MPI_FT_PROCESS_LIFETIME_REFS_LOCK = threading.Lock() + + +def _retain_mpi_ft_comm(comm: Any) -> None: + """Keep a communicator with unsafe cleanup reachable until process exit.""" + with _MPI_FT_PROCESS_LIFETIME_REFS_LOCK: + _MPI_FT_PROCESS_LIFETIME_REFS.append(comm) + + +def _is_mpi_ft_int(value: Any) -> bool: + """Reject bool/float values that compare equal to MPI integer fields.""" + return isinstance(value, int) and not isinstance(value, bool) + + +def _mpi_ft_local_topology(mapping: Mapping, parent_rank: int, parent_size: int, + provided_thread_level: int, + health_size: Optional[int]) -> dict[str, Any]: + """Build a serializable startup record without raising locally.""" + world_size = None + mapping_rank = None + ep_group = None + ep_size = None + ep_rank = None + mapping_error = None + try: + world_size = mapping.world_size + mapping_rank = mapping.rank + ep_group = tuple(mapping.moe_ep_group) + ep_size = mapping.moe_ep_size + ep_rank = mapping.moe_ep_rank + except Exception as error: + mapping_error = ("invalid WideEP FT mapping topology: " + f"{type(error).__name__}: {error}") + + error_kind = None + error_message = None + if provided_thread_level < MPI.THREAD_MULTIPLE: + error_kind = "runtime" + error_message = ( + "WideEP FT requires MPI.THREAD_MULTIPLE because its control-plane " + "thread overlaps other MPI traffic " + f"(provided={provided_thread_level}, required={MPI.THREAD_MULTIPLE})" + ) + elif mapping_error is not None: + error_kind = "value" + error_message = mapping_error + elif not ep_group: + error_kind = "value" + error_message = "mapping.moe_ep_group must not be empty" + elif len(ep_group) != ep_size: + error_kind = "value" + error_message = ( + "mapping.moe_ep_group size must match mapping.moe_ep_size, " + f"got {len(ep_group)} and {ep_size}") + else: + try: + has_duplicate_ranks = len(set(ep_group)) != len(ep_group) + has_invalid_rank = any( + not isinstance(rank, int) or rank < 0 or rank >= parent_size + for rank in ep_group) + except TypeError as error: + error_kind = "value" + error_message = f"mapping.moe_ep_group contains invalid ranks: {error}" + else: + if has_duplicate_ranks: + error_kind = "value" + error_message = ( + "mapping.moe_ep_group contains duplicate ranks: " + f"{ep_group}") + elif parent_size != world_size: + error_kind = "runtime" + error_message = ( + "WideEP FT parent communicator size must match " + f"mapping.world_size, got {parent_size} and {world_size}") + elif parent_rank != mapping_rank: + error_kind = "runtime" + error_message = ( + "WideEP FT parent communicator rank must match mapping.rank, " + f"got {parent_rank} and {mapping_rank}") + elif has_invalid_rank: + error_kind = "value" + error_message = ( + "mapping.moe_ep_group contains a rank outside the parent " + f"communicator: {ep_group}") + elif ep_size != world_size or set(ep_group) != set( + range(parent_size)): + error_kind = "value" + error_message = ( + "WideEP FT MVP requires one MoE EP group spanning the " + "full MPI world; " + f"got world_size={world_size}, moe_ep_group={ep_group}") + elif health_size is not None and health_size != ep_size: + error_kind = "value" + error_message = ( + "EPGroupHealth size must match mapping.moe_ep_size, " + f"got {health_size} and {ep_size}") + elif parent_rank not in ep_group: + error_kind = "value" + error_message = ( + f"local rank {parent_rank} is not in its MoE EP group " + f"{ep_group}") + elif ep_group.index(parent_rank) != ep_rank: + error_kind = "value" + error_message = ( + "mapping.moe_ep_group order must match mapping.moe_ep_rank, " + f"got group index {ep_group.index(parent_rank)} and EP rank " + f"{ep_rank}") + + return { + "parent_rank": parent_rank, + "parent_size": parent_size, + "world_size": world_size, + "mapping_rank": mapping_rank, + "ep_group": ep_group, + "ep_size": ep_size, + "ep_rank": ep_rank, + "health_size": health_size, + "error_kind": error_kind, + "error_message": error_message, + } + + +def _validate_mpi_ft_topologies(topologies: List[Any], + parent_size: int) -> None: + """Raise the same validation error on every parent-communicator rank.""" + if len(topologies) != parent_size: + raise RuntimeError( + "WideEP FT startup allgather returned an unexpected number of " + f"topologies: got {len(topologies)}, expected {parent_size}") + + for gather_rank, topology in enumerate(topologies): + if not isinstance(topology, dict): + raise RuntimeError( + "WideEP FT startup allgather returned an invalid topology for " + f"parent rank {gather_rank}: {topology!r}") + + # Object allgather returns records in communicator-rank order. Selecting the + # first error makes every rank raise the same exception before Split. + for gather_rank, topology in enumerate(topologies): + error_message = topology.get("error_message") + if error_message is None: + continue + reporting_rank = topology.get("parent_rank", gather_rank) + message = ("WideEP FT startup validation failed on parent rank " + f"{reporting_rank}: {error_message}") + if topology.get("error_kind") == "value": + raise ValueError(message) + raise RuntimeError(message) + + for gather_rank, topology in enumerate(topologies): + reported_parent_rank = topology.get("parent_rank") + if (not _is_mpi_ft_int(reported_parent_rank) + or reported_parent_rank != gather_rank): + raise RuntimeError( + "WideEP FT startup topology is inconsistent: allgather slot " + f"{gather_rank} reports parent rank " + f"{reported_parent_rank!r}") + reported_parent_size = topology.get("parent_size") + if (not _is_mpi_ft_int(reported_parent_size) + or reported_parent_size != parent_size): + raise RuntimeError( + "WideEP FT startup topology is inconsistent: parent rank " + f"{gather_rank} reports communicator size " + f"{reported_parent_size!r}, expected {parent_size}") + + ep_group = topology.get("ep_group") + if not isinstance(ep_group, (list, tuple)): + raise RuntimeError( + "WideEP FT startup topology is inconsistent: parent rank " + f"{gather_rank} reports invalid EP group {ep_group!r}") + if len(ep_group) != parent_size: + raise RuntimeError( + "WideEP FT startup topology is inconsistent: parent rank " + f"{gather_rank} reports EP group size {len(ep_group)}, " + f"expected {parent_size}") + if any( + isinstance(peer_rank, bool) or not _is_mpi_ft_int(peer_rank) + or peer_rank < 0 or peer_rank >= parent_size + for peer_rank in ep_group): + raise RuntimeError( + "WideEP FT startup topology is inconsistent: parent rank " + f"{gather_rank} reports invalid EP group ranks {ep_group!r}") + if len(set(ep_group)) != parent_size: + raise RuntimeError( + "WideEP FT startup topology is inconsistent: parent rank " + f"{gather_rank} reports duplicate EP group ranks {ep_group!r}") + world_size = topology.get("world_size") + if not _is_mpi_ft_int(world_size) or world_size != parent_size: + raise RuntimeError( + "WideEP FT startup topology is inconsistent: parent rank " + f"{gather_rank} reports world size " + f"{world_size!r}, expected {parent_size}") + mapping_rank = topology.get("mapping_rank") + if not _is_mpi_ft_int(mapping_rank) or mapping_rank != gather_rank: + raise RuntimeError( + "WideEP FT startup topology is inconsistent: parent rank " + f"{gather_rank} reports mapping rank " + f"{mapping_rank!r}") + ep_size = topology.get("ep_size") + if not _is_mpi_ft_int(ep_size) or ep_size != parent_size: + raise RuntimeError( + "WideEP FT startup topology is inconsistent: parent rank " + f"{gather_rank} reports EP size {ep_size!r}, " + f"expected {parent_size}") + ep_rank = topology.get("ep_rank") + if (not _is_mpi_ft_int(ep_rank) or ep_rank < 0 or ep_rank >= parent_size + or ep_group[ep_rank] != gather_rank): + raise RuntimeError( + "WideEP FT startup topology is inconsistent: parent rank " + f"{gather_rank} reports invalid EP rank {ep_rank!r} for " + f"group {ep_group!r}") + health_size = topology.get("health_size") + if health_size is not None and (not _is_mpi_ft_int(health_size) + or health_size != parent_size): + raise RuntimeError( + "WideEP FT startup topology is inconsistent: parent rank " + f"{gather_rank} reports health size {health_size}, " + f"expected {parent_size}") + # Only dereference peer groups after every gathered record has passed the + # schema/range checks above. This keeps malformed records rank-attributed and + # prevents a valid earlier record from indexing into unchecked peer data. + for gather_rank, topology in enumerate(topologies): + ep_group = topology["ep_group"] + for peer_rank in ep_group: + peer_group = topologies[peer_rank].get("ep_group") + if peer_group != ep_group: + raise RuntimeError( + "WideEP FT startup topology is inconsistent: parent rank " + f"{gather_rank} reports EP group {ep_group}, but member " + f"rank {peer_rank} reports {peer_group}") + + +def _mpi_ft_post_split_status(ft_comm: Any, parent_rank: int, ep_rank: int, + ep_size: int) -> dict[str, Any]: + """Configure a derived communicator without raising before reconciliation.""" + error_message = None + ulfm_available = False + try: + # Install the non-fatal handler before any other operation on the new + # communicator. In particular, do not inspect a possibly broken handle + # while it still inherits the parent's fatal error policy. + ft_comm.Set_errhandler(MPI.ERRORS_RETURN) + actual_rank = ft_comm.Get_rank() + actual_size = ft_comm.Get_size() + if actual_size != ep_size or actual_rank != ep_rank: + error_message = ( + "MPI_Comm_split created an unexpected WideEP FT communicator: " + f"rank={actual_rank}, size={actual_size}, " + f"expected rank={ep_rank}, size={ep_size}") + except Exception as error: + error_message = ("WideEP FT communicator setup raised " + f"{type(error).__name__}: {error}") + + if error_message is None and callable(getattr( + ft_comm, "Is_revoked", None)) and callable( + getattr(ft_comm, "Revoke", None)): + try: + already_revoked = ft_comm.Is_revoked() + except Exception as error: + if not _is_unsupported_mpi_ulfm_error(error): + error_message = ("WideEP FT ULFM probe raised " + f"{type(error).__name__}: {error}") + else: + if already_revoked: + error_message = ( + "WideEP FT communicator is already revoked at construction") + else: + ulfm_available = True + return { + "parent_rank": parent_rank, + "error_message": error_message, + "ulfm_available": ulfm_available, + } + + +def _is_unsupported_mpi_ulfm_error(error: BaseException) -> bool: + if isinstance(error, NotImplementedError): + return True + mpi_exception = getattr(MPI, "Exception", None) + if not isinstance(mpi_exception, type) or not isinstance( + error, mpi_exception): + return False + get_error_class = getattr(error, "Get_error_class", None) + if not callable(get_error_class): + return False + unsupported_classes = { + value + for value in ( + getattr(MPI, "ERR_NOT_SUPPORTED", None), + getattr(MPI, "ERR_UNSUPPORTED_OPERATION", None), + ) if value is not None + } + try: + error_class = get_error_class() + except Exception: + # This helper runs while building a rank-local status that must reach + # the parent allgather. A broken exception classifier is not evidence of + # an unsupported operation, but it must not strand peer ranks either. + return False + return error_class in unsupported_classes + + +def _mpi_ft_post_split_result( + statuses: List[Any], + parent_size: int, +) -> tuple[Optional[str], bool]: + """Reconcile setup errors and the global ULFM capability.""" + if len(statuses) != parent_size: + return ( + "WideEP FT post-split allgather returned an unexpected number of " + f"statuses: got {len(statuses)}, expected {parent_size}", + False, + ) + for gather_rank, status in enumerate(statuses): + if not isinstance(status, dict): + return ( + "WideEP FT post-split allgather returned an invalid status for " + f"parent rank {gather_rank}: {status!r}", + False, + ) + if status.get("parent_rank") != gather_rank: + return ( + "WideEP FT post-split status is inconsistent: allgather slot " + f"{gather_rank} reports parent rank " + f"{status.get('parent_rank')}", + False, + ) + if not isinstance(status.get("ulfm_available"), bool): + return ( + "WideEP FT post-split allgather returned incomplete capability " + f"state for parent rank {gather_rank}: {status!r}", + False, + ) + + ulfm_available = all(status["ulfm_available"] for status in statuses) + for gather_rank, status in enumerate(statuses): + error_message = status.get("error_message") + if error_message is not None: + return ( + "WideEP FT communicator setup failed on parent rank " + f"{gather_rank}: {error_message}", + False, + ) + return None, ulfm_available + + +def create_mpi_ft_subcomm( + mapping: Mapping, + parent_comm: Optional[Any] = None, + health_size: Optional[int] = None) -> MpiFtSubcommSetup: + """Create the dedicated MPI control-plane communicator for WideEP FT. + + This operation is collective across ``parent_comm`` and must run on every + rank during startup, before any background failure-broadcast threads are + launched. The MVP requires one MoE EP group spanning the parent world, + ordered by the EP-local rank used by the caller's detected-rank state. + + Before splitting, ranks exchange their local validation outcome and EP + topology on the healthy parent communicator. This prevents one rank from + raising locally while its peers block forever inside ``MPI_Comm_split``. + + Args: + mapping: Distributed topology for the local rank. + parent_comm: Parent MPI communicator. Defaults to TRT-LLM's active + ``mpi_comm()``, which wraps ``MPI.COMM_WORLD`` in the standard + launch path and preserves custom communicator sessions. + health_size: Optional local detected-rank-state size to validate + collectively before splitting. + + Returns: + The world-spanning EP MPI communicator and its collectively validated + rank, size, and ULFM capability. + + Raises: + RuntimeError: If MPI is unavailable, does not provide + ``MPI.THREAD_MULTIPLE``, or creates an unexpected communicator. + ValueError: If the mapping's EP group is inconsistent. + """ + if MPI is None: + raise RuntimeError( + "mpi4py is required to create the WideEP FT communicator") + + parent_comm = mpi_comm() if parent_comm is None else parent_comm + parent_rank = parent_comm.Get_rank() + parent_size = parent_comm.Get_size() + local_topology = _mpi_ft_local_topology(mapping, parent_rank, parent_size, + MPI.Query_thread(), health_size) + topologies = parent_comm.allgather(local_topology) + _validate_mpi_ft_topologies(topologies, parent_size) + + ep_group = local_topology["ep_group"] + ep_rank = local_topology["ep_rank"] + ep_size = local_topology["ep_size"] + + # The MVP topology gate above makes the world-spanning group's first rank + # a stable color shared by every process. + ft_comm = parent_comm.Split(color=ep_group[0], key=ep_rank) + try: + setup_status = _mpi_ft_post_split_status(ft_comm, parent_rank, ep_rank, + ep_size) + setup_statuses = parent_comm.allgather(setup_status) + setup_error, ulfm_available = _mpi_ft_post_split_result( + setup_statuses, parent_size) + if setup_error is None: + setup = MpiFtSubcommSetup( + comm=ft_comm, + local_rank=ep_rank, + ep_size=ep_size, + ulfm_available=ulfm_available, + ) + # Enforce process-lifetime ownership at the creation boundary. A + # direct or future caller must not be able to drop the last Python + # reference and trigger rank-local collective MPI_Comm_free. + _retain_mpi_ft_comm(ft_comm) + return setup + except Exception: + # Once Split succeeds, never let an unexpected post-Split exception + # drop the last reference and trigger rank-local communicator cleanup. + _retain_mpi_ft_comm(ft_comm) + raise + + # MPI_Comm_free is collective. Once setup has reported any communicator + # invariant failure, attempting collective cleanup on that same handle is + # not provably bounded, even with MPI_ERRORS_RETURN installed. Keep the + # handle reachable and let MPI reclaim it at process teardown. + _retain_mpi_ft_comm(ft_comm) + logger.warning("WideEP FT retained a communicator after startup failure: " + f"{setup_error}") + raise RuntimeError(setup_error) + + class Distributed(ABC): def __init__(self, mapping: Mapping): @@ -1177,16 +1640,67 @@ class PPCommNCCL: def __init__(self, global_mapping: Mapping): self.mapping = global_mapping + self._topology = self._mapping_topology(global_mapping) + self._reconfigure_lock = threading.Lock() + self._reconfigure_generation = 0 + self._completed_recovery = None self.nccl_comm = torch.classes.trtllm.NcclCommunicatorOp( self.mapping.world_size, self.mapping.rank, ) + if NCCL_FAULT_TOLERANCE_ENABLED: + # The FT native communicator survives process-local Python wrapper + # churn. Mirror its actual survivor membership instead of assuming + # a newly-created wrapper always represents the initial world. + self._active_ranks = tuple( + int(rank) for rank in self.nccl_comm.get_active_ranks()) + else: + # Preserve the legacy constructor path when FT is disabled. + self._active_ranks = tuple(range(self.mapping.world_size)) self.tensor_ready_event = torch.cuda.Event() self.send_stream = torch.cuda.Stream() + @staticmethod + def _mapping_topology(mapping: Mapping) -> tuple: + return ( + int(mapping.world_size), + int(mapping.rank), + tuple(mapping.pp_group), + get_helix_cp_mnnvl_topology(mapping), + ) + + def is_compatible(self, mapping: Mapping) -> bool: + """Whether another engine can share this process-local communicator.""" + return self._topology == self._mapping_topology(mapping) + + def _validate_peer(self, peer: int) -> int: + peer = int(peer) + if peer not in self._topology[2]: + raise RuntimeError( + f"NCCL error: peer world rank {peer} is not in this rank's PP group " + f"{list(self._topology[2])}") + if peer not in self._active_ranks: + raise RuntimeError( + f"NCCL error: peer world rank {peer} is not active in the current PP communicator" + ) + return peer + + def _required_pp_peer(self, *, next_peer: bool) -> int: + # Do not silently skip a failed stage: connecting non-adjacent stages + # would bypass model layers. A higher-level topology reconstruction may + # pass an explicit remapped peer once it has reassigned those layers. + peer = (self.mapping.next_pp_rank() + if next_peer else self.mapping.prev_pp_rank()) + return self._validate_peer(peer) + def send(self, tensor: torch.Tensor, dest: Optional[int] = None): - if dest is None: - dest = self.mapping.next_pp_rank() + if not NCCL_FAULT_TOLERANCE_ENABLED: + if dest is None: + dest = self.mapping.next_pp_rank() + elif dest is None: + dest = self._required_pp_peer(next_peer=True) + else: + dest = self._validate_peer(dest) # NCCL send kernel in send_stream cannot be captured, # so we send in the current stream instead in CUDA graph cases. @@ -1204,10 +1718,110 @@ def send(self, tensor: torch.Tensor, dest: Optional[int] = None): self.nccl_comm.send(tensor, dest) def recv(self, tensor: torch.Tensor, src: Optional[int] = None): - if src is None: - src = self.mapping.prev_pp_rank() + if not NCCL_FAULT_TOLERANCE_ENABLED: + if src is None: + src = self.mapping.prev_pp_rank() + elif src is None: + src = self._required_pp_peer(next_peer=False) + else: + src = self._validate_peer(src) self.nccl_comm.recv(tensor, src) + def abort(self, generation: Optional[int] = None) -> None: + """Abort the PP communicator on this survivor.""" + # Capture the generation before waiting for a concurrent rebuild. If + # that rebuild succeeds, this abort belongs to the old communicator + # and must not tear down its replacement. An abort that starts after + # the rebuild observes the new generation and still takes effect. + if generation is None: + generation = self._reconfigure_generation + with self._reconfigure_lock: + if generation != self._reconfigure_generation: + return + self.nccl_comm.abort() + + def abort_and_reinit(self, + active_ranks: List[int], + generation: Optional[int] = None) -> None: + """Rebuild the PP communicator using only surviving world ranks. + + Every survivor must call this with the same rank set. The C++ wrapper + exchanges a fresh NCCL unique ID point-to-point, so excluded ranks do + not participate in communicator initialization. + + ``generation`` is the coordinator's shared monotonic recovery-event + generation, advanced for every distinct attempt. It is required for + same-membership recovery so every rank enters the native rendezvous + even when local watchdog observations differ. Repeated callbacks with + the same generation and target are idempotent. A retry after any native + failure must use a newly advanced generation so stale rendezvous traffic + cannot pair different recovery attempts. + """ + canonical_ranks = tuple(sorted(int(rank) for rank in active_ranks)) + if not canonical_ranks: + raise ValueError("active_ranks must not be empty") + if len(canonical_ranks) != len(set(canonical_ranks)): + raise ValueError("active_ranks must not contain duplicates") + if self.mapping.rank not in canonical_ranks: + raise ValueError( + f"current world rank {self.mapping.rank} is not active") + generation = _canonical_recovery_generation(generation) + rendezvous_id = _recovery_rendezvous_id(generation) + + # Failure notifications can be duplicated or delivered concurrently by + # independent control paths. Keep the native rebuild and its Python + # membership mirror in one transaction so a stale callback cannot + # widen membership after a newer shrink has completed. + with self._reconfigure_lock: + current_ranks = tuple( + int(rank) for rank in self.nccl_comm.get_active_ranks()) + self._active_ranks = current_ranks + if generation is not None and self._completed_recovery is not None: + completed_generation, completed_target = self._completed_recovery + if generation < completed_generation: + return + if generation == completed_generation: + if canonical_ranks != completed_target: + raise RuntimeError( + "NCCL error: conflicting PP communicator recovery " + f"target for generation {generation}: completed " + f"{list(completed_target)}, requested {list(canonical_ranks)}" + ) + return + if canonical_ranks == current_ranks and generation is None: + raise ValueError( + "generation is required for same-membership PP communicator " + "recovery so every survivor makes the same rebuild decision" + ) + if not set(canonical_ranks).issubset(current_ranks): + raise ValueError( + "abort_and_reinit cannot reactivate a removed rank") + + self.nccl_comm.abort_and_reinit(list(canonical_ranks), + rendezvous_id) + native_ranks = tuple( + int(rank) for rank in self.nccl_comm.get_active_ranks()) + self._active_ranks = native_ranks + self._reconfigure_generation += 1 + if native_ranks != canonical_ranks: + raise RuntimeError( + "NCCL error: rebuilt PP communicator returned unexpected " + f"membership {list(native_ranks)}; expected {list(canonical_ranks)}" + ) + if generation is not None: + self._completed_recovery = (generation, canonical_ranks) + + def get_async_error(self) -> str: + """Return the C++ wrapper's latched NCCL abort/error reason, if any.""" + return self.nccl_comm.get_async_error() + + def get_active_ranks(self) -> List[int]: + """Return original world-rank IDs in the current PP communicator.""" + with self._reconfigure_lock: + self._active_ranks = tuple( + int(rank) for rank in self.nccl_comm.get_active_ranks()) + return list(self._active_ranks) + class PPCommTorch: @@ -1239,16 +1853,149 @@ def recv(self, tensor: torch.Tensor, src: Optional[int] = None): _pp_comm = None +_pp_comm_refcount = 0 +_pp_comm_lock = threading.Lock() +_pp_comm_condition = threading.Condition(_pp_comm_lock) +_pp_comm_control_refcount = 0 +_pp_comm_final_release_pending = False def init_pp_comm(mapping): - """Initialize PPComm once at startup""" - global _pp_comm - if mpi_disabled(): - _pp_comm = PPCommTorch(mapping) - else: - _pp_comm = PPCommNCCL(mapping) - init_helix_cp_comm(mapping) + """Acquire the process-local PP communicator for one model engine.""" + global _pp_comm, _pp_comm_refcount + with _pp_comm_condition: + while (_pp_comm_control_refcount > 0 or _pp_comm_final_release_pending): + _pp_comm_condition.wait() + created = False + if mpi_disabled(): + if _pp_comm is None: + _pp_comm = PPCommTorch(mapping) + created = True + elif not isinstance(_pp_comm, PPCommTorch): + raise RuntimeError( + "PP communicator is already initialized with a different backend" + ) + elif _pp_comm.mapping != mapping: + raise RuntimeError( + "PP communicator is already initialized for a different topology" + ) + else: + if _pp_comm is None: + _pp_comm = PPCommNCCL(mapping) + created = True + elif not isinstance(_pp_comm, PPCommNCCL): + raise RuntimeError( + "NCCL error: PP communicator is already initialized with a different backend" + ) + elif not _pp_comm.is_compatible(mapping): + raise RuntimeError( + "NCCL error: PP communicator is already initialized for a " + "different topology (PP or Helix MNNVL CP topology)") + try: + init_helix_cp_comm(mapping) + except Exception: + if created and _pp_comm_refcount == 0: + _pp_comm = None + raise + _pp_comm_refcount += 1 + + +def release_pp_comm() -> None: + """Release one model engine's ownership of the shared PP communicator.""" + global _pp_comm, _pp_comm_refcount, _pp_comm_final_release_pending + with _pp_comm_condition: + if _pp_comm_refcount <= 0: + raise RuntimeError( + "PP communicator release has no matching acquisition") + _pp_comm_refcount -= 1 + if _pp_comm_refcount == 0: + _pp_comm_final_release_pending = True + while _pp_comm_control_refcount > 0: + _pp_comm_condition.wait() + if not (NCCL_FAULT_TOLERANCE_ENABLED + and isinstance(_pp_comm, PPCommNCCL)): + # Preserve the legacy teardown behavior when FT is disabled. + _pp_comm = None + # In FT mode both the native communicator and its Python recovery + # generation state are process-lifetime. A rank-local final release + # is not a distributed teardown barrier, and dropping either state + # could make only one rank enter a duplicate-generation rendezvous. + _pp_comm_final_release_pending = False + _pp_comm_condition.notify_all() + + +def _get_pp_nccl_comm_locked() -> PPCommNCCL: + """Return the NCCL PP communicator while ``_pp_comm_lock`` is held.""" + if not isinstance(_pp_comm, PPCommNCCL): + raise RuntimeError("NCCL error: PP communicator is not initialized") + return _pp_comm + + +def _acquire_pp_nccl_control( + *, + capture_generation: bool = False) -> tuple[PPCommNCCL, Optional[int]]: + """Keep the shared communicator alive without holding its lifecycle lock.""" + global _pp_comm_control_refcount + with _pp_comm_lock: + comm = _get_pp_nccl_comm_locked() + if _pp_comm_refcount <= 0 or _pp_comm_final_release_pending: + raise RuntimeError("NCCL error: PP communicator is being released") + _pp_comm_control_refcount += 1 + generation = (comm._reconfigure_generation + if capture_generation else None) + return comm, generation + + +def _release_pp_nccl_control() -> None: + global _pp_comm_control_refcount + with _pp_comm_condition: + if _pp_comm_control_refcount <= 0: + raise RuntimeError( + "NCCL error: PP communicator control reference underflow") + _pp_comm_control_refcount -= 1 + if _pp_comm_control_refcount == 0: + _pp_comm_condition.notify_all() + + +def pp_comm_abort() -> None: + """Abort the process-local NCCL PP communicator.""" + comm, generation = _acquire_pp_nccl_control(capture_generation=True) + try: + comm.abort(generation) + finally: + comm = None + _release_pp_nccl_control() + + +def pp_comm_abort_and_reinit(active_ranks: List[int], + generation: Optional[int] = None) -> None: + """Rebuild the NCCL PP communicator using surviving original world ranks.""" + comm, _ = _acquire_pp_nccl_control() + try: + comm.abort_and_reinit(active_ranks, generation) + finally: + comm = None + _release_pp_nccl_control() + + +def pp_comm_get_async_error() -> str: + """Return the NCCL PP communicator's latched abort/error reason.""" + comm, _ = _acquire_pp_nccl_control() + try: + return comm.get_async_error() + finally: + comm = None + _release_pp_nccl_control() + + +def pp_comm_get_active_ranks() -> List[int]: + """Return original world-rank IDs in the NCCL PP communicator.""" + comm, _ = _acquire_pp_nccl_control() + try: + return comm.get_active_ranks() + finally: + comm = None + _release_pp_nccl_control() @TorchDist.log_op diff --git a/tensorrt_llm/_torch/distributed/moe_alltoall.py b/tensorrt_llm/_torch/distributed/moe_alltoall.py index ca3a50dcfd12..2bf31ef680a4 100644 --- a/tensorrt_llm/_torch/distributed/moe_alltoall.py +++ b/tensorrt_llm/_torch/distributed/moe_alltoall.py @@ -8,12 +8,18 @@ # ruff: noqa: E501 import os +import sys from dataclasses import dataclass -from typing import Dict, Optional +from typing import Callable, Dict, Optional import torch from tensorrt_llm._mnnvl_utils import MnnvlMemory +from tensorrt_llm._torch.alltoall_watchdog import ( + DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, + DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S, ActiveRankMaskSnapshot, + AlltoAllWatchdog, AlltoAllWatchdogCoordinator, AlltoAllWatchdogTimeout, + EPGroupHealthLike) from tensorrt_llm.bindings import internal as _tllm_internal from tensorrt_llm.logger import logger as tllm_logger from tensorrt_llm.mapping import Mapping @@ -26,6 +32,7 @@ class _A2AState: local_num_tokens: int | None = None combine_payload_offset: int | None = None eplb_gathered_stats: torch.Tensor | None = None + active_rank_mask_snapshot: ActiveRankMaskSnapshot | None = None class MoeAlltoAll: @@ -126,7 +133,13 @@ def __init__( num_slots: int, workspace_size_per_rank: int, num_experts: Optional[int] = None, - ): + ep_group_health: Optional[EPGroupHealthLike] = None, + alltoall_watchdog_timeout_s: Optional[float] = None, + alltoall_watchdog_poll_interval_s: + float = DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, + alltoall_watchdog_on_timeout: Optional[Callable[ + [AlltoAllWatchdogTimeout], None]] = None, + ) -> None: """ Initialize MoeAlltoAll with workspace allocation. @@ -138,6 +151,12 @@ def __init__( Note: The terminology is mapped to `num_experts` in this class and the kernels. num_experts: (Optional) Number of experts for EPLB stats (must be <= num_slots). DO NOT provide this parameter if EPLB is not enabled. Note: The terminology is mapped to `eplb_stats_num_experts` in this class and the kernels. + ep_group_health: Optional read-only committed EP membership. When present, its mask is passed to the + CUDA kernels and defines the peers expected by the watchdog. Timeout detection never mutates it. + alltoall_watchdog_timeout_s: Optional timeout for the host-side AlltoAll watchdog. If None, the + watchdog is disabled. + alltoall_watchdog_poll_interval_s: Poll interval for the watchdog thread. + alltoall_watchdog_on_timeout: Optional callback invoked when the watchdog reports suspects. """ # Check for environment variable override workspace_mb_env = os.environ.get("TRTLLM_MOE_A2A_WORKSPACE_MB") @@ -214,6 +233,45 @@ def __init__( self.metainfo = self._WORKSPACE["metainfo"] # Internal state self._state: _A2AState = _A2AState() + self.ep_group_health = ep_group_health + workspace_state = self._WORKSPACE + assert workspace_state is not None + metainfo_index = self._METAINFO_INDEX + assert metainfo_index is not None + self._watchdog_coordinator = AlltoAllWatchdogCoordinator( + workspace_state=workspace_state, + workspace=self.workspace, + metainfo=self.metainfo, + metainfo_index=metainfo_index, + ep_rank=self.ep_rank, + health=self.ep_group_health, + ) + self._destroyed = False + self._alltoall_watchdog: AlltoAllWatchdog | None = None + if (alltoall_watchdog_timeout_s is None + and self.ep_group_health is not None): + alltoall_watchdog_timeout_s = DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S + if alltoall_watchdog_timeout_s is not None: + self._alltoall_watchdog = self._watchdog_coordinator.acquire_watchdog( + ep_size=self.ep_size, + timeout_s=alltoall_watchdog_timeout_s, + poll_interval_s=alltoall_watchdog_poll_interval_s, + on_timeout=alltoall_watchdog_on_timeout, + ) + + def destroy(self) -> None: + """Stop background watchdog resources owned by this wrapper.""" + if getattr(self, "_destroyed", False): + return + self._destroyed = True + watchdog = getattr(self, "_alltoall_watchdog", None) + if watchdog is not None: + self._watchdog_coordinator.release_watchdog(watchdog) + self._alltoall_watchdog = None + + def __del__(self) -> None: + if not sys.is_finalizing(): + self.destroy() def dispatch(self, token_selected_experts: torch.Tensor, @@ -221,7 +279,8 @@ def dispatch(self, runtime_max_tokens_per_rank: int, invalid_token_expert_id: Optional[int] = None, expert_id_payload_index: Optional[int] = None, - eplb_local_stats: Optional[torch.Tensor] = None): + eplb_local_stats: Optional[torch.Tensor] = None, + active_rank_mask: Optional[torch.Tensor] = None): """ Perform MoE all-to-all dispatch operation. @@ -232,6 +291,9 @@ def dispatch(self, invalid_token_expert_id: If not None, set the token_selected_experts of the invalid tokens to this expert id. This is used to notify the MoE to skip these tokens for GroupGEMM. expert_id_payload_index: The index of token_selected_experts in the input_payloads. Must be provided if invalid_token_expert_id is not None. eplb_local_stats: (Optional) [num_experts] tensor containing local statistics for EPLB + active_rank_mask: Optional uint64 CPU tensor overriding committed membership for this dispatch. When + omitted, the committed mask and generation are captured together. Combine reuses that mask and + fails closed if the committed generation changes first. Returns: recv_tensors: List of tensors received, each has shape [ep_size, max_tokens_per_rank, payload_num_elements_per_token] @@ -246,6 +308,9 @@ def dispatch(self, 0 ) == self.eplb_stats_num_experts, "eplb_local_stats size must match eplb_stats_num_experts" + active_rank_mask_snapshot = self._watchdog_coordinator.capture_active_rank_mask( + active_rank_mask) + active_rank_mask = active_rank_mask_snapshot.active_rank_mask recv_tensors, combine_payload_offset, eplb_gathered_stats = torch.ops.trtllm.moe_a2a_dispatch( token_selected_experts, input_payloads, @@ -257,7 +322,11 @@ def dispatch(self, self.top_k, self.num_experts, eplb_local_stats, + active_rank_mask, ) + self._watchdog_coordinator.watch_collective(self._alltoall_watchdog, + "dispatch", + active_rank_mask) if eplb_gathered_stats.numel() == 0: eplb_gathered_stats = None @@ -265,6 +334,7 @@ def dispatch(self, self._state.local_num_tokens = token_selected_experts.size(0) self._state.combine_payload_offset = combine_payload_offset self._state.eplb_gathered_stats = eplb_gathered_stats + self._state.active_rank_mask_snapshot = active_rank_mask_snapshot self._state.phase = "dispatched" if invalid_token_expert_id is not None: @@ -287,6 +357,7 @@ def combine( runtime_max_tokens_per_rank: int, payload_in_workspace: bool = False, use_low_precision_combine: bool = False, + active_rank_mask: Optional[torch.Tensor] = None, ): """ Perform MoE all-to-all combine operation. @@ -296,6 +367,8 @@ def combine( runtime_max_tokens_per_rank: Maximum of the number of tokens of each DP rank's local batch. payload_in_workspace: If True, 'payload' is a view into 'workspace' at 'combine_payload_offset' and no staging copy is needed. If False, the op stages 'payload' into the workspace region before combining. use_low_precision_combine: If True, quantize the combine payload to FP8 for NVLink transfer (halves NVLink bandwidth usage, output precision is preserved). + active_rank_mask: Optional uint64 CPU tensor. If supplied, it must match the mask captured by dispatch + for this collective. A committed-generation change since dispatch aborts the collective epoch. Returns: combined_output: [local_num_tokens, num_elements_per_token] tensor of combined results @@ -303,11 +376,17 @@ def combine( assert self._state.phase == "dispatched", "combine called before a successful dispatch" assert runtime_max_tokens_per_rank <= self.max_num_tokens, "runtime_max_tokens_per_rank must not exceed max_num_tokens" + active_rank_mask_snapshot = self._state.active_rank_mask_snapshot + assert active_rank_mask_snapshot is not None + active_rank_mask = self._watchdog_coordinator.active_rank_mask_for_combine( + active_rank_mask_snapshot, active_rank_mask) output = torch.ops.trtllm.moe_a2a_combine( payload, self._state.local_num_tokens, self.workspace, self.metainfo, runtime_max_tokens_per_rank, self.ep_rank, self.ep_size, self.top_k, self._state.combine_payload_offset, - payload_in_workspace, use_low_precision_combine) + payload_in_workspace, use_low_precision_combine, active_rank_mask) + self._watchdog_coordinator.watch_collective(self._alltoall_watchdog, + "combine", active_rank_mask) # Reset state for next round self.reset_state() diff --git a/tensorrt_llm/_torch/distributed/nccl_fault_tolerance.py b/tensorrt_llm/_torch/distributed/nccl_fault_tolerance.py new file mode 100644 index 000000000000..e2e2dc7e46d3 --- /dev/null +++ b/tensorrt_llm/_torch/distributed/nccl_fault_tolerance.py @@ -0,0 +1,229 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Process-local survivor membership for raw NCCL communicators. + +The C++ raw-NCCL registry keys communicators by their global-rank set. After +``nccl_comm_abort_and_reinit`` replaces a communicator, explicitly dynamic +callers must use the replacement set, while statically sharded TP/CP callers +must stop until their missing state is redistributed. This module keeps that +membership and fail-fast decision in one place instead of attaching it to the +MoE communication object that initiated recovery. + +ProcessGroup-backed collectives intentionally do not use this registry: their +groups need to be reconstructed by the ProcessGroup owner. +""" + +from __future__ import annotations + +import os +import threading +from typing import Callable, Iterable, List, Tuple + +_Group = Tuple[int, ...] +_MAX_RECOVERY_GENERATION = (1 << 63) - 3 + +# Fault tolerance is a startup-only mode. Cache the flag once so the default +# collective hot path pays only a cheap boolean branch and never consults the +# survivor registry. +NCCL_FAULT_TOLERANCE_ENABLED = os.environ.get("TLLM_FAULT_TOLERANCE_MODE") == "1" + +_registry_lock = threading.RLock() +_active_groups: dict[_Group, _Group] = {} +_completed_reconfigurations: dict[_Group, tuple[int, _Group]] = {} + + +def _canonical_group(group: Iterable[int], name: str) -> _Group: + ranks = tuple(int(rank) for rank in group) + if not ranks: + raise ValueError(f"{name} must not be empty") + if len(ranks) != len(set(ranks)): + raise ValueError(f"{name} must not contain duplicate ranks") + return ranks + + +def _canonical_recovery_generation(generation: int | None) -> int | None: + """Validate a coordinator generation accepted by the Torch ``int`` ABI.""" + if generation is None: + return None + canonical_generation = int(generation) + if ( + canonical_generation < 0 + or canonical_generation != generation + or canonical_generation > _MAX_RECOVERY_GENERATION + ): + raise ValueError( + f"generation must be a nonnegative integer no greater than {_MAX_RECOVERY_GENERATION}" + ) + return canonical_generation + + +def _recovery_rendezvous_id(generation: int | None) -> int: + """Map a recovery generation to its wire ID. + + ID 0 is reserved for initial bootstrap and ID 1 for the legacy one-shot + membership shrink without a coordinator generation. Explicit generations + start at ID 2, keeping every namespace distinct on the retained MPI control + communicator. + """ + return 1 if generation is None else generation + 2 + + +def resolve_nccl_group(group: Iterable[int]) -> List[int]: + """Return the latest survivor set for a raw NCCL rank group.""" + if not NCCL_FAULT_TOLERANCE_ENABLED: + return list(group) + ranks = _canonical_group(group, "group") + # Keep the hot-path lookup free of synchronization primitives so Dynamo + # can guard the dictionary value and recompile when recovery publishes a + # new group. Recovery quiesces collective submission before publishing. + return list(_active_groups.get(ranks, ranks)) + + +def resolve_nccl_group_and_rank(group: Iterable[int], world_rank: int) -> tuple[List[int], int]: + """Resolve a raw NCCL group and the caller's compact communicator rank.""" + active_group = resolve_nccl_group(group) + try: + compact_rank = active_group.index(int(world_rank)) + except ValueError as error: + raise RuntimeError( + f"NCCL error: world rank {world_rank} is not in active communicator {active_group}" + ) from error + return active_group, compact_rank + + +def is_nccl_group_reconfigured(group: Iterable[int]) -> bool: + """Cheap hot-path membership-change check for a trusted rank group.""" + if not NCCL_FAULT_TOLERANCE_ENABLED: + return False + original_group = group if isinstance(group, tuple) else tuple(group) + active_group = _active_groups.get(original_group) + return active_group is not None and active_group != original_group + + +def assert_nccl_group_not_reconfigured(group: Iterable[int], operation: str) -> None: + """Reject static rank sharding after transport-only recovery. + + This hot-path check trusts Mapping-owned groups to already be canonical. A + single dictionary lookup is enough before recovery; full validation stays + in the control-path APIs that publish membership. + """ + if not NCCL_FAULT_TOLERANCE_ENABLED: + return + original_group = group if isinstance(group, tuple) else tuple(group) + active_group = _active_groups.get(original_group) + if active_group is not None and active_group != original_group: + raise RuntimeError( + "NCCL error: " + f"{operation} cannot use survivor-only communicator {list(active_group)} " + f"for statically sharded group {list(original_group)}; redistribute the " + "missing rank's state and rebuild the mapping before resuming" + ) + + +def publish_nccl_group_reconfiguration( + original_group: Iterable[int], + current_group: Iterable[int], + active_group: Iterable[int], +) -> None: + """Publish a successful native raw-NCCL communicator replacement. + + ``current_group`` makes stale or racing reconfiguration attempts fail + before they can overwrite newer membership. Publishing happens only + after the coordinated native rebuild succeeds. + """ + original = _canonical_group(original_group, "original_group") + current = _canonical_group(current_group, "current_group") + active = _canonical_group(active_group, "active_group") + + if not set(active).issubset(current): + raise ValueError("active_group must be a subset of current_group") + + with _registry_lock: + registered = _active_groups.get(original, original) + if registered == active: + return + if registered != current: + raise RuntimeError( + "NCCL error: stale communicator membership update: " + f"expected {list(registered)}, got {list(current)}" + ) + + # Preserve aliases for the model's original mapping and for any + # survivor-only mapping views already handed to communication helpers. + aliases = [key for key, value in _active_groups.items() if value == current] + aliases.extend((original, current, active)) + for alias in aliases: + _active_groups[alias] = active + + +def reconfigure_nccl_group( + original_group: Iterable[int], + active_group: Iterable[int], + rebuild: Callable[[List[int], List[int], int], None], + generation: int | None = None, +) -> List[int]: + """Serialize a native communicator rebuild and membership publication. + + Recovery is a process-wide safe-point operation. Holding the registry + lock across ``rebuild`` prevents two local layers or worker threads from + rebuilding the same native communicator to different survivor sets before + either update is published. ``generation`` is the coordinator's monotonic + recovery-event generation, advanced for every distinct attempt (including + transport-only retries). It is required to distinguish a same-membership + transport rebuild from a duplicate callback without consulting rank-local + watchdog state, which can differ transiently across survivors. A first + membership shrink may omit it for compatibility, but every retry after a + native failure must provide a newly advanced generation so the wire + rendezvous cannot pair different attempts. + """ + if not NCCL_FAULT_TOLERANCE_ENABLED: + raise RuntimeError( + "NCCL error: communicator reinitialization requires TLLM_FAULT_TOLERANCE_MODE=1" + ) + + original = _canonical_group(original_group, "original_group") + requested = _canonical_group(active_group, "active_group") + generation = _canonical_recovery_generation(generation) + rendezvous_id = _recovery_rendezvous_id(generation) + + with _registry_lock: + current = _active_groups.get(original, original) + completed = _completed_reconfigurations.get(original) + if generation is not None and completed is not None: + completed_generation, completed_target = completed + if generation < completed_generation: + return list(current) + if generation == completed_generation: + if requested != completed_target: + raise RuntimeError( + "NCCL error: conflicting communicator recovery target " + f"for generation {generation}: completed " + f"{list(completed_target)}, requested {list(requested)}" + ) + return list(current) + + if not set(requested).issubset(current): + raise ValueError("active_group must be a subset of the current_group") + if requested == current: + if generation is None: + raise ValueError( + "generation is required for same-membership communicator " + "recovery so every survivor makes the same rebuild decision" + ) + + rebuild(list(current), list(requested), rendezvous_id) + + aliases = [key for key, value in _active_groups.items() if value == current] + aliases.extend((original, current, requested)) + for alias in aliases: + _active_groups[alias] = requested + if generation is not None: + _completed_reconfigurations[original] = (generation, requested) + return list(requested) + + +def _reset_nccl_group_registry_for_tests() -> None: + """Clear process-local membership. Tests only; native state is untouched.""" + with _registry_lock: + _active_groups.clear() + _completed_reconfigurations.clear() diff --git a/tensorrt_llm/_torch/distributed/ops.py b/tensorrt_llm/_torch/distributed/ops.py index 63bf9d76c2c0..5e46309b701e 100644 --- a/tensorrt_llm/_torch/distributed/ops.py +++ b/tensorrt_llm/_torch/distributed/ops.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + import math import os import platform @@ -8,6 +11,9 @@ from torch import nn from tensorrt_llm._mnnvl_utils import HelixCpMnnvlMemory, MnnvlMemory +from tensorrt_llm._torch.distributed.nccl_fault_tolerance import ( + NCCL_FAULT_TOLERANCE_ENABLED, assert_nccl_group_not_reconfigured, + is_nccl_group_reconfigured, resolve_nccl_group) from tensorrt_llm._torch.distributed.symm_mem_allreduce import \ SymmetricMemoryAllReduce from tensorrt_llm._torch.utils import get_model_extra_attrs @@ -160,6 +166,10 @@ def get_or_scale_allreduce_mnnvl_workspace( def userbuffers_allreduce_finalize( input: torch.Tensor, force_applying_finalize: bool = False) -> torch.Tensor: + if NCCL_FAULT_TOLERANCE_ENABLED and not mpi_disabled(): + raise RuntimeError( + "NCCL error: userbuffers allreduce has no peer-failure abort path " + "and is unavailable while TLLM_FAULT_TOLERANCE_MODE=1") output = torch.ops.trtllm.userbuffers_allreduce_finalize( input, force_applying_finalize) return output @@ -307,9 +317,17 @@ def allgather( Returns: The gathered tensor or tensor list. ''' - group_boxed = mapping.tp_group_pg.boxed() if mpi_disabled() else None - return _allgather(input, mapping.tp_group, mapping.tp_rank, group_boxed, - dim, sizes) + if mpi_disabled(): + group = mapping.tp_group + rank = mapping.tp_rank + group_boxed = mapping.tp_group_pg.boxed() + else: + group = mapping.tp_group + if NCCL_FAULT_TOLERANCE_ENABLED: + assert_nccl_group_not_reconfigured(group, "allgather") + rank = mapping.tp_rank + group_boxed = None + return _allgather(input, group, rank, group_boxed, dim, sizes) def cp_allgather( @@ -331,9 +349,17 @@ def cp_allgather( Returns: The gathered tensor or tensor list. ''' - group_boxed = mapping.cp_group_pg.boxed() if mpi_disabled() else None - return _allgather(input, mapping.cp_group, mapping.cp_rank, group_boxed, - dim, sizes) + if mpi_disabled(): + group = mapping.cp_group + rank = mapping.cp_rank + group_boxed = mapping.cp_group_pg.boxed() + else: + group = mapping.cp_group + if NCCL_FAULT_TOLERANCE_ENABLED: + assert_nccl_group_not_reconfigured(group, "CP allgather") + rank = mapping.cp_rank + group_boxed = None + return _allgather(input, group, rank, group_boxed, dim, sizes) def alltoall_helix( @@ -357,6 +383,16 @@ def alltoall_helix( For each group of input tensors (of size group size), there is one output tensor with shape (group size, *input shape). ''' + if NCCL_FAULT_TOLERANCE_ENABLED and not mpi_disabled(): + original_group = list(group) + group = resolve_nccl_group(original_group) + if group != original_group: + raise RuntimeError( + "NCCL error: Helix alltoall cannot use survivor-only " + f"communicator {group} for statically sharded CP group " + f"{original_group}; redistribute the missing rank's state " + "and rebuild the mapping before resuming") + n_ranks = len(group) if n_ranks == 1: return inputs @@ -459,11 +495,20 @@ def reducescatter( dim: int = -1, sizes: Optional[List[int]] = None, ) -> Union[torch.Tensor, List[torch.Tensor]]: - if mapping.tp_size == 1: + if mpi_disabled(): + group = mapping.tp_group + group_size = mapping.tp_size + else: + group = mapping.tp_group + if NCCL_FAULT_TOLERANCE_ENABLED: + assert_nccl_group_not_reconfigured(group, "reducescatter") + group_size = len(group) + + if group_size == 1: return input if sizes is not None: - assert len(sizes) == len(mapping.tp_group) + assert len(sizes) == group_size sum_split_size = sum(sizes) if isinstance(input, torch.Tensor): assert input.shape[dim] == sum_split_size @@ -479,7 +524,7 @@ def convert_input(x, x_info): x = x.contiguous().view(-1, x_info['numel_base']) else: if sizes is None: - x_list = x.chunk(mapping.tp_size, dim=dim) + x_list = x.chunk(group_size, dim=dim) else: x_list = x.split(sizes, dim=dim) x = torch.cat([x.reshape(-1, x_info['numel_base']) for x in x_list]) @@ -505,10 +550,9 @@ def convert_input(x, x_info): ] if mpi_disabled(): - output = torch_op(input, sizes, mapping.tp_group, - mapping.tp_group_pg.boxed()) + output = torch_op(input, sizes, group, mapping.tp_group_pg.boxed()) else: - output = torch_op(input, sizes, mapping.tp_group) + output = torch_op(input, sizes, group) if isinstance(input, torch.Tensor): output = output.view(output_info['output_shape']) @@ -701,11 +745,22 @@ def __init__(self, """ self.mapping = mapping + self._tp_group_tuple = tuple(mapping.tp_group) self.workspace = None self.strategy = strategy self.mnnvl_allreduce = None self.symm_mem_allreduce = None self._disable_mpi = mpi_disabled() + self._fault_tolerant_raw_nccl = (NCCL_FAULT_TOLERANCE_ENABLED + and not self._disable_mpi) + + if (self._fault_tolerant_raw_nccl and self.strategy + not in (AllReduceStrategy.AUTO, AllReduceStrategy.NCCL)): + strategy_name = AllReduceStrategy(self.strategy).name + raise RuntimeError( + "NCCL error: TLLM_FAULT_TOLERANCE_MODE=1 requires a " + "watchdog-managed NCCL allreduce; strategy " + f"{strategy_name} has no peer-failure abort path") self.all_reduce_op = torch.ops.trtllm.allreduce_pg if self._disable_mpi else torch.ops.trtllm.allreduce @@ -729,7 +784,8 @@ def __init__(self, if self.mapping.tp_size > 1 and not self.mapping.enable_attention_dp: # Initialize Symmetric Memory AllReduce if needed (before workspace allocation) - if self.strategy == AllReduceStrategy.SYMM_MEM: + if (not self._fault_tolerant_raw_nccl + and self.strategy == AllReduceStrategy.SYMM_MEM): try: symm_mem = SymmetricMemoryAllReduce( self.mapping, @@ -758,7 +814,8 @@ def __init__(self, # Allocate workspace for strategies that need it # Note: SYMM_MEM now also needs workspace for fallback scenarios (fused ops, etc.) # Only UB doesn't need workspace - if self.strategy != AllReduceStrategy.UB: + if (not self._fault_tolerant_raw_nccl + and self.strategy != AllReduceStrategy.UB): if self.strategy == AllReduceStrategy.LOWPRECISION: allocate_low_presicion_allreduce_workspace(self.mapping) if self.strategy not in (AllReduceStrategy.UB, @@ -767,8 +824,8 @@ def __init__(self, self.workspace = get_allreduce_workspace(self.mapping) # Initialize MNNVL if using AUTO or MNNVL strategy - if self.strategy in (AllReduceStrategy.AUTO, - AllReduceStrategy.MNNVL): + if (not self._fault_tolerant_raw_nccl and self.strategy + in (AllReduceStrategy.AUTO, AllReduceStrategy.MNNVL)): # Try to initialize MNNVL if MNNVLAllReduce.is_mnnvl(self.mapping, dtype): # ALWAYS capture the exception when creating this instance @@ -791,6 +848,12 @@ def uses_nccl_symmetric_memory_window(self) -> bool: Requires TLLM_NCCL_SYMMETRIC_ZERO_COPY=1 AND NCCL_SYMMETRIC/NCCL/AUTO strategy AND tp_size > 1 AND MPI not disabled. """ + if (NCCL_FAULT_TOLERANCE_ENABLED and not self._disable_mpi + and is_nccl_group_reconfigured(self._tp_group_tuple)): + # Window registration and its allocation workspace belong to the + # original communicator. Forward will reject the incomplete static + # shard; avoid allocating another stale window before that check. + return False return (_NCCL_SYMMETRIC_ZERO_COPY and self.strategy in (AllReduceStrategy.NCCL_SYMMETRIC, AllReduceStrategy.NCCL, AllReduceStrategy.AUTO) and self.mapping.tp_size > 1 @@ -802,8 +865,9 @@ def output_buffer_kind(self) -> int: be passed into this allreduce. Returns int(BufferKind.NCCL_WINDOW) when zero-copy window output is - active, int(BufferKind.DEFAULT) otherwise. The value depends solely on - compile-time constants so it is safe to branch on inside torch.compile. + active, int(BufferKind.DEFAULT) otherwise. Dynamo guards the survivor + registry lookup and recompiles after membership is published at the + recovery safe point, so this remains safe inside ``torch.compile``. """ return (int(BufferKind.NCCL_WINDOW) if self.uses_nccl_symmetric_memory_window() else int( @@ -844,23 +908,36 @@ def forward( == False): return input + active_group = self.mapping.tp_group + if NCCL_FAULT_TOLERANCE_ENABLED and not self._disable_mpi: + assert_nccl_group_not_reconfigured(self._tp_group_tuple, + "allreduce") + input = input.contiguous() # Underlying op requires contiguous input allreduce_strategy = self.strategy + if (self._fault_tolerant_raw_nccl + and allreduce_strategy == AllReduceStrategy.AUTO): + # AUTO includes one-shot/two-shot and other custom transports with + # no abort hook. NCCL_SYMMETRIC also performs MPI-local topology + # discovery, whose communicator may still contain the dead rank. + allreduce_strategy = AllReduceStrategy.NCCL if all_reduce_params is None: all_reduce_params = AllReduceParams() # Try Symmetric Memory AllReduce first if available # Note: Currently only supports NONE fusion op (plain allreduce) - if self.symm_mem_allreduce and all_reduce_params.fusion_op == AllReduceFusionOp.NONE: + if (self.symm_mem_allreduce + and all_reduce_params.fusion_op == AllReduceFusionOp.NONE): symm_mem_output = self.symm_mem_allreduce(input) if symm_mem_output is not None: logger.debug( f"Using SymmetricMemoryAllReduce (MULTIMEM) for input shape {input.shape}" ) return symm_mem_output - elif self.symm_mem_allreduce and all_reduce_params.fusion_op != AllReduceFusionOp.NONE: + elif (self.symm_mem_allreduce + and all_reduce_params.fusion_op != AllReduceFusionOp.NONE): # Log once per rank that we're skipping symm_mem due to fusion logger.debug_once( f"Skipping SymmetricMemoryAllReduce for fused operation (fusion_op={all_reduce_params.fusion_op}), using regular allreduce", @@ -880,7 +957,6 @@ def forward( if allreduce_strategy in (AllReduceStrategy.MNNVL, AllReduceStrategy.SYMM_MEM): allreduce_strategy = AllReduceStrategy.AUTO - additional_args = {} if self._disable_mpi: # Get ProcessGroup from mapping @@ -896,7 +972,8 @@ def forward( disable_allreduce_autotune = os.environ.get( "TLLM_DISABLE_ALLREDUCE_AUTOTUNE", "0") == "1" - if allreduce_strategy == AllReduceStrategy.AUTO and not disable_allreduce_autotune and not self._disable_mpi: + if (allreduce_strategy == AllReduceStrategy.AUTO + and not disable_allreduce_autotune and not self._disable_mpi): output = torch.ops.trtllm.tunable_allreduce( input=input, residual=all_reduce_params.residual, @@ -904,7 +981,7 @@ def forward( scale=all_reduce_params.scale, bias=all_reduce_params.bias, workspace=self.workspace, - group=self.mapping.tp_group, + group=active_group, strategy=allreduce_strategy, op=all_reduce_params.fusion_op, eps=all_reduce_params.eps, @@ -919,7 +996,7 @@ def forward( scale=all_reduce_params.scale, bias=all_reduce_params.bias, workspace=self.workspace, - group=self.mapping.tp_group, + group=active_group, strategy=allreduce_strategy, op=all_reduce_params.fusion_op, eps=all_reduce_params.eps, @@ -961,6 +1038,11 @@ def __init__(self, mapping: Mapping): output_hidden_states = rms_norm(output_residual, norm_weight, eps) """ super().__init__() + if NCCL_FAULT_TOLERANCE_ENABLED and not mpi_disabled(): + raise RuntimeError( + "NCCL error: MoEAllReduce uses a custom collective without a " + "peer-failure abort path and is unavailable while " + "TLLM_FAULT_TOLERANCE_MODE=1") self.mapping = mapping self.workspace = get_allreduce_workspace(self.mapping) # Pls keep this value in sync with the kOneShotMaxToken in moeAllReduceFusionKernels.h @@ -1224,6 +1306,11 @@ class MiniMaxAllReduceRMS(nn.Module): def __init__(self, mapping: Mapping): super().__init__() + if NCCL_FAULT_TOLERANCE_ENABLED and not mpi_disabled(): + raise RuntimeError( + "NCCL error: MiniMaxAllReduceRMS uses a custom collective " + "without a peer-failure abort path and is unavailable while " + "TLLM_FAULT_TOLERANCE_MODE=1") self.mapping = mapping self.workspace = get_allreduce_workspace(self.mapping) diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/allgather_reducescatter.py b/tensorrt_llm/_torch/modules/fused_moe/communication/allgather_reducescatter.py index 706de4247a84..f5484865cc9b 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/communication/allgather_reducescatter.py +++ b/tensorrt_llm/_torch/modules/fused_moe/communication/allgather_reducescatter.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -12,7 +12,6 @@ # 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. - """ AllGather + ReduceScatter Communication Strategy @@ -27,11 +26,47 @@ import torch from tensorrt_llm._torch.distributed import allgather, reducescatter +from tensorrt_llm._torch.distributed.nccl_fault_tolerance import ( + NCCL_FAULT_TOLERANCE_ENABLED, + reconfigure_nccl_group, + resolve_nccl_group, +) +from tensorrt_llm._utils import mpi_disabled from tensorrt_llm.mapping import Mapping from .base import Communication +class _ActiveGroupMapping: + """Read-only mapping view for a survivor-only NCCL communicator. + + The model-wide :class:`Mapping` remains immutable because TP/PP users may + still be inspecting it while the MoE fallback switches membership. + Attributes unrelated to the TP communicator delegate to the original + mapping. + """ + + def __init__(self, mapping: Mapping, group: List[int], rank: int) -> None: + self._mapping = mapping + self._group = tuple(group) + self._rank = rank + + @property + def tp_group(self) -> List[int]: + return list(self._group) + + @property + def tp_size(self) -> int: + return len(self._group) + + @property + def tp_rank(self) -> int: + return self._rank + + def __getattr__(self, name): + return getattr(self._mapping, name) + + class AllGatherReduceScatter(Communication): def __init__( self, @@ -41,6 +76,11 @@ def __init__( # Initialize dispatch state self._dispatch_state = {} + self._original_group = tuple(mapping.tp_group) + self._active_local_ranks = tuple(range(len(self._original_group))) + self._active_global_group = self._original_group + self._active_mapping = mapping + self._raw_nccl_fault_tolerance = NCCL_FAULT_TOLERANCE_ENABLED and not mpi_disabled() @staticmethod def is_platform_supported() -> bool: @@ -57,6 +97,113 @@ def is_workload_feasible(self, all_rank_num_tokens: List[int], num_chunks: int) """ return True + def abort_and_reinit(self, active_ranks: List[int], generation: Optional[int] = None) -> None: + """Abort the current NCCL communicator and rebuild it for survivors. + + ``active_ranks`` contains global/world rank IDs from the original TP + communication group. A global-rank API is intentional: when MoE TP is + enabled, one TP communicator spans several MoE-EP slices, so an + EP-local rank number is ambiguous without the failure coordinator's + slice identity. The coordinator must translate its health snapshot to + global ranks before calling this method. + + Reconfiguration is monotonic: a removed rank cannot be reintroduced by + this Phase-1 API. Replacement-rank joins belong to process-group + reconstruction. The failure coordinator must pass a shared monotonic + recovery-event generation for same-membership recovery and advance it + for every distinct attempt, including transport-only retries. Passing + it for every recovery also deduplicates callbacks on every survivor. + """ + if not NCCL_FAULT_TOLERANCE_ENABLED: + raise RuntimeError( + "NCCL error: communicator reinitialization requires TLLM_FAULT_TOLERANCE_MODE=1" + ) + if not self._raw_nccl_fault_tolerance: + raise NotImplementedError( + "AllGatherReduceScatter fault tolerance requires the raw NCCL/MPI path; " + "ProcessGroup reconstruction is not implemented" + ) + + if not active_ranks: + raise ValueError("active_ranks must not be empty") + if len(active_ranks) != len(set(active_ranks)): + raise ValueError("active_ranks must not contain duplicates") + requested = {int(rank) for rank in active_ranks} + if not requested.issubset(self._original_group): + raise ValueError( + "active_ranks must be global ranks from the original TP group " + f"{list(self._original_group)}, got {active_ranks}" + ) + if self.mapping.rank not in requested: + raise ValueError(f"current world rank {self.mapping.rank} is not active") + + active_global_group = tuple(rank for rank in self._original_group if rank in requested) + + # Commit Python metadata only after the coordinated native rebuild + # succeeds. The old transport is terminal once native recovery starts, + # even if construction of its replacement fails. + reconfigure_nccl_group( + self._original_group, + active_global_group, + torch.ops.trtllm.nccl_comm_abort_and_reinit, + generation, + ) + self._refresh_active_mapping() + + def _refresh_active_mapping(self) -> None: + """Synchronize this layer with process-wide survivor membership.""" + active_global_group = tuple(resolve_nccl_group(self._original_group)) + if active_global_group == self._active_global_group: + return + global_rank = self.mapping.rank + if global_rank not in active_global_group: + raise RuntimeError( + f"NCCL error: world rank {global_rank} is not in the active communicator" + ) + active_global_ranks = set(active_global_group) + active_local_ranks = tuple( + index + for index, world_rank in enumerate(self._original_group) + if world_rank in active_global_ranks + ) + active_mapping = _ActiveGroupMapping( + self.mapping, + list(active_global_group), + active_global_group.index(global_rank), + ) + self._active_local_ranks = active_local_ranks + self._active_mapping = active_mapping + # Publish the value used by the fast-path guard last. A concurrent + # reader that observes this generation must also observe its mapping. + self._active_global_group = active_global_group + + def check_async_error(self) -> None: + """Raise a classifier-friendly exception for a watchdog failure.""" + if not NCCL_FAULT_TOLERANCE_ENABLED: + raise RuntimeError( + "NCCL error: async-error monitoring requires TLLM_FAULT_TOLERANCE_MODE=1" + ) + if not self._raw_nccl_fault_tolerance: + raise NotImplementedError( + "AllGatherReduceScatter fault tolerance cannot inspect ProcessGroup errors; " + "the raw NCCL/MPI watchdog is not active" + ) + self._refresh_active_mapping() + error = torch.ops.trtllm.nccl_comm_get_async_error(list(self._active_global_group)) + if error is not None: + raise RuntimeError(f"NCCL error: communicator was aborted: {error}") + + def _active_sizes(self, all_rank_num_tokens: List[int]) -> List[int]: + self._refresh_active_mapping() + if len(all_rank_num_tokens) != len(self._original_group): + raise ValueError( + "all_rank_num_tokens must remain indexed by the original communicator " + f"({len(self._original_group)} entries), got {len(all_rank_num_tokens)}" + ) + if self._active_global_group == self._original_group: + return all_rank_num_tokens + return [all_rank_num_tokens[rank] for rank in self._active_local_ranks] + def dispatch( self, hidden_states: torch.Tensor, @@ -70,11 +217,20 @@ def dispatch( """ AllGather dispatch (always post-quant dispatch) """ - sizes = None if use_dp_padding else all_rank_num_tokens + if self._raw_nccl_fault_tolerance: + if use_dp_padding: + self._refresh_active_mapping() + sizes = None + else: + sizes = self._active_sizes(all_rank_num_tokens) + mapping = self._active_mapping + else: + sizes = None if use_dp_padding else all_rank_num_tokens + mapping = self.mapping hidden_states, hidden_states_sf, token_selected_slots, token_final_scales = allgather( [hidden_states, hidden_states_sf, token_selected_slots, token_final_scales], - self.mapping, + mapping, dim=0, sizes=sizes, ) @@ -92,7 +248,15 @@ def combine( """ ReduceScatter combine phase """ + if self._raw_nccl_fault_tolerance: + self._refresh_active_mapping() + mapping = self._active_mapping + else: + mapping = self.mapping outputs = reducescatter( - final_hidden_states, self.mapping, dim=0, sizes=self._dispatch_state.get("sizes") + final_hidden_states, + mapping, + dim=0, + sizes=self._dispatch_state.get("sizes"), ) return outputs diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/communication_factory.py b/tensorrt_llm/_torch/modules/fused_moe/communication/communication_factory.py index a4ec2ceefe44..17a67deee219 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/communication/communication_factory.py +++ b/tensorrt_llm/_torch/modules/fused_moe/communication/communication_factory.py @@ -28,6 +28,7 @@ from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm.logger import logger +from ..wide_ep_ft import get_wide_ep_ft_options from .allgather_reducescatter import AllGatherReduceScatter from .base import Communication from .deep_ep import DeepEP @@ -133,6 +134,9 @@ def create_strategy( try: enable_eplb = model_config.moe_load_balancer is not None + ep_group_health, watchdog_timeout_s, watchdog_poll_interval_s = get_wide_ep_ft_options( + model_config + ) strategy = NVLinkOneSided( mapping, num_slots, @@ -143,6 +147,9 @@ def create_strategy( dtype=act_dtype, num_experts=num_experts if enable_eplb else None, use_low_precision_combine=use_low_precision_combine, + ep_group_health=ep_group_health, + alltoall_watchdog_timeout_s=watchdog_timeout_s, + alltoall_watchdog_poll_interval_s=watchdog_poll_interval_s, ) logger.info("Selected communication strategy: NVLinkOneSided") return strategy @@ -285,6 +292,9 @@ def _create_forced_method( ) elif method in ["NVLINK_ONE_SIDED"]: enable_eplb = model_config.moe_load_balancer is not None + ep_group_health, watchdog_timeout_s, watchdog_poll_interval_s = get_wide_ep_ft_options( + model_config + ) return NVLinkOneSided( mapping, num_slots, @@ -295,6 +305,9 @@ def _create_forced_method( dtype=act_dtype, num_experts=num_experts if enable_eplb else None, use_low_precision_combine=use_low_precision_combine, + ep_group_health=ep_group_health, + alltoall_watchdog_timeout_s=watchdog_timeout_s, + alltoall_watchdog_poll_interval_s=watchdog_poll_interval_s, ) elif method == "DEEPEP": return DeepEP( diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_one_sided.py b/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_one_sided.py index f9068f71d717..8f0b3c9486aa 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_one_sided.py +++ b/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_one_sided.py @@ -25,11 +25,20 @@ """ import os -from typing import Dict, List, Optional, Tuple +from typing import Callable, Dict, List, Optional, Tuple import torch from tensorrt_llm._mnnvl_utils import MnnvlMemory +from tensorrt_llm._torch.alltoall_watchdog import ( + DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, + DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S, + ActiveRankMaskSnapshot, + AlltoAllWatchdog, + AlltoAllWatchdogCoordinator, + AlltoAllWatchdogTimeout, + EPGroupHealthLike, +) from tensorrt_llm.bindings import internal as _tllm_internal from tensorrt_llm.logger import logger as tllm_logger from tensorrt_llm.mapping import Mapping @@ -151,7 +160,11 @@ def __init__( dtype: Optional[torch.dtype] = None, num_experts: Optional[int] = None, use_low_precision_combine: bool = False, - ): + ep_group_health: EPGroupHealthLike | None = None, + alltoall_watchdog_timeout_s: Optional[float] = None, + alltoall_watchdog_poll_interval_s: float = DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, + alltoall_watchdog_on_timeout: Optional[Callable[[AlltoAllWatchdogTimeout], None]] = None, + ) -> None: """ Initialize NVLinkOneSided with workspace allocation. @@ -169,6 +182,12 @@ def __init__( use_low_precision_combine: If True, quantize the combine payload to FP8 for NVLink transfer (halves NVLink bandwidth usage, output precision is preserved). Corresponds to model_config.use_low_precision_moe_combine. + ep_group_health: Optional read-only committed EP membership. When present, its mask is passed to + the CUDA kernels and defines the peers expected by the watchdog. Timeout detection never mutates it. + alltoall_watchdog_timeout_s: Optional timeout for the host-side AlltoAll watchdog. If None, the + watchdog is disabled. + alltoall_watchdog_poll_interval_s: Poll interval for the watchdog thread. + alltoall_watchdog_on_timeout: Optional callback invoked when the watchdog reports suspects. """ super().__init__(mapping) @@ -296,10 +315,34 @@ def __init__( NVLinkOneSided._WORKSPACE_REFCOUNTS.get(self._workspace_key, 0) + 1 ) self._destroyed = False + self._workspace_state = workspace_state self.mnnvl_mem = workspace_state["mnnvl_mem"] self.workspace = workspace_state["workspace"] self.moe_a2a_metainfo = workspace_state["metainfo"] self.max_num_tokens_per_rank = workspace_state["max_num_tokens_per_rank"] + self.ep_group_health = ep_group_health + self._watchdog_coordinator = AlltoAllWatchdogCoordinator( + workspace_state=workspace_state, + workspace=self.workspace, + metainfo=self.moe_a2a_metainfo, + metainfo_index={ + "FLAG_VAL_OFFSET_INDEX": self.FLAG_VAL_OFFSET_INDEX, + "DISPATCH_COMPLETION_FLAGS_OFFSET_INDEX": self.DISPATCH_COMPLETION_FLAGS_OFFSET_INDEX, + "COMBINE_COMPLETION_FLAGS_OFFSET_INDEX": self.COMBINE_COMPLETION_FLAGS_OFFSET_INDEX, + }, + ep_rank=self.ep_rank, + health=self.ep_group_health, + ) + self._alltoall_watchdog: AlltoAllWatchdog | None = None + if alltoall_watchdog_timeout_s is None and self.ep_group_health is not None: + alltoall_watchdog_timeout_s = DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S + if alltoall_watchdog_timeout_s is not None: + self._alltoall_watchdog = self._watchdog_coordinator.acquire_watchdog( + ep_size=self.ep_size, + timeout_s=alltoall_watchdog_timeout_s, + poll_interval_s=alltoall_watchdog_poll_interval_s, + on_timeout=alltoall_watchdog_on_timeout, + ) # Initialize dispatch state self._dispatch_state = {"phase": "idle"} @@ -326,6 +369,9 @@ def destroy(self): return self._destroyed = True + if self._alltoall_watchdog is not None: + self._watchdog_coordinator.release_watchdog(self._alltoall_watchdog) + self._alltoall_watchdog = None workspace_key = getattr(self, "_workspace_key", None) if workspace_key is None: return @@ -347,6 +393,7 @@ def destroy(self): self.mnnvl_mem = None self.workspace = None self.moe_a2a_metainfo = None + self._workspace_state = None self._dispatch_state = {"phase": "destroyed"} def is_workload_feasible(self, all_rank_num_tokens: List[int], num_chunks: int) -> bool: @@ -378,7 +425,9 @@ def dispatch( token_final_scales: Router weights [local_num_tokens, top_k] all_rank_num_tokens: Token counts per rank [ep_size] use_dp_padding: Whether to use DP padding (optional) - **kwargs: Strategy-specific arguments (unused) + **kwargs: Strategy-specific arguments. ``active_rank_mask`` may override the committed membership + for dispatch. Without an override, the committed mask and generation are captured together; + combine reuses that mask and fails closed if the generation changes first. Returns: Tuple of (hidden_states, hidden_states_sf, token_selected_slots, token_final_scales) @@ -409,6 +458,10 @@ def dispatch( assert eplb_local_stats.size(0) == self.eplb_stats_num_experts, ( "eplb_local_stats size must match eplb_stats_num_experts" ) + active_rank_mask_snapshot = self._watchdog_coordinator.capture_active_rank_mask( + kwargs.get("active_rank_mask") + ) + active_rank_mask = active_rank_mask_snapshot.active_rank_mask recv_buffers, combine_payload_offset, eplb_gathered_stats = ( torch.ops.trtllm.moe_a2a_dispatch( @@ -422,14 +475,19 @@ def dispatch( self.top_k, self.num_experts, eplb_local_stats, + active_rank_mask, ) ) + self._watchdog_coordinator.watch_collective( + self._alltoall_watchdog, "dispatch", active_rank_mask + ) if eplb_gathered_stats.numel() == 0: eplb_gathered_stats = None self._dispatch_state["eplb_gathered_stats"] = eplb_gathered_stats self._dispatch_state["combine_payload_offset"] = int(combine_payload_offset) self._dispatch_state["local_num_tokens"] = token_selected_slots.size(0) self._dispatch_state["runtime_max_tokens_per_rank"] = runtime_max_tokens_per_rank + self._dispatch_state["active_rank_mask_snapshot"] = active_rank_mask_snapshot self._dispatch_state["phase"] = "dispatched" # Extract results from recv_buffers @@ -492,6 +550,9 @@ def combine( final_hidden_states: Output from MoE computation Shape: [ep_size, max_tokens_per_rank, hidden_size] or [ep_size * max_tokens_per_rank, hidden_size] (will be reshaped) + **kwargs: Strategy-specific arguments. If ``active_rank_mask`` is supplied, it must match the mask + captured by dispatch for this collective. A committed-generation change since dispatch aborts + the collective epoch. Returns: Combined output tensor [local_num_tokens, hidden_size] @@ -526,6 +587,13 @@ def combine( raise ValueError( f"final_hidden_states must be 2D or 3D, got {final_hidden_states.dim()}D" ) + active_rank_mask_snapshot = self._dispatch_state.get("active_rank_mask_snapshot") + if not isinstance(active_rank_mask_snapshot, ActiveRankMaskSnapshot): + raise RuntimeError("combine called but dispatch rank-mask snapshot is missing") + active_rank_mask = self._watchdog_coordinator.active_rank_mask_for_combine( + active_rank_mask_snapshot, + kwargs.get("active_rank_mask"), + ) output = torch.ops.trtllm.moe_a2a_combine( final_hidden_states, int(local_num_tokens), @@ -538,6 +606,10 @@ def combine( int(combine_payload_offset), bool(self.payload_in_workspace), bool(self.use_low_precision_combine), + active_rank_mask, + ) + self._watchdog_coordinator.watch_collective( + self._alltoall_watchdog, "combine", active_rank_mask ) # Reset state for next round diff --git a/tensorrt_llm/_torch/modules/fused_moe/ep_group_health.py b/tensorrt_llm/_torch/modules/fused_moe/ep_group_health.py index 1864d4ba825d..c449b4bfb059 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/ep_group_health.py +++ b/tensorrt_llm/_torch/modules/fused_moe/ep_group_health.py @@ -13,25 +13,30 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""EP group health tracking for WideEP fault tolerance. +"""Committed EP data-plane membership for WideEP fault tolerance. This module provides :class:`EPGroupHealth`, a process-local, thread-safe data -structure that records which ranks in an Expert Parallel (EP) group are currently -alive vs. failed. It is the single source of truth for EP rank health within one -process and is consumed by: +structure that records which Expert Parallel (EP) ranks the recovery coordinator +has committed as included vs. excluded from the data plane. It is consumed by: * AlltoAll communication backends (rank masking on dispatch / combine) - * The host-side AlltoAll watchdog (failure detection) + * The host-side AlltoAll watchdog (read-only expected-peer snapshot) * The MoE load balancer (emergency-mask reconfiguration) * The model engine and PyExecutor (degraded health reporting) -Cross-process consensus on which ranks are dead (failure broadcast across the EP -group) is the responsibility of higher-layer coordination components and is not -performed here. +Detected or suspected physical liveness is separate evidence. Higher-layer +coordination reconciles that evidence and commits membership; detectors and +telemetry consumers must not mutate this object to drive recovery. """ import threading -from typing import NamedTuple +import uuid +from typing import Any, NamedTuple + +# Canonical ModelConfig.extra_attrs key for coordinator-committed data-plane +# membership. The object is not a detector or a physical-liveness tracker. +EP_GROUP_HEALTH_EXTRA_ATTR = "wide_ep_ft_ep_group_health" +EP_GROUP_HEALTH_LEGACY_EXTRA_ATTR = "ep_group_health" # Default number of uint64 words for EPGroupHealth.get_mask_words(). # Matches the active_rank_mask ABI of the NVLink AlltoAll kernels @@ -50,24 +55,29 @@ class EPGroupHealthSnapshot(NamedTuple): if a mutator runs between calls). """ - # Active-rank bitmask (bit i set iff rank i is active). + # Committed-rank bitmask (bit i set iff rank i is included). mask: int - # Number of ranks currently marked active. + # Number of ranks included in committed data-plane membership. active_count: int - # Immutable snapshot of failed ranks. + # Immutable snapshot of ranks excluded from committed membership. failed_ranks: frozenset[int] - # Monotonic counter; bumps only on effective state changes. + # Monotonic committed-membership version. generation: int class EPGroupHealth: - """Thread-safe health tracker for the ranks of an EP group. + """Thread-safe committed-membership tracker for an EP group. Internally backed by an arbitrary-precision Python ``int`` bitmask - (bit ``i`` set iff rank ``i`` is currently active). The kernel-side - representation as a fixed-width array of ``uint64`` words is exposed via - :meth:`get_mask_words`; this matches the ``active_rank_mask`` parameter - expected by the NVLink AlltoAll kernels. + (bit ``i`` set iff rank ``i`` is included in the committed data plane). + The kernel-side representation as a fixed-width array of ``uint64`` words + is exposed via :meth:`get_mask_words`; this matches the + ``active_rank_mask`` parameter expected by the NVLink AlltoAll kernels. + + This object does not represent uncommitted detector observations or + suspected physical liveness. A higher-level recovery coordinator owns + membership transitions; passive consumers such as telemetry must only + read coherent snapshots. All read and write operations take an internal lock. Read operations that return collection types return defensive snapshots so the caller cannot @@ -99,6 +109,10 @@ def __init__(self, moe_world_size: int) -> None: self._active_count: int = moe_world_size self._failed_ranks: set[int] = set() self._generation: int = 0 + # A generation is monotonic only within one producer lifetime. The + # immutable epoch lets serving distinguish a restarted producer from a + # stale snapshot emitted by the current one. + self._source_epoch = uuid.uuid4().hex self._lock = threading.Lock() @property @@ -108,7 +122,7 @@ def moe_world_size(self) -> int: @property def generation(self) -> int: - """Monotonic counter incremented on every effective state change. + """Monotonic counter incremented on each committed membership change. Idempotent calls (e.g. marking an already-failed rank as failed) do not bump the counter. Consumers that need to react to mask changes can @@ -121,13 +135,22 @@ def generation(self) -> int: with self._lock: return self._generation + @property + def source_epoch(self) -> str: + """Unique identifier for this producer lifetime.""" + return self._source_epoch + def _validate_rank(self, rank: int) -> None: """Raise if ``rank`` is outside the MoE world.""" if not 0 <= rank < self._moe_world_size: raise ValueError(f"rank must be in [0, {self._moe_world_size}), got {rank}") def mark_failed(self, rank: int) -> bool: - """Mark ``rank`` as failed. Idempotent. + """Exclude ``rank`` from committed data-plane membership. Idempotent. + + A higher-level recovery coordinator owns calls to this mutator after + reconciling detector evidence; detection and telemetry must not call it + directly. Args: rank: Index of the rank to mark, in ``[0, moe_world_size)``. @@ -151,7 +174,7 @@ def mark_failed(self, rank: int) -> bool: return True def mark_active(self, rank: int) -> bool: - """Mark ``rank`` as active. Idempotent. + """Include ``rank`` in committed data-plane membership. Idempotent. Used when a replacement rank rejoins the group after a failure. Higher layers may impose a "monotonic failure" policy (do not @@ -180,7 +203,7 @@ def mark_active(self, rank: int) -> bool: return True def is_active(self, rank: int) -> bool: - """Return ``True`` iff ``rank`` is currently marked active. + """Return whether ``rank`` is included in committed membership. Raises: ValueError: If ``rank`` is outside ``[0, moe_world_size)``. @@ -190,15 +213,15 @@ def is_active(self, rank: int) -> bool: return bool(self._active_mask & (1 << rank)) def get_mask(self) -> int: - """Return the active-rank bitmask as a Python ``int``. + """Return the committed active-rank bitmask as a Python ``int``. - Bit ``i`` is set iff rank ``i`` is currently active. + Bit ``i`` is set iff rank ``i`` is included in data-plane membership. """ with self._lock: return self._active_mask def get_mask_words(self, num_words: int = EP_MASK_NUM_WORDS) -> tuple[int, ...]: - """Return the active-rank bitmask split into little-endian uint64 words. + """Return the committed rank mask as little-endian uint64 words. Suitable for passing to CUDA kernels that accept ``uint64_t[num_words]``. Word ``0`` covers ranks ``0..63``, word ``1`` covers ranks ``64..127``, @@ -229,22 +252,22 @@ def get_mask_words(self, num_words: int = EP_MASK_NUM_WORDS) -> tuple[int, ...]: return tuple((mask >> (i * 64)) & word_mask for i in range(num_words)) def get_active_count(self) -> int: - """Return the number of ranks currently marked active.""" + """Return the number of ranks in committed data-plane membership.""" with self._lock: return self._active_count def get_failed_ranks(self) -> frozenset[int]: - """Return an immutable snapshot of the set of failed ranks.""" + """Return ranks excluded from committed data-plane membership.""" with self._lock: return frozenset(self._failed_ranks) def all_active(self) -> bool: - """Return ``True`` iff every rank in the group is currently active.""" + """Return whether committed membership includes every group rank.""" with self._lock: return self._active_mask == self._all_active_mask def snapshot(self) -> EPGroupHealthSnapshot: - """Return an atomic snapshot of mask, active_count, failed_ranks, generation. + """Return an atomic committed-membership snapshot. All four fields reflect a single point in time (one lock acquisition). Prefer this over calling individual accessors back-to-back when the @@ -262,13 +285,18 @@ def __len__(self) -> int: """Return the MoE world size (``moe_world_size``), not the active count. Provided for compatibility with code that treats this object as a - sized container of ranks (alive or dead). Use :meth:`get_active_count` - when you specifically want the surviving-rank count. + sized container of included and excluded ranks. Use + :meth:`get_active_count` for the committed-membership count. """ return self._moe_world_size + def __deepcopy__(self, memo: dict[int, Any]) -> "EPGroupHealth": + """Preserve producer identity when ``ModelConfig`` is deep-copied.""" + memo[id(self)] = self + return self + def __repr__(self) -> str: - """Return a concise debug representation of the current health state.""" + """Return a concise representation of committed membership.""" with self._lock: return ( f"EPGroupHealth(moe_world_size={self._moe_world_size}, " diff --git a/tensorrt_llm/_torch/modules/fused_moe/ep_metrics.py b/tensorrt_llm/_torch/modules/fused_moe/ep_metrics.py new file mode 100644 index 000000000000..121da2757ac1 --- /dev/null +++ b/tensorrt_llm/_torch/modules/fused_moe/ep_metrics.py @@ -0,0 +1,116 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. +"""Prometheus-independent serialization of committed EP membership. + +Workers serialize the recovery coordinator's committed data-plane membership +for a dedicated serving RPC without importing the client before its +multiprocess directory is configured. This module is a passive observer: it +does not expose detected/suspected physical liveness and must not drive +recovery. +""" + +from typing import Any + +from .ep_group_health import ( + EP_GROUP_HEALTH_EXTRA_ATTR, + EP_GROUP_HEALTH_LEGACY_EXTRA_ATTR, + EPGroupHealth, +) + +_EP_HEALTH_METRICS_STATUS_KEY = "status" +_EP_HEALTH_METRICS_PENDING_STATUS = "pending" + + +def validate_ep_health_metrics_topology(model_config: Any) -> None: + """Reject topologies whose rank-0 endpoint cannot expose every EP group.""" + mapping = getattr(model_config, "mapping", None) + if any( + getattr(mapping, axis, 1) > 1 for axis in ("moe_tp_size", "pp_size", "moe_cluster_size") + ): + raise NotImplementedError("EP health telemetry does not support multi-group MoE topologies") + + +def get_ep_group_health(model_config: Any) -> EPGroupHealth | None: + """Return an already-configured coordinator-owned membership tracker. + + ``ModelConfig.extra_attrs`` shares the committed data-plane membership + tracker with communication. It does not carry raw detector observations or + suspected physical liveness. This telemetry consumer never creates, + registers, enables, or mutates it. The legacy prototype key remains + readable during rollout; if both keys are populated they must reference the + same object. + """ + extra_attrs = getattr(model_config, "extra_attrs", None) + if extra_attrs is None: + return None + + health = extra_attrs.get(EP_GROUP_HEALTH_EXTRA_ATTR) + legacy_health = extra_attrs.get(EP_GROUP_HEALTH_LEGACY_EXTRA_ATTR) + if health is not None and legacy_health is not None and health is not legacy_health: + raise ValueError("configured EP group health keys reference different trackers") + health = health if health is not None else legacy_health + if health is None: + return None + if not isinstance(health, EPGroupHealth): + raise TypeError("configured EP group health must be an EPGroupHealth") + + mapping = getattr(model_config, "mapping", None) + moe_world_size = getattr(mapping, "moe_ep_size", None) + if ( + type(moe_world_size) is int + and moe_world_size > 0 + and health.moe_world_size != moe_world_size + ): + raise ValueError("configured EP group health world size does not match model mapping") + return health + + +def is_ep_group_health_registration_pending(model_config: Any) -> bool: + """Return whether a producer reserved a tracker slot for late attachment.""" + extra_attrs = getattr(model_config, "extra_attrs", None) + if extra_attrs is None: + return False + return any( + key in extra_attrs and extra_attrs[key] is None + for key in (EP_GROUP_HEALTH_EXTRA_ATTR, EP_GROUP_HEALTH_LEGACY_EXTRA_ATTR) + ) + + +def pending_ep_health_metrics() -> dict[str, str]: + """Return the JSON-safe marker for an explicitly pending registration.""" + return {_EP_HEALTH_METRICS_STATUS_KEY: _EP_HEALTH_METRICS_PENDING_STATUS} + + +def is_pending_ep_health_metrics(ep_health_stats: dict[str, Any]) -> bool: + """Return whether a worker reported an explicitly pending registration.""" + return ep_health_stats == pending_ep_health_metrics() + + +def serialize_ep_health_metrics(ep_group_health: EPGroupHealth) -> dict[str, Any]: + """Serialize one coherent committed-membership snapshot. + + ``failedRanks`` contains ranks excluded from the committed data-plane mask; + it is not a physical-liveness or suspicion signal. The per-rank gauge is + derived from ``worldSize`` and ``failedRanks`` without mutating recovery + state. + """ + snapshot = ep_group_health.snapshot() + return { + "sourceEpoch": ep_group_health.source_epoch, + "worldSize": ep_group_health.moe_world_size, + "activeCount": snapshot.active_count, + "failedRanks": sorted(snapshot.failed_ranks), + "generation": snapshot.generation, + } diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py index 3a89ec2fb35d..0b619bb31916 100755 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py @@ -20,6 +20,7 @@ Fp4QuantizedTensor) from .interface import AlltoallMethodType, MoE from .quantization import UnquantizedFusedMoEMethod +from .wide_ep_ft import get_wide_ep_ft_options # isort: off from .quantization import ( @@ -331,6 +332,8 @@ def __init__( dtype, self.num_experts if self.layer_load_balancer else None, ) + ep_group_health, watchdog_timeout_s, watchdog_poll_interval_s = ( + get_wide_ep_ft_options(model_config)) self.moe_a2a = MoeAlltoAll( mapping=self.mapping, @@ -340,6 +343,10 @@ def __init__( workspace_size_per_rank=workspace_size, num_experts=self.num_experts if self.layer_load_balancer else None, + ep_group_health=ep_group_health, + alltoall_watchdog_timeout_s=watchdog_timeout_s, + alltoall_watchdog_poll_interval_s= + watchdog_poll_interval_s, ) elif self.alltoall_method_type == AlltoallMethodType.DeepEP or self.alltoall_method_type == AlltoallMethodType.DeepEPLowLatency: raise NotImplementedError( diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py index 820d92a95e4a..6f59870c374b 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py @@ -36,6 +36,7 @@ from ...utils import ActivationType, AuxStreamType, Fp4QuantizedTensor from .interface import AlltoallMethodType, MoE, MoEWeightLoadingMode from .moe_op_backend import MoEOpBackend, get_op_backend +from .wide_ep_ft import get_wide_ep_ft_options # isort: off from .quantization import ( @@ -291,6 +292,8 @@ def __init__( ep_size, self.routing_method.experts_per_token, max_num_tokens, hidden_size, dtype, self.num_experts if self.layer_load_balancer else None) + ep_group_health, watchdog_timeout_s, watchdog_poll_interval_s = ( + get_wide_ep_ft_options(model_config)) self.moe_a2a = MoeAlltoAll( mapping=self.mapping, @@ -299,7 +302,11 @@ def __init__( num_slots=self.num_slots, workspace_size_per_rank=workspace_size, num_experts=self.num_experts - if self.layer_load_balancer else None) + if self.layer_load_balancer else None, + ep_group_health=ep_group_health, + alltoall_watchdog_timeout_s=watchdog_timeout_s, + alltoall_watchdog_poll_interval_s= + watchdog_poll_interval_s) elif self.alltoall_method_type == AlltoallMethodType.DeepEP or self.alltoall_method_type == AlltoallMethodType.DeepEPLowLatency: raise NotImplementedError( "DeepEP and DeepEPLowLatency are not supported for TRTLLMGenFusedMoE yet" diff --git a/tensorrt_llm/_torch/modules/fused_moe/wide_ep_ft.py b/tensorrt_llm/_torch/modules/fused_moe/wide_ep_ft.py new file mode 100644 index 000000000000..fb1586ae1054 --- /dev/null +++ b/tensorrt_llm/_torch/modules/fused_moe/wide_ep_ft.py @@ -0,0 +1,87 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. +"""Shared WideEP fault-tolerance options for MoE communication paths.""" + +from __future__ import annotations + +import os +from typing import Any, Optional + +from tensorrt_llm._torch.alltoall_watchdog import ( + DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, + DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S, +) + +from .ep_group_health import EPGroupHealth + +_ENABLE_ENV = "TLLM_FAULT_TOLERANCE_MODE" +_TIMEOUT_ENV = "TRTLLM_ALLTOALL_WATCHDOG_TIMEOUT_S" +_POLL_INTERVAL_ENV = "TRTLLM_ALLTOALL_WATCHDOG_POLL_INTERVAL_S" + +# This object contains membership committed by higher-layer recovery +# coordination. Detection threads must treat it as read-only. +_HEALTH_KEY = "wide_ep_ft_ep_group_health" +_TIMEOUT_KEY = "alltoall_watchdog_timeout_s" +_POLL_INTERVAL_KEY = "alltoall_watchdog_poll_interval_s" + + +def _env_enabled() -> bool: + return os.environ.get(_ENABLE_ENV) == "1" + + +def _float_option(extra_attrs: dict, key: str, env_name: str, default: float) -> float: + if key in extra_attrs: + return float(extra_attrs[key]) + if env_name in os.environ: + return float(os.environ[env_name]) + return default + + +def get_wide_ep_ft_options( + model_config: Any, +) -> tuple[Optional[EPGroupHealth], Optional[float], float]: + """Return committed EP membership and watchdog timing for a model. + + WideEP FT remains opt-in until the integration PR wires a public model + option. Callers can either inject ``wide_ep_ft_ep_group_health`` through + ``ModelConfig.extra_attrs`` or set ``TLLM_FAULT_TOLERANCE_MODE=1`` to create + one process-local membership object shared by all MoE communication layers. + The AlltoAll watchdog reads this object to determine expected peers and + reports suspects through its ``on_timeout`` seam; it never mutates the + committed membership directly. + """ + + extra_attrs = getattr(model_config, "extra_attrs", {}) + health = extra_attrs.get(_HEALTH_KEY) or extra_attrs.get("ep_group_health") + if health is None and _env_enabled(): + health = EPGroupHealth(model_config.mapping.moe_ep_size) + extra_attrs[_HEALTH_KEY] = health + + poll_interval_s = _float_option( + extra_attrs, + _POLL_INTERVAL_KEY, + _POLL_INTERVAL_ENV, + DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, + ) + if health is None: + return None, None, poll_interval_s + + timeout_s = _float_option( + extra_attrs, + _TIMEOUT_KEY, + _TIMEOUT_ENV, + DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S, + ) + return health, timeout_s, poll_interval_s diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index 9ff6f28efb9f..a106fcfd2ada 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + from __future__ import annotations import enum @@ -14,6 +17,8 @@ import tensorrt_llm.quantization.utils.fp4_utils as fp4_utils from tensorrt_llm._torch.custom_ops.torch_custom_ops import BufferKind +from tensorrt_llm._torch.distributed.nccl_fault_tolerance import ( + NCCL_FAULT_TOLERANCE_ENABLED, assert_nccl_group_not_reconfigured) from tensorrt_llm._torch.peft.lora.layer import LoraLayer from tensorrt_llm._utils import is_device_integrated, mpi_disabled from tensorrt_llm.bindings import ipc_nvls_supported @@ -3188,6 +3193,7 @@ def __init__( ) self.tp_size = self.mapping.tp_size self.tp_rank = self.mapping.tp_rank + self._tp_group_tuple = tuple(self.mapping.tp_group) self.tp_mode = tensor_parallel_mode self.gather_output = gather_output self.force_dynamic_quantization = force_dynamic_quantization @@ -3250,10 +3256,17 @@ def __init__( "TRTLLM_GEMM_ALLREDUCE_FUSION_ENABLED", "0") == "1") self.use_fused_gemm_allreduce = all([ - self.reduce_output, mpi_enabled, dtype_supported, - in_features_aligned, out_features_aligned, tp_valid, quant_valid, - device_supported, enable_gemm_allreduce_fusion, - enable_gemm_allreduce_fusion_env + self.reduce_output, + mpi_enabled, + dtype_supported, + in_features_aligned, + out_features_aligned, + tp_valid, + quant_valid, + device_supported, + enable_gemm_allreduce_fusion, + enable_gemm_allreduce_fusion_env, + not (NCCL_FAULT_TOLERANCE_ENABLED and mpi_enabled), ]) if self.use_fused_gemm_allreduce: self.use_fused_gemm_allreduce = ipc_nvls_supported() @@ -3368,8 +3381,14 @@ def apply_linear_allreduce(self, input, bias, layer_idx: Optional[int] | None = None): + tp_group = self.mapping.tp_group + tp_rank = self.tp_rank + if NCCL_FAULT_TOLERANCE_ENABLED and not mpi_disabled(): + raise RuntimeError( + "NCCL error: fused GEMM allreduce has no peer-failure abort " + "path and is unavailable while TLLM_FAULT_TOLERANCE_MODE=1") output = self.quant_method.apply_linear_allreduce( - self, input, bias, self.tp_rank, self.mapping.tp_group) + self, input, bias, tp_rank, tp_group) return output def _maybe_fuse_bias_into_allreduce( @@ -3397,12 +3416,19 @@ def forward( lora_params: Optional[dict] = None, layer_idx: Optional[int] = None, ) -> torch.Tensor: + if (NCCL_FAULT_TOLERANCE_ENABLED and self.tp_mode + in (TensorParallelMode.ROW, TensorParallelMode.COLUMN) + and self.tp_size > 1 and not mpi_disabled()): + assert_nccl_group_not_reconfigured( + self._tp_group_tuple, + f"{self.tp_mode.name.lower()}-parallel linear") + if self.tp_mode == TensorParallelMode.ROW: use_fused_gemm_allreduce = self.use_fused_gemm_allreduce and lora_params is None if use_fused_gemm_allreduce and all_reduce_params is not None: use_fused_gemm_allreduce = all_reduce_params.enable_allreduce and all_reduce_params.fusion_op == AllReduceFusionOp.NONE - bias = None if (self.tp_rank > 0) else self.bias + bias = None if self.tp_rank > 0 else self.bias if self.reduce_output: if use_fused_gemm_allreduce: output = self.apply_linear_allreduce( diff --git a/tensorrt_llm/_torch/pyexecutor/ep_failure_broadcast.py b/tensorrt_llm/_torch/pyexecutor/ep_failure_broadcast.py new file mode 100644 index 000000000000..69821ae701b9 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/ep_failure_broadcast.py @@ -0,0 +1,1715 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""Out-of-band MPI failure-detection propagation for WideEP fault tolerance. + +The model-forward thread can be stuck in an AlltoAll kernel when a peer fails, +so it cannot participate in recovery consensus. :class:`MpiFtSubcomm` owns a +dedicated EP-scoped MPI communicator and progresses only nonblocking +point-to-point requests on an independent CPU thread. + +The component propagates one rank-failure detection and updates only the +process-local :class:`DetectedRankState`. It never receives or updates the +committed ``EPGroupHealth`` execution mask. Multi-rank suspect/confirm +consensus, commit authorization, and communicator reconstruction are +intentionally left to later WideEP FT phases. + +Every survivor echoes a newly observed failure. Receiving one echo from every +active survivor proves only that they propagated the same detection. Detection +reconciliation does **not** authorize request resume, EPLB changes, NCCL +reconfiguration or use, or committed-mask mutation. The future 1c.4b recovery +coordinator consumes the detection callback and owns that authorization. + +ULFM revoke is reserved for terminal control-plane aborts because MPI does not +carry enough information to distinguish a successful commit revoke from an +emergency abort at remote ranks. Without ULFM, terminal state is echoed to every +survivor; failure to confirm that echo within a bounded interval uses +``MPI_Abort`` on the world-spanning FT communicator as a fail-stop fallback. +""" + +from __future__ import annotations + +import math +import queue +import threading +import time +from collections.abc import Callable +from dataclasses import dataclass +from enum import Enum, IntEnum, auto +from typing import TYPE_CHECKING, Protocol, cast + +import numpy as np + +try: + from mpi4py import MPI +except ImportError: + MPI = None + +from tensorrt_llm.logger import logger + +if TYPE_CHECKING: + from tensorrt_llm.mapping import Mapping + + +_FAILURE_TAG = 0 +_MESSAGE_WORDS = 2 + + +class _MessageKind(IntEnum): + FAILURE = 1 + ABORT = 2 + + +_PROCESS_LIFETIME_REFS: list[object] = [] +_PROCESS_LIFETIME_REFS_LOCK = threading.Lock() + + +def _retain_for_process_lifetime(*references: object) -> None: + """Keep active MPI objects reachable until interpreter teardown.""" + with _PROCESS_LIFETIME_REFS_LOCK: + _PROCESS_LIFETIME_REFS.extend(references) + + +class _MpiRequest(Protocol): + """Subset of mpi4py request operations used by the progress loop.""" + + def Cancel(self) -> None: ... + + def Test(self) -> bool: ... + + +class _MpiComm(Protocol): + """Subset of mpi4py communicator operations used by this component.""" + + def Get_rank(self) -> int: ... + + def Get_size(self) -> int: ... + + def Irecv(self, buffer: np.ndarray, source: int, tag: int) -> _MpiRequest: ... + + def Isend(self, buffer: np.ndarray, dest: int, tag: int) -> _MpiRequest: ... + + def Is_revoked(self) -> bool: ... + + def Revoke(self) -> None: ... + + def Abort(self, errorcode: int) -> None: ... + + def Set_errhandler(self, errhandler: object) -> None: ... + + +class _MpiModule(Protocol): + """Subset of the mpi4py ``MPI`` module used by this component.""" + + ERRORS_RETURN: object + ERR_PROC_FAILED: int + ERR_UNKNOWN: int + THREAD_MULTIPLE: int + Exception: type[BaseException] + + def Query_thread(self) -> int: ... + + +FailureDetectedCallback = Callable[[int, int, float], None] +# Compatibility name for callers of the original draft API. New integrations +# should use ``FailureDetectedCallback`` because this callback carries detected, +# never committed, state. +FailureReceivedCallback = FailureDetectedCallback + + +@dataclass(frozen=True) +class DetectedRankStateSnapshot: + """Atomic snapshot of rank-failure evidence for one communicator epoch.""" + + mask: int + failed_ranks: frozenset[int] + generation: int + + +class DetectedRankState: + """Thread-safe, process-local rank-failure evidence. + + This state is deliberately distinct from ``EPGroupHealth``. It is an + observation owned by :class:`MpiFtSubcomm`, not a committed execution mask, + and it is monotonic for the lifetime of one FT communicator epoch. + + Args: + ep_size: Number of EP-local ranks represented by this state. + + Raises: + ValueError: If ``ep_size`` is not positive. + """ + + def __init__(self, ep_size: int) -> None: + if ep_size <= 0: + raise ValueError(f"ep_size must be > 0, got {ep_size}") + self._ep_size = ep_size + self._all_active_mask = (1 << ep_size) - 1 + self._active_mask = self._all_active_mask + self._failed_ranks: set[int] = set() + self._generation = 0 + self._lock = threading.Lock() + + @property + def ep_size(self) -> int: + """Number of ranks represented by this immutable communicator epoch.""" + return self._ep_size + + @property + def generation(self) -> int: + """Number of effective detected-failure transitions.""" + with self._lock: + return self._generation + + def record_failure(self, rank: int) -> bool: + """Record one failed rank, returning whether evidence changed.""" + self._validate_rank(rank) + bit = 1 << rank + with self._lock: + if not self._active_mask & bit: + return False + self._active_mask &= ~bit + self._failed_ranks.add(rank) + self._generation += 1 + return True + + def is_active(self, rank: int) -> bool: + """Return whether no failure has been recorded for ``rank``.""" + self._validate_rank(rank) + with self._lock: + return bool(self._active_mask & (1 << rank)) + + def get_mask(self) -> int: + """Return the active-rank mask derived from detected evidence.""" + with self._lock: + return self._active_mask + + def get_failed_ranks(self) -> frozenset[int]: + """Return an immutable snapshot of ranks with failure evidence.""" + with self._lock: + return frozenset(self._failed_ranks) + + def all_active(self) -> bool: + """Return whether no rank failure has been detected.""" + with self._lock: + return self._active_mask == self._all_active_mask + + def snapshot(self) -> DetectedRankStateSnapshot: + """Return one coherent mask, failure-set, and generation snapshot.""" + with self._lock: + return DetectedRankStateSnapshot( + mask=self._active_mask, + failed_ranks=frozenset(self._failed_ranks), + generation=self._generation, + ) + + def _validate_rank(self, rank: int) -> None: + if not 0 <= rank < self._ep_size: + raise ValueError(f"rank must be in [0, {self._ep_size}), got {rank}") + + +@dataclass(frozen=True) +class MpiFtSubcommConfig: + """Runtime-independent tuning for the FT progress thread. + + Args: + poll_interval_sec: Maximum delay between MPI progress passes. The + default leaves margin below the 100 ms agreement budget. + startup_timeout_sec: Maximum time to wait for receive requests to be + posted when :meth:`MpiFtSubcomm.start` is called. + stop_timeout_sec: Default bounded join timeout used by + :meth:`MpiFtSubcomm.stop`. + reconcile_timeout_sec: Maximum time for every active survivor to + announce the same failure before the control plane fails closed. + unattributed_error_timeout_sec: Maximum time to wait for the host + watchdog to identify a rank after non-ULFM MPI reports a generic + transport error. The default leaves margin for the design's + five-second watchdog bound while keeping terminal escalation + inside the ten-second recovery budget. + abort_timeout_sec: Maximum time to reconcile a relayed terminal ABORT + before falling back to communicator-wide ``MPI_Abort`` when ULFM + is unavailable. + """ + + poll_interval_sec: float = 0.01 + startup_timeout_sec: float = 2.0 + stop_timeout_sec: float = 2.0 + reconcile_timeout_sec: float = 1.0 + unattributed_error_timeout_sec: float = 6.0 + abort_timeout_sec: float = 0.1 + + def __post_init__(self) -> None: + for name, value in ( + ("poll_interval_sec", self.poll_interval_sec), + ("startup_timeout_sec", self.startup_timeout_sec), + ("stop_timeout_sec", self.stop_timeout_sec), + ("reconcile_timeout_sec", self.reconcile_timeout_sec), + ("unattributed_error_timeout_sec", self.unattributed_error_timeout_sec), + ("abort_timeout_sec", self.abort_timeout_sec), + ): + if not math.isfinite(value) or value <= 0: + raise ValueError(f"{name} must be finite and > 0, got {value}") + + +@dataclass +class _PendingSend: + request: _MpiRequest + buffer: np.ndarray + destination: int + failed_rank: int + kind: _MessageKind + + +@dataclass +class _PendingReceive: + request: _MpiRequest + buffer: np.ndarray + source: int + + +@dataclass(frozen=True) +class _OutboundMessage: + kind: _MessageKind + failed_rank: int + + +class _Lifecycle(Enum): + CREATED = auto() + STARTING = auto() + RUNNING = auto() + STOPPING = auto() + STOPPED = auto() + FAILED = auto() + + +class MpiFtSubcomm: + """Broadcast EP-rank failure detections on a dedicated MPI control plane. + + Construction creates the FT communicator collectively on the caller's + thread. Call :meth:`start` only after every rank has constructed the + component. During normal operation, the progress thread owns MPI calls on + this communicator; :meth:`report_detected_failure` updates only local + detected state and enqueues a report, so it never waits for a peer or MPI + progress. The watchdog must call this method instead of pre-recording evidence. + Because ``MPI.THREAD_MULTIPLE`` is required, a caller whose start/stop + deadline proves that thread is wedged may issue the terminal Revoke/Abort + fallback. + + Args: + mapping: Distributed mapping whose ``moe_ep_group`` defines the FT + communicator and whose ``moe_ep_rank`` defines the local rank. + detected_state: Process-local EP detection evidence updated by sent and + received failure reports. ``EPGroupHealth`` is rejected because it + is the committed execution-mask state owned by the recovery + coordinator. + config: Progress-loop timing configuration. + on_failure_detected: Optional callback invoked as + ``(failed_rank, source_rank, monotonic_time)`` when a local or + received report first changes detected state. This is the handoff + to the future 1c.4b recovery coordinator; it is detection evidence, + not commit authorization. Delivery runs on a dedicated daemon + thread so callback latency cannot block MPI progress. + comm: Pre-created FT communicator. Intended for tests; production + callers should let the component create it from ``mapping``. + mpi_module: MPI module paired with ``comm``. Intended for tests. + + Raises: + RuntimeError: If MPI support or thread support is insufficient, or if + the communicator does not match ``mapping``. + TypeError: If ``detected_state`` is not a :class:`DetectedRankState`. + ValueError: If ``detected_state`` and ``mapping`` describe different + EP groups, or if both callback keyword names are provided. + """ + + def __init__( + self, + mapping: Mapping, + detected_state: DetectedRankState, + config: MpiFtSubcommConfig | None = None, + on_failure_detected: FailureDetectedCallback | None = None, + *, + comm: _MpiComm | None = None, + mpi_module: _MpiModule | None = None, + on_failure_received: FailureReceivedCallback | None = None, + ) -> None: + if not isinstance(detected_state, DetectedRankState): + raise TypeError( + "detected_state must be a DetectedRankState; committed " + "EPGroupHealth cannot be used as failure-detection evidence" + ) + if on_failure_detected is not None and on_failure_received is not None: + raise ValueError( + "Specify only on_failure_detected; on_failure_received is a compatibility alias" + ) + if on_failure_detected is None: + on_failure_detected = on_failure_received + if mpi_module is None: + if MPI is None: + raise RuntimeError("mpi4py is required for WideEP failure propagation") + mpi_module = cast(_MpiModule, MPI) + self._mpi = mpi_module + + self._config = config or MpiFtSubcommConfig() + comm_created_collectively = comm is None + collective_setup = None + if comm_created_collectively: + # Keep the collective startup validation and Split on the + # constructing thread. In particular, do not raise a rank-local + # mapping error before peers have entered the same collective + # validation path. + from tensorrt_llm._torch.distributed.communicator import create_mpi_ft_subcomm + + collective_setup = create_mpi_ft_subcomm( + mapping, + health_size=detected_state.ep_size, + ) + comm = cast(_MpiComm, collective_setup.comm) + + assert comm is not None + self._comm = comm + if comm_created_collectively: + # create_mpi_ft_subcomm already reconciled thread support, mapping, + # health size, communicator rank/size, and ERRORS_RETURN across the + # parent communicator. Repeating any of those MPI operations here + # would reintroduce a rank-local failure after collective startup. + assert collective_setup is not None + self._local_rank = collective_setup.local_rank + self._ep_size = collective_setup.ep_size + else: + if self._mpi.Query_thread() < self._mpi.THREAD_MULTIPLE: + raise RuntimeError( + "WideEP FT requires MPI.THREAD_MULTIPLE because its control-plane " + "thread overlaps other MPI traffic" + ) + + ep_group = tuple(mapping.moe_ep_group) + if len(ep_group) != mapping.moe_ep_size: + raise ValueError( + "mapping.moe_ep_group size must match mapping.moe_ep_size, " + f"got {len(ep_group)} and {mapping.moe_ep_size}" + ) + expected_world = set(range(mapping.world_size)) + if mapping.moe_ep_size != mapping.world_size or set(ep_group) != expected_world: + raise ValueError( + "WideEP FT MVP requires one MoE EP group spanning the full MPI world; " + f"got world_size={mapping.world_size}, moe_ep_group={ep_group}" + ) + if detected_state.ep_size != mapping.moe_ep_size: + raise ValueError( + "DetectedRankState size must match mapping.moe_ep_size, " + f"got {detected_state.ep_size} and {mapping.moe_ep_size}" + ) + self._local_rank = mapping.moe_ep_rank + self._ep_size = mapping.moe_ep_size + # Install the non-fatal handler before inspecting the injected + # communicator so even validation errors are returned to Python. + self._comm.Set_errhandler(self._mpi.ERRORS_RETURN) + actual_rank = self._comm.Get_rank() + actual_size = self._comm.Get_size() + if actual_rank != self._local_rank or actual_size != self._ep_size: + raise RuntimeError( + "WideEP FT communicator does not match the mapping: " + f"rank={actual_rank}, size={actual_size}, " + f"expected rank={self._local_rank}, size={self._ep_size}" + ) + + self._detected_state = detected_state + self._on_failure_detected = on_failure_detected + self._outbound_reports: queue.SimpleQueue[_OutboundMessage] = queue.SimpleQueue() + self._announced_failures: set[int] = set() + self._failure_reporters: dict[int, set[int]] = {} + self._reconcile_deadlines: dict[int, float] = {} + self._unattributed_error_deadlines: dict[int, float] = {} + self._failure_fanout_posted: set[int] = set() + self._failure_fanout_complete: set[int] = set() + self._abort_reporters: set[int] = set() + self._abort_expected_reporters: frozenset[int] = frozenset() + self._abort_deadline: float | None = None + self._abort_payload = -1 + self._abort_announced = False + self._abort_fanout_posted = False + self._abort_fanout_complete = False + self._protocol_lock = threading.Lock() + self._failure_claim_lock = threading.Lock() + self._accepted_failed_rank: int | None = None + self._accepted_failure_generation: int | None = None + self._stop_event = threading.Event() + self._abort_requested = threading.Event() + self._wake_event = threading.Event() + self._ready_event = threading.Event() + self._deadline_monitor_stop_event = threading.Event() + self._deadline_monitor_wake_event = threading.Event() + self._deadline_monitor_ready_event = threading.Event() + self._transport_poisoned = threading.Event() + self._progress_failed = threading.Event() + self._lifecycle = _Lifecycle.CREATED + self._lifecycle_lock = threading.Lock() + self._thread: threading.Thread | None = None + self._deadline_monitor_thread: threading.Thread | None = None + self._last_error: BaseException | None = None + self._error_lock = threading.Lock() + self._retained_requests: list[_PendingSend | _PendingReceive] = [] + self._process_lifetime_retained = comm_created_collectively + if comm_created_collectively: + # The FT communicator is a process-lifetime control-plane + # resource. Letting mpi4py destroy it when one broadcaster is + # garbage-collected would call a collective Free at rank-local + # and nondeterministic times. + _retain_for_process_lifetime(self._comm) + self._shutdown_deadline: float | None = None + self._callback_queue: queue.SimpleQueue[tuple[int, int, float] | None] = queue.SimpleQueue() + self._callback_stop_event = threading.Event() + self._callback_thread: threading.Thread | None = None + self._terminal_action_lock = threading.Lock() + self._revoke_attempted = False + self._world_abort_attempted = False + self._ulfm_available = ( + collective_setup.ulfm_available if collective_setup is not None else self._probe_ulfm() + ) + + @property + def ulfm_available(self) -> bool: + """Whether this MPI build implements the ULFM communicator methods.""" + with self._terminal_action_lock: + return self._ulfm_available + + @property + def last_error(self) -> BaseException | None: + """First terminal control-plane error, if any.""" + with self._error_lock: + return self._last_error + + def start(self) -> None: + """Start the dedicated MPI progress thread. + + This method is not restartable. A second call, including one after + :meth:`stop`, raises ``RuntimeError``. + """ + thread_start_error: BaseException | None = None + with self._lifecycle_lock: + if self._lifecycle is not _Lifecycle.CREATED: + raise RuntimeError( + f"WideEP FT broadcaster cannot start from {self._lifecycle.name.lower()} state" + ) + self._lifecycle = _Lifecycle.STARTING + try: + if self._on_failure_detected is not None: + callback_thread = threading.Thread( + target=self._callback_loop, + name="wide-ep-ft-callback", + daemon=True, + ) + callback_thread.start() + # Store only a successfully started thread. stop() may join + # every stored thread and joining an unstarted Thread raises. + self._callback_thread = callback_thread + deadline_monitor_thread = threading.Thread( + target=self._deadline_monitor_loop, + name="wide-ep-ft-deadline-monitor", + daemon=True, + ) + deadline_monitor_thread.start() + self._deadline_monitor_thread = deadline_monitor_thread + progress_thread = threading.Thread( + target=self._progress_loop, + name="wide-ep-ft-broadcast", + daemon=True, + ) + progress_thread.start() + self._thread = progress_thread + except Exception as error: + # Publish the error before FAILED so concurrent observers never + # see a terminal lifecycle without its cause. + self._record_error(error) + self._progress_failed.set() + self._lifecycle = _Lifecycle.FAILED + thread_start_error = error + + if thread_start_error is not None: + self._request_deadline_monitor_stop() + self._request_callback_stop() + self._retain_requests([]) + # There is no progress thread to relay ABORT, so fail closed on the + # caller thread. This is the only thread that can own the FT comm + # after progress-thread creation itself failed. + self._fail_closed_immediately(thread_start_error) + callback_thread = self._callback_thread + if callback_thread is not None and callback_thread.is_alive(): + callback_thread.join(self._config.stop_timeout_sec) + deadline_monitor_thread = self._deadline_monitor_thread + if deadline_monitor_thread is not None and deadline_monitor_thread.is_alive(): + deadline_monitor_thread.join(self._config.stop_timeout_sec) + raise RuntimeError( + "WideEP FT control-plane thread failed to start" + ) from thread_start_error + + startup_deadline = time.monotonic() + self._config.startup_timeout_sec + progress_ready = self._ready_event.wait(max(0.0, startup_deadline - time.monotonic())) + deadline_monitor_ready = self._deadline_monitor_ready_event.wait( + max(0.0, startup_deadline - time.monotonic()) + ) + if not progress_ready or not deadline_monitor_ready: + timeout_error = TimeoutError( + "WideEP FT control-plane threads did not start before the timeout" + ) + with self._lifecycle_lock: + self._lifecycle = _Lifecycle.FAILED + # Stop callback delivery before the terminal MPI action. Real MPI_Abort + # does not return, while test doubles and Revoke do; either way no + # callback worker should be left waiting for progress-thread cleanup. + self._request_deadline_monitor_stop() + self._request_callback_stop() + # The progress thread may be wedged inside the MPI call whose + # startup deadline just expired, so it cannot execute the terminal + # action itself. MPI.THREAD_MULTIPLE is a construction invariant; + # use the caller thread to issue the terminal communicator action. + self._fail_closed_immediately(timeout_error) + # The progress thread may be stuck inside the first MPI Irecv and + # therefore unable to reach its ``finally`` block. Do not leave the + # independent callback worker alive after start() has failed. + callback_thread = self._callback_thread + if callback_thread is not None and callback_thread is not threading.current_thread(): + callback_thread.join(self._config.stop_timeout_sec) + deadline_monitor_thread = self._deadline_monitor_thread + if ( + deadline_monitor_thread is not None + and deadline_monitor_thread is not threading.current_thread() + ): + deadline_monitor_thread.join(self._config.stop_timeout_sec) + raise timeout_error + with self._lifecycle_lock: + startup_error = self.last_error + lifecycle = self._lifecycle + if startup_error is not None: + self._lifecycle = _Lifecycle.FAILED + elif lifecycle is not _Lifecycle.RUNNING: + raise RuntimeError( + "WideEP FT broadcaster did not reach running state during startup " + f"(state={lifecycle.name.lower()})" + ) + if startup_error is not None: + self._request_deadline_monitor_stop() + deadline_monitor_thread = self._deadline_monitor_thread + if ( + deadline_monitor_thread is not None + and deadline_monitor_thread is not threading.current_thread() + ): + deadline_monitor_thread.join(self._config.stop_timeout_sec) + raise RuntimeError("WideEP FT progress thread failed during startup") from startup_error + + def stop(self, timeout: float | None = None) -> None: + """Stop progress without blocking on MPI requests or freeing the comm. + + Healthy requests are cancelled and polled within the same bounded + timeout. After a peer or transport failure, requests and their NumPy + buffers are retained until process teardown because cancelling, + waiting for, or freeing them can hang inside MPI. + + If a known failure is not yet reconciled, the progress thread keeps + running until agreement or the existing reconciliation deadline drives + the terminal fail-stop path. + + Args: + timeout: Maximum join time. Defaults to + :attr:`MpiFtSubcommConfig.stop_timeout_sec`. + + Raises: + ValueError: If ``timeout`` is not finite and positive. + TimeoutError: If the progress thread does not exit in time. + """ + timeout = self._config.stop_timeout_sec if timeout is None else timeout + if not math.isfinite(timeout) or timeout <= 0: + raise ValueError(f"timeout must be finite and > 0, got {timeout}") + + stop_deadline = time.monotonic() + timeout + with self._lifecycle_lock: + if self._lifecycle is _Lifecycle.CREATED: + self._lifecycle = _Lifecycle.STOPPED + self._request_deadline_monitor_stop() + self._request_callback_stop() + return + callback_thread = self._callback_thread + deadline_monitor_thread = self._deadline_monitor_thread + if ( + self._lifecycle is _Lifecycle.STOPPED + and (callback_thread is None or not callback_thread.is_alive()) + and (deadline_monitor_thread is None or not deadline_monitor_thread.is_alive()) + ): + return + thread = self._thread + if self._lifecycle in (_Lifecycle.STARTING, _Lifecycle.RUNNING): + self._lifecycle = _Lifecycle.STOPPING + join_grace = min(0.01, timeout / 2) + self._shutdown_deadline = stop_deadline - join_grace + self._stop_event.set() + self._wake_event.set() + self._request_callback_stop() + + if thread is not None: + thread.join(max(0.0, stop_deadline - time.monotonic())) + if thread.is_alive(): + timeout_error = TimeoutError( + "WideEP FT progress thread did not stop before the timeout" + ) + with self._lifecycle_lock: + self._lifecycle = _Lifecycle.FAILED + # Do not enqueue fail-stop work onto the thread that this + # timeout proves is not making MPI progress. + self._request_deadline_monitor_stop() + self._fail_closed_immediately(timeout_error) + raise timeout_error + + self._request_deadline_monitor_stop() + if ( + deadline_monitor_thread is not None + and deadline_monitor_thread is not threading.current_thread() + ): + deadline_monitor_thread.join(max(0.0, stop_deadline - time.monotonic())) + if deadline_monitor_thread.is_alive(): + timeout_error = TimeoutError( + "WideEP FT deadline monitor did not stop before the timeout" + ) + self._fail_closed_immediately(timeout_error) + raise timeout_error + + if callback_thread is not None and callback_thread is not threading.current_thread(): + callback_thread.join(max(0.0, stop_deadline - time.monotonic())) + if callback_thread.is_alive(): + timeout_error = TimeoutError( + "WideEP FT callback thread did not stop before the timeout" + ) + # The detection handoff does not own MPI state. A stuck future + # coordinator callback may make this stop call time out, but it + # must not make one rank take the poisoned-world process-exit + # path. + raise timeout_error + + with self._lifecycle_lock: + if self._lifecycle is not _Lifecycle.FAILED: + self._lifecycle = _Lifecycle.STOPPED + + def report_detected_failure(self, failed_rank: int) -> bool: + """Record and asynchronously announce one detected EP-rank failure. + + This is the nonblocking seam invoked by the host watchdog. The + broadcaster owns the detected-state transition; callers must not record + the failure before invoking it. This method performs no MPI + calls and never waits for the progress thread. Repeated reports for the + same rank are coalesced after the first local announcement. + + Args: + failed_rank: EP-local rank in ``[0, ep_size)``. + + Returns: + ``True`` if this call changed local detected state, else ``False``. + """ + self._validate_rank(failed_rank) + with self._lifecycle_lock: + if self._lifecycle is not _Lifecycle.RUNNING or self._progress_failed.is_set(): + raise RuntimeError( + f"WideEP FT broadcaster is not running (state={self._lifecycle.name.lower()})" + ) + start_time = time.monotonic() + try: + self._claim_failure(failed_rank) + except RuntimeError as error: + self._lifecycle = _Lifecycle.FAILED + self._request_terminal_abort(error, failed_rank, self._local_rank) + raise + changed = self._detected_state.record_failure(failed_rank) + enqueued = self._observe_failure(failed_rank, self._local_rank) + detected_time = time.monotonic() + logger.warning( + f"WideEP FT local failure rank={failed_rank} changed={changed} " + f"enqueued={enqueued} elapsed_ms={(detected_time - start_time) * 1000.0:.3f}" + ) + if changed: + self._invoke_detection_callback(failed_rank, self._local_rank, detected_time) + return changed + + def pre_failover(self, failed_rank: int) -> bool: + """Compatibility alias for :meth:`report_detected_failure`. + + Despite the historical name, this method does not commit failover and + callers must not pre-record detected evidence. + """ + return self.report_detected_failure(failed_rank) + + def broadcast_failure(self, failed_rank: int) -> bool: + """Compatibility alias for :meth:`report_detected_failure`.""" + return self.report_detected_failure(failed_rank) + + def world_is_poisoned(self) -> bool: + """Return whether this communicator epoch has ever observed failure.""" + return ( + self._transport_poisoned.is_set() + or self._accepted_failure_rank() is not None + or not self._detected_state.all_active() + ) + + def failure_detection_is_reconciled(self, failed_rank: int) -> bool: + """Return whether every active survivor announced the same detection. + + This is detection-plane evidence only. A ``True`` result does not + authorize request resume, EPLB changes, NCCL reconfiguration or use, or + committed-mask mutation. The future 1c.4b recovery coordinator owns + those decisions. This method performs no MPI calls, and a terminal + progress error always fails closed. + """ + self._validate_rank(failed_rank) + with self._lifecycle_lock: + if self._lifecycle is not _Lifecycle.RUNNING: + return False + if self.last_error is not None or self._progress_failed.is_set(): + return False + snapshot = self._detected_state.snapshot() + if failed_rank not in snapshot.failed_ranks: + return False + accepted_failure, accepted_generation = self._accepted_failure_state() + with self._protocol_lock: + if not self._detection_reconciliation_state_is_valid_locked( + snapshot, accepted_failure, accepted_generation + ): + return False + return self._failure_detection_is_reconciled_locked(failed_rank, snapshot.mask) + + def detected_state_is_reconciled(self) -> bool: + """Return whether survivors propagated every local failure detection. + + This is not a commit gate. A ``True`` result does not authorize request + resume, EPLB changes, NCCL reconfiguration or use, or committed-mask + mutation. The future 1c.4b recovery coordinator owns those decisions. + """ + with self._lifecycle_lock: + if self._lifecycle is not _Lifecycle.RUNNING: + return False + if self.last_error is not None or self._progress_failed.is_set(): + return False + snapshot = self._detected_state.snapshot() + accepted_failure, accepted_generation = self._accepted_failure_state() + if not snapshot.failed_ranks: + return accepted_failure is None and not self._transport_poisoned.is_set() + with self._protocol_lock: + if not self._detection_reconciliation_state_is_valid_locked( + snapshot, accepted_failure, accepted_generation + ): + return False + return all( + self._failure_detection_is_reconciled_locked(failed_rank, snapshot.mask) + for failed_rank in snapshot.failed_ranks + ) + + def detected_health_is_reconciled(self) -> bool: + """Compatibility alias for :meth:`detected_state_is_reconciled`.""" + return self.detected_state_is_reconciled() + + def failure_is_reconciled(self, failed_rank: int) -> bool: + """Compatibility alias for :meth:`failure_detection_is_reconciled`.""" + return self.failure_detection_is_reconciled(failed_rank) + + def health_is_reconciled(self) -> bool: + """Compatibility alias for :meth:`detected_state_is_reconciled`.""" + return self.detected_state_is_reconciled() + + def _validate_rank(self, rank: int) -> None: + if not 0 <= rank < self._ep_size: + raise ValueError(f"failed_rank must be in [0, {self._ep_size}), got {rank}") + + def _claim_failure(self, failed_rank: int) -> None: + """Atomically enforce the single-distinct-failure MVP contract.""" + with self._failure_claim_lock: + snapshot = self._detected_state.snapshot() + known_failures = snapshot.failed_ranks + conflicting_evidence = known_failures - {failed_rank} + accepted = self._accepted_failed_rank + if (accepted is not None and accepted != failed_rank) or conflicting_evidence: + first_failure = accepted + if first_failure is None: + first_failure = min(conflicting_evidence) + raise RuntimeError( + "WideEP FT MVP supports exactly one distinct failed rank; " + f"already accepted rank {first_failure}, received rank {failed_rank}" + ) + if accepted is None: + if failed_rank in known_failures: + raise RuntimeError( + "WideEP FT detected state was mutated before the failure " + "broadcaster accepted the report; the watchdog must call " + "report_detected_failure() without pre-recording evidence" + ) + self._accepted_failed_rank = failed_rank + # Exactly one effective record_failure() must be the next + # detected-state transition. Capturing that generation here + # prevents an independent mutation in the claim-to-observe + # window from becoming the accepted epoch baseline. + self._accepted_failure_generation = snapshot.generation + 1 + + def _accepted_failure_rank(self) -> int | None: + with self._failure_claim_lock: + return self._accepted_failed_rank + + def _accepted_failure_state(self) -> tuple[int | None, int | None]: + with self._failure_claim_lock: + return self._accepted_failed_rank, self._accepted_failure_generation + + def _observe_failure(self, failed_rank: int, reporter: int) -> bool: + """Record one detection report and enqueue this rank's echo once.""" + now = time.monotonic() + snapshot = self._detected_state.snapshot() + message: _OutboundMessage | None = None + with self._protocol_lock: + unattributed_deadline = self._unattributed_error_deadlines.get(failed_rank) + if unattributed_deadline != math.inf: + self._unattributed_error_deadlines.pop(failed_rank, None) + reporters = self._failure_reporters.setdefault(failed_rank, set()) + reporters.update((reporter, self._local_rank)) + self._reconcile_deadlines.setdefault( + failed_rank, now + self._config.reconcile_timeout_sec + ) + failure_enqueued = failed_rank not in self._announced_failures + if failure_enqueued: + self._announced_failures.add(failed_rank) + message = _OutboundMessage(_MessageKind.FAILURE, failed_rank) + if self._failure_detection_is_reconciled_locked(failed_rank, snapshot.mask): + self._reconcile_deadlines.pop(failed_rank, None) + if message is not None: + self._outbound_reports.put(message) + self._wake_event.set() + self._deadline_monitor_wake_event.set() + return failure_enqueued + + def _detection_reconciliation_state_is_valid_locked( + self, + snapshot: DetectedRankStateSnapshot, + accepted_failure: int | None, + accepted_generation: int | None, + ) -> bool: + """Reject unresolved transport errors and same-epoch rank restoration.""" + if self._unattributed_error_deadlines: + return False + if accepted_failure is None: + return not snapshot.failed_ranks and accepted_generation is None + return ( + snapshot.failed_ranks == frozenset({accepted_failure}) + and accepted_generation == snapshot.generation + ) + + def _request_terminal_abort( + self, + error: BaseException, + payload: int, + reporter: int, + ) -> None: + """Relay terminal state before the progress thread fails closed. + + The relay is best effort because it uses the same FT communicator. If + survivor echoes do not converge before the bounded deadline, the + progress thread escalates to ULFM revoke or communicator-wide + ``MPI_Abort``. This method performs no MPI calls. + """ + self._record_error(error) + snapshot = self._detected_state.snapshot() + enqueue_abort = False + now = time.monotonic() + with self._protocol_lock: + if self._abort_deadline is None: + self._abort_payload = payload + self._abort_expected_reporters = frozenset( + rank for rank in range(self._ep_size) if snapshot.mask & (1 << rank) + ) + self._abort_deadline = now + self._config.abort_timeout_sec + self._abort_reporters.update((reporter, self._local_rank)) + if not self._abort_announced: + self._abort_announced = True + enqueue_abort = True + self._abort_requested.set() + if enqueue_abort: + self._outbound_reports.put(_OutboundMessage(_MessageKind.ABORT, self._abort_payload)) + self._wake_event.set() + self._deadline_monitor_wake_event.set() + + def _failure_detection_is_reconciled_locked(self, failed_rank: int, active_mask: int) -> bool: + required_reporters = {rank for rank in range(self._ep_size) if active_mask & (1 << rank)} + if not required_reporters: + # There is no survivor that can participate in later recovery. + # Do not let the usual ``empty set is a subset`` rule claim + # detection propagation completed in this terminal case. + return False + reporters = self._failure_reporters.get(failed_rank, set()) + return failed_rank in self._failure_fanout_complete and required_reporters.issubset( + reporters + ) + + def _probe_ulfm(self) -> bool: + if not callable(getattr(self._comm, "Is_revoked", None)) or not callable( + getattr(self._comm, "Revoke", None) + ): + return False + try: + already_revoked = self._comm.Is_revoked() + except Exception as error: + if isinstance(error, NotImplementedError) or self._is_unsupported_ulfm_error(error): + logger.debug(f"WideEP FT ULFM is unavailable: {error}") + return False + self._record_error(error) + self._retain_requests([]) + raise + if already_revoked: + error = RuntimeError("WideEP FT communicator is already revoked at construction") + self._record_error(error) + self._retain_requests([]) + raise error + return True + + def _is_unsupported_ulfm_error(self, error: BaseException) -> bool: + if not isinstance(error, self._mpi.Exception): + return False + get_error_class = getattr(error, "Get_error_class", None) + if not callable(get_error_class): + return False + unsupported_classes = { + value + for value in ( + getattr(self._mpi, "ERR_NOT_SUPPORTED", None), + getattr(self._mpi, "ERR_UNSUPPORTED_OPERATION", None), + ) + if value is not None + } + try: + error_class = get_error_class() + except Exception: + # Error inspection is advisory. Preserve the original probe error + # as the construction failure when the runtime cannot classify it. + return False + return error_class in unsupported_classes + + def _progress_loop(self) -> None: + pending_sends: list[_PendingSend] = [] + pending_receives: dict[int, _PendingReceive] = {} + try: + for peer in range(self._ep_size): + if peer != self._local_rank: + pending_receives[peer] = self._post_receive(peer) + with self._lifecycle_lock: + if self._lifecycle is _Lifecycle.STARTING: + self._lifecycle = _Lifecycle.RUNNING + self._ready_event.set() + + while True: + if self._progress_failed.is_set(): + terminal_error = self.last_error + if terminal_error is None: + terminal_error = RuntimeError("WideEP FT progress was asked to fail closed") + raise terminal_error + self._wake_event.clear() + self._drain_outbound_reports(pending_sends) + self._progress_sends(pending_sends) + self._progress_receives(pending_receives) + self._check_accepted_failure_epoch() + self._check_unattributed_transport_deadlines() + self._check_reconciliation_deadlines() + with self._lifecycle_lock: + terminal_pending = ( + self._abort_requested.is_set() or self._lifecycle is _Lifecycle.FAILED + ) + stop_requested = self._stop_event.is_set() and not terminal_pending + if terminal_pending: + if self._progress_terminal_abort(): + break + self._wake_event.wait(self._config.poll_interval_sec) + continue + if stop_requested: + safe_to_stop, stop_error = self._stop_reconciliation_state() + if stop_error is not None: + with self._lifecycle_lock: + self._lifecycle = _Lifecycle.FAILED + self._request_terminal_abort( + stop_error, + -1, + self._local_rank, + ) + elif safe_to_stop: + break + self._wake_event.wait(self._config.poll_interval_sec) + except Exception as error: + self._fail_closed_immediately(error) + logger.error(f"WideEP FT progress thread failed: {error}") + finally: + self._ready_event.set() + self._request_deadline_monitor_stop() + self._request_callback_stop() + try: + self._cleanup_requests(pending_sends, pending_receives) + except Exception as error: + self._retain_requests([*pending_sends, *pending_receives.values()]) + logger.error(f"WideEP FT request cleanup failed: {error}") + self._fail_closed_after_cleanup(error) + with self._lifecycle_lock: + if self._last_error is not None: + self._lifecycle = _Lifecycle.FAILED + elif self._stop_event.is_set(): + self._lifecycle = _Lifecycle.STOPPED + + def _post_receive(self, source: int) -> _PendingReceive: + buffer = np.empty(_MESSAGE_WORDS, dtype=np.int64) + request = self._comm.Irecv(buffer, source=source, tag=_FAILURE_TAG) + return _PendingReceive(request=request, buffer=buffer, source=source) + + def _drain_outbound_reports(self, pending_sends: list[_PendingSend]) -> None: + while True: + try: + message = self._outbound_reports.get_nowait() + except queue.Empty: + return + + failed_rank = message.failed_rank + snapshot = self._detected_state.snapshot() + if message.kind is _MessageKind.ABORT: + # Terminal state must reach every rank still believed active. + # If one of them is actually the second failed rank, send + # failure falls through to the fail-stop MPI_Abort fallback. + peers = [ + peer + for peer in range(self._ep_size) + if peer != self._local_rank and snapshot.mask & (1 << peer) + ] + else: + peers = [ + peer + for peer in range(self._ep_size) + if peer != self._local_rank + and peer != failed_rank + and snapshot.mask & (1 << peer) + ] + if not peers: + with self._protocol_lock: + if message.kind is _MessageKind.ABORT: + self._abort_fanout_posted = True + self._abort_fanout_complete = True + else: + self._failure_fanout_posted.add(failed_rank) + self._failure_fanout_complete.add(failed_rank) + continue + for peer in peers: + buffer = np.asarray([int(message.kind), failed_rank], dtype=np.int64) + request = self._comm.Isend(buffer, dest=peer, tag=_FAILURE_TAG) + pending_sends.append( + _PendingSend( + request=request, + buffer=buffer, + destination=peer, + failed_rank=failed_rank, + kind=message.kind, + ) + ) + with self._protocol_lock: + if message.kind is _MessageKind.ABORT: + # Posting is not completion: rendezvous sends can still + # require this thread to call Test before peers can receive + # the relay. + self._abort_fanout_posted = True + else: + self._failure_fanout_posted.add(failed_rank) + + def _progress_sends(self, pending_sends: list[_PendingSend]) -> None: + incomplete: list[_PendingSend] = [] + for pending in pending_sends: + completed = pending.request.Test() + if not completed: + incomplete.append(pending) + pending_sends[:] = incomplete + incomplete_failure_ranks = { + pending.failed_rank for pending in incomplete if pending.kind is _MessageKind.FAILURE + } + abort_incomplete = any(pending.kind is _MessageKind.ABORT for pending in incomplete) + with self._protocol_lock: + self._failure_fanout_complete.update( + self._failure_fanout_posted - incomplete_failure_ranks + ) + if self._abort_fanout_posted and not abort_incomplete: + self._abort_fanout_complete = True + + def _progress_receives(self, pending_receives: dict[int, _PendingReceive]) -> None: + for source, pending in list(pending_receives.items()): + try: + completed = pending.request.Test() + except Exception as error: + if not isinstance(error, self._mpi.Exception): + raise + self._transport_poisoned.set() + # A request that raised may still own its receive buffer. Keep + # both alive until process teardown instead of assuming the + # implementation completed it while returning the error. + self._retain_requests([pending]) + if self._is_failed_peer_error(error): + del pending_receives[source] + self._handle_failed_peer(source, source) + continue + if not self.ulfm_available: + # Without ULFM, MPI implementations may surface the dead + # source as a generic error. Do not guess the failed rank; + # drop only this fixed-source receive and keep polling the + # other survivors for a rank-attributed detection report. + del pending_receives[source] + self._track_unattributed_transport_error(source) + logger.warning( + f"WideEP FT receive from EP rank {source} failed without ULFM: {error}" + ) + continue + raise + if not completed: + continue + + del pending_receives[source] + message_kind_value = int(pending.buffer[0]) + failed_rank = int(pending.buffer[1]) + try: + message_kind = _MessageKind(message_kind_value) + except ValueError: + logger.warning( + f"WideEP FT ignored invalid message kind {message_kind_value} " + f"from EP rank {source}" + ) + else: + if message_kind is _MessageKind.ABORT and -1 <= failed_rank < self._ep_size: + self._handle_received_abort(failed_rank, source) + elif message_kind is _MessageKind.FAILURE and 0 <= failed_rank < self._ep_size: + self._handle_received_report(failed_rank, source) + else: + logger.warning( + f"WideEP FT ignored invalid {message_kind.name.lower()} payload " + f"{failed_rank} from EP rank {source}" + ) + + if self._detected_state.is_active(source) and not self._stop_event.is_set(): + try: + pending_receives[source] = self._post_receive(source) + except Exception as error: + if isinstance(error, self._mpi.Exception) and not self.ulfm_available: + self._transport_poisoned.set() + self._track_unattributed_transport_error(source) + logger.warning( + f"WideEP FT could not repost receive from EP rank {source} " + f"without ULFM: {error}" + ) + else: + raise + + def _handle_received_report(self, failed_rank: int, source: int) -> None: + start_time = time.monotonic() + changed = self._accept_received_failure( + failed_rank, + reporter=source, + transport_poisoned=False, + ) + if changed is None: + return + if not changed: + return + received_time = time.monotonic() + logger.warning( + f"WideEP FT received failure rank={failed_rank} source={source} " + f"elapsed_ms={(received_time - start_time) * 1000.0:.3f}" + ) + self._invoke_detection_callback(failed_rank, source, received_time) + + def _handle_failed_peer(self, failed_rank: int, source: int) -> None: + changed = self._accept_received_failure( + failed_rank, + reporter=self._local_rank, + transport_poisoned=True, + ) + if changed is None: + return + if changed: + detected_time = time.monotonic() + logger.warning( + f"WideEP FT observed failed EP rank {failed_rank} through the FT communicator" + ) + self._invoke_detection_callback(failed_rank, source, detected_time) + + def _accept_received_failure( + self, + failed_rank: int, + *, + reporter: int, + transport_poisoned: bool, + ) -> bool | None: + """Atomically claim, mark, and observe one progress-thread failure.""" + claim_error: RuntimeError | None = None + with self._lifecycle_lock: + try: + self._claim_failure(failed_rank) + except RuntimeError as error: + self._lifecycle = _Lifecycle.FAILED + claim_error = error + else: + changed = self._detected_state.record_failure(failed_rank) + if transport_poisoned: + self._transport_poisoned.set() + self._observe_failure(failed_rank, reporter) + return changed + + assert claim_error is not None + self._request_terminal_abort( + claim_error, + failed_rank, + self._local_rank, + ) + return None + + def _handle_received_abort(self, payload: int, source: int) -> None: + error = RuntimeError( + f"WideEP FT received terminal ABORT from EP rank {source} (payload={payload})" + ) + with self._lifecycle_lock: + self._lifecycle = _Lifecycle.FAILED + self._request_terminal_abort(error, payload, source) + + def _invoke_detection_callback(self, failed_rank: int, source: int, event_time: float) -> None: + if self._on_failure_detected is None or self._callback_stop_event.is_set(): + return + self._callback_queue.put((failed_rank, source, event_time)) + + def _callback_loop(self) -> None: + """Deliver detection handoffs without blocking MPI progress.""" + while True: + event = self._callback_queue.get() + if event is None: + return + failed_rank, source, event_time = event + try: + callback = self._on_failure_detected + if callback is not None: + callback(failed_rank, source, event_time) + except Exception as error: + # A coordinator callback must never terminate either + # control-plane thread. + logger.warning(f"WideEP FT failure callback raised: {error}") + + def _request_callback_stop(self) -> None: + if self._callback_stop_event.is_set(): + return + self._callback_stop_event.set() + self._callback_queue.put(None) + + def _request_deadline_monitor_stop(self) -> None: + self._deadline_monitor_stop_event.set() + self._deadline_monitor_wake_event.set() + + def _deadline_monitor_loop(self) -> None: + """Enforce Python-side deadlines even when an MPI call is wedged. + + Normal communicator traffic remains owned by the progress thread. This + monitor only observes protected protocol state. If a terminal relay + itself reaches its deadline, ``MPI.THREAD_MULTIPLE`` permits this thread + to issue the existing one-shot Revoke/Abort fail-stop action. + """ + self._deadline_monitor_ready_event.set() + try: + while True: + # Clear before checking stop/deadline state so a concurrent + # setter cannot be erased between the check and wait below. + self._deadline_monitor_wake_event.clear() + if self._deadline_monitor_stop_event.is_set(): + return + if self._progress_failed.is_set(): + return + + self._check_accepted_failure_epoch() + self._check_unattributed_transport_deadlines() + self._check_reconciliation_deadlines() + + if self._progress_failed.is_set(): + return + now = time.monotonic() + with self._protocol_lock: + abort_deadline = self._abort_deadline + abort_reconciled = ( + self._abort_fanout_complete + and self._abort_expected_reporters.issubset(self._abort_reporters) + ) + if self._abort_requested.is_set() and abort_deadline is not None: + if abort_reconciled: + # The relay reached every expected survivor. The progress + # thread will break cleanly when it resumes. Do not mark + # progress failed here: its exception path would otherwise + # escalate a reconciled non-ULFM terminal relay to + # MPI_Abort. + self._wake_event.set() + return + if now >= abort_deadline: + self._fail_closed_immediately( + TimeoutError( + "WideEP FT terminal relay did not complete before " + "the abort deadline" + ) + ) + self._wake_event.set() + return + + self._deadline_monitor_wake_event.wait(self._config.poll_interval_sec) + except Exception as error: + self._fail_closed_immediately(error) + self._wake_event.set() + finally: + # start() also uses this event to distinguish a scheduled monitor + # from a thread that failed before its loop became observable. + self._deadline_monitor_ready_event.set() + + def _is_failed_peer_error(self, error: BaseException) -> bool: + get_error_class = getattr(error, "Get_error_class", None) + if not callable(get_error_class): + return False + try: + error_class = get_error_class() + except Exception: + # Preserve the original transport failure as the terminal cause (or + # let the non-ULFM watchdog identify the source) when an MPI runtime + # cannot even classify its own exception. + return False + unknown = getattr(self._mpi, "ERR_UNKNOWN", None) + failure_classes = { + value + for value in (getattr(self._mpi, "ERR_PROC_FAILED", None),) + if value is not None and value != unknown + } + return error_class in failure_classes + + def _track_unattributed_transport_error(self, source: int) -> None: + """Bound how long a generic non-ULFM peer error can remain unresolved.""" + # The watchdog report and the MPI error race independently. If this + # source is already the accepted failed rank, the transport error is + # explained and must not arm a stale deadline after _observe_failure() + # has already cleared the error-first ordering. + deadline = time.monotonic() + self._config.unattributed_error_timeout_sec + with self._protocol_lock: + if self._accepted_failure_rank() == source: + return + self._unattributed_error_deadlines.setdefault(source, deadline) + self._deadline_monitor_wake_event.set() + + def _check_accepted_failure_epoch(self) -> None: + """Fail closed when detected state leaves the accepted failure epoch.""" + error: RuntimeError | None = None + accepted_failure: int | None = None + with self._lifecycle_lock: + if self._lifecycle not in (_Lifecycle.RUNNING, _Lifecycle.STOPPING): + return + snapshot = self._detected_state.snapshot() + accepted_failure, accepted_generation = self._accepted_failure_state() + if accepted_failure is None: + if not snapshot.failed_ranks: + return + error = RuntimeError( + "WideEP FT detected state changed outside the failure " + "broadcaster; the watchdog must call " + "report_detected_failure() without pre-recording evidence " + f"(snapshot={snapshot})" + ) + elif ( + snapshot.failed_ranks == frozenset({accepted_failure}) + and snapshot.generation == accepted_generation + ): + return + else: + error = RuntimeError( + "WideEP FT accepted failure epoch changed before communicator " + "reconstruction " + f"(accepted_rank={accepted_failure}, " + f"expected_generation={accepted_generation}, snapshot={snapshot})" + ) + self._lifecycle = _Lifecycle.FAILED + + assert error is not None + payload = accepted_failure if accepted_failure is not None else -1 + self._request_terminal_abort(error, payload, self._local_rank) + + def _check_unattributed_transport_deadlines(self) -> None: + now = time.monotonic() + # Match the public reconciliation gate's lifecycle -> protocol lock + # order. Publish FAILED while both locks exclude a concurrent gate or + # local broadcast, then latch the expired poison before either can run. + with self._lifecycle_lock: + if self._lifecycle not in (_Lifecycle.RUNNING, _Lifecycle.STOPPING): + return + with self._protocol_lock: + expired_sources = [ + source + for source, deadline in self._unattributed_error_deadlines.items() + if now >= deadline + ] + if expired_sources: + self._lifecycle = _Lifecycle.FAILED + # Retaining these few per-source entries for this + # communicator epoch is harmless and permanently keeps the + # protocol gate closed. The infinity latch also prevents + # repeated terminal requests on later progress passes. + for source in expired_sources: + self._unattributed_error_deadlines[source] = math.inf + if not expired_sources: + return + error = TimeoutError( + "WideEP FT could not attribute non-ULFM transport errors before " + f"the watchdog deadline: {expired_sources}" + ) + self._request_terminal_abort(error, -1, self._local_rank) + + def _check_reconciliation_deadlines(self) -> None: + now = time.monotonic() + expired: list[int] = [] + with self._lifecycle_lock: + if self._lifecycle not in (_Lifecycle.RUNNING, _Lifecycle.STOPPING): + return + snapshot = self._detected_state.snapshot() + with self._protocol_lock: + for failed_rank, deadline in list(self._reconcile_deadlines.items()): + if self._failure_detection_is_reconciled_locked(failed_rank, snapshot.mask): + del self._reconcile_deadlines[failed_rank] + elif now >= deadline: + expired.append(failed_rank) + if expired: + self._lifecycle = _Lifecycle.FAILED + if expired: + error = TimeoutError( + f"WideEP FT did not reconcile failed EP ranks before the timeout: {expired}" + ) + self._request_terminal_abort(error, expired[0], self._local_rank) + + def _stop_reconciliation_state(self) -> tuple[bool, BaseException | None]: + """Return whether shutdown can stop MPI progress without splitting ranks.""" + snapshot = self._detected_state.snapshot() + failed_ranks = snapshot.failed_ranks + accepted_failure, accepted_generation = self._accepted_failure_state() + with self._protocol_lock: + reconciliation_state_is_valid = self._detection_reconciliation_state_is_valid_locked( + snapshot, + accepted_failure, + accepted_generation, + ) + if not reconciliation_state_is_valid: + return False, RuntimeError( + "WideEP FT cannot stop after the accepted failure epoch changed " + f"(accepted_rank={accepted_failure}, " + f"expected_generation={accepted_generation}, " + f"snapshot={snapshot})" + ) + if not failed_ranks: + if self._transport_poisoned.is_set(): + return False, RuntimeError( + "WideEP FT cannot stop safely after an unattributed transport error" + ) + return True, None + + with self._protocol_lock: + unreconciled = [ + failed_rank + for failed_rank in failed_ranks + if not self._failure_detection_is_reconciled_locked(failed_rank, snapshot.mask) + ] + untracked = [ + failed_rank + for failed_rank in unreconciled + if failed_rank not in self._reconcile_deadlines + ] + if not unreconciled: + return True, None + if untracked: + return False, RuntimeError( + "WideEP FT cannot stop with locally failed ranks that were never announced: " + f"{untracked}" + ) + # Their existing reconciliation deadlines bound how long stop keeps + # progressing. Expiry enters the normal terminal ABORT path. + return False, None + + def _progress_terminal_abort(self) -> bool: + """Return ``True`` after terminal state is globally safe to stop.""" + if self.ulfm_available: + # Revoke is itself a communicator-wide terminal notification; it + # does not depend on rendezvous Isends reaching completion. + self._progress_failed.set() + with self._lifecycle_lock: + self._lifecycle = _Lifecycle.FAILED + try: + self._try_revoke() + except Exception as error: + logger.warning(f"WideEP FT terminal revoke failed: {error}") + self._try_world_abort() + return True + if self.ulfm_available: + return True + # A runtime unsupported result disables ULFM and falls through to + # the relayed-ABORT protocol. Clear the flag so the next loop pass + # can continue progressing those requests. + self._progress_failed.clear() + + now = time.monotonic() + with self._protocol_lock: + reconciled = self._abort_fanout_complete and self._abort_expected_reporters.issubset( + self._abort_reporters + ) + deadline = self._abort_deadline + expired = deadline is not None and now >= deadline + if not reconciled and not expired: + return False + + self._progress_failed.set() + with self._lifecycle_lock: + self._lifecycle = _Lifecycle.FAILED + # A final echo can race the deadline check. Once every expected + # survivor has acknowledged ABORT, fail-stop escalation is no longer + # needed even if the wall-clock deadline has just elapsed. + if expired and not reconciled: + self._try_world_abort() + return True + + def _fail_closed_immediately(self, error: BaseException) -> None: + """Abort globally after an error makes bounded relay unsafe.""" + with self._lifecycle_lock: + self._record_error(error) + self._progress_failed.set() + self._lifecycle = _Lifecycle.FAILED + if self.ulfm_available: + try: + self._try_revoke() + if self.ulfm_available: + return + except Exception as revoke_error: + logger.warning(f"WideEP FT could not revoke after terminal error: {revoke_error}") + self._try_world_abort() + + def _fail_closed_after_cleanup(self, error: BaseException) -> None: + """Terminate the world when healthy MPI requests cannot be quiesced. + + At this point the progress loop has already stopped, so a relayed ABORT + cannot be made reliable and ULFM revoke alone would not update peers' + process-local shutdown state. Communicator-wide abort is the only + deterministic way to prevent rank-asymmetric MPI_Finalize behavior. + """ + self._record_error(error) + self._progress_failed.set() + with self._lifecycle_lock: + self._lifecycle = _Lifecycle.FAILED + self._try_world_abort() + + def _try_world_abort(self) -> None: + """Last-resort fail-stop when terminal state cannot be reconciled.""" + with self._terminal_action_lock: + if self._world_abort_attempted: + return + # MPI_Abort normally never returns. Mark the attempt before calling + # it so test doubles, broken runtimes, or a later-resuming progress + # thread cannot issue the communicator-wide action twice. + self._world_abort_attempted = True + abort = getattr(self._comm, "Abort", None) + if not callable(abort): + logger.error( + "WideEP FT communicator has no MPI_Abort fallback; " + "survivors must terminate through the poisoned-world shutdown hook" + ) + return + try: + abort(1) + except Exception as error: + # Real MPI_Abort does not return. Test doubles and broken runtimes + # may do so or raise; the local process remains terminal either way. + logger.error(f"WideEP FT communicator-wide abort failed: {error}") + + def _retain_requests(self, requests: list[_PendingSend | _PendingReceive]) -> None: + self._retained_requests.extend(requests) + references: list[object] = list(requests) + if not self._process_lifetime_retained: + self._process_lifetime_retained = True + references.append(self._comm) + _retain_for_process_lifetime(*references) + + def _cleanup_requests( + self, + pending_sends: list[_PendingSend], + pending_receives: dict[int, _PendingReceive], + ) -> None: + requests: list[_PendingSend | _PendingReceive] = [ + *pending_sends, + *pending_receives.values(), + ] + # Never cancel, wait for, or free MPI operations after a peer or the + # communicator has failed. Keeping each request together with its + # NumPy buffer is the only bounded shutdown behavior in that state; + # the later model-engine shutdown hook skips MPI_Finalize entirely. + if self.last_error is not None or self.world_is_poisoned(): + self._retain_requests(requests) + return + + if not requests: + return + + # Healthy shutdown can locally cancel the preposted receives/sends. + # Poll Test rather than calling Wait so stop() remains bounded even on + # an MPI implementation with surprising cancellation behavior. + for pending in requests: + try: + pending.request.Cancel() + except Exception as error: + if not isinstance(error, self._mpi.Exception): + raise + logger.warning(f"WideEP FT could not cancel a healthy MPI request: {error}") + + remaining = requests + deadline = self._shutdown_deadline or time.monotonic() + while remaining and time.monotonic() < deadline: + next_remaining: list[_PendingSend | _PendingReceive] = [] + for pending in remaining: + try: + if not pending.request.Test(): + next_remaining.append(pending) + except Exception as error: + if not isinstance(error, self._mpi.Exception): + raise + next_remaining.append(pending) + remaining = next_remaining + if remaining: + time.sleep(min(0.001, max(0.0, deadline - time.monotonic()))) + + if remaining: + raise TimeoutError( + f"{len(remaining)} WideEP FT MPI requests remained active after " + "the healthy shutdown timeout" + ) + + def _record_error(self, error: BaseException) -> None: + self._transport_poisoned.set() + with self._error_lock: + if self._last_error is None: + self._last_error = error + + def _try_revoke(self) -> None: + # Claim the one permitted Revoke attempt, but never hold a Python lock + # across MPI. A caller-thread timeout and a later-resuming progress + # thread may enter terminal handling concurrently; the claimant owns + # any unsupported/error fallback. + with self._terminal_action_lock: + if not self._ulfm_available or self._revoke_attempted: + return + self._revoke_attempted = True + try: + self._comm.Revoke() + except Exception as error: + if isinstance(error, NotImplementedError) or self._is_unsupported_ulfm_error(error): + with self._terminal_action_lock: + self._ulfm_available = False + logger.warning(f"WideEP FT revoke is unavailable at runtime: {error}") + return + raise diff --git a/tensorrt_llm/_torch/pyexecutor/error_classification.py b/tensorrt_llm/_torch/pyexecutor/error_classification.py index c356e6ca3a62..be10331f6a9e 100644 --- a/tensorrt_llm/_torch/pyexecutor/error_classification.py +++ b/tensorrt_llm/_torch/pyexecutor/error_classification.py @@ -26,9 +26,26 @@ IMMEDIATE_FATAL_PATTERNS: list[str] = [ "cudaerrorillegaladdress", "cudaerrorlaunchfailure", + "cudaerrorcontextisdestroyed", + "cudaerrornodevice", "illegal memory access", "device-side assert", - "unrecoverable", + "cuda context destroyed", + "cuda context is destroyed", + "cuda context was destroyed", + "cuda context is unrecoverable", + "unrecoverable error in the engine", + "cuda error no device", + "gpu has fallen off the bus", + "xid 79", + "nccl communicator abort", + "nccl communicator was aborted", + "ncclcommabort", + "nvshmem peer unreachable", + "mpi rank terminated", + "mpi worker terminated", + "mpi rank exited unexpectedly", + "mpi worker exited unexpectedly", ] # Patterns that are serious but may be transient (e.g. a single OOM @@ -38,6 +55,20 @@ "cuda out of memory", "cuda error", "nccl error", + "nccl timeout", + "nccl operation timed out", + "alltoall timeout", + "all-to-all timeout", + "alltoall watchdog", + "completion_flags timeout", + "deep_ep buffer barrier hang", + "deepep buffer barrier hang", + "intranode::barrier", + "symmetric memory access violation", + "rdma timeout", + "nixl transfer failed", + "nixl transfer entered error state", + "nixl transfer wait timed out", ] diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index c9cf887562c7..ebc8fc46acbd 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + import bisect import contextlib import functools @@ -44,7 +47,7 @@ from ..compilation.backend import Backend from ..compilation.utils import capture_piecewise_cuda_graph from ..distributed import Distributed -from ..distributed.communicator import init_pp_comm +from ..distributed.communicator import init_pp_comm, release_pp_comm from ..expert_statistic import ExpertStatistic from ..memory_buffer_utils import clear_memory_buffers, with_shared_pool from ..metadata import KVCacheParams @@ -108,6 +111,40 @@ def warmup(self, resource_manager: ResourceManager) -> None: return +def _force_eager_mode_for_nccl_fault_tolerance( + llm_args: TorchLlmArgs) -> TorchLlmArgs: + """Disable CUDA capture until communicator-aware recapture lands. + + A captured NCCL graph node retains the communicator handle used during + capture. PR 1a.7 can replace that handle after a rank failure, but PR + 1a.11 owns graph-cache invalidation and recapture. Keep the MVP usable by + running both whole-model and piecewise paths eagerly while FT mode is on. + Work on the engine-local Pydantic copy so callers can reuse their original + arguments for a non-FT engine. + """ + if os.environ.get("TLLM_FAULT_TOLERANCE_MODE") != "1": + return llm_args + + updates = {} + if llm_args.cuda_graph_config is not None: + updates["cuda_graph_config"] = None + + torch_compile_config = llm_args.torch_compile_config + if (torch_compile_config is not None + and torch_compile_config.enable_piecewise_cuda_graph): + updates["torch_compile_config"] = torch_compile_config.model_copy( + update={"enable_piecewise_cuda_graph": False}) + + if not updates: + return llm_args + + logger.warning( + "TLLM_FAULT_TOLERANCE_MODE=1 disables CUDA graph capture until " + "communicator-aware invalidation and recapture are available; this " + "engine will serve in eager mode.") + return llm_args.model_copy(update=updates) + + def _filter_piecewise_capture_num_tokens( candidate_num_tokens: list[int], max_num_tokens: int, @@ -240,8 +277,10 @@ def __init__( model_weights_memory_tag: Optional[str] = None, model_weights_restore_mode=None, ): + llm_args = _force_eager_mode_for_nccl_fault_tolerance(llm_args) self.forward_pass_callable = None self.ub_buffers = None + self._pp_comm_acquired = False if llm_args.encode_only and llm_args.mm_encoder_only: raise ValueError( "encode_only and mm_encoder_only are mutually exclusive.") @@ -269,6 +308,7 @@ def __init__( self.mapping = mapping if mapping.has_pp(): init_pp_comm(mapping) + self._pp_comm_acquired = True self.dist = dist if dist is not None: ExpertStatistic.create(self.dist.rank) @@ -1911,7 +1951,9 @@ def cleanup(self) -> None: 5. Userbuffers (``ub.ub_deallocate`` per buffer); on per-buffer failure the unfreed buffers are kept attached so a deterministic retry doesn't double-free already-released ones, and the - collected errors are re-raised after the loop. + collected errors are re-raised after unrelated global ownership is + released. + 6. This engine's reference to the process-local PP communicator. Idempotency: Subsequent calls are no-ops (guarded by ``_cleanup_done``). @@ -1948,9 +1990,9 @@ def cleanup(self) -> None: self.input_processor_with_hash = None ub_buffers = getattr(self, 'ub_buffers', None) + ub_errors = [] if ub_buffers: remaining_ub_buffers = [] - ub_errors = [] for u in ub_buffers: try: ub.ub_deallocate(u.addr) @@ -1961,13 +2003,21 @@ def cleanup(self) -> None: remaining_ub_buffers.append(u) ub_errors.append(e) self.ub_buffers = remaining_ub_buffers or None - if ub_errors: - raise RuntimeError( - "Failed to deallocate one or more userbuffers during " - "PyTorchModelEngine cleanup") from ub_errors[0] + + if getattr(self, "_pp_comm_acquired", False): + release_pp_comm() + self._pp_comm_acquired = False # Release model weights. release_gc() + if ub_errors: + # Terminal teardown may not get another deterministic retry. Defer + # the userbuffer error until unrelated process-global ownership is + # released, while retaining only the failed buffers for a caller + # that can retry cleanup. + raise RuntimeError( + "Failed to deallocate one or more userbuffers during " + "PyTorchModelEngine cleanup") from ub_errors[0] self._cleanup_done = True def __del__(self) -> None: diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 188ab3302dda..507116f8571f 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -51,6 +51,10 @@ maybe_prefetch_mm_encoder_for_next_iter from ..models.modeling_utils import DecoderModelForCausalLM from ..modules.decoder_layer import DecoderLayer +from ..modules.fused_moe.ep_metrics import ( + get_ep_group_health, is_ep_group_health_registration_pending, + pending_ep_health_metrics, serialize_ep_health_metrics, + validate_ep_health_metrics_topology) from ..speculative.drafter import Drafter from ..speculative.spec_sampler_base import SampleStateTensorsSpec from ..speculative.speculation_gate import SpeculationGate @@ -472,6 +476,8 @@ def __init__( self.resource_manager = resource_manager self.scheduler = scheduler self.model_engine = model_engine + model = getattr(model_engine, "model", None) + self._ep_health_model_config = getattr(model, "model_config", None) self.enable_attention_dp = model_engine.enable_attention_dp self.dist = dist self.sampler = sampler @@ -1193,6 +1199,19 @@ def can_enqueue_requests(self) -> bool: """ return self.executor_request_queue.can_enqueue_request() + def _get_ep_health_stats(self) -> Optional[dict]: + """Return committed EP membership for passive internal telemetry.""" + ep_group_health = get_ep_group_health(self._ep_health_model_config) + if ep_group_health is None: + if is_ep_group_health_registration_pending( + self._ep_health_model_config): + validate_ep_health_metrics_topology( + self._ep_health_model_config) + return pending_ep_health_metrics() + return None + validate_ep_health_metrics_topology(self._ep_health_model_config) + return serialize_ep_health_metrics(ep_group_health) + def get_latest_iteration_stats(self): """ Returns the per-iterations statistics computed since last call to this method. diff --git a/tensorrt_llm/executor/base_worker.py b/tensorrt_llm/executor/base_worker.py index a03fbb4a3c56..16edc6afa1b6 100644 --- a/tensorrt_llm/executor/base_worker.py +++ b/tensorrt_llm/executor/base_worker.py @@ -280,6 +280,17 @@ def fetch_stats(self) -> list: else: return self.engine.get_latest_iteration_stats() + def _get_ep_health_stats(self, timeout: float = 1.0) -> Optional[dict]: + """Return committed EP membership when the backend supports it.""" + health_reader = getattr(self.engine, "_get_ep_health_stats", None) + if health_reader is None: + return None + return health_reader() + + def fetch_ep_health_stats(self) -> Optional[dict]: + """RPC entry point for passive committed-membership telemetry.""" + return self._get_ep_health_stats() + def fetch_kv_cache_events(self) -> list: if isinstance(self.engine, tllm.Executor): return self.engine.get_latest_kv_cache_events() diff --git a/tensorrt_llm/executor/executor.py b/tensorrt_llm/executor/executor.py index fdd3dff44a48..864310834c6f 100644 --- a/tensorrt_llm/executor/executor.py +++ b/tensorrt_llm/executor/executor.py @@ -406,6 +406,10 @@ def get_stats(self, timeout: float) -> List[dict]: self._iter_stats_result.set_timeout(timeout) return self._iter_stats_result.get_results() + def _get_ep_health_stats(self, timeout: float = 1.0) -> Optional[dict]: + """Return passive committed EP membership telemetry when supported.""" + return None + def aget_stats(self, timeout: float) -> IterationResult: """ Get iteration statistics from the runtime. diff --git a/tensorrt_llm/executor/proxy.py b/tensorrt_llm/executor/proxy.py index ffe793bc5b48..c9294e30f631 100644 --- a/tensorrt_llm/executor/proxy.py +++ b/tensorrt_llm/executor/proxy.py @@ -624,6 +624,19 @@ def get_stats(self, timeout: float) -> List[dict]: stats = self.rpc_client.fetch_stats_wait_async(timeout=timeout).remote() return [json.loads(s) if isinstance(s, str) else s for s in stats] + def _get_ep_health_stats(self, timeout: float = 1.0) -> Optional[dict]: + """Fetch best-effort committed EP membership from the rank-0 worker. + + The rank-0 RPC is the only MPI worker control endpoint today. Callers + must treat ``ep_health_available`` as the validity bit if that endpoint + is lost. The snapshot is not detected/suspected physical liveness, and + this passive metrics hook must not drive recovery. Survivor-side + observation belongs to the fault-tolerant control-plane work. + """ + if self.rpc_client is None: + return None + return self.rpc_client.fetch_ep_health_stats().remote(timeout=timeout) + def get_disaggregated_params(self) -> dict: """Get disaggregated params from worker runtime via RPC.""" if self.rpc_client is None: diff --git a/tensorrt_llm/executor/rpc_proxy_mixin.py b/tensorrt_llm/executor/rpc_proxy_mixin.py index ecbb86e25ead..a432ffea1fa5 100644 --- a/tensorrt_llm/executor/rpc_proxy_mixin.py +++ b/tensorrt_llm/executor/rpc_proxy_mixin.py @@ -165,6 +165,10 @@ async def _fetch_responses_loop_async(self): method_name="_fetch_responses_loop_async", ) + def _get_ep_health_stats(self, timeout: float = 1.0) -> Optional[dict]: + """Fetch passive committed EP membership from the rank-0 RPC worker.""" + return self.rpc_client.fetch_ep_health_stats().remote(timeout=timeout) + # NOTE: _fetch_stats_loop_async and _fetch_kv_cache_events_loop_async have been removed. # Stats and kv_events are now fetched on-demand via direct RPC calls # (get_stats, aget_stats, get_kv_events, aget_kv_events) instead of streaming loops. diff --git a/tensorrt_llm/llmapi/llm.py b/tensorrt_llm/llmapi/llm.py index b637c1738dc8..309fd78ca1e3 100644 --- a/tensorrt_llm/llmapi/llm.py +++ b/tensorrt_llm/llmapi/llm.py @@ -1181,6 +1181,10 @@ def get_stats_async(self, timeout: Optional[float] = 2) -> IterationResult: "Use llm.encode() for encoder-only models.") return self._executor.aget_stats(timeout=timeout) + def _get_ep_health_stats(self) -> Optional[dict]: + """Return committed EP membership for the passive metrics hook.""" + return self._executor._get_ep_health_stats() + @set_api_status("beta") def get_kv_cache_events(self, timeout: Optional[float] = 2) -> List[dict]: """Get iteration KV events from the runtime. diff --git a/tensorrt_llm/metrics/collector.py b/tensorrt_llm/metrics/collector.py index cd5a91123ddf..7fb52393bd46 100644 --- a/tensorrt_llm/metrics/collector.py +++ b/tensorrt_llm/metrics/collector.py @@ -14,12 +14,131 @@ # limitations under the License. """Utilities for Prometheus Metrics Collection.""" +import logging import math +import threading import time -from typing import Dict, List, Optional, Union +from collections import deque +from typing import Dict, List, NamedTuple, Optional, Union from .enums import MetricNames +_LOGGER = logging.getLogger(__name__) + +# Phase-1 WideEP uses the two-word uint64 active-rank-mask ABI. Reject a +# malformed payload before it can create unbounded Prometheus cardinality on +# the serving event loop. +_MAX_EP_HEALTH_RANKS = 128 +# Prometheus gauges are IEEE-754 doubles. Keep the monotonic generation exact +# so stale/conflict comparisons have the same meaning before and after export. +_MAX_EP_HEALTH_GENERATION = (1 << 53) - 1 +_MAX_EP_HEALTH_SOURCE_EPOCH_LENGTH = 128 +_MAX_RETIRED_EP_HEALTH_SOURCE_EPOCHS = 32 + + +class _EPHealthSnapshot(NamedTuple): + """One immutable snapshot exported by the local Prometheus collector.""" + + world_size: int + active_count: int + failed_ranks: frozenset[int] + generation: int + + +class _EPHealthPrometheusCollector: + """Export a coherent EP snapshot without multiprocess gauge merging. + + ``prometheus_client`` combines each multiprocess gauge independently. A + rank vector, its aggregate counts, and its validity bit therefore cannot + form an atomic snapshot when more than one process writes the same label + set. EP health is polled by the OpenAI server itself, so keep its latest + immutable value in that process and register this collector directly on + the server's scrape registry instead. + """ + + def __init__(self, labels: Dict[str, str], metric_prefix: str) -> None: + self._label_names = list(labels) + self._label_values = [labels[name] for name in self._label_names] + self._metric_prefix = metric_prefix + self._lock = threading.Lock() + self._snapshot: Optional[_EPHealthSnapshot] = None + self._available = False + + def publish(self, snapshot: _EPHealthSnapshot) -> None: + """Atomically replace the snapshot and mark it available.""" + with self._lock: + self._snapshot = snapshot + self._available = True + + def mark_unavailable(self) -> None: + """Atomically invalidate the retained snapshot.""" + with self._lock: + self._available = False + + def collect(self) -> list: + """Build all metric families from one locked state read.""" + with self._lock: + snapshot = self._snapshot + available = self._available + + # Do not advertise EP health for non-EP or non-FT deployments. Once a + # producer has published a valid snapshot, retain its samples during an + # outage and use the availability family as their validity bit. + if snapshot is None: + return [] + + from prometheus_client.core import GaugeMetricFamily + + metrics = [] + rank_active = GaugeMetricFamily( + self._metric_prefix + "ep_rank_active", + "Whether an Expert Parallel rank is included in coordinator-committed " + "data-plane membership (1) or excluded (0); not physical liveness.", + labels=[*self._label_names, MetricsCollector.labelname_ep_rank], + ) + for rank in range(snapshot.world_size): + rank_active.add_metric( + [*self._label_values, str(rank)], + int(rank not in snapshot.failed_ranks), + ) + metrics.append(rank_active) + + for name, documentation, value in ( + ( + "ep_active_ranks", + "Number of ranks in coordinator-committed EP data-plane membership.", + snapshot.active_count, + ), + ( + "ep_failed_ranks", + "Number of ranks excluded from coordinator-committed EP data-plane " + "membership; not a physical-failure detector.", + snapshot.world_size - snapshot.active_count, + ), + ( + "ep_health_generation", + "Committed EP membership version; increments on each effective transition.", + snapshot.generation, + ), + ): + metric = GaugeMetricFamily( + self._metric_prefix + name, + documentation, + labels=self._label_names, + ) + metric.add_metric(self._label_values, value) + metrics.append(metric) + + availability = GaugeMetricFamily( + self._metric_prefix + "ep_health_available", + "Whether rank-0 committed EP membership telemetry is available. " + "When 0, ignore all other trtllm_ep_* health samples.", + labels=self._label_names, + ) + availability.add_metric(self._label_values, int(available)) + metrics.append(availability) + return metrics + # Adapted from https://github.com/vllm-project/vllm/blob/v0.10.0rc1/vllm/engine/metrics.py#L30 class MetricsCollector: @@ -52,6 +171,13 @@ class MetricsCollector: trtllm_generation_perplexity trtllm_request_error_total + Expert Parallel health metrics: + trtllm_ep_rank_active + trtllm_ep_active_ranks + trtllm_ep_failed_ranks + trtllm_ep_health_generation + trtllm_ep_health_available + Iteration-level metrics: trtllm_kv_cache_hit_rate trtllm_kv_cache_utilization @@ -102,6 +228,7 @@ class MetricsCollector: trtllm_kv_cache_config_info """ labelname_finish_reason = "finished_reason" + labelname_ep_rank = "ep_rank" def __init__( self, @@ -315,6 +442,17 @@ def __init__( documentation="Maximum number of active requests", labelnames=self.labels.keys()) + # Committed Expert Parallel membership is passively observed from the + # rank-0 worker through a dedicated internal RPC. Unlike independent + # multiprocess gauges, this local collector preserves the rank vector, + # aggregate counts, generation, and availability as one snapshot. + self._last_ep_health_state = None + self._last_ep_health_rejection = None + self._retired_ep_health_source_epochs = set() + self._retired_ep_health_source_epoch_order = deque() + self._ep_health_prometheus_collector = _EPHealthPrometheusCollector( + self.labels, self.metric_prefix) + # Iteration latency self.iteration_latency_seconds = Gauge( name=self.metric_prefix + "iteration_latency_seconds", @@ -868,6 +1006,113 @@ def log_iteration_stats(self, iteration_stats: dict) -> None: self._log_counter(self.kv_cache_intra_device_copy_bytes_total, {}, total_intra_device_copy_bytes) + def _reject_ep_health_stats(self, message: str, *args: object) -> bool: + """Mark telemetry unavailable and warn once per repeated rejection.""" + signature = (message, tuple(str(arg) for arg in args)) + if signature != self._last_ep_health_rejection: + _LOGGER.warning(message, *args) + self._last_ep_health_rejection = signature + self.log_ep_health_unavailable() + return False + + def log_ep_health_stats(self, ep_health_stats: dict) -> bool: + """Materialize committed ``EPGroupHealth`` membership as gauges. + + Returns ``False`` when the payload is invalid, stale, or conflicts with + the last accepted snapshot. Callers may retry because a later generation + or a never-before-seen producer epoch can restore a coherent stream. + ``sourceEpoch`` is required so a producer restart cannot be confused + with delayed state from a retired producer. This passive consumer does + not interpret the snapshot as physical liveness or mutate recovery + state. + """ + try: + world_size = ep_health_stats["worldSize"] + active_count = ep_health_stats["activeCount"] + generation = ep_health_stats["generation"] + failed_rank_list = ep_health_stats["failedRanks"] + source_epoch = ep_health_stats["sourceEpoch"] + scalar_values = (world_size, active_count, generation) + if any(type(value) is not int for value in scalar_values): + raise TypeError( + "worldSize, activeCount, and generation must be integers") + if (world_size <= 0 or world_size > _MAX_EP_HEALTH_RANKS + or not 0 <= active_count <= world_size or generation < 0 + or generation > _MAX_EP_HEALTH_GENERATION): + raise ValueError("EP health scalar values are out of range") + if not isinstance(failed_rank_list, list): + raise TypeError("failedRanks must be a list") + if (not isinstance(source_epoch, str) or not source_epoch + or len(source_epoch) > _MAX_EP_HEALTH_SOURCE_EPOCH_LENGTH): + raise TypeError( + "sourceEpoch must be a bounded non-empty string") + if len(failed_rank_list) > world_size: + raise ValueError("failedRanks contains too many ranks") + if any( + type(rank) is not int or not 0 <= rank < world_size + for rank in failed_rank_list): + raise ValueError("failedRanks contains an invalid rank") + failed_ranks = set(failed_rank_list) + if (len(failed_ranks) != len(failed_rank_list) + or active_count != world_size - len(failed_ranks)): + raise ValueError("EP health counts are inconsistent") + except (KeyError, TypeError, ValueError) as error: + return self._reject_ep_health_stats( + "Ignoring invalid epHealthStats payload: %s", error) + + state = (source_epoch, world_size, generation, + tuple(sorted(failed_ranks))) + last_source_epoch = None + if self._last_ep_health_state is not None: + last_source_epoch, last_world_size, last_generation, last_failed_ranks = ( + self._last_ep_health_state) + if world_size != last_world_size: + return self._reject_ep_health_stats( + "Ignoring epHealthStats world-size change from %s to %s", + last_world_size, world_size) + if source_epoch == last_source_epoch and generation < last_generation: + return self._reject_ep_health_stats( + "Ignoring stale epHealthStats generation %s after %s", + generation, last_generation) + if (source_epoch == last_source_epoch + and generation == last_generation + and state[3] != last_failed_ranks): + return self._reject_ep_health_stats( + "Conflicting epHealthStats payloads at generation %s", + generation) + if source_epoch != last_source_epoch: + if source_epoch in self._retired_ep_health_source_epochs: + return self._reject_ep_health_stats( + "Ignoring epHealthStats from retired source epoch %s", + source_epoch) + + if last_source_epoch is not None and source_epoch != last_source_epoch: + if (len(self._retired_ep_health_source_epoch_order) == + _MAX_RETIRED_EP_HEALTH_SOURCE_EPOCHS): + expired_epoch = self._retired_ep_health_source_epoch_order.popleft( + ) + self._retired_ep_health_source_epochs.remove(expired_epoch) + self._retired_ep_health_source_epoch_order.append(last_source_epoch) + self._retired_ep_health_source_epochs.add(last_source_epoch) + self._last_ep_health_state = state + self._last_ep_health_rejection = None + self._ep_health_prometheus_collector.publish( + _EPHealthSnapshot( + world_size=world_size, + active_count=active_count, + failed_ranks=frozenset(failed_ranks), + generation=generation, + )) + return True + + def log_ep_health_unavailable(self) -> None: + """Mark the last EP health telemetry read as unavailable.""" + self._ep_health_prometheus_collector.mark_unavailable() + + def register_ep_health_metrics(self, registry: object) -> None: + """Register coherent EP health metrics on a scrape registry.""" + registry.register(self._ep_health_prometheus_collector) + def log_request_error(self, http_code: Union[int, str] = "") -> None: """Increment the error counter, labeled by HTTP status code.""" labels = {**self.labels, self.labelname_http_code: str(http_code)} diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index 1424e409fc98..84f558239b1f 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -9,6 +9,7 @@ import time import traceback import uuid +from builtins import anext from collections import deque from contextlib import asynccontextmanager from datetime import datetime @@ -18,6 +19,7 @@ Optional, Union) import uvicorn +import zmq from fastapi import Body, FastAPI, Request from fastapi.exceptions import RequestValidationError from fastapi.responses import (FileResponse, JSONResponse, Response, @@ -28,10 +30,13 @@ from tensorrt_llm._tensorrt_engine import LLM from tensorrt_llm._torch.async_llm import AsyncLLM +from tensorrt_llm._torch.modules.fused_moe.ep_metrics import \ + is_pending_ep_health_metrics from tensorrt_llm._utils import EnergyMonitor # yapf: disable from tensorrt_llm.executor import CppExecutorError from tensorrt_llm.executor.postproc_worker import PostprocParams +from tensorrt_llm.executor.rpc import RPCCancelled, RPCError, RPCTimeout from tensorrt_llm.inputs import prompt_inputs from tensorrt_llm.inputs.data import TokensPrompt from tensorrt_llm.inputs.multimodal import MultimodalServerConfig @@ -93,6 +98,51 @@ TIMEOUT_KEEP_ALIVE = 5 # seconds. +def _is_transient_ep_health_error(error: Exception) -> bool: + """Return whether an EP health read can reasonably succeed on retry.""" + if isinstance( + error, + (RPCCancelled, RPCTimeout, TimeoutError, OSError, zmq.ZMQError)): + return True + if isinstance(error, RPCError) and error.cause is not None: + return _is_transient_ep_health_error(error.cause) + return False + + +async def _await_task_cancellation_safe(task: asyncio.Task) -> Any: + """Await ``task`` without forwarding repeated owner cancellation to it. + + ``asyncio.shield`` protects a child only from the cancellation that reaches + that particular shield future. After catching the first cancellation, an + unshielded ``await task`` lets a later ``Task.cancel()`` cancel the child. + That is especially unsafe for ``asyncio.to_thread``: cancelling its asyncio + wrapper does not stop the worker thread. Drain through a separately + shielded task instead, retrying the shield after every owner cancellation, + then preserve the first cancellation request. + """ + try: + return await asyncio.shield(task) + except asyncio.CancelledError as cancellation: + + async def drain_task() -> None: + try: + await task + except (Exception, asyncio.CancelledError): + pass + + drain = asyncio.create_task(drain_task()) + while not drain.done(): + try: + await asyncio.shield(drain) + except asyncio.CancelledError: + continue + # Retrieve an unexpected BaseException from the drain before restoring + # cancellation. Ordinary task errors and child cancellation were + # consumed above to preserve the owner's cancellation semantics. + drain.result() + raise cancellation + + def _build_tool_strict_guided_decoding_params(tools, tool_parser_name): """Build GuidedDecodingParams with structural tags for tools with strict=True. @@ -212,6 +262,9 @@ def __init__( self.perf_metrics = None self.perf_metrics_lock = None self._iteration_stats_collector_task = None + self._ep_health_collector_task = None + self._ep_health_stop_event = asyncio.Event() + self._ep_health_outage_logged = False self._iteration_stats_wakeup_event = asyncio.Event() # Bounded snapshot of iteration stats for the GET /metrics handler. # When the background Prometheus collector loop is active, it is the @@ -292,25 +345,34 @@ async def lifespan(app: FastAPI): logger.info( "Started background iteration stats collector task") - yield + await self._start_ep_health_collector() - # Stop background iteration stats collector - if self._iteration_stats_collector_task is not None: - self._iteration_stats_collector_task.cancel() + try: + yield + finally: try: - await self._iteration_stats_collector_task - except asyncio.CancelledError: - pass - logger.info("Stopped background iteration stats collector task") - - if self.metadata_server is not None: - self.metadata_server.remove(f"trtllm/{self.generator.llm_id}") - logger.info(f"trtllm/{self.generator.llm_id} is unregistered") - if self.disagg_cluster_worker: - await self.disagg_cluster_worker.deregister_worker() - if self.resource_governor is not None: - self.resource_governor.close() - self.generator.shutdown() + await self._stop_ep_health_collector() + finally: + # Stop background iteration stats collector + if self._iteration_stats_collector_task is not None: + self._iteration_stats_collector_task.cancel() + try: + await self._iteration_stats_collector_task + except asyncio.CancelledError: + pass + logger.info( + "Stopped background iteration stats collector task") + + if self.metadata_server is not None: + self.metadata_server.remove( + f"trtllm/{self.generator.llm_id}") + logger.info( + f"trtllm/{self.generator.llm_id} is unregistered") + if self.disagg_cluster_worker: + await self.disagg_cluster_worker.deregister_worker() + if self.resource_governor is not None: + self.resource_governor.close() + self.generator.shutdown() self.app = FastAPI(lifespan=lifespan) @@ -750,6 +812,7 @@ def mount_metrics(self): from prometheus_fastapi_instrumentator import Instrumentator registry = CollectorRegistry() multiprocess.MultiProcessCollector(registry) + self.metrics_collector.register_ep_health_metrics(registry) Instrumentator( should_group_status_codes=False, should_respect_env_var=True, @@ -1170,6 +1233,113 @@ async def _iteration_stats_collector_loop(self): logger.info("Iteration stats collector loop cancelled") raise + def _log_ep_health_outage_once(self, + message: str, + error: Optional[Exception] = None, + terminal: bool = False) -> None: + """Log one retry diagnostic, while always logging a terminal state.""" + if (getattr(self, "_ep_health_outage_logged", False) and not terminal): + return + detail = f": {error}" if error is not None else "" + logger.warning(f"{message}{detail}") + self._ep_health_outage_logged = True + + def _get_ep_health_stop_event(self) -> asyncio.Event: + """Return the lazily initialized health-collector stop event.""" + stop_event = getattr(self, "_ep_health_stop_event", None) + if stop_event is None: + stop_event = asyncio.Event() + self._ep_health_stop_event = stop_event + return stop_event + + async def _wait_for_ep_health_stop(self, timeout: float) -> bool: + """Wait up to ``timeout`` seconds for a collector stop request.""" + stop_event = self._get_ep_health_stop_event() + if stop_event.is_set(): + return True + try: + await asyncio.wait_for(stop_event.wait(), timeout=timeout) + except asyncio.TimeoutError: + return False + return True + + async def _read_ep_health_stats(self) -> Optional[dict]: + """Run one synchronous health read without abandoning its worker thread.""" + read_task = asyncio.create_task( + asyncio.to_thread(self.generator._get_ep_health_stats)) + # asyncio.to_thread cancellation does not stop the underlying thread. + # Drain it before allowing the owner to close the RPC client. + return await _await_task_cancellation_safe(read_task) + + async def _stop_ep_health_collector(self) -> None: + """Cooperatively stop and cancellation-safely drain the collector.""" + task = self._ep_health_collector_task + if task is None: + return + self._get_ep_health_stop_event().set() + try: + await _await_task_cancellation_safe(task) + finally: + self._ep_health_collector_task = None + logger.info("Stopped EP health metrics collector task") + + async def _ep_health_collector_loop(self) -> None: + """Passively poll rank-0's committed EP membership for metrics. + + This loop observes coordinator-committed state, not detected or + suspected physical liveness, and never drives recovery. + """ + try: + while not self._get_ep_health_stop_event().is_set(): + try: + stats = await self._read_ep_health_stats() + if stats is None: + self.metrics_collector.log_ep_health_unavailable() + logger.info( + "EP health telemetry is unsupported; polling disabled" + ) + return + if is_pending_ep_health_metrics(stats): + self.metrics_collector.log_ep_health_unavailable() + self._log_ep_health_outage_once( + "EP health telemetry registration is pending; will retry" + ) + elif not self.metrics_collector.log_ep_health_stats(stats): + self._log_ep_health_outage_once( + "EP health telemetry returned an invalid, stale, or conflicting payload; will retry", + ) + else: + self._ep_health_outage_logged = False + except Exception as e: + self.metrics_collector.log_ep_health_unavailable() + if not _is_transient_ep_health_error(e): + self._log_ep_health_outage_once( + "EP health telemetry failed deterministically; stopping polling", + e, + terminal=True) + return + self._log_ep_health_outage_once( + "Transient error collecting EP health telemetry; will retry", + e) + if await self._wait_for_ep_health_stop(timeout=1.0): + return + except asyncio.CancelledError: + logger.info("EP health metrics collector loop cancelled") + raise + + async def _start_ep_health_collector(self) -> None: + """Start EP health polling without delaying server readiness.""" + if not self.metrics_collector: + return + if self._ep_health_collector_task is not None: + if not self._ep_health_collector_task.done(): + return + self._ep_health_collector_task = None + self._get_ep_health_stop_event().clear() + self._ep_health_collector_task = asyncio.create_task( + self._ep_health_collector_loop()) + logger.info("Started EP health metrics collector task") + async def openai_chat(self, request: ChatCompletionRequest, raw_request: Request) -> Response: diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index 6c603e43a381..b4d7dc3f1dc2 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -34,11 +34,20 @@ l0_a10: - unittest/_torch/executor/test_kv_cache_budget_split.py - unittest/_torch/executor/test_kv_pool_rebalance.py - unittest/_torch/executor/test_disagg_index_mapper_early_release.py + - unittest/_torch/distributed/test_communicator.py + - unittest/_torch/pyexecutor/test_ep_failure_broadcast.py + # MPI prerequisites are required by default. Local environments can opt out + # with TLLM_ALLOW_MPI_FT_SMOKE_SKIP=1. The three-minute outer timeout covers + # four 30-second launcher modes plus bounded cleanup. + - unittest/_torch/pyexecutor/test_ep_failure_broadcast_mpi.py TIMEOUT (3) ISOLATION - unittest/_torch/pyexecutor/test_kv_cache_compression_manager.py + - unittest/_torch/pyexecutor/test_error_classification.py - unittest/_torch/modules/dwdp/test_dwdp_fixup_moe_backends.py - unittest/_torch/modules/dwdp/test_dwdp_manager.py - unittest/_torch/modules/dwdp/test_dwdp_mapping.py - unittest/_torch/modules/dwdp/test_dwdp_peer_ranges.py + - unittest/_torch/modules/moe/test_allgather_reducescatter_ft.py + - unittest/_torch/modules/test_pp_nccl_communicator.py # NOTE: this is a CPU-only test, but we do not have a dedicated job for this (and therefore no # test list either). - unittest/_torch/models/checkpoints/hf/test_weight_loader.py diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index a958f8b24951..cf28b7e684da 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -193,6 +193,8 @@ l0_b200: # ------------- KV Cache Iteration Stats --------------- - unittest/executor/test_stats_serializer.py - unittest/metrics/test_collector.py + - unittest/_torch/modules/test_ep_metrics.py + - unittest/_torch/pyexecutor/test_ep_metrics_stats.py - kv_cache/test_kv_cache_iteration_stats.py::TestKvCacheIterationStats::test_cold_start - kv_cache/test_kv_cache_iteration_stats.py::TestKvCacheIterationStats::test_partial_block_reuse - kv_cache/test_kv_cache_iteration_stats.py::TestKvCacheIterationStats::test_full_block_reuse diff --git a/tests/unittest/_torch/distributed/test_communicator.py b/tests/unittest/_torch/distributed/test_communicator.py new file mode 100644 index 000000000000..0ed029297a00 --- /dev/null +++ b/tests/unittest/_torch/distributed/test_communicator.py @@ -0,0 +1,681 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +from tensorrt_llm._torch.distributed import communicator + + +def _mapping(**overrides: object) -> SimpleNamespace: + values = { + "world_size": 8, + "rank": 5, + "moe_ep_group": list(range(8)), + "moe_ep_size": 8, + "moe_ep_rank": 5, + } + values.update(overrides) + return SimpleNamespace(**values) + + +def _valid_topologies(parent_size: int = 8) -> list[dict[str, object]]: + topologies = [] + for rank in range(parent_size): + ep_group = tuple(range(parent_size)) + topologies.append( + { + "parent_rank": rank, + "parent_size": parent_size, + "world_size": parent_size, + "mapping_rank": rank, + "ep_group": ep_group, + "ep_size": len(ep_group), + "ep_rank": ep_group.index(rank), + "health_size": None, + "error_kind": None, + "error_message": None, + } + ) + return topologies + + +def _valid_setup_statuses(parent_size: int = 8) -> list[dict[str, object]]: + return [ + { + "parent_rank": rank, + "error_message": None, + "ulfm_available": True, + } + for rank in range(parent_size) + ] + + +def _communicators(*, parent_rank: int = 5, parent_size: int = 8) -> tuple[Mock, Mock]: + ft_comm = Mock() + ft_comm.Get_rank.return_value = 5 + ft_comm.Get_size.return_value = 8 + ft_comm.Is_revoked.return_value = False + + parent_comm = Mock() + parent_comm.Get_rank.return_value = parent_rank + parent_comm.Get_size.return_value = parent_size + parent_comm.Split.return_value = ft_comm + + def _allgather(local_record: dict[str, object]) -> list[dict[str, object]]: + if "ep_group" in local_record: + records = _valid_topologies(parent_size) + else: + records = _valid_setup_statuses(parent_size) + records[parent_rank] = local_record + return records + + parent_comm.allgather.side_effect = _allgather + return parent_comm, ft_comm + + +def _fake_mpi(parent_comm: Mock, *, thread_level: int = 3) -> SimpleNamespace: + return SimpleNamespace( + COMM_WORLD=parent_comm, + ERRORS_RETURN=object(), + THREAD_MULTIPLE=3, + Exception=Exception, + Query_thread=Mock(return_value=thread_level), + ) + + +def test_create_mpi_ft_subcomm_splits_ep_group_and_sets_error_handler( + monkeypatch: pytest.MonkeyPatch, +) -> None: + parent_comm, ft_comm = _communicators() + fake_mpi = _fake_mpi(parent_comm) + monkeypatch.setattr(communicator, "MPI", fake_mpi) + active_comm = Mock(return_value=parent_comm) + monkeypatch.setattr(communicator, "mpi_comm", active_comm) + initial_refs = len(communicator._MPI_FT_PROCESS_LIFETIME_REFS) + + try: + result = communicator.create_mpi_ft_subcomm(_mapping()) + + assert result.comm is ft_comm + assert result.local_rank == 5 + assert result.ep_size == 8 + assert result.ulfm_available is True + active_comm.assert_called_once_with() + assert parent_comm.allgather.call_count == 2 + parent_comm.Split.assert_called_once_with(color=0, key=5) + ft_comm.Set_errhandler.assert_called_once_with(fake_mpi.ERRORS_RETURN) + ft_comm.Free.assert_not_called() + assert communicator._MPI_FT_PROCESS_LIFETIME_REFS[initial_refs:] == [ft_comm] + finally: + del communicator._MPI_FT_PROCESS_LIFETIME_REFS[initial_refs:] + + +def test_create_mpi_ft_subcomm_requires_thread_multiple( + monkeypatch: pytest.MonkeyPatch, +) -> None: + parent_comm, _ = _communicators() + fake_mpi = _fake_mpi(parent_comm, thread_level=2) + monkeypatch.setattr(communicator, "MPI", fake_mpi) + + with pytest.raises(RuntimeError, match=r"requires MPI\.THREAD_MULTIPLE"): + communicator.create_mpi_ft_subcomm(_mapping(), parent_comm) + + parent_comm.allgather.assert_called_once() + parent_comm.Split.assert_not_called() + + +@pytest.mark.parametrize( + ("mapping_overrides", "error_match"), + [ + ({"moe_ep_group": []}, "must not be empty"), + ({"moe_ep_group": list(range(7))}, "size must match mapping.moe_ep_size"), + ({"moe_ep_group": [0, 1, 2, 3, 4, 5, 6, 6]}, "contains duplicate ranks"), + ({"moe_ep_group": [0, 1, 2, 3, 4, 5, 6, 8]}, "outside the parent communicator"), + ( + {"moe_ep_group": [0, 1, 2, 3], "moe_ep_size": 4}, + "spanning the full MPI world", + ), + ( + {"moe_ep_group": [5, 0, 1, 2, 3, 4, 6, 7]}, + "group order must match mapping.moe_ep_rank", + ), + ], +) +def test_create_mpi_ft_subcomm_validates_ep_group( + monkeypatch: pytest.MonkeyPatch, + mapping_overrides: dict[str, object], + error_match: str, +) -> None: + parent_comm, _ = _communicators() + monkeypatch.setattr(communicator, "MPI", _fake_mpi(parent_comm)) + + with pytest.raises(ValueError, match=error_match): + communicator.create_mpi_ft_subcomm(_mapping(**mapping_overrides), parent_comm) + + parent_comm.allgather.assert_called_once() + parent_comm.Split.assert_not_called() + + +@pytest.mark.parametrize( + ("parent_rank", "parent_size", "error_match"), + [ + (5, 7, "size must match mapping.world_size"), + (6, 8, "rank must match mapping.rank"), + ], +) +def test_create_mpi_ft_subcomm_validates_parent_mapping( + monkeypatch: pytest.MonkeyPatch, + parent_rank: int, + parent_size: int, + error_match: str, +) -> None: + parent_comm, _ = _communicators(parent_rank=parent_rank, parent_size=parent_size) + monkeypatch.setattr(communicator, "MPI", _fake_mpi(parent_comm)) + + with pytest.raises(RuntimeError, match=error_match): + communicator.create_mpi_ft_subcomm(_mapping(), parent_comm) + + parent_comm.allgather.assert_called_once() + parent_comm.Split.assert_not_called() + + +@pytest.mark.parametrize( + ("mapping_error", "error_match"), + [ + (IndexError("missing EP group"), "IndexError: missing EP group"), + (ValueError("invalid EP group"), "ValueError: invalid EP group"), + ], + ids=["index-error", "value-error"], +) +def test_create_mpi_ft_subcomm_collectively_reports_mapping_accessor_error( + monkeypatch: pytest.MonkeyPatch, + mapping_error: Exception, + error_match: str, +) -> None: + class FailingMapping: + world_size = 8 + rank = 5 + moe_ep_size = 8 + moe_ep_rank = 5 + + @property + def moe_ep_group(self) -> list[int]: + raise mapping_error + + parent_comm, _ = _communicators() + monkeypatch.setattr(communicator, "MPI", _fake_mpi(parent_comm)) + + with pytest.raises(ValueError, match=error_match): + communicator.create_mpi_ft_subcomm(FailingMapping(), parent_comm) + + parent_comm.allgather.assert_called_once() + parent_comm.Split.assert_not_called() + + +def test_create_mpi_ft_subcomm_uses_collectively_validated_ep_size( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class SingleReadMapping: + world_size = 8 + rank = 5 + moe_ep_group = list(range(8)) + + def __init__(self) -> None: + self.ep_size_reads = 0 + self.ep_rank_reads = 0 + + @property + def moe_ep_size(self) -> int: + self.ep_size_reads += 1 + if self.ep_size_reads > 1: + raise ValueError("moe_ep_size was read after collective validation") + return 8 + + @property + def moe_ep_rank(self) -> int: + self.ep_rank_reads += 1 + if self.ep_rank_reads > 1: + raise ValueError("moe_ep_rank was read after collective validation") + return 5 + + parent_comm, ft_comm = _communicators() + fake_mpi = _fake_mpi(parent_comm) + monkeypatch.setattr(communicator, "MPI", fake_mpi) + mapping = SingleReadMapping() + + result = communicator.create_mpi_ft_subcomm(mapping, parent_comm) + + assert result.comm is ft_comm + assert result.local_rank == 5 + assert result.ep_size == 8 + assert result.ulfm_available is True + assert mapping.ep_size_reads == 1 + assert mapping.ep_rank_reads == 1 + assert parent_comm.allgather.call_count == 2 + ft_comm.Set_errhandler.assert_called_once_with(fake_mpi.ERRORS_RETURN) + + +@pytest.mark.parametrize( + ("failure_kind", "error_match"), + [ + ("rank-mismatch", "unexpected WideEP FT communicator"), + ("size-mismatch", "unexpected WideEP FT communicator"), + ("get-rank-error", "RuntimeError: cannot inspect rank"), + ("get-size-error", "RuntimeError: cannot inspect size"), + ("errhandler", "RuntimeError: cannot set error handler"), + ], +) +def test_create_mpi_ft_subcomm_retains_comm_after_post_split_failure( + monkeypatch: pytest.MonkeyPatch, + failure_kind: str, + error_match: str, +) -> None: + parent_comm, ft_comm = _communicators() + monkeypatch.setattr(communicator, "MPI", _fake_mpi(parent_comm)) + initial_refs = len(communicator._MPI_FT_PROCESS_LIFETIME_REFS) + if failure_kind == "rank-mismatch": + ft_comm.Get_rank.return_value = 4 + elif failure_kind == "size-mismatch": + ft_comm.Get_size.return_value = 7 + elif failure_kind == "get-rank-error": + ft_comm.Get_rank.side_effect = RuntimeError("cannot inspect rank") + elif failure_kind == "get-size-error": + ft_comm.Get_size.side_effect = RuntimeError("cannot inspect size") + else: + ft_comm.Set_errhandler.side_effect = RuntimeError("cannot set error handler") + + try: + with pytest.raises(RuntimeError, match=error_match): + communicator.create_mpi_ft_subcomm(_mapping(), parent_comm) + + assert parent_comm.allgather.call_count == 2 + ft_comm.Free.assert_not_called() + assert communicator._MPI_FT_PROCESS_LIFETIME_REFS[initial_refs:] == [ft_comm] + finally: + del communicator._MPI_FT_PROCESS_LIFETIME_REFS[initial_refs:] + + +def test_create_mpi_ft_subcomm_installs_error_handler_before_inspection( + monkeypatch: pytest.MonkeyPatch, +) -> None: + parent_comm, ft_comm = _communicators() + fake_mpi = _fake_mpi(parent_comm) + monkeypatch.setattr(communicator, "MPI", fake_mpi) + calls = [] + ft_comm.Set_errhandler.side_effect = lambda _handler: calls.append("errhandler") + ft_comm.Get_rank.side_effect = lambda: calls.append("rank") or 5 + ft_comm.Get_size.side_effect = lambda: calls.append("size") or 8 + + communicator.create_mpi_ft_subcomm(_mapping(), parent_comm) + + assert calls[:3] == ["errhandler", "rank", "size"] + + +def test_create_mpi_ft_subcomm_collectively_disables_ulfm( + monkeypatch: pytest.MonkeyPatch, +) -> None: + parent_comm, ft_comm = _communicators() + monkeypatch.setattr(communicator, "MPI", _fake_mpi(parent_comm)) + default_allgather = parent_comm.allgather.side_effect + + def _allgather(local_record: dict[str, object]) -> list[dict[str, object]]: + records = default_allgather(local_record) + if "ep_group" not in local_record: + records[2]["ulfm_available"] = False + return records + + parent_comm.allgather.side_effect = _allgather + + result = communicator.create_mpi_ft_subcomm(_mapping(), parent_comm) + + assert result.comm is ft_comm + assert result.ulfm_available is False + + +def test_create_mpi_ft_subcomm_collectively_reports_ulfm_probe_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + parent_comm, ft_comm = _communicators() + monkeypatch.setattr(communicator, "MPI", _fake_mpi(parent_comm)) + ft_comm.Is_revoked.side_effect = RuntimeError("probe failed") + initial_refs = len(communicator._MPI_FT_PROCESS_LIFETIME_REFS) + + try: + with pytest.raises( + RuntimeError, + match="setup failed on parent rank 5.*ULFM probe raised.*probe failed", + ): + communicator.create_mpi_ft_subcomm(_mapping(), parent_comm) + + ft_comm.Free.assert_not_called() + assert communicator._MPI_FT_PROCESS_LIFETIME_REFS[initial_refs:] == [ft_comm] + finally: + del communicator._MPI_FT_PROCESS_LIFETIME_REFS[initial_refs:] + + +def test_create_mpi_ft_subcomm_reconciles_ulfm_error_classifier_failure( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class ClassifierFailureMpiException(Exception): + def Get_error_class(self) -> int: + raise RuntimeError("cannot classify MPI error") + + parent_comm, ft_comm = _communicators() + fake_mpi = _fake_mpi(parent_comm) + fake_mpi.Exception = ClassifierFailureMpiException + monkeypatch.setattr(communicator, "MPI", fake_mpi) + ft_comm.Is_revoked.side_effect = ClassifierFailureMpiException("probe failed") + initial_refs = len(communicator._MPI_FT_PROCESS_LIFETIME_REFS) + + try: + with pytest.raises( + RuntimeError, + match="setup failed on parent rank 5.*ULFM probe raised.*probe failed", + ): + communicator.create_mpi_ft_subcomm(_mapping(), parent_comm) + + assert parent_comm.allgather.call_count == 2 + ft_comm.Free.assert_not_called() + assert communicator._MPI_FT_PROCESS_LIFETIME_REFS[initial_refs:] == [ft_comm] + finally: + del communicator._MPI_FT_PROCESS_LIFETIME_REFS[initial_refs:] + + +def test_create_mpi_ft_subcomm_treats_unsupported_ulfm_probe_as_fallback( + monkeypatch: pytest.MonkeyPatch, +) -> None: + parent_comm, ft_comm = _communicators() + monkeypatch.setattr(communicator, "MPI", _fake_mpi(parent_comm)) + ft_comm.Is_revoked.side_effect = NotImplementedError("ULFM is not compiled") + + result = communicator.create_mpi_ft_subcomm(_mapping(), parent_comm) + + assert result.ulfm_available is False + ft_comm.Free.assert_not_called() + + +def test_create_mpi_ft_subcomm_collectively_rejects_revoked_comm( + monkeypatch: pytest.MonkeyPatch, +) -> None: + parent_comm, ft_comm = _communicators() + monkeypatch.setattr(communicator, "MPI", _fake_mpi(parent_comm)) + ft_comm.Is_revoked.return_value = True + initial_refs = len(communicator._MPI_FT_PROCESS_LIFETIME_REFS) + + try: + with pytest.raises(RuntimeError, match="already revoked at construction"): + communicator.create_mpi_ft_subcomm(_mapping(), parent_comm) + + ft_comm.Free.assert_not_called() + assert communicator._MPI_FT_PROCESS_LIFETIME_REFS[initial_refs:] == [ft_comm] + finally: + del communicator._MPI_FT_PROCESS_LIFETIME_REFS[initial_refs:] + + +def test_create_mpi_ft_subcomm_retains_comm_after_remote_post_split_failure( + monkeypatch: pytest.MonkeyPatch, +) -> None: + parent_comm, ft_comm = _communicators() + monkeypatch.setattr(communicator, "MPI", _fake_mpi(parent_comm)) + default_allgather = parent_comm.allgather.side_effect + initial_refs = len(communicator._MPI_FT_PROCESS_LIFETIME_REFS) + + def _allgather(local_record: dict[str, object]) -> list[dict[str, object]]: + records = default_allgather(local_record) + if "ep_group" not in local_record: + records[2]["error_message"] = "remote communicator setup failed" + return records + + parent_comm.allgather.side_effect = _allgather + + try: + with pytest.raises( + RuntimeError, + match="setup failed on parent rank 2: remote communicator setup failed", + ): + communicator.create_mpi_ft_subcomm(_mapping(), parent_comm) + + ft_comm.Set_errhandler.assert_called_once() + ft_comm.Free.assert_not_called() + assert communicator._MPI_FT_PROCESS_LIFETIME_REFS[initial_refs:] == [ft_comm] + finally: + del communicator._MPI_FT_PROCESS_LIFETIME_REFS[initial_refs:] + + +def test_create_mpi_ft_subcomm_retains_comm_when_post_split_allgather_raises( + monkeypatch: pytest.MonkeyPatch, +) -> None: + parent_comm, ft_comm = _communicators() + monkeypatch.setattr(communicator, "MPI", _fake_mpi(parent_comm)) + default_allgather = parent_comm.allgather.side_effect + allgather_error = RuntimeError("post-split allgather failed") + initial_refs = len(communicator._MPI_FT_PROCESS_LIFETIME_REFS) + + def _allgather(local_record: dict[str, object]) -> list[dict[str, object]]: + if "ep_group" in local_record: + return default_allgather(local_record) + raise allgather_error + + parent_comm.allgather.side_effect = _allgather + + try: + with pytest.raises(RuntimeError, match="post-split allgather failed") as raised: + communicator.create_mpi_ft_subcomm(_mapping(), parent_comm) + + assert raised.value is allgather_error + assert parent_comm.allgather.call_count == 2 + ft_comm.Free.assert_not_called() + assert communicator._MPI_FT_PROCESS_LIFETIME_REFS[initial_refs:] == [ft_comm] + finally: + del communicator._MPI_FT_PROCESS_LIFETIME_REFS[initial_refs:] + + +def test_create_mpi_ft_subcomm_retains_comm_after_malformed_post_split_status( + monkeypatch: pytest.MonkeyPatch, +) -> None: + parent_comm, ft_comm = _communicators() + monkeypatch.setattr(communicator, "MPI", _fake_mpi(parent_comm)) + default_allgather = parent_comm.allgather.side_effect + initial_refs = len(communicator._MPI_FT_PROCESS_LIFETIME_REFS) + + def _allgather(local_record: dict[str, object]) -> list[object]: + records = default_allgather(local_record) + if "ep_group" not in local_record: + records[2] = None + return records + + parent_comm.allgather.side_effect = _allgather + + try: + with pytest.raises( + RuntimeError, + match="invalid status for parent rank 2: None", + ): + communicator.create_mpi_ft_subcomm(_mapping(), parent_comm) + + ft_comm.Free.assert_not_called() + assert communicator._MPI_FT_PROCESS_LIFETIME_REFS[initial_refs:] == [ft_comm] + finally: + del communicator._MPI_FT_PROCESS_LIFETIME_REFS[initial_refs:] + + +def test_create_mpi_ft_subcomm_propagates_remote_validation_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + parent_comm, _ = _communicators() + topologies = _valid_topologies() + topologies[2]["error_kind"] = "value" + topologies[2]["error_message"] = "mapping.moe_ep_group must not be empty" + parent_comm.allgather.side_effect = None + parent_comm.allgather.return_value = topologies + monkeypatch.setattr(communicator, "MPI", _fake_mpi(parent_comm)) + + with pytest.raises( + ValueError, match="startup validation failed on parent rank 2.*must not be empty" + ): + communicator.create_mpi_ft_subcomm(_mapping(), parent_comm) + + parent_comm.Split.assert_not_called() + + +def test_create_mpi_ft_subcomm_collectively_validates_health_size( + monkeypatch: pytest.MonkeyPatch, +) -> None: + parent_comm, _ = _communicators() + monkeypatch.setattr(communicator, "MPI", _fake_mpi(parent_comm)) + + with pytest.raises(ValueError, match="EPGroupHealth size must match"): + communicator.create_mpi_ft_subcomm( + _mapping(), + parent_comm, + health_size=7, + ) + + parent_comm.allgather.assert_called_once() + parent_comm.Split.assert_not_called() + + +def test_create_mpi_ft_subcomm_rejects_inconsistent_remote_ep_group( + monkeypatch: pytest.MonkeyPatch, +) -> None: + parent_comm, _ = _communicators() + topologies = _valid_topologies() + topologies[6]["ep_group"] = (6, 0, 1, 2, 3, 4, 5, 7) + parent_comm.allgather.side_effect = None + parent_comm.allgather.return_value = topologies + monkeypatch.setattr(communicator, "MPI", _fake_mpi(parent_comm)) + + with pytest.raises(RuntimeError, match="startup topology is inconsistent"): + communicator.create_mpi_ft_subcomm(_mapping(), parent_comm) + + parent_comm.Split.assert_not_called() + + +@pytest.mark.parametrize( + ("field", "value", "error_match"), + [ + ("ep_group", None, "invalid EP group"), + ("ep_group", "01234567", "invalid EP group"), + ("ep_group", (0, 1, True, 3, 4, 5, 6, 7), "invalid EP group ranks"), + ("ep_group", (0, 1, 2, 3, 4, 5, 6, 8), "invalid EP group ranks"), + ("ep_group", (0, 1, 2, 3, 4, 5, 6, 6), "duplicate EP group ranks"), + ("ep_size", None, "reports EP size"), + ("parent_rank", 2.0, "reports parent rank"), + ("parent_size", 8.0, "reports communicator size"), + ("world_size", 8.0, "reports world size"), + ("mapping_rank", 2.0, "reports mapping rank"), + ("ep_size", 8.0, "reports EP size"), + ("ep_rank", 2.0, "reports invalid EP rank"), + ("health_size", 8.0, "reports health size"), + ], + ids=[ + "missing-group", + "string-group", + "bool-rank", + "out-of-range-rank", + "duplicate-rank", + "missing-ep-size", + "float-parent-rank", + "float-parent-size", + "float-world-size", + "float-mapping-rank", + "float-ep-size", + "float-ep-rank", + "float-health-size", + ], +) +def test_create_mpi_ft_subcomm_rejects_malformed_remote_topology_schema( + monkeypatch: pytest.MonkeyPatch, + field: str, + value: object, + error_match: str, +) -> None: + parent_comm, _ = _communicators() + topologies = _valid_topologies() + topologies[2][field] = value + parent_comm.allgather.side_effect = None + parent_comm.allgather.return_value = topologies + monkeypatch.setattr(communicator, "MPI", _fake_mpi(parent_comm)) + + with pytest.raises(RuntimeError, match=error_match): + communicator.create_mpi_ft_subcomm(_mapping(), parent_comm) + + parent_comm.Split.assert_not_called() + + +def test_create_mpi_ft_subcomm_rejects_float_local_ep_size_before_split( + monkeypatch: pytest.MonkeyPatch, +) -> None: + parent_comm, _ = _communicators() + monkeypatch.setattr(communicator, "MPI", _fake_mpi(parent_comm)) + + with pytest.raises(RuntimeError, match=r"reports EP size 8\.0"): + communicator.create_mpi_ft_subcomm( + _mapping(moe_ep_size=8.0), + parent_comm, + ) + + parent_comm.allgather.assert_called_once() + parent_comm.Split.assert_not_called() + + +def test_create_mpi_ft_subcomm_uses_pkl5_compatible_object_allgather( + monkeypatch: pytest.MonkeyPatch, +) -> None: + ft_comm = Mock() + ft_comm.Get_rank.return_value = 5 + ft_comm.Get_size.return_value = 8 + ft_comm.Is_revoked.return_value = False + + class Pkl5LikeParent: + def __init__(self) -> None: + self.gathered_records = [] + self.split_args = None + + def Get_rank(self) -> int: + return 5 + + def Get_size(self) -> int: + return 8 + + def allgather(self, local_record: dict[str, object]) -> list[dict[str, object]]: + self.gathered_records.append(local_record) + if "ep_group" in local_record: + records = _valid_topologies() + else: + records = _valid_setup_statuses() + records[5] = local_record + return records + + def Split(self, *, color: int, key: int) -> Mock: + self.split_args = (color, key) + return ft_comm + + parent_comm = Pkl5LikeParent() + fake_mpi = _fake_mpi(Mock()) + monkeypatch.setattr(communicator, "MPI", fake_mpi) + + result = communicator.create_mpi_ft_subcomm(_mapping(), parent_comm) + + assert result.comm is ft_comm + assert result.local_rank == 5 + assert result.ep_size == 8 + assert result.ulfm_available is True + assert len(parent_comm.gathered_records) == 2 + assert parent_comm.split_args == (0, 5) + ft_comm.Set_errhandler.assert_called_once_with(fake_mpi.ERRORS_RETURN) diff --git a/tests/unittest/_torch/modules/moe/test_allgather_reducescatter_ft.py b/tests/unittest/_torch/modules/moe/test_allgather_reducescatter_ft.py new file mode 100644 index 000000000000..8555d542850f --- /dev/null +++ b/tests/unittest/_torch/modules/moe/test_allgather_reducescatter_ft.py @@ -0,0 +1,861 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Unit tests for the NCCL fault-tolerance hooks on AllGatherReduceScatter. + +These tests intentionally mock the C++ custom op. They verify the Python-side +contract (global-rank validation, survivor publication, and dispatch-size +bookkeeping) without requiring NCCL, MPI, or GPUs. +""" + +from __future__ import annotations + +import threading +from dataclasses import dataclass, field +from typing import Callable + +import pytest + +from tensorrt_llm._torch.custom_ops import torch_custom_ops +from tensorrt_llm._torch.distributed import nccl_fault_tolerance +from tensorrt_llm._torch.distributed import ops as distributed_ops +from tensorrt_llm._torch.modules import linear as linear_module +from tensorrt_llm._torch.modules.fused_moe.communication import ( + allgather_reducescatter as allgather_rs_module, +) +from tensorrt_llm._torch.modules.fused_moe.communication.allgather_reducescatter import ( + AllGatherReduceScatter, +) +from tensorrt_llm._torch.modules.linear import Linear, TensorParallelMode + + +@pytest.fixture(autouse=True) +def _enable_and_clear_nccl_survivor_membership(monkeypatch): + monkeypatch.setattr(nccl_fault_tolerance, "NCCL_FAULT_TOLERANCE_ENABLED", True) + monkeypatch.setattr(distributed_ops, "NCCL_FAULT_TOLERANCE_ENABLED", True) + monkeypatch.setattr(torch_custom_ops, "NCCL_FAULT_TOLERANCE_ENABLED", True) + monkeypatch.setattr(allgather_rs_module, "NCCL_FAULT_TOLERANCE_ENABLED", True) + monkeypatch.setattr(linear_module, "NCCL_FAULT_TOLERANCE_ENABLED", True) + nccl_fault_tolerance._reset_nccl_group_registry_for_tests() + yield + nccl_fault_tolerance._reset_nccl_group_registry_for_tests() + + +@dataclass +class _FakeMapping: + """A PP slice whose EP-local ranks map to non-zero global ranks.""" + + world_size: int = 8 + rank: int = 6 + tp_size: int = 4 + tp_rank: int = 2 + moe_ep_size: int = 4 + moe_ep_rank: int = 2 + enable_attention_dp: bool = False + _group: list[int] = field(default_factory=lambda: [4, 5, 6, 7]) + _moe_group: list[int] | None = None + + @property + def tp_group(self) -> list[int]: + return list(self._group) + + @property + def moe_ep_group(self) -> list[int]: + return list(self._group if self._moe_group is None else self._moe_group) + + +def _patch_reinit_op( + monkeypatch: pytest.MonkeyPatch, + implementation: Callable[[list[int], list[int], int], None], +) -> None: + monkeypatch.setattr( + allgather_rs_module.torch.ops.trtllm, + "nccl_comm_abort_and_reinit", + implementation, + raising=False, + ) + + +def _patch_async_error_op( + monkeypatch: pytest.MonkeyPatch, + implementation: Callable[[list[int]], str | None], +) -> None: + monkeypatch.setattr( + allgather_rs_module.torch.ops.trtllm, + "nccl_comm_get_async_error", + implementation, + raising=False, + ) + + +def _assert_active_mapping(mapping, expected_group: list[int], expected_rank: int) -> None: + """Assert the distributed op sees compact ranks for the rebuilt comm.""" + assert list(mapping.tp_group) == expected_group + assert mapping.tp_size == len(expected_group) + assert mapping.tp_rank == expected_rank + + +def test_abort_and_reinit_canonicalizes_global_ranks(monkeypatch): + calls: list[tuple[list[int], list[int], int]] = [] + _patch_reinit_op( + monkeypatch, + lambda group, active_group, rendezvous_id: calls.append( + (list(group), list(active_group), rendezvous_id) + ), + ) + comm = AllGatherReduceScatter(_FakeMapping()) + + # Input order must not affect NCCL rank assignment: the C++ op receives the + # original TP-group order in the global/world-rank domain. + comm.abort_and_reinit([7, 4, 6]) + + assert calls == [([4, 5, 6, 7], [4, 6, 7], 1)] + + +def test_reconfiguration_rejects_static_sharded_allgather(monkeypatch): + _patch_reinit_op(monkeypatch, lambda group, active_group, rendezvous_id: None) + observed = [] + + def fake_allgather(input, group, rank, group_boxed, dim, sizes): + observed.append((input, group, rank, group_boxed, dim, sizes)) + return "gathered" + + monkeypatch.setattr(distributed_ops, "_allgather", fake_allgather) + + mapping = _FakeMapping() + comm = AllGatherReduceScatter(mapping) + comm.abort_and_reinit([4, 6, 7]) + + # Transport recovery alone cannot recreate a TP shard. A normal model-wide + # call using the original Mapping must fail instead of returning fewer + # output features. AllGatherReduceScatter uses its explicit survivor mapping below. + with pytest.raises(RuntimeError, match="statically sharded|redistribute"): + distributed_ops.allgather("input", mapping, dim=0, sizes=[5, 7, 11, 13]) + assert observed == [] + + +def test_reconfiguration_rejects_static_sharded_helix_alltoall(monkeypatch): + _patch_reinit_op(monkeypatch, lambda group, active_group, rendezvous_id: None) + AllGatherReduceScatter(_FakeMapping()).abort_and_reinit([4, 6, 7]) + calls = [] + + def fake_alltoall(inputs, group, num_lists): + calls.append((list(inputs), list(group), num_lists)) + return ["output"] + + monkeypatch.setattr( + distributed_ops.torch.ops.trtllm, + "alltoall_helix", + fake_alltoall, + raising=False, + ) + + inputs = [allgather_rs_module.torch.tensor([rank]) for rank in range(4)] + with pytest.raises(RuntimeError, match="statically sharded|redistribute"): + distributed_ops.alltoall_helix(inputs, [4, 5, 6, 7]) + assert calls == [] + + +def test_reconfiguration_is_shared_and_idempotent_across_moe_layers(monkeypatch): + calls = [] + _patch_reinit_op( + monkeypatch, + lambda group, active_group, rendezvous_id: calls.append((list(group), list(active_group))), + ) + first_layer = AllGatherReduceScatter(_FakeMapping()) + second_layer = AllGatherReduceScatter(_FakeMapping()) + + first_layer.abort_and_reinit([4, 6, 7], generation=1) + second_layer.abort_and_reinit([4, 6, 7], generation=1) + + assert calls == [([4, 5, 6, 7], [4, 6, 7])] + _assert_active_mapping(second_layer._active_mapping, [4, 6, 7], 1) + + +def test_reconfiguration_transaction_serializes_same_target(monkeypatch): + entered_native = threading.Event() + release_native = threading.Event() + calls = [] + + def native_rebuild(group, active_group, rendezvous_id): + calls.append((list(group), list(active_group))) + entered_native.set() + assert release_native.wait(timeout=5) + + _patch_reinit_op(monkeypatch, native_rebuild) + first_layer = AllGatherReduceScatter(_FakeMapping()) + second_layer = AllGatherReduceScatter(_FakeMapping()) + errors = [] + + def reconfigure(layer): + try: + layer.abort_and_reinit([4, 6, 7], generation=1) + except Exception as error: # pragma: no cover - assertion reports details + errors.append(error) + + first = threading.Thread(target=reconfigure, args=(first_layer,)) + second = threading.Thread(target=reconfigure, args=(second_layer,)) + first.start() + assert entered_native.wait(timeout=5) + second.start() + release_native.set() + first.join(timeout=5) + second.join(timeout=5) + + assert not first.is_alive() + assert not second.is_alive() + assert errors == [] + assert calls == [([4, 5, 6, 7], [4, 6, 7])] + + +def test_shared_generation_deduplicates_and_advances_same_membership_rebuild(monkeypatch): + calls = [] + _patch_reinit_op( + monkeypatch, + lambda group, active_group, rendezvous_id: calls.append( + (list(group), list(active_group), rendezvous_id) + ), + ) + + comm = AllGatherReduceScatter(_FakeMapping()) + comm.abort_and_reinit([4, 5, 6, 7], generation=1) + comm.abort_and_reinit([4, 5, 6, 7], generation=1) + comm.abort_and_reinit([4, 5, 6, 7], generation=2) + + assert calls == [ + ([4, 5, 6, 7], [4, 5, 6, 7], 3), + ([4, 5, 6, 7], [4, 5, 6, 7], 4), + ] + + +def test_same_membership_recovery_requires_shared_generation(monkeypatch): + calls = [] + _patch_reinit_op( + monkeypatch, + lambda group, active_group, rendezvous_id: calls.append((list(group), list(active_group))), + ) + + comm = AllGatherReduceScatter(_FakeMapping()) + with pytest.raises(ValueError, match="generation is required"): + comm.abort_and_reinit([4, 5, 6, 7]) + + assert calls == [] + + +@pytest.mark.parametrize("generation", [-1, 1.5, (1 << 63) - 2]) +def test_raw_recovery_generation_must_fit_reserved_torch_int_range(monkeypatch, generation): + calls = [] + _patch_reinit_op( + monkeypatch, + lambda group, active_group, rendezvous_id: calls.append( + (group, active_group, rendezvous_id) + ), + ) + + with pytest.raises(ValueError, match="generation must be a nonnegative integer"): + AllGatherReduceScatter(_FakeMapping()).abort_and_reinit([4, 6, 7], generation=generation) + + assert calls == [] + + +def test_generation_rejects_conflicting_raw_recovery_targets(monkeypatch): + calls = [] + _patch_reinit_op( + monkeypatch, + lambda group, active_group, rendezvous_id: calls.append((list(group), list(active_group))), + ) + + comm = AllGatherReduceScatter(_FakeMapping()) + comm.abort_and_reinit([4, 6, 7], generation=9) + with pytest.raises(RuntimeError, match="conflicting.*generation 9"): + comm.abort_and_reinit([6, 7], generation=9) + + assert calls == [([4, 5, 6, 7], [4, 6, 7])] + + +def test_fault_tolerance_control_path_requires_startup_mode(monkeypatch): + monkeypatch.setattr(allgather_rs_module, "NCCL_FAULT_TOLERANCE_ENABLED", False) + calls = [] + _patch_reinit_op( + monkeypatch, lambda group, active_group, rendezvous_id: calls.append((group, active_group)) + ) + + comm = AllGatherReduceScatter(_FakeMapping()) + with pytest.raises(RuntimeError, match="TLLM_FAULT_TOLERANCE_MODE=1"): + comm.abort_and_reinit([4, 6, 7]) + + assert calls == [] + + +def test_default_off_collective_skips_survivor_registry(monkeypatch): + monkeypatch.setattr(distributed_ops, "NCCL_FAULT_TOLERANCE_ENABLED", False) + monkeypatch.setattr(distributed_ops, "mpi_disabled", lambda: False) + + def unexpected_registry_lookup(group, operation): + raise AssertionError("default-off allgather consulted survivor state") + + observed = [] + monkeypatch.setattr( + distributed_ops, "assert_nccl_group_not_reconfigured", unexpected_registry_lookup + ) + monkeypatch.setattr( + distributed_ops, + "_allgather", + lambda input, group, rank, group_boxed, dim, sizes: observed.append( + (group, rank, group_boxed, dim, sizes) + ) + or "gathered", + ) + + class ReferenceMapping: + tp_group = [4, 5, 6, 7] + tp_rank = 2 + + mapping = ReferenceMapping() + result = distributed_ops.allgather("input", mapping, dim=0, sizes=[1, 2, 3, 4]) + + assert result == "gathered" + assert observed == [([4, 5, 6, 7], 2, None, 0, [1, 2, 3, 4])] + assert observed[0][0] is mapping.tp_group + + +def test_ft_auto_allreduce_uses_only_watchdog_managed_nccl(monkeypatch): + monkeypatch.setattr(distributed_ops, "mpi_disabled", lambda: False) + calls = [] + + def fake_allreduce(**kwargs): + calls.append(kwargs) + return [kwargs["input"]] + + monkeypatch.setattr( + distributed_ops.torch.ops.trtllm, "allreduce", fake_allreduce, raising=False + ) + monkeypatch.setattr( + distributed_ops, + "get_allreduce_workspace", + lambda mapping: (_ for _ in ()).throw( + AssertionError("FT AUTO allocated a custom-allreduce workspace") + ), + ) + monkeypatch.setattr( + distributed_ops.MNNVLAllReduce, + "is_mnnvl", + lambda mapping, dtype: (_ for _ in ()).throw( + AssertionError("FT AUTO probed the MNNVL allreduce backend") + ), + ) + + allreduce = distributed_ops.AllReduce( + _FakeMapping(), + strategy=distributed_ops.AllReduceStrategy.AUTO, + dtype=distributed_ops.torch.float16, + ) + input_tensor = distributed_ops.torch.ones(2) + output = allreduce(input_tensor) + + assert output is input_tensor + assert allreduce.workspace is None + assert allreduce.mnnvl_allreduce is None + assert len(calls) == 1 + assert calls[0]["strategy"] == distributed_ops.AllReduceStrategy.NCCL + + +@pytest.mark.parametrize( + "strategy", + [ + distributed_ops.AllReduceStrategy.UB, + distributed_ops.AllReduceStrategy.ONESHOT, + distributed_ops.AllReduceStrategy.MNNVL, + distributed_ops.AllReduceStrategy.NCCL_SYMMETRIC, + distributed_ops.AllReduceStrategy.SYMM_MEM, + ], +) +def test_ft_allreduce_rejects_custom_transport_strategies(monkeypatch, strategy): + monkeypatch.setattr(distributed_ops, "mpi_disabled", lambda: False) + with pytest.raises(RuntimeError, match="watchdog-managed NCCL allreduce"): + distributed_ops.AllReduce(_FakeMapping(), strategy=strategy) + + +def test_ft_tunable_allreduce_exposes_only_nccl_tactics(): + runner = torch_custom_ops.AllReduceRunner( + 4, + [4, 5, 6, 7], + int(distributed_ops.AllReduceFusionOp.NONE), + 1e-6, + False, + ) + inputs = [distributed_ops.torch.ones(2)] + assert runner.get_valid_tactics(inputs, torch_custom_ops.OptimizationProfile()) == [ + distributed_ops.AllReduceStrategy.NCCL.value, + ] + + +def test_ft_custom_allreduce_helpers_fail_before_workspace_allocation(monkeypatch): + monkeypatch.setattr(distributed_ops, "mpi_disabled", lambda: False) + monkeypatch.setattr( + distributed_ops, + "get_allreduce_workspace", + lambda mapping: (_ for _ in ()).throw( + AssertionError("unsupported FT helper allocated a workspace") + ), + ) + + with pytest.raises(RuntimeError, match="MoEAllReduce.*custom collective"): + distributed_ops.MoEAllReduce(_FakeMapping()) + with pytest.raises(RuntimeError, match="MiniMaxAllReduceRMS.*custom collective"): + distributed_ops.MiniMaxAllReduceRMS(_FakeMapping()) + monkeypatch.setattr( + distributed_ops.torch.ops.trtllm, + "userbuffers_allreduce_finalize", + lambda *args: (_ for _ in ()).throw( + AssertionError("unsupported FT userbuffers op was invoked") + ), + raising=False, + ) + with pytest.raises(RuntimeError, match="userbuffers allreduce"): + distributed_ops.userbuffers_allreduce_finalize(distributed_ops.torch.ones(1)) + + +def test_ft_disables_fused_nvfp4_gemm_allreduce(monkeypatch): + class QuantMode: + @staticmethod + def has_nvfp4(): + return True + + class QuantConfig: + layer_quant_mode = QuantMode() + + monkeypatch.setattr(linear_module, "mpi_disabled", lambda: False) + monkeypatch.setattr(linear_module, "get_sm_version", lambda: 100) + monkeypatch.setenv("TRTLLM_GEMM_ALLREDUCE_FUSION_ENABLED", "1") + nvls_queries = [] + monkeypatch.setattr( + linear_module, "ipc_nvls_supported", lambda: nvls_queries.append(True) or True + ) + + linear = Linear( + 128, + 64, + bias=False, + dtype=distributed_ops.torch.float16, + mapping=_FakeMapping(), + tensor_parallel_mode=TensorParallelMode.ROW, + quant_config=QuantConfig(), + skip_create_weights_in_init=True, + enable_gemm_allreduce_fusion=True, + ) + + assert not linear.use_fused_gemm_allreduce + assert nvls_queries == [] + with pytest.raises(RuntimeError, match="fused GEMM allreduce"): + linear.apply_linear_allreduce(None, None) + + +def test_active_mapping_and_sizes_are_reused_until_membership_changes(monkeypatch): + calls = [] + _patch_reinit_op( + monkeypatch, + lambda group, active_group, rendezvous_id: calls.append((list(group), list(active_group))), + ) + comm = AllGatherReduceScatter(_FakeMapping()) + original_mapping = comm._active_mapping + original_sizes = [5, 7, 11, 13] + + assert comm._active_sizes(original_sizes) is original_sizes + assert comm._active_mapping is original_mapping + + comm.abort_and_reinit([4, 6, 7]) + survivor_mapping = comm._active_mapping + assert comm._active_sizes(original_sizes) == [5, 11, 13] + assert comm._active_mapping is survivor_mapping + comm._refresh_active_mapping() + assert comm._active_mapping is survivor_mapping + + +def test_process_group_backend_skips_raw_nccl_membership_refresh(monkeypatch): + monkeypatch.setattr(allgather_rs_module, "mpi_disabled", lambda: True) + comm = AllGatherReduceScatter(_FakeMapping()) + monkeypatch.setattr( + allgather_rs_module, + "resolve_nccl_group", + lambda group: (_ for _ in ()).throw( + AssertionError("ProcessGroup dispatch consulted raw NCCL state") + ), + ) + observed = [] + monkeypatch.setattr( + allgather_rs_module, + "allgather", + lambda inputs, mapping, *, dim, sizes: observed.append((mapping, sizes)) or inputs, + ) + sizes = [5, 7, 11, 13] + + comm.dispatch("hidden", None, "slots", None, sizes) + + assert observed == [(comm.mapping, sizes)] + + +def test_reconfiguration_rejects_static_sharded_allreduce(monkeypatch): + _patch_reinit_op(monkeypatch, lambda group, active_group, rendezvous_id: None) + mapping = _FakeMapping() + comm = AllGatherReduceScatter(mapping) + comm.abort_and_reinit([4, 6, 7]) + + calls = [] + + def fake_allreduce(**kwargs): + calls.append(kwargs) + return ("reduced",) + + allreduce = object.__new__(distributed_ops.AllReduce) + allreduce.mapping = mapping + allreduce._tp_group_tuple = tuple(mapping.tp_group) + allreduce.strategy = distributed_ops.AllReduceStrategy.AUTO + allreduce.workspace = None + allreduce.symm_mem_allreduce = None + allreduce.mnnvl_allreduce = None + allreduce._disable_mpi = False + allreduce.all_reduce_op = fake_allreduce + + tensor = allgather_rs_module.torch.ones(1) + with pytest.raises(RuntimeError, match="statically sharded|redistribute"): + allreduce.forward(tensor) + assert calls == [] + assert not allreduce.uses_nccl_symmetric_memory_window() + + +@pytest.mark.parametrize( + "active_ranks", + [ + pytest.param([], id="empty"), + pytest.param([4, 6, 6], id="duplicate"), + pytest.param([-1, 6], id="negative"), + pytest.param([4, 6, 8], id="outside_tp_group"), + pytest.param([4, 5, 7], id="world_rank_missing"), + ], +) +def test_abort_and_reinit_rejects_invalid_active_ranks(monkeypatch, active_ranks): + calls = [] + _patch_reinit_op( + monkeypatch, lambda group, active_group, rendezvous_id: calls.append((group, active_group)) + ) + comm = AllGatherReduceScatter(_FakeMapping()) + + with pytest.raises(ValueError): + comm.abort_and_reinit(active_ranks) + + assert calls == [] + + +def test_abort_and_reinit_supports_mixed_tp_and_ep_topology(monkeypatch): + calls = [] + _patch_reinit_op( + monkeypatch, lambda group, active_group, rendezvous_id: calls.append((group, active_group)) + ) + mapping = _FakeMapping( + rank=6, + tp_rank=2, + moe_ep_size=2, + moe_ep_rank=1, + _group=[4, 5, 6, 7], + _moe_group=[4, 6], + ) + comm = AllGatherReduceScatter(mapping) + + # Global ranks are unambiguous even when the TP communicator spans several + # MoE-EP slices. The higher-level health coordinator performs its local to + # global translation before entering this API. + comm.abort_and_reinit([4, 6, 7]) + + assert calls == [([4, 5, 6, 7], [4, 6, 7])] + + +def test_abort_and_reinit_rejects_process_group_backend(monkeypatch): + calls = [] + _patch_reinit_op( + monkeypatch, lambda group, active_group, rendezvous_id: calls.append((group, active_group)) + ) + monkeypatch.setattr(allgather_rs_module, "mpi_disabled", lambda: True, raising=False) + comm = AllGatherReduceScatter(_FakeMapping()) + + # The raw communicator cache is not used in Ray/ProcessGroup mode, and the + # active mapping would otherwise delegate to the stale tp_group_pg. + with pytest.raises(NotImplementedError, match="ProcessGroup|MPI-disabled"): + comm.abort_and_reinit([4, 6, 7]) + + assert calls == [] + + +def test_check_async_error_returns_cleanly_when_watchdog_is_healthy(monkeypatch): + queried_groups = [] + _patch_async_error_op( + monkeypatch, + lambda group: queried_groups.append(list(group)), + ) + comm = AllGatherReduceScatter(_FakeMapping()) + + comm.check_async_error() + + assert queried_groups == [[4, 5, 6, 7]] + + +def test_check_async_error_rejects_process_group_backend(monkeypatch): + queried_groups = [] + _patch_async_error_op( + monkeypatch, + lambda group: queried_groups.append(list(group)), + ) + monkeypatch.setattr(allgather_rs_module, "mpi_disabled", lambda: True, raising=False) + comm = AllGatherReduceScatter(_FakeMapping()) + + with pytest.raises(NotImplementedError, match="ProcessGroup|watchdog"): + comm.check_async_error() + + assert queried_groups == [] + + +def test_check_async_error_raises_stable_nccl_error_for_active_group(monkeypatch): + _patch_reinit_op(monkeypatch, lambda group, active_group, rendezvous_id: None) + queried_groups = [] + + def get_async_error(group): + queried_groups.append(list(group)) + return "ncclSystemError: peer disconnected" + + _patch_async_error_op(monkeypatch, get_async_error) + comm = AllGatherReduceScatter(_FakeMapping()) + comm.abort_and_reinit([4, 6, 7]) + + with pytest.raises( + RuntimeError, + match="NCCL error: communicator was aborted: ncclSystemError: peer disconnected", + ): + comm.check_async_error() + + assert queried_groups == [[4, 6, 7]] + + +def test_other_layer_queries_shared_active_group_after_recovery(monkeypatch): + _patch_reinit_op(monkeypatch, lambda group, active_group, rendezvous_id: None) + queried_groups = [] + _patch_async_error_op(monkeypatch, lambda group: queried_groups.append(list(group))) + first_layer = AllGatherReduceScatter(_FakeMapping()) + second_layer = AllGatherReduceScatter(_FakeMapping()) + + first_layer.abort_and_reinit([4, 6, 7]) + second_layer.check_async_error() + + assert queried_groups == [[4, 6, 7]] + _assert_active_mapping(second_layer._active_mapping, [4, 6, 7], 1) + + +def test_reconfigured_row_linear_fails_before_using_incomplete_tp_shards(monkeypatch): + _patch_reinit_op(monkeypatch, lambda group, active_group, rendezvous_id: None) + AllGatherReduceScatter(_FakeMapping()).abort_and_reinit([6, 7]) + calls = [] + + class QuantMethod: + supports_nccl_symmetric_memory_window_output = False + + class FakeRowLinear: + tp_mode = TensorParallelMode.ROW + mapping = _FakeMapping(rank=6, tp_rank=2) + _tp_group_tuple = tuple(mapping.tp_group) + tp_rank = 2 + tp_size = 4 + use_fused_gemm_allreduce = True + reduce_output = True + bias = "bias" + quant_method = QuantMethod() + lora = None + + def apply_linear_allreduce(self, *args, **kwargs): + raise AssertionError("stale fused GEMM-allreduce path was used") + + def _maybe_fuse_bias_into_allreduce(self, bias, all_reduce_params): + return False + + def apply_linear(self, input, bias, lora_params, layer_idx): + calls.append(("gemm", input, bias)) + return "gemm-output" + + def all_reduce(self, output, *, all_reduce_params): + calls.append(("allreduce", output)) + return "reduced" + + with pytest.raises(RuntimeError, match="statically sharded|redistribute"): + Linear.forward(FakeRowLinear(), "input") + assert calls == [] + + +def test_reconfigured_column_linear_without_gather_fails_before_gemm(monkeypatch): + _patch_reinit_op(monkeypatch, lambda group, active_group, rendezvous_id: None) + AllGatherReduceScatter(_FakeMapping()).abort_and_reinit([6, 7]) + calls = [] + + class FakeColumnLinear: + tp_mode = TensorParallelMode.COLUMN + mapping = _FakeMapping(rank=6, tp_rank=2) + _tp_group_tuple = tuple(mapping.tp_group) + tp_size = 4 + gather_output = False + bias = None + + def apply_linear(self, *args, **kwargs): + calls.append((args, kwargs)) + return "partial-output" + + with pytest.raises(RuntimeError, match="statically sharded|redistribute"): + Linear.forward(FakeColumnLinear(), "input") + assert calls == [] + + +def test_dispatch_and_combine_filter_sizes_to_active_ranks(monkeypatch): + _patch_reinit_op(monkeypatch, lambda group, active_group, rendezvous_id: None) + gather_calls = [] + scatter_calls = [] + + def fake_allgather(inputs, mapping, *, dim, sizes): + gather_calls.append((inputs, mapping, dim, sizes)) + return inputs + + def fake_reducescatter(input_tensor, mapping, *, dim, sizes): + scatter_calls.append((input_tensor, mapping, dim, sizes)) + return "combined" + + monkeypatch.setattr(allgather_rs_module, "allgather", fake_allgather) + monkeypatch.setattr(allgather_rs_module, "reducescatter", fake_reducescatter) + + mapping = _FakeMapping() + comm = AllGatherReduceScatter(mapping) + comm.abort_and_reinit([4, 6, 7]) + + inputs = ["hidden", None, "slots", "scales"] + assert comm.dispatch(*inputs, all_rank_num_tokens=[5, 7, 11, 13]) == tuple(inputs) + assert comm.combine("moe-output") == "combined" + + assert len(gather_calls) == 1 + gathered_inputs, gather_mapping, gather_dim, gather_sizes = gather_calls[0] + assert gathered_inputs == inputs + assert gather_dim == 0 + assert gather_sizes == [5, 11, 13] + _assert_active_mapping(gather_mapping, [4, 6, 7], 1) + + assert len(scatter_calls) == 1 + scatter_input, scatter_mapping, scatter_dim, scatter_sizes = scatter_calls[0] + assert scatter_input == "moe-output" + assert scatter_dim == 0 + assert scatter_sizes == [5, 11, 13] + _assert_active_mapping(scatter_mapping, [4, 6, 7], 1) + + # Reconfiguration does not mutate Mapping, but the shared raw-NCCL + # membership registry reroutes other collectives that consume it. + assert mapping.tp_group == [4, 5, 6, 7] + assert mapping.tp_rank == 2 + + +def test_dp_padding_keeps_sizes_none_after_reconfiguration(monkeypatch): + _patch_reinit_op(monkeypatch, lambda group, active_group, rendezvous_id: None) + observed = [] + + def fake_allgather(inputs, mapping, *, dim, sizes): + observed.append((mapping, sizes)) + return inputs + + monkeypatch.setattr(allgather_rs_module, "allgather", fake_allgather) + + comm = AllGatherReduceScatter(_FakeMapping()) + comm.abort_and_reinit([4, 6, 7]) + comm.dispatch("hidden", None, "slots", None, [5, 7, 11, 13], use_dp_padding=True) + + assert len(observed) == 1 + active_mapping, sizes = observed[0] + assert sizes is None + _assert_active_mapping(active_mapping, [4, 6, 7], 1) + + +def test_dp_padding_refreshes_mapping_after_another_layer_recovers(monkeypatch): + _patch_reinit_op(monkeypatch, lambda group, active_group, rendezvous_id: None) + observed = [] + + def fake_allgather(inputs, mapping, *, dim, sizes): + observed.append((mapping, sizes)) + return inputs + + monkeypatch.setattr(allgather_rs_module, "allgather", fake_allgather) + recovering_layer = AllGatherReduceScatter(_FakeMapping()) + stale_layer = AllGatherReduceScatter(_FakeMapping()) + + recovering_layer.abort_and_reinit([4, 6, 7]) + stale_layer.dispatch("hidden", None, "slots", None, [5, 7, 11, 13], use_dp_padding=True) + + assert len(observed) == 1 + active_mapping, sizes = observed[0] + assert sizes is None + _assert_active_mapping(active_mapping, [4, 6, 7], 1) + + +def test_later_reconfiguration_replaces_active_rank_filter(monkeypatch): + calls = [] + _patch_reinit_op( + monkeypatch, + lambda group, active_group, rendezvous_id: calls.append((list(group), list(active_group))), + ) + observed = [] + + def fake_allgather(inputs, mapping, *, dim, sizes): + observed.append((mapping, sizes)) + return inputs + + monkeypatch.setattr(allgather_rs_module, "allgather", fake_allgather) + + comm = AllGatherReduceScatter(_FakeMapping()) + comm.abort_and_reinit([4, 6, 7]) + comm.abort_and_reinit([6, 7]) + comm.dispatch("hidden", None, "slots", None, [5, 7, 11, 13]) + + assert calls == [ + ([4, 5, 6, 7], [4, 6, 7]), + ([4, 6, 7], [6, 7]), + ] + assert len(observed) == 1 + active_mapping, sizes = observed[0] + assert sizes == [11, 13] + _assert_active_mapping(active_mapping, [6, 7], 0) + + +def test_reconfiguration_cannot_reactivate_a_removed_rank(monkeypatch): + calls = [] + _patch_reinit_op( + monkeypatch, + lambda group, active_group, rendezvous_id: calls.append((list(group), list(active_group))), + ) + + comm = AllGatherReduceScatter(_FakeMapping()) + comm.abort_and_reinit([4, 6, 7]) + + with pytest.raises(ValueError, match="reactivate|subset"): + comm.abort_and_reinit([4, 5, 6, 7]) + + assert calls == [([4, 5, 6, 7], [4, 6, 7])] + + +def test_failed_reinit_does_not_commit_active_rank_filter(monkeypatch): + def fail_reinit(group, active_group, rendezvous_id): + raise RuntimeError("NCCL communicator rebuild failed") + + _patch_reinit_op(monkeypatch, fail_reinit) + comm = AllGatherReduceScatter(_FakeMapping()) + old_local_ranks = comm._active_local_ranks + old_global_group = comm._active_global_group + old_mapping = comm._active_mapping + + with pytest.raises(RuntimeError, match="NCCL communicator rebuild failed"): + comm.abort_and_reinit([4, 6, 7]) + + # The native call aborts the old communicator before initializing its + # replacement, so transport is terminal after this failure. The only + # transactional guarantee is that Python membership metadata does not + # claim the failed replacement was committed. + assert comm._active_local_ranks is old_local_ranks + assert comm._active_global_group is old_global_group + assert comm._active_mapping is old_mapping diff --git a/tests/unittest/_torch/modules/moe/test_moe_comm.py b/tests/unittest/_torch/modules/moe/test_moe_comm.py index 057f87a95f7a..ac75c895cf6d 100644 --- a/tests/unittest/_torch/modules/moe/test_moe_comm.py +++ b/tests/unittest/_torch/modules/moe/test_moe_comm.py @@ -225,14 +225,31 @@ def _read_nvlink_topk_target_ranks( return raw.view(torch.int32).view(max_num_tokens, top_k).cpu() -def _run_nvlink_rank_mask_dispatch_combine( +def _read_nvlink_topk_send_indices( + comm: NVLinkOneSided, + max_num_tokens: int, + top_k: int, +) -> torch.Tensor: + """Read topk_send_indices[max_num_tokens, top_k] from NVLinkOneSided workspace.""" + from tensorrt_llm.bindings import internal as _tllm_internal + + offset_index = int(_tllm_internal.thop.MOE_A2A_TOPK_SEND_INDICES_OFFSET_INDEX) + offset = comm.moe_a2a_metainfo[offset_index].item() + raw = comm.workspace[ + comm.ep_rank, + offset : offset + max_num_tokens * top_k * 4, + ] + return raw.view(torch.int32).view(max_num_tokens, top_k).cpu() + + +def _run_nvlink_rank_mask_dispatch( comm: NVLinkOneSided, token_selected_experts: torch.Tensor, payload: torch.Tensor, runtime_max_tokens_per_rank: int, active_rank_mask: Optional[torch.Tensor], -) -> Tuple[torch.Tensor, torch.Tensor]: - """Run raw NVLink one-sided dispatch/combine with an optional active rank mask.""" +) -> Tuple[List[torch.Tensor], int, torch.Tensor, torch.Tensor]: + """Run raw NVLink one-sided dispatch with an optional active rank mask.""" recv_tensors, combine_payload_offset, _ = torch.ops.trtllm.moe_a2a_dispatch( token_selected_experts, [payload], @@ -252,24 +269,105 @@ def _run_nvlink_rank_mask_dispatch_combine( runtime_max_tokens_per_rank, comm.top_k, ) + topk_send_indices = _read_nvlink_topk_send_indices( + comm, + runtime_max_tokens_per_rank, + comm.top_k, + ) + return recv_tensors, int(combine_payload_offset), topk_target_ranks, topk_send_indices - combined = torch.ops.trtllm.moe_a2a_combine( - recv_tensors[0], - token_selected_experts.size(0), + +def _run_nvlink_rank_mask_combine( + comm: NVLinkOneSided, + combine_payload: torch.Tensor, + local_num_tokens: int, + runtime_max_tokens_per_rank: int, + combine_payload_offset: int, + active_rank_mask: Optional[torch.Tensor], +) -> torch.Tensor: + """Run raw NVLink one-sided combine with an optional active rank mask.""" + return torch.ops.trtllm.moe_a2a_combine( + combine_payload, + local_num_tokens, comm.workspace, comm.moe_a2a_metainfo, runtime_max_tokens_per_rank, comm.ep_rank, comm.ep_size, comm.top_k, - int(combine_payload_offset), + combine_payload_offset, False, # payload_in_workspace False, # use_low_precision active_rank_mask, ) + + +def _run_nvlink_rank_mask_dispatch_combine( + comm: NVLinkOneSided, + token_selected_experts: torch.Tensor, + payload: torch.Tensor, + runtime_max_tokens_per_rank: int, + active_rank_mask: Optional[torch.Tensor], +) -> Tuple[torch.Tensor, torch.Tensor]: + """Run raw NVLink one-sided dispatch/combine with an optional active rank mask.""" + recv_tensors, combine_payload_offset, topk_target_ranks, _ = _run_nvlink_rank_mask_dispatch( + comm, + token_selected_experts, + payload, + runtime_max_tokens_per_rank, + active_rank_mask, + ) + combined = _run_nvlink_rank_mask_combine( + comm, + recv_tensors[0], + token_selected_experts.size(0), + runtime_max_tokens_per_rank, + combine_payload_offset, + active_rank_mask, + ) return combined.cpu(), topk_target_ranks +def _expected_nvlink_rank_mask_combine_output( + comm: NVLinkOneSided, + payload: torch.Tensor, + topk_target_ranks: torch.Tensor, + topk_send_indices: torch.Tensor, + local_num_tokens: int, + runtime_max_tokens_per_rank: int, + dead_ranks: Set[int], +) -> torch.Tensor: + """Compute combine output from dispatched workspace while skipping dead ranks.""" + from tensorrt_llm.bindings import internal as _tllm_internal + + hidden_size = payload.shape[-1] + expected = torch.zeros( + (local_num_tokens, hidden_size), + dtype=payload.dtype, + device=payload.device, + ) + payload_offset_index = int(_tllm_internal.thop.MOE_A2A_PAYLOAD_DATA_OFFSET_INDEX) + payload_offset = comm.moe_a2a_metainfo[payload_offset_index].item() + bytes_per_rank = ( + comm.ep_size * runtime_max_tokens_per_rank * hidden_size * payload.element_size() + ) + + for token_idx in range(local_num_tokens): + for k in range(comm.top_k): + target_rank = int(topk_target_ranks[token_idx, k].item()) + dst_idx = int(topk_send_indices[token_idx, k].item()) + if dst_idx < 0 or target_rank in dead_ranks: + continue + raw = comm.workspace[target_rank, payload_offset : payload_offset + bytes_per_rank] + recv_payload = raw.view(payload.dtype).view( + comm.ep_size, + runtime_max_tokens_per_rank, + hidden_size, + ) + expected[token_idx] += recv_payload[comm.ep_rank, dst_idx] + return expected.cpu() + + # ============================================================================ # Source Encoding Utilities # ============================================================================ @@ -992,6 +1090,7 @@ def _worker_rank_mask_all_active_matches_no_mask(config: CommTestConfig) -> dict ) return { + "rank": rank, "output_eq": torch.equal(out_no_mask, out_all_active), "topk_eq": torch.equal(topk_no_mask, topk_all_active), } @@ -1040,7 +1139,7 @@ def _worker_rank_mask_one_rank_masked( if rank == dead_rank: MPI.COMM_WORLD.barrier() - return {"status": "dead"} + return {"rank": rank, "status": "dead"} local_num_tokens = config.all_num_tokens[rank] torch.manual_seed(0xA2A + rank) @@ -1069,6 +1168,7 @@ def _worker_rank_mask_one_rank_masked( MPI.COMM_WORLD.barrier() return { + "rank": rank, "status": "alive", "combined": combined, "topk_target_ranks": topk_target_ranks, @@ -1082,6 +1182,83 @@ def _worker_rank_mask_one_rank_masked( comm.destroy() +def _worker_rank_mask_inactive_before_combine( + config: CommTestConfig, + dead_rank: int, +) -> dict: + """Dispatch with all ranks active, then omit one rank from combine's active mask.""" + rank = tllm.mpi_rank() + torch.cuda.set_device(rank) + + comm = None + try: + mapping = Mapping( + rank=rank, + tp_size=config.ep_size, + moe_ep_size=config.ep_size, + world_size=config.ep_size, + ) + comm = create_comm_object(config.comm_type, mapping, config) + + local_num_tokens = config.all_num_tokens[rank] + torch.manual_seed(0xA2A + rank) + token_selected_experts = torch.randint( + 0, + config.num_experts, + (local_num_tokens, config.top_k), + dtype=torch.int32, + device="cuda", + ) + payload = _make_rank_mask_payload(local_num_tokens, config.hidden_size, rank) + + recv_tensors, combine_payload_offset, topk_target_ranks, topk_send_indices = ( + _run_nvlink_rank_mask_dispatch( + comm, + token_selected_experts, + payload, + local_num_tokens, + active_rank_mask=_ep_mask_words(config.ep_size, dead_ranks=set()), + ) + ) + + if rank == dead_rank: + MPI.COMM_WORLD.barrier() + return {"rank": rank, "status": "dead"} + + dead_ranks = {dead_rank} + expected = _expected_nvlink_rank_mask_combine_output( + comm, + payload, + topk_target_ranks, + topk_send_indices, + local_num_tokens, + local_num_tokens, + dead_ranks, + ) + combined = _run_nvlink_rank_mask_combine( + comm, + recv_tensors[0], + local_num_tokens, + local_num_tokens, + combine_payload_offset, + active_rank_mask=_ep_mask_words(config.ep_size, dead_ranks=dead_ranks), + ).cpu() + + MPI.COMM_WORLD.barrier() + return { + "rank": rank, + "status": "alive", + "combined": combined, + "expected": expected, + } + except Exception: + traceback.print_exc() + raise + finally: + if comm is not None and hasattr(comm, "destroy"): + comm.destroy() + + # ============================================================================ # Verification Functions # ============================================================================ @@ -1807,7 +1984,7 @@ def _make_postquant_test_params(): @pytest.fixture(autouse=True) -def setup_test(): +def setup_test() -> None: torch.manual_seed(0x1234) tllm.logger.set_level("error") @@ -1886,7 +2063,8 @@ def _run_rank_mask_all_active_test( ) ) - for rank, result in enumerate(results): + for result in results: + rank = result["rank"] assert result["output_eq"], ( f"rank {rank}: combine output differs between no-mask and all-active mask" ) @@ -1915,7 +2093,8 @@ def _run_rank_mask_one_rank_masked_test( ) saw_dead = False - for rank, result in enumerate(results): + for result in results: + rank = result["rank"] if result["status"] == "dead": assert rank == dead_rank saw_dead = True @@ -1952,6 +2131,45 @@ def _run_rank_mask_one_rank_masked_test( assert saw_dead, f"dead rank {dead_rank} did not appear in results" +def _run_rank_mask_inactive_before_combine_test( + mpi_pool_executor, + dead_rank: int, + local_num_tokens: int, + top_k: int, +) -> None: + ep_size = mpi_pool_executor.num_workers + config = _make_rank_mask_config(ep_size, local_num_tokens, top_k) + _skip_if_rank_mask_config_unsupported(config) + assert 0 <= dead_rank < ep_size + + worker_args = [(config, dead_rank)] * config.ep_size + results = list( + mpi_pool_executor.map( + _worker_rank_mask_inactive_before_combine, + *zip(*worker_args), + ) + ) + + saw_dead = False + for result in results: + rank = result["rank"] + if result["status"] == "dead": + assert rank == dead_rank + saw_dead = True + continue + + assert result["status"] == "alive" + combined = result["combined"] + expected = result["expected"] + assert combined is not None + assert expected is not None + assert torch.equal(combined, expected), ( + f"rank {rank}: combine output included a rank masked inactive before combine" + ) + + assert saw_dead, f"dead rank {dead_rank} did not appear in results" + + # ============================================================================ # Test Class # ============================================================================ @@ -2019,7 +2237,7 @@ def test_moe_comm_rank_mask_all_active_matches_no_mask( mpi_pool_executor, local_num_tokens: int, top_k: int, - ): + ) -> None: """Verify all-active active_rank_mask matches omitted mask for NVLinkOneSided.""" _run_rank_mask_all_active_test(mpi_pool_executor, local_num_tokens, top_k) @@ -2039,7 +2257,7 @@ def test_moe_comm_rank_mask_one_rank_masked_completes( dead_rank: int, local_num_tokens: int, top_k: int, - ): + ) -> None: """Verify masked-dead rank is skipped by raw NVLinkOneSided moe_a2a ops.""" _run_rank_mask_one_rank_masked_test( mpi_pool_executor, @@ -2047,3 +2265,26 @@ def test_moe_comm_rank_mask_one_rank_masked_completes( local_num_tokens, top_k, ) + + @pytest.mark.threadleak(enabled=False) + @pytest.mark.parametrize( + "mpi_pool_executor,dead_rank,local_num_tokens,top_k", + [ + (4, 2, 16, 2), + ], + indirect=["mpi_pool_executor"], + ) + def test_moe_comm_rank_mask_inactive_before_combine_skips_stale_dispatch_slots( + self, + mpi_pool_executor, + dead_rank: int, + local_num_tokens: int, + top_k: int, + ) -> None: + """Verify combine skips slots from a rank masked inactive after dispatch.""" + _run_rank_mask_inactive_before_combine_test( + mpi_pool_executor, + dead_rank, + local_num_tokens, + top_k, + ) diff --git a/tests/unittest/_torch/modules/test_alltoall_watchdog.py b/tests/unittest/_torch/modules/test_alltoall_watchdog.py new file mode 100644 index 000000000000..cdff07cbcea4 --- /dev/null +++ b/tests/unittest/_torch/modules/test_alltoall_watchdog.py @@ -0,0 +1,643 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. +"""Unit tests for AlltoAllWatchdog (WideEP fault tolerance, PR 1a.4).""" + +import threading +import time +from collections.abc import Callable +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.alltoall_watchdog import ( + DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, + DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S, + UNKNOWN_COMPLETION_FLAG, + AlltoAllWatchdog, + AlltoAllWatchdogCoordinator, + AlltoAllWatchdogTimeout, + CompletionFlagReadTimeout, +) +from tensorrt_llm._torch.modules.fused_moe.ep_group_health import EPGroupHealth +from tensorrt_llm._torch.modules.fused_moe.wide_ep_ft import get_wide_ep_ft_options + + +class FakeCompletionFlagReader: + """Thread-safe completion flag reader for pure-Python watchdog tests.""" + + def __init__(self, ep_size: int) -> None: + self._lock = threading.Lock() + self._flags = { + "dispatch": [0 for _ in range(ep_size)], + "combine": [0 for _ in range(ep_size)], + } + + def set_flags(self, phase: str, flags: list[int]) -> None: + with self._lock: + self._flags[phase] = list(flags) + + def read_completion_flags(self, phase: str) -> tuple[int, ...]: + with self._lock: + return tuple(self._flags[phase]) + + +class TimeoutCompletionFlagReader: + def read_completion_flags(self, phase: str) -> tuple[int, ...]: + raise CompletionFlagReadTimeout("blocked") + + +class OneGoodReadThenTimeoutReader: + def __init__(self, flags: tuple[int, ...]) -> None: + self._flags = flags + self._read_count = 0 + + def read_completion_flags(self, phase: str) -> tuple[int, ...]: + self._read_count += 1 + if self._read_count == 1: + return self._flags + raise CompletionFlagReadTimeout("blocked") + + +def _wait_for(predicate: Callable[[], bool], timeout_s: float = 1.0) -> None: + deadline = time.monotonic() + timeout_s + while time.monotonic() < deadline: + if predicate(): + return + time.sleep(0.005) + raise AssertionError("condition was not reached before timeout") + + +def test_watchdog_completes_when_all_active_flags_arrive() -> None: + health = EPGroupHealth(4) + reader = FakeCompletionFlagReader(ep_size=4) + events: list[AlltoAllWatchdogTimeout] = [] + + with AlltoAllWatchdog( + ep_size=4, + ep_rank=0, + completion_reader=reader, + timeout_s=0.2, + poll_interval_s=0.005, + health=health, + on_timeout=events.append, + ) as watchdog: + watchdog.watch(phase="dispatch", expected_flag=1) + reader.set_flags("dispatch", [1, 1, 1, 1]) + assert watchdog.wait_until_idle(timeout_s=1.0) + + assert events == [] + assert health.all_active() is True + + +def test_watchdog_completes_when_flags_advance_past_expected_generation() -> None: + health = EPGroupHealth(4) + reader = FakeCompletionFlagReader(ep_size=4) + reader.set_flags("dispatch", [2, 1, 5, 3]) + events: list[AlltoAllWatchdogTimeout] = [] + + with AlltoAllWatchdog( + ep_size=4, + ep_rank=0, + completion_reader=reader, + timeout_s=0.02, + poll_interval_s=0.005, + health=health, + on_timeout=events.append, + ) as watchdog: + watchdog.watch(phase="dispatch", expected_flag=1) + assert watchdog.wait_until_idle(timeout_s=1.0) + + assert events == [] + assert health.all_active() is True + + +def test_watchdog_handles_signed_uint32_generation_boundaries() -> None: + reader = FakeCompletionFlagReader(ep_size=2) + events: list[AlltoAllWatchdogTimeout] = [] + + with AlltoAllWatchdog( + ep_size=2, + ep_rank=0, + completion_reader=reader, + timeout_s=0.05, + poll_interval_s=0.005, + on_timeout=events.append, + ) as watchdog: + reader.set_flags("dispatch", [-(1 << 31), -(1 << 31)]) + watchdog.watch(phase="dispatch", expected_flag=1 << 31) + assert watchdog.wait_until_idle(timeout_s=1.0) + + reader.set_flags("dispatch", [-1, -1]) + watchdog.watch(phase="dispatch", expected_flag=(1 << 32) - 1) + assert watchdog.wait_until_idle(timeout_s=1.0) + + reader.set_flags("dispatch", [0, 0]) + watchdog.watch(phase="dispatch", expected_flag=0) + assert watchdog.wait_until_idle(timeout_s=1.0) + + assert events == [] + + +def test_watchdog_defaults_match_design_doc() -> None: + reader = FakeCompletionFlagReader(ep_size=1) + watchdog = AlltoAllWatchdog(ep_size=1, ep_rank=0, completion_reader=reader) + + assert watchdog._timeout_s == DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S + assert watchdog._poll_interval_s == DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S + + +def test_watchdog_stop_is_terminal() -> None: + reader = FakeCompletionFlagReader(ep_size=1) + watchdog = AlltoAllWatchdog( + ep_size=1, + ep_rank=0, + completion_reader=reader, + timeout_s=0.2, + poll_interval_s=0.005, + ) + watchdog.start() + watchdog.stop(timeout_s=1.0) + + with pytest.raises(RuntimeError, match="stopped AlltoAllWatchdog"): + watchdog.start() + with pytest.raises(RuntimeError, match="stopped AlltoAllWatchdog"): + watchdog.watch(phase="dispatch", expected_flag=1) + + +def test_wide_ep_ft_options_create_shared_health_when_enabled( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("TLLM_FAULT_TOLERANCE_MODE", "1") + model_config = SimpleNamespace( + extra_attrs={}, + mapping=SimpleNamespace(moe_ep_size=4), + ) + + health, timeout_s, poll_interval_s = get_wide_ep_ft_options(model_config) + health_again, timeout_again_s, poll_again_s = get_wide_ep_ft_options(model_config) + + assert isinstance(health, EPGroupHealth) + assert health_again is health + assert timeout_s == DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S + assert timeout_again_s == timeout_s + assert poll_interval_s == DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S + assert poll_again_s == poll_interval_s + + +def test_wide_ep_ft_options_ignore_legacy_enable_flag(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("TLLM_FAULT_TOLERANCE_MODE", raising=False) + monkeypatch.setenv("TRTLLM_ENABLE_WIDE_EP_FT", "1") + model_config = SimpleNamespace( + extra_attrs={}, + mapping=SimpleNamespace(moe_ep_size=4), + ) + + health, timeout_s, _ = get_wide_ep_ft_options(model_config) + + assert health is None + assert timeout_s is None + + +def test_watchdog_coordinator_reuses_committed_mask_when_generation_is_unchanged() -> None: + health = EPGroupHealth(4) + coordinator = AlltoAllWatchdogCoordinator( + workspace_state={}, + workspace=torch.zeros((4, 1), dtype=torch.uint8), + metainfo=torch.zeros((1,), dtype=torch.int64), + metainfo_index={}, + ep_rank=0, + health=health, + ) + dispatch_snapshot = coordinator.capture_active_rank_mask(None) + assert dispatch_snapshot.active_rank_mask is not None + assert dispatch_snapshot.committed_generation == 0 + + combine_mask = coordinator.active_rank_mask_for_combine(dispatch_snapshot, None) + + assert combine_mask is dispatch_snapshot.active_rank_mask + assert combine_mask.tolist() == [0b1111, 0] + + +def test_watchdog_coordinator_fails_closed_on_committed_generation_change() -> None: + health = EPGroupHealth(4) + coordinator = AlltoAllWatchdogCoordinator( + workspace_state={}, + workspace=torch.zeros((4, 1), dtype=torch.uint8), + metainfo=torch.zeros((1,), dtype=torch.int64), + metainfo_index={}, + ep_rank=0, + health=health, + ) + dispatch_snapshot = coordinator.capture_active_rank_mask(None) + + health.mark_failed(2) + health.mark_active(2) + assert health.get_mask() == 0b1111 + + with pytest.raises( + RuntimeError, + match="committed EP membership changed between dispatch and combine", + ): + coordinator.active_rank_mask_for_combine(dispatch_snapshot, None) + + +def test_watchdog_coordinator_converts_atomic_snapshot_to_two_mask_words() -> None: + health = EPGroupHealth(72) + health.mark_failed(70) + coordinator = AlltoAllWatchdogCoordinator( + workspace_state={}, + workspace=torch.zeros((72, 1), dtype=torch.uint8), + metainfo=torch.zeros((1,), dtype=torch.int64), + metainfo_index={}, + ep_rank=0, + health=health, + ) + + dispatch_snapshot = coordinator.capture_active_rank_mask(None) + + assert dispatch_snapshot.committed_generation == 1 + assert dispatch_snapshot.active_rank_mask is not None + assert dispatch_snapshot.active_rank_mask.tolist() == [(1 << 64) - 1, 0xBF] + + +def test_watchdog_coordinator_explicit_mask_is_not_bound_to_health_generation() -> None: + health = EPGroupHealth(4) + coordinator = AlltoAllWatchdogCoordinator( + workspace_state={}, + workspace=torch.zeros((4, 1), dtype=torch.uint8), + metainfo=torch.zeros((1,), dtype=torch.int64), + metainfo_index={}, + ep_rank=0, + health=health, + ) + explicit_mask = torch.tensor([0b1101, 0], dtype=torch.uint64) + dispatch_snapshot = coordinator.capture_active_rank_mask(explicit_mask) + assert dispatch_snapshot.committed_generation is None + + health.mark_failed(2) + combine_mask = coordinator.active_rank_mask_for_combine(dispatch_snapshot, None) + + assert combine_mask is dispatch_snapshot.active_rank_mask + assert combine_mask.tolist() == explicit_mask.tolist() + with pytest.raises(ValueError, match="mask captured at dispatch"): + coordinator.active_rank_mask_for_combine( + dispatch_snapshot, + torch.tensor(health.get_mask_words(), dtype=torch.uint64), + ) + + +def test_watchdog_no_detected_failure_publication_to_committed_health() -> None: + health = EPGroupHealth(4) + committed_before = health.snapshot() + reader = FakeCompletionFlagReader(ep_size=4) + reader.set_flags("dispatch", [1, 0, 1, 0]) + events: list[AlltoAllWatchdogTimeout] = [] + + with AlltoAllWatchdog( + ep_size=4, + ep_rank=0, + completion_reader=reader, + timeout_s=0.02, + poll_interval_s=0.005, + health=health, + on_timeout=events.append, + ) as watchdog: + watchdog.watch(phase="dispatch", expected_flag=1) + _wait_for(lambda: len(events) == 1) + assert watchdog.wait_until_idle(timeout_s=1.0) + + event = events[0] + assert event.phase == "dispatch" + assert event.expected_flag == 1 + assert event.observed_flags == (1, 0, 1, 0) + assert event.missing_ranks == (1, 3) + assert not hasattr(event, "marked_failed_ranks") + assert health.snapshot() == committed_before + + +def test_watchdog_ignores_ranks_already_failed_in_health_mask() -> None: + health = EPGroupHealth(4) + assert health.mark_failed(2) is True + reader = FakeCompletionFlagReader(ep_size=4) + reader.set_flags("dispatch", [1, 1, 0, 1]) + events: list[AlltoAllWatchdogTimeout] = [] + + with AlltoAllWatchdog( + ep_size=4, + ep_rank=0, + completion_reader=reader, + timeout_s=0.05, + poll_interval_s=0.005, + health=health, + on_timeout=events.append, + ) as watchdog: + watchdog.watch(phase="dispatch", expected_flag=1) + assert watchdog.wait_until_idle(timeout_s=1.0) + + assert events == [] + assert health.get_failed_ranks() == frozenset({2}) + + +def test_watchdog_reports_local_missing_without_changing_committed_health() -> None: + health = EPGroupHealth(4) + reader = FakeCompletionFlagReader(ep_size=4) + reader.set_flags("combine", [0, 2, 2, 2]) + events: list[AlltoAllWatchdogTimeout] = [] + + with AlltoAllWatchdog( + ep_size=4, + ep_rank=0, + completion_reader=reader, + timeout_s=0.02, + poll_interval_s=0.005, + health=health, + on_timeout=events.append, + ) as watchdog: + watchdog.watch(phase="combine", expected_flag=2) + _wait_for(lambda: len(events) == 1) + + event = events[0] + assert event.missing_ranks == (0,) + assert health.get_failed_ranks() == frozenset() + + +def test_watchdog_poll_timeout_without_snapshot_fails_closed() -> None: + health = EPGroupHealth(3) + events: list[AlltoAllWatchdogTimeout] = [] + + with AlltoAllWatchdog( + ep_size=3, + ep_rank=0, + completion_reader=TimeoutCompletionFlagReader(), + timeout_s=0.02, + poll_interval_s=0.005, + health=health, + on_timeout=events.append, + ) as watchdog: + watchdog.watch(phase="dispatch", expected_flag=1) + _wait_for(lambda: len(events) == 1) + + event = events[0] + assert event.poll_timed_out is True + assert event.observed_flags == (UNKNOWN_COMPLETION_FLAG,) * 3 + assert event.missing_ranks == (0, 1, 2) + assert health.all_active() is True + + +def test_watchdog_poll_timeout_with_prior_snapshot_does_not_mark_failed_rank() -> None: + health = EPGroupHealth(3) + events: list[AlltoAllWatchdogTimeout] = [] + + with AlltoAllWatchdog( + ep_size=3, + ep_rank=0, + completion_reader=OneGoodReadThenTimeoutReader((1, 0, 1)), + timeout_s=0.02, + poll_interval_s=0.005, + health=health, + on_timeout=events.append, + ) as watchdog: + watchdog.watch(phase="dispatch", expected_flag=1) + _wait_for(lambda: len(events) == 1) + + event = events[0] + assert event.poll_timed_out is True + assert event.observed_flags == (1, 0, 1) + assert event.missing_ranks == (1,) + assert health.all_active() is True + + +def test_watchdog_callback_error_stops_and_clears_queue() -> None: + health = EPGroupHealth(2) + reader = FakeCompletionFlagReader(ep_size=2) + reader.set_flags("dispatch", [1, 0]) + + def raise_from_callback(event: AlltoAllWatchdogTimeout) -> None: + raise RuntimeError(f"callback failed for {event.phase}") + + with AlltoAllWatchdog( + ep_size=2, + ep_rank=0, + completion_reader=reader, + timeout_s=0.02, + poll_interval_s=0.005, + health=health, + on_timeout=raise_from_callback, + ) as watchdog: + watchdog.watch(phase="dispatch", expected_flag=1) + _wait_for(lambda: watchdog.last_error is not None) + assert watchdog.wait_until_idle(timeout_s=1.0) + assert isinstance(watchdog.last_error, RuntimeError) + with pytest.raises(RuntimeError, match="stopped AlltoAllWatchdog"): + watchdog.watch(phase="dispatch", expected_flag=2) + + +def test_watchdog_preserves_fifo_order_and_clears_followups_after_timeout() -> None: + health = EPGroupHealth(3) + reader = FakeCompletionFlagReader(ep_size=3) + reader.set_flags("dispatch", [1, 0, 1]) + reader.set_flags("combine", [0, 0, 0]) + events: list[AlltoAllWatchdogTimeout] = [] + + with AlltoAllWatchdog( + ep_size=3, + ep_rank=0, + completion_reader=reader, + timeout_s=0.02, + poll_interval_s=0.005, + health=health, + on_timeout=events.append, + ) as watchdog: + watchdog.watch(phase="dispatch", expected_flag=1) + watchdog.watch(phase="combine", expected_flag=2) + _wait_for(lambda: len(events) == 1) + assert watchdog.wait_until_idle(timeout_s=1.0) + time.sleep(0.05) + + assert len(events) == 1 + assert events[0].phase == "dispatch" + assert events[0].missing_ranks == (1,) + assert health.get_failed_ranks() == frozenset() + + +def test_watchdog_from_workspace_reads_phase_specific_offsets() -> None: + ep_size = 3 + ep_rank = 1 + workspace = torch.zeros((ep_size, 64), dtype=torch.uint8) + metainfo = torch.zeros((10,), dtype=torch.int64) + metainfo_index = { + "DISPATCH_COMPLETION_FLAGS_OFFSET_INDEX": 4, + "COMBINE_COMPLETION_FLAGS_OFFSET_INDEX": 5, + } + metainfo[4] = 4 + metainfo[5] = 20 + workspace[ep_rank, 4:16].view(torch.int32).copy_(torch.tensor([7, 7, 7], dtype=torch.int32)) + workspace[ep_rank, 20:32].view(torch.int32).copy_(torch.tensor([0, 8, 8], dtype=torch.int32)) + health = EPGroupHealth(ep_size) + events: list[AlltoAllWatchdogTimeout] = [] + + with AlltoAllWatchdog.from_workspace( + workspace=workspace, + metainfo=metainfo, + metainfo_index=metainfo_index, + ep_rank=ep_rank, + ep_size=ep_size, + timeout_s=0.02, + poll_interval_s=0.005, + health=health, + on_timeout=events.append, + ) as watchdog: + watchdog.watch(phase="dispatch", expected_flag=7) + assert watchdog.wait_until_idle(timeout_s=1.0) + + watchdog.watch(phase="combine", expected_flag=8) + _wait_for(lambda: len(events) == 1) + + assert events[0].phase == "combine" + assert events[0].missing_ranks == (0,) + assert health.get_failed_ranks() == frozenset() + + +def test_workspace_coordinators_share_fifo_watchdog() -> None: + ep_size = 3 + ep_rank = 0 + workspace_state: dict[str, object] = {} + workspace = torch.zeros((ep_size, 64), dtype=torch.uint8) + metainfo = torch.tensor([0, 4, 16], dtype=torch.int64) + metainfo_index = { + "FLAG_VAL_OFFSET_INDEX": 0, + "DISPATCH_COMPLETION_FLAGS_OFFSET_INDEX": 1, + "COMBINE_COMPLETION_FLAGS_OFFSET_INDEX": 2, + } + workspace[ep_rank, 4:16].view(torch.int32).copy_(torch.tensor([1, 0, 1])) + health = EPGroupHealth(ep_size) + events: list[AlltoAllWatchdogTimeout] = [] + on_timeout = events.append + coordinators = [ + AlltoAllWatchdogCoordinator( + workspace_state=workspace_state, + workspace=workspace, + metainfo=metainfo, + metainfo_index=metainfo_index, + ep_rank=ep_rank, + health=health, + ) + for _ in range(2) + ] + watchdogs = [ + coordinator.acquire_watchdog( + ep_size=ep_size, + timeout_s=0.02, + poll_interval_s=0.005, + on_timeout=on_timeout, + ) + for coordinator in coordinators + ] + + try: + assert watchdogs[0] is watchdogs[1] + coordinators[0].watch_collective(watchdogs[0], "dispatch", None) + coordinators[1].watch_collective(watchdogs[1], "combine", None) + _wait_for(lambda: len(events) == 1) + assert watchdogs[0].wait_until_idle(timeout_s=1.0) + time.sleep(0.05) + + assert len(events) == 1 + assert events[0].phase == "dispatch" + assert events[0].missing_ranks == (1,) + assert health.get_failed_ranks() == frozenset() + finally: + for coordinator, watchdog in zip(coordinators, watchdogs): + coordinator.release_watchdog(watchdog) + + +def test_workspace_coordinator_wraps_shared_generation() -> None: + workspace_state: dict[str, object] = {} + workspace = torch.zeros((1, 32), dtype=torch.uint8) + workspace[0, 0:4].view(torch.int32).fill_(-1) + metainfo = torch.tensor([0, 4, 8], dtype=torch.int64) + metainfo_index = { + "FLAG_VAL_OFFSET_INDEX": 0, + "DISPATCH_COMPLETION_FLAGS_OFFSET_INDEX": 1, + "COMBINE_COMPLETION_FLAGS_OFFSET_INDEX": 2, + } + coordinator = AlltoAllWatchdogCoordinator( + workspace_state=workspace_state, + workspace=workspace, + metainfo=metainfo, + metainfo_index=metainfo_index, + ep_rank=0, + ) + watchdog = coordinator.acquire_watchdog( + ep_size=1, + timeout_s=0.05, + poll_interval_s=0.005, + ) + + try: + coordinator.watch_collective(watchdog, "dispatch", None) + assert watchdog.wait_until_idle(timeout_s=1.0) + finally: + coordinator.release_watchdog(watchdog) + + +def test_unmonitored_coordinator_advances_shared_generation() -> None: + workspace_state: dict[str, object] = {} + workspace = torch.zeros((2, 32), dtype=torch.uint8) + metainfo = torch.tensor([0, 4, 12], dtype=torch.int64) + metainfo_index = { + "FLAG_VAL_OFFSET_INDEX": 0, + "DISPATCH_COMPLETION_FLAGS_OFFSET_INDEX": 1, + "COMBINE_COMPLETION_FLAGS_OFFSET_INDEX": 2, + } + coordinators = [ + AlltoAllWatchdogCoordinator( + workspace_state=workspace_state, + workspace=workspace, + metainfo=metainfo, + metainfo_index=metainfo_index, + ep_rank=0, + ) + for _ in range(2) + ] + events: list[AlltoAllWatchdogTimeout] = [] + watchdog = coordinators[0].acquire_watchdog( + ep_size=2, + timeout_s=0.02, + poll_interval_s=0.005, + on_timeout=events.append, + ) + + try: + coordinators[1].watch_collective(None, "dispatch", None) + coordinators[0].watch_collective(watchdog, "dispatch", None) + _wait_for(lambda: len(events) == 1) + assert events[0].expected_flag == 2 + finally: + coordinators[0].release_watchdog(watchdog) + + +def test_watchdog_rejects_active_mask_without_local_rank() -> None: + reader = FakeCompletionFlagReader(ep_size=4) + with AlltoAllWatchdog( + ep_size=4, + ep_rank=2, + completion_reader=reader, + timeout_s=0.1, + poll_interval_s=0.005, + ) as watchdog: + with pytest.raises(ValueError, match="local ep_rank"): + watchdog.watch(phase="dispatch", expected_flag=1, active_mask=0b1011) diff --git a/tests/unittest/_torch/modules/test_ep_metrics.py b/tests/unittest/_torch/modules/test_ep_metrics.py new file mode 100644 index 000000000000..25f95004f19c --- /dev/null +++ b/tests/unittest/_torch/modules/test_ep_metrics.py @@ -0,0 +1,219 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. +"""Unit tests for passive committed WideEP membership serialization.""" + +import copy +import threading +from types import SimpleNamespace +from typing import Any +from unittest.mock import MagicMock + +import pytest + +from tensorrt_llm._torch.modules.fused_moe.ep_group_health import ( + EP_GROUP_HEALTH_EXTRA_ATTR, + EP_GROUP_HEALTH_LEGACY_EXTRA_ATTR, + EPGroupHealth, +) +from tensorrt_llm._torch.modules.fused_moe.ep_metrics import ( + get_ep_group_health, + is_ep_group_health_registration_pending, + is_pending_ep_health_metrics, + pending_ep_health_metrics, + serialize_ep_health_metrics, + validate_ep_health_metrics_topology, +) + + +def test_serialize_initial_health() -> None: + health = EPGroupHealth(4) + + assert serialize_ep_health_metrics(health) == { + "sourceEpoch": health.source_epoch, + "worldSize": 4, + "activeCount": 4, + "failedRanks": [], + "generation": 0, + } + + +def test_deepcopy_preserves_shared_producer_identity() -> None: + health = EPGroupHealth(4) + + assert copy.deepcopy(health) is health + + +def test_serialization_reads_one_coherent_snapshot() -> None: + health: Any = SimpleNamespace( + moe_world_size=4, + source_epoch="source-a", + snapshot=MagicMock( + return_value=SimpleNamespace( + active_count=3, + failed_ranks=frozenset({2}), + generation=1, + ) + ), + ) + + assert serialize_ep_health_metrics(health) == { + "sourceEpoch": "source-a", + "worldSize": 4, + "activeCount": 3, + "failedRanks": [2], + "generation": 1, + } + health.snapshot.assert_called_once_with() + + +def test_serialize_committed_exclusion_and_reinclusion() -> None: + health = EPGroupHealth(72) + health.mark_failed(70) + health.mark_failed(3) + health.mark_failed(3) + + assert serialize_ep_health_metrics(health) == { + "sourceEpoch": health.source_epoch, + "worldSize": 72, + "activeCount": 70, + "failedRanks": [3, 70], + "generation": 2, + } + + health.mark_active(3) + assert serialize_ep_health_metrics(health) == { + "sourceEpoch": health.source_epoch, + "worldSize": 72, + "activeCount": 71, + "failedRanks": [70], + "generation": 3, + } + + +@pytest.mark.parametrize( + "key", + [EP_GROUP_HEALTH_EXTRA_ATTR, EP_GROUP_HEALTH_LEGACY_EXTRA_ATTR], +) +def test_get_ep_group_health_passively_reads_registered_tracker(key: str) -> None: + health = EPGroupHealth(4) + extra_attrs = {key: health} + model_config = SimpleNamespace( + extra_attrs=extra_attrs, + mapping=SimpleNamespace(moe_ep_size=4), + ) + + assert get_ep_group_health(model_config) is health + assert model_config.extra_attrs == extra_attrs + + +def test_get_ep_group_health_does_not_create_or_enable_tracker(monkeypatch) -> None: + monkeypatch.setenv("TRTLLM_ENABLE_WIDE_EP_FT", "1") + model_config = SimpleNamespace( + extra_attrs={}, + mapping=SimpleNamespace(moe_ep_size=8), + ) + + assert get_ep_group_health(model_config) is None + assert not is_ep_group_health_registration_pending(model_config) + assert model_config.extra_attrs == {} + + +def test_explicit_empty_registration_is_pending_without_mutation() -> None: + model_config = SimpleNamespace(extra_attrs={EP_GROUP_HEALTH_EXTRA_ATTR: None}) + + assert get_ep_group_health(model_config) is None + assert is_ep_group_health_registration_pending(model_config) + pending = pending_ep_health_metrics() + assert pending == {"status": "pending"} + assert is_pending_ep_health_metrics(pending) + assert not is_pending_ep_health_metrics({"status": "unsupported"}) + assert model_config.extra_attrs == {EP_GROUP_HEALTH_EXTRA_ATTR: None} + + +def test_get_ep_group_health_rejects_split_trackers() -> None: + model_config = SimpleNamespace( + extra_attrs={ + EP_GROUP_HEALTH_EXTRA_ATTR: EPGroupHealth(4), + EP_GROUP_HEALTH_LEGACY_EXTRA_ATTR: EPGroupHealth(4), + }, + mapping=SimpleNamespace(moe_ep_size=4), + ) + + with pytest.raises(ValueError, match="different trackers"): + get_ep_group_health(model_config) + + +def test_get_ep_group_health_rejects_invalid_shared_object() -> None: + wrong_type = SimpleNamespace( + extra_attrs={EP_GROUP_HEALTH_EXTRA_ATTR: object()}, + mapping=SimpleNamespace(moe_ep_size=8), + ) + with pytest.raises(TypeError, match="must be an EPGroupHealth"): + get_ep_group_health(wrong_type) + + wrong_size = SimpleNamespace( + extra_attrs={EP_GROUP_HEALTH_EXTRA_ATTR: EPGroupHealth(4)}, + mapping=SimpleNamespace(moe_ep_size=8), + ) + with pytest.raises(ValueError, match="world size does not match"): + get_ep_group_health(wrong_size) + + +def test_topology_validation_is_deferred_to_metric_read() -> None: + health = EPGroupHealth(4) + model_config = SimpleNamespace( + extra_attrs={EP_GROUP_HEALTH_EXTRA_ATTR: health}, + mapping=SimpleNamespace( + moe_ep_size=4, + moe_tp_size=1, + pp_size=2, + moe_cluster_size=1, + ), + ) + + assert get_ep_group_health(model_config) is health + with pytest.raises(NotImplementedError, match="multi-group MoE topologies"): + validate_ep_health_metrics_topology(model_config) + + +def test_concurrent_serialization_is_coherent() -> None: + health = EPGroupHealth(8) + start = threading.Barrier(2) + worker_errors: list[BaseException] = [] + + def toggle_rank() -> None: + try: + start.wait(timeout=10.0) + for _ in range(1_000): + health.mark_failed(3) + health.mark_active(3) + except BaseException as error: + worker_errors.append(error) + + thread = threading.Thread(target=toggle_rank) + thread.start() + start.wait(timeout=10.0) + + for _ in range(1_000): + stats = serialize_ep_health_metrics(health) + assert stats["activeCount"] == stats["worldSize"] - len(stats["failedRanks"]) + assert stats["failedRanks"] in ([], [3]) + assert bool(stats["generation"] % 2) is bool(stats["failedRanks"]) + + thread.join(timeout=10.0) + assert not thread.is_alive() + assert not worker_errors + assert health.generation == 2_000 + assert health.get_failed_ranks() == frozenset() diff --git a/tests/unittest/_torch/modules/test_pp_nccl_communicator.py b/tests/unittest/_torch/modules/test_pp_nccl_communicator.py new file mode 100644 index 000000000000..71c9006a8cf8 --- /dev/null +++ b/tests/unittest/_torch/modules/test_pp_nccl_communicator.py @@ -0,0 +1,782 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import threading + +import pytest + +from tensorrt_llm import _mnnvl_utils as mnnvl_module +from tensorrt_llm._torch.distributed import communicator as communicator_module + + +@pytest.fixture(autouse=True) +def _reset_pp_comm_lifecycle(monkeypatch): + previous_helix_comm = mnnvl_module.HelixCpMnnvlMemory.comm + previous_helix_topology = mnnvl_module.HelixCpMnnvlMemory.comm_topology + communicator_module._pp_comm = None + communicator_module._pp_comm_refcount = 0 + communicator_module._pp_comm_control_refcount = 0 + communicator_module._pp_comm_final_release_pending = False + monkeypatch.setattr(communicator_module, "NCCL_FAULT_TOLERANCE_ENABLED", True) + mnnvl_module.HelixCpMnnvlMemory.comm = None + mnnvl_module.HelixCpMnnvlMemory.comm_topology = None + yield + communicator_module._pp_comm = None + communicator_module._pp_comm_refcount = 0 + communicator_module._pp_comm_control_refcount = 0 + communicator_module._pp_comm_final_release_pending = False + mnnvl_module.HelixCpMnnvlMemory.comm = previous_helix_comm + mnnvl_module.HelixCpMnnvlMemory.comm_topology = previous_helix_topology + + +class _FakeNcclCommunicatorOp: + def __init__(self): + self.aborted = False + self.reinitialized_with = None + self.reinit_calls = [] + self.rendezvous_ids = [] + self.active_ranks = [0, 1, 2, 3] + self.async_error = "NCCL error: communicator was aborted after injected async failure" + self.received_from = None + + def abort(self): + self.aborted = True + + def abort_and_reinit(self, active_ranks, rendezvous_id): + self.reinitialized_with = list(active_ranks) + self.reinit_calls.append(list(active_ranks)) + self.rendezvous_ids.append(rendezvous_id) + self.active_ranks = list(active_ranks) + self.async_error = "" + + def get_async_error(self): + return self.async_error + + def get_active_ranks(self): + return list(self.active_ranks) + + def recv(self, tensor, src): + self.received_from = src + + def send(self, tensor, dest): + self.sent_to = dest + + +class _FakeMapping: + world_size = 4 + rank = 2 + pp_group = [0, 1, 2, 3] + cp_config = {} + + def has_cp_helix(self): + return False + + def prev_pp_rank(self): + return 1 + + def next_pp_rank(self): + return 3 + + +class _CompatibleMapping(_FakeMapping): + pass + + +class _HelixMapping(_FakeMapping): + cp_config = {"use_nccl_for_alltoall": False} + cp_group = [0, 2] + cp_rank = 1 + pp_rank = 0 + tp_size = 2 + tp_rank = 0 + + def has_cp_helix(self): + return True + + +class _OtherHelixMapping(_HelixMapping): + cp_group = [1, 2] + + +def _make_pp_comm(mapping=None): + pp_comm = object.__new__(communicator_module.PPCommNCCL) + pp_comm.mapping = mapping or _FakeMapping() + pp_comm._topology = pp_comm._mapping_topology(pp_comm.mapping) + pp_comm._active_ranks = tuple(range(pp_comm.mapping.world_size)) + pp_comm._reconfigure_lock = threading.Lock() + pp_comm._reconfigure_generation = 0 + pp_comm._completed_recovery = None + pp_comm.nccl_comm = _FakeNcclCommunicatorOp() + return pp_comm + + +def test_pp_wrapper_initializes_membership_from_persistent_native_comm(monkeypatch): + native_comm = _FakeNcclCommunicatorOp() + native_comm.active_ranks = [0, 2, 3] + monkeypatch.setattr( + communicator_module.torch.classes.trtllm, + "NcclCommunicatorOp", + lambda world_size, rank: native_comm, + raising=False, + ) + monkeypatch.setattr(communicator_module.torch.cuda, "Event", object) + monkeypatch.setattr(communicator_module.torch.cuda, "Stream", object) + + pp_comm = communicator_module.PPCommNCCL(_FakeMapping()) + + assert pp_comm.get_active_ranks() == [0, 2, 3] + + +def test_default_off_wrapper_does_not_query_native_membership(monkeypatch): + native_comm = _FakeNcclCommunicatorOp() + native_comm.get_active_ranks = lambda: (_ for _ in ()).throw( + AssertionError("default-off PP constructor queried native membership") + ) + monkeypatch.setattr(communicator_module, "NCCL_FAULT_TOLERANCE_ENABLED", False) + monkeypatch.setattr( + communicator_module.torch.classes.trtllm, + "NcclCommunicatorOp", + lambda world_size, rank: native_comm, + raising=False, + ) + monkeypatch.setattr(communicator_module.torch.cuda, "Event", object) + monkeypatch.setattr(communicator_module.torch.cuda, "Stream", object) + + pp_comm = communicator_module.PPCommNCCL(_FakeMapping()) + + assert pp_comm._active_ranks == (0, 1, 2, 3) + + +def test_pp_nccl_fault_tolerance_methods_forward_to_custom_class(): + pp_comm = _make_pp_comm() + + assert "NCCL error" in pp_comm.get_async_error() + pp_comm.abort_and_reinit((0, 2, 3)) + assert pp_comm.nccl_comm.reinitialized_with == [0, 2, 3] + assert pp_comm.nccl_comm.rendezvous_ids == [1] + assert pp_comm.get_async_error() == "" + assert pp_comm.get_active_ranks() == [0, 2, 3] + + pp_comm.abort() + assert pp_comm.nccl_comm.aborted + + +def test_module_level_pp_nccl_fault_tolerance_api(monkeypatch): + pp_comm = _make_pp_comm() + monkeypatch.setattr(communicator_module, "_pp_comm", pp_comm) + monkeypatch.setattr(communicator_module, "_pp_comm_refcount", 1) + + assert "NCCL error" in communicator_module.pp_comm_get_async_error() + communicator_module.pp_comm_abort_and_reinit([0, 2, 3], generation=3) + assert pp_comm.nccl_comm.reinitialized_with == [0, 2, 3] + assert pp_comm.nccl_comm.rendezvous_ids == [5] + assert communicator_module.pp_comm_get_async_error() == "" + assert communicator_module.pp_comm_get_active_ranks() == [0, 2, 3] + + communicator_module.pp_comm_abort() + assert pp_comm.nccl_comm.aborted + + +def test_module_level_pp_nccl_api_requires_nccl_backend(monkeypatch): + monkeypatch.setattr(communicator_module, "_pp_comm", object()) + + with pytest.raises(RuntimeError, match="NCCL error: PP communicator is not initialized"): + communicator_module.pp_comm_get_async_error() + + +def test_reinit_rejects_stale_default_peer_and_accepts_explicit_remap(): + pp_comm = _make_pp_comm() + pp_comm.abort_and_reinit([0, 2, 3]) + + # Static Mapping.prev_pp_rank() names failed rank 1. Silently skipping it + # would bypass that stage's model layers, so default routing must fail. + with pytest.raises(RuntimeError, match="peer world rank 1 is not active"): + pp_comm.recv(object()) + + # A higher-level layer/topology reconstruction may provide an explicit + # replacement peer after it has reassigned the failed stage. + pp_comm.recv(object(), src=0) + assert pp_comm.nccl_comm.received_from == 0 + assert pp_comm._required_pp_peer(next_peer=True) == 3 + + +def test_explicit_peer_must_belong_to_the_callers_pp_lane(): + class LaneMapping: + world_size = 8 + rank = 2 + pp_group = [2, 6] + cp_config = {} + + def has_cp_helix(self): + return False + + def prev_pp_rank(self): + return 6 + + def next_pp_rank(self): + return 6 + + pp_comm = _make_pp_comm(LaneMapping()) + + with pytest.raises(RuntimeError, match="not in this rank's PP group"): + pp_comm.recv(object(), src=3) + + pp_comm.recv(object(), src=6) + assert pp_comm.nccl_comm.received_from == 6 + + +def test_default_off_preserves_legacy_peer_routing(monkeypatch): + pp_comm = _make_pp_comm() + monkeypatch.setattr(communicator_module, "NCCL_FAULT_TOLERANCE_ENABLED", False) + monkeypatch.setattr( + communicator_module.torch.cuda, + "is_current_stream_capturing", + lambda: True, + ) + + # Before FT wiring, explicit peers were forwarded to the native wrapper + # without a Python PP-lane scan. Preserve that behavior and cost when the + # startup mode is disabled. + tensor = object() + pp_comm.send(tensor, dest=7) + pp_comm.recv(tensor, src=7) + + assert pp_comm.nccl_comm.sent_to == 7 + assert pp_comm.nccl_comm.received_from == 7 + + +def test_reinit_rejects_reactivating_removed_pp_rank(): + pp_comm = _make_pp_comm() + pp_comm.abort_and_reinit([0, 2, 3]) + + with pytest.raises(ValueError, match="reactivate"): + pp_comm.abort_and_reinit([0, 1, 2, 3]) + + +def test_duplicate_reinit_target_is_idempotent(): + pp_comm = _make_pp_comm() + + pp_comm.abort_and_reinit([0, 2, 3], generation=1) + pp_comm.abort_and_reinit([3, 2, 0], generation=1) + + assert pp_comm.nccl_comm.reinit_calls == [[0, 2, 3]] + assert pp_comm.nccl_comm.rendezvous_ids == [3] + assert pp_comm.get_active_ranks() == [0, 2, 3] + + +def test_shared_generation_forces_same_membership_rebuild(): + pp_comm = _make_pp_comm() + + pp_comm.abort_and_reinit([0, 1, 2, 3], generation=1) + pp_comm.abort_and_reinit([0, 1, 2, 3], generation=1) + pp_comm.abort_and_reinit([0, 1, 2, 3], generation=2) + + assert pp_comm.nccl_comm.reinit_calls == [[0, 1, 2, 3], [0, 1, 2, 3]] + assert pp_comm.nccl_comm.rendezvous_ids == [3, 4] + assert pp_comm.get_async_error() == "" + + +def test_same_membership_recovery_requires_shared_generation(): + pp_comm = _make_pp_comm() + + with pytest.raises(ValueError, match="generation is required"): + pp_comm.abort_and_reinit([0, 1, 2, 3]) + + assert pp_comm.nccl_comm.reinit_calls == [] + + +@pytest.mark.parametrize("generation", [-1, 1.5, (1 << 63) - 2]) +def test_recovery_generation_must_fit_reserved_torch_int_range(generation): + pp_comm = _make_pp_comm() + + with pytest.raises(ValueError, match="generation must be a nonnegative integer"): + pp_comm.abort_and_reinit([0, 2, 3], generation=generation) + + assert pp_comm.nccl_comm.reinit_calls == [] + + +def test_generation_rejects_conflicting_pp_recovery_targets(): + pp_comm = _make_pp_comm() + pp_comm.abort_and_reinit([0, 2, 3], generation=7) + + with pytest.raises(RuntimeError, match="conflicting.*generation 7"): + pp_comm.abort_and_reinit([0, 2], generation=7) + + assert pp_comm.nccl_comm.reinit_calls == [[0, 2, 3]] + + +def test_concurrent_nested_reinit_requests_commit_newest_membership(): + pp_comm = _make_pp_comm() + first_entered_native = threading.Event() + second_entered_native = threading.Event() + release_first = threading.Event() + calls = [] + + def blocking_reinit(active_ranks, rendezvous_id): + active_ranks = list(active_ranks) + calls.append((active_ranks, rendezvous_id)) + pp_comm.nccl_comm.active_ranks = active_ranks + if active_ranks == [0, 2, 3]: + first_entered_native.set() + assert release_first.wait(timeout=5) + else: + second_entered_native.set() + + pp_comm.nccl_comm.abort_and_reinit = blocking_reinit + errors = [] + + def reconfigure(active_ranks): + try: + pp_comm.abort_and_reinit(active_ranks) + except Exception as error: # pragma: no cover - assertion reports details + errors.append(error) + + first = threading.Thread(target=reconfigure, args=([0, 2, 3],)) + second = threading.Thread(target=reconfigure, args=([0, 2],)) + first.start() + assert first_entered_native.wait(timeout=5) + second.start() + + # The second callback must not enter the native rendezvous while the first + # callback is between native completion and Python membership publication. + serialized = not second_entered_native.wait(timeout=0.1) + release_first.set() + first.join(timeout=5) + second.join(timeout=5) + + assert serialized + assert not first.is_alive() + assert not second.is_alive() + assert errors == [] + assert calls == [([0, 2, 3], 1), ([0, 2], 1)] + assert pp_comm.get_active_ranks() == [0, 2] + + +def test_abort_waiting_for_reinit_does_not_abort_replacement(): + pp_comm = _make_pp_comm() + abort_generation = pp_comm._reconfigure_generation + entered_reinit = threading.Event() + release_reinit = threading.Event() + abort_started = threading.Event() + + def blocking_reinit(active_ranks, rendezvous_id): + assert rendezvous_id == 1 + entered_reinit.set() + assert release_reinit.wait(timeout=5) + pp_comm.nccl_comm.active_ranks = list(active_ranks) + pp_comm.nccl_comm.async_error = "" + + pp_comm.nccl_comm.abort_and_reinit = blocking_reinit + errors = [] + + def reconfigure(): + try: + pp_comm.abort_and_reinit([0, 2, 3]) + except Exception as error: # pragma: no cover - assertion reports details + errors.append(error) + + def abort(): + abort_started.set() + try: + pp_comm.abort(abort_generation) + except Exception as error: # pragma: no cover - assertion reports details + errors.append(error) + + rebuild_thread = threading.Thread(target=reconfigure) + abort_thread = threading.Thread(target=abort) + rebuild_thread.start() + assert entered_reinit.wait(timeout=5) + abort_thread.start() + assert abort_started.wait(timeout=5) + release_reinit.set() + rebuild_thread.join(timeout=5) + abort_thread.join(timeout=5) + + assert not rebuild_thread.is_alive() + assert not abort_thread.is_alive() + assert errors == [] + assert not pp_comm.nccl_comm.aborted + assert pp_comm.get_active_ranks() == [0, 2, 3] + + # A genuinely new abort request after the rebuild must still take effect. + pp_comm.abort() + assert pp_comm.nccl_comm.aborted + + +def test_module_abort_waiting_for_reinit_does_not_abort_replacement(monkeypatch): + pp_comm = _make_pp_comm() + entered_reinit = threading.Event() + abort_captured_generation = threading.Event() + release_reinit = threading.Event() + errors = [] + + def blocking_reinit(active_ranks, rendezvous_id): + assert rendezvous_id == 1 + entered_reinit.set() + assert release_reinit.wait(timeout=5) + pp_comm.nccl_comm.active_ranks = list(active_ranks) + pp_comm.nccl_comm.async_error = "" + + pp_comm.nccl_comm.abort_and_reinit = blocking_reinit + original_abort = pp_comm.abort + + def observed_abort(generation=None): + # The module wrapper has already captured the epoch before it invokes + # the class method. Do not let the rebuild finish until that point. + abort_captured_generation.set() + return original_abort(generation) + + pp_comm.abort = observed_abort + monkeypatch.setattr(communicator_module, "_pp_comm", pp_comm) + monkeypatch.setattr(communicator_module, "_pp_comm_refcount", 1) + + def reconfigure(): + try: + communicator_module.pp_comm_abort_and_reinit([0, 2, 3]) + except Exception as error: # pragma: no cover - assertion reports details + errors.append(error) + + def abort(): + try: + communicator_module.pp_comm_abort() + except Exception as error: # pragma: no cover - assertion reports details + errors.append(error) + + rebuild_thread = threading.Thread(target=reconfigure) + abort_thread = threading.Thread(target=abort) + rebuild_thread.start() + assert entered_reinit.wait(timeout=5) + abort_thread.start() + assert abort_captured_generation.wait(timeout=5) + release_reinit.set() + rebuild_thread.join(timeout=5) + abort_thread.join(timeout=5) + + assert not rebuild_thread.is_alive() + assert not abort_thread.is_alive() + assert errors == [] + assert not pp_comm.nccl_comm.aborted + assert pp_comm.get_active_ranks() == [0, 2, 3] + + communicator_module.pp_comm_abort() + assert pp_comm.nccl_comm.aborted + + +def test_module_control_call_blocks_final_release(monkeypatch): + pp_comm = _make_pp_comm() + entered_query = threading.Event() + release_query = threading.Event() + final_release_completed = threading.Event() + errors = [] + + def blocking_get_async_error(): + entered_query.set() + assert release_query.wait(timeout=5) + return "" + + pp_comm.nccl_comm.get_async_error = blocking_get_async_error + monkeypatch.setattr(communicator_module, "_pp_comm", pp_comm) + monkeypatch.setattr(communicator_module, "_pp_comm_refcount", 1) + + def query(): + try: + communicator_module.pp_comm_get_async_error() + except Exception as error: # pragma: no cover - assertion reports details + errors.append(error) + + def final_release(): + try: + communicator_module.release_pp_comm() + final_release_completed.set() + except Exception as error: # pragma: no cover - assertion reports details + errors.append(error) + + query_thread = threading.Thread(target=query) + release_thread = threading.Thread(target=final_release) + query_thread.start() + assert entered_query.wait(timeout=5) + release_thread.start() + + # The control call owns a transient reference under _pp_comm_lock, so a + # final engine release cannot complete its ownership transition yet. + assert not final_release_completed.wait(timeout=0.1) + release_query.set() + query_thread.join(timeout=5) + release_thread.join(timeout=5) + + assert not query_thread.is_alive() + assert not release_thread.is_alive() + assert errors == [] + assert final_release_completed.is_set() + assert communicator_module._pp_comm is pp_comm + assert communicator_module._pp_comm_refcount == 0 + + +def test_init_pp_comm_reuses_compatible_nccl_communicator(monkeypatch): + created = [] + + class FakePPComm: + def __init__(self, mapping): + self.topology = (mapping.world_size, mapping.rank, tuple(mapping.pp_group)) + created.append(self) + + def is_compatible(self, mapping): + return self.topology == (mapping.world_size, mapping.rank, tuple(mapping.pp_group)) + + monkeypatch.setattr(communicator_module, "PPCommNCCL", FakePPComm) + monkeypatch.setattr(communicator_module, "_pp_comm", None) + monkeypatch.setattr(communicator_module, "mpi_disabled", lambda: False) + monkeypatch.setattr(communicator_module, "init_helix_cp_comm", lambda mapping: None) + + communicator_module.init_pp_comm(_FakeMapping()) + first = communicator_module._pp_comm + # A target and draft engine construct equivalent Mapping objects. They + # must share the sole process-local native PP communicator. + communicator_module.init_pp_comm(_CompatibleMapping()) + + assert communicator_module._pp_comm is first + assert len(created) == 1 + assert communicator_module._pp_comm_refcount == 2 + + communicator_module.release_pp_comm() + assert communicator_module._pp_comm is first + communicator_module.release_pp_comm() + assert communicator_module._pp_comm is first + assert communicator_module._pp_comm_refcount == 0 + + +def test_ft_wrapper_churn_preserves_completed_recovery_generation(monkeypatch): + pp_comm = _make_pp_comm() + monkeypatch.setattr(communicator_module, "_pp_comm", pp_comm) + monkeypatch.setattr(communicator_module, "_pp_comm_refcount", 1) + monkeypatch.setattr(communicator_module, "mpi_disabled", lambda: False) + monkeypatch.setattr(communicator_module, "init_helix_cp_comm", lambda mapping: None) + + communicator_module.pp_comm_abort_and_reinit([0, 2, 3], generation=7) + communicator_module.release_pp_comm() + assert communicator_module._pp_comm is pp_comm + assert communicator_module._pp_comm_refcount == 0 + + communicator_module.init_pp_comm(_CompatibleMapping()) + communicator_module.pp_comm_abort_and_reinit([0, 2, 3], generation=7) + + assert communicator_module._pp_comm is pp_comm + assert pp_comm.nccl_comm.reinit_calls == [[0, 2, 3]] + assert pp_comm.nccl_comm.rendezvous_ids == [9] + + +def test_init_pp_comm_rejects_incompatible_live_topology(monkeypatch): + existing = _make_pp_comm() + monkeypatch.setattr(communicator_module, "_pp_comm", existing) + monkeypatch.setattr(communicator_module, "mpi_disabled", lambda: False) + monkeypatch.setattr(communicator_module, "init_helix_cp_comm", lambda mapping: None) + + class OtherMapping(_FakeMapping): + pp_group = [0, 2] + + with pytest.raises(RuntimeError, match="different topology"): + communicator_module.init_pp_comm(OtherMapping()) + + +def test_init_pp_comm_rejects_incompatible_live_helix_topology(monkeypatch): + existing = _make_pp_comm(_HelixMapping()) + monkeypatch.setattr(communicator_module, "_pp_comm", existing) + monkeypatch.setattr(communicator_module, "mpi_disabled", lambda: False) + + with pytest.raises(RuntimeError, match="Helix MNNVL CP topology"): + communicator_module.init_pp_comm(_OtherHelixMapping()) + + +def test_helix_mnnvl_comm_rejects_incompatible_cached_topology(monkeypatch): + split_calls = [] + comm = object() + + class FakeMpiComm: + def Split(self, color, key): + split_calls.append((color, key)) + return comm + + monkeypatch.setattr(mnnvl_module, "mpi_comm", lambda: FakeMpiComm()) + + assert mnnvl_module.HelixCpMnnvlMemory.get_comm(_HelixMapping()) is comm + assert mnnvl_module.HelixCpMnnvlMemory.get_comm(_HelixMapping()) is comm + with pytest.raises(RuntimeError, match="cannot reuse.*incompatible topology"): + mnnvl_module.HelixCpMnnvlMemory.get_comm(_OtherHelixMapping()) + + assert split_calls == [(0, 1)] + + +def test_final_release_preserves_helix_comm_topology(monkeypatch): + original_pp_comm_class = communicator_module.PPCommNCCL + split_calls = [] + + class FakeMpiComm: + def Split(self, color, key): + split_calls.append((color, key)) + return object() + + class FakePPComm: + def __init__(self, mapping): + self.topology = original_pp_comm_class._mapping_topology(mapping) + + def is_compatible(self, mapping): + return self.topology == original_pp_comm_class._mapping_topology(mapping) + + monkeypatch.setattr(mnnvl_module, "mpi_comm", lambda: FakeMpiComm()) + monkeypatch.setattr(communicator_module, "PPCommNCCL", FakePPComm) + monkeypatch.setattr(communicator_module, "NCCL_FAULT_TOLERANCE_ENABLED", False) + monkeypatch.setattr(communicator_module, "mpi_disabled", lambda: False) + + # Reacquiring the ordinary PP wrapper after its final release is valid when + # the persistent Helix MPI communicator has the same CP rank topology. + communicator_module.init_pp_comm(_HelixMapping()) + communicator_module.release_pp_comm() + communicator_module.init_pp_comm(_HelixMapping()) + communicator_module.release_pp_comm() + + # A later engine may create a fresh PP wrapper, but it cannot reinterpret + # the process-lifetime Helix MPI communicator with a different CP group. + with pytest.raises(RuntimeError, match="cannot reuse.*incompatible topology"): + communicator_module.init_pp_comm(_OtherHelixMapping()) + + assert split_calls == [(0, 1)] + assert communicator_module._pp_comm is None + assert communicator_module._pp_comm_refcount == 0 + + +def test_default_off_final_release_allows_a_later_incompatible_topology(monkeypatch): + created = [] + + class FakePPComm: + def __init__(self, mapping): + self.topology = (mapping.world_size, mapping.rank, tuple(mapping.pp_group)) + created.append(self.topology) + + def is_compatible(self, mapping): + return self.topology == (mapping.world_size, mapping.rank, tuple(mapping.pp_group)) + + class OtherMapping(_FakeMapping): + pp_group = [0, 2] + + monkeypatch.setattr(communicator_module, "PPCommNCCL", FakePPComm) + monkeypatch.setattr(communicator_module, "NCCL_FAULT_TOLERANCE_ENABLED", False) + monkeypatch.setattr(communicator_module, "mpi_disabled", lambda: False) + monkeypatch.setattr(communicator_module, "init_helix_cp_comm", lambda mapping: None) + + communicator_module.init_pp_comm(_FakeMapping()) + communicator_module.release_pp_comm() + assert communicator_module._pp_comm is None + communicator_module.init_pp_comm(OtherMapping()) + + assert len(created) == 2 + assert created[0] != created[1] + + +def test_model_engine_cleanup_releases_pp_before_reporting_ub_failure(monkeypatch): + from types import SimpleNamespace + + from tensorrt_llm._torch.pyexecutor import model_engine + + engine = object.__new__(model_engine.PyTorchModelEngine) + engine._cleanup_done = False + engine._pp_comm_acquired = True + engine.model_loader = None + engine.model = object() + engine.input_processor = object() + engine.input_processor_with_hash = object() + engine.ub_buffers = [SimpleNamespace(addr=123)] + engine._release_cuda_graphs = lambda: None + + released = [] + gc_calls = [] + monkeypatch.setattr(model_engine, "release_pp_comm", lambda: released.append(True)) + monkeypatch.setattr(model_engine, "release_gc", lambda: gc_calls.append(True)) + monkeypatch.setattr( + model_engine.ub, + "ub_deallocate", + lambda addr: (_ for _ in ()).throw(RuntimeError("injected UB failure")), + ) + + with pytest.raises(RuntimeError, match="userbuffers"): + engine.cleanup() + + assert released == [True] + assert gc_calls == [True] + assert not engine._pp_comm_acquired + assert not engine._cleanup_done + assert [buffer.addr for buffer in engine.ub_buffers] == [123] + + monkeypatch.setattr(model_engine.ub, "ub_deallocate", lambda addr: None) + engine.cleanup() + assert released == [True] + assert engine._cleanup_done + assert engine.ub_buffers is None + + +def test_ft_mode_uses_engine_local_eager_graph_configuration(monkeypatch): + from tensorrt_llm._torch.pyexecutor import model_engine + + class FakeTorchCompileConfig: + def __init__(self, enable_piecewise_cuda_graph): + self.enable_piecewise_cuda_graph = enable_piecewise_cuda_graph + + def model_copy(self, *, update): + copied = FakeTorchCompileConfig(self.enable_piecewise_cuda_graph) + for name, value in update.items(): + setattr(copied, name, value) + return copied + + class FakeLlmArgs: + def __init__(self, cuda_graph_config, torch_compile_config): + self.cuda_graph_config = cuda_graph_config + self.torch_compile_config = torch_compile_config + + def model_copy(self, *, update): + copied = FakeLlmArgs(self.cuda_graph_config, self.torch_compile_config) + for name, value in update.items(): + setattr(copied, name, value) + return copied + + original = FakeLlmArgs(object(), FakeTorchCompileConfig(True)) + monkeypatch.setenv("TLLM_FAULT_TOLERANCE_MODE", "1") + + effective = model_engine._force_eager_mode_for_nccl_fault_tolerance(original) + + assert effective is not original + assert effective.cuda_graph_config is None + assert not effective.torch_compile_config.enable_piecewise_cuda_graph + # The caller-owned arguments remain reusable by a non-FT engine. + assert original.cuda_graph_config is not None + assert original.torch_compile_config.enable_piecewise_cuda_graph + + +def test_ft_mode_uses_eager_allreduce_autotuning(monkeypatch): + from tensorrt_llm._torch.custom_ops import torch_custom_ops + + monkeypatch.delenv("TLLM_FAULT_TOLERANCE_MODE", raising=False) + assert torch_custom_ops._use_cuda_graph_for_allreduce_tuning() + + monkeypatch.setenv("TLLM_FAULT_TOLERANCE_MODE", "1") + assert not torch_custom_ops._use_cuda_graph_for_allreduce_tuning() + + +def test_partially_constructed_engine_cleanup_releases_pp(monkeypatch): + from tensorrt_llm._torch.pyexecutor import model_engine + + # Simulate a constructor failure immediately after init_pp_comm(), before + # model/model_loader/input-processor attributes have been populated. + engine = object.__new__(model_engine.PyTorchModelEngine) + engine._cleanup_done = False + engine._pp_comm_acquired = True + + released = [] + gc_calls = [] + monkeypatch.setattr(model_engine, "release_pp_comm", lambda: released.append(True)) + monkeypatch.setattr(model_engine, "release_gc", lambda: gc_calls.append(True)) + + engine.cleanup() + + assert released == [True] + assert gc_calls == [True] + assert not engine._pp_comm_acquired + assert engine._cleanup_done diff --git a/tests/unittest/_torch/pyexecutor/_ep_failure_broadcast_mpi_worker.py b/tests/unittest/_torch/pyexecutor/_ep_failure_broadcast_mpi_worker.py new file mode 100644 index 000000000000..d9c9bd732443 --- /dev/null +++ b/tests/unittest/_torch/pyexecutor/_ep_failure_broadcast_mpi_worker.py @@ -0,0 +1,626 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. +"""Logical WideEP failure-injection worker using a real MPI transport. + +The smoke keeps every process alive so ``COMM_WORLD`` can coordinate portable +results. Actual peer death and MPI runtime error classification are outside its +scope because common launchers terminate the whole job when a worker exits +without ``MPI_Finalize``. +""" + +import atexit +import os +import sys +import time +import traceback +from types import SimpleNamespace + +import numpy as np +from mpi4py import MPI + +from tensorrt_llm._torch.pyexecutor.ep_failure_broadcast import ( + DetectedRankState, + MpiFtSubcomm, + MpiFtSubcommConfig, +) +from tensorrt_llm.mapping import MpiTopology + +_SKIP_MARKER = "WIDEEP_MPI_SMOKE_SKIP:" +_PROPAGATION_MARKER = "WIDEEP_MPI_PROPAGATION" +_INVALID_TOPOLOGY_MARKER = "WIDEEP_MPI_INVALID_TOPOLOGY_OK" +_HEALTHY_LIFECYCLE_MARKER = "WIDEEP_MPI_HEALTHY_LIFECYCLE_OK" +_ABORT_WORLD_MARKER = "WIDEEP_MPI_ABORT_WORLD_OK" +_TERMINAL_READY_MARKER = "WIDEEP_MPI_TERMINAL_READY" +_TERMINAL_COMPLETE_MARKER = "WIDEEP_MPI_TERMINAL_COMPLETE" +_TERMINAL_DONE_MARKER = "WIDEEP_MPI_TERMINAL_DONE" +_TERMINAL_ERROR_MARKER = "WIDEEP_MPI_TERMINAL_ERROR" +_TERMINAL_ATEXIT_MARKER = "WIDEEP_MPI_TERMINAL_ATEXIT_RAN" +_PROPAGATION_TARGET_SEC = 0.1 +_CONVERGENCE_TIMEOUT_SEC = 2.0 +_ALLOW_SKIP_ENV = "TLLM_ALLOW_MPI_FT_SMOKE_SKIP" +_HEALTHY_MODE = "healthy" +_TERMINAL_MODE = "terminal" +_WORKER_SKIPPED = False + + +class _NonUlfmComm: + """Typed MPI communicator view that makes smoke behavior independent of ULFM.""" + + def __init__(self, comm: MPI.Intracomm) -> None: + self._comm = comm + self.abort_calls = 0 + + def Get_rank(self) -> int: + return self._comm.Get_rank() + + def Get_size(self) -> int: + return self._comm.Get_size() + + def Set_errhandler(self, errhandler: MPI.Errhandler) -> None: + self._comm.Set_errhandler(errhandler) + + def Irecv(self, buffer: np.ndarray, source: int, tag: int) -> MPI.Request: + return self._comm.Irecv(buffer, source=source, tag=tag) + + def Isend(self, buffer: np.ndarray, dest: int, tag: int) -> MPI.Request: + return self._comm.Isend(buffer, dest=dest, tag=tag) + + def Abort(self, errorcode: int) -> None: + self.abort_calls += 1 + self._comm.Abort(errorcode) + + +def _mapping(world_size: int, rank: int) -> MpiTopology: + return MpiTopology( + world_size=world_size, + rank=rank, + tp_size=world_size, + moe_tp_size=1, + moe_ep_size=world_size, + ) + + +def _unsupported_environment(message: str, *, rank: int) -> int: + """Skip optional local runs, but fail a designated required smoke run.""" + global _WORKER_SKIPPED + _WORKER_SKIPPED = True + required = os.environ.get(_ALLOW_SKIP_ENV) != "1" + if rank == 0: + if required: + print(f"WideEP MPI FT smoke requirement failed: {message}", file=sys.stderr) + else: + print(f"{_SKIP_MARKER} {message}") + return 2 if required else 0 + + +def _synchronize_phase_error( + world: MPI.Intracomm, + phase: str, + local_error: str | None, +) -> str | None: + """Return the same rank-attributed phase failure before the next collective.""" + errors = world.allgather(local_error) + failures = [(rank, error) for rank, error in enumerate(errors) if error is not None] + if not failures: + return None + details = "\n".join(f"rank {rank}:\n{error}" for rank, error in failures) + return f"{phase} failed before the next MPI phase:\n{details}" + + +def _run_invalid_topology_smoke(world: MPI.Intracomm) -> str | None: + """Prove a rank-local topology error is reconciled before ``Split``.""" + rank = world.Get_rank() + world_size = world.Get_size() + ep_group = tuple(range(world_size)) + if rank == 0: + ep_group = ep_group[:-1] + mapping = SimpleNamespace( + world_size=world_size, + rank=rank, + moe_ep_size=world_size, + moe_ep_rank=rank, + moe_ep_group=ep_group, + ) + caught_error: str | None = None + unexpected_broadcaster: MpiFtSubcomm | None = None + try: + unexpected_broadcaster = MpiFtSubcomm( + mapping, + DetectedRankState(world_size), + MpiFtSubcommConfig(startup_timeout_sec=5.0, stop_timeout_sec=5.0), + ) + except Exception as error: + caught_error = f"{type(error).__name__}: {error}" + + caught_errors = world.allgather(caught_error) + expected_fragment = "startup validation failed on parent rank 0" + if any(error is None for error in caught_errors): + return f"invalid topology unexpectedly constructed on some ranks: {caught_errors}" + if len(set(caught_errors)) != 1 or expected_fragment not in caught_errors[0]: + return f"invalid topology errors were not identical across ranks: {caught_errors}" + + # This barrier is intentionally after the failed constructor. It proves the + # parent communicator did not strand any rank in MPI_Comm_split. + world.Barrier() + if rank == 0: + print(_INVALID_TOPOLOGY_MARKER, flush=True) + if unexpected_broadcaster is not None: + return "invalid topology unexpectedly returned a broadcaster" + return None + + +def _run_healthy_lifecycle_smoke(world: MPI.Intracomm) -> str | None: + """Exercise production construction, progress startup, and clean shutdown.""" + rank = world.Get_rank() + world_size = world.Get_size() + mapping = _mapping(world_size, rank) + broadcaster: MpiFtSubcomm | None = None + error: str | None = None + + try: + config = MpiFtSubcommConfig( + startup_timeout_sec=5.0, + stop_timeout_sec=5.0, + reconcile_timeout_sec=5.0, + ) + # This path obtains TRT-LLM's pkl5 parent through mpi_comm(), performs + # collective validation and Split, and owns MPI progress directly. + broadcaster = MpiFtSubcomm( + mapping, + DetectedRankState(world_size), + config, + ) + ft_comm = broadcaster._comm + if ft_comm.Get_rank() != rank or ft_comm.Get_size() != world_size: + raise AssertionError( + "MPI_Comm_split returned an unexpected FT communicator: " + f"rank={ft_comm.Get_rank()}, size={ft_comm.Get_size()}" + ) + if ft_comm.Get_errhandler() != MPI.ERRORS_RETURN: + raise AssertionError("FT communicator does not use MPI.ERRORS_RETURN") + broadcaster.start() + broadcaster.stop() + if broadcaster.last_error is not None: + raise AssertionError( + "Production-construction progress failed during healthy shutdown: " + f"{broadcaster.last_error}" + ) + except Exception: + error = traceback.format_exc() + finally: + if broadcaster is not None: + try: + broadcaster.stop(timeout=5.0) + if broadcaster.last_error is not None: + raise AssertionError( + "MPI progress failed during healthy lifecycle smoke: " + f"{broadcaster.last_error}" + ) + except Exception: + cleanup_error = traceback.format_exc() + error = f"{error or ''}\ncleanup failure:\n{cleanup_error}" + # Production-created FT communicators are intentionally retained for + # process lifetime. Freeing one here would violate the same ownership + # contract this smoke is meant to exercise. + + return error + + +def _run_failure_propagation_smoke(world: MPI.Intracomm) -> str | None: + """Exercise logical failure fanout and sticky state over real MPI traffic.""" + rank = world.Get_rank() + world_size = world.Get_size() + mapping = _mapping(world_size, rank) + health = DetectedRankState(world_size) + owner: MpiFtSubcomm | None = None + broadcaster: MpiFtSubcomm | None = None + detector_rank = 0 + failure_rank = world_size - 1 + error: str | None = None + propagation_converged = False + + setup_error: str | None = None + try: + config = MpiFtSubcommConfig( + startup_timeout_sec=5.0, + stop_timeout_sec=5.0, + reconcile_timeout_sec=5.0, + ) + # Create the real FT communicator collectively, then inject a non-ULFM + # view so this smoke deterministically exercises typed Isend/Irecv+Test. + owner = MpiFtSubcomm( + mapping, + DetectedRankState(world_size), + config, + ) + ft_comm = owner._comm + owner.stop() + broadcaster = MpiFtSubcomm( + mapping, + health, + config, + comm=_NonUlfmComm(ft_comm), + ) + broadcaster.start() + except Exception: + setup_error = traceback.format_exc() + error = _synchronize_phase_error(world, "failure-propagation setup", setup_error) + + detector_elapsed_sec: float | None = None + local_converged = False + if error is None: + phase_errors: list[str] = [] + try: + # Every broadcaster is running before rank 0 starts the latency + # clock. This synchronization is deliberately outside the measured + # interval. + world.Barrier() + except Exception: + phase_errors.append(f"readiness:\n{traceback.format_exc()}") + + trigger_error: str | None = None + propagation_start: float | None = None + if not phase_errors and rank == detector_rank: + propagation_start = time.monotonic() + try: + assert broadcaster is not None + broadcaster.report_detected_failure(failure_rank) + except Exception: + trigger_error = traceback.format_exc() + + observation_error: str | None = None + try: + # The victim process remains alive only so the smoke can coordinate + # its result on COMM_WORLD. Survivors exclude its logical rank and + # exercise the same detection/reconciliation fanout as production. + # This smoke does not model recovery commit or request resumption. + assert broadcaster is not None + deadline = time.monotonic() + _CONVERGENCE_TIMEOUT_SEC + while not phase_errors and rank != failure_rank and time.monotonic() < deadline: + failure_observed = not health.is_active(failure_rank) + failure_reconciled = broadcaster.failure_detection_is_reconciled(failure_rank) + if failure_observed and failure_reconciled: + if rank == detector_rank: + assert propagation_start is not None + detector_elapsed_sec = time.monotonic() - propagation_start + break + time.sleep(0.005) + local_converged = rank == failure_rank or ( + not health.is_active(failure_rank) + and broadcaster.failure_detection_is_reconciled(failure_rank) + and broadcaster.detected_state_is_reconciled() + and broadcaster.world_is_poisoned() + ) + except Exception: + observation_error = traceback.format_exc() + + if trigger_error is not None: + phase_errors.append(f"trigger:\n{trigger_error}") + if observation_error is not None: + phase_errors.append(f"observation:\n{observation_error}") + error = _synchronize_phase_error( + world, + "failure-propagation trigger/observation", + "\n".join(phase_errors) or None, + ) + + if error is None: + all_converged = world.allreduce(local_converged, op=MPI.LAND) + # Share only rank 0's completed duration after the timed interval; raw + # monotonic timestamps never cross rank or host boundaries. + detector_elapsed_sec = world.bcast(detector_elapsed_sec, root=detector_rank) + if all_converged and detector_elapsed_sec is not None: + propagation_converged = True + else: + error = ( + f"failure rank {failure_rank} did not converge within " + f"{_CONVERGENCE_TIMEOUT_SEC:.0f}s" + ) + if rank == detector_rank and detector_elapsed_sec is not None: + elapsed_ms = detector_elapsed_sec * 1000.0 + target_ms = _PROPAGATION_TARGET_SEC * 1000.0 + print( + f"{_PROPAGATION_MARKER} world_size={world_size} " + f"elapsed_ms={elapsed_ms!r} target_ms={target_ms!r} " + f"target_met={elapsed_ms < target_ms}", + flush=True, + ) + + stop_succeeded = False + if broadcaster is not None: + # Deliberately do not reactivate failure_rank. Survivors must take the + # poisoned shutdown path and retain their unmatched receive requests. + try: + broadcaster.stop(timeout=5.0) + if broadcaster.last_error is not None: + raise AssertionError( + f"MPI progress failed during propagation smoke: {broadcaster.last_error}" + ) + stop_succeeded = True + except Exception: + cleanup_error = traceback.format_exc() + error = f"{error or ''}\ncleanup failure:\n{cleanup_error}" + if propagation_converged and stop_succeeded: + try: + if rank == failure_rank: + if broadcaster.world_is_poisoned(): + raise AssertionError("logical victim unexpectedly entered poisoned state") + if broadcaster._retained_requests: + raise AssertionError("logical victim retained healthy MPI requests") + else: + if not broadcaster.world_is_poisoned(): + raise AssertionError("survivor lost sticky poisoned state during stop") + if not broadcaster._retained_requests: + raise AssertionError("survivor did not retain poisoned MPI requests") + except Exception: + sticky_error = traceback.format_exc() + error = f"{error or ''}\nsticky-poison failure:\n{sticky_error}" + if owner is not None: + try: + owner.stop(timeout=5.0) + except Exception: + cleanup_error = traceback.format_exc() + error = f"{error or ''}\nowner cleanup failure:\n{cleanup_error}" + + return error + + +def _write_terminal_manifest(lines: list[str]) -> None: + """Emit canonical rank evidence from the single rank-0 output stream.""" + sys.stdout.write("".join(f"{line}\n" for line in lines)) + sys.stdout.flush() + + +def _terminal_atexit_sentinel() -> None: + """Negative evidence: this must not run in intentional no-Finalize mode.""" + marker = f"{_TERMINAL_ATEXIT_MARKER} pid={os.getpid()}\n".encode() + try: + os.write(sys.stderr.fileno(), marker) + except BaseException: + pass + + +def _run_terminal_abort_smoke(world: MPI.Intracomm) -> str | None: + """Exercise terminal ABORT relay while keeping ``COMM_WORLD`` healthy.""" + rank = world.Get_rank() + world_size = world.Get_size() + mapping = _mapping(world_size, rank) + broadcaster: MpiFtSubcomm | None = None + wrapped_comm: _NonUlfmComm | None = None + error: str | None = None + + setup_error: str | None = None + try: + # Construct a fresh production communicator, then hide ULFM so this + # scenario specifically validates the bounded ABORT echo fallback. + owner = MpiFtSubcomm( + mapping, + DetectedRankState(world_size), + MpiFtSubcommConfig(startup_timeout_sec=5.0, stop_timeout_sec=5.0), + ) + ft_comm = owner._comm + owner.stop() + wrapped_comm = _NonUlfmComm(ft_comm) + broadcaster = MpiFtSubcomm( + mapping, + DetectedRankState(world_size), + MpiFtSubcommConfig( + startup_timeout_sec=5.0, + stop_timeout_sec=5.0, + reconcile_timeout_sec=5.0, + abort_timeout_sec=3.0, + ), + comm=wrapped_comm, + ) + broadcaster.start() + except Exception: + setup_error = traceback.format_exc() + error = _synchronize_phase_error(world, "terminal-ABORT setup", setup_error) + + if error is None: + trigger_error: str | None = None + try: + world.Barrier() + if rank == 0: + # A transport-independent terminal request is the protocol seam + # used by local second-failure and timeout paths. + assert broadcaster is not None + broadcaster._request_terminal_abort( + RuntimeError("real-MPI terminal smoke"), -1, rank + ) + except Exception: + trigger_error = traceback.format_exc() + error = _synchronize_phase_error( + world, + "terminal-ABORT trigger", + trigger_error, + ) + + local_converged = False + if error is None: + observation_error: str | None = None + try: + assert broadcaster is not None + assert wrapped_comm is not None + deadline = time.monotonic() + _CONVERGENCE_TIMEOUT_SEC + while time.monotonic() < deadline: + if broadcaster._progress_failed.is_set(): + break + time.sleep(0.005) + local_converged = ( + broadcaster._progress_failed.is_set() + and broadcaster.last_error is not None + and broadcaster.world_is_poisoned() + and wrapped_comm.abort_calls == 0 + ) + except Exception: + observation_error = traceback.format_exc() + error = _synchronize_phase_error( + world, + "terminal-ABORT observation", + observation_error, + ) + + if error is None: + if not world.allreduce(local_converged, op=MPI.LAND): + error = "terminal ABORT did not converge on every MPI rank" + else: + # The ABORT protocol is confined to the FT subcommunicator. A + # parent barrier after convergence proves COMM_WORLD remains usable. + world.Barrier() + if rank == 0: + print(_ABORT_WORLD_MARKER, flush=True) + if broadcaster is not None: + try: + broadcaster.stop(timeout=5.0) + if not broadcaster.world_is_poisoned(): + raise AssertionError("terminal FT state lost its poisoned-world latch") + except Exception: + cleanup_error = traceback.format_exc() + error = f"{error or ''}\nterminal cleanup failure:\n{cleanup_error}" + + # Do not free the terminal communicator. Its component intentionally + # retains potentially active requests until process teardown. + return error + + +def _report_scenario_errors( + world: MPI.Intracomm, + scenario: str, + local_error: str | None, +) -> bool: + rank = world.Get_rank() + errors = world.gather(local_error, root=0) + failed = None + if rank == 0: + failed = any(error is not None for error in errors) + if failed: + for error_rank, error in enumerate(errors): + if error is not None: + print( + f"{scenario} failed on rank {error_rank}:\n{error}", + file=sys.stderr, + ) + return world.bcast(failed, root=0) + + +def main() -> int: + world = MPI.COMM_WORLD + expected_world_size = int(sys.argv[1]) if len(sys.argv) > 1 else 0 + mode = sys.argv[2] if len(sys.argv) > 2 else "" + world_size = world.Get_size() + rank = world.Get_rank() + if world_size < 2 or world_size != expected_world_size: + return _unsupported_environment( + f"expected {expected_world_size} MPI ranks, got {world_size}", + rank=rank, + ) + if MPI.Query_thread() < MPI.THREAD_MULTIPLE: + return _unsupported_environment( + f"MPI.THREAD_MULTIPLE is unavailable (provided={MPI.Query_thread()})", + rank=rank, + ) + + if mode == _HEALTHY_MODE: + scenarios = ( + ("invalid-topology", _run_invalid_topology_smoke), + ("healthy-lifecycle", _run_healthy_lifecycle_smoke), + ) + elif mode == _TERMINAL_MODE: + scenarios = ( + ("failure-broadcast", _run_failure_propagation_smoke), + ("terminal-abort", _run_terminal_abort_smoke), + ) + else: + if rank == 0: + print(f"unknown WideEP MPI smoke mode: {mode!r}", file=sys.stderr) + return 2 + for scenario, run_scenario in scenarios: + if _report_scenario_errors(world, scenario, run_scenario(world)): + return 1 + if mode == _HEALTHY_MODE and rank == 0: + print(_HEALTHY_LIFECYCLE_MARKER, flush=True) + return 0 + + +if __name__ == "__main__": + worker_mode = sys.argv[2] if len(sys.argv) > 2 else "" + if worker_mode == _TERMINAL_MODE: + # A normal interpreter shutdown would run mpi4py's Finalize hook. The + # launcher rejects this sentinel, so the smoke independently proves that + # the terminal path below bypassed Python atexit via os._exit(). + atexit.register(_terminal_atexit_sentinel) + exit_code = main() + if worker_mode == _TERMINAL_MODE and not _WORKER_SKIPPED: + # This mode deliberately leaves the FT control plane terminal and + # retains its requests. Skip mpi4py's collective MPI_Finalize just as + # the production poisoned-world shutdown hook will do. + sys.stdout.flush() + sys.stderr.flush() + if exit_code == 0: + world = MPI.COMM_WORLD + rank = world.Get_rank() + world_size = world.Get_size() + try: + # Gather evidence first, then let rank 0 emit a canonical + # manifest. MPI launchers forward separate per-rank pipes and do + # not preserve cross-rank write atomicity in the merged stream. + ready_ranks = world.gather(rank, root=0) + if rank == 0: + if sorted(ready_ranks) != list(range(world_size)): + raise AssertionError(f"invalid terminal READY ranks: {ready_ranks}") + _write_terminal_manifest( + [ + f"{_TERMINAL_READY_MARKER} rank={ready_rank} world_size={world_size}" + for ready_rank in ready_ranks + ] + ) + world.Barrier() + done_ranks = world.gather(rank, root=0) + if rank == 0: + if sorted(done_ranks) != list(range(world_size)): + raise AssertionError(f"invalid terminal DONE ranks: {done_ranks}") + _write_terminal_manifest( + [ + f"{_TERMINAL_COMPLETE_MARKER} world_size={world_size}", + *( + f"{_TERMINAL_DONE_MARKER} rank={done_rank} world_size={world_size}" + for done_rank in done_ranks + ), + ] + ) + # No rank can trigger launcher fail-fast until rank 0 has copied + # the complete gathered manifest to stdout. + world.Barrier() + except BaseException: + try: + _write_terminal_manifest([f"{_TERMINAL_ERROR_MARKER} rank={rank} exit_code=1"]) + except BaseException: + pass + traceback.print_exc() + sys.stderr.flush() + exit_code = 1 + else: + rank = MPI.COMM_WORLD.Get_rank() + try: + _write_terminal_manifest( + [f"{_TERMINAL_ERROR_MARKER} rank={rank} exit_code={exit_code}"] + ) + except BaseException: + pass + os._exit(exit_code) + # Healthy mode returns normally so mpi4py's MPI_Finalize path remains part + # of the smoke test. + sys.exit(exit_code) diff --git a/tests/unittest/_torch/pyexecutor/test_ep_failure_broadcast.py b/tests/unittest/_torch/pyexecutor/test_ep_failure_broadcast.py new file mode 100644 index 000000000000..7e1cab9b5d30 --- /dev/null +++ b/tests/unittest/_torch/pyexecutor/test_ep_failure_broadcast.py @@ -0,0 +1,2166 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""Unit tests for the WideEP MPI failure-broadcast control plane. + +The tests use in-memory MPI, communicator, and request doubles. They exercise +the real progress thread without requiring mpi4py, GPUs, or a multi-rank +launcher. +""" + +from __future__ import annotations + +import gc +import sys +import threading +import time +import types +import weakref +from collections import defaultdict +from collections.abc import Callable +from dataclasses import dataclass + +import numpy as np +import pytest + +from tensorrt_llm._torch.modules.fused_moe.ep_group_health import EPGroupHealth +from tensorrt_llm._torch.pyexecutor import ep_failure_broadcast +from tensorrt_llm._torch.pyexecutor.ep_failure_broadcast import ( + DetectedRankState, + MpiFtSubcomm, + MpiFtSubcommConfig, +) + +_TEST_CONFIG = MpiFtSubcommConfig( + poll_interval_sec=0.001, + startup_timeout_sec=0.5, + stop_timeout_sec=0.5, + reconcile_timeout_sec=2.0, + unattributed_error_timeout_sec=2.0, + abort_timeout_sec=2.0, +) + +_FAILURE_MESSAGE = int(ep_failure_broadcast._MessageKind.FAILURE) +_ABORT_MESSAGE = int(ep_failure_broadcast._MessageKind.ABORT) + + +@pytest.mark.parametrize("value", [0.0, -1.0, float("nan"), float("inf"), float("-inf")]) +def test_config_requires_positive_finite_timeouts(value: float) -> None: + with pytest.raises(ValueError, match="must be finite and > 0"): + MpiFtSubcommConfig(poll_interval_sec=value) + with pytest.raises(ValueError, match="must be finite and > 0"): + MpiFtSubcommConfig(unattributed_error_timeout_sec=value) + with pytest.raises(ValueError, match="must be finite and > 0"): + MpiFtSubcommConfig(abort_timeout_sec=value) + + +def test_default_poll_interval_leaves_margin_below_agreement_budget() -> None: + assert MpiFtSubcommConfig().poll_interval_sec == 0.01 + + +@dataclass(frozen=True) +class _FakeMapping: + world_size: int + rank: int + moe_ep_size: int + moe_ep_rank: int + moe_ep_group: tuple[int, ...] + + @classmethod + def ep_world(cls, size: int, rank: int = 0) -> "_FakeMapping": + return cls( + world_size=size, + rank=rank, + moe_ep_size=size, + moe_ep_rank=rank, + moe_ep_group=tuple(range(size)), + ) + + +class _FakeMpiException(RuntimeError): + def __init__(self, error_class: int, message: str = "fake MPI failure") -> None: + super().__init__(message) + self._error_class = error_class + + def Get_error_class(self) -> int: + return self._error_class + + +class _FakeMPI: + ERRORS_RETURN = object() + ERR_PROC_FAILED = 101 + ERR_REVOKED = 102 + ERR_UNSUPPORTED_OPERATION = 103 + ERR_UNKNOWN = -1 + THREAD_MULTIPLE = 3 + Exception = _FakeMpiException + + def __init__(self, thread_level: int = THREAD_MULTIPLE) -> None: + self._thread_level = thread_level + + def Query_thread(self) -> int: + return self._thread_level + + +class _FakeRequest: + """Controllable nonblocking request that never owns its MPI buffer.""" + + def __init__(self, *, completed: bool = False) -> None: + self._completed = threading.Event() + if completed: + self._completed.set() + self._error: BaseException | None = None + self._lock = threading.Lock() + self.cancel_calls = 0 + self.test_calls = 0 + + def Cancel(self) -> None: + with self._lock: + self.cancel_calls += 1 + self._completed.set() + + def Test(self) -> bool: + with self._lock: + self.test_calls += 1 + error = self._error + if error is not None: + raise error + return self._completed.is_set() + + def complete(self) -> None: + self._completed.set() + + def fail(self, error: BaseException) -> None: + with self._lock: + self._error = error + + +class _BlockingRequest(_FakeRequest): + """Request whose Test call can hold the progress thread for stop tests.""" + + def __init__(self) -> None: + super().__init__() + self.entered_test = threading.Event() + self.release_test = threading.Event() + + def Test(self) -> bool: + self.entered_test.set() + self.release_test.wait() + return super().Test() + + +class _NonCompletingCancelRequest(_FakeRequest): + """Cancellation request that remains incomplete until explicitly released.""" + + def Cancel(self) -> None: + with self._lock: + self.cancel_calls += 1 + + +@dataclass(frozen=True) +class _ReceiveRecord: + request: _FakeRequest + buffer: np.ndarray + thread_id: int + + +@dataclass(frozen=True) +class _SendRecord: + request: _FakeRequest + buffer_ref: weakref.ReferenceType[np.ndarray] + destination: int + message_kind: int + payload: int + tag: int + thread_id: int + + +class _FakeComm: + def __init__( + self, + *, + size: int, + rank: int = 0, + receive_factory: Callable[[int], _FakeRequest] | None = None, + send_factory: Callable[[int], _FakeRequest] | None = None, + ulfm_probe_error: BaseException | None = None, + ulfm_revoked: bool = False, + revoke_error: BaseException | None = None, + abort_error: BaseException | None = None, + ) -> None: + self._size = size + self._rank = rank + self._receive_factory = receive_factory or (lambda _source: _FakeRequest()) + self._send_factory = send_factory or (lambda _destination: _FakeRequest()) + self._ulfm_probe_error = ulfm_probe_error + self._ulfm_revoked = ulfm_revoked + self._revoke_error = revoke_error + self._abort_error = abort_error + self._lock = threading.Lock() + self._receives: dict[int, list[_ReceiveRecord]] = defaultdict(list) + self._sends: list[_SendRecord] = [] + self.errhandler: object | None = None + self.errhandler_calls = 0 + self.is_revoked_calls = 0 + self.revoke_calls = 0 + self.abort_calls = 0 + + def Get_rank(self) -> int: + return self._rank + + def Get_size(self) -> int: + return self._size + + def Set_errhandler(self, errhandler: object) -> None: + self.errhandler_calls += 1 + self.errhandler = errhandler + + def Is_revoked(self) -> bool: + self.is_revoked_calls += 1 + if self._ulfm_probe_error is not None: + raise self._ulfm_probe_error + return self._ulfm_revoked + + def Revoke(self) -> None: + self.revoke_calls += 1 + if self._revoke_error is not None: + raise self._revoke_error + + def Abort(self, _errorcode: int) -> None: + self.abort_calls += 1 + if self._abort_error is not None: + raise self._abort_error + + def Irecv(self, buffer: np.ndarray, source: int, tag: int) -> _FakeRequest: + del tag + request = self._receive_factory(source) + record = _ReceiveRecord( + request=request, + buffer=buffer, + thread_id=threading.get_ident(), + ) + with self._lock: + self._receives[source].append(record) + return request + + def Isend(self, buffer: np.ndarray, dest: int, tag: int) -> _FakeRequest: + request = self._send_factory(dest) + record = _SendRecord( + request=request, + buffer_ref=weakref.ref(buffer), + destination=dest, + message_kind=int(buffer[0]), + payload=int(buffer[1]), + tag=tag, + thread_id=threading.get_ident(), + ) + with self._lock: + self._sends.append(record) + return request + + def receive_count(self, source: int) -> int: + with self._lock: + return len(self._receives[source]) + + def latest_receive(self, source: int) -> _ReceiveRecord: + with self._lock: + return self._receives[source][-1] + + def deliver( + self, + source: int, + failed_rank: int, + *, + message_kind: int = _FAILURE_MESSAGE, + ) -> None: + record = self.latest_receive(source) + record.buffer[:] = (message_kind, failed_rank) + record.request.complete() + + @property + def sends(self) -> tuple[_SendRecord, ...]: + with self._lock: + return tuple(self._sends) + + +def _wait_until( + predicate: Callable[[], bool], + *, + timeout: float = 1.0, + description: str = "condition", +) -> None: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if predicate(): + return + time.sleep(0.001) + raise AssertionError(f"timed out waiting for {description}") + + +def _make_broadcaster( + *, + size: int = 4, + rank: int = 0, + detected_state: DetectedRankState | None = None, + comm: _FakeComm | None = None, + mpi: _FakeMPI | None = None, + callback: Callable[[int, int, float], None] | None = None, +) -> tuple[MpiFtSubcomm, DetectedRankState, _FakeComm, _FakeMPI]: + detected_state = detected_state or DetectedRankState(size) + comm = comm or _FakeComm(size=size, rank=rank) + mpi = mpi or _FakeMPI() + broadcaster = MpiFtSubcomm( + _FakeMapping.ep_world(size, rank), + detected_state, + _TEST_CONFIG, + callback, + comm=comm, + mpi_module=mpi, + ) + return broadcaster, detected_state, comm, mpi + + +def test_committed_ep_group_health_is_rejected_before_mpi_setup() -> None: + comm = _FakeComm(size=2) + + with pytest.raises(TypeError, match="committed EPGroupHealth cannot be used"): + MpiFtSubcomm( + _FakeMapping.ep_world(2), + EPGroupHealth(2), + _TEST_CONFIG, + comm=comm, + mpi_module=_FakeMPI(), + ) + + assert comm.errhandler_calls == 0 + assert comm.is_revoked_calls == 0 + + +def test_watchdog_callback_records_and_fans_out_detected_failure() -> None: + detected_state = DetectedRankState(4) + broadcaster, _, comm, _ = _make_broadcaster(detected_state=detected_state) + broadcaster.start() + caller_thread = threading.get_ident() + try: + watchdog_callback = broadcaster.report_detected_failure + assert watchdog_callback(2) is True + assert detected_state.get_failed_ranks() == frozenset({2}) + _wait_until(lambda: len(comm.sends) == 2, description="failure fanout") + _wait_until( + lambda: all(record.request.test_calls > 0 for record in comm.sends), + description="send request polling", + ) + + assert {record.destination for record in comm.sends} == {1, 3} + assert {record.payload for record in comm.sends} == {2} + assert all(record.thread_id != caller_thread for record in comm.sends) + assert all(record.request.test_calls > 0 for record in comm.sends) + assert all(record.buffer_ref() is not None for record in comm.sends) + + # A repeated detector report is locally coalesced. + assert watchdog_callback(2) is False + time.sleep(2 * _TEST_CONFIG.poll_interval_sec) + assert len(comm.sends) == 2 + + for send in comm.sends: + send.request.complete() + comm.deliver(source=1, failed_rank=2) + comm.deliver(source=3, failed_rank=2) + _wait_until( + broadcaster.detected_state_is_reconciled, + description="detection reconciliation", + ) + finally: + broadcaster.stop() + + +def test_externally_premarked_detected_state_fails_closed() -> None: + detected_state = DetectedRankState(2) + assert detected_state.record_failure(1) is True + broadcaster, _, comm, _ = _make_broadcaster( + size=2, + detected_state=detected_state, + ) + + try: + try: + broadcaster.start() + except RuntimeError: + # The progress thread can detect the invalid initial state before + # start() observes RUNNING, or immediately after start() returns. + pass + _wait_until( + lambda: isinstance(broadcaster.last_error, RuntimeError), + description="pre-recorded detected-state rejection", + ) + assert "without pre-recording evidence" in str(broadcaster.last_error) + _wait_until(lambda: comm.revoke_calls == 1, description="pre-mark fail-closed revoke") + finally: + broadcaster.stop() + + +def test_detection_reconciliation_does_not_mutate_committed_health() -> None: + callbacks: list[tuple[int, int, float]] = [] + detected_state = DetectedRankState(2) + committed_health = EPGroupHealth(2) + initial_committed_snapshot = committed_health.snapshot() + broadcaster, _, _, _ = _make_broadcaster( + size=2, + detected_state=detected_state, + callback=lambda failed, source, when: callbacks.append((failed, source, when)), + ) + broadcaster.start() + try: + assert broadcaster.report_detected_failure(1) is True + _wait_until( + lambda: len(callbacks) == 1, + description="local detection coordinator handoff", + ) + _wait_until( + broadcaster.detected_state_is_reconciled, + description="detection reconciliation", + ) + + assert detected_state.get_failed_ranks() == frozenset({1}) + assert detected_state.generation == 1 + assert committed_health.snapshot() == initial_committed_snapshot + assert committed_health.all_active() is True + assert callbacks[0][:2] == (1, 0) + finally: + broadcaster.stop() + + +def test_survivor_echoes_reconcile_failure_detection() -> None: + broadcaster, _, comm, _ = _make_broadcaster(size=4) + broadcaster.start() + try: + assert broadcaster.report_detected_failure(2) is True + assert broadcaster.failure_detection_is_reconciled(2) is False + _wait_until(lambda: len(comm.sends) == 2, description="initial failure fanout") + + comm.deliver(source=1, failed_rank=2) + _wait_until( + lambda: comm.receive_count(1) == 2, + description="source 1 receive repost", + ) + assert broadcaster.failure_detection_is_reconciled(2) is False + + comm.deliver(source=3, failed_rank=2) + for send in comm.sends: + send.request.complete() + _wait_until( + lambda: broadcaster.failure_detection_is_reconciled(2), + description="all survivor echoes", + ) + time.sleep(2 * _TEST_CONFIG.poll_interval_sec) + assert len(comm.sends) == 2 + assert {send.destination for send in comm.sends} == {1, 3} + assert all(send.message_kind == _FAILURE_MESSAGE for send in comm.sends) + assert broadcaster.detected_health_is_reconciled() is True + assert comm.revoke_calls == 0 + finally: + broadcaster.stop() + + +def test_detected_state_is_monotonic_for_one_communicator_epoch() -> None: + state = DetectedRankState(3) + + assert state.snapshot() == ep_failure_broadcast.DetectedRankStateSnapshot( + mask=0b111, + failed_ranks=frozenset(), + generation=0, + ) + assert state.record_failure(2) is True + assert state.record_failure(2) is False + assert state.get_mask() == 0b011 + assert state.get_failed_ranks() == frozenset({2}) + assert state.generation == 1 + assert not hasattr(state, "mark_active") + + +def test_detected_state_coalesces_concurrent_evidence() -> None: + state = DetectedRankState(4) + barrier = threading.Barrier(8) + results: list[bool] = [] + results_lock = threading.Lock() + + def record_same_failure() -> None: + barrier.wait() + changed = state.record_failure(2) + with results_lock: + results.append(changed) + + threads = [threading.Thread(target=record_same_failure) for _ in range(8)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + assert results.count(True) == 1 + assert results.count(False) == 7 + assert state.snapshot() == ep_failure_broadcast.DetectedRankStateSnapshot( + mask=0b1011, + failed_ranks=frozenset({2}), + generation=1, + ) + + +def test_second_failure_between_claim_and_observe_cannot_open_single_failure_gate( + monkeypatch: pytest.MonkeyPatch, +) -> None: + health = DetectedRankState(3) + broadcaster, _, _, _ = _make_broadcaster(size=3, detected_state=health) + original_claim = broadcaster._claim_failure + mutated = False + + def claim_then_add_second_failure(failed_rank: int) -> None: + nonlocal mutated + original_claim(failed_rank) + if not mutated: + mutated = True + assert health.record_failure(1) is True + + monkeypatch.setattr(broadcaster, "_claim_failure", claim_then_add_second_failure) + broadcaster.start() + try: + assert broadcaster.broadcast_failure(2) is True + _wait_until( + lambda: 2 in broadcaster._failure_fanout_complete, + description="single-survivor fanout completion", + ) + assert health.get_failed_ranks() == frozenset({1, 2}) + _wait_until( + lambda: isinstance(broadcaster.last_error, RuntimeError), + description="second-detected-failure terminal handling", + ) + assert broadcaster.failure_is_reconciled(2) is False + assert broadcaster.health_is_reconciled() is False + finally: + broadcaster.stop() + + +def test_reconciliation_timeout_fails_closed_while_mpi_test_is_wedged() -> None: + config = MpiFtSubcommConfig( + poll_interval_sec=0.001, + startup_timeout_sec=0.5, + stop_timeout_sec=0.5, + reconcile_timeout_sec=0.02, + ) + blocking_request = _BlockingRequest() + comm = _FakeComm( + size=3, + receive_factory=lambda source: blocking_request if source == 1 else _FakeRequest(), + ) + broadcaster = MpiFtSubcomm( + _FakeMapping.ep_world(3), + DetectedRankState(3), + config, + comm=comm, + mpi_module=_FakeMPI(), + ) + broadcaster.start() + _wait_until(blocking_request.entered_test.is_set, description="blocked MPI progress") + assert broadcaster.pre_failover(2) is True + assert blocking_request.release_test.is_set() is False + assert broadcaster._outbound_reports.empty() is False + + _wait_until( + lambda: isinstance(broadcaster.last_error, TimeoutError), + description="terminal reconciliation timeout", + ) + _wait_until(lambda: comm.revoke_calls == 1, description="deadline-monitor revoke") + assert blocking_request.release_test.is_set() is False + assert broadcaster.health_is_reconciled() is False + assert broadcaster.world_is_poisoned() is True + with pytest.raises(RuntimeError, match="not running"): + broadcaster.pre_failover(1) + blocking_request.release_test.set() + broadcaster.stop() + + +def test_reconciliation_timeout_world_aborts_when_mpi_test_is_wedged_without_ulfm() -> None: + config = MpiFtSubcommConfig( + poll_interval_sec=0.001, + startup_timeout_sec=0.5, + stop_timeout_sec=0.5, + reconcile_timeout_sec=0.01, + abort_timeout_sec=0.01, + ) + blocking_request = _BlockingRequest() + comm = _FakeComm( + size=3, + receive_factory=lambda source: blocking_request if source == 1 else _FakeRequest(), + ulfm_probe_error=NotImplementedError(), + ) + broadcaster = MpiFtSubcomm( + _FakeMapping.ep_world(3), + DetectedRankState(3), + config, + comm=comm, + mpi_module=_FakeMPI(), + ) + broadcaster.start() + _wait_until(blocking_request.entered_test.is_set, description="blocked MPI progress") + assert broadcaster.pre_failover(2) is True + + _wait_until(lambda: comm.abort_calls == 1, description="deadline-monitor MPI_Abort") + assert blocking_request.release_test.is_set() is False + assert isinstance(broadcaster.last_error, TimeoutError) + assert broadcaster.world_is_poisoned() is True + + blocking_request.release_test.set() + broadcaster.stop() + + +def test_second_distinct_local_failure_fails_closed_before_health_mutation() -> None: + broadcaster, health, comm, _ = _make_broadcaster(size=4) + broadcaster.start() + assert broadcaster.pre_failover(3) is True + + with pytest.raises(RuntimeError, match="exactly one distinct failed rank") as raised: + broadcaster.pre_failover(2) + + assert health.get_failed_ranks() == frozenset({3}) + _wait_until(lambda: broadcaster.last_error is raised.value, description="terminal failure") + _wait_until(lambda: comm.revoke_calls == 1, description="progress-thread abort revoke") + assert broadcaster.health_is_reconciled() is False + with pytest.raises(RuntimeError, match="not running"): + broadcaster.pre_failover(1) + broadcaster.stop() + + +def test_no_ulfm_terminal_abort_relays_and_converges_without_world_abort() -> None: + comm = _FakeComm(size=3, ulfm_probe_error=NotImplementedError()) + broadcaster, health, _, _ = _make_broadcaster(size=3, comm=comm) + broadcaster.start() + assert broadcaster.pre_failover(2) is True + with pytest.raises(RuntimeError, match="exactly one distinct failed rank"): + broadcaster.pre_failover(1) + + _wait_until( + lambda: any(send.message_kind == _ABORT_MESSAGE for send in comm.sends), + description="terminal ABORT relay", + ) + for send in comm.sends: + if send.message_kind == _ABORT_MESSAGE: + send.request.complete() + comm.deliver(source=1, failed_rank=1, message_kind=_ABORT_MESSAGE) + _wait_until(broadcaster._progress_failed.is_set, description="ABORT convergence") + + assert health.get_failed_ranks() == frozenset({2}) + assert comm.abort_calls == 0 + assert comm.revoke_calls == 0 + broadcaster.stop() + + +def test_remote_conflicting_failure_relays_abort_before_world_abort() -> None: + comm = _FakeComm(size=3, ulfm_probe_error=NotImplementedError()) + broadcaster, health, _, _ = _make_broadcaster(size=3, comm=comm) + broadcaster.start() + assert broadcaster.pre_failover(2) is True + _wait_until( + lambda: any(send.message_kind == _FAILURE_MESSAGE for send in comm.sends), + description="first failure fanout", + ) + _wait_until(lambda: comm.receive_count(1) == 1, description="conflicting report receive") + comm.deliver(source=1, failed_rank=1) + + _wait_until( + lambda: any(send.message_kind == _ABORT_MESSAGE for send in comm.sends), + description="remote-conflict ABORT relay", + ) + assert comm.abort_calls == 0 + for send in comm.sends: + if send.message_kind == _ABORT_MESSAGE: + send.request.complete() + _wait_until(lambda: comm.receive_count(1) == 2, description="survivor ABORT receive") + comm.deliver(source=1, failed_rank=1, message_kind=_ABORT_MESSAGE) + _wait_until(broadcaster._progress_failed.is_set, description="remote-conflict convergence") + + assert health.get_failed_ranks() == frozenset({2}) + assert isinstance(broadcaster.last_error, RuntimeError) + assert "exactly one distinct failed rank" in str(broadcaster.last_error) + assert comm.abort_calls == 0 + broadcaster.stop() + + +def test_reconciled_terminal_abort_at_deadline_does_not_world_abort() -> None: + comm = _FakeComm(size=1, ulfm_probe_error=NotImplementedError()) + broadcaster, _, _, _ = _make_broadcaster(size=1, comm=comm) + broadcaster._request_terminal_abort(RuntimeError("terminal"), -1, 0) + broadcaster._drain_outbound_reports([]) + with broadcaster._protocol_lock: + broadcaster._abort_deadline = time.monotonic() - 1.0 + + assert broadcaster._progress_terminal_abort() is True + assert comm.abort_calls == 0 + broadcaster.stop() + + +def test_monitor_does_not_world_abort_reconciled_relay_when_mpi_test_resumes() -> None: + config = MpiFtSubcommConfig( + poll_interval_sec=0.001, + startup_timeout_sec=0.5, + stop_timeout_sec=0.5, + reconcile_timeout_sec=0.2, + abort_timeout_sec=0.2, + ) + blocking_request = _BlockingRequest() + comm = _FakeComm( + size=2, + receive_factory=lambda _source: blocking_request, + ulfm_probe_error=NotImplementedError(), + ) + broadcaster = MpiFtSubcomm( + _FakeMapping.ep_world(2), + DetectedRankState(2), + config, + comm=comm, + mpi_module=_FakeMPI(), + ) + broadcaster.start() + _wait_until(blocking_request.entered_test.is_set, description="blocked MPI progress") + broadcaster._request_terminal_abort(RuntimeError("terminal"), -1, 0) + with broadcaster._protocol_lock: + broadcaster._abort_reporters.add(1) + broadcaster._abort_fanout_posted = True + broadcaster._abort_fanout_complete = True + broadcaster._deadline_monitor_wake_event.set() + + _wait_until( + lambda: ( + broadcaster._deadline_monitor_thread is not None + and not broadcaster._deadline_monitor_thread.is_alive() + ), + description="monitor-observed terminal reconciliation", + ) + assert broadcaster._progress_failed.is_set() is False + assert comm.abort_calls == 0 + + blocking_request.release_test.set() + _wait_until(broadcaster._progress_failed.is_set, description="clean terminal progress stop") + broadcaster.stop() + assert comm.abort_calls == 0 + + +def test_terminal_abort_waits_until_its_relay_completes() -> None: + comm = _FakeComm(size=3, ulfm_probe_error=NotImplementedError()) + broadcaster, _, _, _ = _make_broadcaster(size=3, comm=comm) + broadcaster._handle_received_abort(-1, 1) + broadcaster._handle_received_abort(-1, 2) + + assert broadcaster._progress_terminal_abort() is False + pending_sends: list[ep_failure_broadcast._PendingSend] = [] + broadcaster._drain_outbound_reports(pending_sends) + assert {send.destination for send in comm.sends} == {1, 2} + broadcaster._progress_sends(pending_sends) + assert broadcaster._progress_terminal_abort() is False + + for send in comm.sends: + send.request.complete() + broadcaster._progress_sends(pending_sends) + assert pending_sends == [] + assert broadcaster._progress_terminal_abort() is True + assert comm.abort_calls == 0 + broadcaster.stop() + + +def test_received_terminal_abort_is_relayed_without_mutating_health() -> None: + comm = _FakeComm(size=3, ulfm_probe_error=NotImplementedError()) + broadcaster, health, _, _ = _make_broadcaster(size=3, comm=comm) + broadcaster.start() + _wait_until(lambda: comm.receive_count(1) == 1, description="source 1 receive") + comm.deliver(source=1, failed_rank=2, message_kind=_ABORT_MESSAGE) + + _wait_until( + lambda: len([send for send in comm.sends if send.message_kind == _ABORT_MESSAGE]) == 2, + description="relayed terminal ABORT", + ) + for send in comm.sends: + if send.message_kind == _ABORT_MESSAGE: + send.request.complete() + _wait_until(lambda: comm.receive_count(2) == 1, description="source 2 receive") + comm.deliver(source=2, failed_rank=2, message_kind=_ABORT_MESSAGE) + _wait_until(broadcaster._progress_failed.is_set, description="peer ABORT convergence") + + assert health.all_active() is True + assert isinstance(broadcaster.last_error, RuntimeError) + assert broadcaster.world_is_poisoned() is True + assert comm.abort_calls == 0 + broadcaster.stop() + assert broadcaster.world_is_poisoned() is True + + +def test_no_ulfm_abort_timeout_uses_world_abort_before_stop_completes() -> None: + config = MpiFtSubcommConfig( + poll_interval_sec=0.001, + startup_timeout_sec=0.5, + stop_timeout_sec=0.5, + reconcile_timeout_sec=0.2, + abort_timeout_sec=0.02, + ) + comm = _FakeComm(size=3, ulfm_probe_error=NotImplementedError()) + broadcaster = MpiFtSubcomm( + _FakeMapping.ep_world(3), + DetectedRankState(3), + config, + comm=comm, + mpi_module=_FakeMPI(), + ) + broadcaster.start() + assert broadcaster.pre_failover(2) is True + with pytest.raises(RuntimeError, match="exactly one distinct failed rank"): + broadcaster.pre_failover(1) + + # stop_event must not let the progress loop skip the pending ABORT fallback. + broadcaster.stop(timeout=0.5) + assert comm.abort_calls == 1 + assert comm.revoke_calls == 0 + + +def test_stop_cannot_bypass_abort_requested_during_blocked_drain() -> None: + send_entered = threading.Event() + release_send = threading.Event() + + def blocking_send(_destination: int) -> _FakeRequest: + send_entered.set() + release_send.wait() + return _FakeRequest() + + config = MpiFtSubcommConfig( + poll_interval_sec=0.001, + startup_timeout_sec=0.5, + stop_timeout_sec=0.5, + reconcile_timeout_sec=0.2, + abort_timeout_sec=0.02, + ) + comm = _FakeComm( + size=3, + send_factory=blocking_send, + ulfm_probe_error=NotImplementedError(), + ) + broadcaster = MpiFtSubcomm( + _FakeMapping.ep_world(3), + DetectedRankState(3), + config, + comm=comm, + mpi_module=_FakeMPI(), + ) + broadcaster.start() + _wait_until(lambda: comm.receive_count(1) == 1, description="terminal receive") + comm.deliver(source=1, failed_rank=-1, message_kind=_ABORT_MESSAGE) + _wait_until(send_entered.is_set, description="blocked ABORT drain") + + stop_errors: list[BaseException] = [] + + def stop_broadcaster() -> None: + try: + broadcaster.stop(timeout=0.5) + except BaseException as error: + stop_errors.append(error) + + stop_thread = threading.Thread(target=stop_broadcaster) + stop_thread.start() + _wait_until(broadcaster._stop_event.is_set, description="stop during ABORT drain") + release_send.set() + stop_thread.join(timeout=1.0) + + assert stop_thread.is_alive() is False + assert stop_errors == [] + assert comm.abort_calls == 1 + assert broadcaster._progress_failed.is_set() + + +def test_concurrent_local_and_remote_distinct_failures_accept_only_one( + monkeypatch: pytest.MonkeyPatch, +) -> None: + broadcaster, health, comm, _ = _make_broadcaster(size=4) + broadcaster.start() + _wait_until(lambda: comm.receive_count(1) == 1, description="remote receive") + + original_claim = broadcaster._claim_failure + local_claimed = threading.Event() + release_local_claim = threading.Event() + + def racing_claim(failed_rank: int) -> None: + original_claim(failed_rank) + if failed_rank == 2: + local_claimed.set() + release_local_claim.wait() + + monkeypatch.setattr(broadcaster, "_claim_failure", racing_claim) + local_errors: list[BaseException] = [] + + def report_local_failure() -> None: + try: + broadcaster.pre_failover(2) + except BaseException as error: + local_errors.append(error) + + local_thread = threading.Thread(target=report_local_failure) + local_thread.start() + _wait_until(local_claimed.is_set, description="local failure claim") + comm.deliver(source=1, failed_rank=3) + # The remote handler must wait for the local lifecycle transaction to commit + # its accepted rank and health generation atomically. + time.sleep(2 * _TEST_CONFIG.poll_interval_sec) + assert health.all_active() is True + release_local_claim.set() + local_thread.join(timeout=1.0) + assert local_thread.is_alive() is False + + _wait_until(lambda: broadcaster.last_error is not None, description="second-failure rejection") + _wait_until(lambda: comm.revoke_calls == 1, description="terminal abort revoke") + assert len(health.get_failed_ranks()) == 1 + assert health.get_failed_ranks() == frozenset({2}) + assert "exactly one distinct failed rank" in str(broadcaster.last_error) + assert local_errors == [] + broadcaster.stop() + + +def test_isend_post_error_to_survivor_is_terminal_without_ulfm() -> None: + error = _FakeMpiException(999, "send post failed") + + def fail_send(_destination: int) -> _FakeRequest: + raise error + + comm = _FakeComm( + size=3, + send_factory=fail_send, + ulfm_probe_error=NotImplementedError(), + ) + broadcaster, _, _, _ = _make_broadcaster(size=3, comm=comm) + broadcaster.start() + assert broadcaster.pre_failover(2) is True + + _wait_until(lambda: broadcaster.last_error is error, description="terminal Isend error") + assert broadcaster.health_is_reconciled() is False + assert comm.revoke_calls == 0 + assert comm.sends == () + broadcaster.stop() + + +def test_send_test_error_is_terminal_and_revokes_with_ulfm() -> None: + error = _FakeMpiException(999, "send Test failed") + failed_request = _FakeRequest() + failed_request.fail(error) + comm = _FakeComm(size=3, send_factory=lambda _destination: failed_request) + broadcaster, _, _, _ = _make_broadcaster(size=3, comm=comm) + broadcaster.start() + assert broadcaster.pre_failover(2) is True + + _wait_until(lambda: broadcaster.last_error is error, description="terminal send Test error") + _wait_until(lambda: comm.revoke_calls == 1, description="ULFM emergency revoke") + assert broadcaster.health_is_reconciled() is False + assert failed_request.cancel_calls == 0 + assert broadcaster._retained_requests + broadcaster.stop() + + +def test_received_report_is_idempotent_and_callback_runs_once() -> None: + callbacks: list[tuple[int, int, float]] = [] + broadcaster, health, comm, _ = _make_broadcaster( + size=3, + callback=lambda failed, source, when: callbacks.append((failed, source, when)), + ) + broadcaster.start() + try: + _wait_until(lambda: comm.receive_count(1) == 1, description="initial receive") + comm.deliver(source=1, failed_rank=2) + _wait_until(lambda: health.is_active(2) is False, description="received failure") + _wait_until(lambda: comm.receive_count(1) == 2, description="reposted receive") + + comm.deliver(source=1, failed_rank=2) + _wait_until(lambda: comm.receive_count(1) == 3, description="duplicate receive progress") + _wait_until(lambda: len(callbacks) == 1, description="failure callback") + + assert health.generation == 1 + assert len(callbacks) == 1 + failed_rank, source, event_time = callbacks[0] + assert (failed_rank, source) == (2, 1) + assert isinstance(event_time, float) + _wait_until(lambda: len(comm.sends) == 1, description="failure relay") + comm.sends[0].request.complete() + _wait_until(broadcaster.health_is_reconciled, description="failure reconciliation") + finally: + broadcaster.stop() + + +def test_monitor_cannot_observe_remote_claim_before_detected_state_update( + monkeypatch: pytest.MonkeyPatch, +) -> None: + broadcaster, health, comm, _ = _make_broadcaster(size=3) + original_claim = broadcaster._claim_failure + claim_entered = threading.Event() + release_claim = threading.Event() + + def claim_then_pause(failed_rank: int) -> None: + original_claim(failed_rank) + claim_entered.set() + release_claim.wait() + + monkeypatch.setattr(broadcaster, "_claim_failure", claim_then_pause) + broadcaster.start() + try: + comm.deliver(source=1, failed_rank=2) + _wait_until(claim_entered.is_set, description="paused remote failure claim") + time.sleep(3 * _TEST_CONFIG.poll_interval_sec) + + assert health.is_active(2) is True + assert broadcaster.last_error is None + assert comm.revoke_calls == 0 + + release_claim.set() + _wait_until(lambda: not health.is_active(2), description="recorded remote detection") + _wait_until(lambda: len(comm.sends) == 1, description="remote failure echo") + comm.sends[0].request.complete() + _wait_until(broadcaster.health_is_reconciled, description="failure reconciliation") + finally: + release_claim.set() + broadcaster.stop() + + +def test_blocking_callback_does_not_block_mpi_progress_and_stop_is_bounded() -> None: + callback_entered = threading.Event() + release_callback = threading.Event() + + def blocking_callback(_failed: int, _source: int, _when: float) -> None: + callback_entered.set() + release_callback.wait() + + broadcaster, health, comm, _ = _make_broadcaster( + size=3, + callback=blocking_callback, + ) + broadcaster.start() + _wait_until(lambda: comm.receive_count(1) == 1, description="initial receive") + comm.deliver(source=1, failed_rank=2) + _wait_until(callback_entered.is_set, description="blocking callback") + _wait_until(lambda: len(comm.sends) == 1, description="failure relay") + comm.sends[0].request.complete() + + # The MPI thread must repost and continue polling while user telemetry is blocked. + _wait_until(lambda: comm.receive_count(1) == 2, description="receive repost") + comm.deliver(source=1, failed_rank=2) + _wait_until(lambda: comm.receive_count(1) == 3, description="duplicate receive progress") + assert health.get_failed_ranks() == frozenset({2}) + + start = time.monotonic() + with pytest.raises(TimeoutError, match="callback thread did not stop"): + broadcaster.stop(timeout=0.02) + assert time.monotonic() - start < 0.2 + assert broadcaster.last_error is None + assert broadcaster.world_is_poisoned() is True + + release_callback.set() + broadcaster.stop(timeout=0.5) + + +def test_invalid_received_rank_is_ignored_and_receive_is_reposted() -> None: + callbacks: list[tuple[int, int, float]] = [] + broadcaster, health, comm, _ = _make_broadcaster( + size=3, + callback=lambda failed, source, when: callbacks.append((failed, source, when)), + ) + broadcaster.start() + try: + _wait_until(lambda: comm.receive_count(1) == 1, description="initial receive") + comm.deliver(source=1, failed_rank=99) + _wait_until(lambda: comm.receive_count(1) == 2, description="receive after invalid payload") + + assert health.all_active() is True + assert callbacks == [] + assert broadcaster.last_error is None + finally: + broadcaster.stop() + + +def test_single_rank_group_has_no_peer_requests_or_sends() -> None: + config = MpiFtSubcommConfig( + poll_interval_sec=0.001, + startup_timeout_sec=0.5, + stop_timeout_sec=0.5, + reconcile_timeout_sec=0.02, + abort_timeout_sec=0.02, + ) + comm = _FakeComm(size=1, ulfm_probe_error=NotImplementedError()) + health = DetectedRankState(1) + broadcaster = MpiFtSubcomm( + _FakeMapping.ep_world(1), + health, + config, + comm=comm, + mpi_module=_FakeMPI(), + ) + broadcaster.start() + try: + assert broadcaster.broadcast_failure(0) is True + _wait_until(broadcaster._outbound_reports.empty, description="empty single-rank report") + assert health.get_failed_ranks() == frozenset({0}) + assert broadcaster.world_is_poisoned() is True + assert broadcaster.failure_is_reconciled(0) is False + assert broadcaster.health_is_reconciled() is False + assert comm.sends == () + assert comm.receive_count(0) == 0 + _wait_until( + lambda: isinstance(broadcaster.last_error, TimeoutError), + description="empty-survivor terminal handling", + ) + finally: + broadcaster.stop() + + +def test_lifecycle_is_not_restartable_and_stop_before_start_is_terminal() -> None: + stopped, _, _, _ = _make_broadcaster(size=1) + stopped.stop() + stopped.stop() + with pytest.raises(RuntimeError, match="cannot start from stopped"): + stopped.start() + + running, _, _, _ = _make_broadcaster(size=1) + running.start() + try: + with pytest.raises(RuntimeError, match="cannot start from running"): + running.start() + finally: + running.stop() + with pytest.raises(RuntimeError, match="cannot start from stopped"): + running.start() + + +@pytest.mark.parametrize( + "failing_thread_name", + [ + "wide-ep-ft-callback", + "wide-ep-ft-deadline-monitor", + "wide-ep-ft-broadcast", + ], + ids=["callback", "deadline-monitor", "progress"], +) +def test_thread_start_failure_is_terminal_and_stoppable( + monkeypatch: pytest.MonkeyPatch, + failing_thread_name: str, +) -> None: + registry = ep_failure_broadcast._PROCESS_LIFETIME_REFS + initial_registry_size = len(registry) + original_start = threading.Thread.start + start_error = OSError(f"cannot start {failing_thread_name}") + + def failing_start(thread: threading.Thread) -> None: + if thread.name == failing_thread_name: + raise start_error + original_start(thread) + + monkeypatch.setattr(threading.Thread, "start", failing_start) + comm = _FakeComm(size=2) + broadcaster, _, _, _ = _make_broadcaster( + size=2, + comm=comm, + callback=lambda _failed, _source, _when: None, + ) + + try: + with pytest.raises(RuntimeError, match="thread failed to start") as raised: + broadcaster.start() + + assert raised.value.__cause__ is start_error + assert broadcaster.last_error is start_error + assert broadcaster._progress_failed.is_set() + assert broadcaster._thread is None + assert comm.revoke_calls == 1 + if failing_thread_name == "wide-ep-ft-callback": + assert broadcaster._callback_thread is None + assert broadcaster._deadline_monitor_thread is None + elif failing_thread_name == "wide-ep-ft-deadline-monitor": + assert broadcaster._callback_thread is not None + assert broadcaster._callback_thread.is_alive() is False + assert broadcaster._deadline_monitor_thread is None + else: + assert broadcaster._callback_thread is not None + assert broadcaster._callback_thread.is_alive() is False + assert broadcaster._deadline_monitor_thread is not None + assert broadcaster._deadline_monitor_thread.is_alive() is False + + # A failed start must not leave an unstarted Thread object that stop() + # later attempts to join. + broadcaster.stop() + finally: + del registry[initial_registry_size:] + + +def test_concurrent_stop_prevents_start_from_publishing_running() -> None: + receive_entered = threading.Event() + release_receive = threading.Event() + + def blocking_receive(_source: int) -> _FakeRequest: + receive_entered.set() + release_receive.wait() + return _FakeRequest() + + comm = _FakeComm(size=2, receive_factory=blocking_receive) + broadcaster, _, _, _ = _make_broadcaster(size=2, comm=comm) + start_errors: list[BaseException] = [] + stop_errors: list[BaseException] = [] + + def start_broadcaster() -> None: + try: + broadcaster.start() + except BaseException as error: + start_errors.append(error) + + def stop_broadcaster() -> None: + try: + broadcaster.stop(timeout=0.5) + except BaseException as error: + stop_errors.append(error) + + start_thread = threading.Thread(target=start_broadcaster) + start_thread.start() + _wait_until(receive_entered.is_set, description="blocked startup receive") + stop_thread = threading.Thread(target=stop_broadcaster) + stop_thread.start() + _wait_until(broadcaster._stop_event.is_set, description="concurrent stop") + release_receive.set() + + start_thread.join(timeout=1.0) + stop_thread.join(timeout=1.0) + assert start_thread.is_alive() is False + assert stop_thread.is_alive() is False + assert len(start_errors) == 1 + assert "did not reach running state" in str(start_errors[0]) + assert stop_errors == [] + assert broadcaster.health_is_reconciled() is False + + +def test_startup_failure_publishes_error_before_failed_flag( + monkeypatch: pytest.MonkeyPatch, +) -> None: + comm = _FakeComm(size=1) + broadcaster, _, _, _ = _make_broadcaster(size=1, comm=comm) + error = _FakeMpiException(999, "forced startup failure") + record_entered = threading.Event() + release_record = threading.Event() + original_record_error = broadcaster._record_error + + def blocked_record_error(recorded_error: BaseException) -> None: + record_entered.set() + release_record.wait() + original_record_error(recorded_error) + + def failing_progress_loop() -> None: + broadcaster._fail_closed_immediately(error) + broadcaster._ready_event.set() + + monkeypatch.setattr(broadcaster, "_record_error", blocked_record_error) + monkeypatch.setattr(broadcaster, "_progress_loop", failing_progress_loop) + start_errors: list[BaseException] = [] + + def start_broadcaster() -> None: + try: + broadcaster.start() + except BaseException as start_error: + start_errors.append(start_error) + + start_thread = threading.Thread(target=start_broadcaster) + start_thread.start() + _wait_until(record_entered.is_set, description="blocked terminal publication") + assert broadcaster._progress_failed.is_set() is False + assert broadcaster.last_error is None + release_record.set() + start_thread.join(timeout=1.0) + + assert start_thread.is_alive() is False + assert len(start_errors) == 1 + assert "failed during startup" in str(start_errors[0]) + assert broadcaster.last_error is error + assert broadcaster._progress_failed.is_set() + assert comm.revoke_calls == 1 + + +def test_startup_timeout_fails_closed_before_progress_resumes() -> None: + receive_entered = threading.Event() + release_receive = threading.Event() + + def blocking_receive(_source: int) -> _FakeRequest: + receive_entered.set() + release_receive.wait() + return _FakeRequest() + + config = MpiFtSubcommConfig( + poll_interval_sec=0.001, + startup_timeout_sec=0.02, + stop_timeout_sec=0.5, + reconcile_timeout_sec=0.2, + abort_timeout_sec=0.02, + ) + comm = _FakeComm( + size=2, + receive_factory=blocking_receive, + ulfm_probe_error=NotImplementedError(), + ) + broadcaster = MpiFtSubcomm( + _FakeMapping.ep_world(2), + DetectedRankState(2), + config, + comm=comm, + mpi_module=_FakeMPI(), + ) + start_errors: list[BaseException] = [] + + def start_broadcaster() -> None: + try: + broadcaster.start() + except BaseException as error: + start_errors.append(error) + + start_thread = threading.Thread(target=start_broadcaster) + start_thread.start() + _wait_until(receive_entered.is_set, description="blocked startup") + start_thread.join(timeout=1.0) + + assert start_thread.is_alive() is False + assert len(start_errors) == 1 + assert isinstance(start_errors[0], TimeoutError) + assert "did not start" in str(start_errors[0]) + assert broadcaster._abort_requested.is_set() is False + assert isinstance(broadcaster.last_error, TimeoutError) + # The progress thread is still blocked in Irecv, so the caller thread must + # execute the no-ULFM fail-stop fallback before start() returns. + assert release_receive.is_set() is False + assert comm.abort_calls == 1 + release_receive.set() + broadcaster.stop(timeout=0.5) + assert comm.abort_calls == 1 + + +def test_startup_timeout_stops_callback_before_blocked_receive_is_released() -> None: + receive_entered = threading.Event() + release_receive = threading.Event() + + def blocking_receive(_source: int) -> _FakeRequest: + receive_entered.set() + release_receive.wait() + return _FakeRequest() + + config = MpiFtSubcommConfig( + poll_interval_sec=0.001, + startup_timeout_sec=0.02, + stop_timeout_sec=0.5, + reconcile_timeout_sec=0.2, + unattributed_error_timeout_sec=0.2, + abort_timeout_sec=0.02, + ) + comm = _FakeComm( + size=2, + receive_factory=blocking_receive, + ulfm_probe_error=NotImplementedError(), + ) + broadcaster = MpiFtSubcomm( + _FakeMapping.ep_world(2), + DetectedRankState(2), + config, + lambda _failed, _source, _when: None, + comm=comm, + mpi_module=_FakeMPI(), + ) + + try: + with pytest.raises(TimeoutError, match="did not start"): + broadcaster.start() + assert receive_entered.is_set() + assert release_receive.is_set() is False + assert broadcaster._callback_thread is not None + assert broadcaster._callback_thread.is_alive() is False + assert comm.abort_calls == 1 + finally: + release_receive.set() + broadcaster.stop(timeout=0.5) + assert comm.abort_calls == 1 + + +def test_reconciliation_gates_fail_closed_outside_running_lifecycle() -> None: + broadcaster, _, _, _ = _make_broadcaster(size=2) + assert broadcaster.health_is_reconciled() is False + assert broadcaster.failure_is_reconciled(1) is False + + broadcaster.start() + assert broadcaster.pre_failover(1) is True + _wait_until( + lambda: broadcaster.failure_is_reconciled(1), + description="single-survivor reconciliation", + ) + assert broadcaster.health_is_reconciled() is True + + broadcaster.stop() + assert broadcaster.failure_is_reconciled(1) is False + assert broadcaster.health_is_reconciled() is False + + +@pytest.mark.parametrize("timeout", [0.0, -1.0, float("nan"), float("inf"), float("-inf")]) +def test_stop_requires_positive_finite_timeout(timeout: float) -> None: + broadcaster, _, _, _ = _make_broadcaster(size=1) + with pytest.raises(ValueError, match="finite and > 0"): + broadcaster.stop(timeout) + broadcaster.start() + broadcaster.stop() + + +def test_stop_is_bounded_when_mpi_test_is_stuck() -> None: + blocking_request = _BlockingRequest() + comm = _FakeComm(size=2, receive_factory=lambda _source: blocking_request) + broadcaster, _, _, _ = _make_broadcaster(size=2, comm=comm) + broadcaster.start() + _wait_until(blocking_request.entered_test.is_set, description="blocked MPI Test") + + start = time.monotonic() + with pytest.raises(TimeoutError, match="did not stop"): + broadcaster.stop(timeout=0.01) + assert time.monotonic() - start < 0.2 + assert blocking_request.release_test.is_set() is False + assert comm.revoke_calls == 1 + assert comm.abort_calls == 0 + + blocking_request.release_test.set() + broadcaster.stop(timeout=0.5) + assert comm.revoke_calls == 1 + assert comm.abort_calls == 0 + + +def test_broadcast_rejects_report_after_stop_begins() -> None: + blocking_request = _BlockingRequest() + comm = _FakeComm(size=2, receive_factory=lambda _source: blocking_request) + broadcaster, health, _, _ = _make_broadcaster(size=2, comm=comm) + broadcaster.start() + _wait_until(blocking_request.entered_test.is_set, description="blocked MPI Test") + + stop_errors: list[BaseException] = [] + + def stop_broadcaster() -> None: + try: + broadcaster.stop(timeout=0.5) + except BaseException as error: + stop_errors.append(error) + + stop_thread = threading.Thread(target=stop_broadcaster) + stop_thread.start() + _wait_until(broadcaster._stop_event.is_set, description="stop transition") + + with pytest.raises(RuntimeError, match="state=stopping"): + broadcaster.broadcast_failure(1) + assert health.all_active() is True + + blocking_request.release_test.set() + stop_thread.join(timeout=1.0) + assert stop_thread.is_alive() is False + assert stop_errors == [] + + +def test_healthy_stop_cancels_preposted_receives() -> None: + broadcaster, health, comm, _ = _make_broadcaster(size=3) + broadcaster.start() + _wait_until( + lambda: comm.receive_count(1) == 1 and comm.receive_count(2) == 1, + description="preposted receives", + ) + receives = [comm.latest_receive(source).request for source in (1, 2)] + + assert health.all_active() is True + broadcaster.stop() + + assert [request.cancel_calls for request in receives] == [1, 1] + assert all(request.Test() is True for request in receives) + assert broadcaster._retained_requests == [] + + +def test_healthy_stop_cleanup_timeout_aborts_world_and_retains_request() -> None: + request = _NonCompletingCancelRequest() + comm = _FakeComm(size=2, receive_factory=lambda _source: request) + config = MpiFtSubcommConfig( + poll_interval_sec=0.001, + startup_timeout_sec=0.5, + stop_timeout_sec=0.05, + reconcile_timeout_sec=0.2, + ) + broadcaster = MpiFtSubcomm( + _FakeMapping.ep_world(2), + DetectedRankState(2), + config, + comm=comm, + mpi_module=_FakeMPI(), + ) + broadcaster.start() + _wait_until(lambda: comm.receive_count(1) == 1, description="preposted receive") + + start = time.monotonic() + broadcaster.stop() + assert time.monotonic() - start < 0.2 + assert request.cancel_calls == 1 + assert isinstance(broadcaster.last_error, TimeoutError) + assert broadcaster._retained_requests + assert comm.abort_calls == 1 + assert comm.revoke_calls == 0 + + +def test_poisoned_stop_keeps_progressing_until_failure_is_reconciled() -> None: + broadcaster, health, comm, _ = _make_broadcaster(size=3) + broadcaster.start() + _wait_until( + lambda: comm.receive_count(1) == 1 and comm.receive_count(2) == 1, + description="preposted receives", + ) + receives = [comm.latest_receive(source).request for source in (1, 2)] + assert broadcaster.broadcast_failure(2) is True + _wait_until(lambda: len(comm.sends) == 1, description="pending failure send") + send = comm.sends[0] + + assert health.get_failed_ranks() == frozenset({2}) + initial_test_calls = send.request.test_calls + stop_errors: list[BaseException] = [] + + def stop_broadcaster() -> None: + try: + broadcaster.stop() + except BaseException as error: + stop_errors.append(error) + + stop_thread = threading.Thread(target=stop_broadcaster) + stop_thread.start() + _wait_until(broadcaster._stop_event.is_set, description="stop request") + _wait_until( + lambda: send.request.test_calls > initial_test_calls, + description="continued send progress during stop", + ) + assert stop_thread.is_alive() is True + + send.request.complete() + comm.deliver(source=1, failed_rank=2) + stop_thread.join(timeout=1.0) + + assert stop_thread.is_alive() is False + assert stop_errors == [] + assert [request.cancel_calls for request in receives] == [0, 0] + assert send.request.cancel_calls == 0 + assert broadcaster.last_error is None + assert broadcaster._retained_requests + + +def test_poisoned_resources_outlive_broadcaster_object() -> None: + registry = ep_failure_broadcast._PROCESS_LIFETIME_REFS + initial_registry_size = len(registry) + + def retain_then_drop_owner() -> tuple[ + weakref.ReferenceType[_FakeRequest], + weakref.ReferenceType[np.ndarray], + weakref.ReferenceType[_FakeComm], + ]: + broadcaster, _, comm, _ = _make_broadcaster(size=3) + broadcaster.start() + _wait_until( + lambda: comm.receive_count(1) == 1 and comm.receive_count(2) == 1, + description="preposted receives", + ) + assert broadcaster.pre_failover(2) is True + _wait_until(lambda: len(comm.sends) == 1, description="pending send") + send = comm.sends[0] + send.request.complete() + comm.deliver(source=1, failed_rank=2) + _wait_until(broadcaster.health_is_reconciled, description="failure reconciliation") + retained_receive = comm.latest_receive(2) + references = ( + weakref.ref(retained_receive.request), + weakref.ref(retained_receive.buffer), + weakref.ref(comm), + ) + broadcaster.stop() + # The fake communicator owns test bookkeeping references that a real + # MPI communicator does not. Drop those so only the process-lifetime + # registry can keep the poisoned request and buffer reachable. + comm._receives.clear() + return references + + request_ref, buffer_ref, comm_ref = retain_then_drop_owner() + gc.collect() + try: + assert request_ref() is not None + assert buffer_ref() is not None + assert comm_ref() is not None + assert len(registry) > initial_registry_size + finally: + del registry[initial_registry_size:] + + +def test_startup_failure_retains_poisoned_comm_without_pending_requests() -> None: + registry = ep_failure_broadcast._PROCESS_LIFETIME_REFS + initial_registry_size = len(registry) + + def fail_first_receive(_source: int) -> _FakeRequest: + raise _FakeMpiException(999, "first Irecv failed") + + def start_then_drop_owner() -> weakref.ReferenceType[_FakeComm]: + comm = _FakeComm(size=2, receive_factory=fail_first_receive) + comm_ref = weakref.ref(comm) + broadcaster, _, _, _ = _make_broadcaster(size=2, comm=comm) + with pytest.raises(RuntimeError, match="failed during startup"): + broadcaster.start() + assert broadcaster._retained_requests == [] + assert isinstance(broadcaster.last_error, _FakeMpiException) + return comm_ref + + comm_ref = start_then_drop_owner() + gc.collect() + try: + assert comm_ref() is not None + assert len(registry) > initial_registry_size + assert any(reference is comm_ref() for reference in registry[initial_registry_size:]) + finally: + del registry[initial_registry_size:] + + +def test_constructor_requires_mpi_thread_multiple() -> None: + mpi = _FakeMPI(thread_level=_FakeMPI.THREAD_MULTIPLE - 1) + comm = _FakeComm(size=2) + with pytest.raises(RuntimeError, match="MPI.THREAD_MULTIPLE"): + MpiFtSubcomm( + _FakeMapping.ep_world(2), + DetectedRankState(2), + _TEST_CONFIG, + comm=comm, + mpi_module=mpi, + ) + assert comm.errhandler is None + + +def test_collectively_created_comm_uses_frozen_setup_and_lives_for_process( + monkeypatch: pytest.MonkeyPatch, +) -> None: + registry = ep_failure_broadcast._PROCESS_LIFETIME_REFS + initial_registry_size = len(registry) + comm = _FakeComm(size=2) + mpi = _FakeMPI() + + class SingleReadMapping: + def __init__(self) -> None: + self.rank_reads = 0 + + @property + def moe_ep_rank(self) -> int: + self.rank_reads += 1 + if self.rank_reads > 1: + raise RuntimeError("mapping rank was read after collective validation") + return 0 + + mapping = SingleReadMapping() + + def create_mpi_ft_subcomm( + setup_mapping: SingleReadMapping, + *, + health_size: int, + ) -> types.SimpleNamespace: + assert health_size == 2 + local_rank = setup_mapping.moe_ep_rank + comm.Set_errhandler(mpi.ERRORS_RETURN) + return types.SimpleNamespace( + comm=comm, + local_rank=local_rank, + ep_size=2, + ulfm_available=False, + ) + + distributed = types.ModuleType("tensorrt_llm._torch.distributed") + distributed.__path__ = [] + communicator = types.ModuleType("tensorrt_llm._torch.distributed.communicator") + communicator.create_mpi_ft_subcomm = create_mpi_ft_subcomm + distributed.communicator = communicator + monkeypatch.setitem(sys.modules, "tensorrt_llm._torch.distributed", distributed) + monkeypatch.setitem( + sys.modules, + "tensorrt_llm._torch.distributed.communicator", + communicator, + ) + + try: + broadcaster = MpiFtSubcomm( + mapping, + DetectedRankState(2), + _TEST_CONFIG, + mpi_module=mpi, + ) + + assert broadcaster._comm is comm + assert broadcaster._local_rank == 0 + assert broadcaster._ep_size == 2 + assert mapping.rank_reads == 1 + assert comm.errhandler_calls == 1 + assert comm.is_revoked_calls == 0 + assert any(reference is comm for reference in registry[initial_registry_size:]) + broadcaster.stop() + del broadcaster + gc.collect() + assert any(reference is comm for reference in registry[initial_registry_size:]) + finally: + del registry[initial_registry_size:] + + +def test_constructor_rejects_non_world_spanning_ep_group_for_mvp() -> None: + mapping = _FakeMapping( + world_size=4, + rank=0, + moe_ep_size=2, + moe_ep_rank=0, + moe_ep_group=(0, 1), + ) + comm = _FakeComm(size=2) + with pytest.raises(ValueError, match="spanning the full MPI world"): + MpiFtSubcomm( + mapping, + DetectedRankState(2), + _TEST_CONFIG, + comm=comm, + mpi_module=_FakeMPI(), + ) + assert comm.errhandler is None + + +@pytest.mark.parametrize( + "probe_error", + [ + NotImplementedError("ULFM not compiled"), + _FakeMpiException(_FakeMPI.ERR_UNSUPPORTED_OPERATION), + ], + ids=["not-implemented", "unsupported-operation"], +) +def test_ulfm_probe_falls_back_when_runtime_does_not_support_it( + probe_error: BaseException, +) -> None: + comm = _FakeComm(size=2, ulfm_probe_error=probe_error) + broadcaster, _, _, _ = _make_broadcaster(size=2, comm=comm) + assert broadcaster.ulfm_available is False + assert comm.is_revoked_calls == 1 + broadcaster.stop() + + +def test_ulfm_probe_propagates_unrelated_mpi_error() -> None: + registry = ep_failure_broadcast._PROCESS_LIFETIME_REFS + initial_registry_size = len(registry) + error = _FakeMpiException(_FakeMPI.ERR_UNKNOWN, "corrupt communicator") + comm = _FakeComm(size=2, ulfm_probe_error=error) + try: + with pytest.raises(_FakeMpiException, match="corrupt communicator") as raised: + _make_broadcaster(size=2, comm=comm) + assert raised.value is error + assert any(reference is comm for reference in registry[initial_registry_size:]) + finally: + del registry[initial_registry_size:] + + +def test_ulfm_probe_retains_original_error_when_error_class_inspection_raises() -> None: + registry = ep_failure_broadcast._PROCESS_LIFETIME_REFS + initial_registry_size = len(registry) + + class UnclassifiableMpiError(_FakeMpiException): + def Get_error_class(self) -> int: + raise RuntimeError("error-class lookup failed") + + probe_error = UnclassifiableMpiError( + _FakeMPI.ERR_UNSUPPORTED_OPERATION, + "unclassifiable ULFM probe failure", + ) + comm = _FakeComm(size=2, ulfm_probe_error=probe_error) + try: + with pytest.raises(UnclassifiableMpiError, match="unclassifiable") as raised: + _make_broadcaster(size=2, comm=comm) + assert raised.value is probe_error + assert any(reference is comm for reference in registry[initial_registry_size:]) + finally: + del registry[initial_registry_size:] + + +def test_constructor_rejects_and_retains_already_revoked_comm() -> None: + registry = ep_failure_broadcast._PROCESS_LIFETIME_REFS + initial_registry_size = len(registry) + + def construct_then_drop_owner() -> weakref.ReferenceType[_FakeComm]: + comm = _FakeComm(size=2, ulfm_revoked=True) + comm_ref = weakref.ref(comm) + with pytest.raises(RuntimeError, match="already revoked"): + _make_broadcaster(size=2, comm=comm) + return comm_ref + + comm_ref = construct_then_drop_owner() + gc.collect() + try: + assert comm_ref() is not None + assert len(registry) > initial_registry_size + finally: + del registry[initial_registry_size:] + + +def test_peer_failure_reconciles_without_revoking_control_plane() -> None: + callbacks: list[tuple[int, int, float]] = [] + broadcaster, health, comm, _ = _make_broadcaster( + size=3, + callback=lambda failed, source, when: callbacks.append((failed, source, when)), + ) + assert broadcaster.ulfm_available is True + broadcaster.start() + try: + _wait_until(lambda: comm.receive_count(1) == 1, description="peer receive") + comm.latest_receive(1).request.fail(_FakeMpiException(_FakeMPI.ERR_PROC_FAILED)) + + _wait_until(lambda: health.is_active(1) is False, description="peer failure classification") + _wait_until(lambda: len(comm.sends) == 1, description="peer failure relay") + + assert health.get_failed_ranks() == frozenset({1}) + _wait_until(lambda: len(callbacks) == 1, description="failed-peer callback") + assert [(failed, source) for failed, source, _ in callbacks] == [(1, 1)] + assert comm.sends[0].destination == 2 + assert comm.sends[0].message_kind == _FAILURE_MESSAGE + assert comm.sends[0].payload == 1 + assert comm.revoke_calls == 0 + assert broadcaster.last_error is None + + comm.sends[0].request.complete() + comm.deliver(source=2, failed_rank=1) + _wait_until(broadcaster.health_is_reconciled, description="failure reconciliation") + _wait_until(lambda: comm.receive_count(2) == 2, description="reposted receive") + assert broadcaster.health_is_reconciled() is True + assert broadcaster.last_error is None + assert comm.revoke_calls == 0 + finally: + broadcaster.stop() + + +def test_remote_revoke_after_local_reconciliation_still_fails_closed() -> None: + error = _FakeMpiException(_FakeMPI.ERR_REVOKED, "remote revoke") + comm = _FakeComm(size=3) + broadcaster, _, _, _ = _make_broadcaster(size=3, comm=comm) + broadcaster.start() + assert broadcaster.pre_failover(2) is True + _wait_until(lambda: len(comm.sends) == 1, description="failure relay") + comm.sends[0].request.complete() + _wait_until(lambda: comm.receive_count(1) == 1, description="initial peer receive") + comm.deliver(source=1, failed_rank=2) + _wait_until(broadcaster.health_is_reconciled, description="local reconciliation") + _wait_until( + lambda: comm.receive_count(1) == 2, + description="receive after failure report", + ) + + comm.latest_receive(1).request.fail(error) + _wait_until(lambda: broadcaster.last_error is error, description="terminal remote revoke") + assert broadcaster.health_is_reconciled() is False + assert comm.revoke_calls == 1 + with pytest.raises(RuntimeError, match="not running"): + broadcaster.pre_failover(1) + assert broadcaster._detected_state.get_failed_ranks() == frozenset({2}) + broadcaster.stop() + + +def test_premature_remote_revoke_fails_closed() -> None: + error = _FakeMpiException(_FakeMPI.ERR_REVOKED, "premature revoke") + comm = _FakeComm(size=3) + broadcaster, _, _, _ = _make_broadcaster(size=3, comm=comm) + broadcaster.start() + assert broadcaster.pre_failover(2) is True + _wait_until(lambda: comm.receive_count(1) == 1, description="initial peer receive") + + comm.latest_receive(1).request.fail(error) + _wait_until(lambda: broadcaster.last_error is error, description="terminal revoke") + assert broadcaster.health_is_reconciled() is False + assert comm.revoke_calls == 1 + broadcaster.stop() + + +def test_generic_mpi_error_is_terminal_and_revokes_once() -> None: + comm = _FakeComm(size=2) + broadcaster, health, _, _ = _make_broadcaster(size=2, comm=comm) + broadcaster.start() + try: + _wait_until(lambda: comm.receive_count(1) == 1, description="peer receive") + error = _FakeMpiException(999, "generic transport failure") + comm.latest_receive(1).request.fail(error) + + _wait_until(lambda: broadcaster.last_error is error, description="terminal MPI error") + _wait_until(lambda: comm.revoke_calls == 1, description="communicator revoke") + assert health.all_active() is True + with pytest.raises(RuntimeError, match="broadcaster is not running"): + broadcaster.broadcast_failure(1) + finally: + broadcaster.stop() + assert comm.revoke_calls == 1 + + +def test_proc_failed_is_classified_without_full_ulfm_support() -> None: + comm = _FakeComm(size=3, ulfm_probe_error=NotImplementedError()) + broadcaster, health, _, _ = _make_broadcaster(size=3, comm=comm) + assert broadcaster.ulfm_available is False + broadcaster.start() + try: + _wait_until(lambda: comm.receive_count(1) == 1, description="peer receive") + comm.latest_receive(1).request.fail(_FakeMpiException(_FakeMPI.ERR_PROC_FAILED)) + _wait_until(lambda: health.is_active(1) is False, description="peer failure classification") + _wait_until(lambda: len(comm.sends) == 1, description="peer failure relay") + + assert health.get_failed_ranks() == frozenset({1}) + assert comm.sends[0].destination == 2 + assert comm.sends[0].payload == 1 + assert broadcaster.last_error is None + assert comm.revoke_calls == 0 + comm.sends[0].request.complete() + comm.deliver(source=2, failed_rank=1) + _wait_until(broadcaster.health_is_reconciled, description="failure reconciliation") + finally: + broadcaster.stop() + + +def test_non_ulfm_generic_receive_error_keeps_other_sources_live() -> None: + callbacks: list[tuple[int, int, float]] = [] + comm = _FakeComm(size=3, ulfm_probe_error=NotImplementedError()) + broadcaster, health, _, _ = _make_broadcaster( + size=3, + comm=comm, + callback=lambda failed, source, when: callbacks.append((failed, source, when)), + ) + broadcaster.start() + try: + _wait_until( + lambda: comm.receive_count(1) == 1 and comm.receive_count(2) == 1, + description="fixed-source receives", + ) + failed_receive = comm.latest_receive(1).request + failed_receive.fail(_FakeMpiException(999, "implementation-specific peer error")) + _wait_until( + lambda: len(broadcaster._retained_requests) == 1, + description="retained failed receive", + ) + + assert health.all_active() is True + assert broadcaster.last_error is None + assert failed_receive.cancel_calls == 0 + + # Another survivor's rank-attributed detection remains usable even + # when the dead source's fixed receive failed with a generic error. + comm.deliver(source=2, failed_rank=1) + _wait_until(lambda: health.is_active(1) is False, description="detection report") + _wait_until(lambda: len(comm.sends) == 1, description="failure relay") + comm.sends[0].request.complete() + _wait_until(lambda: comm.receive_count(2) == 2, description="live-source receive repost") + _wait_until(lambda: len(callbacks) == 1, description="detection callback") + + assert [(failed, source) for failed, source, _ in callbacks] == [(1, 2)] + assert broadcaster.last_error is None + assert comm.revoke_calls == 0 + _wait_until(broadcaster.health_is_reconciled, description="failure reconciliation") + finally: + broadcaster.stop() + + +def test_non_ulfm_generic_error_after_watchdog_report_does_not_arm_stale_deadline() -> None: + config = MpiFtSubcommConfig( + poll_interval_sec=0.001, + startup_timeout_sec=0.5, + stop_timeout_sec=0.5, + reconcile_timeout_sec=0.2, + unattributed_error_timeout_sec=0.01, + abort_timeout_sec=0.02, + ) + comm = _FakeComm(size=3, ulfm_probe_error=NotImplementedError()) + broadcaster = MpiFtSubcomm( + _FakeMapping.ep_world(3), + DetectedRankState(3), + config, + comm=comm, + mpi_module=_FakeMPI(), + ) + broadcaster.start() + try: + _wait_until( + lambda: comm.receive_count(1) == 1 and comm.receive_count(2) == 1, + description="fixed-source receives", + ) + comm.deliver(source=2, failed_rank=1) + _wait_until(lambda: len(comm.sends) == 1, description="failure relay") + + # The transport notification can arrive after the watchdog report. It + # is already attributed to rank 1 and must not create a deadline that + # aborts an otherwise reconciled survivor set later. + failed_receive = comm.latest_receive(1).request + failed_receive.fail(_FakeMpiException(999, "late generic peer error")) + _wait_until( + lambda: ( + failed_receive in [pending.request for pending in broadcaster._retained_requests] + ), + description="retained explained receive", + ) + comm.sends[0].request.complete() + _wait_until(broadcaster.health_is_reconciled, description="failure reconciliation") + time.sleep(3 * config.unattributed_error_timeout_sec) + + assert broadcaster.last_error is None + assert broadcaster.health_is_reconciled() is True + assert comm.abort_calls == 0 + finally: + broadcaster.stop() + + +def test_unattributed_error_from_other_source_keeps_reconciliation_gate_closed() -> None: + config = MpiFtSubcommConfig( + poll_interval_sec=0.001, + startup_timeout_sec=0.5, + stop_timeout_sec=0.5, + reconcile_timeout_sec=0.5, + unattributed_error_timeout_sec=0.05, + abort_timeout_sec=0.02, + ) + comm = _FakeComm(size=4, ulfm_probe_error=NotImplementedError()) + broadcaster = MpiFtSubcomm( + _FakeMapping.ep_world(4), + DetectedRankState(4), + config, + comm=comm, + mpi_module=_FakeMPI(), + ) + broadcaster.start() + try: + _wait_until( + lambda: all(comm.receive_count(source) == 1 for source in (1, 2, 3)), + description="fixed-source receives", + ) + comm.deliver(source=1, failed_rank=3) + _wait_until(lambda: comm.receive_count(1) == 2, description="source 1 receive repost") + _wait_until(lambda: len(comm.sends) == 2, description="failure fanout") + + # Source 1 already reported rank 3, but a later generic error from + # source 1 is not explained by that different failed rank. + comm.latest_receive(1).request.fail(_FakeMpiException(999, "second transport error")) + _wait_until( + lambda: bool(broadcaster._unattributed_error_deadlines), + description="unattributed deadline", + ) + comm.deliver(source=2, failed_rank=3) + for send in comm.sends: + send.request.complete() + _wait_until(lambda: comm.receive_count(2) == 2, description="source 2 report") + _wait_until( + lambda: 3 in broadcaster._failure_fanout_complete, + description="completed failure fanout", + ) + + assert broadcaster.failure_is_reconciled(3) is False + assert broadcaster.health_is_reconciled() is False + _wait_until( + lambda: isinstance(broadcaster.last_error, TimeoutError), + description="unattributed transport deadline", + ) + _wait_until(lambda: comm.abort_calls == 1, description="terminal fallback") + finally: + broadcaster.stop() + + +def test_non_ulfm_unattributed_receive_error_aborts_after_bounded_deadlines( + monkeypatch: pytest.MonkeyPatch, +) -> None: + config = MpiFtSubcommConfig( + poll_interval_sec=0.001, + startup_timeout_sec=0.5, + stop_timeout_sec=0.5, + reconcile_timeout_sec=0.2, + unattributed_error_timeout_sec=0.01, + abort_timeout_sec=0.05, + ) + comm = _FakeComm(size=2, ulfm_probe_error=NotImplementedError()) + health = DetectedRankState(2) + broadcaster = MpiFtSubcomm( + _FakeMapping.ep_world(2), + health, + config, + comm=comm, + mpi_module=_FakeMPI(), + ) + terminal_request_entered = threading.Event() + release_terminal_request = threading.Event() + original_request_terminal_abort = broadcaster._request_terminal_abort + + def blocking_terminal_request( + error: BaseException, + payload: int, + reporter: int, + ) -> None: + terminal_request_entered.set() + assert release_terminal_request.wait(timeout=1.0) + original_request_terminal_abort(error, payload, reporter) + + monkeypatch.setattr(broadcaster, "_request_terminal_abort", blocking_terminal_request) + broadcaster.start() + _wait_until(lambda: comm.receive_count(1) == 1, description="fixed-source receive") + failed_receive = comm.latest_receive(1).request + failed_receive.fail(_FakeMpiException(999, "unattributed peer error")) + + _wait_until(terminal_request_entered.is_set, description="blocked terminal publication") + # The request is deliberately blocked before last_error is recorded. The + # lifecycle and protocol poison must already make the public gate fail + # closed, with no lock held that prevents the iteration thread from polling. + try: + assert broadcaster.last_error is None + assert broadcaster._lifecycle is ep_failure_broadcast._Lifecycle.FAILED + assert broadcaster._unattributed_error_deadlines[1] == float("inf") + assert broadcaster.health_is_reconciled() is False + finally: + release_terminal_request.set() + + _wait_until( + lambda: isinstance(broadcaster.last_error, TimeoutError), + description="unattributed transport deadline", + ) + assert "could not attribute" in str(broadcaster.last_error) + # Expiry remains protocol-visible for the rest of this communicator epoch. + assert 1 in broadcaster._unattributed_error_deadlines + assert broadcaster.health_is_reconciled() is False + _wait_until( + lambda: any(send.message_kind == _ABORT_MESSAGE for send in comm.sends), + description="terminal ABORT relay", + ) + _wait_until(lambda: comm.abort_calls == 1, description="bounded MPI_Abort fallback") + + assert health.all_active() is True + assert failed_receive.cancel_calls == 0 + assert comm.revoke_calls == 0 + broadcaster.stop() + + +def test_send_buffer_lives_until_request_test_completes() -> None: + comm = _FakeComm(size=3) + broadcaster, _, _, _ = _make_broadcaster(size=3, comm=comm) + broadcaster.start() + try: + assert broadcaster.broadcast_failure(2) is True + _wait_until(lambda: len(comm.sends) == 1, description="pending send") + record = comm.sends[0] + _wait_until(lambda: record.request.test_calls > 0, description="send Test polling") + + gc.collect() + assert record.buffer_ref() is not None + + record.request.complete() + + def buffer_released() -> bool: + gc.collect() + return record.buffer_ref() is None + + _wait_until(buffer_released, description="completed send buffer release") + comm.deliver(source=1, failed_rank=2) + _wait_until(broadcaster.health_is_reconciled, description="failure reconciliation") + finally: + broadcaster.stop() diff --git a/tests/unittest/_torch/pyexecutor/test_ep_failure_broadcast_mpi.py b/tests/unittest/_torch/pyexecutor/test_ep_failure_broadcast_mpi.py new file mode 100644 index 000000000000..d158fc8f7b29 --- /dev/null +++ b/tests/unittest/_torch/pyexecutor/test_ep_failure_broadcast_mpi.py @@ -0,0 +1,687 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. +"""Launcher for logical WideEP failure injection over a real MPI transport.""" + +import importlib.util +import math +import os +import re +import shutil +import signal +import subprocess # nosec B404 +import sys +from collections.abc import Mapping +from pathlib import Path +from typing import Protocol + +import pytest + +_SKIP_MARKER = "WIDEEP_MPI_SMOKE_SKIP:" +_PROPAGATION_MARKER = "WIDEEP_MPI_PROPAGATION" +_INVALID_TOPOLOGY_MARKER = "WIDEEP_MPI_INVALID_TOPOLOGY_OK" +_HEALTHY_LIFECYCLE_MARKER = "WIDEEP_MPI_HEALTHY_LIFECYCLE_OK" +_ABORT_WORLD_MARKER = "WIDEEP_MPI_ABORT_WORLD_OK" +_TERMINAL_READY_MARKER = "WIDEEP_MPI_TERMINAL_READY" +_TERMINAL_COMPLETE_MARKER = "WIDEEP_MPI_TERMINAL_COMPLETE" +_TERMINAL_DONE_MARKER = "WIDEEP_MPI_TERMINAL_DONE" +_TERMINAL_ERROR_MARKER = "WIDEEP_MPI_TERMINAL_ERROR" +_TERMINAL_ATEXIT_MARKER = "WIDEEP_MPI_TERMINAL_ATEXIT_RAN" +_WORKER_TIMEOUT_SEC = 30.0 +_REAP_TIMEOUT_SEC = 5.0 +_ALLOW_SKIP_ENV = "TLLM_ALLOW_MPI_FT_SMOKE_SKIP" +_HEALTHY_MODE = "healthy" +_TERMINAL_MODE = "terminal" +_PROPAGATION_PATTERN = re.compile( + rf"^{_PROPAGATION_MARKER} world_size=(\d+) " + r"elapsed_ms=(\S+) target_ms=(\S+) target_met=(True|False)$", + re.MULTILINE, +) +_TERMINAL_READY_PATTERN = re.compile(rf"^{_TERMINAL_READY_MARKER} rank=(\d+) world_size=(\d+)$") +_TERMINAL_COMPLETE_PATTERN = re.compile(rf"^{_TERMINAL_COMPLETE_MARKER} world_size=(\d+)$") +_TERMINAL_DONE_PATTERN = re.compile(rf"^{_TERMINAL_DONE_MARKER} rank=(\d+) world_size=(\d+)$") + + +class _ReapableProcess(Protocol): + def kill(self) -> None: ... + + def wait(self, timeout: float | None = None) -> int: ... + + +def _mpi_launcher() -> str | None: + return shutil.which("mpiexec") or shutil.which("mpirun") + + +def _smoke_is_required(environment: Mapping[str, str] | None = None) -> bool: + """Return whether missing MPI prerequisites must fail this invocation.""" + environment = os.environ if environment is None else environment + return environment.get(_ALLOW_SKIP_ENV) != "1" + + +def _handle_missing_prerequisite(reason: str, *, required: bool) -> None: + if required: + pytest.fail( + f"WideEP MPI FT smoke is required but cannot run: {reason}", + pytrace=False, + ) + pytest.skip(reason) + + +def _subprocess_output_text(output: str | bytes | None) -> str: + if output is None: + return "" + if isinstance(output, bytes): + return output.decode("utf-8", errors="replace") + return output + + +def _merge_subprocess_output(*outputs: str | bytes | None) -> str: + """Preserve timeout output without duplicating communicate's cumulative data.""" + merged = "" + for output in outputs: + text = _subprocess_output_text(output) + if not text or text in merged: + continue + if merged and merged in text: + merged = text + else: + merged += text + return merged + + +def _validate_rank_marker_set( + output: str, + expected_world_size: int, + marker: str, + pattern: re.Pattern[str], +) -> None: + """Require one exact structured marker from every expected rank.""" + marker_lines = [line for line in output.splitlines() if marker in line] + # MPI launchers may prefix a forwarded rank-0 stdout line. The worker emits + # all success evidence from one canonical stream, so strip only the prefix + # and retain strict validation of the marker payload itself. + marker_payloads = [line[line.index(marker) :] for line in marker_lines] + matches = [pattern.fullmatch(payload) for payload in marker_payloads] + malformed = [line for line, match in zip(marker_lines, matches) if match is None] + reported = [ + (int(match.group(1)), int(match.group(2))) for match in matches if match is not None + ] + expected_ranks = list(range(expected_world_size)) + reported_ranks = sorted(rank for rank, _world_size in reported) + invalid_world_sizes = sorted( + {world_size for _rank, world_size in reported if world_size != expected_world_size} + ) + if malformed or invalid_world_sizes or reported_ranks != expected_ranks: + pytest.fail( + f"expected exactly one {marker} from ranks " + f"{expected_ranks} with world_size={expected_world_size}, got " + f"markers={reported}, malformed={malformed}, " + f"invalid_world_sizes={invalid_world_sizes}:\n{output}", + pytrace=False, + ) + + +def _validate_terminal_completion(output: str, expected_world_size: int) -> None: + """Require exact READY/DONE sets around one global completion marker.""" + if _TERMINAL_ERROR_MARKER in output: + pytest.fail(f"terminal MPI worker reported an error:\n{output}", pytrace=False) + if _TERMINAL_ATEXIT_MARKER in output: + pytest.fail( + "terminal MPI worker entered Python atexit/mpi4py Finalize instead " + f"of the intentional os._exit path:\n{output}", + pytrace=False, + ) + + _validate_rank_marker_set( + output, + expected_world_size, + _TERMINAL_READY_MARKER, + _TERMINAL_READY_PATTERN, + ) + + completion_lines = [line for line in output.splitlines() if _TERMINAL_COMPLETE_MARKER in line] + completion_payloads = [ + line[line.index(_TERMINAL_COMPLETE_MARKER) :] for line in completion_lines + ] + completion_matches = [ + _TERMINAL_COMPLETE_PATTERN.fullmatch(payload) for payload in completion_payloads + ] + if ( + len(completion_matches) != 1 + or completion_matches[0] is None + or int(completion_matches[0].group(1)) != expected_world_size + ): + pytest.fail( + f"expected exactly one valid {_TERMINAL_COMPLETE_MARKER} with " + f"world_size={expected_world_size}, got {completion_lines}:\n{output}", + pytrace=False, + ) + _validate_rank_marker_set( + output, + expected_world_size, + _TERMINAL_DONE_MARKER, + _TERMINAL_DONE_PATTERN, + ) + + +def _validate_worker_result( + returncode: int, + output: str, + *, + required: bool, + mode: str = _HEALTHY_MODE, + expected_world_size: int | None = None, +) -> None: + """Validate normal-finalize and intentional no-finalize launcher results.""" + # A skip marker never excuses a launcher failure. Required workers return a + # nonzero status instead of printing this marker, but retain this ordering + # so a malformed optional worker cannot hide a launch error. + if returncode != 0 and (mode != _TERMINAL_MODE or _SKIP_MARKER in output): + pytest.fail(output or f"MPI smoke launcher exited with status {returncode}", pytrace=False) + if _SKIP_MARKER not in output: + if mode == _TERMINAL_MODE: + if expected_world_size is None: + raise ValueError("terminal MPI smoke validation requires expected_world_size") + # Open MPI reports a nonzero launcher status when workers + # intentionally bypass MPI_Finalize via os._exit(). Structured + # rank readiness plus global completion distinguishes that expected + # status from a real worker failure or MPI_Abort. + _validate_terminal_completion(output, expected_world_size) + return + if required: + pytest.fail( + f"WideEP MPI FT smoke is required but the worker skipped:\n{output.strip()}", + pytrace=False, + ) + pytest.skip(output.strip()) + + +def _parse_propagation_metric(output: str, expected_world_size: int) -> tuple[float, bool]: + """Validate and return the functional smoke's structured latency metric.""" + metric_lines = [line for line in output.splitlines() if _PROPAGATION_MARKER in line] + metric_payloads = [line[line.index(_PROPAGATION_MARKER) :] for line in metric_lines] + matches = [_PROPAGATION_PATTERN.fullmatch(payload) for payload in metric_payloads] + if len(matches) != 1 or matches[0] is None: + pytest.fail( + f"expected exactly one valid {_PROPAGATION_MARKER} line, got {metric_lines}:\n{output}", + pytrace=False, + ) + match = matches[0] + assert match is not None + world_size = int(match.group(1)) + try: + elapsed_ms = float(match.group(2)) + target_ms = float(match.group(3)) + except ValueError: + pytest.fail(f"invalid propagation metric:\n{match.group(0)}", pytrace=False) + target_met = match.group(4) == "True" + if world_size != expected_world_size: + pytest.fail( + f"propagation metric reported world_size={world_size}, expected {expected_world_size}", + pytrace=False, + ) + if not math.isfinite(elapsed_ms) or elapsed_ms < 0: + pytest.fail(f"invalid elapsed_ms={elapsed_ms}", pytrace=False) + if not math.isfinite(target_ms) or target_ms != 100.0: + pytest.fail(f"unexpected target_ms={target_ms}, expected 100", pytrace=False) + if target_met is not (elapsed_ms < target_ms): + pytest.fail( + f"target_met={target_met} is inconsistent with " + f"elapsed_ms={elapsed_ms} and target_ms={target_ms}", + pytrace=False, + ) + return elapsed_ms, target_met + + +def _validate_propagation_budget(output: str, expected_world_size: int) -> float: + """Require the real-MPI logical-failure fanout to meet its 100 ms budget.""" + elapsed_ms, target_met = _parse_propagation_metric(output, expected_world_size) + if not target_met: + pytest.fail( + f"WideEP MPI propagation exceeded the <100 ms budget: elapsed_ms={elapsed_ms}", + pytrace=False, + ) + return elapsed_ms + + +def _kill_and_reap_launcher(process: _ReapableProcess, timeout: float) -> bool: + """Kill the launcher and make one bounded attempt to reap it.""" + try: + process.kill() + except ProcessLookupError: + pass + try: + process.wait(timeout=timeout) + except subprocess.TimeoutExpired: + return False + return True + + +@pytest.mark.parametrize( + ("environment", "expected"), + [ + ({}, True), + ({_ALLOW_SKIP_ENV: "0"}, True), + ({_ALLOW_SKIP_ENV: "1"}, False), + ], +) +def test_mpi_ft_smoke_required_mode(environment: dict[str, str], expected: bool) -> None: + assert _smoke_is_required(environment) is expected + + +def test_mpi_ft_smoke_missing_prerequisite_is_optional_locally() -> None: + with pytest.raises(pytest.skip.Exception, match="missing launcher"): + _handle_missing_prerequisite("missing launcher", required=False) + + +def test_mpi_ft_smoke_missing_prerequisite_fails_when_required() -> None: + with pytest.raises(pytest.fail.Exception, match="required.*missing launcher"): + _handle_missing_prerequisite("missing launcher", required=True) + + +def test_mpi_ft_smoke_launcher_failure_precedes_skip_marker() -> None: + with pytest.raises(pytest.fail.Exception, match="launcher failed"): + _validate_worker_result( + 1, + f"launcher failed\n{_SKIP_MARKER} unsupported thread level", + required=False, + ) + + +def test_mpi_ft_smoke_worker_skip_fails_when_required() -> None: + with pytest.raises(pytest.fail.Exception, match="worker skipped"): + _validate_worker_result( + 0, + f"{_SKIP_MARKER} unsupported thread level", + required=True, + ) + + +def test_mpi_ft_smoke_worker_skip_is_optional_when_opted_out() -> None: + with pytest.raises(pytest.skip.Exception, match="unsupported thread level"): + _validate_worker_result( + 0, + f"{_SKIP_MARKER} unsupported thread level", + required=False, + ) + + +def _terminal_completion_output(world_size: int) -> str: + ready = [ + f"{_TERMINAL_READY_MARKER} rank={rank} world_size={world_size}" + for rank in range(world_size) + ] + done = [ + f"{_TERMINAL_DONE_MARKER} rank={rank} world_size={world_size}" for rank in range(world_size) + ] + return "\n".join([*ready, f"{_TERMINAL_COMPLETE_MARKER} world_size={world_size}", *done]) + + +@pytest.mark.parametrize("returncode", [0, 1, 137], ids=["zero", "open-mpi", "signal-style"]) +def test_mpi_ft_terminal_result_accepts_complete_evidence(returncode: int) -> None: + _validate_worker_result( + returncode, + _terminal_completion_output(4), + required=True, + mode=_TERMINAL_MODE, + expected_world_size=4, + ) + + +@pytest.mark.parametrize( + "output", + [ + "", + "\n".join( + [ + f"{_TERMINAL_READY_MARKER} rank=0 world_size=2", + f"{_TERMINAL_COMPLETE_MARKER} world_size=2", + ] + ), + "\n".join( + [ + f"{_TERMINAL_READY_MARKER} rank=0 world_size=2", + f"{_TERMINAL_READY_MARKER} rank=0 world_size=2", + f"{_TERMINAL_READY_MARKER} rank=1 world_size=2", + f"{_TERMINAL_COMPLETE_MARKER} world_size=2", + ] + ), + "\n".join( + [ + f"{_TERMINAL_READY_MARKER} rank=0 world_size=3", + f"{_TERMINAL_READY_MARKER} rank=1 world_size=3", + f"{_TERMINAL_COMPLETE_MARKER} world_size=3", + ] + ), + ], + ids=["missing", "incomplete", "duplicate", "wrong-world-size"], +) +def test_mpi_ft_terminal_result_rejects_incomplete_rank_set(output: str) -> None: + with pytest.raises(pytest.fail.Exception, match="expected exactly one"): + _validate_worker_result( + 1, + output, + required=True, + mode=_TERMINAL_MODE, + expected_world_size=2, + ) + + +@pytest.mark.parametrize( + "completion_suffix", + [ + "", + "\n".join( + [ + f"{_TERMINAL_COMPLETE_MARKER} world_size=2", + f"{_TERMINAL_COMPLETE_MARKER} world_size=2", + ] + ), + f"{_TERMINAL_COMPLETE_MARKER} world_size=3", + f"{_TERMINAL_COMPLETE_MARKER} world_size=2 trailing-garbage", + f"{_TERMINAL_COMPLETE_MARKER} world_size=2\n{_TERMINAL_ERROR_MARKER} rank=0", + ], + ids=["missing", "duplicate", "wrong-world-size", "malformed", "worker-error"], +) +def test_mpi_ft_terminal_result_rejects_invalid_global_completion( + completion_suffix: str, +) -> None: + ready = "\n".join(f"{_TERMINAL_READY_MARKER} rank={rank} world_size=2" for rank in range(2)) + done = "\n".join(f"{_TERMINAL_DONE_MARKER} rank={rank} world_size=2" for rank in range(2)) + output = "\n".join(part for part in (ready, completion_suffix, done) if part) + with pytest.raises(pytest.fail.Exception): + _validate_worker_result( + 1, + output, + required=True, + mode=_TERMINAL_MODE, + expected_world_size=2, + ) + + +def test_mpi_ft_terminal_result_rejects_crash_before_done() -> None: + ready_and_complete = "\n".join( + [ + *(f"{_TERMINAL_READY_MARKER} rank={rank} world_size=2" for rank in range(2)), + f"{_TERMINAL_COMPLETE_MARKER} world_size=2", + ] + ) + with pytest.raises(pytest.fail.Exception, match=_TERMINAL_DONE_MARKER): + _validate_worker_result( + 137, + ready_and_complete, + required=True, + mode=_TERMINAL_MODE, + expected_world_size=2, + ) + + +@pytest.mark.parametrize( + "done_suffix", + [ + f"{_TERMINAL_DONE_MARKER} rank=0 world_size=2", + "\n".join( + [ + f"{_TERMINAL_DONE_MARKER} rank=0 world_size=2", + f"{_TERMINAL_DONE_MARKER} rank=0 world_size=2", + f"{_TERMINAL_DONE_MARKER} rank=1 world_size=2", + ] + ), + "\n".join( + [ + f"{_TERMINAL_DONE_MARKER} rank=0 world_size=3", + f"{_TERMINAL_DONE_MARKER} rank=1 world_size=3", + ] + ), + "\n".join( + [ + f"{_TERMINAL_DONE_MARKER} rank=0 world_size=2", + f"{_TERMINAL_DONE_MARKER} rank=1 world_size=2 trailing-garbage", + ] + ), + ], + ids=["incomplete", "duplicate", "wrong-world-size", "malformed"], +) +def test_mpi_ft_terminal_result_rejects_invalid_done_set(done_suffix: str) -> None: + ready = "\n".join(f"{_TERMINAL_READY_MARKER} rank={rank} world_size=2" for rank in range(2)) + output = "\n".join([ready, f"{_TERMINAL_COMPLETE_MARKER} world_size=2", done_suffix]) + with pytest.raises(pytest.fail.Exception, match=_TERMINAL_DONE_MARKER): + _validate_worker_result( + 1, + output, + required=True, + mode=_TERMINAL_MODE, + expected_world_size=2, + ) + + +def test_mpi_ft_terminal_result_accepts_launcher_prefixes_and_noise() -> None: + output = "\n".join( + [ + "unrelated launcher diagnostic", + *(f"[rank0]: {line}" for line in _terminal_completion_output(2).splitlines()), + "another unrelated diagnostic", + ] + ) + + _validate_worker_result( + 1, + output, + required=True, + mode=_TERMINAL_MODE, + expected_world_size=2, + ) + + +def test_mpi_ft_terminal_result_rejects_python_atexit_evidence() -> None: + output = "\n".join( + [ + _terminal_completion_output(2), + f"{_TERMINAL_ATEXIT_MARKER} pid=123", + ] + ) + + with pytest.raises(pytest.fail.Exception, match="atexit/mpi4py Finalize"): + _validate_worker_result( + 0, + output, + required=True, + mode=_TERMINAL_MODE, + expected_world_size=2, + ) + + +def test_mpi_ft_terminal_launcher_failure_precedes_skip_marker() -> None: + with pytest.raises(pytest.fail.Exception, match="launcher failed"): + _validate_worker_result( + 1, + f"launcher failed\n{_SKIP_MARKER} unsupported thread level", + required=False, + mode=_TERMINAL_MODE, + expected_world_size=2, + ) + + +@pytest.mark.parametrize( + ("outputs", "expected"), + [ + ((None, "complete"), "complete"), + ((b"partial ", b"partial complete"), "partial complete"), + (("first\n", "second\n"), "first\nsecond\n"), + ((b"bad-\xff",), "bad-\ufffd"), + ], + ids=["none", "cumulative", "incremental", "bytes-replacement"], +) +def test_merge_subprocess_output_preserves_timeout_diagnostics( + outputs: tuple[str | bytes | None, ...], expected: str +) -> None: + assert _merge_subprocess_output(*outputs) == expected + + +def test_parse_propagation_metric_validates_target_without_requiring_it() -> None: + elapsed_ms, target_met = _parse_propagation_metric( + f"{_PROPAGATION_MARKER} world_size=4 elapsed_ms=125.5 target_ms=100.0 target_met=False", + 4, + ) + + assert elapsed_ms == 125.5 + assert target_met is False + + +def test_parse_propagation_metric_accepts_launcher_prefix_and_noise() -> None: + metric = f"{_PROPAGATION_MARKER} world_size=4 elapsed_ms=25.5 target_ms=100.0 target_met=True" + output = f"launcher diagnostic\n[rank0]: {metric}\nmore noise" + + assert _parse_propagation_metric(output, 4) == (25.5, True) + + +def test_validate_propagation_budget_rejects_semantically_valid_miss() -> None: + metric = f"{_PROPAGATION_MARKER} world_size=4 elapsed_ms=125.5 target_ms=100.0 target_met=False" + assert _parse_propagation_metric(metric, 4) == (125.5, False) + + with pytest.raises(pytest.fail.Exception, match="exceeded the <100 ms budget"): + _validate_propagation_budget(metric, 4) + + +@pytest.mark.parametrize("wait_times_out", [False, True], ids=["reaped", "still-running"]) +def test_kill_and_reap_launcher_is_bounded(wait_times_out: bool) -> None: + class FakeProcess: + def __init__(self) -> None: + self.kill_calls = 0 + self.wait_timeouts: list[float] = [] + + def kill(self) -> None: + self.kill_calls += 1 + + def wait(self, timeout: float | None = None) -> int: + assert timeout is not None + self.wait_timeouts.append(timeout) + if wait_times_out: + raise subprocess.TimeoutExpired("mpiexec", timeout) + return -signal.SIGKILL + + process = FakeProcess() + assert _kill_and_reap_launcher(process, 0.25) is (not wait_times_out) + assert process.kill_calls == 1 + assert process.wait_timeouts == [0.25] + + +@pytest.mark.parametrize( + "metric", + [ + "", + f"{_PROPAGATION_MARKER} world_size=2 elapsed_ms=nan target_ms=100 target_met=False", + f"{_PROPAGATION_MARKER} world_size=2 elapsed_ms=-1 target_ms=100 target_met=True", + f"{_PROPAGATION_MARKER} world_size=2 elapsed_ms=1 target_ms=99 target_met=True", + f"{_PROPAGATION_MARKER} world_size=2 elapsed_ms=101 target_ms=100 target_met=True", + f"{_PROPAGATION_MARKER} world_size=2 elapsed_ms=1 target_ms=100 target_met=True trailing", + ], +) +def test_parse_propagation_metric_rejects_malformed_semantics(metric: str) -> None: + with pytest.raises(pytest.fail.Exception): + _parse_propagation_metric(metric, 2) + + +@pytest.mark.parametrize("world_size", [2, 4], ids=["2-ranks", "4-ranks"]) +def test_ep_failure_broadcast_real_mpi(world_size: int) -> None: + """Exercise logical failure fanout and shutdown over a real MPI transport. + + The victim remains alive for portable ``COMM_WORLD`` result coordination; + actual process death and MPI error classification are outside this smoke. + """ + required = _smoke_is_required() + if importlib.util.find_spec("mpi4py") is None: + _handle_missing_prerequisite("mpi4py is not installed", required=required) + launcher = _mpi_launcher() + if launcher is None: + _handle_missing_prerequisite("mpiexec/mpirun is not installed", required=required) + + worker = Path(__file__).with_name("_ep_failure_broadcast_mpi_worker.py") + environment = os.environ.copy() + # Open MPI requires explicit opt-in when CI runs inside a root container. + # Other MPI implementations ignore these variables. + environment["OMPI_ALLOW_RUN_AS_ROOT"] = "1" + environment["OMPI_ALLOW_RUN_AS_ROOT_CONFIRM"] = "1" + environment["OMPI_MCA_rmaps_base_oversubscribe"] = "1" + environment["PYTHONUNBUFFERED"] = "1" + if not required: + environment[_ALLOW_SKIP_ENV] = "1" + + outputs: dict[str, str] = {} + for mode in (_HEALTHY_MODE, _TERMINAL_MODE): + command = [ + launcher, + "-n", + str(world_size), + sys.executable, + str(worker), + str(world_size), + mode, + ] + process = subprocess.Popen( # nosec B603 + command, + env=environment, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + start_new_session=True, + text=True, + ) + try: + output, _ = process.communicate(timeout=_WORKER_TIMEOUT_SEC) + except subprocess.TimeoutExpired as timeout_error: + partial_output = _subprocess_output_text(timeout_error.output) + try: + os.killpg(process.pid, signal.SIGKILL) + except ProcessLookupError: + # mpiexec may exit between communicate() timing out and cleanup. + pass + try: + output, _ = process.communicate(timeout=_REAP_TIMEOUT_SEC) + except subprocess.TimeoutExpired as reap_error: + output = _merge_subprocess_output(partial_output, reap_error.output) + reaped_after_kill = _kill_and_reap_launcher(process, _REAP_TIMEOUT_SEC) + if process.stdout is not None: + try: + process.stdout.close() + except OSError: + pass + reap_status = ( + "launcher was reaped after direct SIGKILL" + if reaped_after_kill + else "launcher still could not be reaped after direct SIGKILL" + ) + pytest.fail( + f"{world_size}-rank WideEP MPI {mode} smoke timed out after " + f"{_WORKER_TIMEOUT_SEC:.0f}s and could not be reaped within " + f"{_REAP_TIMEOUT_SEC:.0f}s; {reap_status}\n{output}" + ) + output = _merge_subprocess_output(partial_output, output) + pytest.fail( + f"{world_size}-rank WideEP MPI {mode} smoke timed out after " + f"{_WORKER_TIMEOUT_SEC:.0f}s\n{output}" + ) + + _validate_worker_result( + process.returncode, + output, + required=required, + mode=mode, + expected_world_size=world_size, + ) + outputs[mode] = output + + _validate_propagation_budget(outputs[_TERMINAL_MODE], world_size) + assert _INVALID_TOPOLOGY_MARKER in outputs[_HEALTHY_MODE] + assert _HEALTHY_LIFECYCLE_MARKER in outputs[_HEALTHY_MODE] + assert _ABORT_WORLD_MARKER in outputs[_TERMINAL_MODE] diff --git a/tests/unittest/_torch/pyexecutor/test_ep_metrics_stats.py b/tests/unittest/_torch/pyexecutor/test_ep_metrics_stats.py new file mode 100644 index 000000000000..5f8faa5d56d7 --- /dev/null +++ b/tests/unittest/_torch/pyexecutor/test_ep_metrics_stats.py @@ -0,0 +1,707 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. +"""Tests for the passive PyExecutor committed-membership metrics transport.""" + +import asyncio +import threading +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest +import zmq + +from tensorrt_llm._torch.modules.fused_moe.ep_group_health import ( + EP_GROUP_HEALTH_EXTRA_ATTR, + EPGroupHealth, +) +from tensorrt_llm._torch.modules.fused_moe.ep_metrics import pending_ep_health_metrics +from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor +from tensorrt_llm.executor.base_worker import BaseWorker +from tensorrt_llm.executor.proxy import GenerationExecutorProxy +from tensorrt_llm.executor.rpc import RPCCancelled, RPCClient, RPCError, RPCServer, RPCTimeout +from tensorrt_llm.executor.rpc.rpc_common import get_unique_ipc_addr +from tensorrt_llm.executor.rpc_proxy import GenerationExecutorRpcProxy +from tensorrt_llm.executor.rpc_proxy_mixin import RpcExecutorMixin +from tensorrt_llm.llmapi.llm import BaseLLM +from tensorrt_llm.serve.openai_server import OpenAIServer + +EP_HEALTH_SOURCE_EPOCH = "source-a" + + +def _make_executor(health: EPGroupHealth) -> PyExecutor: + executor = PyExecutor.__new__(PyExecutor) + executor._ep_health_model_config = SimpleNamespace( + extra_attrs={EP_GROUP_HEALTH_EXTRA_ATTR: health} + ) + executor.model_engine = SimpleNamespace( + model=SimpleNamespace(model_config=executor._ep_health_model_config) + ) + return executor + + +def test_pyexecutor_returns_committed_membership_snapshot() -> None: + health = EPGroupHealth(4) + health.mark_failed(2) + executor = _make_executor(health) + + assert executor._get_ep_health_stats() == { + "sourceEpoch": health.source_epoch, + "worldSize": 4, + "activeCount": 3, + "failedRanks": [2], + "generation": 1, + } + executor.enable_iter_perf_stats = False + assert executor.get_latest_iteration_stats() == [] + + +def test_pyexecutor_discovers_tracker_attached_after_initialization() -> None: + executor = PyExecutor.__new__(PyExecutor) + executor._ep_health_model_config = SimpleNamespace( + extra_attrs={EP_GROUP_HEALTH_EXTRA_ATTR: None} + ) + assert executor._get_ep_health_stats() == pending_ep_health_metrics() + + health = EPGroupHealth(4) + health.mark_failed(1) + executor._ep_health_model_config.extra_attrs[EP_GROUP_HEALTH_EXTRA_ATTR] = health + + stats = executor._get_ep_health_stats() + assert stats["sourceEpoch"] == health.source_epoch + assert stats["failedRanks"] == [1] + + +def test_pyexecutor_rejects_unqualified_multi_group_metrics() -> None: + executor = _make_executor(EPGroupHealth(4)) + executor._ep_health_model_config.mapping = SimpleNamespace( + moe_tp_size=1, + pp_size=2, + moe_cluster_size=1, + ) + + with pytest.raises(NotImplementedError, match="multi-group MoE topologies"): + executor._get_ep_health_stats() + + +def test_base_worker_exposes_engine_health_snapshot() -> None: + health = EPGroupHealth(4) + health.mark_failed(1) + worker = BaseWorker.__new__(BaseWorker) + worker.engine = _make_executor(health) + + assert worker.fetch_ep_health_stats()["failedRanks"] == [1] + + +def test_base_worker_reports_unsupported_engine_without_rpc_failure() -> None: + worker = BaseWorker.__new__(BaseWorker) + worker.engine = SimpleNamespace() + + assert worker.fetch_ep_health_stats() is None + + +def test_proxy_fetches_health_over_dedicated_rpc() -> None: + expected = { + "sourceEpoch": EP_HEALTH_SOURCE_EPOCH, + "worldSize": 4, + "activeCount": 3, + "failedRanks": [1], + "generation": 1, + } + proxy = GenerationExecutorProxy.__new__(GenerationExecutorProxy) + proxy.rpc_client = MagicMock() + proxy.rpc_client.fetch_ep_health_stats.return_value.remote.return_value = expected + + assert proxy._get_ep_health_stats() == expected + proxy.rpc_client.fetch_ep_health_stats.return_value.remote.assert_called_once_with(timeout=1.0) + + +def test_proxy_propagates_transient_rpc_failure() -> None: + proxy = GenerationExecutorProxy.__new__(GenerationExecutorProxy) + proxy.rpc_client = MagicMock() + proxy.rpc_client.fetch_ep_health_stats.return_value.remote.side_effect = TimeoutError + + with pytest.raises(TimeoutError): + proxy._get_ep_health_stats() + + +def test_rpc_proxy_fetches_health_over_dedicated_rpc() -> None: + expected = { + "sourceEpoch": EP_HEALTH_SOURCE_EPOCH, + "worldSize": 4, + "activeCount": 4, + "failedRanks": [], + "generation": 0, + } + proxy = GenerationExecutorRpcProxy.__new__(GenerationExecutorRpcProxy) + proxy.rpc_client = MagicMock() + proxy.rpc_client.fetch_ep_health_stats.return_value.remote.return_value = expected + + assert proxy._get_ep_health_stats(timeout=0.25) == expected + proxy.rpc_client.fetch_ep_health_stats.return_value.remote.assert_called_once_with(timeout=0.25) + + +def test_rpc_executor_mixin_forwards_health_for_ray_and_rpc_executors() -> None: + """Every RpcExecutorMixin consumer, including RayExecutor, gets telemetry.""" + expected = { + "sourceEpoch": EP_HEALTH_SOURCE_EPOCH, + "worldSize": 4, + "activeCount": 4, + "failedRanks": [], + "generation": 0, + } + executor = SimpleNamespace(rpc_client=MagicMock()) + executor.rpc_client.fetch_ep_health_stats.return_value.remote.return_value = expected + + assert RpcExecutorMixin._get_ep_health_stats(executor, timeout=0.5) == expected + executor.rpc_client.fetch_ep_health_stats.return_value.remote.assert_called_once_with( + timeout=0.5 + ) + + +def test_ep_health_snapshot_crosses_real_rpc_and_llm_bridge() -> None: + """Exercise the passive committed-membership path without a model worker.""" + + class HealthRpcSurface: + _get_ep_health_stats = BaseWorker._get_ep_health_stats + fetch_ep_health_stats = BaseWorker.fetch_ep_health_stats + + health = EPGroupHealth(4) + health.mark_failed(2) + worker = HealthRpcSurface() + worker.engine = _make_executor(health) + address = get_unique_ipc_addr() + + with RPCServer(worker) as rpc_server: + rpc_server.bind(address) + rpc_server.start() + with RPCClient(address, hmac_key=rpc_server.hmac_key) as rpc_client: + proxy = GenerationExecutorProxy.__new__(GenerationExecutorProxy) + proxy.rpc_client = rpc_client + llm = BaseLLM.__new__(BaseLLM) + llm._executor = proxy + + assert llm._get_ep_health_stats() == { + "sourceEpoch": health.source_epoch, + "worldSize": 4, + "activeCount": 3, + "failedRanks": [2], + "generation": 1, + } + + +def test_server_mount_metrics_registers_local_ep_health_collector(monkeypatch) -> None: + """The Prometheus endpoint registry includes the coherent local snapshot.""" + import prometheus_client + import prometheus_fastapi_instrumentator + from prometheus_client import multiprocess + from prometheus_client.core import GaugeMetricFamily + + class SnapshotCollector: + def collect(self): + metric = GaugeMetricFamily( + "trtllm_ep_health_available", + "Whether EP health telemetry is available.", + ) + metric.add_metric([], 1) + return [metric] + + metrics_collector = MagicMock() + metrics_collector.register_ep_health_metrics.side_effect = lambda registry: registry.register( + SnapshotCollector() + ) + monkeypatch.setattr(multiprocess, "MultiProcessCollector", MagicMock()) + instrumentator = MagicMock() + instrumentator.add.return_value = instrumentator + instrumentator.instrument.return_value = instrumentator + instrumentator.expose.return_value = instrumentator + monkeypatch.setattr( + prometheus_fastapi_instrumentator, "Instrumentator", MagicMock(return_value=instrumentator) + ) + monkeypatch.setattr(prometheus_client, "make_asgi_app", MagicMock(return_value=MagicMock())) + + server = OpenAIServer.__new__(OpenAIServer) + server.app = SimpleNamespace(routes=[]) + server.metrics_collector = metrics_collector + server.mount_metrics() + + registry = metrics_collector.register_ep_health_metrics.call_args.args[0] + output = prometheus_client.generate_latest(registry).decode() + assert "trtllm_ep_health_available 1.0" in output + assert len(server.app.routes) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "transient_error", + [ + TimeoutError("health read timed out"), + RPCTimeout("health RPC timed out"), + RPCCancelled("rank-0 worker is being replaced"), + RPCError("health transport failed", cause=ConnectionError("closed")), + zmq.ZMQError(zmq.ENOTCONN, "socket disconnected"), + RPCError( + "wrapped socket failure", + cause=zmq.ZMQError(zmq.ENOTCONN, "socket disconnected"), + ), + ], +) +async def test_server_retries_after_transient_first_read_failure( + transient_error: Exception, +) -> None: + expected = { + "sourceEpoch": EP_HEALTH_SOURCE_EPOCH, + "worldSize": 4, + "activeCount": 3, + "failedRanks": [1], + "generation": 1, + } + server = OpenAIServer.__new__(OpenAIServer) + server.metrics_collector = MagicMock() + server.metrics_collector.log_ep_health_stats.return_value = True + server._ep_health_collector_task = None + server._ep_health_stop_event = asyncio.Event() + + results = iter((transient_error, expected)) + + def read_health(): + result = next(results) + if isinstance(result, Exception): + raise result + server._ep_health_stop_event.set() + return result + + health_reader = MagicMock(side_effect=read_health) + server.generator = SimpleNamespace(_get_ep_health_stats=health_reader) + server._wait_for_ep_health_stop = AsyncMock(side_effect=[False, True]) + + await server._start_ep_health_collector() + assert server._ep_health_collector_task is not None + await server._ep_health_collector_task + + server.metrics_collector.log_ep_health_unavailable.assert_called_once_with() + server.metrics_collector.log_ep_health_stats.assert_called_once_with(expected) + assert health_reader.call_count == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "transient_error", + [ + TimeoutError("timed out"), + RPCCancelled("rank-0 worker is being replaced"), + zmq.ZMQError(zmq.ENOTCONN, "socket disconnected"), + RPCError( + "wrapped socket failure", + cause=zmq.ZMQError(zmq.ENOTCONN, "socket disconnected"), + ), + ], +) +async def test_server_loop_retries_transient_rank0_loss_and_recovers( + transient_error: Exception, +) -> None: + """Transient rank-0 loss clears availability, then accepts its replacement.""" + before_restart = { + "sourceEpoch": "source-before-restart", + "worldSize": 4, + "activeCount": 3, + "failedRanks": [1], + "generation": 2, + } + after_restart = { + "sourceEpoch": "source-after-restart", + "worldSize": 4, + "activeCount": 4, + "failedRanks": [], + "generation": 0, + } + server = OpenAIServer.__new__(OpenAIServer) + server.metrics_collector = MagicMock() + server.metrics_collector.log_ep_health_stats.return_value = True + server._ep_health_stop_event = asyncio.Event() + + results = iter((before_restart, transient_error, after_restart)) + + def read_health(): + result = next(results) + if isinstance(result, Exception): + raise result + return result + + health_reader = MagicMock(side_effect=read_health) + server.generator = SimpleNamespace(_get_ep_health_stats=health_reader) + server._wait_for_ep_health_stop = AsyncMock(side_effect=[False, False, True]) + + await asyncio.wait_for(server._ep_health_collector_loop(), timeout=1.0) + + assert health_reader.call_count == 3 + server.metrics_collector.log_ep_health_unavailable.assert_called_once_with() + assert [ + call.args[0] for call in server.metrics_collector.log_ep_health_stats.call_args_list + ] == [before_restart, after_restart] + assert server._wait_for_ep_health_stop.await_count == 3 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "read_result", + [ + None, + ValueError("invalid health configuration"), + RPCError("remote health configuration failed", cause=ValueError("invalid")), + RPCError( + "unsupported multi-group topology", + cause=NotImplementedError("multi-group"), + ), + ], +) +async def test_server_loop_stops_after_deterministic_health_loss(read_result) -> None: + """Unsupported and deterministic failures do not create a polling loop.""" + server = OpenAIServer.__new__(OpenAIServer) + server.metrics_collector = MagicMock() + if isinstance(read_result, Exception): + health_reader = MagicMock(side_effect=read_result) + else: + health_reader = MagicMock(return_value=read_result) + server.generator = SimpleNamespace(_get_ep_health_stats=health_reader) + + await asyncio.wait_for(server._ep_health_collector_loop(), timeout=0.5) + + health_reader.assert_called_once_with() + server.metrics_collector.log_ep_health_unavailable.assert_called_once_with() + server.metrics_collector.log_ep_health_stats.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "startup_error", + [ + ValueError("invalid health configuration"), + RPCError("remote health configuration failed", cause=ValueError("invalid")), + RPCError("method unavailable"), + ], +) +async def test_server_stops_after_deterministic_first_read_failure(startup_error) -> None: + server = OpenAIServer.__new__(OpenAIServer) + server.metrics_collector = MagicMock() + server._ep_health_collector_task = None + health_reader = MagicMock(side_effect=startup_error) + server.generator = SimpleNamespace(_get_ep_health_stats=health_reader) + + await server._start_ep_health_collector() + task = server._ep_health_collector_task + assert task is not None + await task + + health_reader.assert_called_once_with() + assert task.done() + server.metrics_collector.log_ep_health_unavailable.assert_called_once_with() + server.metrics_collector.log_ep_health_stats.assert_not_called() + await server._stop_ep_health_collector() + assert server._ep_health_collector_task is None + + +@pytest.mark.asyncio +async def test_server_retries_after_collector_rejects_payload() -> None: + snapshot = { + "sourceEpoch": EP_HEALTH_SOURCE_EPOCH, + "worldSize": 4, + "activeCount": 4, + "failedRanks": [], + "generation": 0, + } + server = OpenAIServer.__new__(OpenAIServer) + server.metrics_collector = MagicMock() + server.metrics_collector.log_ep_health_stats.side_effect = [False, True] + server._ep_health_stop_event = asyncio.Event() + + read_count = 0 + + def read_health(): + nonlocal read_count + read_count += 1 + if read_count == 2: + server._ep_health_stop_event.set() + return snapshot + + server.generator = SimpleNamespace(_get_ep_health_stats=MagicMock(side_effect=read_health)) + server._wait_for_ep_health_stop = AsyncMock(side_effect=[False, True]) + + await server._ep_health_collector_loop() + + assert server.generator._get_ep_health_stats.call_count == 2 + assert [ + args.args[0] for args in server.metrics_collector.log_ep_health_stats.call_args_list + ] == [snapshot, snapshot] + + +@pytest.mark.asyncio +async def test_server_retries_after_initial_payload_rejection() -> None: + snapshot = { + "sourceEpoch": EP_HEALTH_SOURCE_EPOCH, + "worldSize": 4, + "activeCount": 4, + "failedRanks": [], + "generation": 0, + } + server = OpenAIServer.__new__(OpenAIServer) + server.metrics_collector = MagicMock() + server.metrics_collector.log_ep_health_stats.side_effect = [False, True] + server._ep_health_collector_task = None + server._ep_health_stop_event = asyncio.Event() + + read_count = 0 + + def read_health(): + nonlocal read_count + read_count += 1 + if read_count == 2: + server._ep_health_stop_event.set() + return snapshot + + server.generator = SimpleNamespace(_get_ep_health_stats=MagicMock(side_effect=read_health)) + server._wait_for_ep_health_stop = AsyncMock(side_effect=[False, True]) + + await server._start_ep_health_collector() + + assert server._ep_health_collector_task is not None + await server._ep_health_collector_task + assert server.generator._get_ep_health_stats.call_count == 2 + assert server.metrics_collector.log_ep_health_stats.call_count == 2 + + +@pytest.mark.asyncio +async def test_server_does_not_poll_unsupported_backend() -> None: + server = OpenAIServer.__new__(OpenAIServer) + server.metrics_collector = MagicMock() + server._ep_health_collector_task = None + server.generator = SimpleNamespace(_get_ep_health_stats=MagicMock(return_value=None)) + + await server._start_ep_health_collector() + task = server._ep_health_collector_task + assert task is not None + await task + + assert task.done() + server.generator._get_ep_health_stats.assert_called_once_with() + server.metrics_collector.log_ep_health_stats.assert_not_called() + server.metrics_collector.log_ep_health_unavailable.assert_called_once_with() + await server._stop_ep_health_collector() + assert server._ep_health_collector_task is None + + +@pytest.mark.asyncio +async def test_server_retries_explicitly_pending_registration() -> None: + snapshot = { + "sourceEpoch": "source-after-attachment", + "worldSize": 4, + "activeCount": 4, + "failedRanks": [], + "generation": 0, + } + server = OpenAIServer.__new__(OpenAIServer) + server.metrics_collector = MagicMock() + server.metrics_collector.log_ep_health_stats.return_value = True + server._ep_health_collector_task = None + server._ep_health_stop_event = asyncio.Event() + server.generator = SimpleNamespace( + _get_ep_health_stats=MagicMock(side_effect=[pending_ep_health_metrics(), snapshot]) + ) + server._wait_for_ep_health_stop = AsyncMock(side_effect=[False, True]) + + await server._start_ep_health_collector() + assert server._ep_health_collector_task is not None + await server._ep_health_collector_task + + server.metrics_collector.log_ep_health_unavailable.assert_called_once_with() + server.metrics_collector.log_ep_health_stats.assert_called_once_with(snapshot) + + +@pytest.mark.asyncio +async def test_server_lifespan_starts_and_stops_health_collector(monkeypatch) -> None: + """The real server lifespan owns the health collector task.""" + generator = SimpleNamespace( + args=SimpleNamespace(enable_energy_metrics=False, enable_iter_perf_stats=False), + _get_ep_health_stats=MagicMock(), + shutdown=MagicMock(), + ) + monkeypatch.setattr(OpenAIServer, "_init_llm", lambda self, chat_template=None: None) + monkeypatch.setattr(OpenAIServer, "register_routes", lambda self: None) + server = OpenAIServer( + generator=generator, + model="test-model", + tool_parser=None, + server_role=None, + metadata_server_cfg=None, + ) + server.metrics_collector = MagicMock() + server.metrics_collector.log_ep_health_stats.return_value = True + collector_started = asyncio.Event() + + async def controlled_collector_loop() -> None: + collector_started.set() + await server._ep_health_stop_event.wait() + + server._ep_health_collector_loop = controlled_collector_loop + + async with server.app.router.lifespan_context(server.app): + await asyncio.wait_for(collector_started.wait(), timeout=1.0) + collector_task = server._ep_health_collector_task + assert collector_task is not None + assert not collector_task.done() + + assert collector_task.done() + assert not collector_task.cancelled() + generator._get_ep_health_stats.assert_not_called() + generator.shutdown.assert_called_once_with() + + +@pytest.mark.asyncio +async def test_server_lifespan_body_failure_still_shuts_down_once(monkeypatch) -> None: + """An exception thrown at the lifespan yield cannot bypass cleanup.""" + generator = SimpleNamespace( + args=SimpleNamespace(enable_energy_metrics=False, enable_iter_perf_stats=False), + _get_ep_health_stats=MagicMock(return_value=None), + shutdown=MagicMock(), + ) + monkeypatch.setattr(OpenAIServer, "_init_llm", lambda self, chat_template=None: None) + monkeypatch.setattr(OpenAIServer, "register_routes", lambda self: None) + server = OpenAIServer( + generator=generator, + model="test-model", + tool_parser=None, + server_role=None, + metadata_server_cfg=None, + ) + server.metrics_collector = MagicMock() + + with pytest.raises(RuntimeError, match="lifespan body failed"): + async with server.app.router.lifespan_context(server.app): + raise RuntimeError("lifespan body failed") + + assert server._ep_health_collector_task is None + generator.shutdown.assert_called_once_with() + + +@pytest.mark.asyncio +async def test_server_lifespan_drains_inflight_health_poll_before_shutdown(monkeypatch) -> None: + """Startup stays non-blocking, while shutdown drains the health RPC.""" + expected = { + "sourceEpoch": EP_HEALTH_SOURCE_EPOCH, + "worldSize": 4, + "activeCount": 4, + "failedRanks": [], + "generation": 0, + } + poll_started = threading.Event() + release_poll = threading.Event() + + def read_health(): + poll_started.set() + release_poll.wait(timeout=5.0) + return expected + + health_reader = MagicMock(side_effect=read_health) + generator = SimpleNamespace( + args=SimpleNamespace(enable_energy_metrics=False, enable_iter_perf_stats=False), + _get_ep_health_stats=health_reader, + shutdown=MagicMock(), + ) + monkeypatch.setattr(OpenAIServer, "_init_llm", lambda self, chat_template=None: None) + monkeypatch.setattr(OpenAIServer, "register_routes", lambda self: None) + server = OpenAIServer( + generator=generator, + model="test-model", + tool_parser=None, + server_role=None, + metadata_server_cfg=None, + ) + server.metrics_collector = MagicMock() + server.metrics_collector.log_ep_health_stats.return_value = True + + lifespan = server.app.router.lifespan_context(server.app) + await asyncio.wait_for(lifespan.__aenter__(), timeout=1.0) + exit_task = None + try: + assert await asyncio.wait_for(asyncio.to_thread(poll_started.wait, 1.0), timeout=2.0) + exit_task = asyncio.create_task(lifespan.__aexit__(None, None, None)) + await asyncio.wait_for(server._ep_health_stop_event.wait(), timeout=1.0) + + generator.shutdown.assert_not_called() + assert not exit_task.done() + finally: + release_poll.set() + if exit_task is None: + exit_task = asyncio.create_task(lifespan.__aexit__(None, None, None)) + await asyncio.wait_for(exit_task, timeout=1.0) + + health_reader.assert_called_once_with() + generator.shutdown.assert_called_once_with() + + +@pytest.mark.asyncio +async def test_server_lifespan_cancellation_still_drains_health_poll(monkeypatch) -> None: + """Cancelling lifespan cleanup cannot abandon a running health RPC thread.""" + snapshot = { + "sourceEpoch": EP_HEALTH_SOURCE_EPOCH, + "worldSize": 4, + "activeCount": 4, + "failedRanks": [], + "generation": 0, + } + poll_started = threading.Event() + release_poll = threading.Event() + + def read_health(): + poll_started.set() + release_poll.wait(timeout=5.0) + return snapshot + + generator = SimpleNamespace( + args=SimpleNamespace(enable_energy_metrics=False, enable_iter_perf_stats=False), + _get_ep_health_stats=MagicMock(side_effect=read_health), + shutdown=MagicMock(), + ) + monkeypatch.setattr(OpenAIServer, "_init_llm", lambda self, chat_template=None: None) + monkeypatch.setattr(OpenAIServer, "register_routes", lambda self: None) + server = OpenAIServer( + generator=generator, + model="test-model", + tool_parser=None, + server_role=None, + metadata_server_cfg=None, + ) + server.metrics_collector = MagicMock() + server.metrics_collector.log_ep_health_stats.return_value = True + + lifespan = server.app.router.lifespan_context(server.app) + await asyncio.wait_for(lifespan.__aenter__(), timeout=1.0) + exit_task = None + try: + assert await asyncio.wait_for(asyncio.to_thread(poll_started.wait, 1.0), timeout=2.0) + exit_task = asyncio.create_task(lifespan.__aexit__(None, None, None)) + await asyncio.wait_for(server._ep_health_stop_event.wait(), timeout=1.0) + for _ in range(3): + exit_task.cancel() + await asyncio.sleep(0) + generator.shutdown.assert_not_called() + assert not exit_task.done() + finally: + release_poll.set() + + assert exit_task is not None + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(exit_task, timeout=1.0) + generator._get_ep_health_stats.assert_called_once_with() + generator.shutdown.assert_called_once_with() diff --git a/tests/unittest/_torch/pyexecutor/test_error_classification.py b/tests/unittest/_torch/pyexecutor/test_error_classification.py new file mode 100644 index 000000000000..9e90b2b55493 --- /dev/null +++ b/tests/unittest/_torch/pyexecutor/test_error_classification.py @@ -0,0 +1,104 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +import importlib.util +import sys +from pathlib import Path + +import pytest + +_MODULE_PATH = ( + Path(__file__).resolve().parents[4] / "tensorrt_llm/_torch/pyexecutor/error_classification.py" +) +_SPEC = importlib.util.spec_from_file_location("error_classification_under_test", _MODULE_PATH) +assert _SPEC is not None +assert _SPEC.loader is not None +_ERROR_CLASSIFICATION = importlib.util.module_from_spec(_SPEC) +sys.modules[_SPEC.name] = _ERROR_CLASSIFICATION +_SPEC.loader.exec_module(_ERROR_CLASSIFICATION) + +ErrorBudget = _ERROR_CLASSIFICATION.ErrorBudget +classify_error = _ERROR_CLASSIFICATION.classify_error + + +@pytest.mark.parametrize( + "message", + [ + "CUDA context destroyed", + "cudaErrorContextIsDestroyed during collective setup", + "CUDA error: cudaErrorNoDevice on rank 37", + "NVRM: Xid 79, GPU has fallen off the bus", + "CUDA context is destroyed after device reset", + "CUDA context was destroyed after device reset", + "CUDA context is unrecoverable after device reset", + "Unrecoverable error in the engine", + "ncclCommAbort returned during communicator teardown", + "NCCL communicator was aborted during PP send", + "NVSHMEM peer unreachable for EP rank 12", + "MPI rank terminated with signal 9", + "MPI worker exited unexpectedly", + ], +) +def test_wide_ep_immediate_fatal_patterns(message): + assert classify_error(message) == "immediate_fatal" + + +@pytest.mark.parametrize( + "message", + [ + "NCCL timeout", + "AlltoAll timeout waiting for completion_flags[3][7]", + "AlltoAll watchdog timed out waiting for peer rank 7", + "completion_flags timeout", + "NCCL operation timed out during all_reduce", + "deep_ep buffer barrier hang", + "deepep buffer barrier hang", + "DeepEP Buffer.__del__ hung in intranode::barrier", + "Symmetric memory access violation reading dead peer", + "RDMA timeout on cross-node transfer", + "NIXL transfer failed: remote peer closed", + "NIXL transfer entered error state", + "NIXL transfer wait timed out after 5000 ms", + ], +) +def test_wide_ep_severe_patterns(message): + assert classify_error(message) == "severe" + + +@pytest.mark.parametrize( + "message", + [ + "AlltoAll slow path observed for rank 3", + "NCCL retry scheduled after transient transport hiccup", + "ECC correctable error threshold warning", + "Application marked the request unrecoverable but retryable", + ], +) +def test_wide_ep_nonfatal_signals_remain_transient(message): + assert classify_error(message) == "transient" + + +def test_error_budget_charges_wide_ep_severe_patterns(): + budget = ErrorBudget(recovery_rate=0.0) + + assert not budget.consume("NIXL transfer failed: remote peer closed") + assert budget.consume("AlltoAll timeout waiting for completion_flags[0][2]") + + +def test_error_budget_bypasses_budget_for_wide_ep_fatal_patterns(): + budget = ErrorBudget(recovery_rate=0.0) + + assert budget.consume("MPI rank terminated with signal 9") + assert budget.budget == 1.0 diff --git a/tests/unittest/metrics/test_collector.py b/tests/unittest/metrics/test_collector.py index e66975ca6f8e..efc0cd04f6fb 100644 --- a/tests/unittest/metrics/test_collector.py +++ b/tests/unittest/metrics/test_collector.py @@ -14,12 +14,14 @@ # limitations under the License. """Unit tests for MetricsCollector and process_req_perf_metrics.""" +import threading +import time from typing import Dict import pytest -from prometheus_client import REGISTRY +from prometheus_client import REGISTRY, CollectorRegistry, generate_latest, multiprocess -from tensorrt_llm.metrics.collector import MetricsCollector +from tensorrt_llm.metrics.collector import _MAX_RETIRED_EP_HEALTH_SOURCE_EPOCHS, MetricsCollector from tensorrt_llm.metrics.enums import MetricNames, RequestEventTiming from tensorrt_llm.metrics.perf_utils import process_req_perf_metrics @@ -47,6 +49,13 @@ def collector(): def _get_gauge_value(collector, metric_name: str): """Get the current value of a Prometheus gauge.""" + if metric_name.startswith("ep_"): + sample_name = f"trtllm_{metric_name}" + for metric in collector._ep_health_prometheus_collector.collect(): + for sample in metric.samples: + if sample.name == sample_name: + return sample.value + raise AssertionError(f"Missing EP health sample {sample_name}") metric = getattr(collector, metric_name) return metric.labels(**collector.labels)._value.get() @@ -57,6 +66,37 @@ def _get_counter_value(collector, metric_name: str): return metric.labels(**collector.labels)._value.get() +def _get_ep_rank_gauge_value(collector, rank: int): + for metric in collector._ep_health_prometheus_collector.collect(): + for sample in metric.samples: + if sample.name == "trtllm_ep_rank_active" and sample.labels[ + collector.labelname_ep_rank + ] == str(rank): + return sample.value + raise AssertionError(f"Missing EP rank sample {rank}") + + +def _get_ep_rank_samples(collector): + return [ + sample + for metric in collector._ep_health_prometheus_collector.collect() + for sample in metric.samples + if sample.name == "trtllm_ep_rank_active" + ] + + +def _get_ep_metric_samples(collector): + return [ + sample + for metric in collector._ep_health_prometheus_collector.collect() + for sample in metric.samples + ] + + +EP_HEALTH_SOURCE_A = "source-a" +EP_HEALTH_SOURCE_B = "source-b" + + SAMPLE_ITERATION_STATS = { "numActiveRequests": 5, "numQueuedRequests": 3, @@ -93,6 +133,491 @@ def _get_counter_value(collector, metric_name: str): } +class TestEPHealthStats: + """Test committed EPGroupHealth membership is exposed as passive gauges.""" + + def test_no_snapshot_exports_no_ep_health_families(self, collector): + registry = CollectorRegistry() + collector.register_ep_health_metrics(registry) + + assert _get_ep_metric_samples(collector) == [] + assert "trtllm_ep_" not in generate_latest(registry).decode() + + # ``None``/unsupported health is represented by marking a collector + # unavailable before any producer has published a valid snapshot. + collector.log_ep_health_unavailable() + assert _get_ep_metric_samples(collector) == [] + assert "trtllm_ep_" not in generate_latest(registry).decode() + + def test_per_rank_and_group_gauges(self, collector): + assert ( + collector.log_ep_health_stats( + { + "sourceEpoch": EP_HEALTH_SOURCE_A, + "worldSize": 4, + "activeCount": 3, + "failedRanks": [2], + "generation": 1, + } + ) + is True + ) + + assert [_get_ep_rank_gauge_value(collector, rank) for rank in range(4)] == [ + 1, + 1, + 0, + 1, + ] + assert _get_gauge_value(collector, "ep_active_ranks") == 3 + assert _get_gauge_value(collector, "ep_failed_ranks") == 1 + assert _get_gauge_value(collector, "ep_health_generation") == 1 + assert _get_gauge_value(collector, "ep_health_available") == 1 + + def test_reactivation_updates_existing_rank_series(self, collector): + collector.log_ep_health_stats( + { + "sourceEpoch": EP_HEALTH_SOURCE_A, + "worldSize": 2, + "activeCount": 1, + "failedRanks": [1], + "generation": 1, + } + ) + collector.log_ep_health_stats( + { + "sourceEpoch": EP_HEALTH_SOURCE_A, + "worldSize": 2, + "activeCount": 2, + "failedRanks": [], + "generation": 2, + } + ) + + assert _get_ep_rank_gauge_value(collector, 1) == 1 + assert _get_gauge_value(collector, "ep_active_ranks") == 2 + assert _get_gauge_value(collector, "ep_failed_ranks") == 0 + assert _get_gauge_value(collector, "ep_health_generation") == 2 + + def test_conflicting_same_generation_warns_and_is_ignored(self, collector, caplog): + collector.log_ep_health_stats( + { + "sourceEpoch": EP_HEALTH_SOURCE_A, + "worldSize": 2, + "activeCount": 1, + "failedRanks": [1], + "generation": 1, + } + ) + assert ( + collector.log_ep_health_stats( + { + "sourceEpoch": EP_HEALTH_SOURCE_A, + "worldSize": 2, + "activeCount": 1, + "failedRanks": [0], + "generation": 1, + } + ) + is False + ) + + assert "Conflicting epHealthStats payloads" in caplog.text + assert _get_ep_rank_gauge_value(collector, 0) == 1 + assert _get_ep_rank_gauge_value(collector, 1) == 0 + assert _get_gauge_value(collector, "ep_health_available") == 0 + + def test_stale_generation_warns_and_is_ignored(self, collector, caplog): + collector.log_ep_health_stats( + { + "sourceEpoch": EP_HEALTH_SOURCE_A, + "worldSize": 2, + "activeCount": 1, + "failedRanks": [1], + "generation": 2, + } + ) + assert ( + collector.log_ep_health_stats( + { + "sourceEpoch": EP_HEALTH_SOURCE_A, + "worldSize": 2, + "activeCount": 1, + "failedRanks": [0], + "generation": 1, + } + ) + is False + ) + + assert "Ignoring stale epHealthStats generation" in caplog.text + assert _get_ep_rank_gauge_value(collector, 0) == 1 + assert _get_ep_rank_gauge_value(collector, 1) == 0 + assert _get_gauge_value(collector, "ep_health_available") == 0 + + def test_new_producer_epoch_can_restart_generation(self, collector): + assert collector.log_ep_health_stats( + { + "sourceEpoch": "producer-a", + "worldSize": 2, + "activeCount": 1, + "failedRanks": [1], + "generation": 5, + } + ) + collector.log_ep_health_unavailable() + + assert collector.log_ep_health_stats( + { + "sourceEpoch": "producer-b", + "worldSize": 2, + "activeCount": 2, + "failedRanks": [], + "generation": 0, + } + ) + assert [_get_ep_rank_gauge_value(collector, rank) for rank in range(2)] == [1, 1] + assert _get_gauge_value(collector, "ep_active_ranks") == 2 + assert _get_gauge_value(collector, "ep_health_generation") == 0 + assert _get_gauge_value(collector, "ep_health_available") == 1 + + def test_retired_producer_epoch_cannot_become_current_again(self, collector, caplog): + assert collector.log_ep_health_stats( + { + "sourceEpoch": EP_HEALTH_SOURCE_A, + "worldSize": 2, + "activeCount": 1, + "failedRanks": [1], + "generation": 9, + } + ) + assert collector.log_ep_health_stats( + { + "sourceEpoch": EP_HEALTH_SOURCE_B, + "worldSize": 2, + "activeCount": 2, + "failedRanks": [], + "generation": 0, + } + ) + + assert ( + collector.log_ep_health_stats( + { + "sourceEpoch": EP_HEALTH_SOURCE_A, + "worldSize": 2, + "activeCount": 1, + "failedRanks": [1], + "generation": 9, + } + ) + is False + ) + assert "retired source epoch" in caplog.text + assert [_get_ep_rank_gauge_value(collector, rank) for rank in range(2)] == [1, 1] + assert _get_gauge_value(collector, "ep_health_generation") == 0 + assert _get_gauge_value(collector, "ep_health_available") == 0 + + def test_source_epoch_is_required_after_epoch_aware_payload(self, collector): + assert collector.log_ep_health_stats( + { + "sourceEpoch": EP_HEALTH_SOURCE_A, + "worldSize": 1, + "activeCount": 1, + "failedRanks": [], + "generation": 0, + } + ) + + assert ( + collector.log_ep_health_stats( + { + "worldSize": 1, + "activeCount": 1, + "failedRanks": [], + "generation": 1, + } + ) + is False + ) + assert collector._last_ep_health_state[0] == EP_HEALTH_SOURCE_A + assert _get_gauge_value(collector, "ep_health_available") == 0 + + def test_source_epoch_churn_keeps_recent_replay_window_bounded(self, collector): + transitions = _MAX_RETIRED_EP_HEALTH_SOURCE_EPOCHS + 10 + for epoch in range(transitions): + assert collector.log_ep_health_stats( + { + "sourceEpoch": f"source-{epoch}", + "worldSize": 1, + "activeCount": 1, + "failedRanks": [], + "generation": 0, + } + ) + + # Legitimate producer replacements continue after the bounded history + # fills, while recently retired producers are still rejected. + assert ( + collector.log_ep_health_stats( + { + "sourceEpoch": f"source-{transitions - 2}", + "worldSize": 1, + "activeCount": 1, + "failedRanks": [], + "generation": 0, + } + ) + is False + ) + assert ( + len(collector._retired_ep_health_source_epochs) == _MAX_RETIRED_EP_HEALTH_SOURCE_EPOCHS + ) + assert ( + len(collector._retired_ep_health_source_epoch_order) + == _MAX_RETIRED_EP_HEALTH_SOURCE_EPOCHS + ) + assert _get_gauge_value(collector, "ep_health_available") == 0 + + assert collector.log_ep_health_stats( + { + "sourceEpoch": "source-0", + "worldSize": 1, + "activeCount": 1, + "failedRanks": [], + "generation": 0, + } + ) + assert _get_gauge_value(collector, "ep_health_available") == 1 + + def test_oversized_world_is_rejected_before_rank_series_creation(self, collector): + assert _get_ep_rank_samples(collector) == [] + + assert ( + collector.log_ep_health_stats( + { + "sourceEpoch": EP_HEALTH_SOURCE_A, + "worldSize": 129, + "activeCount": 129, + "failedRanks": [], + "generation": 0, + } + ) + is False + ) + assert _get_ep_metric_samples(collector) == [] + + def test_unrepresentable_generation_is_rejected_before_metric_writes(self, collector): + assert _get_ep_rank_samples(collector) == [] + + assert ( + collector.log_ep_health_stats( + { + "sourceEpoch": EP_HEALTH_SOURCE_A, + "worldSize": 1, + "activeCount": 1, + "failedRanks": [], + "generation": 10**1000, + } + ) + is False + ) + assert collector._last_ep_health_state is None + assert _get_ep_metric_samples(collector) == [] + + def test_world_size_change_warns_and_is_ignored(self, collector, caplog): + collector.log_ep_health_stats( + { + "sourceEpoch": EP_HEALTH_SOURCE_A, + "worldSize": 4, + "activeCount": 3, + "failedRanks": [3], + "generation": 1, + } + ) + assert ( + collector.log_ep_health_stats( + { + "sourceEpoch": EP_HEALTH_SOURCE_A, + "worldSize": 2, + "activeCount": 2, + "failedRanks": [], + "generation": 2, + } + ) + is False + ) + + assert "Ignoring epHealthStats world-size change" in caplog.text + assert _get_ep_rank_gauge_value(collector, 3) == 0 + assert _get_gauge_value(collector, "ep_active_ranks") == 3 + assert _get_gauge_value(collector, "ep_failed_ranks") == 1 + assert _get_gauge_value(collector, "ep_health_generation") == 1 + assert _get_gauge_value(collector, "ep_health_available") == 0 + + def test_invalid_payload_is_ignored(self, collector, caplog): + invalid = { + "sourceEpoch": EP_HEALTH_SOURCE_A, + "worldSize": 4, + "activeCount": 4, + "failedRanks": [4], + "generation": 1, + } + assert collector.log_ep_health_stats(invalid) is False + assert collector.log_ep_health_stats(invalid) is False + + assert "Ignoring invalid epHealthStats payload" in caplog.text + assert caplog.text.count("Ignoring invalid epHealthStats payload") == 1 + assert collector._last_ep_health_state is None + assert _get_ep_metric_samples(collector) == [] + + assert collector.log_ep_health_stats( + { + "sourceEpoch": EP_HEALTH_SOURCE_A, + "worldSize": 4, + "activeCount": 4, + "failedRanks": [], + "generation": 0, + } + ) + assert collector.log_ep_health_stats(invalid) is False + assert caplog.text.count("Ignoring invalid epHealthStats payload") == 2 + + def test_rpc_outage_marks_snapshot_unavailable_and_duplicate_recovers(self, collector): + snapshot = { + "sourceEpoch": EP_HEALTH_SOURCE_A, + "worldSize": 2, + "activeCount": 2, + "failedRanks": [], + "generation": 0, + } + assert collector.log_ep_health_stats(snapshot) is True + collector.log_ep_health_unavailable() + + assert _get_ep_rank_gauge_value(collector, 0) == 1 + assert _get_gauge_value(collector, "ep_active_ranks") == 2 + assert _get_gauge_value(collector, "ep_failed_ranks") == 0 + assert _get_gauge_value(collector, "ep_health_generation") == 0 + assert _get_gauge_value(collector, "ep_health_available") == 0 + + assert collector.log_ep_health_stats(snapshot) is True + assert _get_gauge_value(collector, "ep_health_available") == 1 + + def test_health_metrics_register_on_scrape_registry(self, collector): + collector.log_ep_health_stats( + { + "sourceEpoch": EP_HEALTH_SOURCE_A, + "worldSize": 2, + "activeCount": 1, + "failedRanks": [1], + "generation": 3, + } + ) + registry = CollectorRegistry() + collector.register_ep_health_metrics(registry) + output = generate_latest(registry).decode() + + assert 'trtllm_ep_rank_active{ep_rank="0",model_name="test_model"} 1.0' in output + assert 'trtllm_ep_rank_active{ep_rank="1",model_name="test_model"} 0.0' in output + assert 'trtllm_ep_active_ranks{model_name="test_model"} 1.0' in output + assert 'trtllm_ep_failed_ranks{model_name="test_model"} 1.0' in output + assert 'trtllm_ep_health_generation{model_name="test_model"} 3.0' in output + assert 'trtllm_ep_health_available{model_name="test_model"} 1.0' in output + + def test_multiprocess_registry_and_local_health_collector_do_not_duplicate( + self, collector, tmp_path, monkeypatch + ): + monkeypatch.setenv("PROMETHEUS_MULTIPROC_DIR", str(tmp_path)) + collector.log_ep_health_stats( + { + "sourceEpoch": EP_HEALTH_SOURCE_A, + "worldSize": 2, + "activeCount": 1, + "failedRanks": [1], + "generation": 4, + } + ) + + registry = CollectorRegistry() + multiprocess.MultiProcessCollector(registry) + collector.register_ep_health_metrics(registry) + output = generate_latest(registry).decode() + + assert output.count("# HELP trtllm_ep_rank_active") == 1 + assert output.count("trtllm_ep_rank_active{") == 2 + assert output.count("trtllm_ep_health_available{") == 1 + assert "pid=" not in "\n".join( + line for line in output.splitlines() if line.startswith("trtllm_ep_") + ) + + def test_concurrent_scrapes_observe_only_complete_health_snapshots(self, collector): + registry = CollectorRegistry() + collector.register_ep_health_metrics(registry) + assert collector.log_ep_health_stats( + { + "sourceEpoch": EP_HEALTH_SOURCE_A, + "worldSize": 4, + "activeCount": 3, + "failedRanks": [0], + "generation": 0, + } + ) + writer_done = threading.Event() + start = threading.Barrier(2) + writer_errors = [] + + def writer(): + try: + start.wait(timeout=10.0) + for generation in range(1, 501): + failed_ranks = [generation % 4] + assert collector.log_ep_health_stats( + { + "sourceEpoch": EP_HEALTH_SOURCE_A, + "worldSize": 4, + "activeCount": 3, + "failedRanks": failed_ranks, + "generation": generation, + } + ) + time.sleep(0) + except BaseException as error: + writer_errors.append(error) + finally: + writer_done.set() + + thread = threading.Thread(target=writer) + thread.start() + start.wait(timeout=10.0) + scrapes = 0 + while not writer_done.is_set() or scrapes == 0: + families = { + metric.name: metric + for metric in collector._ep_health_prometheus_collector.collect() + } + availability = families["trtllm_ep_health_available"].samples[0].value + if availability: + ranks = families["trtllm_ep_rank_active"].samples + active_count = families["trtllm_ep_active_ranks"].samples[0].value + failed_count = families["trtllm_ep_failed_ranks"].samples[0].value + generation = families["trtllm_ep_health_generation"].samples[0].value + failed_ranks = [ + int(sample.labels[collector.labelname_ep_rank]) + for sample in ranks + if not sample.value + ] + assert len(ranks) == 4 + assert sum(sample.value for sample in ranks) == active_count + assert len(ranks) - active_count == failed_count + assert failed_ranks == [int(generation) % 4] + scrapes += 1 + thread.join(timeout=10.0) + + assert not thread.is_alive() + assert writer_errors == [] + assert collector._last_ep_health_state[2] == 500 + + class TestIterationStatsTopLevel: """Test top-level iteration stats are correctly exposed as Prometheus metrics."""