如何在第二次尝试蛇规则时将内存加倍

Question

我试图弄清楚如何自动将给定规则的内存加倍。我把内存和CPU分配放在配置文件中：

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

这是我的测试规则：

rule Test:
    input:
        "A.txt"
    output:
        "B.txt"
    params:
        rulename = "test",
        batch = config["test"],
     shell: """
         ......
     """

这是我的snakemake命令：

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

我想知道规则或配置文件中是否有一种方法可以在第一次失败后在第二次尝试中自动将内存加倍，类似于配置文件中的以下内容：

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

或者规则中类似以下内容：

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: """
         ......
     """

我已经阅读了snakemake文档和在线论坛，但没有找到答案。

Answer 1

文档有几个使用 lambda 或输入函数通过连续尝试增加内存的示例，但它们有一个固定的起点。我认为您是在询问一种更自动地执行此操作的机制，但该机制似乎尚未实现。您可以像此问题中那样更改 sbatch 命令。

我会使用自定义函数，例如：

def increment_memory(base_memory):
    def mem(wildcards, attempt):
        return base_memory * (2 ** (attempt - 1))
    return mem

rule Test:
    input:
        "A.txt"
    output:
        "B.txt"
    params:
        rulename = "test",
        batch = config["test"],
    resources:
        mem_mb=increment_memory(20_000)
    ...

如果您需要多个规则中的逻辑，这会减少样板文件，并且易于添加/删除。您必须以整数形式提供基本内存，您需要额外的逻辑来解析和处理像

20G

这样的字符串。我选择在每次重试时将内存加倍，例如重试 3 将使用 4 倍的内存。如果您不想重试每个规则，另请注意 retries 指令。

希望有帮助！