frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

[Environment Variable][1/N] Use thread-safe env variable API in c10 (#119449)

Browse Source 
This PR is the beginning of attempts to wrap thread-unsafe getenv and set_env functions inside a RW mutex.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/119449
Approved by: https://github.com/albanD
This commit is contained in:
cyy
2024-04-18 13:35:48 +00:00
committed by PyTorch MergeBot
parent 1542874311
commit b51f66c195
12 changed files with 163 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files

46
c10/util/env.h
View File

@ -1,11 +1,20 @@
#pragma once
#include <c10/util/Exception.h>
#include <cstdlib>
#include <cstring>
#include <c10/macros/Export.h>
#include <optional>
#include <string>
namespace c10::utils {
// Set an environment variable.
C10_API void set_env(
const char* name,
const char* value,
bool overwrite = true);
// Checks an environment variable is set.
C10_API bool has_env(const char* name) noexcept;
// Reads an environment variable and returns
// - optional<true>, if set equal to "1"
// - optional<false>, if set equal to "0"
@ -13,29 +22,10 @@ namespace c10::utils {
//
// NB:
// Issues a warning if the value of the environment variable is not 0 or 1.
inline std::optional<bool> check_env(const char* name) {
#ifdef _MSC_VER
#pragma warning(push)
#pragma warning(disable : 4996)
#endif
auto envar = std::getenv(name);
#ifdef _MSC_VER
#pragma warning(pop)
#endif
if (envar) {
if (strcmp(envar, "0") == 0) {
return false;
}
if (strcmp(envar, "1") == 0) {
return true;
}
TORCH_WARN(
"Ignoring invalid value for boolean flag ",
name,
": ",
envar,
"valid values are 0 or 1.");
}
return std::nullopt;
}
C10_API std::optional<bool> check_env(const char* name);
// Reads the value of an environment variable if it is set.
// However, check_env should be used if the value is assumed to be a flag.
C10_API std::optional<std::string> get_env(const char* name) noexcept;
} // namespace c10::utils