From b94209437439ea23977f4e481d89353e32631159 Mon Sep 17 00:00:00 2001 From: Yan Cui Date: Wed, 11 Mar 2026 15:00:07 -0700 Subject: [PATCH] Add comms Id to trace output JSON (#1286) Summary: This is part of a larger effort to expose comms_id in profiler traces to enable correlating the same communication operation across different ranks. This diff adds the Kineto side: reading "Comms Id" from the profiler metadata and writing it into the Chrome trace JSON output. When the metadata key is present, it will be included in the trace event args, making it visible in trace viewers. The PyTorch side (computing and emitting the comms_id) will be in a follow-up diff after the Kineto submodule is updated. Differential Revision: D96153960 --- libkineto/src/output_json.cpp | 10 ++++++++++ libkineto/test/CuptiActivityProfilerTest.cpp | 4 ++++ 2 files changed, 14 insertions(+) diff --git a/libkineto/src/output_json.cpp b/libkineto/src/output_json.cpp index 539ca606f..dfdba816e 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 ff223603d..4c5d475dc 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 }