forked from JointCloud/SchedulingSimulator
441 lines
18 KiB
Python
441 lines
18 KiB
Python
import requests
|
||
import json
|
||
from datetime import datetime
|
||
|
||
|
||
class HPCTaskSubmitter:
|
||
"""HPC任务提交器，提供完整的HPC任务管理功能"""
|
||
|
||
def __init__(self):
|
||
"""初始化HPC任务提交器"""
|
||
self.token = None
|
||
|
||
def get_token(self):
|
||
"""
|
||
调用登录接口获取 Token
|
||
:return: 返回获取到的 Token，如果获取失败则返回 None
|
||
"""
|
||
login_url = "http://jcc.jointcloud.net/jcc-admin/admin/login"
|
||
login_payload = {"username": "admin", "password": "Nudt@123"}
|
||
try:
|
||
# 发送 POST 请求进行登录
|
||
login_response = requests.post(login_url, json=login_payload)
|
||
# 检查响应状态码
|
||
if login_response.status_code == 200:
|
||
try:
|
||
# 解析响应内容为 JSON 格式
|
||
login_result = login_response.json()
|
||
# 提取 Token
|
||
token = login_result.get("data", {}).get("token")
|
||
if token:
|
||
self.token = token
|
||
return token
|
||
else:
|
||
print("登录响应中未找到有效的 Token")
|
||
except json.JSONDecodeError:
|
||
print("登录响应数据解析失败，不是有效的 JSON 格式")
|
||
else:
|
||
print(f"登录失败，状态码: {login_response.status_code}")
|
||
except requests.RequestException as e:
|
||
print(f"登录请求过程中发生异常: {e}")
|
||
return None
|
||
|
||
def query_uploaded_files(self):
|
||
"""
|
||
查询文件夹接口
|
||
:return: 返回查询结果，如果查询失败则返回 None
|
||
"""
|
||
if not self.token:
|
||
print("未获取到认证Token，请先调用get_token()")
|
||
return None
|
||
|
||
query_url = "http://jcc.jointcloud.net/jsm/jobSet/queryUploaded"
|
||
headers = {'Content-Type': 'application/json', 'Authorization': f'Bearer {self.token}'}
|
||
query_payload = {
|
||
"queryParams": {
|
||
"dataType": "HPCSlurm",
|
||
"userID": 5,
|
||
"packageID": -1,
|
||
"path": "",
|
||
"CurrentPage": 1,
|
||
"pageSize": 99999,
|
||
"orderBy": "time"
|
||
}
|
||
}
|
||
try:
|
||
# 发送 POST 请求
|
||
query_response = requests.post(query_url, headers=headers, json=query_payload)
|
||
# 检查响应状态码
|
||
if query_response.status_code == 200:
|
||
try:
|
||
# 解析响应内容为 JSON 格式
|
||
query_result = query_response.json()
|
||
return query_result
|
||
except json.JSONDecodeError:
|
||
print("查询响应数据解析失败，不是有效的 JSON 格式")
|
||
else:
|
||
print(f"查询失败，状态码: {query_response.status_code}")
|
||
except requests.RequestException as e:
|
||
print(f"查询请求过程中发生异常: {e}")
|
||
return None
|
||
|
||
def get_package_id_by_name(self, packageName):
|
||
"""
|
||
根据 packageName 查询 packageID
|
||
:param packageName: 包名称
|
||
:return: 返回查询到的 packageID，如果未找到则返回 None
|
||
"""
|
||
query_result = self.query_uploaded_files()
|
||
if query_result:
|
||
uploaded_datas = query_result.get("data", {}).get("uploadedDatas", [])
|
||
for data in uploaded_datas:
|
||
if data.get("packageName") == packageName:
|
||
return data.get("packageID")
|
||
return None
|
||
|
||
def get_object_id_by_package_id(self, packageID, hpcFile, hpctype):
|
||
"""
|
||
根据 packageID 和 hpcFile 查询 objectID
|
||
:param packageID: 包 ID
|
||
:param hpcFile: 文件名
|
||
:param hpctype: HPC任务类型
|
||
:return: 返回查询到的 objectID，如果未找到则返回 None
|
||
"""
|
||
if not self.token:
|
||
print("未获取到认证Token，请先调用get_token()")
|
||
return None
|
||
|
||
query_url = "http://jcc.jointcloud.net/jsm/jobSet/queryUploaded"
|
||
headers = {'Content-Type': 'application/json', 'Authorization': f'Bearer {self.token}'}
|
||
query_payload = {
|
||
"queryParams": {
|
||
"dataType": "HPCSlurm",
|
||
"userID": 5,
|
||
"packageID": packageID,
|
||
"path": "",
|
||
"CurrentPage": 1,
|
||
"pageSize": 99999,
|
||
"orderBy": "time"
|
||
}
|
||
}
|
||
try:
|
||
# 发送 POST 请求
|
||
query_response = requests.post(query_url, headers=headers, json=query_payload)
|
||
# 检查响应状态码
|
||
if query_response.status_code == 200:
|
||
try:
|
||
# 解析响应内容为 JSON 格式
|
||
query_result = query_response.json()
|
||
uploaded_datas = query_result.get("data", {}).get("uploadedDatas", [])
|
||
if uploaded_datas:
|
||
objects = uploaded_datas[0].get("objects", [])
|
||
if hpctype == "bwa":
|
||
for obj in objects:
|
||
if obj.get("path").strip('/') == hpcFile:
|
||
return obj.get("objectID")
|
||
elif hpctype == "lammps":
|
||
if objects:
|
||
return objects[0].get("objectID")
|
||
else:
|
||
# 对于其他类型，如hashcat，返回第一个objectID
|
||
if objects:
|
||
return objects[0].get("objectID")
|
||
except json.JSONDecodeError:
|
||
print("查询响应数据解析失败，不是有效的 JSON 格式")
|
||
else:
|
||
print(f"查询失败，状态码: {query_response.status_code}")
|
||
except requests.RequestException as e:
|
||
print(f"查询请求过程中发生异常: {e}")
|
||
return None
|
||
|
||
def submit_hpc_task(self, task_name, partition, ntasks, nodes, objectID, hpctype="bwa", hpcFile=""):
|
||
"""
|
||
提交任务到指定接口
|
||
:param task_name: 任务名称
|
||
:param partition: 分区名称
|
||
:param ntasks: 任务数量
|
||
:param nodes: 节点数量
|
||
:param objectID: 对象 ID
|
||
:param hpctype: HPC任务类型 (bwa, hashcat, lammps)
|
||
:param hpcFile: HPC文件名
|
||
:return: 提交任务后的响应结果，格式为字典。如果提交失败则返回 None
|
||
"""
|
||
if not self.token:
|
||
print("未获取到认证Token，请先调用get_token()")
|
||
return None
|
||
|
||
# 任务提交接口的 URL
|
||
request_url = "http://jcc.jointcloud.net/jsm/jobSet/submit"
|
||
headers = {'Content-Type': 'application/json', 'Authorization': f'Bearer {self.token}'}
|
||
|
||
# 根据任务类型构造请求负载
|
||
if hpctype.lower() == "bwa":
|
||
create_package_payload = {
|
||
"userID": 5,
|
||
"jobSetInfo": {
|
||
"jobs": [
|
||
{
|
||
"localJobID": "1",
|
||
"type": "HPC",
|
||
"name": task_name,
|
||
"clusterId": "1830873578531228942",
|
||
"backend": "slurm",
|
||
"app": "bwa",
|
||
"operateType": "index",
|
||
"parameters": {
|
||
"hpcBindingFiles": [
|
||
{
|
||
"paramName": "indexFile",
|
||
"resource": {
|
||
"type": "object",
|
||
"objectID": objectID
|
||
}
|
||
}
|
||
],
|
||
"partition": partition,
|
||
"ntasks": ntasks,
|
||
"nodes": nodes
|
||
}
|
||
},
|
||
{
|
||
"localJobID": "4",
|
||
"type": "DataReturn",
|
||
"targetJob": [
|
||
{
|
||
"targetJobID": "1",
|
||
"inputParams": {}
|
||
}
|
||
]
|
||
}
|
||
]
|
||
}
|
||
}
|
||
elif hpctype.lower() == "hashcat":
|
||
create_package_payload = {
|
||
"userID": 5,
|
||
"jobSetInfo": {
|
||
"jobs": [
|
||
{
|
||
"localJobID": "1",
|
||
"type": "HPC",
|
||
"name": task_name,
|
||
"clusterId": "1830873578531228942",
|
||
"backend": "slurm",
|
||
"app": "hashcat",
|
||
"operateType": "",
|
||
"parameters": {
|
||
"partition": partition,
|
||
"ntasks": ntasks,
|
||
"nodes": nodes,
|
||
"hashType": "0",
|
||
"attackMode": "3",
|
||
"hashInput": "e2d0272ec941c0fb4021fb1b7104cfd7",
|
||
"mask": "?d?d?d?d?d?d?d?d"
|
||
}
|
||
},
|
||
{
|
||
"localJobID": "4",
|
||
"type": "DataReturn",
|
||
"targetLocalJobID": "1"
|
||
}
|
||
]
|
||
}
|
||
}
|
||
elif hpctype.lower() == "lammps":
|
||
create_package_payload = {
|
||
"userID": 5,
|
||
"jobSetInfo": {
|
||
"jobs": [
|
||
{
|
||
"localJobID": "1",
|
||
"type": "HPC",
|
||
"description": "qq",
|
||
"name": task_name,
|
||
"clusterId": "1830873578531228942",
|
||
"backend": "slurm",
|
||
"app": "lammps",
|
||
"operateType": "",
|
||
"parameters": {
|
||
"inputFile": "in.min",
|
||
"hpcBindingFiles": [
|
||
{
|
||
"paramName": "inputFile",
|
||
"resource": {
|
||
"type": "object",
|
||
"objectID": objectID
|
||
}
|
||
}
|
||
],
|
||
"partition": partition,
|
||
"ntasks": ntasks,
|
||
"nodes": nodes
|
||
}
|
||
},
|
||
{
|
||
"localJobID": "4",
|
||
"type": "DataReturn",
|
||
"targetJob": [
|
||
{
|
||
"targetJobID": "1",
|
||
"inputParams": {}
|
||
}
|
||
]
|
||
}
|
||
]
|
||
}
|
||
}
|
||
else:
|
||
print(f"未知任务类型: {hpctype}")
|
||
return None
|
||
|
||
try:
|
||
# 发送 POST 请求
|
||
submit_response = requests.post(request_url, headers=headers, json=create_package_payload)
|
||
# 检查响应状态码
|
||
if submit_response.status_code == 200:
|
||
try:
|
||
# 解析响应内容为 JSON 格式
|
||
submit_result = submit_response.json()
|
||
print(f"提交超算结果: {submit_result}")
|
||
return submit_result
|
||
except json.JSONDecodeError:
|
||
print("响应数据解析失败，不是有效的 JSON 格式")
|
||
else:
|
||
print(f"提交超算结果失败，状态码: {submit_response.status_code}")
|
||
except requests.RequestException as e:
|
||
print(f"请求过程中发生异常: {e}")
|
||
return None
|
||
|
||
def submit_single_task(self, task_name, hpctype, partition, ntasks, nodes, packageName, hpcFile):
|
||
"""
|
||
统一的HPC任务提交方法
|
||
:param task_name: 任务名称
|
||
:param hpctype: HPC任务类型 (bwa, hashcat, lammps)
|
||
:param partition: 分区名称
|
||
:param ntasks: 任务数量
|
||
:param nodes: 节点数量
|
||
:param packageName: 包名称
|
||
:param hpcFile: 文件名
|
||
:return: 任务提交结果
|
||
"""
|
||
# 获取 Token
|
||
if not self.token:
|
||
self.get_token()
|
||
|
||
if not self.token:
|
||
print("获取token失败")
|
||
return None
|
||
|
||
# 如果是hashcat任务，不需要查询文件ID，直接提交
|
||
if hpctype.lower() == "hashcat":
|
||
result = self.submit_hpc_task(task_name, partition, ntasks, nodes, None, hpctype)
|
||
return result
|
||
|
||
# 其他类型需要查询文件ID
|
||
if not packageName or not hpcFile:
|
||
print(f"{hpctype}任务需要packageName和hpcFile参数")
|
||
return None
|
||
|
||
# 根据 packageName 获取 packageID
|
||
packageID = self.get_package_id_by_name(packageName)
|
||
if not packageID:
|
||
print(f"未找到名为 {packageName} 的包")
|
||
return None
|
||
|
||
# 根据 packageID 和 hpcFile 获取 objectID
|
||
objectID = self.get_object_id_by_package_id(packageID, hpcFile, hpctype)
|
||
if not objectID:
|
||
print(f"未找到对应的 objectID for package {packageName} and file {hpcFile}")
|
||
return None
|
||
|
||
# 提交任务
|
||
result = self.submit_hpc_task(task_name, partition, ntasks, nodes, objectID, hpctype, hpcFile)
|
||
return result
|
||
|
||
def submit_multiple_tasks(self, tasks_config):
|
||
"""
|
||
批量提交HPC任务
|
||
:param tasks_config: 任务配置列表，每个元素包含任务参数的字典
|
||
:return: 提交结果列表
|
||
"""
|
||
results = []
|
||
for task_config in tasks_config:
|
||
task_name = task_config.get("task_name")
|
||
hpctype = task_config.get("hpctype")
|
||
partition = task_config.get("partition", "ft_test")
|
||
ntasks = task_config.get("ntasks", "1")
|
||
nodes = task_config.get("nodes", "1")
|
||
packageName = task_config.get("packageName", "")
|
||
hpcFile = task_config.get("hpcFile", "")
|
||
|
||
print(f"正在提交任务: {task_name}, 类型: {hpctype}")
|
||
result = self.submit_single_task(task_name, hpctype, partition, ntasks, nodes, packageName, hpcFile)
|
||
results.append({
|
||
"task_name": task_name,
|
||
"result": result
|
||
})
|
||
return results
|
||
|
||
def submit_task_with_timestamp(self, task_name, hpctype, partition="ft_test", ntasks="1", nodes="1", packageName="", hpcFile=""):
|
||
"""
|
||
提交带时间戳的任务
|
||
:param task_name: 任务名称
|
||
:param hpctype: HPC任务类型 (bwa, hashcat, lammps)
|
||
:param partition: 分区名称
|
||
:param ntasks: 任务数量
|
||
:param nodes: 节点数量
|
||
:param packageName: 包名称
|
||
:param hpcFile: 文件名
|
||
:return: 任务提交结果
|
||
"""
|
||
# 添加时间戳到任务名称
|
||
timestamped_task_name = task_name + "_" + datetime.now().strftime("%Y%m%d%H%M%S")
|
||
return self.submit_single_task(timestamped_task_name, hpctype, partition, ntasks, nodes, packageName, hpcFile)
|
||
|
||
|
||
def main(task_name, hpctype, partition="ft_test", ntasks="1", nodes="1", packageName="", hpcFile=""):
|
||
"""
|
||
主函数，统一的任务提交入口
|
||
:param task_name: 任务名称
|
||
:param hpctype: HPC任务类型 (bwa, hashcat, lammps)
|
||
:param partition: 分区名称
|
||
:param ntasks: 任务数量
|
||
:param nodes: 节点数量
|
||
:param packageName: 包名称
|
||
:param hpcFile: 文件名
|
||
:return: 任务提交结果
|
||
"""
|
||
submitter = HPCTaskSubmitter()
|
||
result = submitter.submit_task_with_timestamp(task_name, hpctype, partition, ntasks, nodes, packageName, hpcFile)
|
||
return result
|
||
|
||
|
||
if __name__ == "__main__":
|
||
# 示例：提交BWA任务
|
||
print("提交BWA任务示例:")
|
||
bwa_result = main("hpctask_bwa_test", "bwa", "ft_test", "1", "1", "bwa_data", "reads.fq")
|
||
if bwa_result:
|
||
print("BWA任务提交成功:", bwa_result)
|
||
|
||
print("\n提交Hashcat任务示例:")
|
||
hashcat_result = main("hashcat_test", "hashcat", "ft_test", "64", "4")
|
||
if hashcat_result:
|
||
print("Hashcat任务提交成功:", hashcat_result)
|
||
|
||
print("\n提交LAMMPS任务示例:")
|
||
lammps_result = main("hpctask_lammps_test", "lammps", "ft_test", "1", "1", "lammps_data", "in.min")
|
||
if lammps_result:
|
||
print("LAMMPS任务提交成功:", lammps_result)
|
||
|
||
# 批量提交示例
|
||
print("\n批量提交任务示例:")
|
||
tasks = [
|
||
{"task_name": "batch_bwa_task", "hpctype": "bwa", "partition": "ft_test", "ntasks": "1", "nodes": "1", "packageName": "bwa_data", "hpcFile": "reads.fq"},
|
||
{"task_name": "batch_hashcat_task", "hpctype": "hashcat", "partition": "ft_test", "ntasks": "64", "nodes": "4"},
|
||
{"task_name": "batch_lammps_task", "hpctype": "lammps", "partition": "ft_test", "ntasks": "1", "nodes": "1", "packageName": "lammps_data", "hpcFile": "in.min"}
|
||
]
|
||
|
||
submitter = HPCTaskSubmitter()
|
||
batch_results = submitter.submit_multiple_tasks(tasks)
|
||
for res in batch_results:
|
||
print(f"任务 {res['task_name']} 提交结果: {res['result']}") |