forked from JointCloud/SchedulingSimulator
690 lines
26 KiB
Python
690 lines
26 KiB
Python
import os
|
||
import json
|
||
import time
|
||
import uuid
|
||
import requests
|
||
from typing import Dict, Optional, List
|
||
from threading import Lock
|
||
import logging # 导入 logging 模块
|
||
|
||
from config.config import TASK_STATUS as TASK_STATUS_COMMON
|
||
from config.db_manager import update_file_all_fields, update_cluster_resources
|
||
|
||
# 任务状态常量
|
||
TASK_STATUS = {
|
||
"PENDING": "pending",
|
||
"RUNNING": "running",
|
||
"SUCCEED": "Succeeded",
|
||
"FAILED": "failed",
|
||
"RETRY_EXHAUSTED": "retry_exhausted"
|
||
}
|
||
|
||
# 全局变量与锁
|
||
# 初始化日志记录器
|
||
logger = logging.getLogger(__name__)
|
||
dataset_map = {} # 存储数据集信息的映射
|
||
dataset_lock = Lock() # 数据集操作锁
|
||
task_map_lock = Lock() # 任务映射操作锁
|
||
cluster_lock = Lock() # 集群资源操作锁
|
||
cluster_resources = {} # 集群资源信息
|
||
API_CONFIG = {} # API配置
|
||
|
||
|
||
# 移除 YAML 配置文件路径，不再需要
|
||
# sonCode_cluster_mapping_path = "sonCode_cluster_mapping.yaml" # YAML配置文件路径
|
||
|
||
|
||
class DatasetInfo:
|
||
"""数据集信息类"""
|
||
|
||
def __init__(self, file_location: str, name: str, size: int, is_uploaded: bool,
|
||
upload_cluster: List[str], upload_time: str, dataset_target_id: str):
|
||
self.file_location = file_location
|
||
self.name = name
|
||
self.size = size
|
||
self.is_uploaded = is_uploaded
|
||
self.upload_cluster = upload_cluster
|
||
self.upload_time = upload_time
|
||
self.dataset_target_id = dataset_target_id
|
||
|
||
def __dict__(self):
|
||
return {
|
||
"file_location": self.file_location,
|
||
"name": self.name,
|
||
"size": self.size,
|
||
"is_uploaded": self.is_uploaded,
|
||
"upload_cluster": self.upload_cluster,
|
||
"upload_time": self.upload_time,
|
||
"dataset_target_id": self.dataset_target_id
|
||
}
|
||
|
||
|
||
def get_config_from_db() -> Dict:
|
||
"""从数据库获取配置"""
|
||
try:
|
||
from config.cache_manager import get_cached_config
|
||
return get_cached_config()
|
||
except Exception as e:
|
||
logger.error(f"从缓存加载配置失败: {str(e)}")
|
||
# 回退到直接从数据库加载
|
||
try:
|
||
from config.db_manager import load_config_from_db
|
||
return load_config_from_db()
|
||
except Exception as e:
|
||
logger.error(f"从数据库加载配置失败: {str(e)}")
|
||
return {}
|
||
|
||
|
||
UPLOAD_CLUSTER_MAPPING = {
|
||
"1865927992266463180": 1,
|
||
"1865927992266461184": 23,
|
||
"1790300942428540928": 3,
|
||
"1830873903296155648": 8,
|
||
"1830873578531228942": 9,
|
||
"1865927992266462181": 10,
|
||
"1830873903296155649": 11,
|
||
"1865927992266462182": 13,
|
||
"1865927992266462180": 14,
|
||
"1770703902472146944": 12
|
||
}
|
||
|
||
|
||
def upload_dataset(task: Dict, token) -> Optional[str]:
|
||
"""上传数据集并返回dataset_target_id"""
|
||
global task_name, copiedToFullRoots
|
||
try:
|
||
# 获取必要参数
|
||
cluster_id = task.get("cluster_id")
|
||
dataset_id = task.get("dataset_id")
|
||
card_type = task.get("card_type")
|
||
|
||
dataset_info = get_fileinfo_by_dataset_id(dataset_id)
|
||
is_uploaded = dataset_info.get("is_uploaded")
|
||
#如果文件已经上传完成并存在dataset_target_id，则直接返回
|
||
dataset_target_id = dataset_info.get("dataset_target_id")
|
||
if dataset_target_id and is_uploaded:
|
||
logger.info(f"数据集已经上传，数据集id为：{dataset_target_id}")
|
||
return dataset_target_id
|
||
file_location = dataset_info.get("file_location")
|
||
file_name = dataset_info.get("dataset_name")
|
||
file_path = f"{file_location}{file_name}"
|
||
package_name = task.get("package_name")
|
||
if not package_name:
|
||
package_name = f"dataset_{uuid.uuid4().hex[:8]}" # 生成默认包名
|
||
task_name = task.get("task_name", "unknown_task")
|
||
|
||
# 检查文件是否存在
|
||
if not os.path.exists(file_path):
|
||
print(f"[{task_name}] 数据集文件不存在: {file_path}")
|
||
return None
|
||
|
||
headers = {
|
||
'Content-Type': 'application/json',
|
||
'Authorization': f'Bearer {token}'
|
||
}
|
||
|
||
# 从数据库获取API配置
|
||
config_data = get_config_from_db()
|
||
api_config = config_data.get("api_config", {})
|
||
|
||
# 1. 创建数据集文件夹
|
||
create_config = api_config.get("create_package", {})
|
||
create_url = create_config.get("url", "http://jcc.jointcloud.net/jsm/jobSet/createPackage")
|
||
create_timeout = create_config.get("timeout", 10)
|
||
|
||
create_payload = {
|
||
"userID": 5,
|
||
"name": package_name,
|
||
"dataType": "dataset",
|
||
"packageID": 0,
|
||
"uploadPriority": {"type": "specify", "clusters": [cluster_id]},
|
||
"bindingInfo": {
|
||
"clusterIDs": [cluster_id],
|
||
"name": package_name,
|
||
"category": "image",
|
||
"type": "dataset",
|
||
"imageID": "",
|
||
"bias": [],
|
||
"region": [],
|
||
"chip": [card_type],
|
||
"selectedCluster": [],
|
||
"modelType": "",
|
||
"env": "",
|
||
"version": "",
|
||
"packageID": 0,
|
||
"points": 0
|
||
}
|
||
}
|
||
create_resp = requests.post(
|
||
create_url,
|
||
json=create_payload,
|
||
headers=headers,
|
||
timeout=create_timeout
|
||
)
|
||
create_resp.raise_for_status()
|
||
create_result = create_resp.json()
|
||
if create_result.get("code") != "OK":
|
||
raise ValueError(f"创建文件夹失败 | 响应: {create_result}")
|
||
packageID = create_result["data"]["newPackage"]["packageID"]
|
||
print(f"[{task_name}] 已创建数据集文件夹 | packageID: {packageID}")
|
||
|
||
# 2. 查询文件夹内容（验证创建结果）
|
||
query_url = "http://jcc.jointcloud.net/jsm/jobSet/queryUploaded"
|
||
query_payload = {
|
||
"queryParams": {
|
||
"dataType": "dataset",
|
||
"userID": 5,
|
||
"packageID": packageID,
|
||
"path": "",
|
||
"CurrentPage": 1,
|
||
"pageSize": 10,
|
||
"orderBy": "name"
|
||
}
|
||
}
|
||
query_resp = requests.post(query_url, json=query_payload, headers=headers, timeout=15)
|
||
query_resp.raise_for_status()
|
||
query_result = query_resp.json()
|
||
if query_result.get("code") != "OK":
|
||
raise ValueError(f"查询文件夹内容失败 | 响应: {query_result}")
|
||
print(f"[{task_name}] 文件夹内容查询成功")
|
||
|
||
# 3. 获取预签名URL用于上传文件到目标文件夹
|
||
|
||
#获取任务的调度类型
|
||
strategy = task.get("strategy")
|
||
|
||
# 根据调度策略选择集群
|
||
if strategy == "data":
|
||
info_data = {
|
||
"userID": 5,
|
||
"info": {
|
||
"type": "upload",
|
||
"params": {
|
||
"packageID": packageID,
|
||
"copyTo": [],
|
||
"copyToPath": [],
|
||
"path": file_name
|
||
}
|
||
}
|
||
}
|
||
else:
|
||
info_data = {
|
||
"userID": 5,
|
||
"info": {
|
||
"type": "upload",
|
||
"params": {
|
||
"packageID": packageID,
|
||
"copyTo": [UPLOAD_CLUSTER_MAPPING[cluster_id]],
|
||
"copyToPath": [f"/dataset/{package_name}/"],
|
||
"path": file_name
|
||
}
|
||
}
|
||
}
|
||
|
||
upload_url = "http://101.201.215.196:7891/storage/presign"
|
||
|
||
presign_resp = requests.post(
|
||
upload_url,
|
||
json=info_data,
|
||
headers={'Authorization': f'Bearer {token}'},
|
||
timeout=300
|
||
)
|
||
|
||
presign_resp.raise_for_status()
|
||
presign_result = presign_resp.json()
|
||
|
||
if presign_result.get("code") != "OK":
|
||
raise ValueError(f"获取预签名URL失败 | 响应: {presign_result}")
|
||
|
||
# 获取预签名URL
|
||
presign_url = presign_result["data"]["presignUrl"]
|
||
|
||
with open(file_path, "rb") as f:
|
||
file_data = f.read()
|
||
|
||
try:
|
||
upload_resp = requests.post(
|
||
url=presign_url,
|
||
data=file_data
|
||
)
|
||
# 打印响应结果
|
||
upload_resp.raise_for_status()
|
||
print(f"[{task_name}] 文件通过预签名URL上传成功")
|
||
|
||
upload_result = upload_resp.json()
|
||
|
||
if upload_result.get("code") != "OK":
|
||
raise ValueError(f"文件上传失败 | 响应: {upload_result}")
|
||
object_id = upload_result["data"]["object"]["objectID"]
|
||
copiedToFullRoots = upload_result["data"]["copiedToFullRoots"]
|
||
|
||
print(f"[{task_name}] 文件上传成功 | objectID: {object_id}")
|
||
|
||
except Exception as e:
|
||
print(f"上传失败，错误信息：{str(e)}")
|
||
|
||
# 4. 通知上传完成
|
||
notify_config = api_config.get("notify_upload", {})
|
||
notify_url = notify_config.get("url", "http://jcc.jointcloud.net/jsm/jobSet/notifyUploaded")
|
||
notify_timeout = notify_config.get("timeout", 15)
|
||
|
||
notify_payload = {
|
||
"userID": 5,
|
||
"packageID": packageID,
|
||
"uploadParams": {
|
||
"copiedTo": [UPLOAD_CLUSTER_MAPPING[cluster_id]],
|
||
"copiedToFullRoots": copiedToFullRoots,
|
||
"dataType": "dataset",
|
||
"uploadInfo": {
|
||
"type": "local",
|
||
"localPath": file_name,
|
||
"objectIDs": [object_id] # 补充objectID
|
||
}
|
||
}
|
||
}
|
||
notify_resp = requests.post(
|
||
notify_url,
|
||
json=notify_payload,
|
||
headers=headers,
|
||
timeout=notify_timeout
|
||
)
|
||
notify_resp.raise_for_status()
|
||
notify_result = notify_resp.json()
|
||
if notify_result.get("code") != "OK":
|
||
raise ValueError(f"通知上传完成失败 | 响应: {notify_result}")
|
||
print(f"[{task_name}] 已通知上传完成")
|
||
|
||
# 5. 二次查询上传内容
|
||
query_resp2 = requests.post(query_url, json=query_payload, headers=headers, timeout=15)
|
||
query_resp2.raise_for_status()
|
||
query_result2 = query_resp2.json()
|
||
if query_result2.get("code") != "OK":
|
||
raise ValueError(f"二次查询失败 | 响应: {query_result2}")
|
||
print(f"[{task_name}] 二次查询上传内容成功")
|
||
|
||
# 6. 查询上传状态
|
||
status_url = "http://jcc.jointcloud.net/jsm/jobSet/uploadStatus"
|
||
status_payload = {
|
||
"userID": 5,
|
||
"operate": "query",
|
||
"packageID": packageID,
|
||
"dataType": "code" # 修正为dataset类型
|
||
}
|
||
status_resp = requests.post(status_url, json=status_payload, headers=headers, timeout=15)
|
||
status_resp.raise_for_status()
|
||
status_result = status_resp.json()
|
||
# if status_result.get("code") != "OK":
|
||
# raise ValueError(f"查询上传状态失败 | 响应: {status_result}")
|
||
if status_resp.status_code == 200:
|
||
upload_status_result = status_resp.json()
|
||
print(f"查询上传状态: {upload_status_result}")
|
||
else:
|
||
print(f"查询上传状态，状态码: {status_resp.status_code}")
|
||
return
|
||
# print(f"[{task_name}] 上传状态: {status_result['data'].get('status')}")
|
||
|
||
# 7. 绑定数据集到集群
|
||
bind_config = api_config.get("bind_cluster", {})
|
||
bind_url = bind_config.get("url", "http://jcc.jointcloud.net/jsm/jobSet/binding")
|
||
bind_timeout = bind_config.get("timeout", 300)
|
||
|
||
bind_payload = {
|
||
"userID": 5,
|
||
"info": {"type": "dataset", "packageID": packageID, "clusterIDs": [cluster_id]}
|
||
}
|
||
bind_resp = requests.post(
|
||
bind_url,
|
||
json=bind_payload,
|
||
headers=headers,
|
||
timeout=300
|
||
)
|
||
bind_resp.raise_for_status()
|
||
bind_result = bind_resp.json()
|
||
if bind_result.get("code") != "OK":
|
||
raise ValueError(f"绑定集群失败 | 响应: {bind_result}")
|
||
print(f"[{task_name}] 已绑定集群 {cluster_id}")
|
||
|
||
# 8. 获取dataset_target_id
|
||
query_binding_url = "http://jcc.jointcloud.net/jsm/jobSet/queryBinding"
|
||
query_binding_payload = {
|
||
"dataType": "dataset",
|
||
"param": {"userID": 5, "bindingID": -1, "type": "private"}
|
||
}
|
||
binding_resp = requests.post(query_binding_url, json=query_binding_payload, headers=headers, timeout=15)
|
||
binding_resp.raise_for_status()
|
||
binding_result = binding_resp.json()
|
||
if binding_result.get("code") != "OK":
|
||
raise ValueError(f"查询绑定信息失败 | 响应: {binding_result}")
|
||
|
||
# 提取目标ID
|
||
dataset_target_id = None
|
||
for data in binding_result["data"]["datas"]:
|
||
if data["info"]["name"] == package_name:
|
||
dataset_target_id = data["ID"]
|
||
break
|
||
if not dataset_target_id:
|
||
raise ValueError(f"未找到 {package_name} 对应的绑定ID")
|
||
print(f"[{task_name}] 成功获取dataset_target_id: {dataset_target_id}")
|
||
|
||
# 保存数据集信息
|
||
with dataset_lock:
|
||
dataset_map[file_location] = DatasetInfo(
|
||
file_location=file_location,
|
||
name=file_name,
|
||
size=os.path.getsize(file_path),
|
||
is_uploaded=True,
|
||
upload_cluster=[cluster_id],
|
||
upload_time=time.strftime("%Y-%m-%d %H:%M:%S"),
|
||
dataset_target_id=dataset_target_id
|
||
)
|
||
#更新数据库中file_mapping表的数据
|
||
if update_file_all_fields(dataset_id, dataset_target_id=dataset_target_id, is_uploaded=1):
|
||
logger.info(
|
||
f"成功更新文件 {dataset_id} 在数据库中的dataset_target_id为 {dataset_target_id}")
|
||
return dataset_target_id
|
||
|
||
except Exception as e:
|
||
print(f"[{task_name}] 数据集上传失败: {str(e)}")
|
||
return None
|
||
|
||
|
||
def get_file_info(cluster_id, file_location, file_name,card_type):
|
||
""""#根据code_id和cluster_id获取son_code_id"""
|
||
try:
|
||
# 从数据库获取算法映射配置
|
||
config_data = get_config_from_db()
|
||
file_mapping = config_data.get('file_mapping', {})
|
||
|
||
for file_id, info in file_mapping.items():
|
||
if cluster_id == str(info["cluster_id"]) and file_location == info["file_location"] and file_name == info["dataset_name"] and card_type == info["card_type"]:
|
||
return info
|
||
|
||
except Exception as e:
|
||
logger.error(f"从数据库查询 文件信息 失败: {str(e)}")
|
||
return None
|
||
|
||
|
||
def get_fileinfo_by_dataset_id(dataset_id):
|
||
"""根据 dataset_id 从数据库中查询 code_id"""
|
||
try:
|
||
# 从数据库获取算法映射配置
|
||
config_data = get_config_from_db()
|
||
file_mapping = config_data.get('file_mapping', {})
|
||
|
||
for file_id, info in file_mapping.items():
|
||
if dataset_id == file_id:
|
||
return info
|
||
|
||
except Exception as e:
|
||
logger.error(f"从数据库查询 文件信息 失败: {str(e)}")
|
||
return None
|
||
|
||
|
||
def query_resource_info(card_type, cluster_id, token):
|
||
"""调用接口查询资源信息"""
|
||
try:
|
||
# 接口URL
|
||
API_URL = "https://jcc.jointcloud.net/jsm/jobSet/queryResource"
|
||
|
||
# 请求参数
|
||
payload = {
|
||
"queryResource": {
|
||
"cpu": {"min": 0, "max": 0},
|
||
"memory": {"min": 0, "max": 0},
|
||
"gpu": {"min": 0, "max": 0},
|
||
"storage": {"min": 0, "max": 0},
|
||
"type": card_type
|
||
},
|
||
"calcNetworkCost": False,
|
||
"resourceType": "Train",
|
||
"clusterIDs": [cluster_id]
|
||
}
|
||
|
||
headers = {
|
||
'Content-Type': 'application/json',
|
||
'Authorization': f'Bearer {token}'
|
||
}
|
||
|
||
# 发送请求并设置超时
|
||
response = requests.post(
|
||
url=API_URL,
|
||
json=payload,
|
||
headers=headers,
|
||
timeout=10 # 10秒超时设置
|
||
)
|
||
response.raise_for_status() # 触发HTTP错误状态码的异常
|
||
|
||
# 解析响应
|
||
result = response.json()
|
||
|
||
# 验证响应状态
|
||
if result.get("code") != "OK":
|
||
logger.error(f"接口返回非成功状态: {result.get('message', '未知错误')}")
|
||
return None
|
||
|
||
# 提取资源列表
|
||
resources = result.get("data", {}).get("resource", [])
|
||
if not resources:
|
||
logger.info("未获取到任何资源数据")
|
||
return None
|
||
|
||
# 筛选符合条件的资源
|
||
for resource in resources:
|
||
# 严格验证类型和集群ID
|
||
if resource.get("type") == card_type and resource.get("clusterId") in cluster_id:
|
||
# 特殊处理 章鱼部分规格不可用
|
||
if card_type == "ILUVATAR-GPGPU" and cluster_id == "1865927992266462182":
|
||
return "BI-V100"
|
||
resource_name = resource.get("name")
|
||
if resource_name:
|
||
logger.info(f"找到符合条件的资源: {resource_name}")
|
||
return resource_name
|
||
logger.warning("找到符合条件的资源，但名称为空")
|
||
|
||
# 未找到符合条件的资源
|
||
logger.info(f"没有找到类型为{card_type}且属于集群{cluster_id}的资源")
|
||
return None
|
||
|
||
except requests.exceptions.Timeout:
|
||
logger.error("请求接口超时")
|
||
return None
|
||
except requests.exceptions.RequestException as e:
|
||
logger.error(f"请求接口失败: {str(e)}")
|
||
return None
|
||
except ValueError as e:
|
||
logger.error(f"解析响应JSON失败: {str(e)}")
|
||
return None
|
||
|
||
|
||
# 提交单个AI任务到集群
|
||
def submit_ai_task(task: Dict, token) -> Optional[str]:
|
||
# 加载API配置
|
||
global API_CONFIG, TASK_STATUS
|
||
config_data = get_config_from_db()
|
||
API_CONFIG = config_data.get('api_config', {})
|
||
|
||
# 获取选择的集群ID
|
||
cluster_id = task.get("cluster_id")
|
||
resource_name = query_resource_info(task.get("card_type"), cluster_id, token)
|
||
if not cluster_id:
|
||
with task_map_lock:
|
||
task["status"] = TASK_STATUS_COMMON["FAILED"]
|
||
task["error_msg"] = "未指定集群ID"
|
||
print(f"[{task['task_name']}] 提交失败：未指定集群ID")
|
||
return None
|
||
|
||
# 数据集处理（获取dataset_target_id）
|
||
dataset_target_id = upload_dataset(task, token)
|
||
if task.get("type") == "Ai" and task.get("dataset_name") and not dataset_target_id:
|
||
# 构造资源变化字典
|
||
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 update_cluster_resources(cluster_id, resource_changes):
|
||
logger.info(f"成功从数据库恢复集群 {cluster_id} 的资源: {resource_changes}")
|
||
else:
|
||
logger.error(f"从数据库恢复集群 {cluster_id} 的资源失败: {resource_changes}")
|
||
|
||
print(f"[{task['task_name']}] 提交失败：数据集处理失败")
|
||
return None
|
||
|
||
headers = {
|
||
'Content-Type': 'application/json',
|
||
'Authorization': f'Bearer {token}'
|
||
}
|
||
task_name = task["task_name"]
|
||
## code_id = task["code_id"]
|
||
card_type = task["card_type"]
|
||
|
||
# 从数据库中查询 son_code_id
|
||
fileInfo = get_fileinfo_by_dataset_id(task["dataset_id"])
|
||
file_location = fileInfo.get("file_location")
|
||
file_name = fileInfo.get("dataset_name")
|
||
file_info = get_file_info(cluster_id,file_location,file_name,card_type)
|
||
son_code_id = file_info.get("son_code_id")
|
||
image_id = file_info.get("image_id")
|
||
|
||
# 检查son_code_id是否有效
|
||
if not son_code_id:
|
||
error_msg = "无法获取有效的son_code_id，任务提交失败"
|
||
print(f"[{task_name}] {error_msg}")
|
||
|
||
# 构造资源变化字典
|
||
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 update_cluster_resources(cluster_id, resource_changes):
|
||
logger.info(f"成功从数据库恢复集群 {cluster_id} 的资源: {resource_changes}")
|
||
else:
|
||
logger.error(f"从数据库恢复集群 {cluster_id} 的资源失败: {resource_changes}")
|
||
|
||
return None
|
||
|
||
try:
|
||
# 1. 提交训练任务
|
||
submit_config = API_CONFIG.get("submit_task", {})
|
||
config = submit_config.get("url", "http://jcc.jointcloud.net/jsm/jobSet/submit")
|
||
submit_timeout = submit_config.get("timeout", 100)
|
||
|
||
task_res = task["resource"]
|
||
# 获取加速器类型（如NPU）
|
||
acc_type = next((key for key in task_res if key not in ["CPU", "MEMORY"]), None)
|
||
|
||
submit_payload = {
|
||
"userID": 5,
|
||
"jobSetInfo": {
|
||
"jobs": [
|
||
{
|
||
"localJobID": "1",
|
||
"name": task_name,
|
||
"description": "",
|
||
"type": "AI",
|
||
"files": {
|
||
"dataset": {"type": "Binding", "bindingID": dataset_target_id},
|
||
"model": {"type": "Binding", "bindingID": None},
|
||
"image": {"type": "Image", "imageID": image_id}
|
||
},
|
||
"jobResources": {
|
||
"scheduleStrategy": "dataLocality",
|
||
"clusters": [
|
||
{
|
||
"code": {"type": "Binding", "bindingID": son_code_id},
|
||
"clusterID": cluster_id,
|
||
"runtime": {"envs": {}, "params": {}},
|
||
"resources": [
|
||
{"type": "CPU", "name": "CPU", "number": task_res["CPU"]},
|
||
{"type": "MEMORY", "name": "RAM", "number": task_res["MEMORY"]},
|
||
{"type": acc_type, "name": resource_name, "number": task_res.get(acc_type, 0)}
|
||
]
|
||
}
|
||
]
|
||
}
|
||
}
|
||
]
|
||
}
|
||
}
|
||
|
||
response = requests.post(
|
||
config,
|
||
json=submit_payload,
|
||
headers=headers,
|
||
timeout=submit_timeout
|
||
)
|
||
response.raise_for_status()
|
||
submit_resp = response.json()
|
||
if submit_resp.get("code") != "OK":
|
||
raise ValueError(f"任务提交失败 | API返回: {submit_resp}")
|
||
|
||
third_party_task_id = submit_resp.get('data', {}).get('jobSetID')
|
||
print(f"[{task_name}] 任务提交至集群 {cluster_id} 成功 | 云际任务ID: {third_party_task_id}")
|
||
|
||
# 更新任务状态为已提交到数据库
|
||
from config.db_manager import update_task_status
|
||
from config.config import TASK_STATUS
|
||
if not update_task_status(task_name, TASK_STATUS["SUBMITTING"], third_party_task_id):
|
||
logger.error(f"更新任务 {task_name} 状态到数据库失败")
|
||
return third_party_task_id
|
||
|
||
|
||
except Exception as e:
|
||
error_msg = f"提交失败: {str(e)}"
|
||
|
||
# 构造资源变化字典
|
||
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 update_cluster_resources(cluster_id, resource_changes):
|
||
logger.info(f"成功从数据库恢复集群 {cluster_id} 的资源: {resource_changes}")
|
||
else:
|
||
logger.error(f"从数据库恢复集群 {cluster_id} 的资源失败: {resource_changes}")
|
||
|
||
task["fail_count"] = task.get("fail_count", 0) + 1
|
||
max_fail = task.get("max_fail_threshold", 3)
|
||
if task["fail_count"] >= max_fail:
|
||
task["status"] = TASK_STATUS["RETRY_EXHAUSTED"]
|
||
else:
|
||
task["status"] = TASK_STATUS["FAILED"]
|
||
task["error_msg"] = error_msg
|
||
|
||
print(f"[{task_name}] {error_msg}")
|
||
return None
|