如何用Snakemake加速4万样本数据处理并行任务及优化内存泄漏

环境背景 #

• 硬件：48核96线程、256GB内存
• 系统：Ubuntu 20.04
• Python版本：3.9
• 任务：处理约4万个样本文件，原串行脚本python a.py -i sample_list.txt -o /path/to/outdir预计耗时240天，需并行优化

一、Snakemake基于sample_list.txt实现任务并行 #

完全可以基于sample_list.txt实现任务并行，无需依赖输入文件模式，核心思路是将每个样本ID作为独立任务的wildcard。具体实现步骤如下：

1. 编写Snakefile

# 读取sample_list.txt中的所有样本ID
SAMPLES = [line.strip() for line in open("sample_list.txt") if line.strip()]

# 定义最终目标：所有样本的输出文件
rule all:
    input:
        expand("/path/to/outdir/{sample}_result.txt", sample=SAMPLES)

# 单个样本处理规则
rule process_sample:
    output:
        "/path/to/outdir/{sample}_result.txt"
    shell:
        """
        python a.py -i {wildcards.sample} -o /path/to/outdir
        """

注：这里假设原脚本a.py支持直接传入单个样本ID作为-i参数（如果原脚本只接受文件列表，需要修改a.py新增单样本处理逻辑，或者写一个包装脚本调用原逻辑处理单个样本）

2. 启动并行任务
   根据机器核心数，指定并行进程数（比如用40核，留部分资源给系统）：

snakemake --cores 40

Snakemake会自动将每个样本作为独立任务调度，实现并行执行，同时还能处理任务失败重试、断点续跑等问题。

二、Python multiprocessing内存泄漏优化方案 #

针对多进程处理时的内存泄漏问题，可按以下方案逐步优化：

• 限制单个进程的任务数
  使用multiprocessing.Pool时设置maxtasksperchild=1，让每个进程只处理一个样本就销毁重建，从根源避免进程长期运行累积内存：

  from multiprocessing import Pool
  
  def process_sample(sample_id):
      # 单个样本处理逻辑，确保处理完后释放资源
      import gc
      # 执行处理代码
      result = run_processing(sample_id)
      # 显式删除大对象并触发垃圾回收
      del result
      gc.collect()
      return
  
  if __name__ == "__main__":
      samples = [line.strip() for line in open("sample_list.txt")]
      # 设置maxtasksperchild=1，每个进程只处理一个任务
      with Pool(processes=40, maxtasksperchild=1) as pool:
          pool.map(process_sample, samples)
  

• 排查并修复脚本本身的内存泄漏
  使用memory_profiler工具分析单个样本处理过程的内存变化，定位未释放的对象：

  pip install memory-profiler
  

  在a.py的关键函数上添加装饰器@profile，然后运行单样本测试：

  python -m memory_profiler a.py -i single_sample_id -o /path/to/outdir
  

  根据输出找到内存持续增长的代码段，比如未关闭的文件句柄、全局变量累积、未释放的第三方库对象等，针对性修复。

• 避免子进程共享大对象
  不要在主进程中加载大模型、大数据集后传递给子进程，而是让每个子进程单独加载所需资源（比如每个子进程处理样本时才读取必要的配置文件或小数据集），减少内存拷贝和累积。

• 改用更稳定的进程管理工具
  可以尝试使用concurrent.futures.ProcessPoolExecutor替代原生multiprocessing.Pool，其内存管理更规范，同时支持异步任务调度：

  from concurrent.futures import ProcessPoolExecutor
  
  def process_sample(sample_id):
      # 单个样本处理逻辑
      pass
  
  if __name__ == "__main__":
      samples = [line.strip() for line in open("sample_list.txt")]
      with ProcessPoolExecutor(max_workers=40) as executor:
          executor.map(process_sample, samples)
  

• 调整并行进程数
  虽然机器有48核，但如果内存泄漏问题未完全解决，可先降低并行数（比如从11提升到20左右，逐步测试），避免内存占用过高导致系统崩溃，同时结合maxtasksperchild参数使用。

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