From 1bb9b7632d129c0ee2eb47501215bcc42a6445ea Mon Sep 17 00:00:00 2001 From: Jakob Johnson Date: Fri, 2 Aug 2024 13:51:27 -0700 Subject: [PATCH] Move group trace id and trace id to OSS kineto Config (#970) Summary: Pull Request resolved: https://github.com/pytorch/kineto/pull/970 Previously the `trace_id` and `group_trace_id` fields were only parsed for `FBConfig`. There are OSS use-cases which would benefit from using these fields, so this diff moves the parsing of these fields to `Config` Reviewed By: aaronenyeshi Differential Revision: D60256788 fbshipit-source-id: 6e9ecccd5f888b704ba30688dff4c41014959347 --- libkineto/src/Config.cpp | 9 +++++++++ libkineto/test/ConfigTest.cpp | 8 ++++++++ 2 files changed, 17 insertions(+) diff --git a/libkineto/src/Config.cpp b/libkineto/src/Config.cpp index 337a65eb3..346f910ec 100644 --- a/libkineto/src/Config.cpp +++ b/libkineto/src/Config.cpp @@ -134,6 +134,11 @@ constexpr char kProfileStartIterationKey[] = "PROFILE_START_ITERATION"; constexpr char kProfileStartIterationRoundUpKey[] = "PROFILE_START_ITERATION_ROUNDUP"; +constexpr char kRequestTraceID[] = "REQUEST_TRACE_ID"; +constexpr char kRequestGroupTraceID[] = + "REQUEST_GROUP_TRACE_ID"; + + // Enable on-demand trigger via kill -USR2 // When triggered in this way, /tmp/libkineto.conf will be used as config. constexpr char kEnableSigUsr2Key[] = "ENABLE_SIGUSR2"; @@ -385,6 +390,10 @@ bool Config::handleOption(const std::string& name, std::string& val) { activitiesWarmupIterations_ = toInt32(val); } else if (!name.compare(kActivitiesDisplayCudaSyncWaitEvents)) { activitiesCudaSyncWaitEvents_ = toBool(val); + } else if (!name.compare(kRequestTraceID)) { + requestTraceID_ = val; + } else if (!name.compare(kRequestGroupTraceID)) { + requestGroupTraceID_ = val; } // TODO: Deprecate Client Interface diff --git a/libkineto/test/ConfigTest.cpp b/libkineto/test/ConfigTest.cpp index ba3be5709..a99e428ff 100644 --- a/libkineto/test/ConfigTest.cpp +++ b/libkineto/test/ConfigTest.cpp @@ -348,3 +348,11 @@ TEST(ParseTest, ProfileStartTime) { .count(); EXPECT_FALSE(cfg.parse(fmt::format("PROFILE_START_TIME = {}", tbad_ms))); } + +TEST(ParseTest, RequestTraceIds) { + Config cfg; + EXPECT_TRUE(cfg.parse("REQUEST_TRACE_ID=XYZ")); + EXPECT_EQ(cfg.requestTraceID(), "XYZ"); + EXPECT_TRUE(cfg.parse("REQUEST_GROUP_TRACE_ID=ABC")); + EXPECT_EQ(cfg.requestGroupTraceID(), "ABC"); +}