如何在单工作站通过Snakemake正确启用多GPU并行执行任务？

根据你的需求，分两种常见场景给出配置方案：

场景1：单个任务占用2个GPU，依次处理36个任务 #

适用于每个任务需要用到2个GPU加速的场景，同一时间仅运行1个任务，占满2个GPU。

Snakefile配置 #

# 全局声明工作站的GPU总数量
resources:
    nvidia_gpu=2

rule a:
    output:
        "test_{idx}.txt"  # 通过wildcard生成36个任务实例
    resources:
        nvidia_gpu=2  # 当前任务申请2个GPU资源
    shell:
        # 手动指定使用的GPU编号（根据实际硬件调整，比如0和1）
        "CUDA_VISIBLE_DEVICES=0,1 your_command --gpu-count 2 --output {output}"

# 定义总目标：生成36个任务的输出文件
rule all:
    input:
        expand("test_{idx}.txt", idx=range(36))

运行命令 #

snakemake --cores 2 --resources nvidia_gpu=2 --set-resource-scopes nvidia_gpu=global

场景2：36个任务并行执行，每个任务占用1个GPU，同时运行2个任务 #

适用于单个任务仅需1个GPU的场景，最大化GPU利用率，同时跑2个任务。

Snakefile配置 #

# 全局声明工作站的GPU总数量
resources:
    nvidia_gpu=2

rule a:
    output:
        "test_{idx}.txt"
    resources:
        nvidia_gpu=1  # 单个任务申请1个GPU资源
    shell:
        # 直接使用Snakemake分配的GPU编号，无需手动指定
        "CUDA_VISIBLE_DEVICES={resources.nvidia_gpu} your_command --gpu-count 1 --output {output}"

rule all:
    input:
        expand("test_{idx}.txt", idx=range(36))

运行命令 #

snakemake --cores 2 --resources nvidia_gpu=2 --set-resource-scopes nvidia_gpu=global

关键注意事项 #

• 必须在Snakefile开头全局声明resources: nvidia_gpu=2，明确GPU总资源量，否则Snakemake无法正确调度。
• shell命令中需要通过CUDA_VISIBLE_DEVICES指定GPU，或者在你的任务命令中加入GPU编号参数，Snakemake仅负责资源分配，不会自动帮命令绑定GPU。
• 运行前可执行snakemake --list-resources确认nvidia_gpu资源已被识别。

内容的提问来源于stack exchange，提问作者sunxd