1+ import threading
2+ from typing import Dict , Optional
3+ from memos .configs .llm import HFLLMConfig
4+ from memos .llms .hf import HFLLM
5+ from memos .log import get_logger
6+
7+ logger = get_logger (__name__ )
8+
9+
10+ class HFSingletonLLM (HFLLM ):
11+ """
12+ Singleton version of HFLLM that prevents multiple loading of the same model.
13+ This class inherits from HFLLM and adds singleton behavior.
14+ """
15+
16+ _instances : Dict [str , 'HFSingletonLLM' ] = {}
17+ _lock = threading .Lock ()
18+
19+ def __new__ (cls , config : HFLLMConfig ):
20+ """
21+ Singleton pattern implementation.
22+ Returns existing instance if config already exists, otherwise creates new one.
23+ """
24+ config_key = cls ._get_config_key (config )
25+
26+ if config_key in cls ._instances :
27+ logger .debug (f"Reusing existing HF model: { config .model_name_or_path } )
28+ return cls ._instances [config_key ]
29+
30+ with cls ._lock :
31+ # Double-check pattern to prevent race conditions
32+ if config_key in cls ._instances :
33+ logger .debug (f"Reusing existing HF model: { config .model_name_or_path } )
34+ return cls ._instances [config_key ]
35+
36+ logger .info (f"Creating new HF model: { config .model_name_or_path } )
37+ instance = super ().__new__ (cls )
38+ cls ._instances [config_key ] = instance
39+ return instance
40+
41+ def __init__ (self , config : HFLLMConfig ):
42+ """
43+ Initialize the singleton HFLLM instance.
44+ Only initializes if this is a new instance.
45+ """
46+ # Check if already initialized
47+ if hasattr (self , '_initialized' ):
48+ return
49+
50+ # Call parent constructor
51+ super ().__init__ (config )
52+ self ._initialized = True
53+
54+ @classmethod
55+ def _get_config_key (cls , config : HFLLMConfig ) -> str :
56+ """
57+ Generate a unique key for the HF model configuration.
58+
59+ Args:
60+ config: The HFLLM configuration
61+
62+ Returns:
63+ A unique string key representing the configuration
64+ """
65+ # Create a unique key based on model path and key parameters
66+ # str(config.temperature),
67+ # str(config.max_tokens),
68+ # str(config.top_p),
69+ # str(config.top_k),
70+ # str(config.add_generation_prompt),
71+ # str(config.remove_think_prefix),
72+ # str(config.do_sample)
73+ key_parts = [
74+ config .model_name_or_path
75+ ]
76+ return "|" .join (key_parts )
77+
78+ @classmethod
79+ def get_instance_count (cls ) -> int :
80+ """
81+ Get the number of unique HF model instances currently managed.
82+
83+ Returns:
84+ Number of HF model instances
85+ """
86+ return len (cls ._instances )
87+
88+ @classmethod
89+ def get_instance_info (cls ) -> Dict [str , str ]:
90+ """
91+ Get information about all managed HF model instances.
92+
93+ Returns:
94+ Dictionary mapping config keys to model paths
95+ """
96+ return {key : instance .config .model_name_or_path
97+ for key , instance in cls ._instances .items ()}
98+
99+ @classmethod
100+ def clear_all (cls ) -> None :
101+ """
102+ Clear all HF model instances from memory.
103+ This should be used carefully as it will force reloading of models.
104+ """
105+ with cls ._lock :
106+ cls ._instances .clear ()
107+ logger .info ("All HF model instances cleared from singleton manager" )
108+
109+
110+ # Convenience function to get singleton manager info
111+ def get_hf_singleton_info () -> Dict [str , int ]:
112+ """
113+ Get information about the HF singleton manager.
114+
115+ Returns:
116+ Dictionary with instance count and info
117+ """
118+ return {
119+ "instance_count" : HFSingletonLLM .get_instance_count (),
120+ "instance_info" : HFSingletonLLM .get_instance_info ()
121+ }
0 commit comments