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

Fix AllocatorConfig potential SIO issue (#159629)

Browse Source 
# Motivation
As @ScottTodd identified in this [comment](https://github.com/pytorch/pytorch/pull/150312#issuecomment-3141524874), using STL containers like `std::string` and `std::unordered_set` at static init time can cause static initialization order issues. This PR is based on and modified from his original PR: https://github.com/pytorch/pytorch/pull/159607. I’m stacking this PR here to help facilitate the landing and validation process.

Co-authored-by: @ScottTodd
Pull Request resolved: https://github.com/pytorch/pytorch/pull/159629
Approved by: https://github.com/ScottTodd, https://github.com/albanD
This commit is contained in:
Yu, Guangye
2025-08-01 15:17:28 +00:00
committed by PyTorch MergeBot
parent b6c53383fe
commit 56d19a5ced
2 changed files with 17 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
c10/core/AllocatorConfig.cpp
View File

@ -224,7 +224,7 @@ void AcceleratorAllocatorConfig::parseArgs(const std::string& env) {
// check if the key is unrecognized.
if (device_config_parser_hook_) {
TORCH_CHECK(
keys_.find(key) != keys_.end(),
getKeys().find(key) != getKeys().end(),
"Unrecognized key '",
key,
"' in Accelerator allocator config.");

29
c10/core/AllocatorConfig.h
View File

@ -220,11 +220,24 @@ class C10_API AcceleratorAllocatorConfig {
return instance().last_allocator_settings_;
}
// Use `Construct On First Use Idiom` to avoid `Static Initialization Order`
// issue.
static std::unordered_set<std::string>& getMutableKeys() {
static std::unordered_set<std::string> keys{
"max_split_size_mb",
"max_non_split_rounding_mb",
"garbage_collection_threshold",
"roundup_power2_divisions",
"expandable_segments",
"pinned_use_background_threads"};
return keys;
}
// Returns the set of valid keys for the allocator configuration.
// This set is used to validate the presence and correctness of keys in
// device-specific configuration parsers.
static const std::unordered_set<std::string>& getKeys() {
return keys_;
return getMutableKeys();
}
// Registers a device-specific configuration parser hook and its key. This
@ -238,9 +251,10 @@ class C10_API AcceleratorAllocatorConfig {
std::function<void(const std::string&)>&& hook,
const std::unordered_set<std::string>& keys) {
device_config_parser_hook_ = std::move(hook);
auto& mutable_keys = getMutableKeys();
for (auto& key : keys) {
TORCH_CHECK(
keys_.insert(key).second,
mutable_keys.insert(key).second,
"Duplicated key '",
key,
"' found in device-specific configuration parser hook registration");
@ -326,17 +340,6 @@ class C10_API AcceleratorAllocatorConfig {
// their own environment configuration extensions.
inline static std::function<void(const std::string&)>
device_config_parser_hook_{nullptr};
// A set of valid configuration keys, including both common and
// device-specific options. This set is used to validate the presence and
// correctness of keys during parsing.
inline static std::unordered_set<std::string> keys_{
"max_split_size_mb",
"max_non_split_rounding_mb",
"garbage_collection_threshold",
"roundup_power2_divisions",
"expandable_segments",
"pinned_use_background_threads"};
};
C10_API inline void setAllocatorSettings(const std::string& env) {