1636 lines
55 KiB
Python
1636 lines
55 KiB
Python
# server.py
|
||
import hashlib
|
||
import os
|
||
import pickle
|
||
import secrets
|
||
from datetime import datetime
|
||
import json
|
||
import logging
|
||
|
||
import torch
|
||
from flask import Flask, jsonify, request, render_template
|
||
from flask_socketio import SocketIO, emit, join_room, leave_room
|
||
from config import get_default_config
|
||
from mandatory_model_crypto import MandatoryModelEncryptor
|
||
from model_upload_manager import get_upload_manager
|
||
from task_manager import task_manager # 导入任务管理器
|
||
from global_data import gd
|
||
from log import logger
|
||
import time
|
||
import traceback
|
||
from mandatory_model_crypto import ModelEncryptionService, validate_models_before_task, verify_single_model_api
|
||
|
||
# 配置日志
|
||
logging.basicConfig(
|
||
level=logging.INFO,
|
||
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
|
||
handlers=[
|
||
logging.StreamHandler(),
|
||
logging.FileHandler("public_server.log", encoding="utf-8")
|
||
]
|
||
)
|
||
pub_logger = logging.getLogger("PublicServer")
|
||
|
||
# Flask初始化
|
||
app = Flask(__name__, static_url_path='/static')
|
||
CORS(app)
|
||
|
||
# 处理Windows平台兼容性问题
|
||
import platform
|
||
|
||
if platform.system() == 'Windows':
|
||
socketio = SocketIO(app,
|
||
cors_allowed_origins="*",
|
||
async_mode='threading',
|
||
allow_unsafe_werkzeug=True, #
|
||
max_http_buffer_size=5 * 1024 * 1024)
|
||
else:
|
||
socketio = SocketIO(app,
|
||
cors_allowed_origins="*",
|
||
async_mode='threading',
|
||
max_http_buffer_size=5 * 1024 * 1024)
|
||
|
||
_initialized = False
|
||
|
||
|
||
@app.before_request
|
||
def initialize_once():
|
||
global _initialized
|
||
if not _initialized:
|
||
with app.app_context():
|
||
gd.set_value('task_manager', task_manager)
|
||
logger.info("任务管理器初始化完成")
|
||
_initialized = True
|
||
|
||
|
||
# ======================= Flask路由 =======================
|
||
@app.route('/')
|
||
def task_management():
|
||
"""任务管理页面"""
|
||
return render_template("task_management.html")
|
||
|
||
|
||
@app.route('/video_player')
|
||
def video_player():
|
||
"""视频播放页面"""
|
||
return render_template("flv2.html")
|
||
|
||
|
||
# server.py - 修改 create_task 函数
|
||
|
||
@app.route('/api/tasks/create', methods=['POST'])
|
||
def create_task():
|
||
"""创建新任务 - 强制模型加密和密钥验证"""
|
||
try:
|
||
config = get_default_config()
|
||
config['socketIO'] = socketio
|
||
|
||
if not request.json:
|
||
return jsonify({"status": "error", "message": "请求体不能为空"}), 400
|
||
|
||
data = request.json
|
||
import json
|
||
logger.info(f"请求参数: {json.dumps(data, indent=2, ensure_ascii=False)}")
|
||
logger.info(f"收到创建任务请求: {data.get('taskname', '未命名')}")
|
||
|
||
# 验证必须的参数
|
||
if 'rtmp_url' not in data:
|
||
return jsonify({"status": "error", "message": "必须提供rtmp_url"}), 400
|
||
|
||
if 'models' not in data or not isinstance(data['models'], list):
|
||
return jsonify({"status": "error", "message": "必须提供models列表"}), 400
|
||
|
||
# 检查模型数量
|
||
if len(data['models']) == 0:
|
||
return jsonify({"status": "error", "message": "models列表不能为空"}), 400
|
||
|
||
# ================= 关键修改：创建任务前的模型验证 =================
|
||
logger.info("开始创建任务前的模型验证...")
|
||
|
||
# 1. 检查所有模型是否都有加密密钥
|
||
for i, model_data in enumerate(data['models']):
|
||
if 'encryption_key' not in model_data:
|
||
return jsonify({
|
||
"status": "error",
|
||
"message": f"模型 {i} 必须提供encryption_key"
|
||
}), 400
|
||
|
||
# 2. 验证所有模型的密钥
|
||
task_config = {
|
||
'models': data['models']
|
||
}
|
||
|
||
validation_result = validate_models_before_task(task_config)
|
||
|
||
if not validation_result['success']:
|
||
logger.error(f"模型验证失败: {validation_result.get('error', '未知错误')}")
|
||
|
||
# 提供详细的验证结果
|
||
error_details = []
|
||
for result in validation_result.get('validation_results', []):
|
||
if not result.get('key_valid', False):
|
||
error_details.append(f"模型 {result['model_index']}: {result.get('error', '验证失败')}")
|
||
|
||
error_message = validation_result.get('error', '模型验证失败')
|
||
if error_details:
|
||
error_message += f" | 详情: {', '.join(error_details)}"
|
||
|
||
return jsonify({
|
||
"status": "error",
|
||
"message": error_message,
|
||
"data": validation_result
|
||
}), 400
|
||
|
||
logger.info(f"模型验证通过: {validation_result['valid_models']}/{validation_result['total_models']} 个模型有效")
|
||
# ================= 验证结束 =================
|
||
|
||
# 更新配置
|
||
config['rtmp']['url'] = data['rtmp_url']
|
||
|
||
if 'uavType' in data and data['uavType']:
|
||
config['task']['uavType'] = data['uavType']
|
||
|
||
if 'algoInstancesName' in data and data['algoInstancesName']:
|
||
config['task']['algoInstancesName'] = data['algoInstancesName']
|
||
|
||
if 'push_url' in data and data['push_url']:
|
||
config['push']['url'] = data['push_url']
|
||
|
||
if 'taskname' in data:
|
||
config['task']['taskname'] = data['taskname']
|
||
else:
|
||
config['task']['taskname'] = f"task_{int(time.time())}"
|
||
|
||
if 'AlgoId' in data:
|
||
config['task']['aiid'] = data['AlgoId']
|
||
|
||
# 处理多模型配置 - 使用已验证的模型
|
||
config['models'] = []
|
||
|
||
for i, model_data in enumerate(data['models']):
|
||
# 此时密钥已验证通过
|
||
encryption_key = model_data['encryption_key']
|
||
# 使用加密模型文件名（相对路径）
|
||
model_path = model_data.get('path', f'model_{i}.enc')
|
||
if not model_path.startswith('encrypted_models/'):
|
||
model_filename = os.path.basename(model_path)
|
||
if not model_filename.endswith('.enc'):
|
||
model_filename += '.enc'
|
||
model_path = f"{model_filename}"
|
||
|
||
# 构建模型配置
|
||
model_config = {
|
||
'path': model_path,
|
||
'encryption_key': encryption_key,
|
||
'encrypted': True,
|
||
'tags': model_data.get('tags', {}),
|
||
'conf_thres': float(model_data.get('conf_thres', 0.25)),
|
||
'iou_thres': float(model_data.get('iou_thres', 0.45)),
|
||
'imgsz': max(128, min(1920, int(model_data.get('imgsz', 640)))),
|
||
'color': model_data.get('color'),
|
||
'line_width': int(model_data.get('line_width', 1)),
|
||
'device': model_data.get('device', 'cuda:0' if torch.cuda.is_available() else 'cpu'),
|
||
'half': model_data.get('half', True),
|
||
'enabled': model_data.get('enabled', True),
|
||
# 注意：不再需要 download_url，模型从本地加载
|
||
}
|
||
|
||
config['models'].append(model_config)
|
||
logger.info(f"添加已验证的加密模型 {i}: {model_path}")
|
||
|
||
# 在创建任务前清理已停止的任务，释放资源
|
||
logger.info("创建任务前清理已停止的任务...")
|
||
cleaned_count = task_manager.cleanup_stopped_tasks()
|
||
if cleaned_count > 0:
|
||
logger.info(f"已清理 {cleaned_count} 个已停止的任务")
|
||
|
||
# 检查系统资源
|
||
active_tasks = task_manager.get_active_tasks_count()
|
||
max_tasks = task_manager.get_current_max_tasks()
|
||
|
||
if active_tasks >= max_tasks:
|
||
# 尝试强制清理一些资源
|
||
logger.warning(f"当前活动任务数 {active_tasks} 已达到限制 {max_tasks}，尝试强制清理")
|
||
|
||
# 清理所有非运行状态的任务
|
||
all_tasks = task_manager.get_all_tasks()
|
||
force_cleaned = 0
|
||
for task in all_tasks:
|
||
if task['status'] not in ['running', 'starting', 'creating']:
|
||
if task_manager.cleanup_task(task['task_id']):
|
||
force_cleaned += 1
|
||
|
||
if force_cleaned > 0:
|
||
logger.info(f"强制清理了 {force_cleaned} 个非运行状态的任务")
|
||
|
||
# 再次检查
|
||
active_tasks = task_manager.get_active_tasks_count()
|
||
if active_tasks >= max_tasks:
|
||
return jsonify({
|
||
"status": "error",
|
||
"message": f"已达到最大并发任务数限制 ({max_tasks})，请等待其他任务完成或停止部分任务"
|
||
}), 503
|
||
|
||
# 创建任务
|
||
logger.info(f"开始创建任务，包含 {len(config['models'])} 个已验证的加密模型...")
|
||
|
||
try:
|
||
task_id = task_manager.create_task(config, socketio)
|
||
logger.info(f"任务创建成功，ID: {task_id}")
|
||
except Exception as e:
|
||
logger.error(f"任务创建失败: {str(e)}")
|
||
return jsonify({
|
||
"status": "error",
|
||
"message": f"任务创建失败: {str(e)}"
|
||
}), 500
|
||
|
||
# 启动任务
|
||
logger.info(f"启动任务 {task_id}...")
|
||
success = task_manager.start_task(task_id)
|
||
|
||
if success:
|
||
logger.info(f"任务启动成功: {task_id}")
|
||
|
||
# 获取任务详细信息
|
||
task_status = task_manager.get_task_status(task_id)
|
||
|
||
return jsonify({
|
||
"status": "success",
|
||
"message": "任务创建并启动成功",
|
||
"data": {
|
||
"task_id": task_id,
|
||
"models_count": len(config['models']),
|
||
"encryption_required": True,
|
||
"key_validated": True,
|
||
"validation_result": validation_result,
|
||
"task_info": task_status
|
||
}
|
||
})
|
||
else:
|
||
logger.error(f"任务启动失败: {task_id}")
|
||
# 如果启动失败，清理任务
|
||
task_manager.cleanup_task(task_id)
|
||
return jsonify({
|
||
"status": "error",
|
||
"message": "任务创建成功但启动失败",
|
||
"task_id": task_id
|
||
}), 500
|
||
|
||
except Exception as e:
|
||
logger.error(f"创建任务失败: {str(e)}")
|
||
logger.error(traceback.format_exc())
|
||
return jsonify({
|
||
"status": "error",
|
||
"message": f"创建任务失败: {str(e)}"
|
||
}), 500
|
||
|
||
|
||
@app.route('/api/system/resources', methods=['GET'])
|
||
def get_system_resources():
|
||
"""获取系统资源使用情况"""
|
||
try:
|
||
resources = gd.get_value('system_resources')
|
||
max_tasks = gd.get_value('max_concurrent_tasks', 5)
|
||
|
||
if not resources:
|
||
# 实时获取资源
|
||
import psutil
|
||
resources = {
|
||
'cpu_percent': psutil.cpu_percent(),
|
||
'memory_percent': psutil.virtual_memory().percent,
|
||
'memory_used': psutil.virtual_memory().used // (1024 * 1024),
|
||
'memory_total': psutil.virtual_memory().total // (1024 * 1024),
|
||
'timestamp': datetime.now().isoformat()
|
||
}
|
||
|
||
# 获取任务统计
|
||
active_tasks = task_manager.get_active_tasks_count()
|
||
total_tasks = len(task_manager.tasks)
|
||
|
||
return jsonify({
|
||
"status": "success",
|
||
"data": {
|
||
"resources": resources,
|
||
"tasks": {
|
||
"active": active_tasks,
|
||
"total": total_tasks,
|
||
"max_concurrent": max_tasks
|
||
}
|
||
}
|
||
})
|
||
except Exception as e:
|
||
logger.error(f"获取系统资源失败: {str(e)}")
|
||
return jsonify({
|
||
"status": "error",
|
||
"message": f"获取系统资源失败: {str(e)}"
|
||
}), 500
|
||
|
||
|
||
@app.route('/api/models/encrypt', methods=['POST'])
|
||
def encrypt_model():
|
||
"""加密模型文件 - 强制加密"""
|
||
try:
|
||
config = get_default_config()
|
||
data = request.json
|
||
model_path = data.get('model_path')
|
||
output_path = data.get('output_path')
|
||
password = data.get('password')
|
||
download_url = data.get('download_url')
|
||
if not all([model_path, output_path, password, download_url]):
|
||
return jsonify({"status": "error", "message": "缺少必要参数"}), 400
|
||
local_path = os.path.join(config['model_path'], model_path)
|
||
output_path = os.path.join(config['model_path'], output_path)
|
||
# 使用强制加密器
|
||
from mandatory_model_crypto import MandatoryModelEncryptor
|
||
encryptor = MandatoryModelEncryptor()
|
||
|
||
result = encryptor.encrypt_model(local_path, output_path, password, require_encryption=True)
|
||
|
||
if result['success']:
|
||
return jsonify({
|
||
"status": "success",
|
||
"message": f"模型加密成功: {output_path}",
|
||
"data": {
|
||
"model_hash": result.get('model_hash'),
|
||
"key_hash": result.get('key_hash'),
|
||
"output_path": result.get('output_path')
|
||
}
|
||
})
|
||
else:
|
||
return jsonify({
|
||
"status": "error",
|
||
"message": f"模型加密失败: {result.get('error', '未知错误')}"
|
||
}), 500
|
||
|
||
except Exception as e:
|
||
logger.error(f"加密模型失败: {str(e)}")
|
||
return jsonify({
|
||
"status": "error",
|
||
"message": f"加密模型失败: {str(e)}"
|
||
}), 500
|
||
|
||
|
||
@app.route('/api/models/verify_key', methods=['POST'])
|
||
def verify_model_key():
|
||
"""验证模型密钥"""
|
||
try:
|
||
data = request.json
|
||
model_path = data.get('model_path')
|
||
encryption_key = data.get('encryption_key')
|
||
|
||
if not all([model_path, encryption_key]):
|
||
return jsonify({"status": "error", "message": "缺少必要参数"}), 400
|
||
|
||
# 检查模型文件是否存在
|
||
full_path = os.path.join('encrypted_models', os.path.basename(model_path))
|
||
if not os.path.exists(full_path):
|
||
return jsonify({"status": "error", "message": f"模型文件不存在: {full_path}"}), 400
|
||
|
||
# 使用强制加密器验证密钥
|
||
from mandatory_model_crypto import MandatoryModelEncryptor
|
||
encryptor = MandatoryModelEncryptor()
|
||
|
||
# 检查是否为正确加密的模型
|
||
if not encryptor.is_properly_encrypted(full_path):
|
||
return jsonify({
|
||
"status": "error",
|
||
"message": "模型文件未正确加密",
|
||
"valid": False
|
||
}), 400
|
||
|
||
# 验证密钥
|
||
verify_result = encryptor.decrypt_model(full_path, encryption_key, verify_key=True)
|
||
|
||
if verify_result['success']:
|
||
return jsonify({
|
||
"status": "success",
|
||
"message": "密钥验证成功",
|
||
"data": {
|
||
"valid": True,
|
||
"model_hash": verify_result.get('model_hash', '')[:16],
|
||
"model_size": verify_result.get('original_size', 0)
|
||
}
|
||
})
|
||
else:
|
||
return jsonify({
|
||
"status": "error",
|
||
"message": f"密钥验证失败: {verify_result.get('error', '未知错误')}",
|
||
"data": {
|
||
"valid": False,
|
||
"error": verify_result.get('error', '未知错误')
|
||
}
|
||
}), 400
|
||
|
||
except Exception as e:
|
||
logger.error(f"验证密钥失败: {str(e)}")
|
||
return jsonify({
|
||
"status": "error",
|
||
"message": f"验证密钥失败: {str(e)}"
|
||
}), 500
|
||
|
||
|
||
@app.route('/api/models/generate_key', methods=['POST'])
|
||
def generate_secure_encryption_key():
|
||
"""生成安全的加密密钥"""
|
||
try:
|
||
from mandatory_model_crypto import MandatoryModelEncryptor
|
||
|
||
# 生成密钥
|
||
encryptor = MandatoryModelEncryptor()
|
||
key_info = encryptor.generate_secure_key()
|
||
|
||
return jsonify({
|
||
"status": "success",
|
||
"message": "加密密钥生成成功",
|
||
"data": {
|
||
"key": key_info['key'],
|
||
"key_hash": key_info['key_hash'],
|
||
"short_hash": key_info['short_hash'],
|
||
}
|
||
})
|
||
|
||
except Exception as e:
|
||
logger.error(f"生成加密密钥失败: {str(e)}")
|
||
return jsonify({
|
||
"status": "error",
|
||
"message": f"生成加密密钥失败: {str(e)}"
|
||
}), 500
|
||
|
||
|
||
# server.py 中的 stop_detection 路由修改
|
||
|
||
@app.route('/api/tasks/<task_id>/stop', methods=['POST'])
|
||
def stop_task(task_id):
|
||
"""停止指定任务"""
|
||
try:
|
||
logger.info(f"接收到停止任务请求: {task_id}")
|
||
|
||
# 获取请求参数
|
||
data = request.json or {}
|
||
force = data.get('force', False)
|
||
timeout = data.get('timeout', 10) # 默认10秒超时
|
||
|
||
# 调用任务管理器停止任务
|
||
success = task_manager.stop_task(task_id, force=force)
|
||
|
||
if success:
|
||
logger.info(f"任务停止成功: {task_id}")
|
||
return jsonify({
|
||
"status": "success",
|
||
"message": f"任务停止成功: {task_id}",
|
||
"task_id": task_id,
|
||
"stopped": True
|
||
})
|
||
else:
|
||
logger.warning(f"停止任务失败: {task_id}")
|
||
return jsonify({
|
||
"status": "error",
|
||
"message": f"停止任务失败: {task_id}",
|
||
"task_id": task_id,
|
||
"stopped": False
|
||
}), 500
|
||
except Exception as e:
|
||
logger.error(f"停止任务异常: {str(e)}")
|
||
logger.error(traceback.format_exc())
|
||
return jsonify({
|
||
"status": "error",
|
||
"message": f"停止任务异常: {str(e)}",
|
||
"task_id": task_id
|
||
}), 500
|
||
|
||
|
||
@app.route('/api/tasks/<task_id>/status', methods=['GET'])
|
||
def get_task_status(task_id):
|
||
"""获取任务状态（仅支持多模型）"""
|
||
status = task_manager.get_task_status(task_id)
|
||
if status:
|
||
# 增强返回信息，包含模型详情
|
||
enhanced_status = {
|
||
'task_id': status['task_id'],
|
||
'status': status['status'],
|
||
'config': status['config'],
|
||
'models': status.get('models', []), # 直接返回模型列表
|
||
'stats': status['stats'],
|
||
'created_at': status['created_at']
|
||
}
|
||
|
||
return jsonify({
|
||
"status": "success",
|
||
"data": enhanced_status
|
||
})
|
||
else:
|
||
return jsonify({
|
||
"status": "error",
|
||
"message": f"任务不存在: {task_id}"
|
||
}), 404
|
||
|
||
|
||
@app.route('/api/tasks', methods=['GET'])
|
||
def get_all_tasks():
|
||
"""获取所有任务（仅支持多模型）"""
|
||
tasks = task_manager.get_all_tasks()
|
||
# 打印tasks内容
|
||
import json
|
||
print("获取到的任务列表：")
|
||
# print(json.dumps(tasks, indent=2, ensure_ascii=False))
|
||
|
||
# 增强任务信息
|
||
enhanced_tasks = []
|
||
for task in tasks:
|
||
enhanced_task = {
|
||
'task_id': task['task_id'],
|
||
'status': task['status'],
|
||
'algoInstancesName': task['config']['algoInstancesName'],
|
||
'uavType': task['config']['uavType'],
|
||
'config': {
|
||
'rtmp_url': task['config']['rtmp_url'],
|
||
'taskname': task['config']['taskname'],
|
||
'push_url': task['config'].get('push_url', ''),
|
||
'enable_push': task['config'].get('enable_push', False)
|
||
},
|
||
'models': task.get('models', []), # 直接返回模型列表
|
||
'stats': task['stats'],
|
||
'performance': task.get('performance', {
|
||
'fps': 0,
|
||
'avg_process_time': 0,
|
||
'latency': 0,
|
||
'last_fps': 0
|
||
}), # 性能指标
|
||
'created_at': task['created_at']
|
||
}
|
||
enhanced_tasks.append(enhanced_task)
|
||
|
||
return jsonify({
|
||
"status": "success",
|
||
"data": {
|
||
"tasks": enhanced_tasks,
|
||
"total": len(tasks),
|
||
"active": task_manager.get_active_tasks_count(),
|
||
"models_count": sum(len(t.get('models', [])) for t in enhanced_tasks)
|
||
}
|
||
})
|
||
|
||
|
||
@app.route('/api/tasks/<task_id>/cleanup', methods=['POST'])
|
||
def cleanup_task(task_id):
|
||
"""清理任务资源"""
|
||
success = task_manager.cleanup_task(task_id)
|
||
if success:
|
||
return jsonify({
|
||
"status": "success",
|
||
"message": f"任务资源已清理: {task_id}"
|
||
})
|
||
else:
|
||
return jsonify({
|
||
"status": "error",
|
||
"message": f"清理任务失败: {task_id}"
|
||
}), 500
|
||
|
||
|
||
@app.route('/api/tasks/cleanup_all', methods=['POST'])
|
||
def cleanup_all_tasks():
|
||
"""清理所有任务"""
|
||
cleaned_count = task_manager.cleanup_all_tasks()
|
||
return jsonify({
|
||
"status": "success",
|
||
"message": f"已清理 {cleaned_count} 个任务",
|
||
"cleaned_count": cleaned_count
|
||
})
|
||
|
||
|
||
@app.route('/api/tasks/cleanup_stopped', methods=['POST'])
|
||
def cleanup_stopped_tasks():
|
||
"""清理所有已停止的任务"""
|
||
cleaned_count = task_manager.cleanup_stopped_tasks()
|
||
return jsonify({
|
||
"status": "success",
|
||
"message": f"已清理 {cleaned_count} 个已停止的任务",
|
||
"cleaned_count": cleaned_count
|
||
})
|
||
|
||
|
||
@app.route('/api/system/status', methods=['GET'])
|
||
def get_system_status():
|
||
"""获取系统状态"""
|
||
try:
|
||
import psutil
|
||
|
||
system_info = {
|
||
"cpu_percent": psutil.cpu_percent(),
|
||
"memory_percent": psutil.virtual_memory().percent,
|
||
"disk_percent": psutil.disk_usage('/').percent,
|
||
"active_tasks": task_manager.get_active_tasks_count(),
|
||
"total_tasks": len(task_manager.tasks),
|
||
"max_concurrent_tasks": task_manager.get_current_max_tasks()
|
||
}
|
||
|
||
# GPU信息（如果可用）
|
||
try:
|
||
import GPUtil
|
||
gpus = GPUtil.getGPUs()
|
||
gpu_info = []
|
||
for gpu in gpus:
|
||
gpu_info.append({
|
||
"id": gpu.id,
|
||
"name": gpu.name,
|
||
"load": gpu.load * 100,
|
||
"memory_used": gpu.memoryUsed,
|
||
"memory_total": gpu.memoryTotal,
|
||
"temperature": gpu.temperature
|
||
})
|
||
system_info["gpus"] = gpu_info
|
||
except ImportError:
|
||
# 如果没有安装GPUtil，尝试使用torch获取GPU信息
|
||
if torch.cuda.is_available():
|
||
gpu_info = []
|
||
for i in range(torch.cuda.device_count()):
|
||
gpu_info.append({
|
||
"id": i,
|
||
"name": torch.cuda.get_device_name(i),
|
||
"memory_used": torch.cuda.memory_allocated(i) / 1024 ** 2,
|
||
"memory_total": torch.cuda.get_device_properties(i).total_memory / 1024 ** 2
|
||
})
|
||
system_info["gpus"] = gpu_info
|
||
|
||
return jsonify({
|
||
"status": "success",
|
||
"data": system_info
|
||
})
|
||
except Exception as e:
|
||
logger.error(f"获取系统状态失败: {str(e)}")
|
||
return jsonify({
|
||
"status": "error",
|
||
"message": f"获取系统状态失败: {str(e)}"
|
||
}), 500
|
||
|
||
|
||
# WebSocket事件 - 按任务ID发送
|
||
@socketio.on('connect')
|
||
def handle_connect():
|
||
logger.info(f"Socket客户端已连接: {request.sid}")
|
||
|
||
|
||
@socketio.on('disconnect')
|
||
def handle_disconnect():
|
||
logger.info(f"Socket客户端断开: {request.sid}")
|
||
|
||
|
||
@socketio.on('subscribe_task')
|
||
def handle_subscribe_task(data):
|
||
"""订阅特定任务的WebSocket消息"""
|
||
task_id = data.get('task_id')
|
||
if task_id:
|
||
# 这里可以记录客户端订阅关系
|
||
logger.info(f"客户端 {request.sid} 订阅任务: {task_id}")
|
||
return {"status": "subscribed", "task_id": task_id}
|
||
|
||
return {"status": "error", "message": "需要提供task_id"}
|
||
|
||
|
||
@app.route('/api/system/resource_limits', methods=['GET', 'POST'])
|
||
def manage_resource_limits():
|
||
"""获取或设置资源限制"""
|
||
if request.method == 'GET':
|
||
# 获取当前资源限制
|
||
resource_monitor = gd.get_value('resource_monitor')
|
||
if resource_monitor:
|
||
return jsonify({
|
||
"status": "success",
|
||
"data": resource_monitor.resource_limits
|
||
})
|
||
else:
|
||
config = get_default_config()
|
||
return jsonify({
|
||
"status": "success",
|
||
"message": "资源监控器未初始化",
|
||
"data": config['resource_limits']
|
||
})
|
||
|
||
elif request.method == 'POST':
|
||
# 更新资源限制
|
||
try:
|
||
data = request.json
|
||
resource_monitor = gd.get_value('resource_monitor', None)
|
||
|
||
if resource_monitor and data:
|
||
# 更新限制
|
||
resource_monitor.resource_limits.update(data)
|
||
|
||
# 更新任务管理器中的限制
|
||
task_manager.resource_limits.update(data)
|
||
|
||
logger.info(f"更新资源限制: {data}")
|
||
return jsonify({
|
||
"status": "success",
|
||
"message": "资源限制更新成功"
|
||
})
|
||
else:
|
||
return jsonify({
|
||
"status": "error",
|
||
"message": "资源监控器未初始化或请求数据无效"
|
||
}), 400
|
||
except Exception as e:
|
||
logger.error(f"更新资源限制失败: {str(e)}")
|
||
return jsonify({
|
||
"status": "error",
|
||
"message": f"更新失败: {str(e)}"
|
||
}), 500
|
||
|
||
|
||
@app.route('/api/tasks/<task_id>/models', methods=['GET'])
|
||
def get_task_models(task_id):
|
||
"""获取任务中的模型配置"""
|
||
try:
|
||
task = task_manager.get_task_status(task_id)
|
||
if not task:
|
||
return jsonify({
|
||
"status": "error",
|
||
"message": f"任务不存在: {task_id}"
|
||
}), 404
|
||
|
||
# 获取任务中的模型配置
|
||
models_info = []
|
||
if task_id in task_manager.tasks:
|
||
task_info = task_manager.tasks[task_id]
|
||
config = task_info.get('config', {})
|
||
models_config = config.get('models', [])
|
||
|
||
for i, model_config in enumerate(models_config):
|
||
models_info.append({
|
||
'id': i,
|
||
'name': os.path.basename(model_config.get('path', '')).split('.')[0],
|
||
'path': model_config.get('path'),
|
||
'conf_thres': model_config.get('conf_thres'),
|
||
'tags': model_config.get('tags', {}),
|
||
'color': model_config.get('color'),
|
||
'enabled': model_config.get('enabled', True)
|
||
})
|
||
|
||
return jsonify({
|
||
"status": "success",
|
||
"data": {
|
||
"task_id": task_id,
|
||
"models": models_info
|
||
}
|
||
})
|
||
except Exception as e:
|
||
logger.error(f"获取任务模型失败: {str(e)}")
|
||
return jsonify({
|
||
"status": "error",
|
||
"message": f"获取失败: {str(e)}"
|
||
}), 500
|
||
|
||
|
||
# server.py 添加以下路由
|
||
|
||
@app.route('/api/tasks/<task_id>/stream/status', methods=['GET'])
|
||
def get_task_stream_status(task_id):
|
||
"""获取任务推流状态"""
|
||
try:
|
||
from task_stream_manager import task_stream_manager
|
||
|
||
# 获取任务状态
|
||
task_status = task_manager.get_task_status(task_id)
|
||
if not task_status:
|
||
return jsonify({
|
||
"status": "error",
|
||
"message": f"任务不存在: {task_id}"
|
||
}), 404
|
||
|
||
# 获取推流信息
|
||
stream_info = task_stream_manager.get_all_task_streams_info().get(task_id, {})
|
||
|
||
# 合并信息
|
||
result = {
|
||
"task_id": task_id,
|
||
"task_status": task_status['status'],
|
||
"stream_enabled": task_status['config'].get('enable_push', False),
|
||
"stream_info": stream_info,
|
||
"push_url": task_status['config'].get('push_url', '')
|
||
}
|
||
|
||
return jsonify({
|
||
"status": "success",
|
||
"data": result
|
||
})
|
||
except Exception as e:
|
||
logger.error(f"获取任务推流状态失败: {str(e)}")
|
||
return jsonify({
|
||
"status": "error",
|
||
"message": f"获取失败: {str(e)}"
|
||
}), 500
|
||
|
||
|
||
@app.route('/api/system/streams/info', methods=['GET'])
|
||
def get_all_streams_info():
|
||
"""获取所有任务推流信息"""
|
||
try:
|
||
from task_stream_manager import task_stream_manager
|
||
|
||
streams_info = task_stream_manager.get_all_task_streams_info()
|
||
|
||
return jsonify({
|
||
"status": "success",
|
||
"data": {
|
||
"total_streams": len(streams_info),
|
||
"active_streams": sum(1 for info in streams_info.values() if info.get('running', False)),
|
||
"streams": streams_info
|
||
}
|
||
})
|
||
except Exception as e:
|
||
logger.error(f"获取所有推流信息失败: {str(e)}")
|
||
return jsonify({
|
||
"status": "error",
|
||
"message": f"获取失败: {str(e)}"
|
||
}), 500
|
||
|
||
|
||
# 在适当位置初始化上传管理器
|
||
upload_manager = None
|
||
|
||
|
||
def init_upload_manager(config):
|
||
"""初始化上传管理器"""
|
||
global upload_manager
|
||
if not upload_manager:
|
||
upload_manager = get_upload_manager(config['upload'])
|
||
return upload_manager
|
||
|
||
|
||
# 添加上传相关路由
|
||
@app.route('/api/models/upload/start', methods=['POST'])
|
||
def start_model_upload():
|
||
"""开始模型文件上传"""
|
||
try:
|
||
data = request.json
|
||
if not data:
|
||
return jsonify({"status": "error", "message": "请求数据不能为空"}), 400
|
||
|
||
filename = data.get('filename')
|
||
total_size = data.get('total_size')
|
||
encryption_key = data.get('encryption_key') # 可选，但建议提供
|
||
|
||
if not all([filename, total_size]):
|
||
return jsonify({"status": "error", "message": "缺少必要参数"}), 400
|
||
|
||
# 初始化上传管理器
|
||
config = get_default_config()
|
||
upload_mgr = init_upload_manager(config)
|
||
|
||
# 获取分片大小配置
|
||
chunk_size = config['upload'].get('chunk_size', 5 * 1024 * 1024)
|
||
|
||
# 创建上传会话
|
||
result = upload_mgr.create_upload_session(
|
||
filename=filename,
|
||
total_size=total_size,
|
||
chunk_size=chunk_size,
|
||
encryption_key=encryption_key
|
||
)
|
||
|
||
if result['success']:
|
||
return jsonify({
|
||
"status": "success",
|
||
"message": "上传会话创建成功",
|
||
"data": {
|
||
"session_id": result['session_id'],
|
||
"total_chunks": result['total_chunks'],
|
||
"chunk_size": chunk_size
|
||
}
|
||
})
|
||
else:
|
||
return jsonify({
|
||
"status": "error",
|
||
"message": f"创建上传会话失败: {result.get('error', '未知错误')}"
|
||
}), 500
|
||
|
||
except Exception as e:
|
||
logger.error(f"开始上传失败: {str(e)}")
|
||
return jsonify({
|
||
"status": "error",
|
||
"message": f"开始上传失败: {str(e)}"
|
||
}), 500
|
||
|
||
|
||
@app.route('/api/models/upload/chunk', methods=['POST'])
|
||
def upload_model_chunk():
|
||
"""上传模型文件分片"""
|
||
try:
|
||
# 获取表单数据
|
||
session_id = request.form.get('session_id')
|
||
chunk_index = int(request.form.get('chunk_index', 0))
|
||
|
||
if not session_id:
|
||
return jsonify({"status": "error", "message": "缺少session_id"}), 400
|
||
|
||
# 获取文件数据
|
||
if 'chunk' not in request.files:
|
||
return jsonify({"status": "error", "message": "未找到文件分片"}), 400
|
||
|
||
chunk_file = request.files['chunk']
|
||
chunk_data = chunk_file.read()
|
||
|
||
# 获取上传管理器
|
||
config = get_default_config()
|
||
upload_mgr = init_upload_manager(config)
|
||
|
||
# 上传分片
|
||
result = upload_mgr.upload_chunk(session_id, chunk_index, chunk_data)
|
||
|
||
if result['success']:
|
||
return jsonify({
|
||
"status": "success",
|
||
"message": "分片上传成功",
|
||
"data": {
|
||
"progress": result['progress'],
|
||
"received_chunks": result['received_chunks'],
|
||
"total_chunks": result['total_chunks'],
|
||
"chunk_index": chunk_index
|
||
}
|
||
})
|
||
else:
|
||
return jsonify({
|
||
"status": "error",
|
||
"message": f"分片上传失败: {result.get('error', '未知错误')}"
|
||
}), 500
|
||
|
||
except Exception as e:
|
||
logger.error(f"上传分片失败: {str(e)}")
|
||
return jsonify({
|
||
"status": "error",
|
||
"message": f"上传分片失败: {str(e)}"
|
||
}), 500
|
||
|
||
|
||
@app.route('/api/models/upload/status/<session_id>', methods=['GET'])
|
||
def get_upload_status(session_id):
|
||
"""获取上传状态"""
|
||
try:
|
||
config = get_default_config()
|
||
upload_mgr = init_upload_manager(config)
|
||
|
||
result = upload_mgr.get_upload_status(session_id)
|
||
|
||
if result['success']:
|
||
return jsonify({
|
||
"status": "success",
|
||
"data": result['data']
|
||
})
|
||
else:
|
||
return jsonify({
|
||
"status": "error",
|
||
"message": result.get('error', '获取状态失败')
|
||
}), 404
|
||
|
||
except Exception as e:
|
||
logger.error(f"获取上传状态失败: {str(e)}")
|
||
return jsonify({
|
||
"status": "error",
|
||
"message": f"获取上传状态失败: {str(e)}"
|
||
}), 500
|
||
|
||
|
||
@app.route('/api/models/upload/cancel/<session_id>', methods=['POST'])
|
||
def cancel_upload(session_id):
|
||
"""取消上传"""
|
||
try:
|
||
# 实现取消逻辑
|
||
config = get_default_config()
|
||
upload_mgr = init_upload_manager(config)
|
||
|
||
# 这里需要在上传管理器中添加取消功能
|
||
# upload_mgr.cancel_upload(session_id)
|
||
|
||
return jsonify({
|
||
"status": "success",
|
||
"message": f"上传已取消: {session_id}"
|
||
})
|
||
|
||
except Exception as e:
|
||
logger.error(f"取消上传失败: {str(e)}")
|
||
return jsonify({
|
||
"status": "error",
|
||
"message": f"取消上传失败: {str(e)}"
|
||
}), 500
|
||
|
||
|
||
@app.route('/api/models/list', methods=['GET'])
|
||
def list_encrypted_models():
|
||
"""列出所有已加密的模型文件"""
|
||
try:
|
||
config = get_default_config()
|
||
encrypted_dir = config['upload']['encrypted_models_dir']
|
||
|
||
if not os.path.exists(encrypted_dir):
|
||
return jsonify({
|
||
"status": "success",
|
||
"data": {
|
||
"models": [],
|
||
"total": 0
|
||
}
|
||
})
|
||
|
||
models = []
|
||
for filename in os.listdir(encrypted_dir):
|
||
if filename.endswith('.enc'):
|
||
filepath = os.path.join(encrypted_dir, filename)
|
||
stats = os.stat(filepath)
|
||
models.append({
|
||
'filename': filename,
|
||
'path': filepath,
|
||
'size': stats.st_size,
|
||
'modified': stats.st_mtime,
|
||
'encrypted': True
|
||
})
|
||
|
||
return jsonify({
|
||
"status": "success",
|
||
"data": {
|
||
"models": models,
|
||
"total": len(models),
|
||
"encrypted_dir": encrypted_dir
|
||
}
|
||
})
|
||
|
||
except Exception as e:
|
||
logger.error(f"列出模型失败: {str(e)}")
|
||
return jsonify({
|
||
"status": "error",
|
||
"message": f"列出模型失败: {str(e)}"
|
||
}), 500
|
||
|
||
|
||
# 在 server.py 中添加上传页面路由
|
||
@app.route('/model_upload')
|
||
def model_upload_page():
|
||
"""模型上传页面"""
|
||
return render_template("model_upload.html")
|
||
|
||
|
||
# server.py - 添加以下路由
|
||
|
||
from mandatory_model_crypto import ModelEncryptionService, validate_models_before_task, verify_single_model_api
|
||
|
||
|
||
@app.route('/api/models/process/start', methods=['POST'])
|
||
def start_model_processing():
|
||
"""开始模型处理流程：生成密钥 -> 上传模型"""
|
||
try:
|
||
data = request.json
|
||
|
||
# 选项1：客户端提供自己的密钥
|
||
client_key = data.get('encryption_key')
|
||
|
||
# 选项2：服务器生成密钥
|
||
generate_new = data.get('generate_key', False)
|
||
|
||
response_data = {}
|
||
|
||
if generate_new and not client_key:
|
||
# 服务器生成新密钥
|
||
key_info = ModelEncryptionService.generate_secure_key()
|
||
response_data['key_info'] = {
|
||
'key': key_info['key'],
|
||
'key_hash': key_info['key_hash'],
|
||
'short_hash': key_info['short_hash'],
|
||
'generated_by': 'server'
|
||
}
|
||
|
||
elif client_key:
|
||
# 验证客户端提供的密钥
|
||
key_valid, key_msg = ModelEncryptionService.validate_key_strength(client_key)
|
||
if not key_valid:
|
||
return jsonify({
|
||
"status": "error",
|
||
"message": f"密钥强度不足: {key_msg}"
|
||
}), 400
|
||
|
||
response_data['key_info'] = {
|
||
'key': client_key,
|
||
'key_hash': hashlib.sha256(client_key.encode()).hexdigest(),
|
||
'short_hash': hashlib.sha256(client_key.encode()).hexdigest()[:16],
|
||
'generated_by': 'client'
|
||
}
|
||
else:
|
||
return jsonify({
|
||
"status": "error",
|
||
"message": "请提供加密密钥或选择生成新密钥"
|
||
}), 400
|
||
|
||
# 生成上传令牌
|
||
upload_token = secrets.token_urlsafe(32)
|
||
response_data['upload_token'] = upload_token
|
||
response_data['token_expires'] = time.time() + 3600 # 1小时有效期
|
||
|
||
# 存储上传会话（简化版，生产环境应使用数据库）
|
||
upload_sessions = gd.get_or_create_dict('upload_sessions')
|
||
upload_sessions[upload_token] = {
|
||
'key_info': response_data['key_info'],
|
||
'created_at': time.time(),
|
||
'status': 'pending'
|
||
}
|
||
|
||
return jsonify({
|
||
"status": "success",
|
||
"message": "模型处理流程已启动",
|
||
"data": response_data
|
||
})
|
||
|
||
except Exception as e:
|
||
logger.error(f"启动模型处理流程失败: {str(e)}")
|
||
return jsonify({
|
||
"status": "error",
|
||
"message": f"启动流程失败: {str(e)}"
|
||
}), 500
|
||
|
||
|
||
@app.route('/api/models/process/verify_key', methods=['POST'])
|
||
def verify_encryption_key():
|
||
"""验证加密密钥"""
|
||
try:
|
||
data = request.json
|
||
encryption_key = data.get('encryption_key')
|
||
model_path = data.get('model_path') # 可选，如果有具体模型
|
||
|
||
if not encryption_key:
|
||
return jsonify({
|
||
"status": "error",
|
||
"message": "请提供加密密钥"
|
||
}), 400
|
||
|
||
# 验证密钥强度
|
||
key_valid, key_msg = ModelEncryptionService.validate_key_strength(encryption_key)
|
||
|
||
if not key_valid:
|
||
return jsonify({
|
||
"status": "error",
|
||
"message": key_msg,
|
||
"data": {
|
||
"valid": False,
|
||
"strength": "weak"
|
||
}
|
||
}), 400
|
||
|
||
response_data = {
|
||
"valid": True,
|
||
"strength": "strong",
|
||
"key_hash": hashlib.sha256(encryption_key.encode()).hexdigest()[:16]
|
||
}
|
||
|
||
# 如果有具体模型，尝试解密验证
|
||
if model_path:
|
||
verify_result = verify_single_model_api(model_path, encryption_key)
|
||
response_data['model_verification'] = verify_result
|
||
|
||
return jsonify({
|
||
"status": "success",
|
||
"message": "密钥验证成功",
|
||
"data": response_data
|
||
})
|
||
|
||
except Exception as e:
|
||
logger.error(f"验证密钥失败: {str(e)}")
|
||
return jsonify({
|
||
"status": "error",
|
||
"message": f"验证失败: {str(e)}"
|
||
}), 500
|
||
|
||
|
||
@app.route('/api/models/process/validate_task', methods=['POST'])
|
||
def validate_task_models():
|
||
"""创建任务前的模型验证"""
|
||
try:
|
||
data = request.json
|
||
|
||
if not data:
|
||
return jsonify({
|
||
"status": "error",
|
||
"message": "请求数据不能为空"
|
||
}), 400
|
||
|
||
# 提取任务配置
|
||
task_config = {
|
||
'models': data.get('models', [])
|
||
}
|
||
|
||
# 验证所有模型
|
||
validation_result = validate_models_before_task(task_config)
|
||
|
||
if validation_result['success']:
|
||
return jsonify({
|
||
"status": "success",
|
||
"message": f"模型验证通过 ({validation_result['valid_models']}/{validation_result['total_models']})",
|
||
"data": validation_result
|
||
})
|
||
else:
|
||
return jsonify({
|
||
"status": "error",
|
||
"message": f"模型验证失败: {validation_result.get('error', '未知错误')}",
|
||
"data": validation_result
|
||
}), 400
|
||
|
||
except Exception as e:
|
||
logger.error(f"验证任务模型失败: {str(e)}")
|
||
return jsonify({
|
||
"status": "error",
|
||
"message": f"验证失败: {str(e)}"
|
||
}), 500
|
||
|
||
|
||
@app.route('/api/models/encrypted/list_available', methods=['GET'])
|
||
def list_available_encrypted_models():
|
||
"""列出可用的加密模型（用于任务创建选择）"""
|
||
try:
|
||
config = get_default_config()
|
||
encrypted_dir = config['upload']['encrypted_models_dir']
|
||
|
||
if not os.path.exists(encrypted_dir):
|
||
return jsonify({
|
||
"status": "success",
|
||
"data": {
|
||
"models": [],
|
||
"total": 0
|
||
}
|
||
})
|
||
|
||
models = []
|
||
for filename in os.listdir(encrypted_dir):
|
||
if filename.endswith('.enc'):
|
||
filepath = os.path.join(encrypted_dir, filename)
|
||
stats = os.stat(filepath)
|
||
|
||
# 尝试读取模型基本信息（不验证密钥）
|
||
try:
|
||
with open(filepath, 'rb') as f:
|
||
encrypted_data = pickle.load(f)
|
||
|
||
model_info = {
|
||
'filename': filename,
|
||
'path': f"encrypted_models/{filename}", # 相对路径
|
||
'size': stats.st_size,
|
||
'modified': stats.st_mtime,
|
||
'encrypted': True,
|
||
'model_hash': encrypted_data.get('model_hash', '')[:16] if isinstance(encrypted_data,
|
||
dict) else '',
|
||
'version': encrypted_data.get('version', 'unknown') if isinstance(encrypted_data, dict) else '',
|
||
'original_size': encrypted_data.get('original_size', 0) if isinstance(encrypted_data,
|
||
dict) else 0
|
||
}
|
||
|
||
models.append(model_info)
|
||
|
||
except Exception as e:
|
||
logger.warning(f"读取模型信息失败 {filename}: {str(e)}")
|
||
models.append({
|
||
'filename': filename,
|
||
'path': f"encrypted_models/{filename}",
|
||
'size': stats.st_size,
|
||
'modified': stats.st_mtime,
|
||
'encrypted': True,
|
||
'error': '无法读取模型信息'
|
||
})
|
||
|
||
# 按修改时间排序
|
||
models.sort(key=lambda x: x['modified'], reverse=True)
|
||
|
||
return jsonify({
|
||
"status": "success",
|
||
"data": {
|
||
"models": models,
|
||
"total": len(models),
|
||
"directory": encrypted_dir
|
||
}
|
||
})
|
||
|
||
except Exception as e:
|
||
logger.error(f"列出可用模型失败: {str(e)}")
|
||
return jsonify({
|
||
"status": "error",
|
||
"message": f"列出模型失败: {str(e)}"
|
||
}), 500
|
||
|
||
|
||
@app.route('/api/models/process/test_decrypt', methods=['POST'])
|
||
def test_model_decryption():
|
||
"""测试模型解密（不实际加载YOLO模型）"""
|
||
try:
|
||
data = request.json
|
||
model_path = data.get('model_path')
|
||
encryption_key = data.get('encryption_key')
|
||
|
||
if not all([model_path, encryption_key]):
|
||
return jsonify({
|
||
"status": "error",
|
||
"message": "请提供模型路径和加密密钥"
|
||
}), 400
|
||
|
||
# 构建完整路径
|
||
config = get_default_config()
|
||
encrypted_dir = config['upload']['encrypted_models_dir']
|
||
|
||
if not os.path.isabs(model_path):
|
||
model_filename = os.path.basename(model_path)
|
||
full_model_path = os.path.join(encrypted_dir, model_filename)
|
||
else:
|
||
full_model_path = model_path
|
||
|
||
# 检查文件是否存在
|
||
if not os.path.exists(full_model_path):
|
||
return jsonify({
|
||
"status": "error",
|
||
"message": f"模型文件不存在: {full_model_path}"
|
||
}), 404
|
||
|
||
# 测试解密
|
||
from mandatory_model_crypto import MandatoryModelValidator
|
||
validator = MandatoryModelValidator()
|
||
|
||
start_time = time.time()
|
||
decrypt_result = validator.decrypt_and_verify(full_model_path, encryption_key)
|
||
elapsed_time = time.time() - start_time
|
||
|
||
if decrypt_result['success']:
|
||
return jsonify({
|
||
"status": "success",
|
||
"message": "解密测试成功",
|
||
"data": {
|
||
'success': True,
|
||
'model_hash': decrypt_result.get('model_hash', '')[:16],
|
||
'model_size': decrypt_result.get('original_size', 0),
|
||
'decryption_time': elapsed_time,
|
||
'file_path': full_model_path
|
||
}
|
||
})
|
||
else:
|
||
return jsonify({
|
||
"status": "error",
|
||
"message": f"解密测试失败: {decrypt_result.get('error', '未知错误')}",
|
||
"data": {
|
||
'success': False,
|
||
'error': decrypt_result.get('error', '未知错误'),
|
||
'decryption_time': elapsed_time
|
||
}
|
||
}), 400
|
||
|
||
except Exception as e:
|
||
logger.error(f"测试解密失败: {str(e)}")
|
||
return jsonify({
|
||
"status": "error",
|
||
"message": f"测试失败: {str(e)}"
|
||
}), 500
|
||
|
||
|
||
# server.py - 添加任务创建页面路由
|
||
@app.route('/task_create')
|
||
def task_create_page():
|
||
"""任务创建页面"""
|
||
return render_template("task_create.html")
|
||
|
||
|
||
# server.py - 添加资源检查接口
|
||
|
||
@app.route('/api/system/check_resources', methods=['GET'])
|
||
def check_system_resources():
|
||
"""检查系统资源是否足够创建新任务"""
|
||
try:
|
||
import psutil
|
||
import torch
|
||
|
||
# 获取系统资源
|
||
cpu_percent = psutil.cpu_percent(interval=0.1)
|
||
memory_info = psutil.virtual_memory()
|
||
memory_percent = memory_info.percent
|
||
|
||
# GPU信息
|
||
gpu_info = []
|
||
gpu_available = torch.cuda.is_available()
|
||
if gpu_available:
|
||
for i in range(torch.cuda.device_count()):
|
||
gpu_memory_used = torch.cuda.memory_allocated(i) / 1024 ** 2 # MB
|
||
gpu_memory_total = torch.cuda.get_device_properties(i).total_memory / 1024 ** 2 # MB
|
||
gpu_memory_percent = (gpu_memory_used / gpu_memory_total) * 100 if gpu_memory_total > 0 else 0
|
||
|
||
gpu_info.append({
|
||
'id': i,
|
||
'name': torch.cuda.get_device_name(i),
|
||
'memory_used': gpu_memory_used,
|
||
'memory_total': gpu_memory_total,
|
||
'memory_percent': gpu_memory_percent
|
||
})
|
||
|
||
# 获取当前任务信息
|
||
active_tasks = task_manager.get_active_tasks_count()
|
||
max_tasks = task_manager.get_current_max_tasks()
|
||
|
||
# 资源阈值
|
||
cpu_threshold = 80 # CPU使用率阈值
|
||
memory_threshold = 85 # 内存使用率阈值
|
||
gpu_memory_threshold = 90 # GPU内存使用率阈值
|
||
|
||
# 检查资源状态
|
||
resources_ok = True
|
||
warnings = []
|
||
|
||
if cpu_percent > cpu_threshold:
|
||
resources_ok = False
|
||
warnings.append(f"CPU使用率过高: {cpu_percent:.1f}% > {cpu_threshold}%")
|
||
|
||
if memory_percent > memory_threshold:
|
||
resources_ok = False
|
||
warnings.append(f"内存使用率过高: {memory_percent:.1f}% > {memory_threshold}%")
|
||
|
||
if gpu_available and gpu_info:
|
||
max_gpu_memory = max(gpu['memory_percent'] for gpu in gpu_info)
|
||
if max_gpu_memory > gpu_memory_threshold:
|
||
warnings.append(f"GPU内存使用率过高: {max_gpu_memory:.1f}% > {gpu_memory_threshold}%")
|
||
# GPU内存高不是致命错误，只是警告
|
||
|
||
if active_tasks >= max_tasks:
|
||
resources_ok = False
|
||
warnings.append(f"任务数达到上限: {active_tasks}/{max_tasks}")
|
||
|
||
return jsonify({
|
||
"status": "success",
|
||
"data": {
|
||
"resources_available": resources_ok,
|
||
"active_tasks": active_tasks,
|
||
"max_tasks": max_tasks,
|
||
"slots_available": max(0, max_tasks - active_tasks),
|
||
"cpu_percent": cpu_percent,
|
||
"memory_percent": memory_percent,
|
||
"memory_used": memory_info.used / 1024 ** 2, # MB
|
||
"memory_total": memory_info.total / 1024 ** 2, # MB
|
||
"gpu_available": gpu_available,
|
||
"gpu_info": gpu_info,
|
||
"warnings": warnings,
|
||
"thresholds": {
|
||
"cpu": cpu_threshold,
|
||
"memory": memory_threshold,
|
||
"gpu_memory": gpu_memory_threshold
|
||
},
|
||
"timestamp": time.time()
|
||
}
|
||
})
|
||
|
||
except Exception as e:
|
||
logger.error(f"检查系统资源失败: {str(e)}")
|
||
return jsonify({
|
||
"status": "error",
|
||
"message": f"检查资源失败: {str(e)}"
|
||
}), 500
|
||
|
||
|
||
# 初始化函数，可在主程序中调用
|
||
def init_app():
|
||
"""初始化应用程序"""
|
||
with app.app_context():
|
||
gd.set_value('task_manager', task_manager)
|
||
logger.info("任务管理器初始化完成")
|
||
return app
|
||
|
||
# ======================= 公共服务器功能 =======================
|
||
|
||
# 内部系统状态
|
||
internal_system_status = {
|
||
"is_connected": False,
|
||
"devices": {}
|
||
}
|
||
|
||
# 连接管理
|
||
connected_clients = {}
|
||
|
||
|
||
# WebSocket事件 - 公共服务器
|
||
@socketio.on('connect', namespace='/ws/internal')
|
||
def handle_internal_connect():
|
||
"""内部系统连接"""
|
||
client_id = "internal_system"
|
||
connected_clients[client_id] = request.sid
|
||
internal_system_status["is_connected"] = True
|
||
pub_logger.info(f"Internal system connected: {request.sid}")
|
||
|
||
|
||
@socketio.on('disconnect', namespace='/ws/internal')
|
||
def handle_internal_disconnect():
|
||
"""内部系统断开连接"""
|
||
client_id = "internal_system"
|
||
if client_id in connected_clients:
|
||
del connected_clients[client_id]
|
||
internal_system_status["is_connected"] = False
|
||
pub_logger.info("Internal system disconnected")
|
||
|
||
|
||
@socketio.on('message', namespace='/ws/internal')
|
||
def handle_internal_message(data):
|
||
"""处理内部系统消息"""
|
||
try:
|
||
message = json.loads(data)
|
||
if message.get("type") == "status_update":
|
||
# 更新设备状态
|
||
device_id = message.get("device_id")
|
||
if device_id:
|
||
internal_system_status["devices"][device_id] = message
|
||
# 广播状态更新给客户端
|
||
socketio.emit('message', json.dumps(message), namespace='/ws/client')
|
||
pub_logger.info(f"Status updated for device {device_id}: {message}")
|
||
elif message.get("type") == "command_response":
|
||
# 转发命令响应给客户端
|
||
socketio.emit('message', json.dumps(message), namespace='/ws/client')
|
||
pub_logger.info(f"Command response: {message}")
|
||
except json.JSONDecodeError:
|
||
pub_logger.error(f"Invalid JSON from internal system: {data}")
|
||
|
||
|
||
@socketio.on('connect', namespace='/ws/client')
|
||
def handle_client_connect():
|
||
"""客户端连接"""
|
||
client_id = request.args.get('client_id', request.sid)
|
||
connected_clients[client_id] = request.sid
|
||
pub_logger.info(f"Client connected: {client_id} ({request.sid})")
|
||
|
||
# 发送当前状态给新连接的客户端
|
||
for device_id, status in internal_system_status["devices"].items():
|
||
socketio.emit('message', json.dumps(status), namespace='/ws/client', room=request.sid)
|
||
|
||
|
||
@socketio.on('disconnect', namespace='/ws/client')
|
||
def handle_client_disconnect():
|
||
"""客户端断开连接"""
|
||
# 找到对应的客户端ID
|
||
client_id = None
|
||
for id, sid in connected_clients.items():
|
||
if sid == request.sid:
|
||
client_id = id
|
||
break
|
||
if client_id:
|
||
del connected_clients[client_id]
|
||
pub_logger.info(f"Client disconnected: {client_id}")
|
||
|
||
|
||
@socketio.on('message', namespace='/ws/client')
|
||
def handle_client_message(data):
|
||
"""处理客户端消息"""
|
||
try:
|
||
command = json.loads(data)
|
||
# 转发命令给内部系统
|
||
if internal_system_status["is_connected"] and "internal_system" in connected_clients:
|
||
socketio.emit('message', data, namespace='/ws/internal', room=connected_clients["internal_system"])
|
||
pub_logger.info(f"Forwarded command from client: {command}")
|
||
else:
|
||
# 内部系统未连接，返回错误
|
||
error_response = {
|
||
"type": "command_response",
|
||
"status": "error",
|
||
"message": "Internal system not connected",
|
||
"command": command
|
||
}
|
||
socketio.emit('message', json.dumps(error_response), namespace='/ws/client', room=request.sid)
|
||
except json.JSONDecodeError:
|
||
pub_logger.error(f"Invalid JSON from client: {data}")
|
||
|
||
|
||
# API路由 - 公共服务器
|
||
@app.route('/api/devices')
|
||
def get_devices():
|
||
"""获取设备列表"""
|
||
return jsonify({
|
||
"code": 200,
|
||
"msg": "查询成功",
|
||
"data": list(internal_system_status["devices"].values())
|
||
})
|
||
|
||
|
||
@app.route('/api/command', methods=['POST'])
|
||
def send_command():
|
||
"""发送命令"""
|
||
if not internal_system_status["is_connected"]:
|
||
return jsonify({
|
||
"code": 503,
|
||
"msg": "内部系统未连接",
|
||
"data": None
|
||
})
|
||
|
||
try:
|
||
command = request.json
|
||
# 发送命令给内部系统
|
||
if "internal_system" in connected_clients:
|
||
socketio.emit('message', json.dumps(command), namespace='/ws/internal', room=connected_clients["internal_system"])
|
||
pub_logger.info(f"Command sent: {command}")
|
||
return jsonify({
|
||
"code": 200,
|
||
"msg": "命令已发送",
|
||
"data": command
|
||
})
|
||
else:
|
||
return jsonify({
|
||
"code": 503,
|
||
"msg": "内部系统未连接",
|
||
"data": None
|
||
})
|
||
except Exception as e:
|
||
pub_logger.error(f"Error sending command: {str(e)}")
|
||
return jsonify({
|
||
"code": 500,
|
||
"msg": f"发送命令失败: {str(e)}",
|
||
"data": None
|
||
}) |