Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 13 additions & 3 deletions internlm/utils/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
import inspect
import os
import random
import threading
from abc import ABC, abstractmethod
from contextlib import contextmanager
from datetime import datetime
Expand Down Expand Up @@ -197,18 +198,27 @@ def __call__(self, batch_count):

class SingletonMeta(type):
"""
Singleton Meta.
Thread-safe Singleton Meta with double-checked locking.
Reference: https://en.wikipedia.org/wiki/Double-checked_locking
"""

_instances = {}
_lock = threading.Lock()

def __call__(cls, *args, **kwargs):
# First check (without locking) for performance reasons
if cls not in cls._instances:
cls._instances[cls] = super().__call__(*args, **kwargs)
# 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:
instance = super().__call__(*args, **kwargs)
cls._instances[cls] = instance
else:
assert (
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]


Expand Down