Snakemake规则重试时如何自动翻倍内存配置？

需求 #

为指定Snakemake规则实现失败重试时自动翻倍内存的功能，当前已配置基础内存与CPU分配，希望首次失败后第二次重试内存自动翻倍，第三次重试再翻倍，以此类推。

当前配置与规则 #

原配置文件 #

Test: "--mem=20G --cpus-per-task=10"

原测试规则 #

rule Test:
    input:
        "A.txt"
    output:
        "B.txt"
    params:
        rulename = "test",
        batch = config["Test"],
    shell: """
        # 业务执行命令
    """

原执行命令 #

snakemake \
    --cores 100 \
    --retries 2 \
    --cluster 'sbatch {params.batch}'

尝试过的无效方案 #

1. 配置文件中尝试通过set-resources关联重试次数与内存：

set-resources:
    Test:
        mem_mb: attempt * resources.mem_gb

2. 规则内定义资源函数尝试计算内存：

def get_mem_mb(wildcards, attempt):
    return attempt * resources.mem_gb

rule Test:
    input:
        "A.txt"
    output:
        "B.txt"
    params:
        rulename = "test",
        batch = config["Test"],
    resources:
        mem_mb=get_mem_mb
    shell: """
        # 业务执行命令
    """

可行实现方案 #

方案一：动态生成batch参数（适配原配置结构） #

1. 修改配置文件，拆分内存与CPU为独立配置项：

Test:
  base_mem_gb: 20
  cpus: 10

2. 在规则中定义动态生成batch参数的函数，利用Snakemake内置的attempt变量（重试次数，从1开始）计算内存：

def get_batch(wildcards, attempt, config):
    # 每次重试时内存翻倍
    current_mem = config["Test"]["base_mem_gb"] * attempt
    return f"--mem={current_mem}G --cpus-per-task={config['Test']['cpus']}"

rule Test:
    input:
        "A.txt"
    output:
        "B.txt"
    params:
        rulename = "test",
        batch = get_batch  # 绑定动态生成函数
    shell: """
        # 业务执行命令
    """

3. 执行命令保持不变：

snakemake \
    --cores 100 \
    --retries 2 \
    --cluster 'sbatch {params.batch}'

方案二：利用Snakemake资源管理机制（更规范） #

1. 配置文件与set-resources结合：

Test:
  cpus: 10
set-resources:
    Test:
        # 每次重试时内存为基础值乘以重试次数
        mem_gb: lambda wildcards, attempt: 20 * attempt

2. 规则中定义资源项，不再依赖params传递batch：

rule Test:
    input:
        "A.txt"
    output:
        "B.txt"
    params:
        rulename = "test",
    resources:
        mem_gb=20,  # 基础内存值
        cpus=config["Test"]["cpus"]
    shell: """
        # 业务执行命令
    """

3. 修改执行命令，直接从resources渲染集群参数：

snakemake \
    --cores 100 \
    --retries 2 \
    --cluster 'sbatch --mem={resources.mem_gb}G --cpus-per-task={resources.cpus}'

说明 #

• attempt是Snakemake内置变量，第一次运行值为1，第一次重试值为2，第二次重试值为3，以此类推
• 两种方案均能实现“重试时内存自动翻倍”的需求，方案二更贴合Snakemake的资源管理设计，推荐使用

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