diff --git a/libkineto/src/CuptiActivityApi.cpp b/libkineto/src/CuptiActivityApi.cpp index 7d998e5b5..c121284d1 100644 --- a/libkineto/src/CuptiActivityApi.cpp +++ b/libkineto/src/CuptiActivityApi.cpp @@ -28,12 +28,11 @@ constexpr size_t kBufSize(4 * 1024 * 1024); constexpr uint32_t kCuptiBufferRejectionMinVersion = 27; inline bool cuptiTearDown_() { - auto teardown_env = getenv("TEARDOWN_CUPTI"); - return teardown_env != nullptr && strcmp(teardown_env, "1") == 0; + return isEnvVarSetToOne("TEARDOWN_CUPTI"); } inline bool cuptiLazyInit_() { - return cuptiTearDown_() && getenv("DISABLE_CUPTI_LAZY_REINIT") == nullptr; + return cuptiTearDown_() && !isEnvVarSetToOne("DISABLE_CUPTI_LAZY_REINIT"); } inline void reenableCuptiCallbacks_(CuptiCallbackApi& cbapi) { diff --git a/libkineto/src/CuptiActivityApi.h b/libkineto/src/CuptiActivityApi.h index 6a59a8a7c..1286c2e89 100644 --- a/libkineto/src/CuptiActivityApi.h +++ b/libkineto/src/CuptiActivityApi.h @@ -10,10 +10,12 @@ #include #include +#include #include #include #include #include +#include #include @@ -27,6 +29,11 @@ namespace KINETO_NAMESPACE { using namespace libkineto; +inline bool isEnvVarSetToOne(const char* name) { + const auto* value = std::getenv(name); + return value != nullptr && std::string_view{value} == "1"; +} + class CuptiActivityApi { public: enum CorrelationFlowType { Default, User }; diff --git a/libkineto/test/CuptiActivityProfilerTest.cpp b/libkineto/test/CuptiActivityProfilerTest.cpp index 4044c5822..ae094ea95 100644 --- a/libkineto/test/CuptiActivityProfilerTest.cpp +++ b/libkineto/test/CuptiActivityProfilerTest.cpp @@ -345,6 +345,17 @@ class MockCuptiActivities : public CuptiActivityApi { bool available{true}; }; +TEST(CuptiActivityApiTest, ParsesBooleanEnvironmentValues) { + constexpr auto* kTestEnvVar = "KINETO_TEST_BOOLEAN_ENV_VAR"; + EXPECT_EQ(unsetenv(kTestEnvVar), 0); + EXPECT_FALSE(isEnvVarSetToOne(kTestEnvVar)); + EXPECT_EQ(setenv(kTestEnvVar, "0", 1), 0); + EXPECT_FALSE(isEnvVarSetToOne(kTestEnvVar)); + EXPECT_EQ(setenv(kTestEnvVar, "1", 1), 0); + EXPECT_TRUE(isEnvVarSetToOne(kTestEnvVar)); + EXPECT_EQ(unsetenv(kTestEnvVar), 0); +} + TEST(CuptiActivityApiTest, KeepsReturningValidBuffersWhenRejectionUnsupported) { CallbackCuptiActivityApi cupti; cupti.canRejectBuffer_ = false;