182 lines
5.7 KiB
Python
182 lines
5.7 KiB
Python
# 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(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())
|