From 70f735ca626e2881b5881c342e969d078aa04c80 Mon Sep 17 00:00:00 2001
From: Jay Chae <jchae@fb.com>
Date: Tue, 13 Apr 2021 22:39:38 -0700
Subject: [PATCH] add JSON-formatted metadata for ClientTraceActivity (#164)

Summary:
Pull Request resolved: https://github.com/pytorch/kineto/pull/164

Plan is to deprecate ClientTraceActivity and use a single concrete data structure to log all client side activities :`GenericTraceActivity`. CTA has a lot of explicit fields for metadata that are rarely used which results in sparse definition of activities. Instead of adding a new field, encode the key/value as a string. Note, more type-safe and rich alternative is `std::map` but we are not using it for performance reason.

Reviewed By: gdankel

Differential Revision: D27298829

fbshipit-source-id: a1592484333e42f33fbe0805803cf91fa83b52b9
---
 libkineto/CMakeLists.txt                |  1 +
 libkineto/include/ClientTraceActivity.h | 15 +++++++++++++++
 libkineto/libkineto_defs.bzl            |  1 +
 3 files changed, 17 insertions(+)

diff --git a/libkineto/CMakeLists.txt b/libkineto/CMakeLists.txt
index 003de73b0..bed5c251a 100644
--- a/libkineto/CMakeLists.txt
+++ b/libkineto/CMakeLists.txt
@@ -112,6 +112,7 @@ target_include_directories(kineto_base PUBLIC
       $<BUILD_INTERFACE:${CUDA_INCLUDE_DIRS}>)
 
 target_include_directories(kineto_api PUBLIC
+      $<BUILD_INTERFACE:${FMT_INCLUDE_DIR}>
       $<BUILD_INTERFACE:${LIBKINETO_INCLUDE_DIR}>)
 
 if(KINETO_LIBRARY_TYPE STREQUAL "default")
diff --git a/libkineto/include/ClientTraceActivity.h b/libkineto/include/ClientTraceActivity.h
index aefd2f6dd..670558b43 100644
--- a/libkineto/include/ClientTraceActivity.h
+++ b/libkineto/include/ClientTraceActivity.h
@@ -7,8 +7,10 @@
 
 #pragma once
 
+#include <fmt/format.h>
 #include <string>
 #include <thread>
+#include <vector>
 
 #include "TraceActivity.h"
 
@@ -56,6 +58,16 @@ struct ClientTraceActivity : TraceActivity {
     // Unimplemented by default
   }
 
+  // Encode client side metadata as a key/value string.
+  void addMetadata(const std::string& key, const std::string& value) {
+    auto kv = fmt::format("\"{}\": {}", key, value);
+    metadata_.push_back(std::move(kv));
+  }
+
+  [[nodiscard]] const std::string getMetadata() const {
+    return fmt::format("{}", fmt::join(metadata_, ", "));
+  }
+
   int64_t startTime{0};
   int64_t endTime{0};
   int64_t correlation{0};
@@ -72,6 +84,9 @@ struct ClientTraceActivity : TraceActivity {
   std::string inputNames;
   std::string outputNames;
   std::string callStack;
+
+ private:
+  std::vector<std::string> metadata_;
 };
 
 } // namespace libkineto
diff --git a/libkineto/libkineto_defs.bzl b/libkineto/libkineto_defs.bzl
index 333f9f8c9..e6c5d58f0 100644
--- a/libkineto/libkineto_defs.bzl
+++ b/libkineto/libkineto_defs.bzl
@@ -53,6 +53,7 @@ def get_libkineto_public_headers():
     return [
         "include/ActivityProfilerInterface.h",
         "include/ActivityType.h",
+        "include/ClientTraceActivity.h",
         "include/ClientInterface.h",
         "include/TraceActivity.h",
         "include/TraceSpan.h",