forked from JointCloud/SchedulingSimulator
243 lines
9.8 KiB
Python
243 lines
9.8 KiB
Python
import random
|
||
import string
|
||
from typing import Optional, Dict, Any, List
|
||
import requests
|
||
import json
|
||
from config.db_manager import load_config_from_db
|
||
from config.config import logger
|
||
from config.auth import get_token
|
||
|
||
class CloudTaskSubmitter:
|
||
"""云算任务提交器，提供完整的云算任务管理功能"""
|
||
|
||
def __init__(self):
|
||
"""初始化云算任务提交器"""
|
||
self.config = load_config_from_db()
|
||
self.token = None
|
||
|
||
def authenticate(self) -> bool:
|
||
"""获取认证Token"""
|
||
self.token = get_token()
|
||
return self.token is not None
|
||
|
||
def query_adapter_list(self) -> Optional[List[Dict]]:
|
||
"""查询适配器列表"""
|
||
if not self.token:
|
||
logger.error("未获取到认证Token，请先调用authenticate()")
|
||
return None
|
||
|
||
headers = {'Content-Type': 'application/json', 'Authorization': f'Bearer {self.token}'}
|
||
|
||
# 检查配置中是否有adapter_list配置
|
||
if "adapter_list" not in self.config["api_config"]:
|
||
logger.error("API配置中缺少adapter_list配置")
|
||
return None
|
||
|
||
api_url = self.config["api_config"]["adapter_list"]["url"]
|
||
|
||
try:
|
||
response = requests.get(api_url, headers=headers)
|
||
response.raise_for_status()
|
||
result = response.json()
|
||
|
||
if result.get("code") == "OK" and "data" in result:
|
||
logger.info("查询适配器列表成功")
|
||
return result["data"]
|
||
else:
|
||
logger.error(f"查询适配器列表失败 | 响应: {result}")
|
||
return None
|
||
except requests.exceptions.RequestException as e:
|
||
logger.error(f"查询适配器列表时发生异常: {str(e)}", exc_info=True)
|
||
return None
|
||
|
||
def query_cluster_list(self, adapter_id: str) -> Optional[List[Dict]]:
|
||
"""查询集群列表"""
|
||
if not self.token:
|
||
logger.error("未获取到认证Token，请先调用authenticate()")
|
||
return None
|
||
|
||
headers = {'Content-Type': 'application/json', 'Authorization': f'Bearer {self.token}'}
|
||
|
||
# 检查配置中是否有cluster_list配置
|
||
if "cluster_list" not in self.config["api_config"]:
|
||
logger.error("API配置中缺少cluster_list配置")
|
||
return None
|
||
|
||
api_url = self.config["api_config"]["cluster_list"]["url"].format(adapter_id=adapter_id)
|
||
|
||
try:
|
||
response = requests.get(api_url, headers=headers)
|
||
response.raise_for_status()
|
||
result = response.json()
|
||
|
||
if result.get("code") == "OK" and "data" in result:
|
||
logger.info("查询集群列表成功")
|
||
return result["data"]
|
||
else:
|
||
logger.error(f"查询集群列表失败 | 响应: {result}")
|
||
return None
|
||
except requests.exceptions.RequestException as e:
|
||
logger.error(f"查询集群列表时发生异常: {str(e)}", exc_info=True)
|
||
return None
|
||
|
||
def generate_random_string(self, length: int = 7) -> str:
|
||
"""生成随机字符串作为任务名称后缀"""
|
||
characters = string.ascii_lowercase # 只使用小写字母
|
||
random_string = ''.join(random.choice(characters) for _ in range(length))
|
||
return random_string
|
||
|
||
def create_task(self, cluster_id: str, task_name: str, replicas: int = 1, cpu: int = 1, memory: int = 1, port: int = 80, image: str = "nginx:latest", strategy: str = "replication") -> Optional[str]:
|
||
"""创建云算任务，返回 jobSetID 或 None"""
|
||
if not self.token:
|
||
logger.error("未获取到认证Token，请先调用authenticate()")
|
||
return None
|
||
|
||
# 生成带随机后缀的任务名称
|
||
task_name = f"{task_name}-{self.generate_random_string()}"
|
||
|
||
# 根据策略创建不同的任务负载
|
||
if strategy == "replication":
|
||
# 使用标准Deployment方式创建任务
|
||
payload = {
|
||
"userID": 5,
|
||
"jobSetInfo": {
|
||
"jobs": [{
|
||
"localJobID": "1",
|
||
"name": task_name,
|
||
"type": "PCM_Container",
|
||
"clusterId": cluster_id,
|
||
"containerGroupName": task_name,
|
||
"image": image,
|
||
"cpu": str(cpu),
|
||
"memory": str(memory),
|
||
"port": port,
|
||
"capacity": replicas,
|
||
"mountPath": "/models"
|
||
}]
|
||
}
|
||
}
|
||
else:
|
||
# 其他策略可以在此扩展
|
||
payload = {
|
||
"userID": 5,
|
||
"jobSetInfo": {
|
||
"jobs": [{
|
||
"localJobID": "1",
|
||
"name": task_name,
|
||
"type": "PCM_Container",
|
||
"clusterId": cluster_id,
|
||
"containerGroupName": task_name,
|
||
"image": image,
|
||
"cpu": str(cpu),
|
||
"memory": str(memory),
|
||
"port": port,
|
||
"capacity": replicas,
|
||
"mountPath": "/models"
|
||
}]
|
||
}
|
||
}
|
||
|
||
headers = {
|
||
'Content-Type': 'application/json',
|
||
'Authorization': f'Bearer {self.token}'
|
||
}
|
||
|
||
# 尝试从配置中获取create_task的URL，如果不存在则使用默认URL
|
||
if "submit_task" not in self.config["api_config"]:
|
||
# 尝试使用其他可能的键名
|
||
possible_keys = ["create_task", "task_create", "submit_task", "create"]
|
||
api_url = None
|
||
for key in possible_keys:
|
||
if key in self.config["api_config"]:
|
||
api_url = self.config["api_config"][key]["url"]
|
||
break
|
||
|
||
# 如果仍然没有找到，使用默认URL
|
||
if api_url is None:
|
||
logger.warning("API配置中缺少create_task配置，使用默认URL")
|
||
api_url = "https://jcc.jointcloud.net/jsm/jobSet/submit" # 默认URL
|
||
else:
|
||
api_url = self.config["api_config"]["submit_task"]["url"]
|
||
|
||
try:
|
||
response = requests.post(api_url, headers=headers, json=payload)
|
||
response.raise_for_status()
|
||
result = response.json()
|
||
|
||
# 提取并返回 jobSetID
|
||
if result.get("code") == "OK" and "data" in result:
|
||
job_set_id = result["data"].get("jobSetID")
|
||
logger.info(f"任务创建成功 | 任务名称: {task_name} | JobSetID: {job_set_id}")
|
||
return job_set_id
|
||
else:
|
||
logger.error(f"任务创建失败 | 响应: {result}")
|
||
return None
|
||
except requests.exceptions.RequestException as e:
|
||
logger.error(f"创建任务时发生异常: {str(e)}", exc_info=True)
|
||
return None
|
||
|
||
def submit_cloud_task(self, task_name: str, cluster_id: str, replicas: int = 1, cpu: int = 1, memory: int = 1, port: int = 80, image: str = "nginx:latest", strategy: str = "replication") -> Optional[str]:
|
||
"""提交云算任务的主方法"""
|
||
# 首先进行身份验证
|
||
if not self.authenticate():
|
||
logger.error("认证失败，无法提交任务")
|
||
return None
|
||
|
||
# 创建任务并返回JobSetID
|
||
return self.create_task(cluster_id, task_name, replicas, cpu, memory, port, image, strategy)
|
||
|
||
def get_available_clusters(self) -> List[Dict]:
|
||
"""获取所有可用的云集群"""
|
||
available_clusters = []
|
||
for cluster_id, cluster_info in self.config["cluster_resources"].items():
|
||
if cluster_info["cluster_type"] == "Cloud":
|
||
available_clusters.append({
|
||
"cluster_id": cluster_id,
|
||
"cluster_info": cluster_info
|
||
})
|
||
return available_clusters
|
||
|
||
def get_cluster_resources(self, cluster_id: str) -> Optional[Dict]:
|
||
"""获取指定集群的资源信息"""
|
||
return self.config["cluster_resources"].get(cluster_id)
|
||
|
||
def query_task_status(self, job_set_id: str) -> Optional[Dict]:
|
||
"""查询任务状态"""
|
||
if not self.token:
|
||
logger.error("未获取到认证Token，请先调用authenticate()")
|
||
return None
|
||
|
||
headers = {'Authorization': f'Bearer {self.token}'}
|
||
|
||
# 检查配置中是否有task_status配置
|
||
if "task_status" not in self.config["api_config"]:
|
||
logger.error("API配置中缺少task_status配置")
|
||
return None
|
||
|
||
api_url = self.config["api_config"]["task_status"]["url"]
|
||
|
||
try:
|
||
params = {
|
||
"jobSetID": job_set_id,
|
||
"localJobID": 1
|
||
}
|
||
|
||
response = requests.get(api_url, headers=headers, params=params)
|
||
response.raise_for_status()
|
||
result = response.json()
|
||
|
||
if result.get("code") == "OK" and "data" in result:
|
||
logger.info(f"任务状态查询成功 | JobSetID: {job_set_id}")
|
||
return result["data"]
|
||
else:
|
||
logger.error(f"任务状态查询失败 | 响应: {result}")
|
||
return None
|
||
except requests.exceptions.RequestException as e:
|
||
logger.error(f"查询任务状态时发生异常: {str(e)}", exc_info=True)
|
||
return None
|
||
|
||
# 便捷函数，兼容原有接口
|
||
def submit_cloud_task(task_name: str, cluster_id: str, replicas: int = 1, cpu: int = 1, memory: int = 1, port: int = 80) -> Optional[str]:
|
||
"""便捷函数：提交云算任务"""
|
||
submitter = CloudTaskSubmitter()
|
||
return submitter.submit_cloud_task(task_name, cluster_id, replicas, cpu, memory, port) |