Support for trace metadata

libkineto/include/ActivityProfilerInterface.h

@@ -81,6 +81,10 @@ class ActivityProfilerInterface {
   // Maps kernel thread id -> pthread id for CPU ops.
   // Client must record any new kernel thread where the activity has occured.
   virtual void recordThreadInfo(pid_t tid, pthread_t pthreadId) {}
+
+  // Record trace metadata, currently supporting only string key and values,
+  // values with the same key are overwritten
+  virtual void addMetadata(const std::string& key, const std::string& value) = 0;
 };
 
 } // namespace libkineto

libkineto/src/ActivityProfiler.cpp

@@ -127,6 +127,7 @@ void ActivityProfiler::processTraceInternal(ActivityLogger& logger) {
       << " CPU buffers";
   VLOG(0) << "Profile time range: " << captureWindowStartTime_ << " - "
           << captureWindowEndTime_;
+  logger.handleTraceStart(metadata_);
   for (auto& cpu_trace : traceBuffers_->cpu) {
     string trace_name = cpu_trace->span.name;
     VLOG(0) << "Processing CPU buffer for " << trace_name << " ("
@@ -696,6 +697,7 @@ void ActivityProfiler::finalizeTrace(const Config& config, ActivityLogger& logge
       }
     }
   }
+
   // Thread info
   for (auto pair : threadInfo_) {
     const auto& thread_info = pair.second;
@@ -728,6 +730,7 @@ void ActivityProfiler::resetTraceData() {
   clientActivityTraceMap_.clear();
   disabledTraceSpans_.clear();
   traceBuffers_ = nullptr;
+  metadata_.clear();
 }
 
 

libkineto/src/ActivityProfiler.h

@@ -111,6 +111,11 @@ class ActivityProfiler {
     }
   }
 
+  void addMetadata(const std::string& key, const std::string& value) {
+    std::lock_guard<std::mutex> guard(mutex_);
+    metadata_[key] = value;
+  }
+
  private:
   class ExternalEventMap {
    public:
@@ -366,6 +371,9 @@ class ActivityProfiler {
   // Buffers where trace data is stored
   std::unique_ptr<ActivityBuffers> traceBuffers_;
 
+  // Trace metadata
+  std::unordered_map<std::string, std::string> metadata_;
+
 };
 
 } // namespace KINETO_NAMESPACE

libkineto/src/ActivityProfilerController.cpp

@@ -147,4 +147,9 @@ std::unique_ptr<ActivityTraceInterface> ActivityProfilerController::stopTrace()
   return std::make_unique<ActivityTrace>(std::move(logger), CuptiActivityInterface::singleton());
 }
 
+void ActivityProfilerController::addMetadata(
+    const std::string& key, const std::string& value) {
+  profiler_->addMetadata(key, value);
+}
+
 } // namespace KINETO_NAMESPACE

libkineto/src/ActivityProfilerController.h

@@ -58,9 +58,11 @@ class ActivityProfilerController {
   }
 
   void recordThreadInfo(pid_t tid, pthread_t pthreadId) {
-    profiler_->recordThreadInfo(tid, pthreadId); 
+    profiler_->recordThreadInfo(tid, pthreadId);
   }
 
+  void addMetadata(const std::string& key, const std::string& value);
+
  private:
   void profilerLoop();
 

libkineto/src/ActivityProfilerProxy.cpp

@@ -85,7 +85,12 @@ bool ActivityProfilerProxy::enableForRegion(const std::string& match) {
 }
 
 void ActivityProfilerProxy::recordThreadInfo(pid_t tid, pthread_t pthreadId) {
-  controller_->recordThreadInfo(tid, pthreadId); 
+  controller_->recordThreadInfo(tid, pthreadId);
+}
+
+void ActivityProfilerProxy::addMetadata(
+    const std::string& key, const std::string& value) {
+  controller_->addMetadata(key, value);
 }
 
 } // namespace libkineto

libkineto/src/ActivityProfilerProxy.h

@@ -60,6 +60,8 @@ class ActivityProfilerProxy : public ActivityProfilerInterface {
 
   bool enableForRegion(const std::string& match) override;
 
+  void addMetadata(const std::string& key, const std::string& value) override;
+
  private:
   bool cpuOnly_{true};
   ActivityProfilerController* controller_{nullptr};

libkineto/src/output_base.h

@@ -67,6 +67,13 @@ class ActivityLogger {
       const GpuActivity<CUpti_ActivityMemset>& activity) = 0;
 #endif // HAS_CUPTI
 
+  virtual void handleTraceStart(
+      const std::unordered_map<std::string, std::string>& metadata) = 0;
+
+  void handleTraceStart() {
+    handleTraceStart(std::unordered_map<std::string, std::string>());
+  }
+
   virtual void finalizeTrace(
       const KINETO_NAMESPACE::Config& config,
       std::unique_ptr<ActivityBuffers> buffers,

libkineto/src/output_json.cpp

@@ -31,28 +31,48 @@ namespace KINETO_NAMESPACE {
 
 static constexpr int kSchemaVersion = 1;
 
-static void writeHeader(std::ofstream& stream) {
-  stream << fmt::format(R"JSON(
+void ChromeTraceLogger::handleTraceStart(
+    const std::unordered_map<std::string, std::string>& metadata) {
+  traceOf_ << fmt::format(R"JSON(
 {{
   "schemaVersion": {},
-  "traceEvents": [
   )JSON", kSchemaVersion);
+
+  if (!metadata.empty()) {
+    traceOf_ << R"JSON(
+  "metadata": {
+  )JSON";
+    bool first = true;
+    for (const auto& kv : metadata) {
+      if (!first) {
+        traceOf_ << ",\n";
+      }
+      traceOf_ << fmt::format(R"(    "{}": "{}")", kv.first, kv.second);
+      first = false;
+    }
+    traceOf_ << R"JSON(
+  },
+  )JSON";
+  }
+
+  traceOf_ << R"JSON(
+  "traceEvents": [
+  )JSON";
 }
 
-static void openTraceFile(std::string& name, std::ofstream& stream) {
-  stream.open(name, std::ofstream::out | std::ofstream::trunc);
-  if (!stream) {
-    PLOG(ERROR) << "Failed to open '" << name << "'";
+void ChromeTraceLogger::openTraceFile() {
+  traceOf_.open(fileName_, std::ofstream::out | std::ofstream::trunc);
+  if (!traceOf_) {
+    PLOG(ERROR) << "Failed to open '" << fileName_ << "'";
   } else {
-    LOG(INFO) << "Tracing to " << name;
-    writeHeader(stream);
+    LOG(INFO) << "Tracing to " << fileName_;
   }
 }
 
 ChromeTraceLogger::ChromeTraceLogger(const std::string& traceFileName, int smCount)
     : fileName_(traceFileName), pid_(getpid()) {
   traceOf_.clear(std::ios_base::badbit);
-  openTraceFile(fileName_, traceOf_);
+  openTraceFile();
 #ifdef HAS_CUPTI
   smCount_ = CuptiActivityInterface::singleton().smCount();
 #endif

libkineto/src/output_json.h

@@ -60,6 +60,9 @@ class ChromeTraceLogger : public libkineto::ActivityLogger {
   void handleGpuActivity(const GpuActivity<CUpti_ActivityMemset>& activity) override;
 #endif // HAS_CUPTI
 
+  void handleTraceStart(
+      const std::unordered_map<std::string, std::string>& metadata) override;
+
   void finalizeTrace(
       const Config& config,
       std::unique_ptr<ActivityBuffers> buffers,
@@ -73,6 +76,8 @@ class ChromeTraceLogger : public libkineto::ActivityLogger {
   void handleLinkEnd(const TraceActivity& e);
 #endif // HAS_CUPTI
 
+  void openTraceFile();
+
   std::string fileName_;
   std::ofstream traceOf_;
 

libkineto/src/output_membuf.h

@@ -87,6 +87,11 @@ class MemoryTraceLogger : public ActivityLogger {
   }
 #endif // HAS_CUPTI
 
+  void handleTraceStart(
+      const std::unordered_map<std::string, std::string>& metadata) override {
+    metadata_ = metadata;
+  }
+
   void finalizeTrace(
       const Config& config,
       std::unique_ptr<ActivityBuffers> buffers,
@@ -100,6 +105,7 @@ class MemoryTraceLogger : public ActivityLogger {
   }
 
   void log(ActivityLogger& logger) {
+    logger.handleTraceStart(metadata_);
     for (auto& activity : activities_) {
       activity->log(logger);
     }
@@ -151,6 +157,7 @@ class MemoryTraceLogger : public ActivityLogger {
   std::vector<TraceSpan> traceSpanList_;
   std::vector<TraceSpan> iterationList_;
   std::unique_ptr<ActivityBuffers> buffers_;
+  std::unordered_map<std::string, std::string> metadata_;
   int64_t endTime_{0};
 };