如何用Python脚本SSH连接SLURM调度器并执行命令及传参

我在Ubuntu 22.04环境下，目前通过以下方式操作SLURM作业调度器：

• 在终端中用SSH凭据连接SLURM调度器
• 提交SLURM命令sbatch mysbatchscript.sh
• 其他SLURM命令（如squeue）也通过终端提交

其中mysbatchscript.sh脚本末尾包含python3 main_deeplearning.py语句，用于基于CUDA运行深度学习模型。

我的疑问：

1. 能不能在同一个Python脚本里发起SSH连接并执行这些SLURM命令，替代终端操作？
2. 执行sbatch mysbatchscript.sh时，能不能传递Python脚本里的参数（比如arg1、arg2），让这些参数最终传入main_deeplearning.py？

我找到过Pyslurm模块，但不知道怎么把SSH连接和SLURM调度器关联起来，现在还是用终端的传统方式操作。

1. 用Python脚本实现SSH连接+SLURM命令执行 #

完全可以，不需要依赖Pyslurm（Pyslurm需要在SLURM集群节点上安装配置，且直接和SLURM控制器通信，没法通过SSH远程调用），推荐用paramiko这个Python库实现SSH连接并执行命令，步骤如下：

• 先安装paramiko：pip install paramiko
• 示例代码：

import paramiko

# 建立SSH连接
ssh = paramiko.SSHClient()
ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
ssh.connect(
    hostname="你的SLURM调度器地址",
    username="你的用户名",
    password="你的密码"  # 也可以用密钥文件：key_filename="/path/to/your/private/key"
)

# 执行sbatch提交作业
stdin, stdout, stderr = ssh.exec_command("sbatch mysbatchscript.sh")
print("sbatch输出:", stdout.read().decode())
print("sbatch错误信息:", stderr.read().decode())

# 执行squeue查看个人作业状态
stdin, stdout, stderr = ssh.exec_command("squeue -u 你的用户名")
print("squeue输出:", stdout.read().decode())

# 关闭SSH连接
ssh.close()

如果需要持续执行多个命令或交互式操作，也可以用paramiko的invoke_shell方法创建持久化shell会话。

2. 向sbatch脚本传递Python参数并传入深度学习脚本 #

有两种实用方法：

方法一：让sbatch脚本直接接受命令行参数 #

先修改mysbatchscript.sh，让它接收传入的参数并传递给Python脚本：

#!/bin/bash
# 这里保留你的SLURM配置（比如#SBATCH --gres=gpu:1、#SBATCH --nodes=1这类）

# 将sbatch命令传入的所有参数传递给Python脚本
python3 main_deeplearning.py "$@"

然后在Python的SSH命令中，把参数直接追加到sbatch命令后：

# 假设要传递的参数是arg1="cifar10"、arg2="resnet18"
arg1 = "cifar10"
arg2 = "resnet18"
exec_command = f"sbatch mysbatchscript.sh {arg1} {arg2}"
stdin, stdout, stderr = ssh.exec_command(exec_command)

方法二：通过环境变量传递参数 #

修改mysbatchscript.sh，读取环境变量并传递给Python脚本：

#!/bin/bash
# 保留你的SLURM配置

# 读取环境变量作为Python脚本的参数
python3 main_deeplearning.py $ARG1 $ARG2

然后在Python执行sbatch时，先设置环境变量再提交：

exec_command = f"ARG1={arg1} ARG2={arg2} sbatch mysbatchscript.sh"
stdin, stdout, stderr = ssh.exec_command(exec_command)

内容的提问来源于stack exchange，提问作者AAA7