利用 btrun 管理量子化学任务运行队列
前言
现状
多数量子化学程序都需要较长的运行时间。只要一次计算的任务数多起来,就会遇到一个很实际的问题:怎样让任务按照可用资源自动排队,而不是靠人守着终端一个个启动。
常见的处理方式大致有两类。
第一类是使用 Slurm、SGE、PBS 等成熟的作业调度系统。它们的资源分配和队列管理都很完善,也是集群环境中的常规选择。不过在个人工作站或单机服务器上,安装和配置一套完整调度系统略显繁琐,初次接触时还要理解节点、分区、队列、资源申请等概念。Windows 下如果只是想简单排几个量化任务,也很难直接套用这套方案。
即使已经有作业调度系统,任务脚本本身仍然需要管理。一个 Slurm 脚本里往往有相当一部分固定内容,如果每次计算都重新写提交头、环境加载和运行命令,实际使用起来仍然比较麻烦。
第二类是自己写一个调度脚本。这样可以把日常命令压得很短,启动脚本后自动依次运行任务。但这类脚本通常和某一种程序、某一种输入格式或某一种工作流绑定得很紧。只要进一步加入并行、资源检查以及多程序支持,脚本本身很快就会变复杂。
设计思路
如果只考虑量子化学任务的排队运行,每个任务真正需要变化的信息大致可以归纳为三类:
- 任务需要多少核心、内存和 GPU 等资源;
- 运行程序前需要加载什么环境;
- 最终执行什么命令。
像 Slurm 提交脚本中大量固定的 header,对同一台集群上的大多数任务并不是每次都要重新提供的信息。
资源需求通常可以先设置一个默认值,遇到特殊任务时再从命令行覆盖。程序的环境加载方式和主运行命令则天然适合按程序分类保存,例如 Gaussian 使用一份配置,ORCA 使用另一份配置。
这样一来,提交任务时只需要把“程序配置”“本次任务的资源需求”和“执行端的队列配置”组合起来:目标是 Slurm、SGE 或 PBS 时生成对应的任务脚本并提交;目标是本机时,则交给一个简单的本地队列执行。
笔者早期尝试过 Bash 脚本,也单独写过 Python 队列程序,最后把这些功能整合成了 btrun。
btrun是banetask的一个组件,因此程序中还存在一些与本文无关的功能。本文只讨论量子化学任务运行队列相关的部分。
btrun
介绍
btrun 是一个任务队列管理工具,可以统一生成并提交 Slurm、PBS、SGE 的任务脚本,也自带一个 local queue,供没有安装作业调度系统的单机环境进行串行或并行计算。Linux 和 Windows 都可以使用。
对于已经配置好的量子化学程序,日常提交可以缩短成类似下面的命令:
1
btrun task -t g16
这条命令会根据 g16 对应的程序配置寻找输入文件,再结合任务资源和当前执行端配置决定如何提交。
下载程序后,需要先让命令行能够找到 btrun。
Linux 下可以把程序所在目录加入 ~/.bashrc:
1
export PATH=path/to/btrun:$PATH
刷新终端后即可直接执行 btrun。
Windows 下可以把 btrun.exe 所在目录加入系统 Path。例如程序位于 D:\program\banetask\btrun.exe,则把 D:\program\banetask 加入 Path。如果不希望直接在 cmd 中修改环境变量,可以按 Win+R,输入 sysdm.cpl,再通过“环境变量”界面编辑 Path。保存后重新打开终端即可。
下载
Linux、Windows 二进制文件以及配置文件示例均可在计算化学公社原帖下载:
如果遇到 bug,也可以直接在上述帖子中回复反馈。
程序配置文件
btrun 使用配置文件描述不同量子化学程序应该怎样寻找输入、申请默认资源、加载环境并执行。配置目录默认为:
1
~/.bane/task/envs/
例如 Gaussian 可以建立 ~/.bane/task/envs/g16.conf:
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
[main]
# 资源配置
CONF_CORES=32
CONF_MEMORY=96000
SUFFIX=gjf
CONF_SPAN_NODES=false
[ENV_SETUP]
# g16 environment
# module load g16
source "$HOME/apprepo/gaussian/16-hy/scripts/env.sh"
export PGI_FASTMATH_CPU=sandybridge
export PATH=~/scripts/bin:$PATH
[RUN_CMD_TEMPLATE]
g16loop "${ACTUAL_INPUT_FILE}" --crash-handle all > "${ACTUAL_OUTPUT_FILE}"
配置文件整体采用类 TOML 的形式,最核心的是 [main]、[ENV_SETUP] 和 [RUN_CMD_TEMPLATE] 三个块。
[main] 用来描述输入类型和默认资源。常见参数包括(完整列表见btrun软件页):
CONF_CORES:命令行没有另外指定时使用的默认核心数;CONF_MEMORY:命令行没有另外指定时使用的默认内存;SUFFIX:扫描任务时寻找的输入文件后缀;INPUT_FILENAME:MRCC 这类使用固定输入文件名的程序可以在这里指定文件名;RAW_MODE=true:不按固定输入文件名或后缀扫描,而是把整个工作目录作为一个任务;OUTPUT_SUFFIX:指定预期输出文件的后缀,可用于判断任务是否已经运行;USE_ABSOLUTE_PATH=false:对于不接受绝对路径输入的程序,可以关闭绝对路径传递。
[ENV_SETUP] 中直接写程序运行前需要执行的环境设置,例如 module load、source 环境脚本或手动设置环境变量。
Windows 用户需要注意,这一部分写的是 cmd 能执行的命令,而不是 PowerShell 语法。
[RUN_CMD_TEMPLATE] 中写真正启动量化程序的命令。常用占位符包括 ${ACTUAL_INPUT_FILE}、${ACTUAL_OUTPUT_FILE}、${CORES} 和 ${EXTRA_ARGS},分别表示实际输入文件、实际输出文件、任务最终使用的核心数以及命令行临时传入的额外参数。
例如某个程序需要在启动命令中显式写入核心数,就可以直接使用 ${CORES}。如果某次提交还需要临时增加参数,可以运行:
1
btrun task -t some_program -x "args"
此时 args 会替换配置中的 ${EXTRA_ARGS}。
资源初始化
程序配置文件解决的是具体量化程序应该怎么运行,而本机有多少资源、是否存在作业调度系统、集群有哪些队列,则由 btrun init 初始化。
直接运行:
1
btrun init
btrun 会先探查当前机器,再交互式询问执行端配置。以笔者使用的 Slurm 登录节点为例,初始化过程中会涉及这些信息:
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
Profile/cluster name [default]:
Scheduler backend (slurm/sge/pbs/local) [slurm]:
Detected local machine:
cores: 64
memory_mb: 257529
gpus: 0
scheduler commands: slurm,sge/pbs
inferred role: login_node
Enable local execution target? [y/N]: y
Local max CPU cores [64]:
Local max memory, MB [257529]:
Reserve CPU cores for OS/interactive use [4]:
Reserve memory for OS/interactive use, MB [16384]:
Discovered scheduler queues: whhcnormal
Use discovered queues? [Y/n]:
Physical CPU cores per node [32]:
Max CPU cores per node [32]:
Max memory per node, MB [255500]:
Default cores per job [16]: 32
Default memory per job, MB [64000]: 96000
High-memory threshold, MB/core [7984]:
Slurm header style (ntasks_per_node/ntasks/cpus_per_task) [ntasks_per_node]:
初始化时既可以描述本机可用于 local queue 的资源,也可以描述调度系统中的队列、单节点资源上限、默认任务资源以及提交脚本格式。完成这些配置后,btrun 才具备在不同执行端之间组织任务的基础信息。
使用方法
作业调度系统
假设已经配置好 g16.conf,当前目录中存在:
1
2
3
4
1.gjf
2.gjf
3.gjf
1.log
运行:
1
btrun task -t g16
btrun 会先读取 g16.conf,根据其中的 SUFFIX=gjf 找到 1.gjf、2.gjf 和 3.gjf。
随后它会检查对应输出是否已经存在。如果配置的输出后缀为 log,那么现有的 1.log 会使 1.gjf 被视为已经运行过,从而避免重复提交;最终只处理 2.gjf 和 3.gjf。
如果当前执行端配置了 Slurm、SGE 或 PBS,btrun 会把调度系统的提交头、[ENV_SETUP] 中的环境加载以及 [RUN_CMD_TEMPLATE] 中的运行命令组合成完整任务脚本,再提交给调度系统。
不同集群对提交脚本经常有不同要求。例如某些集群要求使用 --ntasks-per-node,另一些集群可能推荐不同的 CPU 参数,也有集群不允许用户显式指定某些内存选项。为此,btrun 允许直接修改队列对应的提交脚本模板。
btrun init 生成的配置位于:
1
~/.bane/task/btrun/queues/default.yaml
其中与任务脚本直接相关的部分类似:
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
queues:
whhcnormal:
scheduler_queue_name: "whhcnormal"
resources:
physical_cpu_num: 32
max_nodes: 455
max_cores_per_node: 32
max_memory_mb_per_node: 255500
default_cores: 32
default_memory_mb: 96000
high_memory_threshold_mb_per_core: 6000
header_lines:
- "#!/bin/bash"
- "#SBATCH -J "
- "#SBATCH -p "
- "#SBATCH -N "
- "#SBATCH --ntasks-per-node="
- "#SBATCH --mem=M"
- ""
- "export CORES="
- "export BTRUN_QUEUE="
- "cd || exit 97"
header_lines 中的内容就是最终任务脚本头部的模板。所在集群有特殊要求时,直接修改这里即可,不必为每一种量化程序再单独维护一套 Slurm 八股文。
本地队列
如果没有配置作业调度系统,btrun 可以把任务交给自带的 local queue。也可以在提交时使用 -L,把执行目标限制为本地队列。
local queue 会根据配置的总核心数、内存和各任务资源需求判断哪些任务能够同时运行。例如本机可用 64 核,一批任务各需要 32 核,并且两个任务同时运行时的总内存不超过上限,那么 btrun 可以一次启动两个任务。
这里的资源管理与 Slurm 这类完整调度系统并不相同。btrun 会在自己的队列逻辑中检查资源是否足够,但不会像真正的作业调度系统一样负责 CPU 绑核,也不会在操作系统层面强制限制进程实际使用的内存。因此 local queue 更适合个人工作站和简单计算服务器上的任务排队。
常用的队列管理命令包括:
1
2
3
btrun queue list
btrun queue history
btrun queue cancel <JOBID>
例如:
1
2
3
4
5
6
7
8
$ btrun queue list
JOBID STATE BACKEND PROFILE PID EXIT JOB_NAME WORKDIR
260823-170000_aaaaaaaa running local - 16995 - live-job /tmp/live
$ btrun queue history
JOBID STATE BACKEND PROFILE PID EXIT JOB_NAME WORKDIR
260823-160000_bbbbbbbb done local - - 0 done-job /tmp/done
260823-150000_cccccccc failed local - - 127 failed-job /tmp/failed
选择性后处理
如果某种后处理对所有任务都要执行,可以直接把它写进 [RUN_CMD_TEMPLATE]。但有些操作只需要针对一部分任务执行,例如 Gaussian 几何优化后检查虚频,而普通能量计算并不需要,这时可以为后处理单独增加一个配置块。
例如在 g16.conf 中加入:
1
2
3
[IMAG]
value=$(banedata ${ACTUAL_OUTPUT_FILE} -i)
[ "$value" -gt 0 ] && touch IMAG_FOUND
提交需要检查虚频的结构优化任务时使用:
1
btrun task -t g16 --post-var IMAG
主计算结束后,btrun 就会继续执行 [IMAG] 中的逻辑。这里的块名可以自行定义,只要提交时通过 --post-var 使用同一个名字即可。后处理块中也可以使用与 [RUN_CMD_TEMPLATE] 相同的占位符。
示例
以Windows下运行 xTB 为例,可以准备一个 xtb.conf:
1
2
3
4
5
6
7
8
9
10
11
[main]
# Generated by the BaneTask installer.
CONF_CORES=8
CONF_MEMORY=4000
SUFFIX=xyz
[ENV_SETUP]
REM You can set xtb's env here or use absolute path in [RUN_CMD_TEMPLATE].
[RUN_CMD_TEMPLATE]
"D:\ChemicalPrograms\xtb-6.7.1\bin\xtb.exe" "${ACTUAL_INPUT_FILE}" ${EXTRA_ARGS} > "${ACTUAL_OUTPUT_FILE}" 2>&1
如果每次都使用固定参数,可以把 --opt --gfn 2 之类的选项直接写进 [RUN_CMD_TEMPLATE]。如果希望提交时灵活改变,则保留 ${EXTRA_ARGS},再通过 -x 传入。例如:
1
btrun task -t xtb -x "--gfn 1 --ohess"
xTB 使用 XYZ 输入时还要注意其工作目录中会生成一些固定文件名,因此不适合把多个彼此独立的 XYZ 任务同时放在同一个目录中运行。更合适的方式是让每个任务使用独立目录,再通过 -a 把当前路径下的子目录一并纳入扫描:
1
btrun task -t xtb -x "--gfn 1 --ohess" -a
如果不同分子的电荷和自旋不同,可以按照 xTB 官方文档所述,在各自目录中通过 .CHRG 和 .UHF 文件分别指定。
一个 Windows 本地队列的运行过程例如:
1
2
3
4
5
6
7
8
(base) PS D:\tests> btrun task -t xtb -x "--gfn 1 --ohess" -a
[btrun] local queue: profile=default pid=30228 capacity=8C/1G[0] poll=5s
[btrun] START 1 [561882c3] 8C pid=38180
[btrun] DONE 1 [561882c3] rc=0 5s
[btrun] START 2 [fdfe97fa] 8C pid=21712
[btrun] DONE 2 [fdfe97fa] rc=0 5s
[btrun] queue drained
(base) PS D:\tests>
其他
还有几项默认行为比较容易在实际使用时遇到。
对于 Gaussian 任务,配置文件必须是
g16.conf;对于 ORCA 任务,配置文件必须是orca.conf。如果命令行没有使用-m、-n等选项显式指定资源,btrun会尝试解析输入文件中的核数和内存设置,并把它们作为提交任务时的资源需求。btrun队列配置文件中的
High-memory threshold用于描述队列的 memory-per-core 限制。某些 post-HF 任务实际运行时可能只需要较少核心,却需要较大内存,而集群又可能规定每核最多申请一定内存。遇到这种情况时,btrun会按照总内存和 memory-per-core 阈值计算满足限制所需的申请核数。例如:某个队列的 memory-per-core 上限为 3000 MB,一个任务实际希望使用 16 核和 96000 MB 内存,则:
1
96000 / 3000 = 32
为了满足队列规则,
btrun会尝试向调度系统申请 32 核和 96000 MB 内存,前提是该队列允许申请这么多核心。如果还配置了其他内存更充足的队列,btrun会优先尝试把高内存任务转交到更合适的队列,而不是直接在当前队列增加核心数。目前
btrun task对 GPU 资源的管理还不完整,后续仍可能调整。更多提交参数和队列管理选项可以直接查看程序内置帮助:
1 2 3
btrun task -h btrun queue -h btrun help conf
其中
btrun task -h还列出了后端筛选、profile、queue、递归扫描、任务核数和内存、dry-run、local queue、额外参数以及后处理等选项。实际使用中遇到本文没有覆盖的提交方式时,以当前版本的帮助信息为准。