Support child profilers adding their deviceProperties to logger (#1008)

Summary: This patch is intended to support our out-of-tree Kineto plugin in adding its deviceProperties. Reviewed By: aaronenyeshi Differential Revision: D65496183 Pulled By: sraikund16
pytorch · Nov 6, 2024 · 22c9ff8 · 22c9ff8
1 parent 956c0f0
commit 22c9ff8
Show file tree

Hide file tree

Showing 6 changed files with 30 additions and 8 deletions.
diff --git a/libkineto/include/IActivityProfiler.h b/libkineto/include/IActivityProfiler.h
@@ -126,6 +126,10 @@ class IActivityProfilerSession {
   virtual void pushUserCorrelationId(uint64_t /*id*/) {}
   virtual void popUserCorrelationId() {}
 
+  virtual std::string getDeviceProperties() {
+    return "";
+  }
+
  protected:
   TraceStatus status_ = TraceStatus::READY;
 };

diff --git a/libkineto/include/output_base.h b/libkineto/include/output_base.h
@@ -56,10 +56,11 @@ class ActivityLogger {
       const libkineto::GenericTraceActivity& activity) = 0;
 
   virtual void handleTraceStart(
-      const std::unordered_map<std::string, std::string>& metadata) = 0;
+      const std::unordered_map<std::string, std::string>& metadata,
+      const std::string& device_properties) = 0;
 
   void handleTraceStart() {
-    handleTraceStart(std::unordered_map<std::string, std::string>());
+    handleTraceStart(std::unordered_map<std::string, std::string>(), "");
   }
 
   virtual void finalizeTrace(

diff --git a/libkineto/src/CuptiActivityProfiler.cpp b/libkineto/src/CuptiActivityProfiler.cpp
@@ -8,6 +8,7 @@
 
 #include "CuptiActivityProfiler.h"
 #include <fmt/format.h>
+#include <fmt/ranges.h>
 #include <time.h>
 #include <atomic>
 #include <cstdint>
@@ -30,6 +31,7 @@
 
 #include "Config.h"
 #include "DeviceUtil.h"
+#include "DeviceProperties.h"
 #include "time_since_epoch.h"
 #ifdef HAS_CUPTI
 #include "CuptiActivity.cpp"
@@ -306,7 +308,17 @@ void CuptiActivityProfiler::processTraceInternal(ActivityLogger& logger) {
   for (auto& pair : versionMetadata_) {
     addMetadata(pair.first, pair.second);
   }
-  logger.handleTraceStart(metadata_);
+  std::vector<std::string> device_properties;
+  if (auto props = devicePropertiesJson(); !props.empty()) {
+    device_properties.push_back(props);
+  }
+  for (const auto& session : sessions_) {
+    if (auto props = session->getDeviceProperties(); !props.empty()) {
+      device_properties.push_back(props);
+    }
+  }
+  logger.handleTraceStart(
+      metadata_, fmt::format("{}", fmt::join(device_properties, ",")));
   setCpuActivityPresent(false);
   setGpuActivityPresent(false);
   for (auto& cpu_trace : traceBuffers_->cpu) {

diff --git a/libkineto/src/output_json.cpp b/libkineto/src/output_json.cpp
@@ -105,7 +105,8 @@ void ChromeTraceLogger::metadataToJSON(
 }
 
 void ChromeTraceLogger::handleTraceStart(
-    const std::unordered_map<std::string, std::string>& metadata) {
+    const std::unordered_map<std::string, std::string>& metadata,
+    const std::string& device_properties) {
   traceOf_ << fmt::format(
       R"JSON(
 {{
@@ -116,7 +117,7 @@ void ChromeTraceLogger::handleTraceStart(
       R"JSON(
   "deviceProperties": [{}
   ],)JSON",
-      devicePropertiesJson());
+      device_properties);
 
   metadataToJSON(metadata);
   traceOf_ << R"JSON(

diff --git a/libkineto/src/output_json.h b/libkineto/src/output_json.h
@@ -71,7 +71,8 @@ class ChromeTraceLogger : public libkineto::ActivityLogger {
   void handleGenericActivity(const GenericTraceActivity& activity) override;
 
   void handleTraceStart(
-      const std::unordered_map<std::string, std::string>& metadata) override;
+      const std::unordered_map<std::string, std::string>& metadata,
+      const std::string& device_properties) override;
 
   void finalizeTrace(
       const Config& config,

diff --git a/libkineto/src/output_membuf.h b/libkineto/src/output_membuf.h
@@ -62,8 +62,10 @@ class MemoryTraceLogger : public ActivityLogger {
   }
 
   void handleTraceStart(
-      const std::unordered_map<std::string, std::string>& metadata) override {
+      const std::unordered_map<std::string, std::string>& metadata,
+      const std::string& device_properties) override {
     metadata_ = metadata;
+    device_properties_ = device_properties;
   }
 
   void finalizeTrace(
@@ -81,7 +83,7 @@ class MemoryTraceLogger : public ActivityLogger {
   }
 
   void log(ActivityLogger& logger) {
-    logger.handleTraceStart(metadata_);
+    logger.handleTraceStart(metadata_, device_properties_);
     for (auto& activity : activities_) {
       activity->log(logger);
     }
@@ -121,6 +123,7 @@ class MemoryTraceLogger : public ActivityLogger {
   std::unique_ptr<ActivityBuffers> buffers_;
   std::unordered_map<std::string, std::string> metadata_;
   std::unordered_map<std::string, std::vector<std::string>> loggerMetadata_;
+  std::string device_properties_;
   int64_t endTime_{0};
   std::shared_ptr<ActivityLogger> chrome_logger_;
 };