diff --git a/libkineto/src/output_json.cpp b/libkineto/src/output_json.cpp index cf26b5e8d..d7d9a250f 100644 --- a/libkineto/src/output_json.cpp +++ b/libkineto/src/output_json.cpp @@ -49,6 +49,7 @@ static constexpr const std::string_view kOutTensorsStart = static constexpr const std::string_view kRank = "Rank"; static constexpr const std::string_view kP2pSrc = "Src Rank"; static constexpr const std::string_view kP2pDst = "Dst Rank"; +static constexpr const std::string_view kCommsId = "Comms Id"; #ifdef __linux__ static constexpr std::string_view kDefaultLogFileFmt = @@ -555,6 +556,15 @@ void ChromeTraceLogger::handleActivity(const libkineto::ITraceActivity& op) { arg_values.append(fmt::format(", \"{}\": {}", kP2pSrc, srcRank)); } + const auto& commsId = + collectiveRecord->getMetadataValue(std::string(kCommsId)); + if (!commsId.empty()) { + if (!arg_values.empty()) { + arg_values.append(","); + } + arg_values.append(fmt::format(" \"{}\": {}", kCommsId, commsId)); + } + if (distInfo_.backend.empty() && processGroupDesc == "\"default_pg\"") { distInfo_.backend = "nccl"; distInfo_.rank = collectiveRecord->getMetadataValue(std::string(kRank)); diff --git a/libkineto/test/CuptiActivityProfilerTest.cpp b/libkineto/test/CuptiActivityProfilerTest.cpp index bd886dcd2..dabe29a37 100644 --- a/libkineto/test/CuptiActivityProfilerTest.cpp +++ b/libkineto/test/CuptiActivityProfilerTest.cpp @@ -52,6 +52,7 @@ static constexpr auto kGroupSize = "Group size"; static constexpr const char* kProcessGroupName = "Process Group Name"; static constexpr const char* kProcessGroupDesc = "Process Group Description"; static constexpr const char* kGroupRanks = "Process Group Ranks"; +static constexpr const char* kCommsId = "Comms Id"; static constexpr int32_t kTruncatLength = 30; #define CUDA_LAUNCH_KERNEL CUPTI_RUNTIME_TRACE_CBID_cudaLaunchKernel_v7000 @@ -663,6 +664,7 @@ TEST_F(CuptiActivityProfilerTest, GpuNCCLCollectiveTest) { metadataMap.emplace(kGroupSize, "2"); metadataMap.emplace(kProcessGroupName, fmt::format("\"{}\"", "12341234")); metadataMap.emplace(kProcessGroupDesc, fmt::format("\"{}\"", "test_purpose")); + metadataMap.emplace(kCommsId, "12345678"); std::vector inSplitSizes(50, 0); std::string inSplitSizesStr; @@ -804,6 +806,8 @@ TEST_F(CuptiActivityProfilerTest, GpuNCCLCollectiveTest) { EXPECT_EQ(2, countSubstrings(jsonString, "test_purpose")); EXPECT_EQ(2, countSubstrings(jsonString, kGroupRanks)); EXPECT_EQ(2, countSubstrings(jsonString, expectedGroupRanksStr)); + EXPECT_EQ(2, countSubstrings(jsonString, kCommsId)); + EXPECT_EQ(2, countSubstrings(jsonString, "12345678")); #endif }