mirror of https://github.com/hpcaitech/ColossalAI
[Fix]: implement thread-safety singleton to avoid deadlock for very large-scale training scenarios (#5625)
* implement thread-safety singleton * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * refactor singleton implementation * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>pull/5654/head
parent
bbb2c21f16
commit
7ef91606e1
|
@ -1,22 +1,27 @@
|
||||||
|
import threading
|
||||||
|
|
||||||
|
|
||||||
class SingletonMeta(type):
|
class SingletonMeta(type):
|
||||||
"""
|
"""
|
||||||
The Singleton class can be implemented in different ways in Python. Some
|
Thread-safe Singleton Meta with double-checked locking.
|
||||||
possible methods include: base class, decorator, metaclass. We will use the
|
Reference: https://en.wikipedia.org/wiki/Double-checked_locking
|
||||||
metaclass because it is best suited for this purpose.
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
_instances = {}
|
_instances = {}
|
||||||
|
_lock = threading.Lock()
|
||||||
|
|
||||||
def __call__(cls, *args, **kwargs):
|
def __call__(cls, *args, **kwargs):
|
||||||
"""
|
# First check (without locking) for performance reasons
|
||||||
Possible changes to the value of the `__init__` argument do not affect
|
if cls not in cls._instances:
|
||||||
the returned instance.
|
# Acquire a lock before proceeding to the second check
|
||||||
"""
|
with cls._lock:
|
||||||
|
# Second check with lock held to ensure thread safety
|
||||||
if cls not in cls._instances:
|
if cls not in cls._instances:
|
||||||
instance = super().__call__(*args, **kwargs)
|
instance = super().__call__(*args, **kwargs)
|
||||||
cls._instances[cls] = instance
|
cls._instances[cls] = instance
|
||||||
else:
|
else:
|
||||||
assert (
|
assert (
|
||||||
len(args) == 0 and len(kwargs) == 0
|
len(args) == 0 and len(kwargs) == 0
|
||||||
), f"{cls.__name__} is a singleton class and a instance has been created."
|
), f"{cls.__name__} is a singleton class and an instance has been created."
|
||||||
|
|
||||||
return cls._instances[cls]
|
return cls._instances[cls]
|
||||||
|
|
Loading…
Reference in New Issue