Yolov/main.py

# main.py
# 主程序入口
import os
import signal
import threading
import time
import schedule
import torch
from log import logger
from global_data import gd
from server import socketio, app
from task_manager import task_manager

# 初始化全局数据
gd._init()

# 设置全局变量
gd.set_value('task_manager', task_manager)  # 直接在这里初始化
gd.set_value('latest_drone_data', None)
gd.set_value('mqtt_data_lock', threading.Lock())

# 初始化任务管理器
logger.info("任务管理器初始化...")
logger.info(f"任务管理器实例: {task_manager}")


def cleanup_stopped_tasks():
    """定期清理已停止的任务"""
    try:
        if not task_manager:
            return
        # 取得所有任务
        all_tasks = task_manager.get_all_tasks()
        # 计数
        cleaned_count = 0

        for task in all_tasks:
            # 取得任务ID
            task_id = task.get('task_id')
            # 取得任务状态
            status = task.get('status', 'unknown')

            # 清理已停止、失败或错误的任务
            if status in ['stopped', 'failed', 'error']:
                # 检查任务线程是否真的已经停止
                thread = task_manager.tasks.get(task_id, {}).get('thread')
                if thread and hasattr(thread, 'is_alive'):
                    if not thread.is_alive():
                        if task_manager.cleanup_task(task_id):
                            cleaned_count += 1
                else:
                    if task_manager.cleanup_task(task_id):
                        cleaned_count += 1

        if cleaned_count > 0:
            logger.info(f"定时清理已完成，清理了 {cleaned_count} 个任务")

    except Exception as e:
        logger.error(f"定时清理任务失败: {str(e)}")


def start_cleanup_scheduler():
    """启动定时清理调度器"""
    if not schedule:
        logger.warning("schedule模块不可用，跳过定时清理调度器")
        return

    try:
        # 每5分钟清理一次
        schedule.every(5).minutes.do(cleanup_stopped_tasks)

        # 每分钟检查一次待清理任务
        schedule.every(1).minutes.do(check_pending_cleanup)

        def run_scheduler():
            while True:
                try:
                    schedule.run_pending()
                except Exception as e:
                    logger.error(f"定时任务执行失败: {str(e)}")
                #time.sleep(60)  # 每分钟检查一次
                time.sleep(1)  # 缩短为1秒检查频率

        scheduler_thread = threading.Thread(target=run_scheduler, daemon=True)
        scheduler_thread.start()
        logger.info("定时清理调度器已启动")
    except Exception as e:
        logger.error(f"启动定时清理调度器失败: {str(e)}")


def check_pending_cleanup():
    """检查待清理的任务"""
    try:
        if not task_manager:
            return

        # 检查是否有长时间处于停止状态但未清理的任务
        for task_id, task_info in task_manager.tasks.items():
            status = task_info.get('status', 'unknown')
            last_updated = task_info.get('last_updated', 0)

            if status == 'stopped' and time.time() - last_updated > 60:  # 停止超过60秒
                logger.info(f"发现长时间停止未清理的任务: {task_id}")
                task_manager.cleanup_task(task_id)

    except Exception as e:
        logger.error(f"检查待清理任务失败: {str(e)}")


# main.py 修改部分

if __name__ == '__main__':
    logger.info("启动多任务版YOLOv8服务")
    # 打印PyTorch版本和CUDA可用性
    logger.info(f"PyTorch版本: {torch.__version__}, CUDA可用: {torch.cuda.is_available()}")

    # 启动定时清理
    start_cleanup_scheduler()

    # 初始化任务推流管理器
    try:
        from task_stream_manager import task_stream_manager

        task_stream_manager.start_health_monitor()
        logger.info("任务推流管理器初始化完成")
    except Exception as e:
        logger.error(f"初始化推流管理器失败: {str(e)}")


    # 退出服务
    def graceful_exit(signum, frame):
        logger.info("收到退出信号，停止所有服务...")

        # 停止所有任务推流
        try:
            from task_stream_manager import task_stream_manager
            if 'task_stream_manager' in globals():
                task_stream_manager.cleanup_all()
        except Exception as e:
            logger.error(f"清理推流管理器失败: {str(e)}")

        # 停止所有任务
        try:
            if task_manager:
                cleaned_count = task_manager.cleanup_all_tasks()
                logger.info(f"清理了 {cleaned_count} 个任务")
        except Exception as e:
            logger.error(f"清理任务管理器失败: {str(e)}")

        logger.info("服务已退出")
        os._exit(0)


    # Windows平台特殊处理信号
    try:
        if os.name != 'nt':  # Unix/Linux/Mac
            signal.signal(signal.SIGINT, graceful_exit)
            signal.signal(signal.SIGTERM, graceful_exit)
        else:  # Windows
            # Windows下使用键盘中断处理
            import atexit

            atexit.register(lambda: graceful_exit(None, None))
    except Exception as e:
        logger.error(f"注册退出处理函数失败: {str(e)}")

    # 启动服务
    try:
        socketio.run(app,
                     host='0.0.0.0',
                     port=9309,
                     debug=True,
                     use_reloader=False,
                     allow_unsafe_werkzeug=True)
    except KeyboardInterrupt:
        logger.info("收到键盘中断信号")
        graceful_exit(None, None)
    except Exception as e:
        logger.error(f"启动服务失败: {str(e)}")
        import traceback

        logger.error(traceback.format_exc())