Post

利用 btrun 管理量子化学任务运行队列

利用 btrun 管理量子化学任务运行队列

前言

现状

多数量子化学程序都需要较长的运行时间。只要一次计算的任务数多起来,就会遇到一个很实际的问题:怎样让任务按照可用资源自动排队,而不是靠人守着终端一个个启动。

常见的处理方式大致有两类。

第一类是使用 Slurm、SGE、PBS 等成熟的作业调度系统。它们的资源分配和队列管理都很完善,也是集群环境中的常规选择。不过在个人工作站或单机服务器上,安装和配置一套完整调度系统略显繁琐,初次接触时还要理解节点、分区、队列、资源申请等概念。Windows 下如果只是想简单排几个量化任务,也很难直接套用这套方案。

即使已经有作业调度系统,任务脚本本身仍然需要管理。一个 Slurm 脚本里往往有相当一部分固定内容,如果每次计算都重新写提交头、环境加载和运行命令,实际使用起来仍然比较麻烦。

第二类是自己写一个调度脚本。这样可以把日常命令压得很短,启动脚本后自动依次运行任务。但这类脚本通常和某一种程序、某一种输入格式或某一种工作流绑定得很紧。只要进一步加入并行、资源检查以及多程序支持,脚本本身很快就会变复杂。

设计思路

如果只考虑量子化学任务的排队运行,每个任务真正需要变化的信息大致可以归纳为三类:

  1. 任务需要多少核心、内存和 GPU 等资源;
  2. 运行程序前需要加载什么环境;
  3. 最终执行什么命令。

像 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 二进制文件以及配置文件示例均可在计算化学公社原帖下载:

利用 btrun 管理任务运行队列——计算化学公社

如果遇到 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 loadsource 环境脚本或手动设置环境变量。

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.gjf2.gjf3.gjf

随后它会检查对应输出是否已经存在。如果配置的输出后缀为 log,那么现有的 1.log 会使 1.gjf 被视为已经运行过,从而避免重复提交;最终只处理 2.gjf3.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、额外参数以及后处理等选项。实际使用中遇到本文没有覆盖的提交方式时,以当前版本的帮助信息为准。

This post is licensed under CC BY 4.0 by the author.
Total hits!