Multiprocessing Pool用map/async时最后几个进程挂起问题求助

可能的原因 #

• 任务负载不均：最后几个分片的数据量远大于前面的任务，导致单个任务耗时远超1-2分钟的平均水平。
• Polars资源残留：spawn上下文虽不共享进程资源，但Polars的DataFrame或内部状态未被正确释放，进程复用后积累内存/文件句柄泄漏，最后几个任务时进程资源耗尽。
• 结果处理阻塞：主进程在处理任务结果时（如imap_unordered循环内的统计操作、apply_async回调里的glob遍历）执行耗时操作，导致无法及时响应任务完成通知，看起来像是任务卡住。
• 信号量泄漏与进程异常：任务中存在未捕获的异常、未关闭的资源，导致进程池的信号量资源未正确释放，进程状态异常，后续任务无法正常调度。

针对性解决办法 #

1. 验证任务负载是否均衡 #

给任务函数添加日志，记录每个分片的大小和执行时间，确认最后几个任务是否真的是数据量过大：

import logging
import os
import time

logging.basicConfig(level=logging.INFO)

def get_edit_info_for_barcode_in_contig_wrapper(params):
    task_id = params.get("task_id", "unknown")
    start_time = time.perf_counter()
    # 记录分片文件大小（如果参数包含文件路径）
    if "input_path" in params:
        file_size = os.path.getsize(params["input_path"])
        logging.info(f"Task {task_id} started, input size: {file_size / 1024 / 1024:.2f} MB")
    # 原有业务逻辑
    # ...
    elapsed = time.perf_counter() - start_time
    logging.info(f"Task {task_id} finished, elapsed: {elapsed:.2f}s")
    return None

如果确实是分片大小不均，可以将大分片拆分为更小的子任务，或者单独分配更多资源处理。

2. 清理Polars资源，避免泄漏 #

在任务函数中显式释放Polars对象并触发垃圾回收，防止进程复用后资源积累：

import gc
import polars as pl

def get_edit_info_for_barcode_in_contig_wrapper(params):
    # 原有逻辑生成DataFrame
    df = pl.read_csv(params["input_path"])
    # 处理逻辑
    # ...
    # 写入文件后立即释放资源
    df.write_csv(params["output_path"])
    del df
    gc.collect()
    return None

同时避免在任务函数中使用全局的Polars配置或对象，每个任务独立初始化必要设置。

3. 优化主进程的结果处理逻辑 #

• 简化imap_unordered循环操作：将非必要的统计操作移到任务全部完成后执行，避免阻塞主进程接收结果：

with get_context("spawn").Pool(processes=processes) as p:
    max_ = len(coverage_counting_job_params)
    print(f"max_ is {max_}")
    start_time = time.perf_counter()
    with tqdm(total=max_) as pbar:
        # 仅处理进度条更新
        for _ in p.imap_unordered(get_edit_info_for_barcode_in_contig_wrapper, coverage_counting_job_params):
            pbar.update()
# 任务全部完成后再做统计
total_time = time.perf_counter() - start_time
print(f"Total elapsed time: {total_time / 60:.2f} mins")

• 替换glob为高效文件检查：在apply_async回调中，预先生成所有预期的输出文件路径，避免每次遍历目录：

# 预先生成所有预期输出文件路径
expected_files = [
    os.path.join(coverage_processing_folder, f"{param['task_id']}.tsv")
    for param in coverage_counting_job_params
]
max_ = len(expected_files)

def update(result):
    pbar.update()

# 提交所有任务
for param in coverage_counting_job_params:
    pool.apply_async(
        get_edit_info_for_barcode_in_contig_wrapper,
        args=(param,),
        callback=update
    )

# 等待所有任务完成后再检查文件
pool.close()
pool.join()
completed = sum(1 for f in expected_files if os.path.exists(f))
if completed == max_:
    print(f"All {max_} expected files are present!")

4. 修复信号量泄漏与进程异常 #

• 限制进程复用次数：给进程池添加maxtasksperchild参数，让每个进程处理一定数量的任务后重启，避免资源积累：

with get_context("spawn").Pool(processes=processes, maxtasksperchild=10) as p:
    # 原有逻辑

• 捕获任务异常：在任务函数中加入异常捕获，避免未处理的异常导致进程状态异常：

def get_edit_info_for_barcode_in_contig_wrapper(params):
    try:
        # 原有业务逻辑
        return True
    except Exception as e:
        task_id = params.get("task_id", "unknown")
        logging.error(f"Task {task_id} failed with error: {str(e)}")
        return False

• 检查文件操作：确保所有手动打开的文件都通过with上下文关闭，避免文件句柄泄漏：

# 正确的文件操作方式
with open(params["some_file"], "r") as f:
    content = f.read()

5. 尝试替代多进程实现 #

改用concurrent.futures.ProcessPoolExecutor，其在spawn上下文下的资源管理可能更稳定：

from concurrent.futures import ProcessPoolExecutor
import time
from tqdm import tqdm

start_time = time.perf_counter()
max_ = len(coverage_counting_job_params)

with ProcessPoolExecutor(max_workers=processes, mp_context=get_context("spawn")) as executor:
    with tqdm(total=max_) as pbar:
        for _ in executor.map(get_edit_info_for_barcode_in_contig_wrapper, coverage_counting_job_params):
            pbar.update()

total_time = time.perf_counter() - start_time
print(f"Total time: {total_time / 60:.2f} mins")

6. 系统资源监控 #

用htop或vmstat监控任务执行期间的CPU、内存、磁盘IO情况，确认是否存在内存耗尽导致swap、磁盘IO瓶颈或CPU被其他进程占用的情况。

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