forked from JointCloud/SchedulingSimulator
279 lines
12 KiB
Python
279 lines
12 KiB
Python
import time
|
||
import uuid
|
||
from datetime import datetime
|
||
from typing import Dict, List, Optional
|
||
import requests
|
||
import sys
|
||
import os
|
||
import importlib
|
||
|
||
from task.ai_task_submit import submit_ai_task
|
||
from config.db_manager import load_config_from_db, update_task_all_fields
|
||
|
||
# 添加TaskSubmissionScript到Python路径
|
||
sys.path.append(os.path.join(os.path.dirname(__file__), '..', 'TaskSubmissionScript'))
|
||
|
||
from task.hpc_task_submit import HPCTaskSubmitter
|
||
from config.database import load_db_config
|
||
from config.db_manager import update_cluster_resources
|
||
from config.config import logger, TASK_STATUS, task_map, task_map_lock
|
||
from config.cache_manager import get_cached_config
|
||
|
||
from task.cloud_task_submit import submit_cloud_task
|
||
|
||
def generate_task_templates() -> List[Dict]:
|
||
"""从数据库加载并返回任务模板"""
|
||
# 使用缓存的配置而不是每次都从数据库加载
|
||
config = get_cached_config()
|
||
return config.get("task_templates", [])
|
||
|
||
|
||
def load_tasks_to_queue(templates: List[Dict]) -> None:
|
||
with task_map_lock:
|
||
task_map.clear()
|
||
for idx, template in enumerate(templates):
|
||
try:
|
||
# 检查必填字段
|
||
required_fields = ["type", "task_name", "resource"]
|
||
missing_fields = [f for f in required_fields if f not in template]
|
||
if missing_fields:
|
||
logger.warning(f"跳过无效任务模板（索引{idx}）：缺少字段 {missing_fields}")
|
||
continue
|
||
|
||
# 生成任务名称（使用数据库中的task_name）
|
||
task_name = template["task_name"] or f"task-{int(time.time())}"
|
||
task_type = template["type"]
|
||
resource = template["resource"]
|
||
strategy = template.get("strategy")
|
||
|
||
# 数据集ID设为可选
|
||
dataset_id = template.get("dataset_id", "")
|
||
code_id = template.get("code_id", "")
|
||
|
||
# 创建任务实例
|
||
task = {
|
||
"target_id": str(uuid.uuid4()), # 任务唯一ID
|
||
"task_name": task_name,
|
||
"package_name": f"{task_name.lower()}-pkg" if task_name else f"task-{int(time.time())}-pkg", # 文件夹名称
|
||
"type": task_type,
|
||
"status": TASK_STATUS["SUBMITTED"],
|
||
"submit_time": time.strftime("%Y-%m-%d %H:%M:%S"),
|
||
"success_time": None,
|
||
"third_party_task_id": "",
|
||
"dataset_id": dataset_id, # 数据集ID
|
||
"code_id":code_id,
|
||
"error_msg": "",
|
||
"fail_count": 0,
|
||
"max_fail_threshold": 3,
|
||
"cluster_id": "",
|
||
"strategy": strategy,
|
||
"resource": resource
|
||
}
|
||
|
||
# 根据任务类型设置特定字段
|
||
if task_type == "Ai":
|
||
# AI任务
|
||
task["dataset_id"] = dataset_id
|
||
task["resource"] = resource or {"CPU": 1, "MEMORY": 4, "NPU": 0}
|
||
# AI任务需要file_location字段，但数据库中没有直接提供，需要通过其他方式获取
|
||
|
||
elif task_type == "Hpc":
|
||
task["partition"] = "default" # 默认分区
|
||
task["ntasks"] = "1"
|
||
task["nodes"] = "1"
|
||
task["packageName"] = f"{task_name.lower()}-hpc" if task_name else f"task-{int(time.time())}-hpc"
|
||
task["hpcFile"] = "" # 数据库中没有提供
|
||
task["dataset_id"] = dataset_id
|
||
|
||
elif task_type == "Cloud":
|
||
task["replicas"] = 1
|
||
task["strategy"] = strategy or "resource" # 调度策略
|
||
task["cloud_strategy"] = "replication" # 部署策略
|
||
task["adapter_name"] = "" # 适配器名称
|
||
task["resource"] = resource or {"CPU": 0.1, "MEMORY": 6, "port": 80}
|
||
task["port"] = task["resource"].get("port", 80)
|
||
task["dataset_id"] = dataset_id
|
||
|
||
# 将任务添加到全局任务映射中
|
||
task_map[task["target_id"]] = task
|
||
|
||
logger.info(
|
||
f"任务入队 | task_name: {task_name} | 任务类型: {task_type} | "
|
||
f"{'含数据集' if dataset_id else '无数据集'}")
|
||
|
||
except Exception as e:
|
||
logger.error(f"处理任务模板时发生异常（索引{idx}）: {e}", exc_info=True)
|
||
continue
|
||
|
||
logger.info(f"任务队列加载完成 | 总任务数: {len(task_map)}")
|
||
|
||
|
||
def get_task_submit_function(module_name):
|
||
"""动态导入任务提交函数"""
|
||
try:
|
||
module = importlib.import_module(module_name)
|
||
return getattr(module, 'main')
|
||
except (ImportError, AttributeError) as e:
|
||
logger.error(f"无法导入模块 {module_name} 或其 main 函数: {e}")
|
||
return None
|
||
|
||
def submit_single_task(task: Dict,token) -> bool:
|
||
"""提交单个任务到集群"""
|
||
|
||
# 获取选择的集群ID
|
||
global resource_changes
|
||
cluster_id = task.get("cluster_id")
|
||
if not cluster_id:
|
||
logger.error(f"[{task['task_name']}] 提交失败：未指定集群ID")
|
||
return False
|
||
|
||
task_type = task["type"]
|
||
try:
|
||
if task_type == "Ai":
|
||
third_party_task_id = submit_ai_task(task,token)
|
||
elif task_type == "Hpc":
|
||
task_name = task["task_name"]
|
||
partition = task["partition"]
|
||
ntasks = task["ntasks"]
|
||
nodes = task["nodes"]
|
||
packageName = task["package_name"]
|
||
hpcFile = task.get("hpc_file", "") # 修正字段名，使用hpcFile而非hpc_file
|
||
|
||
# 确定HPC任务类型（根据任务名称推断）
|
||
hpc_type = "bwa" # 默认类型
|
||
if "_bwa_" in task_name or "_bwa2_" in task_name:
|
||
hpc_type = "bwa"
|
||
elif "_hashcat_" in task_name:
|
||
hpc_type = "hashcat"
|
||
elif "_lammps_" in task_name:
|
||
hpc_type = "lammps"
|
||
|
||
# 创建HPC任务提交器实例
|
||
hpc_submitter = HPCTaskSubmitter()
|
||
|
||
# 提交HPC任务
|
||
third_party_task_id = hpc_submitter.submit_single_task(
|
||
task_name, hpc_type, partition, ntasks, nodes, packageName, hpcFile
|
||
)
|
||
|
||
# 如果返回的是字典格式，提取jobSetID
|
||
if isinstance(third_party_task_id, dict) and "data" in third_party_task_id:
|
||
third_party_task_id = third_party_task_id["data"].get("jobSetID")
|
||
|
||
elif task_type == "Cloud":
|
||
# 提取云任务所需的参数
|
||
replicas = task.get("replicas")
|
||
task_res = task.get("resource")
|
||
cpu = task_res["CPU"]
|
||
memory = task_res["MEMORY"]
|
||
port = task.get("port")
|
||
task_name = task["task_name"]
|
||
cluster_id = task.get("cluster_id")
|
||
|
||
# 传递所有必要的参数给云任务提交函数
|
||
third_party_task_id = submit_cloud_task(
|
||
task_name,
|
||
cluster_id,
|
||
replicas=replicas,
|
||
cpu=cpu,
|
||
memory=memory,
|
||
port=port,
|
||
|
||
)
|
||
else:
|
||
raise ValueError(f"未知任务类型: {task_type}")
|
||
|
||
# 构造资源变化字典
|
||
task_resource = task.get("resource", {})
|
||
resource_changes = {}
|
||
|
||
# 计算资源变化量（负值表示扣减）
|
||
for res_type in ["CPU", "MEMORY"]:
|
||
required = task_resource.get(res_type, 0)
|
||
if required > 0:
|
||
resource_changes[res_type] = -required
|
||
|
||
# 处理加速器资源
|
||
for acc_type in [key for key in task_resource if key not in ["CPU", "MEMORY"]]:
|
||
required_count = task_resource.get(acc_type, 0)
|
||
if required_count > 0:
|
||
resource_changes[acc_type] = -required_count
|
||
|
||
if third_party_task_id:
|
||
# 使用批量操作更新数据库
|
||
if update_cluster_resources(cluster_id, resource_changes):
|
||
logger.info(f"成功从数据库扣减集群 {cluster_id} 的资源: {resource_changes}")
|
||
else:
|
||
logger.error(f"从数据库扣减集群 {cluster_id} 的资源失败: {resource_changes}")
|
||
|
||
|
||
# 更新任务状态为成功
|
||
logger.info(f"[{task['task_name']}] 任务提交至集群 {cluster_id} 成功 | 云际任务ID: {third_party_task_id}")
|
||
# 任务提交成功，更新数据库中的状态为SUBMITTING
|
||
if update_task_all_fields(task['task_name'], status=TASK_STATUS['SUBMITTING'],third_party_task_id=third_party_task_id,commit_time= datetime.now()):
|
||
logger.info(f"成功更新任务 {task['task_name']} 在数据库中的状态为 {TASK_STATUS['SUBMITTING']}")
|
||
else:
|
||
logger.error(f"更新任务 {task['task_name']} 在数据库中的状态失败")
|
||
return False
|
||
return True
|
||
else:
|
||
# 如果任务提交失败且资源已更新，需要恢复资源
|
||
# 这里使用原子操作避免多次数据库访问
|
||
if resource_changes:
|
||
recovery_changes = {res: -change for res, change in resource_changes.items()}
|
||
update_cluster_resources(cluster_id, recovery_changes)
|
||
logger.info(f"任务失败，已从数据库恢复集群 {cluster_id} 的资源: {recovery_changes}")
|
||
|
||
raise ValueError("任务提交返回ID为空")
|
||
|
||
except Exception as e:
|
||
# 使用更高效的日志记录方式
|
||
logger.error(f"[{task['task_name']}] 提交失败: {str(e)}", exc_info=True, extra={
|
||
'cluster_id': cluster_id,
|
||
'task_type': task_type,
|
||
'resource_changes': resource_changes if 'resource_changes' in locals() else {}
|
||
})
|
||
return False
|
||
|
||
|
||
def query_third_party_task_status(third_party_task_id: str,token) -> Optional[Dict]:
|
||
"""查询云际平台任务状态"""
|
||
if not third_party_task_id:
|
||
logger.warning("云际任务ID为空，无法查询状态")
|
||
return None
|
||
try:
|
||
config = load_config_from_db()
|
||
url = config["api_config"]["task_detail"]["url"]
|
||
params = {
|
||
"jobSetID": third_party_task_id, # 使用传入的任务ID作为参数
|
||
"localJobID": 1
|
||
}
|
||
|
||
headers = {"Authorization": f"Bearer {token}"} if token else {}
|
||
|
||
response = requests.get(
|
||
url,
|
||
params=params,
|
||
headers=headers,
|
||
timeout=config["api_config"]["task_detail"]["timeout"]
|
||
)
|
||
response.raise_for_status()
|
||
result = response.json()
|
||
|
||
# 适配新的接口返回结构，检查 code 是否为 "OK"
|
||
if result.get("code") != "OK":
|
||
logger.error(f"查询任务状态接口返回异常，响应码非 OK | 响应: {result}")
|
||
return None
|
||
# 检查 data 字段是否存在
|
||
if "data" not in result:
|
||
logger.error(f"查询任务状态接口返回异常，无 data 字段 | 响应: {result}")
|
||
return None
|
||
data = result["data"]
|
||
|
||
# 从数据库重新加载集群资源信息，确保获取最新的资源状态
|
||
load_db_config()
|
||
|
||
return data
|
||
except Exception as e:
|
||
logger.error(f"查询任务状态时发生异常: {str(e)}", exc_info=True)
|
||
return None |