4. 工作流自动化：Shell脚本基础、Python脚本入门、任务调度器（Slurm/PBS）使用、批量任务提交

做高通量计算，说白了就是让计算机替你干活。你想想看，几百个结构优化任务，一个个手动提交？那不得累死。我刚开始接触这行时，就干过这种傻事——熬夜盯着终端，一个任务跑完赶紧提交下一个。后来我师父看到，骂了我一句：「你这是在浪费生命。」

嗯，从那天起，我开始认真研究工作流自动化。今天这章，我就把这几年的经验梳理出来，帮你彻底告别手动操作。

4.1 Shell脚本基础：你的第一个自动化工具

Shell脚本，是Linux环境下的「胶水语言」。它不擅长复杂计算，但特别适合把零散的命令串起来。我在项目中遇到过最典型的场景：批量创建计算目录、复制输入文件、调用计算程序、检查输出文件——这些用Shell脚本写，几行就搞定。

核心要点：Shell脚本的核心是「命令组合」和「流程控制」。你不需要学得很深，掌握变量、循环、条件判断就够了。

4.1.1 变量与参数传递

#!/bin/bash
# 批量创建计算目录
base_dir="/home/project/calc"
elements=("Fe" "Co" "Ni" "Cu")

for elem in "${elements[@]}"; do
    mkdir -p "${base_dir}/${elem}"
    echo "创建目录: ${base_dir}/${elem}"
done

我个人习惯把目录路径、文件名这些「会变的东西」都写成变量。这样脚本复用性高，改一个地方就行。曾经有个同事把所有路径写死在代码里，换台机器跑就全崩了——这就是教训。

4.1.2 循环与批量操作

#!/bin/bash
# 批量提交VASP任务
for i in {1..10}; do
    cd "struct_${i}"
    cp ../INCAR .
    cp ../POTCAR .
    cp ../KPOINTS .
    mpirun -np 16 vasp_std > vasp.log
    cd ..
done

我的小技巧：在循环里加个 echo 打印当前进度，跑长任务时心里有底。我曾经没加打印，跑了一夜发现第二个结构就报错了——白白浪费8小时。

4.2 Python脚本入门：从「能用」到「好用」

Shell脚本虽然方便，但处理复杂逻辑就力不从心了。比如解析输出文件、生成输入文件、做数据可视化——这些还是Python顺手。我建议你把Python当作「瑞士军刀」，Shell当作「螺丝刀」，各司其职。

4.2.1 文件解析与生成

#!/usr/bin/env python3
# 解析VASP的OUTCAR，提取最终能量
import re

def extract_energy(outcar_file):
    with open(outcar_file, 'r') as f:
        content = f.read()
    # 匹配最后一个"energy without entropy"行
    pattern = r"energy without entropy\s*=\s*([-\d.]+)"
    matches = re.findall(pattern, content)
    if matches:
        return float(matches[-1])
    return None

# 批量处理
import glob
for file in glob.glob("*/OUTCAR"):
    energy = extract_energy(file)
    print(f"{file}: {energy} eV")

你看，Python的正则表达式处理文本，比Shell的grep/sed灵活多了。我刚开始学Python时，觉得正则表达式好难，后来发现——其实记住几个常用模式就够了，不用全背。

4.2.2 批量生成输入文件

#!/usr/bin/env python3
# 批量生成不同k点密度的KPOINTS文件
kpoints_template = """Automatic generation
0
Gamma
{k1} {k2} {k3}
0 0 0
"""

k_values = [(2,2,2), (3,3,3), (4,4,4), (5,5,5)]
for i, (k1, k2, k3) in enumerate(k_values):
    dir_name = f"kpoint_{k1}x{k2}x{k3}"
    os.makedirs(dir_name, exist_ok=True)
    with open(f"{dir_name}/KPOINTS", 'w') as f:
        f.write(kpoints_template.format(k1=k1, k2=k2, k3=k3))
    print(f"生成 {dir_name}/KPOINTS")

注意：Python脚本第一行最好加上 #!/usr/bin/env python3，这样可以直接 ./script.py 运行。我曾经忘了加，在集群上死活跑不起来，排查了半天才发现是shebang的问题。

4.3 任务调度器：Slurm/PBS使用

到了集群上，就不能直接 mpirun 了。集群有调度器，负责分配计算资源。最常用的是Slurm和PBS。我最早用的是PBS，后来换了Slurm，说实话Slurm更友好一些。

4.3.1 Slurm脚本模板

#!/bin/bash
#SBATCH --job-name=Fe_bulk        # 任务名称
#SBATCH --output=Fe_bulk_%j.log   # 输出日志
#SBATCH --error=Fe_bulk_%j.err    # 错误日志
#SBATCH --nodes=2                 # 节点数
#SBATCH --ntasks-per-node=16      # 每节点核心数
#SBATCH --time=24:00:00           # 运行时间限制
#SBATCH --partition=compute       # 分区名称

module load vasp/5.4.4
mpirun vasp_std

提交命令很简单：sbatch job.sh。我个人习惯在脚本里加上 module load 加载环境，这样换集群时只需要改这一行。

4.3.2 PBS脚本模板

#!/bin/bash
#PBS -N Fe_bulk
#PBS -l nodes=2:ppn=16
#PBS -l walltime=24:00:00
#PBS -q compute
#PBS -j oe

cd $PBS_O_WORKDIR
module load vasp/5.4.4
mpirun vasp_std

PBS提交用 qsub job.sh。注意 $PBS_O_WORKDIR 这个变量，它自动切换到提交目录。我曾经忘了加，结果任务在home目录下跑，输出文件散落一地——那叫一个乱。

调度器常用命令对比：

功能 Slurm PBS
提交任务 sbatch qsub
查看队列 squeue qstat
取消任务 scancel qdel
查看节点 sinfo pbsnodes

4.4 批量任务提交：从手动到自动

单个任务会提交了，那100个呢？总不能一个个 sbatch 吧。这时候就需要批量提交脚本了。我给大家展示一个我常用的方案。

4.4.1 批量生成并提交Slurm任务

#!/bin/bash
# 批量提交脚本：batch_submit.sh
base_dir="/home/project/structures"
for dir in "${base_dir}"/struct_*; do
    if [ -d "$dir" ]; then
        cd "$dir"
        # 生成Slurm脚本
        cat > job.slurm << EOF
#!/bin/bash
#SBATCH --job-name=$(basename $dir)
#SBATCH --output=$(basename $dir)_%j.log
#SBATCH --ntasks=16
#SBATCH --time=12:00:00

cd $dir
mpirun vasp_std
EOF
        sbatch job.slurm
        echo "已提交: $(basename $dir)"
        cd "$base_dir"
    fi
done

避坑指南：我曾经一次性提交了500个任务，结果调度器直接把我拉黑了——因为超过了用户配额。后来我学乖了，在脚本里加个 sleep 1 或者分批提交，每次只提50个。

4.4.2 使用Python实现更智能的批量提交

#!/usr/bin/env python3
# 智能批量提交：支持断点续传、错误重试
import os
import subprocess
import time

def submit_job(work_dir):
    """提交单个任务，返回job_id"""
    os.chdir(work_dir)
    result = subprocess.run(
        ["sbatch", "job.slurm"],
        capture_output=True, text=True
    )
    if result.returncode == 0:
        job_id = result.stdout.strip().split()[-1]
        return job_id
    return None

def batch_submit(base_dir, max_jobs=50, interval=2):
    """批量提交，控制并发数"""
    dirs = [d for d in os.listdir(base_dir) 
            if d.startswith("struct_")]
    submitted = 0
    
    for d in dirs:
        work_dir = os.path.join(base_dir, d)
        if not os.path.exists(f"{work_dir}/DONE"):
            # 检查是否已完成
            job_id = submit_job(work_dir)
            if job_id:
                submitted += 1
                print(f"提交 {d}, job_id={job_id}")
                
                if submitted % max_jobs == 0:
                    print(f"已达上限 {max_jobs}，等待...")
                    time.sleep(60)  # 等一分钟再继续
            time.sleep(interval)
    
    print(f"共提交 {submitted} 个任务")

if __name__ == "__main__":
    batch_submit("/home/project/structures")

这个脚本有个好处：它会检查 DONE 标记文件，已经算完的结构不会重复提交。我在做高通量筛选时，经常要反复提交——有些任务因为超时被杀了，重新提交时这个机制就特别有用。

4.5 本章知识体系

为了让你更直观地理解工作流自动化的整体结构，我画了一张图：

工作流自动化知识体系 目标：高通量计算自动化 Shell脚本 命令组合·循环·变量 Python脚本 文件解析·批量生成 任务调度器 Slurm·PBS·队列管理 批量提交 并发控制·断点续传 • 变量与参数传递 • for/while循环 • if条件判断 • 文件操作命令 • 正则表达式 • os/sys模块 • 文件读写 • subprocess调用 • sbatch/qsub提交 • squeue/qstat查看 • scancel/qdel取消 • 资源申请参数 • 循环提交 • 并发控制 • 错误重试 • 状态检查 产出：自动化高通量计算流水线

这张图把工作流自动化的四个模块串起来了。你从Shell和Python学起，掌握基本脚本能力；然后学习调度器，知道怎么在集群上跑任务；最后用批量提交把所有东西整合起来。每一步都是下一步的基础。

我的建议：别想着一次性学完所有内容。先拿一个实际的计算任务练手——比如批量优化10个晶体结构。从手动提交开始，然后写Shell脚本自动化，再升级到Python脚本。走完这一轮，你就真正掌握了。

好了，这章的内容就到这里。记住：自动化的目的不是炫技，而是把时间省下来做更有价值的事——比如分析结果、设计新材料。工具是死的，人是活的，用好工具才是本事。

公众号:蓝海资料掘金营，微信deep3321

← 上一章 📖 返回目录 下一章 →