fix: ray plugin allocate_unit f-string syntax

This commit is contained in:
yumoqing 2026-08-07 13:14:33 +08:00
parent b643b6de0b
commit 66f06e302a
5 changed files with 177 additions and 0 deletions

View File

@ -174,3 +174,48 @@ async def cluster_status(env, cluster_id):
import json
from sqlor.dbpools import DBPools
async def allocate_unit(env, cluster_id, unit_name, cpu, memory, gpu=0):
"""K8s 算力单元分配: 创建 namespace + ResourceQuota"""
from pcc import ssh_exec as _ssh
node = await _find_control_node(env, cluster_id)
if not node:
return {'status': 'error', 'message': 'No control node found'}
quota_yaml = (
f"apiVersion: v1\\nkind: ResourceQuota\\nmetadata:\\n name: {unit_name}-quota\\n namespace: {unit_name}\\n"
f"spec:\\n hard:\\n requests.cpu: \\\"{cpu}\\\"\\n requests.memory: \\\"{memory}Gi\\\"\\n requests.nvidia.com/gpu: \\\"{gpu}\\\""
)
cmds = [
f"kubectl create namespace {unit_name}",
f"echo '{quota_yaml}' | kubectl apply -f -",
]
for cmd in cmds:
rc, out, err = await _ssh(node.ip_address, node.ssh_port or 22, node.ssh_user or 'root', cmd)
if rc != 0:
return {'status': 'error', 'message': f'kubectl failed: {err[:200]}'}
return {'status': 'ok', 'message': f'Unit {unit_name}: ns + {cpu}C/{memory}G/{gpu}GPU'}
async def release_unit(env, cluster_id, unit_name):
"""K8s 算力单元回收: 删除 namespace"""
node = await _find_control_node(env, cluster_id)
if not node:
return {'status': 'error', 'message': 'No control node found'}
from pcc import ssh_exec as _ssh
rc, _, err = await _ssh(node.ip_address, node.ssh_port or 22, node.ssh_user or 'root',
f"kubectl delete namespace {unit_name} --timeout=60s")
if rc != 0:
return {'status': 'error', 'message': f'delete ns failed: {err[:200]}'}
return {'status': 'ok', 'message': f'Unit {unit_name} released'}
async def _find_control_node(env, cluster_id):
from sqlor.dbpools import DBPools
dbname = env.get_module_dbname('pccs')
async with DBPools().sqlorContext(dbname) as sor:
recs = await sor.R('cluster_node', {'cluster_id': cluster_id, 'role': 'control', 'status': 'active'})
if not recs:
return None
return await _node_info(env, recs[0].node_id)

View File

@ -132,3 +132,43 @@ async def cluster_status(env, cluster_id):
import json
from sqlor.dbpools import DBPools
async def allocate_unit(env, cluster_id, unit_name, cpu, memory, gpu=0):
"""Ray 算力单元分配: 创建 detached actor / placement group 绑定资源标签"""
node = await _find_control_node(env, cluster_id)
if not node:
return {'status': 'error', 'message': 'No control node found'}
from pcc import ssh_exec as _ssh
# 用 ray submit 提交一个资源占位 job
script = (
"import ray; ray.init(address='auto'); "
"pg = ray.util.placement_group([{'CPU': " + str(cpu) + ", 'GPU': " + str(gpu) + "}], strategy='STRICT_SPREAD', name='" + unit_name + "'); "
"ray.get(pg.ready()); print('PG ready: " + unit_name + "')"
)
py_cmd = f"python3 -c '{script}'"
rc, out, err = await _ssh(node.ip_address, node.ssh_port or 22, node.ssh_user or 'root', py_cmd, timeout=60)
if rc != 0:
return {'status': 'error', 'message': f'ray pg failed: {err[:200]}'}
return {'status': 'ok', 'message': f'PlacementGroup {unit_name}: {cpu}C/{gpu}GPU'}
async def release_unit(env, cluster_id, unit_name):
"""Ray 算力单元回收: 删除 placement group"""
node = await _find_control_node(env, cluster_id)
if not node:
return {'status': 'error', 'message': 'No control node found'}
from pcc import ssh_exec as _ssh
py_cmd = f"python3 -c 'import ray; ray.init(address=\"auto\"); ray.util.remove_placement_group(ray.util.get_placement_group(\"{unit_name}\")); print(\"removed\")'"
await _ssh(node.ip_address, node.ssh_port or 22, node.ssh_user or 'root', py_cmd, timeout=30)
return {'status': 'ok', 'message': f'PlacementGroup {unit_name} released'}
async def _find_control_node(env, cluster_id):
from sqlor.dbpools import DBPools
dbname = env.get_module_dbname('pccs')
async with DBPools().sqlorContext(dbname) as sor:
recs = await sor.R('cluster_node', {'cluster_id': cluster_id, 'role': 'control', 'status': 'active'})
if not recs:
return None
return await _node_info(env, recs[0].node_id)

View File

@ -176,3 +176,54 @@ async def cluster_status(env, cluster_id):
import json
from sqlor.dbpools import DBPools
async def allocate_unit(env, cluster_id, unit_name, cpu, memory, gpu=0):
"""Slurm 算力单元分配: 创建 partition + 关联节点"""
node = await _find_control_node(env, cluster_id)
if not node:
return {'status': 'error', 'message': 'No control node found'}
from pcc import ssh_exec as _ssh
# 创建 partition
cmd = f"scontrol create PartitionName={unit_name} MaxNodes=UNLIMITED Default=NO MaxTime=UNLIMITED State=UP"
rc, out, err = await _ssh(node.ip_address, node.ssh_port or 22, node.ssh_user or 'root', cmd)
if rc != 0:
return {'status': 'error', 'message': f'scontrol failed: {err[:200]}'}
# 分配 node 到 partition (通过 cluster_node 表找该集群下可用的 compute 节点)
dbname = env.get_module_dbname('pccs')
async with DBPools().sqlorContext(dbname) as sor:
assigned = await sor.sqlExe(
"""SELECT cn.name FROM cluster_node cnn
JOIN compute_node cn ON cn.id = cnn.node_id
WHERE cnn.cluster_id=${cid}$ AND cnn.role='compute' AND cnn.status='active'
LIMIT 1""",
{'cid': cluster_id}
)
if assigned:
node_name = assigned[0].name
await _ssh(node.ip_address, node.ssh_port or 22, node.ssh_user or 'root',
f"scontrol update NodeName={node_name} Partition={unit_name}")
return {'status': 'ok', 'message': f'Partition {unit_name} created'}
async def release_unit(env, cluster_id, unit_name):
"""Slurm 算力单元回收: 删除 partition"""
node = await _find_control_node(env, cluster_id)
if not node:
return {'status': 'error', 'message': 'No control node found'}
from pcc import ssh_exec as _ssh
await _ssh(node.ip_address, node.ssh_port or 22, node.ssh_user or 'root',
f"scontrol delete PartitionName={unit_name}")
return {'status': 'ok', 'message': f'Partition {unit_name} deleted'}
async def _find_control_node(env, cluster_id):
from sqlor.dbpools import DBPools
dbname = env.get_module_dbname('pccs')
async with DBPools().sqlorContext(dbname) as sor:
recs = await sor.R('cluster_node', {'cluster_id': cluster_id, 'role': 'control', 'status': 'active'})
if not recs:
return None
return await _node_info(env, recs[0].node_id)

View File

@ -0,0 +1,22 @@
# 算力单元分配 (集群内部资源: K8s ns+quota / Slurm partition / Ray PlacementGroup)
# params: cluster_id, unit_name, cpu, memory, gpu(可选)
cluster_id = params_kw.get('cluster_id', '')
unit_name = params_kw.get('unit_name', '')
cpu = int(params_kw.get('cpu', 1))
memory = int(params_kw.get('memory', 4))
gpu = int(params_kw.get('gpu', 0))
if not cluster_id or not unit_name:
return {'widgettype': 'Error', 'options': {'title': '失败', 'message': 'Missing cluster_id or unit_name'}}
env = request._run_ns
dbname = env.get_module_dbname('pccs')
async with DBPools().sqlorContext(dbname) as sor:
recs = await sor.R('cluster', {'id': cluster_id})
if not recs:
return {'widgettype': 'Error', 'options': {'title': '失败', 'message': 'Cluster not found'}}
plugin = _get_plugin(recs[0].cluster_type)
result = await plugin.allocate_unit(env, cluster_id, unit_name, cpu, memory, gpu)
if result.get('status') == 'ok':
return {'widgettype': 'Message', 'options': {'title': '分配成功', 'message': result['message'], 'type': 'success', 'timeout': 5}}
return {'widgettype': 'Error', 'options': {'title': '分配失败', 'message': result.get('message', ''), 'cwidth': 16, 'cheight': 9, 'timeout': 3}}

View File

@ -0,0 +1,19 @@
# 算力单元回收
# params: cluster_id, unit_name
cluster_id = params_kw.get('cluster_id', '')
unit_name = params_kw.get('unit_name', '')
if not cluster_id or not unit_name:
return {'widgettype': 'Error', 'options': {'title': '失败', 'message': 'Missing cluster_id or unit_name'}}
env = request._run_ns
dbname = env.get_module_dbname('pccs')
async with DBPools().sqlorContext(dbname) as sor:
recs = await sor.R('cluster', {'id': cluster_id})
if not recs:
return {'widgettype': 'Error', 'options': {'title': '失败', 'message': 'Cluster not found'}}
plugin = _get_plugin(recs[0].cluster_type)
result = await plugin.release_unit(env, cluster_id, unit_name)
if result.get('status') == 'ok':
return {'widgettype': 'Message', 'options': {'title': '回收成功', 'message': result['message'], 'type': 'success', 'timeout': 5}}
return {'widgettype': 'Error', 'options': {'title': '回收失败', 'message': result.get('message', ''), 'cwidth': 16, 'cheight': 9, 'timeout': 3}}