From 6f52299852769555337872442dc63e77b0dce28d Mon Sep 17 00:00:00 2001 From: Darshan Sanghani Date: Wed, 12 Aug 2026 14:51:32 -0700 Subject: [PATCH] Remove eager CUPTI resource callback subscription (#1250) Summary: **Problem** Kineto reserves CUPTI's one callback subscriber at process startup, blocking tools such as wprof even when Kineto is idle. **Why** The `RESOURCE_CONTEXT_CREATED` callback only delays Kineto initialization. Daemon selection was also coupled to `KINETO_USE_DAEMON`, so the OSS setting could replace or reroute the internal Dyno loader. **Fix** - **Internal initialization:** Register `DynoConfigLoader`, then initialize Kineto and start config polling from a late `RegisterInitFunc`. - **OSS daemon:** Register and start the IPC-only `DaemonConfigLoader` only when `KINETO_USE_DAEMON` is set and no platform loader exists. - **Loader transport:** Keep the internal Dyno loader on Thrift when IPC is unavailable, regardless of the OSS daemon environment variable or config. - **Forked workers:** When launch initialization is disabled, repair inherited client state and initialize Kineto synchronously in the child. Defer the CUDA and CUPTI availability probe until the first trace so child initialization does not initialize CUDA. If the probe fails, continue with CPU-only profiling. - **CUPTI ownership:** Remove the startup callback subscription while preserving trace-time CUPTI activity setup and explicit CUDA-injection initialization. Reviewed By: scotts Differential Revision: D92570679 --- libkineto/src/ConfigLoader.cpp | 4 + libkineto/src/ConfigLoader.h | 1 + libkineto/src/CuptiActivityApi.cpp | 5 + libkineto/src/CuptiActivityApi.h | 1 + libkineto/src/CuptiActivityProfiler.cpp | 35 ++++-- libkineto/src/CuptiActivityProfiler.h | 8 +- libkineto/src/init.cpp | 110 ++++--------------- libkineto/test/CuptiActivityProfilerTest.cpp | 29 +++++ 8 files changed, 95 insertions(+), 98 deletions(-) diff --git a/libkineto/src/ConfigLoader.cpp b/libkineto/src/ConfigLoader.cpp index 3ec813560..e117f1a8e 100644 --- a/libkineto/src/ConfigLoader.cpp +++ b/libkineto/src/ConfigLoader.cpp @@ -64,6 +64,10 @@ void ConfigLoader::setDaemonConfigLoaderFactory( daemonConfigLoaderFactory() = std::move(factory); } +bool ConfigLoader::hasDaemonConfigLoaderFactory() { + return daemonConfigLoaderFactory() != nullptr; +} + ConfigLoader& ConfigLoader::instance() { static ConfigLoader config_loader; return config_loader; diff --git a/libkineto/src/ConfigLoader.h b/libkineto/src/ConfigLoader.h index 0d9249cec..85c99196f 100644 --- a/libkineto/src/ConfigLoader.h +++ b/libkineto/src/ConfigLoader.h @@ -88,6 +88,7 @@ class ConfigLoader { static void setDaemonConfigLoaderFactory( std::function()> factory); + static bool hasDaemonConfigLoaderFactory(); std::string getConfString(); diff --git a/libkineto/src/CuptiActivityApi.cpp b/libkineto/src/CuptiActivityApi.cpp index f66a0da74..fafe85253 100644 --- a/libkineto/src/CuptiActivityApi.cpp +++ b/libkineto/src/CuptiActivityApi.cpp @@ -60,6 +60,11 @@ CuptiActivityApi& CuptiActivityApi::singleton() { return *instance; } +bool CuptiActivityApi::isAvailable(uint32_t& version) const { + return isGpuAvailable() && + CUPTI_CALL_NOWARN(cuptiGetVersion(&version)) == CUPTI_SUCCESS; +} + void CuptiActivityApi::pushCorrelationID(int id, CorrelationFlowType type) { if (!singleton().externalCorrelationEnabled_.load( std::memory_order_relaxed)) { diff --git a/libkineto/src/CuptiActivityApi.h b/libkineto/src/CuptiActivityApi.h index b54fd43c3..9f18772a1 100644 --- a/libkineto/src/CuptiActivityApi.h +++ b/libkineto/src/CuptiActivityApi.h @@ -46,6 +46,7 @@ class CuptiActivityApi { static void pushCorrelationID(int id, CorrelationFlowType type); static void popCorrelationID(CorrelationFlowType type); + virtual bool isAvailable(uint32_t& version) const; void enableCuptiActivities( const std::set& selected_activities, bool enablePerThreadBuffers = false); diff --git a/libkineto/src/CuptiActivityProfiler.cpp b/libkineto/src/CuptiActivityProfiler.cpp index d608a16de..be8cbc8a6 100644 --- a/libkineto/src/CuptiActivityProfiler.cpp +++ b/libkineto/src/CuptiActivityProfiler.cpp @@ -75,17 +75,11 @@ namespace KINETO_NAMESPACE { CuptiActivityProfiler::CuptiActivityProfiler( CuptiActivityApi& cupti, bool cpuOnly) - : GenericActivityProfiler(cpuOnly), cupti_(cupti) { - if (isGpuAvailable()) { - logGpuVersions(); - } -} + : GenericActivityProfiler(cpuOnly), cupti_(cupti) {} -void CuptiActivityProfiler::logGpuVersions() { - uint32_t cuptiVersion = 0; +void CuptiActivityProfiler::recordGpuVersions(uint32_t cuptiVersion) { int cudaRuntimeVersion = 0; int cudaDriverVersion = 0; - CUPTI_CALL(cuptiGetVersion(&cuptiVersion)); CUDA_CALL(cudaRuntimeGetVersion(&cudaRuntimeVersion)); CUDA_CALL(cudaDriverGetVersion(&cudaDriverVersion)); LOG(INFO) << "CUDA versions. CUPTI: " << cuptiVersion @@ -109,6 +103,22 @@ void CuptiActivityProfiler::setMaxGpuBufferSize(int64_t size) { } void CuptiActivityProfiler::enableGpuTracing() { + // @lint-ignore CLANGTIDY facebook-hte-std::call_once + std::call_once(cuptiInitializationOnce_, [this] { + uint32_t cuptiVersion = 0; + cuptiAvailable_ = cupti_.isAvailable(cuptiVersion); + if (cuptiAvailable_) { + recordGpuVersions(cuptiVersion); + } else { + cpuOnly_ = true; + VLOG(0) << "CUPTI unavailable; continuing with CPU-only profiling"; + } + }); + if (!cuptiAvailable_) { + toggleState_.store(false); + return; + } + configureCuptiTimestampSource(config().getTSCTimestampFlag()); cupti_.enableCuptiActivities( derivedConfig_->profileActivityTypes(), @@ -116,10 +126,16 @@ void CuptiActivityProfiler::enableGpuTracing() { } void CuptiActivityProfiler::disableGpuTracing() { + if (!cuptiAvailable_) { + return; + } cupti_.disableCuptiActivities(derivedConfig_->profileActivityTypes()); } void CuptiActivityProfiler::clearGpuActivities() { + if (!cuptiAvailable_) { + return; + } cupti_.clearActivities(); } @@ -128,6 +144,9 @@ bool CuptiActivityProfiler::isGpuCollectionStopped() const { } void CuptiActivityProfiler::synchronizeGpuDevice() { + if (!cuptiAvailable_) { + return; + } CUDA_CALL(cudaDeviceSynchronize()); cupti_.flushActivities(); } diff --git a/libkineto/src/CuptiActivityProfiler.h b/libkineto/src/CuptiActivityProfiler.h index 927eb9ba9..9a79c9402 100644 --- a/libkineto/src/CuptiActivityProfiler.h +++ b/libkineto/src/CuptiActivityProfiler.h @@ -9,6 +9,7 @@ #pragma once #include +#include #include "CuptiActivity.h" #include "CuptiActivityApi.h" #include "GenericActivityProfiler.h" @@ -23,7 +24,6 @@ class CuptiActivityProfiler : public GenericActivityProfiler { ~CuptiActivityProfiler() override = default; protected: - void logGpuVersions() override; void setMaxGpuBufferSize(int64_t size) override; void enableGpuTracing() override; void disableGpuTracing() override; @@ -63,6 +63,12 @@ class CuptiActivityProfiler : public GenericActivityProfiler { template void handleGpuActivity(const T* act, ActivityLogger* logger); void logDeferredEvents(); + void recordGpuVersions(uint32_t cuptiVersion); + + // Defer the CUDA/CUPTI probe until tracing starts, after any worker fork. + // Cache the result because an unavailable backend stays CPU-only. + std::once_flag cuptiInitializationOnce_; + bool cuptiAvailable_{false}; // Calls to CUPTI is encapsulated behind this interface CuptiActivityApi& cupti_; diff --git a/libkineto/src/init.cpp b/libkineto/src/init.cpp index a090e2d44..4e9fa662a 100644 --- a/libkineto/src/init.cpp +++ b/libkineto/src/init.cpp @@ -7,19 +7,14 @@ */ #include -#include // TODO(T90238193) // @lint-ignore-every CLANGTIDY facebook-hte-RelativeInclude #include "ActivityProfilerProxy.h" -#include "Config.h" #include "ConfigLoader.h" #include "DaemonConfigLoader.h" -#include "DeviceUtil.h" -#include "ThreadUtil.h" #ifdef HAS_CUPTI #include "CuptiActivityApi.h" -#include "CuptiCallbackApi.h" #endif #ifdef HAS_CUPTI_PM_SAMPLING #include "CuptiPMSamplingProfiler.h" @@ -39,33 +34,7 @@ namespace KINETO_NAMESPACE { -#if __linux__ || defined(HAS_CUPTI) -static bool initialized = false; - -static void initProfilers() { - if (!initialized) { - // Caution: `initProfilerIfRegistered` spawns the `updateConfigThread`, so - // either: - // 1. Ensure nothing the main thread does not do anything which could race - // with the `updateConfigThread` (current invariant) - // 2. Guard the raceable data appropriately - libkineto::api().initProfilerIfRegistered(); - initialized = true; - VLOG(0) << "libkineto profilers activated"; - } -} - -#endif // __linux__ || defined(HAS_CUPTI) - #ifdef HAS_CUPTI -static void initProfilersCallback( - [[maybe_unused]] CUpti_CallbackDomain domain, - [[maybe_unused]] CUpti_CallbackId cbid, - [[maybe_unused]] const CUpti_CallbackData* cbInfo) { - VLOG(0) << "CUDA Context created"; - initProfilers(); -} - // Some models suffer from excessive instrumentation code gen // on dynamic attach which can hang for more than 5+ seconds. // If the workload was meant to be traced, preload the CUPTI @@ -79,45 +48,11 @@ static bool shouldPreloadCuptiInstrumentation() { #endif } -bool setupCuptiInitCallback(bool logOnError) { - // libcupti will be lazily loaded on this call. - // If it is not available (e.g. CUDA is not installed), - // then this call will return an error and we just abort init. - auto& cbapi = CuptiCallbackApi::singleton(); - cbapi.initCallbackApi(); - - bool status = false; - - if (cbapi.initSuccess()) { - const CUpti_CallbackDomain domain = CUPTI_CB_DOMAIN_RESOURCE; - status = cbapi.registerCallback( - domain, - CuptiCallbackApi::CuptiCallBackID::RESOURCE_CONTEXT_CREATED, - initProfilersCallback); - if (status) { - status = - cbapi.enableCallback(domain, CUPTI_CBID_RESOURCE_CONTEXT_CREATED); - } - } - - if (!cbapi.initSuccess() || !status) { - if (logOnError) { - CUPTI_CALL(cbapi.getCuptiStatus()); - LOG(WARNING) << "CUPTI initialization failed - " - << "CUDA profiler activities will be missing"; - LOG(INFO) - << "If you see CUPTI_ERROR_INSUFFICIENT_PRIVILEGES, refer to " - << "https://developer.nvidia.com/nvidia-development-tools-solutions-err-nvgpuctrperm-cupti"; - } - } - - return status; -} #endif // HAS_CUPTI } // namespace KINETO_NAMESPACE -// Callback interface with CUPTI and library constructors +// Library initialization entry points using namespace KINETO_NAMESPACE; extern "C" { @@ -130,20 +65,18 @@ void libkineto_init(bool cpuOnly, [[maybe_unused]] bool logOnError) { SET_LOG_SEVERITY_LEVEL(atoi(logLevelEnv)); } - // Factory to connect to open source daemon if present #if __linux__ - if (libkineto::isDaemonEnvVarSet()) { - LOG(INFO) << "Registering daemon config loader, cpuOnly = " << cpuOnly; + // KINETO_USE_DAEMON selects the OSS IPC loader only when the platform has + // not already registered its own loader. + const bool registerOssDaemon = libkineto::isDaemonEnvVarSet() && + !ConfigLoader::hasDaemonConfigLoaderFactory(); + if (registerOssDaemon) { + LOG(INFO) << "Registering daemon config loader, cpuOnly = " << cpuOnly; DaemonConfigLoader::registerFactory(); } #endif #ifdef HAS_CUPTI - if (!cpuOnly && !libkineto::isDaemonEnvVarSet()) { - bool success = setupCuptiInitCallback(logOnError); - cpuOnly = !success; - } - if (!cpuOnly && shouldPreloadCuptiInstrumentation()) { CuptiActivityApi::forceLoadCupti(); } @@ -155,8 +88,9 @@ void libkineto_init(bool cpuOnly, [[maybe_unused]] bool logOnError) { } #endif - ConfigLoader& config_loader = libkineto::api().configLoader(); - libkineto::api().registerProfiler( + auto& api = libkineto::api(); + ConfigLoader& config_loader = api.configLoader(); + api.registerProfiler( std::make_unique(cpuOnly, config_loader)); #ifdef HAS_CUPTI_PM_SAMPLING @@ -170,23 +104,20 @@ void libkineto_init(bool cpuOnly, [[maybe_unused]] bool logOnError) { #ifdef HAS_XPUPTI if (!cpuOnly) { - // register xpu pti profiler - libkineto::api().registerProfilerFactory( - []() -> std::unique_ptr { - XPUPTI_CALL( - ptiViewGPULocalAvailable(), - /*error_message=*/"Failed to enable Kineto Profiler on XPU."); - XpuptiScopeProfilerConfig::registerFactory(); - return std::make_unique(); - }); + api.registerProfilerFactory([]() -> std::unique_ptr { + XPUPTI_CALL( + ptiViewGPULocalAvailable(), + /*error_message=*/"Failed to enable Kineto Profiler on XPU."); + XpuptiScopeProfilerConfig::registerFactory(); + return std::make_unique(); + }); } #endif // HAS_XPUPTI #if __linux__ - // For open source users that would like to connect to a profiling daemon - // we should always initialize the profiler at this point. - if (libkineto::isDaemonEnvVarSet()) { - initProfilers(); + if (registerOssDaemon) { + api.initProfilerIfRegistered(); + VLOG(0) << "libkineto profilers activated"; } #endif } @@ -196,6 +127,7 @@ void libkineto_init(bool cpuOnly, [[maybe_unused]] bool logOnError) { int InitializeInjection(void) { LOG(INFO) << "Injection mode: Initializing libkineto"; libkineto_init(false /*cpuOnly*/, true /*logOnError*/); + libkineto::api().initProfilerIfRegistered(); return 1; } diff --git a/libkineto/test/CuptiActivityProfilerTest.cpp b/libkineto/test/CuptiActivityProfilerTest.cpp index 8433d36dd..480788828 100644 --- a/libkineto/test/CuptiActivityProfilerTest.cpp +++ b/libkineto/test/CuptiActivityProfilerTest.cpp @@ -288,6 +288,11 @@ struct MockCuptiActivityBuffer { // Mock parts of the CuptiActivityApi class MockCuptiActivities : public CuptiActivityApi { public: + bool isAvailable(uint32_t& version) const override { + version = 0; + return available; + } + const std::pair processActivities( [[maybe_unused]] CuptiActivityBufferMap& bufferMap, const std::function& handler) override { @@ -313,6 +318,7 @@ class MockCuptiActivities : public CuptiActivityApi { } std::unique_ptr activityBuffer; + bool available{true}; }; // Common setup / teardown and helper functions @@ -334,6 +340,29 @@ class CuptiActivityProfilerTest : public ::testing::Test { ActivityLoggerFactory loggerFactory; }; +TEST_F(CuptiActivityProfilerTest, UnavailableCuptiFallsBackToCpuOnly) { + cuptiActivities_.available = false; + const auto startTime = std::chrono::system_clock::now(); + constexpr auto duration = std::chrono::nanoseconds(300); + + profiler_->configure(*cfg_, startTime); + profiler_->startTrace(startTime); + + const auto startTimeNs = libkineto::timeSinceEpoch(startTime); + auto cpuOps = std::make_unique( + startTimeNs, startTimeNs + duration.count()); + cpuOps->addOp("cpu_op", startTimeNs + 20, startTimeNs + 50, 1); + profiler_->transferCpuTrace(std::move(cpuOps)); + profiler_->stopTrace(startTime + duration); + + auto logger = std::make_unique(*cfg_); + profiler_->processTrace(*logger); + + ActivityTrace trace(std::move(logger), loggerFactory); + ASSERT_EQ(trace.activities()->size(), 1); + EXPECT_EQ(trace.activities()->front()->name(), "cpu_op"); +} + TEST_F(CuptiActivityProfilerTest, SyncTrace) { // Verbose logging is useful for debugging std::vector log_modules({"CuptiActivityProfiler.cpp"});