Skip to content

Expose comms_id in traces #1286

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ycui1984 wants to merge 1 commit into
base: main
Choose a base branch
from
+14 −0
Open

Expose comms_id in traces #1286

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions libkineto/src/output_json.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 =
Expand Down Expand Up @@ -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));
Expand Down
4 changes: 4 additions & 0 deletions libkineto/test/CuptiActivityProfilerTest.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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<int64_t> inSplitSizes(50, 0);
std::string inSplitSizesStr;
Expand Down Expand Up @@ -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
}

Expand Down
Loading