From dfb68bd73db66ac8636903bf0a6fb5a00720aa8f Mon Sep 17 00:00:00 2001
From: Darshan Sanghani <dsang@meta.com>
Date: Wed, 23 Sep 2026 19:52:14 -0700
Subject: [PATCH] Parse DISABLE_CUPTI_LAZY_REINIT by value (#1568)

Summary:
Pull Request resolved: https://github.com/pytorch/kineto/pull/1568

**Problem**

Setting `DISABLE_CUPTI_LAZY_REINIT=0` unexpectedly disables lazy CUPTI reinitialization.

**Why**

The code checks only whether the environment variable exists, so every value is treated as enabled.

**Fix**

Disable lazy reinitialization only when the variable is set to `1`, matching `TEARDOWN_CUPTI` parsing.

Reviewed By: ryanzhang22

Differential Revision: D120578573
---
 libkineto/src/CuptiActivityApi.cpp           |  5 ++---
 libkineto/src/CuptiActivityApi.h             |  7 +++++++
 libkineto/test/CuptiActivityProfilerTest.cpp | 11 +++++++++++
 3 files changed, 20 insertions(+), 3 deletions(-)

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 <atomic>
 #include <condition_variable>
+#include <cstdlib>
 #include <functional>
 #include <memory>
 #include <mutex>
 #include <set>
+#include <string_view>
 
 #include <cupti.h>
 
@@ -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;