pytorch · cyyever · Oct 2, 2024 · Oct 11, 2024
diff --git a/aten/src/ATen/ParallelCommon.cpp b/aten/src/ATen/ParallelCommon.cpp
@@ -3,6 +3,7 @@
 #include <ATen/Config.h>
 #include <ATen/PTThreadPool.h>
 #include <ATen/Version.h>
+#include <c10/util/env.h>
 
 #include <sstream>
 #include <thread>
@@ -23,17 +24,17 @@ namespace at {
 
 namespace {
 
-const char* get_env_var(
+std::string get_env_var(
     const char* var_name, const char* def_value = nullptr) {
-  const char* value = std::getenv(var_name);
-  return value ? value : def_value;
+  auto env = c10::utils::get_env(var_name);
+  return env.has_value() ? env.value() : def_value;
 }
 
 #ifndef C10_MOBILE
 size_t get_env_num_threads(const char* var_name, size_t def_value = 0) {
   try {
-    if (auto* value = std::getenv(var_name)) {
-      int nthreads = std::stoi(value);
+    if (auto value = c10::utils::get_env(var_name)) {
+      int nthreads = std::stoi(value.value());
       TORCH_CHECK(nthreads > 0);
       return nthreads;
     }

@@ -11,6 +11,7 @@
 #include <c10/cuda/CUDACachingAllocator.h>
 #include <c10/cuda/CUDAFunctions.h>
 #include <c10/macros/Export.h>
+#include <c10/util/env.h>
 #include <c10/util/irange.h>
 
 #ifdef USE_ROCM
@@ -180,17 +181,17 @@ uint32_t _getAlignment(uintptr_t address) {
 #endif
 
 size_t _parseChosenWorkspaceSize() {
-  const char * val = getenv("CUBLASLT_WORKSPACE_SIZE");
+  auto val = c10::utils::get_env("CUBLASLT_WORKSPACE_SIZE");
 #ifdef USE_ROCM
-  if (!val) {
+  if (!val.has_value()) {
     // accept either env var
-    val = getenv("HIPBLASLT_WORKSPACE_SIZE");
+    val = c10::utils::get_env("HIPBLASLT_WORKSPACE_SIZE");
   }
 #endif
   size_t workspace_size = 1024; /* default size in KiB according to #73328 */
-  if (val) {
+  if (val.has_value()) {
     try {
-      workspace_size = std::stoi(val);
+      workspace_size = std::stoi(val.value());
     } catch(std::invalid_argument const& e) {
       TORCH_WARN("invalid CUBLASLT_WORKSPACE_SIZE,",
                  " using default workspace size of ", workspace_size, " KiB.");

diff --git a/c10/core/impl/alloc_cpu.cpp b/c10/core/impl/alloc_cpu.cpp
@@ -3,6 +3,7 @@
 #include <c10/core/alignment.h>
 #include <c10/util/Flags.h>
 #include <c10/util/Logging.h>
+#include <c10/util/env.h>
 #include <c10/util/irange.h>
 #include <c10/util/numa.h>
 
@@ -53,8 +54,8 @@ void memset_junk(void* data, size_t num) {
 #if defined(__linux__) && !defined(__ANDROID__)
 static inline bool is_thp_alloc_enabled() {
   static bool value = [&] {
-    const char* ptr = std::getenv("THP_MEM_ALLOC_ENABLE");
-    return ptr != nullptr ? std::atoi(ptr) : 0;
+    auto env = c10::utils::check_env("THP_MEM_ALLOC_ENABLE");
+    return env.has_value() ? env.value() : 0;
   }();
   return value;
 }

@@ -9,6 +9,7 @@
 #include <c10/util/Gauge.h>
 #include <c10/util/ScopeExit.h>
 #include <c10/util/UniqueVoidPtr.h>
+#include <c10/util/env.h>
 #include <c10/util/flat_hash_map.h>
 #include <c10/util/hash.h>
 #include <c10/util/llvmMathExtras.h>
@@ -3130,9 +3131,11 @@ class DeviceCachingAllocator {
 static bool forceUncachedAllocator() {
   // Allow either CUDA or HIP name for env var for maximum user comfort
   // the CUDA env var avoids being hipified in cuda_to_hip_mappings.py
-  static const char* cuda_env = getenv("PYTORCH_NO_CUDA_MEMORY_CACHING");
-  static const char* rocm_env = getenv("PYTORCH_NO_HIP_MEMORY_CACHING");
-  static bool force_uncached = (cuda_env != nullptr) || (rocm_env != nullptr);
+  static bool has_cuda_env =
+      c10::utils::has_env("PYTORCH_NO_CUDA_MEMORY_CACHING");
+  static bool has_rocm_env =
+      c10::utils::has_env("PYTORCH_NO_HIP_MEMORY_CACHING");
+  static bool force_uncached = has_cuda_env || has_rocm_env;
   return force_uncached;
 }
 
@@ -3779,9 +3782,9 @@ struct BackendStaticInitializer {
   // version checks, to CUDAAllocatorConfig's runtime doublecheck. If this
   // works, maybe we should move all of CUDAAllocatorConfig here?
   CUDAAllocator* parseEnvForBackend() {
-    const char* val = getenv("PYTORCH_CUDA_ALLOC_CONF");
-    if (val != nullptr) {
-      const std::string config(val);
+    const auto val = c10::utils::get_env("PYTORCH_CUDA_ALLOC_CONF");
+    if (val.has_value()) {
+      const std::string& config = val.value();
 
       std::regex exp("[\\s,]+");
       std::sregex_token_iterator it(config.begin(), config.end(), exp, -1);

@@ -3,6 +3,7 @@
 #include <c10/cuda/CUDAFunctions.h>
 #include <c10/util/Backtrace.h>
 #include <c10/util/Exception.h>
+#include <c10/util/env.h>
 #include <c10/util/irange.h>
 #include <cuda_runtime.h>
 
@@ -80,8 +81,8 @@ bool dsa_check_if_all_devices_support_managed_memory() {
 }
 
 bool env_flag_set(const char* env_var_name) {
-  const char* const env_string = std::getenv(env_var_name);
-  return (env_string == nullptr) ? false : std::strcmp(env_string, "0");
+  const auto env_flag = c10::utils::check_env(env_var_name);
+  return env_flag.has_value() && env_flag.value();
 }
 
 /// Deleter for UVM/managed memory pointers

@@ -1,12 +1,14 @@
 #include <c10/cuda/CUDAMiscFunctions.h>
-#include <cstdlib>
+#include <c10/util/env.h>
 
 namespace c10::cuda {
 
+// NOLINTNEXTLINE(bugprone-exception-escape,-warnings-as-errors)
 const char* get_cuda_check_suffix() noexcept {
-  static char* device_blocking_flag = getenv("CUDA_LAUNCH_BLOCKING");
+  static auto device_blocking_flag =
+      c10::utils::check_env("CUDA_LAUNCH_BLOCKING");
   static bool blocking_enabled =
-      (device_blocking_flag && atoi(device_blocking_flag));
+      (device_blocking_flag.has_value() && device_blocking_flag.value());
   if (blocking_enabled) {
     return "";
   } else {

diff --git a/torch/csrc/distributed/c10d/ProcessGroupUCC.cpp b/torch/csrc/distributed/c10d/ProcessGroupUCC.cpp
@@ -1,6 +1,7 @@
 #ifdef USE_C10D_UCC
 
 #include <ATen/cuda/nvrtc_stub/ATenNVRTC.h>
+#include <c10/util/env.h>
 #include <torch/csrc/distributed/c10d/ProcessGroupUCC.hpp>
 #include <torch/csrc/distributed/c10d/UCCTracing.hpp>
 #include <torch/csrc/distributed/c10d/UCCUtils.hpp>
@@ -157,11 +158,10 @@ void read_config() {
   torch_ucc_config.enable_comms_logger = false;
 
   // read all torch_ucc env. variables and update the map
-  char* env;
  for (auto& torch_ucc_env : torch_ucc_envs_map) {
-    env = std::getenv(torch_ucc_env.first.c_str());
-    if (env) {
-      torch_ucc_envs_map[torch_ucc_env.first] = std::string(env);
+  for (auto& [env_name, value] : torch_ucc_envs_map) {
+    auto env = c10::utils::get_env(env_name.c_str());
+    if (env.has_value()) {
+      value = std::move(env.value());
     }
   }
 

diff --git a/torch/csrc/distributed/c10d/UCCTracing.cpp b/torch/csrc/distributed/c10d/UCCTracing.cpp
@@ -1,5 +1,6 @@
 #ifdef USE_C10D_UCC
 
+#include <c10/util/env.h>
 #include <torch/csrc/distributed/c10d/UCCTracing.hpp>
 #include <torch/csrc/distributed/c10d/UCCUtils.hpp>
 
@@ -32,9 +33,9 @@ void ProcessGroupUCCLogger::flushComms(int rank, int world_size) {
   }
 
   std::string fullpath = "/tmp/" + dirname;
-  char* user_path = std::getenv("TORCH_UCC_COMMS_TRACE_OUTPUT_DIR");
-  if (user_path) {
-    fullpath = user_path;
+  auto user_path = c10::utils::get_env("TORCH_UCC_COMMS_TRACE_OUTPUT_DIR");
+  if (user_path.has_value()) {
+    fullpath = std::move(user_path.value());
   }
   std::string trace_filename = c10::str(fullpath, "/rank", rank, ".json");
   std::ofstream _outfile;

diff --git a/torch/csrc/distributed/c10d/debug.cpp b/torch/csrc/distributed/c10d/debug.cpp
@@ -4,6 +4,7 @@
 // This source code is licensed under the BSD-style license found in the
 // LICENSE file in the root directory of this source tree.
 
+#include <c10/util/env.h>
 #include <torch/csrc/distributed/c10d/debug.h>
 
 #include <algorithm>
@@ -19,15 +20,15 @@ namespace detail {
 namespace {
 
 DebugLevel loadDebugLevelFromEnvironment() {
-  char* env_value = std::getenv("TORCH_DISTRIBUTED_DEBUG");
+  auto env_value = c10::utils::get_env("TORCH_DISTRIBUTED_DEBUG");
 
-  if (env_value == nullptr) {
+  if (!env_value.has_value()) {
     return DebugLevel::Off;
   }
 
   DebugLevel level{};
 
-  std::string level_str{env_value};
+  std::string level_str = std::move(env_value.value());
 
   std::transform(
       level_str.begin(),

@@ -161,14 +161,14 @@ C10_DEFINE_REGISTRY_WITHOUT_WARNING(
 
 const std::string& TensorPipeAgent::guessAddress() {
   static const std::string uvAddress = []() {
-    char* ifnameEnv = std::getenv(kSocketIfnameEnvVar.c_str());
-    if (ifnameEnv != nullptr) {
+    auto ifnameEnv = c10::utils::get_env(kSocketIfnameEnvVar.c_str());
+    if (ifnameEnv.has_value()) {
       auto [error, result] =
-          tensorpipe::transport::uv::lookupAddrForIface(ifnameEnv);
+          tensorpipe::transport::uv::lookupAddrForIface(ifnameEnv.value());
       if (error) {
         LOG(WARNING) << "Failed to look up the IP address for interface "
-                     << ifnameEnv << " (" << error.what() << "), defaulting to "
-                     << kDefaultUvAddress;
+                     << ifnameEnv.value() << " (" << error.what()
+                     << "), defaulting to " << kDefaultUvAddress;
         return kDefaultUvAddress;
       }
       return result;