diff --git a/libkineto/src/plugin/xpupti/XpuptiActivityProfilerSession.cpp b/libkineto/src/plugin/xpupti/XpuptiActivityProfilerSession.cpp index 94a43bea0..b864739c1 100644 --- a/libkineto/src/plugin/xpupti/XpuptiActivityProfilerSession.cpp +++ b/libkineto/src/plugin/xpupti/XpuptiActivityProfilerSession.cpp @@ -11,6 +11,7 @@ #include "time_since_epoch.h" +#include #include #include @@ -116,6 +117,14 @@ std::vector XpuptiActivityProfilerSession:: return result; } +std::unordered_map XpuptiActivityProfilerSession:: + getMetadata() { + const char* version = ptiVersionString(); + return { + {"xpupti_version", + fmt::format("\"{}\"", version != nullptr ? version : "unknown")}}; +} + void XpuptiActivityProfilerSession::pushCorrelationId(uint64_t id) { xpti_.pushCorrelationID(id, XpuptiActivityApi::CorrelationFlowType::Default); } diff --git a/libkineto/src/plugin/xpupti/XpuptiActivityProfilerSession.h b/libkineto/src/plugin/xpupti/XpuptiActivityProfilerSession.h index 4fabfdd80..c00cd8917 100644 --- a/libkineto/src/plugin/xpupti/XpuptiActivityProfilerSession.h +++ b/libkineto/src/plugin/xpupti/XpuptiActivityProfilerSession.h @@ -56,6 +56,7 @@ class XpuptiActivityProfilerSession : public libkineto::IActivityProfilerSession } std::vector getResourceInfos() override; std::unique_ptr getTraceBuffer() override; + std::unordered_map getMetadata() override; void pushCorrelationId(uint64_t id) override; void popCorrelationId() override; diff --git a/libkineto/test/xpupti/XpuptiActivityHandlersTest.cpp b/libkineto/test/xpupti/XpuptiActivityHandlersTest.cpp index db62a7787..9c567d138 100644 --- a/libkineto/test/xpupti/XpuptiActivityHandlersTest.cpp +++ b/libkineto/test/xpupti/XpuptiActivityHandlersTest.cpp @@ -90,6 +90,17 @@ class XpuptiActivityHandlersTest : public ::testing::Test { } }; +TEST_F(XpuptiActivityHandlersTest, SessionMetadataIncludesPtiVersion) { + Config config; + std::set activity_types = {ActivityType::COLLECTIVE_COMM}; + KN::XpuptiActivityProfilerSession session( + mockApi_, "__test_profiler__", config, activity_types); + + const auto metadata = session.getMetadata(); + ASSERT_TRUE(metadata.contains("xpupti_version")); + EXPECT_FALSE(metadata.at("xpupti_version").empty()); +} + // --- Communication Activity Tests --- TEST_F(XpuptiActivityHandlersTest, CommunicationActivityHasXcclPrefix) {