""" slurm_plugin — Slurm 集群管理 - 控制节点: slurmctld + munge + NFS server - 算力节点: slurmd + munge + NFS client - 移除: drain node → scontrol delete → reset - 状态: sinfo """ from pcc import ssh_exec async def _node_info(env, node_id): from sqlor.dbpools import DBPools dbname = env.get_module_dbname('pccs') async with DBPools().sqlorContext(dbname) as sor: recs = await sor.R('compute_node', {'id': node_id}) return recs[0] if recs else None async def _install_slurm_deps(node): """安装 Slurm 依赖 (munge + slurm)""" script = ( "apt-get update -qq && " "apt-get install -y -qq munge slurm-wlm slurm-client nfs-common nfs-kernel-server " "&& systemctl enable munge && systemctl start munge" ) return await ssh_exec(node.ip_address, node.ssh_port or 22, node.ssh_user or 'root', script, timeout=300) async def _setup_munge(control_node, compute_node): """复制 munge key 从控制节点到算力节点""" cmd = f"scp -o StrictHostKeyChecking=no /etc/munge/munge.key {compute_node.ssh_user or 'root'}@{compute_node.ip_address}:/etc/munge/munge.key" rc, _, err = await ssh_exec(control_node.ip_address, control_node.ssh_port or 22, control_node.ssh_user or 'root', cmd, timeout=30) if rc == 0: await ssh_exec(compute_node.ip_address, compute_node.ssh_port or 22, compute_node.ssh_user or 'root', "chown munge:munge /etc/munge/munge.key && chmod 400 /etc/munge/munge.key && systemctl restart munge", timeout=30) async def deploy_cluster(env, cluster_id, control_nodes, compute_nodes, config): """Slurm 集群部署：控制节点安装 slurmctld + munge""" cluster_name = config.get('cluster_name', 'pccs-cluster') results = {'control_nodes': {}, 'endpoint': ''} for node in control_nodes: # 1. 安装依赖 rc, out, err = await _install_slurm_deps(node) if rc != 0: results['control_nodes'][node.id] = {'error': f'install failed: {err[:200]}'} continue # 2. 生成 slurm.conf cpu = node.cpu_cores or 1 mem = node.memory_gb or 4 conf = f"""ClusterName={cluster_name} ControlMachine={node.ip_address} SlurmUser=root SlurmctldPort=6817 SlurmdPort=6818 AuthType=auth/munge StateSaveLocation=/var/spool/slurmctld SlurmdSpoolDir=/var/spool/slurmd ReturnToService=1 SchedulerType=sched/backfill SelectType=select/cons_tres SelectTypeParameters=CR_Core AccountingStorageType=accounting_storage/none JobCompType=jobcomp/none NodeName={node.name} CPUs={cpu} RealMemory={mem * 1024} State=UNKNOWN PartitionName=debug Nodes={node.name} Default=YES MaxTime=INFINITE State=UP """ # 写配置到控制节点 write_cmd = f"cat > /etc/slurm/slurm.conf << 'SLURM_EOF'\n{conf}\nSLURM_EOF" rc2, _, err2 = await ssh_exec(node.ip_address, node.ssh_port or 22, node.ssh_user or 'root', write_cmd, timeout=30) if rc2 != 0: results['control_nodes'][node.id] = {'error': f'config write failed: {err2}'} continue # 3. 启动 slurmctld await ssh_exec(node.ip_address, node.ssh_port or 22, node.ssh_user or 'root', "mkdir -p /var/spool/slurmctld /var/spool/slurmd && " "systemctl enable slurmctld && systemctl start slurmctld", timeout=60) results['control_nodes'][node.id] = {'status': 'ok', 'slurm_conf': conf} results['endpoint'] = node.ip_address async with DBPools().sqlorContext(env.get_module_dbname('pccs')) as sor: await sor.U('cluster', {'id': cluster_id}, {'control_config': json.dumps(results)}) return results async def add_node(env, cluster_id, node_id, role='compute'): """向 Slurm 集群添加算力节点：安装 slurmd + munge → 更新 slurm.conf""" node = await _node_info(env, node_id) if not node: return {'status': 'error', 'message': f'Node {node_id} not found'} # 获取控制节点 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 {'status': 'error', 'message': 'No active control node'} cnode = await _node_info(env, recs[0].node_id) if not cnode: return {'status': 'error', 'message': 'Control node not found'} # 1. 安装依赖 rc, _, err = await _install_slurm_deps(node) if rc != 0: return {'status': 'error', 'message': f'install failed: {err[:200]}'} # 2. 复制 munge key await _setup_munge(cnode, node) # 3. 更新控制节点 slurm.conf 添加此节点 cpu = node.cpu_cores or 1 mem = node.memory_gb or 4 add_line = f"NodeName={node.name} CPUs={cpu} RealMemory={mem * 1024} State=UNKNOWN" await ssh_exec(cnode.ip_address, cnode.ssh_port or 22, cnode.ssh_user or 'root', f"echo '{add_line}' >> /etc/slurm/slurm.conf && scontrol reconfigure", timeout=30) # 4. 启动 slurmd await ssh_exec(node.ip_address, node.ssh_port or 22, node.ssh_user or 'root', "systemctl enable slurmd && systemctl start slurmd", timeout=60) return {'status': 'ok', 'message': f'Node {node.name} joined Slurm cluster'} async def remove_node(env, cluster_id, node_id): """从 Slurm 集群移除节点""" node = await _node_info(env, node_id) if not node: return {'status': 'error', 'message': 'Node not found'} # 获取控制节点 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'}) if recs: cnode = await _node_info(env, recs[0].node_id) if cnode: # drain → remove from slurm.conf → scontrol reconfigure await ssh_exec(cnode.ip_address, cnode.ssh_port or 22, cnode.ssh_user or 'root', f"scontrol update NodeName={node.name} State=DOWN Reason=removing && " f"sed -i '/NodeName={node.name}/d' /etc/slurm/slurm.conf && " "scontrol reconfigure", timeout=30) # stop slurmd on the node await ssh_exec(node.ip_address, node.ssh_port or 22, node.ssh_user or 'root', "systemctl stop slurmd && systemctl disable slurmd", timeout=30) return {'status': 'ok', 'message': f'Node {node.name} removed from Slurm'} async def cluster_status(env, cluster_id): """sinfo""" 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 {'error': 'No active control node'} cnode = await _node_info(env, recs[0].node_id) if not cnode: return {'error': 'Control node not found'} rc, out, _ = await ssh_exec(cnode.ip_address, cnode.ssh_port or 22, cnode.ssh_user or 'root', 'sinfo', timeout=30) return {'status': 'ok', 'sinfo': out if rc == 0 else 'sinfo failed'} 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)