Merge branch 'main' into zejun/add_unified_runtime_op_for_flow_in_jso…

…n_trace_file
pytorch · Jan 10, 2025 · 539879c · 539879c
2 parents 6ff5f1b + 7ae1ce8
commit 539879c
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 0 deletions.
diff --git a/libkineto/include/ActivityType.h b/libkineto/include/ActivityType.h
@@ -41,6 +41,7 @@ enum class ActivityType {
   XPU_RUNTIME, // host side xpu runtime events
   COLLECTIVE_COMM, // collective communication
   MTIA_WORKLOADD, // MTIA workloadd events
+  MTIA_COLLECTIVE_COMM, // MTIA collective communication events
 
   // PRIVATEUSE1 Activity types are used for custom backends.
   // The corresponding device type is `DeviceType::PrivateUse1` in PyTorch.

diff --git a/libkineto/src/ActivityType.cpp b/libkineto/src/ActivityType.cpp
@@ -39,6 +39,7 @@ static constexpr std::array<ActivityTypeName, activityTypeCount + 1> map{
      {"xpu_runtime", ActivityType::XPU_RUNTIME},
      {"collective_comm", ActivityType::COLLECTIVE_COMM},
      {"mtia_workloadd", ActivityType::MTIA_WORKLOADD},
+     {"mtia_collective_comm", ActivityType::MTIA_COLLECTIVE_COMM},
      {"privateuse1_runtime", ActivityType::PRIVATEUSE1_RUNTIME},
      {"privateuse1_driver", ActivityType::PRIVATEUSE1_DRIVER},
      {"ENUM_COUNT", ActivityType::ENUM_COUNT}}};

diff --git a/libkineto/src/output_json.cpp b/libkineto/src/output_json.cpp
@@ -79,6 +79,17 @@ void ChromeTraceLogger::sanitizeStrForJSON(std::string& value) {
   value.erase(std::remove(value.begin(), value.end(), '\n'), value.end());
 }
 
+static void sanitizeForNonReadableChars(std::string& value) {
+  for (auto& c : value) {
+    if (!std::isprint(c)) {
+      LOG(WARNING) << "Non JSON compliant character found in string: " << value
+                   << " Replacing with 'unknown'";
+      value = "unknown";
+      break;
+    }
+  }
+}
+
 static inline int32_t sanitizeTid(int32_t tid) {
   // Convert all negative tids to its positive value. Create a specific case
   // for INT_MIN so it is obvious how it is being handled.
@@ -521,6 +532,7 @@ void ChromeTraceLogger::handleActivity(const libkineto::ITraceActivity& op) {
   // TODO: Remove this once legacy tools are updated.
   std::string op_name = op.name() == "kernel" ? "Kernel" : op.name();
   sanitizeStrForJSON(op_name);
+  sanitizeForNonReadableChars(op_name);
 
   // clang-format off
   ts = transToRelativeTime(ts);