deploy_pickup/predictor/multi_model_predictor.py

# -*- coding: utf-8 -*-
"""
multi_model_predictor.py - 多设备多模型预测器
=============================================

支持每个设备（device_code）加载独立的模型目录。
支持模型热加载（检测文件变化后自动重载）。

使用示例:
    predictor = MultiModelPredictor()
    predictor.register_device("LT-2", "LT-2")
    predictor.register_device("LT-5", "LT-5")
"""

import os
import time
import logging
from pathlib import Path
from typing import Dict, Optional, Tuple
import numpy as np
import torch

from .config import CFG, DeployConfig
from .model_def import ConvAutoencoder
from .utils import get_device

# --- BM1684X NPU 推理适配（预埋，暂未启用） ---
# 取消以下注释以启用 BM1684X NPU 推理（需先用 convert_to_bmodel.py 生成 .bmodel）
from .bm1684x_engine import BM1684XEngine, is_bm1684x_available

logger = logging.getLogger('MultiModelPredictor')


class DevicePredictor:
    """
    单设备预测器

    封装一个设备的模型、Min-Max 标准化参数和阈值。
    """

    def __init__(self, device_code: str, model_subdir: str):
        self.device_code = device_code
        self.model_subdir = model_subdir

        # 模型目录路径
        self.model_dir = CFG.MODEL_ROOT / model_subdir
        self.model_path = self.model_dir / "ae_model.pth"
        self.scale_path = self.model_dir / "global_scale.npy"
        self.threshold_dir = self.model_dir / "thresholds"

        # 加载资源
        self.torch_device = get_device()
        self.model = self._load_model()
        # Min-Max 参数 (min, max)
        self.global_min, self.global_max = self._load_scale()
        # 阈值（标量）
        self.threshold = self._load_threshold()

        # --- BM1684X NPU 推理适配（预埋，暂未启用） ---
        # 如果 .bmodel 文件存在且 BM1684X 硬件可用，优先使用 NPU 推理
        # 启用后 self.bm_engine 不为 None，推理时调用 bm_engine.predict() 而非 self.model()
        self.bm_engine = None
        self.bm_engine = self._load_bmodel()

        # 记录文件 mtime（用于热加载检测）
        self._model_mtime = self._get_mtime(self.model_path)
        self._scale_mtime = self._get_mtime(self.scale_path)

        engine_type = "BM1684X BModel" if self.bm_engine else "PyTorch (.pth)"
        logger.info(f"设备 {device_code} 模型加载完成 | 引擎: {engine_type} | 目录: {model_subdir} | "
                    f"阈值: {self.threshold:.6f}")

    def _get_mtime(self, path: Path) -> float:
        # 获取文件修改时间，不存在返回 0
        try:
            return os.path.getmtime(path)
        except OSError:
            return 0.0
    
    def has_files_changed(self) -> bool:
        # 检查模型或标准化参数文件是否有更新
        new_model_mtime = self._get_mtime(self.model_path)
        new_scale_mtime = self._get_mtime(self.scale_path)
        return (new_model_mtime != self._model_mtime or
                new_scale_mtime != self._scale_mtime)

    def _load_model(self) -> ConvAutoencoder:
        if not self.model_path.exists():
            raise FileNotFoundError(f"模型不存在: {self.model_path}")

        model = ConvAutoencoder().to(self.torch_device)
        state = torch.load(self.model_path, map_location=self.torch_device)
        model.load_state_dict(state)
        model.eval()
        return model

    # --- BM1684X NPU 推理适配（预埋，暂未启用） ---
    # 取消注释以启用 BModel 加载逻辑
    def _load_bmodel(self):
        """尝试加载 BModel，成功返回 BM1684XEngine 实例，否则返回 None"""
        bmodel_path = self.model_dir / "ae_model.bmodel"
        if not bmodel_path.exists():
            return None
        if not is_bm1684x_available():
            logger.info(f"设备 {self.device_code}: 存在 .bmodel 但 BM1684X 不可用，回退 PyTorch")
            return None
        try:
            engine = BM1684XEngine(str(bmodel_path))
            logger.info(f"设备 {self.device_code}: 已加载 BM1684X BModel 推理引擎")
            return engine
        except Exception as e:
            logger.warning(f"设备 {self.device_code}: BModel 加载失败，回退 PyTorch | {e}")
            return None

    def _load_scale(self) -> Tuple[float, float]:
        # 加载 Min-Max 标准化参数 [min, max]
        if not self.scale_path.exists():
            raise FileNotFoundError(f"标准化参数不存在: {self.scale_path}")

        scale = np.load(self.scale_path)
        return float(scale[0]), float(scale[1])

    def _load_threshold(self) -> float:
        """
        加载阈值文件

        Returns:
            overall_threshold
        """
        # 按 device_code 查找
        threshold_file = self.threshold_dir / f"threshold_{self.device_code}.npy"
        if not threshold_file.exists():
            # 尝试 default
            threshold_file = self.threshold_dir / "threshold_default.npy"

        if threshold_file.exists():
            data = np.load(threshold_file)
            # 兼容标量和数组格式
            return float(data.flat[0])

        logger.warning(f"设备 {self.device_code} 无阈值文件，使用默认值 0.01")
        return 0.01


class MultiModelPredictor:
    """
    多设备多模型预测器

    支持：
    - 每设备独立模型、标准化参数和阈值
    - 模型热加载（检测文件更新后自动重载）
    - 冷启动设备定期重试
    """

    # 热加载检查间隔（秒）
    HOT_RELOAD_INTERVAL = 60
    # 失败设备重试间隔（秒）
    FAILED_RETRY_INTERVAL = 300

    def __init__(self):
        # device_code -> model_subdir 映射
        self.device_model_map: Dict[str, str] = {}

        # device_code -> DevicePredictor 实例（懒加载）
        self.predictors: Dict[str, DevicePredictor] = {}

        # 加载失败的设备记录 {device_code: 失败时间}
        self._failed_devices: Dict[str, float] = {}

        # 上次热加载检查时间
        self._last_reload_check: float = 0.0

        logger.info("MultiModelPredictor 初始化完成")

    def register_device(self, device_code: str, model_subdir: str):
        self.device_model_map[device_code] = model_subdir
        logger.debug(f"注册设备: {device_code} -> models/{model_subdir}/")

    def _check_hot_reload(self):
        """
        检查所有已加载设备的模型文件是否有更新

        检查频率由 HOT_RELOAD_INTERVAL 控制（默认60秒），避免频繁 stat 调用。
        如果检测到文件变化，重新创建 DevicePredictor 实例替换旧的。
        同时对失败设备定期重试加载。
        """
        now = time.time()
        if now - self._last_reload_check < self.HOT_RELOAD_INTERVAL:
            return
        self._last_reload_check = now

        # 1. 检查已加载设备的模型是否更新
        for device_code in list(self.predictors.keys()):
            predictor = self.predictors[device_code]
            if predictor.has_files_changed():
                logger.info(f"检测到模型文件更新: {device_code}，执行热加载")
                model_subdir = self.device_model_map.get(device_code, device_code)
                try:
                    new_predictor = DevicePredictor(device_code, model_subdir)
                    self.predictors[device_code] = new_predictor
                    logger.info(f"热加载成功: {device_code}")
                except Exception as e:
                    logger.error(f"热加载失败: {device_code} | {e}")
                    # 保留旧的 predictor 继续使用

        # 2. 对失败设备定期重试
        for device_code in list(self._failed_devices.keys()):
            fail_time = self._failed_devices[device_code]
            if now - fail_time > self.FAILED_RETRY_INTERVAL:
                model_subdir = self.device_model_map.get(device_code, device_code)
                try:
                    predictor = DevicePredictor(device_code, model_subdir)
                    self.predictors[device_code] = predictor
                    del self._failed_devices[device_code]
                    logger.info(f"失败设备重试成功: {device_code}")
                except Exception:
                    # 重试仍然失败，更新时间戳
                    self._failed_devices[device_code] = now

    def get_predictor(self, device_code: str) -> Optional[DevicePredictor]:
        # 先执行热加载检查
        self._check_hot_reload()

        # 已加载
        if device_code in self.predictors:
            return self.predictors[device_code]

        # 已失败且未到重试时间
        if device_code in self._failed_devices:
            return None

        # 首次懒加载
        model_subdir = self.device_model_map.get(device_code)
        if not model_subdir:
            model_subdir = device_code

        try:
            predictor = DevicePredictor(device_code, model_subdir)
            self.predictors[device_code] = predictor
            return predictor
        except Exception as e:
            logger.error(f"加载设备 {device_code} 模型失败: {e}")
            self._failed_devices[device_code] = time.time()
            return None

    def get_threshold(self, device_code: str) -> Optional[float]:
        predictor = self.get_predictor(device_code)
        if predictor:
            return predictor.threshold
        return None

    def get_scale(self, device_code: str) -> Tuple[Optional[float], Optional[float]]:
        # 获取 Z-score 参数 (mean, std)
        predictor = self.get_predictor(device_code)
        if predictor:
            return predictor.global_mean, predictor.global_std
        return None, None

    def get_model(self, device_code: str) -> Optional[ConvAutoencoder]:
        predictor = self.get_predictor(device_code)
        if predictor:
            return predictor.model
        return None

    def reload_device(self, device_code: str) -> bool:
        """
        手动触发指定设备的模型重载

        用于外部调用（如 API 接口更新模型后通知重载）

        Returns:
            bool: 是否重载成功
        """
        model_subdir = self.device_model_map.get(device_code, device_code)
        try:
            new_predictor = DevicePredictor(device_code, model_subdir)
            self.predictors[device_code] = new_predictor
            # 清除失败标记
            self._failed_devices.pop(device_code, None)
            logger.info(f"手动重载成功: {device_code}")
            return True
        except Exception as e:
            logger.error(f"手动重载失败: {device_code} | {e}")
            return False

    @property
    def registered_devices(self) -> list:
        return list(self.device_model_map.keys())

    @property
    def loaded_devices(self) -> list:
        return list(self.predictors.keys())