Snakemake 中的资源管理：线程、内存、集群提交

资源管理这回事，我刚开始做 pipeline 时根本没当回事。想着反正本地跑，8 核 CPU 全给一个 rule 不就行了？结果隔壁进程一卡，整个机器直接 OOM…… 后来才明白，资源不是越多越好，得精准控制。说白了，Snakemake 的资源管理就是给每个规则贴上标签——告诉调度器：这个步骤要几颗 CPU、多少内存、需不需要大显存 GPU。下面我把线程、内存、集群提交这三块掰开细讲。

1. 规则里的资源声明

Snakemake 核心就一句话：threads 和 resources 两个指令，写在每个 rule 里。

rule align:
    input:
        "samples/{sample}.fastq"
    output:
        "samples/{sample}.bam"
    threads: 8
    resources:
        mem_mb=16000,
        time_min=120
    shell:
        "bwa mem -t {threads} genome.fa {input} > {output}"

• threads：声明并行线程数，shell 里用 {threads} 传给工具。
• resources：自由定义的名字 – 值对，常见的有 mem_mb、disk_mb、time_min、gpu。

2. 线程管理 —— 别让 CPU 闲着也别抢

Snakemake 用 threads 控制同一规则内的并行度。默认每个规则最多占用 1 核，如果你不写，它就单线程跑。我习惯按工具最推荐的线程数来定：比如 BWA mem 建议用 4-8 线程，我就写 8。

但要注意，threads 只代表“最多可用”，并不是强制占满。实际分配受全局 --cores 限制。比如你设了 --cores 16，Snakemake 会同时调度多个规则，但每个规则不会超过自身 threads 声明。

# 运行 Snakemake，最多使用 16 核
snakemake --cores 16

这里有个坑：如果你在同一集群节点上同时跑两个规则，每个都声明 threads 8，但 --cores 只给了 8，那么它们会排队，不会挤在一起。嗯，Snakemake 调度挺智能的。

3. 内存管理 —— 别让你的 job 被 OOM 杀掉

内存声明在 resources 里，最常用的是 mem_mb。为什么不是 mem_gb？因为 Snakemake 内部统一按 MB 计算，方便做算术。

rule sort_bam:
    input:
        "alignment/{sample}.bam"
    output:
        "sorted/{sample}.bam"
    resources:
        mem_mb=8000      # 8 GB
    shell:
        "samtools sort -m 4G -o {output} {input}"

有的 pipeline 还会声明 disk_mb 和 time_min，方便集群调度器预留资源。我建议至少把 mem_mb 写清楚，否则 Slurm 会按默认值分配，经常导致超限被杀。

4. 集群提交 —— 从本地到 Slurm/PBS/LSF

本地跑小样本还行，一上生产环境就要对接集群。Snakemake 通过 --cluster 或 --profile 实现。最直接的是直接在命令行写集群提交命令：

snakemake --cluster "sbatch --mem={resources.mem_mb}M --cpus-per-task={threads} --time={resources.time_min}" --jobs 20

大括号里的 {resources.mem_mb} 会自动替换成你规则里定义的值。这种方式简单粗暴，但我不太推荐——因为每次启动新 job 都要解析一遍规则，性能差，且参数拼写容易出错。

更好的做法是用 Snakemake Profile。写一个 config.yaml 放在 ~/.config/snakemake/mycluster 下：

cluster:
  mkdir -p logs/{rule} &&
  sbatch
    --partition=defq
    --cpus-per-task={threads}
    --mem={resources.mem_mb}M
    --time={resources.time_min}
    --output=logs/{rule}/{wildcards}.out
    --error=logs/{rule}/{wildcards}.err
jobs: 20
latency-wait: 60

然后直接 snakemake --profile mycluster 即可。这样配置文件干净，还能版本管理。

5. 迁移到 Nextflow 的资源映射

Nextflow 的资源控制更“进程级”。在 process 块里用 cpus、memory、time 声明。下面是一个对照表：

举个例子，把上面 Snakemake 的 align 规则迁移到 Nextflow：

process ALIGN {
    cpus 8
    memory 16.GB
    time 2.h

    input:
    path fastq

    output:
    path "*.bam"

    script:
    """
    bwa mem -t ${task.cpus} genome.fa ${fastq} > output.bam
    """
}

看到了吧？完全是一一对应的关系。但有个区别：Nextflow 的 ${task.cpus} 是自动注入的，不像 Snakemake 要手动用 {threads} 传一遍。这减少了我不少拼写错误。

6. 避坑指南 —— 我踩过的那些雷

• 线程和内存写反：有个同事把 threads: 16000、mem_mb=8，集群直接爆掉。我建议每个上线前先用 snakemake --stats 检查一下资源声明。
• 集群下 --jobs 设太大：Snakemake 默认提交任务到集群，但如果你设 --jobs 200，而集群队列只有 50 个 slot，调度器会排队，但 Snakemake 的提交进程本身会占用大量内存。我记得有一次集群主节点直接被我们的提交进程打垮了…… 后来改成 --jobs 20 就没事了。
• 资源覆盖优先级：规则里的 resources 会覆盖全局默认值，但 --cluster 命令行参数优先级最高。我曾经不小心在命令行里加了 --mem=2G，结果所有 job 都只拿到 2G，samtools 直接 OOM。排查半小时才发现是命令行的锅。