bettafish-company/utils/model_manager.py

import os
import time
import threading
import logging
import gc
import torch
import numpy as np
from collections import OrderedDict
from datetime import datetime, timedelta

logger = logging.getLogger('model_manager')
logger.setLevel(logging.INFO)

class ModelManager:
    """
    模型管理器 - 实现模型预加载和按需卸载技术
    
    功能：
    1. 预加载经常使用的模型，减少加载等待时间
    2. 使用LRU (Least Recently Used) 策略管理内存中加载的模型
    3. 支持模型的异步加载和监控
    4. 自动检测并释放长时间未使用的模型内存
    5. 提供模型使用统计
    """
    
    _instance = None
    _lock = threading.Lock()
    
    def __new__(cls):
        with cls._lock:
            if cls._instance is None:
                cls._instance = super(ModelManager, cls).__new__(cls)
            return cls._instance
    
    def __init__(self):
        if hasattr(self, 'initialized'):
            return
            
        # 已加载模型的缓存，使用OrderedDict实现LRU
        self.loaded_models = OrderedDict()
        # 模型使用统计
        self.model_stats = {}
        # 模型预热配置
        self.preload_config = {}
        # 最大内存占用（GB）
        self.max_memory_usage = float(os.getenv('MAX_MODEL_MEMORY_USAGE', '4.0'))
        # 模型加载中的锁
        self.loading_locks = {}
        # 模型卸载超时（分钟）
        self.unload_timeout = int(os.getenv('MODEL_UNLOAD_TIMEOUT', '30'))
        
        # 启动模型监控线程
        self.monitor_thread = threading.Thread(target=self._monitor_models, daemon=True)
        self.monitor_thread.start()
        
        self.initialized = True
        logger.info(f"模型管理器初始化完成，最大内存占用: {self.max_memory_usage}GB")
    
    def register_model(self, model_id, model_path, preload=False, model_size_gb=0.5, 
                      load_function=None, device='cuda' if torch.cuda.is_available() else 'cpu'):
        """
        注册模型，可选设置为预加载
        
        参数:
            model_id: 模型唯一标识符
            model_path: 模型路径
            preload: 是否预加载
            model_size_gb: 模型估计大小（GB）
            load_function: 自定义加载函数，签名为 load_function(model_path, device) -> model
            device: 加载模型的设备
        """
        self.preload_config[model_id] = {
            'model_path': model_path,
            'preload': preload,
            'model_size_gb': model_size_gb,
            'load_function': load_function,
            'device': device
        }
        
        self.model_stats[model_id] = {
            'load_count': 0,
            'use_count': 0,
            'total_load_time': 0,
            'last_used': None,
            'avg_load_time': 0
        }
        
        if preload:
            logger.info(f"模型 {model_id} 已注册并标记为预加载")
            # 启动预加载线程
            threading.Thread(target=self._preload_model, args=(model_id,), daemon=True).start()
        else:
            logger.info(f"模型 {model_id} 已注册")
        
        return True
    
    def get_model(self, model_id):
        """
        获取模型，如果未加载则加载
        
        参数:
            model_id: 模型唯一标识符
            
        返回:
            加载好的模型对象
        """
        if model_id not in self.preload_config:
            raise ValueError(f"模型 {model_id} 未注册")
            
        # 更新最后使用时间
        self.model_stats[model_id]['last_used'] = datetime.now()
        self.model_stats[model_id]['use_count'] += 1
        
        # 检查模型是否已加载
        if model_id in self.loaded_models:
            # 将模型移至OrderedDict末尾，表示最近使用
            model = self.loaded_models.pop(model_id)
            self.loaded_models[model_id] = model
            logger.debug(f"使用已加载的模型: {model_id}")
            return model
            
        # 获取模型加载锁，防止并发加载同一模型
        if model_id not in self.loading_locks:
            self.loading_locks[model_id] = threading.Lock()
            
        # 加锁加载模型
        with self.loading_locks[model_id]:
            # 再次检查模型是否已被其他线程加载
            if model_id in self.loaded_models:
                return self.loaded_models[model_id]
                
            # 检查是否有足够内存
            self._ensure_memory_available(self.preload_config[model_id]['model_size_gb'])
            
            # 加载模型
            start_time = time.time()
            model = self._load_model(model_id)
            load_time = time.time() - start_time
            
            # 更新统计
            self.model_stats[model_id]['load_count'] += 1
            self.model_stats[model_id]['total_load_time'] += load_time
            self.model_stats[model_id]['avg_load_time'] = (
                self.model_stats[model_id]['total_load_time'] / 
                self.model_stats[model_id]['load_count']
            )
            
            logger.info(f"模型 {model_id} 加载完成，耗时: {load_time:.2f}秒")
            
            # 存储模型
            self.loaded_models[model_id] = model
            return model
    
    def unload_model(self, model_id):
        """
        手动卸载模型
        
        参数:
            model_id: 模型唯一标识符
        """
        if model_id in self.loaded_models:
            logger.info(f"手动卸载模型: {model_id}")
            del self.loaded_models[model_id]
            # 强制垃圾回收
            gc.collect()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            return True
        return False
    
    def get_model_stats(self):
        """获取所有模型的使用统计"""
        result = {}
        for model_id, stats in self.model_stats.items():
            is_loaded = model_id in self.loaded_models
            result[model_id] = {
                **stats,
                'is_loaded': is_loaded,
                'preload': self.preload_config[model_id]['preload'],
                'model_size_gb': self.preload_config[model_id]['model_size_gb'],
                'device': self.preload_config[model_id]['device'],
            }
        return result
    
    def preload_all(self):
        """预加载所有标记为预加载的模型"""
        for model_id, config in self.preload_config.items():
            if config['preload'] and model_id not in self.loaded_models:
                threading.Thread(target=self._preload_model, args=(model_id,), daemon=True).start()
    
    def _preload_model(self, model_id):
        """预加载单个模型的内部方法"""
        try:
            logger.info(f"开始预加载模型: {model_id}")
            # 确保有足够内存
            self._ensure_memory_available(self.preload_config[model_id]['model_size_gb'])
            
            # 加载模型
            start_time = time.time()
            model = self._load_model(model_id)
            load_time = time.time() - start_time
            
            # 更新统计
            self.model_stats[model_id]['load_count'] += 1
            self.model_stats[model_id]['total_load_time'] += load_time
            self.model_stats[model_id]['avg_load_time'] = (
                self.model_stats[model_id]['total_load_time'] / 
                self.model_stats[model_id]['load_count']
            )
            
            # 存储模型
            self.loaded_models[model_id] = model
            logger.info(f"模型 {model_id} 预加载完成，耗时: {load_time:.2f}秒")
            
        except Exception as e:
            logger.error(f"预加载模型 {model_id} 失败: {e}")
    
    def _load_model(self, model_id):
        """加载模型的内部方法"""
        config = self.preload_config[model_id]
        
        if config['load_function'] is not None:
            # 使用自定义加载函数
            return config['load_function'](config['model_path'], config['device'])
        
        # 默认加载逻辑 - 根据文件扩展名确定加载方式
        model_path = config['model_path']
        device = config['device']
        
        if model_path.endswith('.pt') or model_path.endswith('.pth'):
            # PyTorch模型
            return torch.load(model_path, map_location=device)
        elif model_path.endswith('.pkl'):
            # Pickle模型
            import pickle
            with open(model_path, 'rb') as f:
                return pickle.load(f)
        else:
            # 尝试作为目录加载
            if os.path.isdir(model_path):
                # 如果是目录，尝试加载预训练模型
                try:
                    from transformers import AutoModel, AutoTokenizer
                    model = AutoModel.from_pretrained(model_path)
                    tokenizer = AutoTokenizer.from_pretrained(model_path)
                    return {'model': model.to(device), 'tokenizer': tokenizer}
                except ImportError:
                    logger.error("transformers库未安装，无法加载预训练模型")
                    raise
                except Exception as e:
                    logger.error(f"加载预训练模型失败: {e}")
                    raise
            
            raise ValueError(f"无法确定如何加载模型: {model_path}")
    
    def _ensure_memory_available(self, required_gb):
        """确保有足够的内存来加载新模型"""
        # 如果当前没有加载的模型，直接返回
        if not self.loaded_models:
            return
            
        # 计算当前已加载模型的总内存
        current_usage = sum(
            self.preload_config[model_id]['model_size_gb'] 
            for model_id in self.loaded_models
        )
        
        # 如果添加新模型后超过限制，需要卸载一些模型
        while current_usage + required_gb > self.max_memory_usage and self.loaded_models:
            # 卸载最久未使用的模型（OrderedDict的首项）
            oldest_model_id, _ = next(iter(self.loaded_models.items()))
            # 检查是否是预加载且最近使用过的模型
            if (self.preload_config[oldest_model_id]['preload'] and
                self.model_stats[oldest_model_id]['last_used'] and
                (datetime.now() - self.model_stats[oldest_model_id]['last_used']) < 
                timedelta(minutes=self.unload_timeout)):
                # 跳过预加载且最近使用过的模型
                # 将该模型移至OrderedDict末尾
                model = self.loaded_models.pop(oldest_model_id)
                self.loaded_models[oldest_model_id] = model
                # 如果所有模型都是预加载的且最近使用过，允许超过限制
                if len(self.loaded_models) <= 1:
                    break
                continue
                
            # 卸载模型并更新内存使用
            model_size = self.preload_config[oldest_model_id]['model_size_gb']
            del self.loaded_models[oldest_model_id]
            current_usage -= model_size
            logger.info(f"自动卸载模型以释放内存: {oldest_model_id} ({model_size}GB)")
            
        # 强制垃圾回收
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            
    def _monitor_models(self):
        """监控并管理模型的内部线程方法"""
        while True:
            try:
                # 检查长时间未使用的非预加载模型
                current_time = datetime.now()
                for model_id in list(self.loaded_models.keys()):
                    if (not self.preload_config[model_id]['preload'] and
                        self.model_stats[model_id]['last_used'] and
                        (current_time - self.model_stats[model_id]['last_used']) > 
                        timedelta(minutes=self.unload_timeout)):
                        # 卸载长时间未使用的非预加载模型
                        logger.info(f"卸载长时间未使用的模型: {model_id}")
                        del self.loaded_models[model_id]
                        # 强制垃圾回收
                        gc.collect()
                        if torch.cuda.is_available():
                            torch.cuda.empty_cache()
                            
                # 每5分钟检查一次
                time.sleep(300)
            except Exception as e:
                logger.error(f"模型监控线程出错: {e}")
                time.sleep(300)

# 创建全局模型管理器实例
model_manager = ModelManager()

# 注册示例函数
def register_sentiment_model():
    """注册情感分析模型示例"""
    from utils.model_loader import load_sentiment_model  # 假设您有一个加载情感模型的函数
    
    try:
        model_path = os.path.join('model', 'sentiment.marshal.3')
        model_manager.register_model(
            model_id='sentiment_basic',
            model_path=model_path,
            preload=True,
            model_size_gb=0.2,
            load_function=load_sentiment_model
        )
        return True
    except Exception as e:
        logger.error(f"注册情感分析模型失败: {e}")
        return False

def register_bert_model():
    """注册BERT模型示例"""
    try:
        model_path = os.path.join('model_pro', 'bert_model')
        model_manager.register_model(
            model_id='bert_classifier',
            model_path=model_path,
            preload=True,
            model_size_gb=0.8
        )
        return True
    except Exception as e:
        logger.error(f"注册BERT模型失败: {e}")
        return False

# 自动注册常用模型（在导入时执行）
try:
    register_sentiment_model()
    register_bert_model()
except Exception as e:
    logger.error(f"自动注册模型失败: {e}")