pcc/build/lib/pcc/slurm_plugin/__init__.py

"""
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)