slurm-job-script-generator

SLURM 作业脚本生成器

目标

为运行模拟生成正确且可直接复制粘贴的 SLURM 作业脚本（.sbatch），并指出常见的配置错误（墙钟时间格式不正确、内存标志冲突、超额订阅提示）。

要求

• Python 3.10+
• 无外部依赖（仅使用 Python 标准库）
• 适用于 Linux、macOS 和 Windows（仅生成脚本）

需要收集的输入

输入	说明	示例
作业名称	作业的简短标识符	phasefield-strong-scaling
墙钟时间	SLURM 时间限制	00:30:00
分区	集群分区/队列（如有要求）	compute
账户	项目/账户（如有要求）	matsim
节点数	要分配的节点数	2
MPI 任务数	总任务数或每个节点的任务数	128 或每节点 64
线程数	每个任务的 CPU 数（OpenMP 线程数）	2
内存	--mem 或 --mem-per-cpu（取决于集群策略）	32G
GPU	每个节点的 GPU 数（可选）	4
工作目录	运行作业的目录	$SLURM_SUBMIT_DIR
模块	要加载的环境模块（可选）	gcc/12、openmpi/4.1
运行命令	在 SLURM 下启动的命令	./simulate --config cfg.json

决策指南

MPI 与 MPI+OpenMP 布局

Does the code use OpenMP / threading?
├── NO  → Use MPI-only: cpus-per-task=1
└── YES → Use hybrid: set cpus-per-task = threads per MPI rank
          and export OMP_NUM_THREADS = cpus-per-task

**经验法则：**如果发现 MPI 秩数较高时强扩展效率开始递减，可以尝试减少秩数、增加每个秩的线程数（并进行测量）。

内存标志选择

• 只能使用 --mem（每个节点）或 --mem-per-cpu（每个 CPU）其中之一，不能同时使用。
• 遵循集群文档；有些站点会强制要求采用其中一种方式。
• SLURM 的 --mem 默认单位为整数 MB，也可以使用带后缀 K/M/G/T 的整数（并且 --mem=0 通常表示“使用节点上的全部内存”）。

启动器选择（避免重复包装）

• 默认情况下（--launcher srun），生成器会在运行命令前添加 srun --ntasks=N --cpus-per-task=T，使其沿用 SLURM 的任务放置方式。
• 如果运行命令已以 srun、mpirun 或 mpiexec 开头，请传入 --launcher none，以免生成器重复包装。用 srun 包装 mpirun 会启动 N 个独立的 mpirun 副本（每个副本都会启动自己的 MPI 环境）；用 srun 包装 srun 则会导致格式错误。
• 作为安全措施，生成器会自动检测开头的启动器（srun、mpirun、mpiexec、mpiexec.hydra、orterun、aprun、jsrun），并改为不包装，同时发出建议使用 --launcher none 的警告。

GPU 布局

• 每个 GPU 的秩数 = total_ranks / (nodes * gpus_per_node)。设置 --gpus-per-node 后，生成器会报告 results.derived.total_gpus 和 results.derived.ranks_per_gpu。
• 尽量让秩均匀映射到设备。如果 ntasks 不能被 GPU 总数整除，生成器会发出“任务与 GPU 比例不是整数”的警告；可调整 ntasks/GPU 数，或说明有意共享设备（例如 NVIDIA MPS）。
• 可以考虑使用 --gpu-bind=closest 或 --ntasks-per-gpu 等绑定选项（通过运行命令/--srun-extra 传入）。GPU 和 QoS 策略因站点而异，请查阅集群文档。

脚本输出（JSON 字段）

脚本	主要输出
scripts/slurm_script_generator.py	results.script、results.directives、results.derived、results.warnings、results.run_line

results.derived 会报告 ntasks、ntasks_per_node、cpus_total_requested，以及（如适用）cores_per_node、cpus_per_node_requested、total_gpus 和 ranks_per_gpu。results.warnings 可能包含 CPU 超额订阅、任务与 GPU 比例以及重复启动器相关的警告。

工作流

1. 收集集群约束（分区／账户、GPU 策略、内存策略）。
2. 选择进程布局（仅 MPI，或 MPI+OpenMP 混合模式）。
3. 使用 slurm_script_generator.py 生成脚本。
4. 检查警告（冲突、可疑布局）。
5. 将生成的脚本保存为 job.sbatch。
6. 使用 sbatch job.sbatch 提交，并通过 squeue 监控。

CLI 示例

# Preview a job script (prints to stdout)
python3 skills/hpc-deployment/slurm-job-script-generator/scripts/slurm_script_generator.py \
  --job-name phasefield \
  --time 00:10:00 \
  --partition compute \
  --nodes 1 \
  --ntasks-per-node 8 \
  --cpus-per-task 2 \
  --mem 16G \
  --module gcc/12 \
  --module openmpi/4.1 \
  -- \
  ./simulate --config config.json

# Write to a file and also emit structured JSON
python3 skills/hpc-deployment/slurm-job-script-generator/scripts/slurm_script_generator.py \
  --job-name phasefield \
  --time 00:10:00 \
  --nodes 1 \
  --ntasks 16 \
  --cpus-per-task 1 \
  --out job.sbatch \
  --json \
  -- \
  /bin/echo hello

对话工作流示例

用户：我需要为 MPI 模拟编写一个 sbatch 脚本。要求使用 2 个节点，每个节点 64 个进程，每个进程 2 个 OpenMP 线程，运行 2 小时。

Agent 工作流：

1. 确认分区／账户，以及是否需要 GPU。
2. 生成混合型作业脚本：

   python3 scripts/slurm_script_generator.py --job-name run --time 02:00:00 --nodes 2 --ntasks-per-node 64 --cpus-per-task 2 -- ./simulate
   

3. 说明映射关系：
   • 进程总数 = 128
   • 每个进程的线程数 = 2（OMP_NUM_THREADS=2）
4. 如果用户提供了节点核心数，使用 --cores-per-node 检查是否超额分配。

错误处理

错误	原因	解决方法
time must be HH:MM:SS or D-HH:MM:SS	墙钟时间格式错误	使用 00:30:00 或 1-00:00:00
nodes must be positive	节点数不是正数	提供 --nodes >= 1
Provide either --mem or --mem-per-cpu, not both	内存指令冲突	选择一种内存配置方式
Provide a run command after --	缺少启动命令	添加 -- ./simulate ...
--partition must match /^[A-Za-z0-9]...	分区／账户／qos／约束／预留项包含空格或 shell 元字符	使用普通标识符
module must match /^[A-Za-z0-9]...	模块名称包含 shell 元字符	使用 gcc/12、openmpi/4.1 等
nodes must be <= 100000 (got ...)	整数请求超出合理性检查上限	重新检查请求值

验证清单

• 确认生成的 results.script 将每条 #SBATCH 指令放在 shebang 之后、set -euo pipefail 之前（打开脚本，检查第一条实际命令行）——第一条命令之后的指令会被 SLURM 静默忽略。
• 检查 results.warnings，确认其为空；或者记录每条警告（CPU 超额分配、任务与 GPU 的比例不是整数、双重启动器），并说明有意忽略该警告的理由。
• 记录 results.derived.cpus_total_requested（= ntasks * cpus-per-task）；如果提供了 --cores-per-node，确认 cpus_per_node_requested <= cores_per_node，避免节点超额分配。
• 对于混合运行，确认脚本中的 export OMP_NUM_THREADS 值等于 results.derived 的 cpus_per_task（生成器会将两者设为相同值——确认这与预期的每个进程线程数一致）。
• 对于 GPU 作业，记录 results.derived.total_gpus 和 ranks_per_gpu，并确认 ranks_per_gpu 是预期的整数（或记录有意共享 GPU 的情况，例如 MPS）。
• 确认 results.run_line 未被重复包装：如果运行命令已以 srun／mpirun／mpiexec／orterun／aprun／jsrun 开头，确认使用了 --launcher none（或触发了自动检测警告），避免 SLURM 启动 N 个独立副本。
• 对照集群的实际文档，核实分区、账户、QoS、内存配置方式（--mem 与 --mem-per-cpu）以及 GPU 指令——生成器只检查内部一致性，不会检查站点策略。

常见陷阱与自我安慰

看似省事的做法	为什么不对／应该怎么做
“它生成了一个没有错误的脚本，所以资源配置正确。”	生成器只检查内部一致性，不会查询集群。提交前，请根据你所在站点的实际文档验证分区、账户、QoS、内存样式和 GPU 指令。
“我会保留运行命令中已有的 srun/mpirun，让生成器再包一层。”	在 mpirun 外套一层 srun 会启动 N 个独立的 mpirun 进程（每个进程都有自己的 MPI 世界）；在 srun 外再套一层 srun 则格式错误。请传入 --launcher none，并确认 results.warnings 中出现了自动检测警告。
“我申请了想要的节点，所以作业会用上所有节点。”	如果任何 #SBATCH 指令落在第一条命令之后，该指令会被静默丢弃，作业会退回使用集群默认值。请重新检查生成的脚本，确认所有指令都位于 set -euo pipefail 之前。
“ntasks 不能被 GPU 数量整除，但作业运行了，所以没问题。”	非整数的 ranks_per_gpu 意味着各设备上的 rank 分配不均（GPU 闲置或过度订阅）。生成器会发出任务与 GPU 映射警告——请修正 ntasks/GPU 数量，或明确说明 MPS 共享配置。
“我把 ntasks-per-node 设得很高，因为 rank 越多速度越快。”	不指定 --cores-per-node 时，生成器无法检测过度订阅；而且 ntasks-per-node*cpus-per-task 超过物理核心数会降低性能。请传入 --cores-per-node，并检查 cpus_per_node_requested。
“我同时设置 --mem 和 --mem-per-cpu，这样更保险。”	这两个选项互斥；生成器会拒绝同时提供它们。请选择集群策略要求的那个选项。
“对于纯 MPI 代码，OMP_NUM_THREADS 无关紧要。”	生成器始终会导出 OMP_NUM_THREADS=cpus-per-task；对于纯 MPI 运行，请将 --cpus-per-task=1，以免多线程库在不提示的情况下过度订阅核心。

安全性

输入验证

• --time 会通过正则表达式验证，必须符合严格的 HH:MM:SS 或 D-HH:MM:SS 格式（分钟／秒必须在 [00,59] 范围内）
• --nodes、--ntasks、--ntasks-per-node、--cpus-per-task、--gpus-per-node、--cores-per-node 会验证为正整数，并设有较宽松的上限（例如，nodes ≤ 100000、ntasks ≤ 10000000、cpus-per-task ≤ 4096、gpus-per-node ≤ 64）；推导出的总 ntasks = nodes * ntasks-per-node 也会进行上限检查
• --mem 和 --mem-per-cpu 会根据 SLURM 接受的格式（^[0-9]+([KMGT])?$）进行验证；同时提供两者会被拒绝
• --job-name 会根据 ^[A-Za-z0-9][A-Za-z0-9._-]{0,63}$ 进行验证（不允许空格或 shell 元字符）
• --partition、--account、--qos、--constraint、--reservation 和 --gpu-type 在写入 #SBATCH 指令之前，会根据安全字符允许列表 ^[A-Za-z0-9][A-Za-z0-9._:,+-]{0,127}$ 进行验证
• --module 的值会根据 ^[A-Za-z0-9][A-Za-z0-9._/+-]{0,127}$ 进行验证，以防止 shell 注入（不允许 ;、|、&、反引号、$ 或空白字符）
• --env 的键必须是有效的 shell 标识符（^[A-Za-z_][A-Za-z0-9_]*$）；生成的 export 行会对值进行 shell 引号转义
• --srun-extra 会使用 shlex.split 进行分词，并使用 shlex.quote 对每个词重新加引号，因此无法将 shell 语法（;、|、&、$ 等）注入运行命令行
• 输入无效时，CLI 会向 stderr 输出消息，并以代码 2 退出

文件访问

• 脚本不会读取外部文件；所有输入均通过 CLI 参数提供
• --out 会将生成的 sbatch 脚本写入指定的单个文件路径
• 生成的脚本是包含 #SBATCH 指令的纯文本 shell 脚本，不包含动态生成的代码

工具限制

• Read：用于检查脚本源代码、参考资料和现有作业脚本
• Bash：用于通过显式参数列表执行 slurm_script_generator.py；Agent 不会执行生成的脚本本身
• Write：用于保存生成的 .sbatch 文件；写入范围仅限于用户的工作目录
• Grep/Glob：用于查找现有脚本、配置和集群文档

安全措施

• 不使用 eval()、exec() 或动态代码生成
• 所有子进程调用均使用显式参数列表（不使用 shell=True）
• 运行命令（-- 之后的部分）会按令牌逐个进行 shell 引用并写入生成的脚本，但技能本身绝不会执行该命令
• 模块名称、--srun-extra 令牌和标识符样式的字段都会经过清理或引用处理，以防止向 module load、srun 调用或 #SBATCH 指令中注入内容
• 生成的脚本会将所有 #SBATCH 指令紧接在 shebang 之后（在任何可执行命令之前，以免 SLURM 停止解析），并使用 set -euo pipefail，确保在集群上安全执行 shell 命令

限制

• 不会查询集群硬件或站点策略；只能验证内部一致性。
• 不同 SLURM 安装的配置各不相同（GPU 指令、QoS 规则、分区）。请根据所在站点调整指令。

参考资料

• references/slurm_directives.md - 常见 #SBATCH 指令及映射提示

版本历史

• v1.2.2 (2026-06-24)：新增“验证清单”和“常见陷阱与合理化说辞”章节，涵盖指令顺序、超额订阅、任务到 GPU 的映射、启动器双重封装，以及生成器仅验证内部一致性的范围。
• v1.2.0 (2026-06-23)：修复关键的指令顺序错误（所有 #SBATCH 行现在都位于 set -euo pipefail 之前）；避免运行命令已以 srun/mpirun 开头时重复封装启动器；新增 GPU 任务与 GPU 比例警告及布局指南；强化输入验证（整数上限、分区/账户/qos/约束/预留/GPU 类型允许列表、模块清理、--srun-extra 引用处理）；在 --help 中转义 %j；并修正文档中的安全与输出说明。
• v1.1.0 (2026-03-26)：优化描述以便于发现；新增评估套件、安全审查、标准化元数据和变更日志。
• v1.0.0 (2026-02-25)：初始版本的 SLURM 作业脚本生成器