deploy_pickup/run_with_auto_training.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
run_with_auto_training.py
-------------------------
带自动训练的主程序

功能：
1. RTSP拾音器监控（通过 PickupMonitoringSystem）
2. 冷启动检测（无模型时等待数据收集后自动训练）
3. 每日增量训练（默认02:00，可配置）
4. 每日数据清理（默认00:00，可配置）
5. 训练完成后主动通知 MultiModelPredictor 重载，无需等待60秒轮询
"""

import sys
import signal
import logging
import threading
import time
from pathlib import Path
from datetime import datetime

# 确保项目根目录在 sys.path 中
sys.path.insert(0, str(Path(__file__).parent))

try:
    from apscheduler.schedulers.background import BackgroundScheduler
    from apscheduler.triggers.cron import CronTrigger
except ImportError:
    print("错误：缺少依赖库")
    print("请运行：pip install apscheduler")
    sys.exit(1)

from config.config_manager import ConfigManager


def setup_logging():
    # 配置日志系统（按文件大小轮转），与 run_pickup_monitor.py 共用同一日志文件
    from logging.handlers import RotatingFileHandler

    # 如果根 logger 已经被上层调用者配置过，则直接复用
    root = logging.getLogger()
    if root.handlers:
        return

    log_dir = Path(__file__).parent / "logs"
    log_dir.mkdir(parents=True, exist_ok=True)

    log_file = log_dir / "system.log"

    formatter = logging.Formatter(
        '%(asctime)s | %(levelname)-8s | %(name)-20s | %(message)s',
        datefmt='%Y-%m-%d %H:%M:%S'
    )

    # 10MB 单文件上限，保留 2 个备份（总计约 30MB）
    file_handler = RotatingFileHandler(
        log_file,
        maxBytes=10 * 1024 * 1024,
        backupCount=2,
        encoding='utf-8'
    )
    file_handler.setFormatter(formatter)

    console_handler = logging.StreamHandler(sys.stdout)
    console_handler.setFormatter(formatter)

    logging.basicConfig(
        level=logging.INFO,
        handlers=[file_handler, console_handler]
    )


logger = logging.getLogger('MainSystem')


class ColdStartManager:
    """
    冷启动管理器

    检查注册的设备是否都有对应的模型文件，
    缺少模型的设备会被标记为冷启动状态，等待数据收集后自动训练。
    """

    def __init__(self, deploy_root: Path, config: dict):
        self.deploy_root = deploy_root
        self.config = config

        # 模型根目录（每设备子目录）
        self.model_root = deploy_root / "models"
        # 音频数据根目录
        self.audio_root = deploy_root / "data" / "audio"
        # self.audio_root = "/Volumes/mo/水厂正常音频/龙亭"

        # 冷启动配置
        cold_start_cfg = config.get('auto_training', {}).get('cold_start', {})
        self.enabled = cold_start_cfg.get('enabled', True)
        # 等待收集数据的最短时长（小时），避免数据量过少就开始训练
        self.wait_hours = cold_start_cfg.get('wait_hours', 2)
        # 每设备最少样本数
        self.min_samples = cold_start_cfg.get('min_samples', 100)

        # 运行时状态
        self.is_cold_start = False
        self.cold_start_time = None
        # 缺少模型的设备列表
        self.missing_devices = []

    def check_cold_start_needed(self, registered_devices: dict) -> bool:
        """
        检查是否有设备缺少模型

        参数:
            registered_devices: {device_code: model_subdir} 已注册的设备映射

        返回:
            True 表示至少有一个设备缺少模型，需要冷启动
        """
        if not self.enabled:
            return False

        self.missing_devices = []
        for device_code, model_subdir in registered_devices.items():
            model_path = self.model_root / model_subdir / "ae_model.pth"
            if not model_path.exists():
                self.missing_devices.append(device_code)
                logger.warning(f"模型不存在: {device_code} -> {model_path}")

        if self.missing_devices:
            logger.warning(f"以下设备缺少模型，进入冷启动模式: {self.missing_devices}")
            return True

        return False

    def start_cold_start_mode(self):
        # 进入冷启动模式，开始计时
        self.is_cold_start = True
        self.cold_start_time = datetime.now()

        logger.info("=" * 70)
        logger.info("冷启动模式")
        logger.info("=" * 70)
        logger.info(f"缺少模型的设备: {self.missing_devices}")
        logger.info(f"等待时长: {self.wait_hours} 小时")
        logger.info(f"每设备最少样本: {self.min_samples}")
        logger.info("系统将只采集音频，不进行检测（缺模型的设备）")
        logger.info("=" * 70)

    def check_ready_for_training(self) -> bool:
        """
        检查是否已收集到足够的数据可以开始首次训练

        条件：
        1. 已等待超过 wait_hours 小时
        2. 所有缺模型设备的音频文件数 >= min_samples

        返回:
            True 表示数据已就绪
        """
        if not self.is_cold_start:
            return False

        # 条件1：等待时间
        elapsed = (datetime.now() - self.cold_start_time).total_seconds() / 3600
        if elapsed < self.wait_hours:
            return False

        # 条件2：逐设备检查数据量
        for device_code in self.missing_devices:
            device_dir = self.audio_root / device_code
            if not device_dir.exists():
                logger.info(f"数据目录不存在: {device_code}")
                return False

            # 统计 current 目录和日期目录下的 wav 文件数
            total_samples = 0
            for sub_dir in device_dir.iterdir():
                if sub_dir.is_dir():
                    # 新结构：{date}/normal/ 子目录
                    normal_dir = sub_dir / "normal"
                    if normal_dir.exists():
                        total_samples += len(list(normal_dir.glob("*.wav")))
                    # 兼容旧结构 + current 目录：直接存放的 wav
                    total_samples += len(list(sub_dir.glob("*.wav")))

            if total_samples < self.min_samples:
                logger.info(f"数据量不足: {device_code} -> {total_samples}/{self.min_samples}")
                return False

        logger.info("所有冷启动设备数据收集完成")
        return True

    def run_initial_training(self, on_device_trained=None) -> bool:
        """
        执行首次训练（冷启动）

        使用 IncrementalTrainer 的 cold_start_mode，收集所有可用数据进行全量训练。

        Args:
            on_device_trained: 可选回调 fn(device_code)，单设备训练完成后调用

        返回:
            True 表示训练成功
        """
        logger.info("=" * 70)
        logger.info("开始首次训练（冷启动）")
        logger.info("=" * 70)

        try:
            from auto_training.incremental_trainer import IncrementalTrainer

            # 从当前内存中的配置 dict 初始化训练器（配置来源为数据库）
            trainer = IncrementalTrainer(config=self.config)

            # 冷启动模式：收集所有目录的数据，用全量训练
            trainer.cold_start_mode = True
            trainer.use_days_ago = 0
            trainer.sample_hours = 0
            # 首次训练用更多轮数和稍高学习率
            trainer.epochs = 90
            trainer.learning_rate = 0.001

            success = trainer.run_daily_training(
                on_device_trained=on_device_trained
            )

            if success:
                logger.info("=" * 70)
                logger.info("首次训练完成，切换到正常检测模式")
                logger.info("=" * 70)
                self.is_cold_start = False

            return success

        except Exception as e:
            logger.error(f"首次训练失败: {e}", exc_info=True)
            return False


class IntegratedSystem:
    """
    集成系统：冷启动 + 拾音器监控 + 自动增量训练 + 数据清理

    作为 run_pickup_monitor.py 的上层封装，在其基础上增加：
    1. 冷启动自动训练
    2. APScheduler 定时增量训练和数据清理
    3. 训练完成后主动触发模型热更新
    """

    def __init__(self):
        self.deploy_root = Path(__file__).parent

        # =========================================================================
        # 配置加载来源：自动检测
        #   优先使用 YAML（config/rtsp_config.yaml 存在时）
        #   否则使用 SQLite 数据库
        # =========================================================================
        yaml_path = self.deploy_root / "config" / "rtsp_config.yaml"

        if yaml_path.exists():
            import yaml
            with open(yaml_path, 'r', encoding='utf-8') as f:
                full_config = yaml.safe_load(f)
            self.full_yaml_config = full_config
            self.auto_config = {'auto_training': full_config.get('auto_training', {})}
            logger.info(f"已从 YAML ({yaml_path.name}) 加载配置")
        else:
            self.full_yaml_config = None
            mgr = ConfigManager()
            self.auto_config = {'auto_training': mgr.get_system_config('auto_training')}
            mgr.close()
            logger.info(f"已从数据库加载 auto_training 配置 ({len(self.auto_config.get('auto_training', {}))} 项)")

        # 运行时对象
        self.scheduler = None
        self.pickup_system = None
        self.cold_start_manager = None
        self.cold_start_thread = None

    def _check_and_handle_cold_start(self) -> bool:
        """
        检查并处理冷启动

        利用 PickupMonitoringSystem 中已注册的设备列表,
        检查哪些设备缺少模型文件，对缺失设备启动冷启动训练。

        返回:
            True 表示处于冷启动模式
        """
        self.cold_start_manager = ColdStartManager(
            self.deploy_root,
            self.auto_config
        )

        # 从 pickup_system 的 multi_predictor 获取已注册设备
        registered = self.pickup_system.multi_predictor.device_model_map

        if not self.cold_start_manager.check_cold_start_needed(registered):
            logger.info("所有设备模型完整，进入正常检测模式")
            return False

        # 进入冷启动模式
        self.cold_start_manager.start_cold_start_mode()

        # 启动后台线程监控冷启动状态
        self.cold_start_thread = threading.Thread(
            target=self._cold_start_monitor_loop,
            daemon=True,
            name="cold-start-monitor"
        )
        self.cold_start_thread.start()

        return True

    def _cold_start_monitor_loop(self):
        # 后台线程：每分钟检查冷启动数据是否就绪，就绪后触发首次训练
        while self.cold_start_manager.is_cold_start:
            time.sleep(60)

            if self.cold_start_manager.check_ready_for_training():
                success = self.cold_start_manager.run_initial_training(
                    on_device_trained=self._reload_single_device
                )

                if success:
                    break
                else:
                    # 训练失败，等待10分钟后重试
                    logger.warning("首次训练失败，将在10分钟后重试")
                    time.sleep(600)

    def _reload_single_device(self, device_code: str):
        """训练回调：单个设备训练完成后即时重载该设备模型"""
        if not self.pickup_system:
            return
        mp = self.pickup_system.multi_predictor
        try:
            if mp.reload_device(device_code):
                logger.info(f"模型即时重载成功: {device_code}")
            else:
                logger.warning(f"模型即时重载失败: {device_code}")
        except Exception as e:
            logger.error(f"模型即时重载异常: {device_code} | {e}")

    def _reload_models_after_training(self):
        """
        训练完成后主动触发 MultiModelPredictor 重载所有设备模型

        相比被动等待60秒轮询检测 mtime 变化，主动调用 reload_device() 可以：
        1. 即时生效（零延迟）
        2. 有明确的成功/失败反馈
        """
        if not self.pickup_system:
            return

        mp = self.pickup_system.multi_predictor
        success_count = 0
        fail_count = 0

        for device_code in mp.registered_devices:
            try:
                if mp.reload_device(device_code):
                    success_count += 1
                else:
                    fail_count += 1
            except Exception as e:
                logger.error(f"重载设备模型失败: {device_code} | {e}")
                fail_count += 1

        logger.info(f"模型重载完成: 成功={success_count}, 失败={fail_count}")

    def _setup_auto_training_tasks(self):
        # 配置 APScheduler 定时任务：增量训练 + 数据清理
        if not self.auto_config.get('auto_training', {}).get('enabled', False):
            logger.info("自动训练已禁用（auto_training.enabled=false），跳过定时任务配置")
            return

        logger.info("=" * 70)
        logger.info("配置定时任务")
        logger.info("=" * 70)

        self.scheduler = BackgroundScheduler()

        # 定时增量训练
        incremental_cfg = self.auto_config['auto_training'].get('incremental', {})
        if incremental_cfg.get('enabled', False):
            schedule_time = incremental_cfg.get('schedule_time', '02:00')
            hour, minute = map(int, schedule_time.split(':'))

            self.scheduler.add_job(
                self._run_incremental_training,
                trigger=CronTrigger(hour=hour, minute=minute),
                id='incremental_training',
                name='每日增量训练',
                misfire_grace_time=3600  # 错过1小时内仍执行
            )
            logger.info(f"每日增量训练: 每天 {schedule_time}")

        # 定时数据清理
        data_cfg = self.auto_config['auto_training'].get('data', {})
        cleanup_time = data_cfg.get('cleanup_time', '00:00')
        hour, minute = map(int, cleanup_time.split(':'))

        self.scheduler.add_job(
            self._run_data_cleanup,
            trigger=CronTrigger(hour=hour, minute=minute),
            id='data_cleanup',
            name='每日数据清理',
            misfire_grace_time=3600
        )
        logger.info(f"每日数据清理: 每天 {cleanup_time}")

        self.scheduler.start()
        logger.info("定时任务调度器已启动")
        logger.info("=" * 70)

    def _run_incremental_training(self):
        # 定时任务回调：执行增量训练（逐设备串行，每完成一个即时重载）
        try:
            logger.info("定时任务触发：增量训练开始")

            from auto_training.incremental_trainer import IncrementalTrainer
            # 传 config dict 而非 YAML 路径，配置来源为数据库
            trainer = IncrementalTrainer(config=self.auto_config)
            success = trainer.run_daily_training(
                on_device_trained=self._reload_single_device
            )

            if not success:
                logger.warning("增量训练返回失败，保持当前推理模型不变")

        except Exception as e:
            logger.error(f"增量训练异常: {e}", exc_info=True)



    def _run_data_cleanup(self):
        # 定时任务回调：执行每日数据清理
        try:
            logger.info("定时任务触发：数据清理开始")
            from auto_training.data_cleanup import DataCleaner

            # 传 config dict 而非 YAML 路径，配置来源为数据库
            cleaner = DataCleaner(config=self.auto_config)
            cleaner.run_cleanup()
        except Exception as e:
            logger.error(f"数据清理异常: {e}", exc_info=True)

    def start(self):
        # 主启动流程
        logger.info("=" * 70)
        logger.info("拾音器异响检测系统（带自动训练）")
        logger.info("=" * 70)

        # 1. 创建 PickupMonitoringSystem（会初始化 multi_predictor + 注册设备）
        logger.info("初始化监控系统...")
        from run_pickup_monitor import PickupMonitoringSystem
        self.pickup_system = PickupMonitoringSystem(yaml_config=self.full_yaml_config)

        # 2. 检查冷启动（需要在 pickup_system 初始化之后，因为需要设备注册信息）
        is_cold_start = self._check_and_handle_cold_start()

        # 3. 设置定时任务
        self._setup_auto_training_tasks()

        # 4. 覆盖信号处理（确保关闭 scheduler）
        signal.signal(signal.SIGINT, self._signal_handler)
        signal.signal(signal.SIGTERM, self._signal_handler)

        # 5. 启动拾音器监控（这是阻塞调用，包含主循环）
        logger.info("启动拾音器监控...")
        self.pickup_system.start()

    def stop(self):
        # 关闭所有组件
        logger.info("停止系统...")

        # 先关 scheduler，避免训练任务在关停过程中触发
        if self.scheduler and self.scheduler.running:
            self.scheduler.shutdown(wait=False)
            logger.info("定时任务调度器已停止")

        if self.pickup_system:
            self.pickup_system.stop()

        logger.info("系统已停止")

    def _signal_handler(self, signum, frame):
        logger.info(f"收到信号 {signum}")
        self.stop()
        sys.exit(0)


def main():
    setup_logging()
    system = IntegratedSystem()
    system.start()


if __name__ == "__main__":
    main()