You've already forked ModuleLLM-OpenAI-Plugin
mirror of
https://github.com/m5stack/ModuleLLM-OpenAI-Plugin.git
synced 2026-05-20 11:37:26 -07:00
[update] update model load.
This commit is contained in:
+35
-13
@@ -63,16 +63,27 @@ class ModelDispatcher:
|
||||
port=config.data["server"]["port"]
|
||||
)
|
||||
self.lock = asyncio.Lock()
|
||||
self.total_memory = None
|
||||
self.current_used_memory = 0
|
||||
|
||||
async def _ensure_memory_available(self, required_mem: int):
|
||||
if required_mem <= 0:
|
||||
return
|
||||
|
||||
try:
|
||||
cmm_info = await self.memory_checker.get_cmminfo()
|
||||
remain_mem = cmm_info["data"]["remain"]
|
||||
external_remain = cmm_info["data"]["remain"]
|
||||
|
||||
logger.debug(f"Memory Check | Required: {required_mem} | Available: {remain_mem}")
|
||||
if self.total_memory is None:
|
||||
self.total_memory = cmm_info["data"].get("total", external_remain)
|
||||
logger.info(f"Memory Manager Initialized | Total Capacity: {self.total_memory}")
|
||||
|
||||
internal_remain = self.total_memory - self.current_used_memory
|
||||
|
||||
remain_mem = min(internal_remain, external_remain)
|
||||
|
||||
logger.debug(f"Memory Check | Required: {required_mem} | "
|
||||
f"External Remain: {external_remain} | Internal Remain: {internal_remain} | "
|
||||
f"Effective Available: {remain_mem}")
|
||||
|
||||
if remain_mem >= required_mem:
|
||||
return
|
||||
@@ -103,11 +114,17 @@ class ModelDispatcher:
|
||||
for model_name in models_to_unload:
|
||||
logger.info(f"Unloading model '{model_name}' to free memory...")
|
||||
backend = self.backends.pop(model_name)
|
||||
|
||||
model_conf = config.data["models"].get(model_name, {})
|
||||
mem_freed = model_conf.get("memory_required", 0)
|
||||
self.current_used_memory -= mem_freed
|
||||
if self.current_used_memory < 0:
|
||||
self.current_used_memory = 0
|
||||
|
||||
if backend:
|
||||
await backend.close()
|
||||
|
||||
# await asyncio.sleep(0.1)
|
||||
|
||||
except Exception as e:
|
||||
if isinstance(e, HTTPException):
|
||||
raise e
|
||||
@@ -120,29 +137,34 @@ class ModelDispatcher:
|
||||
backend = self.backends.pop(model_name)
|
||||
self.backends[model_name] = backend
|
||||
return backend
|
||||
|
||||
|
||||
model_config = config.data["models"].get(model_name)
|
||||
if model_config is None:
|
||||
return None
|
||||
|
||||
|
||||
required_mem = model_config.get("memory_required", 0)
|
||||
|
||||
await self._ensure_memory_available(required_mem)
|
||||
|
||||
|
||||
logger.info(f"Loading model: {model_name} (Mem Required: {required_mem})")
|
||||
|
||||
backend_instance = None
|
||||
if model_config["type"] == "openai_proxy":
|
||||
self.backends[model_name] = OpenAIProxyBackend(model_config)
|
||||
backend_instance = OpenAIProxyBackend(model_config)
|
||||
elif model_config["type"] in ("llm", "vlm"):
|
||||
self.backends[model_name] = LlmClientBackend(model_config)
|
||||
backend_instance = LlmClientBackend(model_config)
|
||||
elif model_config["type"] == "vision_model":
|
||||
self.backends[model_name] = VisionModelBackend(model_config)
|
||||
backend_instance = VisionModelBackend(model_config)
|
||||
elif model_config["type"] == "tts":
|
||||
self.backends[model_name] = TtsClientBackend(model_config)
|
||||
backend_instance = TtsClientBackend(model_config)
|
||||
elif model_config["type"] == "asr":
|
||||
self.backends[model_name] = ASRClientBackend(model_config)
|
||||
backend_instance = ASRClientBackend(model_config)
|
||||
else:
|
||||
return None
|
||||
|
||||
|
||||
self.backends[model_name] = backend_instance
|
||||
self.current_used_memory += required_mem
|
||||
|
||||
return self.backends.get(model_name)
|
||||
|
||||
async def initialize():
|
||||
|
||||
Reference in New Issue
Block a user