8000 Use thread-safe getenv wrapper · pytorch/pytorch@d749e3a · GitHub
[go: up one dir, main page]
Skip to content

Commit d749e3a

Browse files
cyyevercyyever
committed
Use thread-safe getenv wrapper
1 parent d4abf0a commit d749e3a

File tree

3 files changed

+17
-15
lines changed

3 files changed

+17
-15
lines changed

‎aten/src/ATen/ParallelCommon.cpp

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
#include <ATen/Config.h>
44
#include <ATen/PTThreadPool.h>
55
#include <ATen/Version.h>
6+
#include <c10/util/env.h>
67

78
#include <sstream>
89
#include <thread>
@@ -23,17 +24,17 @@ namespace at {
2324

2425
namespace {
2526

26-
const char* get_env_var(
27+
std::string get_env_var(
2728
const char* var_name, const char* def_value = nullptr) {
28-
const char* value = std::getenv(var_name);
29-
return value ? value : def_value;
29+
auto env = c10::utils::get_env(var_name);
30+
return env.has_value() ? env.value() : def_value;
3031
}
3132

3233
#ifndef C10_MOBILE
3334
size_t get_env_num_threads(const char* var_name, size_t def_value = 0) {
3435
try {
35-
if (auto* value = std::getenv(var_name)) {
36-
int nthreads = std::stoi(value);
36+
if (auto value = c10::utils::get_env(var_name)) {
37+
int nthreads = std::stoi(value.value());
3738
TORCH_CHECK(nthreads > 0);
3839
return nthreads;
3940
}

‎aten/src/ATen/cuda/CUDABlas.cpp

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
#include <c10/cuda/CUDACachingAllocator.h>
1212
#include <c10/cuda/CUDAFunctions.h>
1313
#include <c10/macros/Export.h>
14+
#include <c10/util/env.h>
1415
#include <c10/util/irange.h>
1516

1617
#ifdef USE_ROCM
@@ -180,17 +181,17 @@ uint32_t _getAlignment(uintptr_t address) {
180181
#endif
181182

182183
static size_t _parseChosenWorkspaceSize() {
183-
const char * val = getenv("CUBLASLT_WORKSPACE_SIZE");
184+
auto< 8000 /span> val = c10::utils::get_env("CUBLASLT_WORKSPACE_SIZE");
184185
#ifdef USE_ROCM
185-
if (!val) {
186+
if (!val.has_value()) {
186187
// accept either env var
187-
val = getenv("HIPBLASLT_WORKSPACE_SIZE");
188+
val = c10::utils::get_env("HIPBLASLT_WORKSPACE_SIZE");
188189
}
189190
#endif
190191
size_t workspace_size = 1024; /* default size in KiB according to #73328 */
191-
if (val) {
192+
if (val.has_value()) {
192193
try {
193-
workspace_size = std::stoi(val);
194+
workspace_size = std::stoi(val.value());
194195
} catch(std::invalid_argument const& e) {
195196
TORCH_WARN("invalid CUBLASLT_WORKSPACE_SIZE,",
196197
" using default workspace size of ", workspace_size, " KiB.");

‎torch/csrc/distributed/c10d/ProcessGroupUCC.cpp

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
#ifdef USE_C10D_UCC
22

33
#include <ATen/cuda/nvrtc_stub/ATenNVRTC.h>
4+
#include <c10/util/env.h>
45
#include <torch/csrc/distributed/c10d/ProcessGroupUCC.hpp>
56
#include <torch/csrc/distributed/c10d/UCCTracing.hpp>
67
#include <torch/csrc/distributed/c10d/UCCUtils.hpp>
@@ -157,11 +158,10 @@ void read_config() {
157158
torch_ucc_config.enable_comms_logger = false;
158159

159160
// read all torch_ucc env. variables and update the map
160-
char* env;
161-
for (auto& torch_ucc_env : torch_ucc_envs_map) {
162-
env = std::getenv(torch_ucc_env.first.c_str());
163-
if (env) {
164-
torch_ucc_envs_map[torch_ucc_env.first] = std::string(env);
161+
for (auto& [env_name, value] : torch_ucc_envs_map) {
162+
auto env = c10::utils::get_env(env_name.c_str());
163+
if (env.has_value()) {
164+
value = std::move(env.value());
165165
}
166166
}
167167

0 commit comments

Comments
 (0)
0