Skip to content

Commit 47a78da

Browse files
cyyeverpytorchmergebot
cyyever
authored and
pytorchmergebot
committed
[Environment Variable][1/N] Use thread-safe env variable API in c10 (#119449)
This PR is the beginning of attempts to wrap thread-unsafe getenv and set_env functions inside a RW mutex. Pull Request resolved: #119449 Approved by: https://github.com/malfet, https://github.com/albanD, https://github.com/eqy
1 parent be169f7 commit 47a78da

File tree

6 files changed

+126
-42
lines changed

6 files changed

+126
-42
lines changed

‎c10/test/util/DeadlockDetection_test.cpp‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,8 @@
11
#include <c10/util/DeadlockDetection.h>
2+
#include <c10/util/env.h>
23

34
#include <gtest/gtest.h>
45

5-
#include <cstdlib>
6-
76
using namespace ::testing;
87
using namespace c10::impl;
98

@@ -23,7 +22,7 @@ TEST(DeadlockDetection, basic) {
2322

2423
#ifndef _WIN32
2524
TEST(DeadlockDetection, disable) {
26-
setenv("TORCH_DISABLE_DEADLOCK_DETECTION", "1", 1);
25+
c10::utils::set_env("TORCH_DISABLE_DEADLOCK_DETECTION", "1");
2726
DummyPythonGILHooks hooks;
2827
SetPythonGILHooks(&hooks);
2928
SetPythonGILHooks(&hooks);

‎c10/util/DeadlockDetection.cpp‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,13 @@
11
#include <c10/util/DeadlockDetection.h>
2-
3-
#include <cstdlib>
2+
#include <c10/util/env.h>
43

54
namespace c10::impl {
65

76
namespace {
87
PythonGILHooks* python_gil_hooks = nullptr;
98

109
bool disable_detection() {
11-
return std::getenv("TORCH_DISABLE_DEADLOCK_DETECTION") != nullptr;
10+
return c10::utils::has_env("TORCH_DISABLE_DEADLOCK_DETECTION");
1211
}
1312
} // namespace
1413

‎c10/util/Logging.cpp‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
#include <c10/util/Flags.h>
44
#include <c10/util/Lazy.h>
55
#include <c10/util/Logging.h>
6+
#include <c10/util/env.h>
67
#ifdef FBCODE_CAFFE2
78
#include <folly/synchronization/SanitizeThread.h>
89
#endif
@@ -12,7 +13,6 @@
1213
#endif
1314

1415
#include <algorithm>
15-
#include <cstdlib>
1616
#include <iostream>
1717

1818
// Common code that we use regardless of whether we use glog or not.
@@ -122,8 +122,8 @@ using DDPUsageLoggerType = std::function<void(const DDPLoggingData&)>;
122122

123123
namespace {
124124
bool IsAPIUsageDebugMode() {
125-
const char* val = getenv("PYTORCH_API_USAGE_STDERR");
126-
return val && *val; // any non-empty value
125+
auto val = c10::utils::get_env("PYTORCH_API_USAGE_STDERR");
126+
return val.has_value() && !val.value().empty(); // any non-empty value
127127
}
128128

129129
void APIUsageDebug(const string& event) {
@@ -504,10 +504,10 @@ namespace c10::detail {
504504
namespace {
505505

506506
void setLogLevelFlagFromEnv() {
507-
const char* level_str = std::getenv("TORCH_CPP_LOG_LEVEL");
507+
auto level_env = c10::utils::get_env("TORCH_CPP_LOG_LEVEL");
508508

509509
// Not set, fallback to the default level (i.e. WARNING).
510-
std::string level{level_str != nullptr ? level_str : ""};
510+
std::string level{level_env.has_value() ? level_env.value() : ""};
511511
if (level.empty()) {
512512
return;
513513
}

‎c10/util/env.cpp‎

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,94 @@
1+
#include <c10/util/Exception.h>
2+
#include <c10/util/env.h>
3+
#include <fmt/format.h>
4+
#include <cstdlib>
5+
#include <shared_mutex>
6+
7+
namespace c10::utils {
8+
9+
static std::shared_mutex env_mutex;
10+
11+
// Set an environment variable.
12+
void set_env(const char* name, const char* value, bool overwrite) {
13+
std::lock_guard lk(env_mutex);
14+
#ifdef _MSC_VER
15+
#pragma warning(push)
16+
#pragma warning(disable : 4996)
17+
if (!overwrite) {
18+
// NOLINTNEXTLINE(concurrency-mt-unsafe)
19+
if (std::getenv(name) != nullptr) {
20+
return;
21+
}
22+
}
23+
auto full_env_variable = fmt::format("{}={}", name, value);
24+
// NOLINTNEXTLINE(concurrency-mt-unsafe)
25+
auto err = putenv(full_env_variable.c_str());
26+
TORCH_INTERNAL_ASSERT(
27+
err == 0,
28+
"putenv failed for environment \"",
29+
name,
30+
"\", the error is: ",
31+
err);
32+
#pragma warning(pop)
33+
#else
34+
// NOLINTNEXTLINE(concurrency-mt-unsafe)
35+
auto err = setenv(name, value, static_cast<int>(overwrite));
36+
TORCH_INTERNAL_ASSERT(
37+
err == 0,
38+
"setenv failed for environment \"",
39+
name,
40+
"\", the error is: ",
41+
err);
42+
#endif
43+
return;
44+
}
45+
46+
// Reads an environment variable and returns the content if it is set
47+
std::optional<std::string> get_env(const char* name) noexcept {
48+
std::shared_lock lk(env_mutex);
49+
#ifdef _MSC_VER
50+
#pragma warning(push)
51+
#pragma warning(disable : 4996)
52+
#endif
53+
// NOLINTNEXTLINE(concurrency-mt-unsafe)
54+
auto envar = std::getenv(name);
55+
#ifdef _MSC_VER
56+
#pragma warning(pop)
57+
#endif
58+
if (envar != nullptr) {
59+
return std::string(envar);
60+
}
61+
return std::nullopt;
62+
}
63+
64+
// Checks an environment variable is set.
65+
bool has_env(const char* name) noexcept {
66+
return get_env(name).has_value();
67+
}
68+
69+
// Reads an environment variable and returns
70+
// - optional<true>, if set equal to "1"
71+
// - optional<false>, if set equal to "0"
72+
// - nullopt, otherwise
73+
//
74+
// NB:
75+
// Issues a warning if the value of the environment variable is not 0 or 1.
76+
std::optional<bool> check_env(const char* name) {
77+
auto env_opt = get_env(name);
78+
if (env_opt.has_value()) {
79+
if (*env_opt == "0") {
80+
return false;
81+
}
82+
if (*env_opt == "1") {
83+
return true;
84+
}
85+
TORCH_WARN(
86+
"Ignoring invalid value for boolean flag ",
87+
name,
88+
": ",
89+
*env_opt,
90+
"valid values are 0 or 1.");
91+
}
92+
return std::nullopt;
93+
}
94+
} // namespace c10::utils

‎c10/util/env.h‎

Lines changed: 18 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -1,41 +1,31 @@
11
#pragma once
22

3-
#include <c10/util/Exception.h>
4-
#include <cstdlib>
5-
#include <cstring>
3+
#include <c10/macros/Export.h>
64
#include <optional>
5+
#include <string>
76

87
namespace c10::utils {
8+
9+
// Set an environment variable.
10+
C10_API void set_env(
11+
const char* name,
12+
const char* value,
13+
bool overwrite = true);
14+
15+
// Checks an environment variable is set.
16+
C10_API bool has_env(const char* name) noexcept;
17+
918
// Reads an environment variable and returns
1019
// - std::optional<true>, if set equal to "1"
1120
// - std::optional<false>, if set equal to "0"
1221
// - nullopt, otherwise
1322
//
1423
// NB:
1524
// Issues a warning if the value of the environment variable is not 0 or 1.
16-
inline std::optional<bool> check_env(const char* name) {
17-
#ifdef _MSC_VER
18-
#pragma warning(push)
19-
#pragma warning(disable : 4996)
20-
#endif
21-
auto envar = std::getenv(name);
22-
#ifdef _MSC_VER
23-
#pragma warning(pop)
24-
#endif
25-
if (envar) {
26-
if (strcmp(envar, "0") == 0) {
27-
return false;
28-
}
29-
if (strcmp(envar, "1") == 0) {
30-
return true;
31-
}
32-
TORCH_WARN(
33-
"Ignoring invalid value for boolean flag ",
34-
name,
35-
": ",
36-
envar,
37-
"valid values are 0 or 1.");
38-
}
39-
return std::nullopt;
40-
}
25+
C10_API std::optional<bool> check_env(const char* name);
26+
27+
// Reads the value of an environment variable if it is set.
28+
// However, check_env should be used if the value is assumed to be a flag.
29+
C10_API std::optional<std::string> get_env(const char* name) noexcept;
30+
4131
} // namespace c10::utils

‎c10/util/tempfile.cpp‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
#include <c10/util/Exception.h>
2+
#include <c10/util/env.h>
23
#include <c10/util/tempfile.h>
34
#include <fmt/format.h>
45

@@ -22,10 +23,11 @@ static std::string make_filename(std::string_view name_prefix) {
2223
// We see if any of these environment variables is set and use their value, or
2324
// else default the temporary directory to `/tmp`.
2425

25-
const char* tmp_directory = "/tmp";
26+
std::string tmp_directory = "/tmp";
2627
for (const char* variable : {"TMPDIR", "TMP", "TEMP", "TEMPDIR"}) {
27-
if (const char* path = getenv(variable)) {
28-
tmp_directory = path;
28+
auto path_opt = c10::utils::get_env(variable);
29+
if (path_opt.has_value()) {
30+
tmp_directory = path_opt.value();
2931
break;
3032
}
3133
}

0 commit comments

Comments
 (0)