Remove eager CUPTI resource callback subscription (#1250)

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;

libkineto/src/ConfigLoader.h

@@ -88,6 +88,7 @@ class ConfigLoader {
 
   static void setDaemonConfigLoaderFactory(
       std::function<std::unique_ptr<IDaemonConfigLoader>()> factory);
+  static bool hasDaemonConfigLoaderFactory();
 
   std::string getConfString();
 

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)) {

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<ActivityType>& selected_activities,
       bool enablePerThreadBuffers = false);

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,17 +103,39 @@ 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(),
       derivedConfig_->isPerThreadBufferEnabled());
 }
 
 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();
 }

libkineto/src/CuptiActivityProfiler.h

@@ -9,6 +9,7 @@
 #pragma once
 
 #include <cupti.h>
+#include <mutex>
 #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 <class T>
   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_;

libkineto/src/init.cpp

@@ -7,19 +7,14 @@
  */
 
 #include <memory>
-#include <mutex>
 
 // 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<ActivityProfilerProxy>(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<IActivityProfiler> {
-          XPUPTI_CALL(
-              ptiViewGPULocalAvailable(),
-              /*error_message=*/"Failed to enable Kineto Profiler on XPU.");
-          XpuptiScopeProfilerConfig::registerFactory();
-          return std::make_unique<XPUActivityProfiler>();
-        });
+    api.registerProfilerFactory([]() -> std::unique_ptr<IActivityProfiler> {
+      XPUPTI_CALL(
+          ptiViewGPULocalAvailable(),
+          /*error_message=*/"Failed to enable Kineto Profiler on XPU.");
+      XpuptiScopeProfilerConfig::registerFactory();
+      return std::make_unique<XPUActivityProfiler>();
+    });
   }
 #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;
 }
 

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<int, size_t> processActivities(
       [[maybe_unused]] CuptiActivityBufferMap& bufferMap,
       const std::function<void(const CUpti_Activity*)>& handler) override {
@@ -313,6 +318,7 @@ class MockCuptiActivities : public CuptiActivityApi {
   }
 
   std::unique_ptr<MockCuptiActivityBuffer> 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<MockCpuActivityBuffer>(
+      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<MemoryTraceLogger>(*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<std::string> log_modules({"CuptiActivityProfiler.cpp"});