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