147 lines
4.4 KiB
Python
147 lines
4.4 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:
|
|
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():
|
|
"""启动定时清理调度器"""
|
|
# 每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) # 每分钟检查一次
|
|
|
|
scheduler_thread = threading.Thread(target=run_scheduler, daemon=True)
|
|
scheduler_thread.start()
|
|
logger.info("定时清理调度器已启动")
|
|
|
|
|
|
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服务")
|
|
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
|
|
task_stream_manager.cleanup_all()
|
|
except:
|
|
pass
|
|
|
|
# 停止所有任务
|
|
if task_manager:
|
|
cleaned_count = task_manager.cleanup_all_tasks()
|
|
logger.info(f"清理了 {cleaned_count} 个任务")
|
|
|
|
logger.info("服务已退出")
|
|
os._exit(0)
|
|
|
|
|
|
signal.signal(signal.SIGINT, graceful_exit)
|
|
signal.signal(signal.SIGTERM, graceful_exit)
|
|
|
|
# 启动服务
|
|
socketio.run(app,
|
|
host='0.0.0.0',
|
|
port=9309,
|
|
debug=True,
|
|
use_reloader=False,
|
|
allow_unsafe_werkzeug=True) |