Support child profilers add their deviceProperties to logger

pytorch · Oct 31, 2024 · 54fe0ed · 54fe0ed
1 parent 5f5dc26
commit 54fe0ed
Show file tree

Hide file tree

Showing 6 changed files with 50 additions and 2 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
@@ -58,6 +58,10 @@ class ActivityLogger {
   virtual void handleTraceStart(
       const std::unordered_map<std::string, std::string>& metadata) = 0;
 
+  virtual void handleTraceStart(
+      const std::unordered_map<std::string, std::string>& metadata,
+      const std::string& device_properties) {}
+
   void handleTraceStart() {
     handleTraceStart(std::unordered_map<std::string, std::string>());
   }

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
@@ -123,6 +123,22 @@ void ChromeTraceLogger::handleTraceStart(
   "traceEvents": [)JSON";
 }
 
+void ChromeTraceLogger::handleTraceStart(
+    const std::unordered_map<std::string, std::string>& metadata,
+    const std::string& device_properties) {
+  traceOf_ << fmt::format(R"JSON(
+{{
+  "schemaVersion": {},)JSON", kSchemaVersion);
+
+  traceOf_ << fmt::format(R"JSON(
+  "deviceProperties": [{}
+  ],)JSON", device_properties);
+
+  metadataToJSON(metadata);
+  traceOf_ << R"JSON(
+  "traceEvents": [)JSON";
+}
+
 static std::string defaultFileName() {
   return fmt::format(kDefaultLogFileFmt, processId());
 }

diff --git a/libkineto/src/output_json.h b/libkineto/src/output_json.h
@@ -73,6 +73,10 @@ class ChromeTraceLogger : public libkineto::ActivityLogger {
   void handleTraceStart(
       const std::unordered_map<std::string, std::string>& metadata) override;
 
+  void handleTraceStart(
+      const std::unordered_map<std::string, std::string>& metadata,
+      const std::string& device_properties) override;
+
   void finalizeTrace(
       const Config& config,
       std::unique_ptr<ActivityBuffers> buffers,

diff --git a/libkineto/src/output_membuf.h b/libkineto/src/output_membuf.h
@@ -66,6 +66,13 @@ class MemoryTraceLogger : public ActivityLogger {
     metadata_ = metadata;
   }
 
+  void handleTraceStart(
+      const std::unordered_map<std::string, std::string>& metadata,
+      const std::string& device_properties) override {
+    metadata_ = metadata;
+    device_properties_ = device_properties;
+  }
+
   void finalizeTrace(
       const Config& config,
       std::unique_ptr<ActivityBuffers> buffers,
@@ -81,7 +88,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 +128,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_;
 };