diff --git a/src/plugin/net.cc b/src/plugin/net.cc index 717f1a92c1..c46c69c995 100644 --- a/src/plugin/net.cc +++ b/src/plugin/net.cc @@ -9,6 +9,7 @@ #include "bootstrap.h" #include "checks.h" #include "plugin.h" +#include "plugin_cleanup.h" #include "nccl_net.h" #include @@ -165,6 +166,7 @@ ncclResult_t ncclNetCheckDeviceVersion(struct ncclComm* comm, ncclNet_t* net, in static ncclResult_t ncclNetPluginInit(struct ncclComm* comm, netPluginLib_t* pluginLib) { int ndev; bool initCompleted = false; + bool collNetInitCompleted = false; // Init must be called for each new comm to set the right context if (pluginLib->ncclNetPluginState >= ncclNetPluginStateInitReady && pluginLib->ncclNet) { ncclNetCommConfig_t commConfig = {}; @@ -189,11 +191,19 @@ static ncclResult_t ncclNetPluginInit(struct ncclComm* comm, netPluginLib_t* plu if (pluginLib->ncclCollNetPluginState >= ncclNetPluginStateInitReady && pluginLib->ncclCollNet) { if (pluginLib->ncclCollNet->init(&comm->collNetContext, comm->commHash, ncclDebugLog) != ncclSuccess) { pluginLib->ncclCollNetPluginState = ncclNetPluginStateDisabled; + } else { + collNetInitCompleted = true; } } // Detection of the devices is only done when the plugin is being initialized the first time if (pluginLib->ncclCollNetPluginState == ncclNetPluginStateInitReady && pluginLib->ncclCollNet) { if (pluginLib->ncclCollNet->devices(&ndev) != ncclSuccess || ndev <= 0) { + if (collNetInitCompleted) { + ncclResult_t finalizeRet = ncclPluginFinalizeContext(pluginLib->ncclCollNet->finalize, &comm->collNetContext); + if (finalizeRet != ncclSuccess) { + WARN("NET/Plugin: CollNet finalize failed during init fallback cleanup: %d", finalizeRet); + } + } pluginLib->ncclCollNetPluginState = ncclNetPluginStateDisabled; } else { pluginLib->collNetPhysDevs = ndev; diff --git a/src/plugin/plugin_cleanup.h b/src/plugin/plugin_cleanup.h new file mode 100644 index 0000000000..42e1a77f02 --- /dev/null +++ b/src/plugin/plugin_cleanup.h @@ -0,0 +1,20 @@ +/************************************************************************* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * See LICENSE.txt for more license information + *************************************************************************/ + +#ifndef NCCL_PLUGIN_CLEANUP_H_ +#define NCCL_PLUGIN_CLEANUP_H_ + +typedef ncclResult_t (*ncclPluginFinalizeFn_t)(void* ctx); + +static inline ncclResult_t ncclPluginFinalizeContext(ncclPluginFinalizeFn_t finalize, void** ctx) { + if (*ctx == nullptr) return ncclSuccess; + ncclResult_t ret = finalize(*ctx); + *ctx = nullptr; + return ret; +} + +#endif diff --git a/src/plugin/rma.cc b/src/plugin/rma.cc index 3ffa22f516..4f1ba7bbe7 100644 --- a/src/plugin/rma.cc +++ b/src/plugin/rma.cc @@ -9,6 +9,7 @@ #include "bootstrap.h" #include "checks.h" #include "plugin.h" +#include "plugin_cleanup.h" #include "nccl_rma.h" #include "gin/gin_host_proxy.h" @@ -105,14 +106,23 @@ static ncclResult_t ncclRmaPluginLoad(rmaPluginLib_t* pluginLib) { static ncclResult_t ncclRmaPluginInit(struct ncclComm* comm, rmaPluginLib_t* pluginLib) { int ndev; + bool rmaInitCompleted = false; // Init must be called for each new comm to set the right context if (pluginLib->state >= ncclRmaPluginStateInitReady && pluginLib->ncclRma) { if (pluginLib->ncclRma->init(&comm->rmaContext, comm->commHash, ncclDebugLog) != ncclSuccess) { pluginLib->state = ncclRmaPluginStateDisabled; + } else { + rmaInitCompleted = true; } } if (pluginLib->state == ncclRmaPluginStateInitReady && pluginLib->ncclRma) { if (pluginLib->ncclRma->devices(&ndev) != ncclSuccess || ndev <= 0) { + if (rmaInitCompleted) { + ncclResult_t finalizeRet = ncclPluginFinalizeContext(pluginLib->ncclRma->finalize, &comm->rmaContext); + if (finalizeRet != ncclSuccess) { + WARN("RMA/Plugin: finalize failed during init fallback cleanup: %d", finalizeRet); + } + } pluginLib->state = ncclRmaPluginStateDisabled; } else { pluginLib->physDevs = ndev; diff --git a/src/plugin/rma/rma_v13.cc b/src/plugin/rma/rma_v13.cc index ec711f8c59..c1fa9c3ea1 100644 --- a/src/plugin/rma/rma_v13.cc +++ b/src/plugin/rma/rma_v13.cc @@ -9,22 +9,37 @@ #include "nccl_rma.h" #include "checks.h" #include "os.h" +#include "plugin_cleanup.h" #include static ncclRma_v13_t* ncclRma_v13; static ncclRma_t ncclRma; static ncclResult_t ncclRma_init(void** ctx, uint64_t commId, ncclDebugLogger_t logFunction) { - NCCLCHECK(ncclRma_v13->init(ctx, commId, logFunction)); + ncclResult_t ret = ncclSuccess; + void* tmpCtx = nullptr; + *ctx = nullptr; + NCCLCHECK(ncclRma_v13->init(&tmpCtx, commId, logFunction)); // RMA plugin must report GIN proxy type ncclNetProperties_t props; - NCCLCHECK(ncclRma_v13->getProperties(0, &props)); + NCCLCHECKGOTO(ncclRma_v13->getProperties(0, &props), ret, fail); if (props.netDeviceType != NCCL_NET_DEVICE_GIN_PROXY) { WARN("RMA v13 (%s) requires GIN PROXY type, got netDeviceType %d", ncclRma_v13->name, props.netDeviceType); - return ncclInternalError; + ret = ncclInternalError; + goto fail; } + *ctx = tmpCtx; return ncclSuccess; +fail: + if (tmpCtx) { + ncclResult_t finalizeRet = ncclPluginFinalizeContext(ncclRma_v13->finalize, &tmpCtx); + if (finalizeRet != ncclSuccess) { + WARN("RMA v13 (%s) finalize failed during init fallback cleanup: %d", ncclRma_v13->name, finalizeRet); + } + if (ret == ncclSuccess) ret = finalizeRet; + } + return ret; } static ncclResult_t ncclRma_createContext(void* collComm, ncclRmaConfig_v14_t* config, void** rmaCtx) { diff --git a/tests/plugin_lifecycle/Makefile b/tests/plugin_lifecycle/Makefile new file mode 100644 index 0000000000..67e84cf09c --- /dev/null +++ b/tests/plugin_lifecycle/Makefile @@ -0,0 +1,22 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +CXX ?= g++ +CXXFLAGS ?= -std=c++17 -O2 -Wall -Wextra -Werror + +BUILD_DIR ?= ../../build/test/plugin_lifecycle +TEST_BIN := $(BUILD_DIR)/plugin_lifecycle_test + +.PHONY: all run clean + +all: $(TEST_BIN) + +$(TEST_BIN): plugin_lifecycle_test.cc + @mkdir -p $(BUILD_DIR) + $(CXX) $(CXXFLAGS) $< -o $@ + +run: $(TEST_BIN) + $(TEST_BIN) + +clean: + rm -rf $(BUILD_DIR) diff --git a/tests/plugin_lifecycle/plugin_lifecycle_test.cc b/tests/plugin_lifecycle/plugin_lifecycle_test.cc new file mode 100644 index 0000000000..2fda751764 --- /dev/null +++ b/tests/plugin_lifecycle/plugin_lifecycle_test.cc @@ -0,0 +1,281 @@ +/************************************************************************* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * See LICENSE.txt for more license information + *************************************************************************/ + +#include +#include +#include + +enum ncclResult_t { + ncclSuccess = 0, + ncclInternalError = 3, +}; + +#include "../../src/plugin/plugin_cleanup.h" + +enum PluginState { + PluginStateDisabled = -2, + PluginStateInitReady = 1, + PluginStateEnabled = 2, +}; + +enum NetDeviceType { + NetDeviceHost = 0, + NetDeviceGinProxy = 2, +}; + +struct Comm { + void* rmaContext = nullptr; + void* collNetContext = nullptr; +}; + +struct MockPlugin { + int initCalls = 0; + int devicesCalls = 0; + int getPropertiesCalls = 0; + int finalizeCalls = 0; + int devicesResult = 0; + ncclResult_t getPropertiesResult = ncclSuccess; + NetDeviceType deviceType = NetDeviceGinProxy; + + ncclResult_t init(void** ctx) { + initCalls++; + *ctx = std::malloc(1); + return *ctx ? ncclSuccess : ncclInternalError; + } + + ncclResult_t devices(int* ndev) { + devicesCalls++; + *ndev = devicesResult; + return ncclSuccess; + } + + ncclResult_t getProperties(NetDeviceType* type) { + getPropertiesCalls++; + if (getPropertiesResult != ncclSuccess) return getPropertiesResult; + *type = deviceType; + return ncclSuccess; + } + + ncclResult_t finalize(void* ctx) { + finalizeCalls++; + std::free(ctx); + return ncclSuccess; + } +}; + +static MockPlugin* activeFinalizePlugin = nullptr; + +static ncclResult_t mockFinalize(void* ctx) { + assert(activeFinalizePlugin != nullptr); + return activeFinalizePlugin->finalize(ctx); +} + +static void rmaPluginInitBeforeFix(Comm* comm, MockPlugin* plugin, PluginState* state) { + int ndev; + if (*state >= PluginStateInitReady) { + if (plugin->init(&comm->rmaContext) != ncclSuccess) *state = PluginStateDisabled; + } + if (*state == PluginStateInitReady) { + if (plugin->devices(&ndev) != ncclSuccess || ndev <= 0) *state = PluginStateDisabled; + else *state = PluginStateEnabled; + } +} + +static void rmaPluginInitAfterFix(Comm* comm, MockPlugin* plugin, PluginState* state) { + int ndev; + bool rmaInitCompleted = false; + if (*state >= PluginStateInitReady) { + if (plugin->init(&comm->rmaContext) != ncclSuccess) { + *state = PluginStateDisabled; + } else { + rmaInitCompleted = true; + } + } + if (*state == PluginStateInitReady) { + if (plugin->devices(&ndev) != ncclSuccess || ndev <= 0) { + if (rmaInitCompleted) { + activeFinalizePlugin = plugin; + ncclPluginFinalizeContext(mockFinalize, &comm->rmaContext); + activeFinalizePlugin = nullptr; + } + *state = PluginStateDisabled; + } else { + *state = PluginStateEnabled; + } + } +} + +static void collNetPluginInitBeforeFix(Comm* comm, MockPlugin* plugin, PluginState* state) { + int ndev; + if (*state >= PluginStateInitReady) { + if (plugin->init(&comm->collNetContext) != ncclSuccess) *state = PluginStateDisabled; + } + if (*state == PluginStateInitReady) { + if (plugin->devices(&ndev) != ncclSuccess || ndev <= 0) *state = PluginStateDisabled; + else *state = PluginStateEnabled; + } +} + +static void collNetPluginInitAfterFix(Comm* comm, MockPlugin* plugin, PluginState* state) { + int ndev; + bool collNetInitCompleted = false; + if (*state >= PluginStateInitReady) { + if (plugin->init(&comm->collNetContext) != ncclSuccess) { + *state = PluginStateDisabled; + } else { + collNetInitCompleted = true; + } + } + if (*state == PluginStateInitReady) { + if (plugin->devices(&ndev) != ncclSuccess || ndev <= 0) { + if (collNetInitCompleted) { + activeFinalizePlugin = plugin; + ncclPluginFinalizeContext(mockFinalize, &comm->collNetContext); + activeFinalizePlugin = nullptr; + } + *state = PluginStateDisabled; + } else { + *state = PluginStateEnabled; + } + } +} + +static ncclResult_t rmaV13InitBeforeFix(void** ctx, MockPlugin* plugin) { + ncclResult_t ret = plugin->init(ctx); + if (ret != ncclSuccess) return ret; + + NetDeviceType type; + ret = plugin->getProperties(&type); + if (ret != ncclSuccess) return ret; + if (type != NetDeviceGinProxy) return ncclInternalError; + return ncclSuccess; +} + +static ncclResult_t rmaV13InitAfterFix(void** ctx, MockPlugin* plugin) { + ncclResult_t ret = ncclSuccess; + void* tmpCtx = nullptr; + *ctx = nullptr; + + ret = plugin->init(&tmpCtx); + if (ret != ncclSuccess) return ret; + + NetDeviceType type; + ret = plugin->getProperties(&type); + if (ret != ncclSuccess) goto fail; + if (type != NetDeviceGinProxy) { + ret = ncclInternalError; + goto fail; + } + *ctx = tmpCtx; + return ncclSuccess; + +fail: + if (tmpCtx) { + activeFinalizePlugin = plugin; + ncclResult_t finalizeRet = ncclPluginFinalizeContext(mockFinalize, &tmpCtx); + activeFinalizePlugin = nullptr; + if (ret == ncclSuccess) ret = finalizeRet; + } + return ret; +} + +static void testRmaDevicesFailureCleanup() { + { + Comm comm; + MockPlugin plugin; + PluginState state = PluginStateInitReady; + rmaPluginInitBeforeFix(&comm, &plugin, &state); + assert(state == PluginStateDisabled); + assert(plugin.finalizeCalls == 0); + assert(comm.rmaContext != nullptr); + std::free(comm.rmaContext); + } + { + Comm comm; + MockPlugin plugin; + PluginState state = PluginStateInitReady; + rmaPluginInitAfterFix(&comm, &plugin, &state); + assert(state == PluginStateDisabled); + assert(plugin.finalizeCalls == 1); + assert(comm.rmaContext == nullptr); + } +} + +static void testCollNetDevicesFailureCleanup() { + { + Comm comm; + MockPlugin plugin; + PluginState state = PluginStateInitReady; + collNetPluginInitBeforeFix(&comm, &plugin, &state); + assert(state == PluginStateDisabled); + assert(plugin.finalizeCalls == 0); + assert(comm.collNetContext != nullptr); + std::free(comm.collNetContext); + } + { + Comm comm; + MockPlugin plugin; + PluginState state = PluginStateInitReady; + collNetPluginInitAfterFix(&comm, &plugin, &state); + assert(state == PluginStateDisabled); + assert(plugin.finalizeCalls == 1); + assert(comm.collNetContext == nullptr); + } +} + +static void testRmaV13ValidationFailureCleanup() { + { + void* ctx = nullptr; + MockPlugin plugin; + plugin.deviceType = NetDeviceHost; + ncclResult_t ret = rmaV13InitBeforeFix(&ctx, &plugin); + assert(ret == ncclInternalError); + assert(plugin.finalizeCalls == 0); + assert(ctx != nullptr); + std::free(ctx); + } + { + void* ctx = nullptr; + MockPlugin plugin; + plugin.deviceType = NetDeviceHost; + ncclResult_t ret = rmaV13InitAfterFix(&ctx, &plugin); + assert(ret == ncclInternalError); + assert(plugin.finalizeCalls == 1); + assert(ctx == nullptr); + } +} + +static void testRmaV13GetPropertiesFailureCleanup() { + { + void* ctx = nullptr; + MockPlugin plugin; + plugin.getPropertiesResult = ncclInternalError; + ncclResult_t ret = rmaV13InitBeforeFix(&ctx, &plugin); + assert(ret == ncclInternalError); + assert(plugin.finalizeCalls == 0); + assert(ctx != nullptr); + std::free(ctx); + } + { + void* ctx = reinterpret_cast(0x1); + MockPlugin plugin; + plugin.getPropertiesResult = ncclInternalError; + ncclResult_t ret = rmaV13InitAfterFix(&ctx, &plugin); + assert(ret == ncclInternalError); + assert(plugin.finalizeCalls == 1); + assert(ctx == nullptr); + } +} + +int main() { + testRmaDevicesFailureCleanup(); + testCollNetDevicesFailureCleanup(); + testRmaV13ValidationFailureCleanup(); + testRmaV13GetPropertiesFailureCleanup(); + std::cout << "plugin lifecycle cleanup tests passed\n"; + return 0; +}