From 68d4a32bcae9888799a76b964a6679e2bc2012a8 Mon Sep 17 00:00:00 2001
From: Darshan Sanghani <dsang@meta.com>
Date: Fri, 22 Nov 2024 12:14:19 -0800
Subject: [PATCH] Setters and getters for trace IDs (#967)

Summary:

This commit moves the trace ID initialization logic to be inline with how ActivityProfilers interact with LoggerObservers by setting a default local trace ID when a profiler config does not contain a trace ID.

The USTLoggerCollector, which is our internal LoggerObserver will record the trace ID for a given environment (which is a PID today) using the `setTraceID` method that all LoggerObservers must staisfy. Additionally any internal calls to read this trace ID, for example from our ManifoldChromeTrace logger, may use a special `getTraceID` method that ships with `USTLoggerCollector`.

Group trace IDs are handled accordingly.

Reviewed By: aaronenyeshi

Differential Revision: D60267172
---
 libkineto/src/ActivityProfilerController.cpp | 5 +++++
 libkineto/src/ActivityProfilerController.h   | 1 +
 libkineto/src/CuptiActivityProfiler.cpp      | 6 ++----
 3 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/libkineto/src/ActivityProfilerController.cpp b/libkineto/src/ActivityProfilerController.cpp
index 99c629cfb..19e204868 100644
--- a/libkineto/src/ActivityProfilerController.cpp
+++ b/libkineto/src/ActivityProfilerController.cpp
@@ -40,6 +40,11 @@ void ActivityProfilerController::setLoggerCollectorFactory(
     std::function<std::shared_ptr<LoggerCollector>()> factory) {
   loggerCollectorFactory() = factory();
 }
+
+std::shared_ptr<LoggerCollector>
+ActivityProfilerController::getLoggerCollector() {
+  return loggerCollectorFactory();
+}
 #endif // !USE_GOOGLE_LOG
 
 ActivityProfilerController::ActivityProfilerController(
diff --git a/libkineto/src/ActivityProfilerController.h b/libkineto/src/ActivityProfilerController.h
index 47d21b245..b4f826a7d 100644
--- a/libkineto/src/ActivityProfilerController.h
+++ b/libkineto/src/ActivityProfilerController.h
@@ -40,6 +40,7 @@ class ActivityProfilerController : public ConfigLoader::ConfigHandler {
   ~ActivityProfilerController();
 
 #if !USE_GOOGLE_LOG
+  static std::shared_ptr<LoggerCollector> getLoggerCollector();
   static void setLoggerCollectorFactory(
       std::function<std::shared_ptr<LoggerCollector>()> factory);
 #endif // !USE_GOOGLE_LOG
diff --git a/libkineto/src/CuptiActivityProfiler.cpp b/libkineto/src/CuptiActivityProfiler.cpp
index 4eadec84e..bf0491235 100644
--- a/libkineto/src/CuptiActivityProfiler.cpp
+++ b/libkineto/src/CuptiActivityProfiler.cpp
@@ -1054,13 +1054,11 @@ void CuptiActivityProfiler::configure(
 
   // Set useful metadata into the logger.
   LOGGER_OBSERVER_SET_TRACE_DURATION_MS(config_->activitiesDuration().count());
+  LOGGER_OBSERVER_SET_TRACE_ID(config_->requestTraceID());
+  LOGGER_OBSERVER_SET_GROUP_TRACE_ID(config_->requestGroupTraceID());
   if (!config_->requestTraceID().empty()) {
-    LOGGER_OBSERVER_SET_TRACE_ID(config_->requestTraceID());
     addMetadata("trace_id", "\"" + config_->requestTraceID() + "\"");
   }
-  if (!config_->requestGroupTraceID().empty()) {
-    LOGGER_OBSERVER_SET_GROUP_TRACE_ID(config_->requestGroupTraceID());
-  }
 
 #if defined(HAS_CUPTI) || defined(HAS_ROCTRACER)
   if (!cpuOnly_) {