如何预判MPI进程被杀死？SLURM+mpi4py场景解决方案咨询

我太懂你这种头疼了——手动拆分大任务虽然能解决问题，但每次都要盯着作业状态、手动调整，效率太低了。下面我结合SLURM和mpi4py的特性，给你一套能自动预判风险、触发任务拆分的方案。

先搞清楚信号9（SIGKILL）的本质 #

首先得明确：SIGKILL是系统级的强制终止信号，程序根本没法直接捕获或忽略它。它通常是因为你的作业占用的资源（内存、CPU时间）超过了SLURM分配的限额，或者节点资源紧张被系统强制杀掉。所以我们没法等它发出来再处理，只能提前预判风险。

方案1：用SLURM的预警信号提前触发拆分 #

SLURM自带一个实用功能：可以在作业快要达到资源限额时，提前给进程发一个自定义预警信号（默认是SIGUSR1）。我们可以捕获这个信号，在作业被强制杀掉前主动拆分任务。

具体步骤： #

1. 在SLURM提交脚本里加一行，设置预警时间（比如作业结束前5分钟发信号）：

   #SBATCH --signal=B:USR1@300  # 300秒=5分钟，可根据你的作业时长调整
   

2. 在mpi4py代码里注册SIGUSR1的处理函数，标记需要拆分：

   import signal
   import sys
   from mpi4py import MPI
   
   # 全局标记，用来通知所有进程要拆分了
   need_split = False
   
   def handle_usr1(signum, frame):
       global need_split
       need_split = True
       rank = MPI.COMM_WORLD.Get_rank()
       print(f"Rank {rank} 收到预警信号，即将触发任务拆分", file=sys.stderr)
   
   # 注册信号处理逻辑
   signal.signal(signal.SIGUSR1, handle_usr1)
   

方案2：实时监控资源使用主动触发 #

如果SLURM的预警信号不够灵活，还可以在代码里定期检查当前进程的资源使用情况，当接近SLURM分配的限额时主动拆分。

示例：监控内存使用（需要安装psutil库） #

import psutil
import os
from mpi4py import MPI

def check_memory_limit():
    comm = MPI.COMM_WORLD
    rank = comm.Get_rank()
    # 从SLURM环境变量获取分配的内存限额（单位MB，转成KB方便计算）
    slurm_mem_limit = int(os.environ.get('SLURM_MEM_PER_NODE', 0)) * 1024
    if slurm_mem_limit == 0:
        return False  # 没获取到限额就不检查
    
    # 获取当前进程的内存使用量（KB）
    current_mem = psutil.Process().memory_info().rss
    # 当使用超过80%限额时，触发拆分（比例可自己调整）
    if current_mem / slurm_mem_limit > 0.8:
        print(f"Rank {rank} 内存使用已达80%限额，准备拆分任务", file=sys.stderr)
        return True
    return False

自动拆分的核心逻辑实现 #

不管用哪种预警方式，接下来要在作业执行过程中定期检查标记，一旦触发就收集未完成任务，自动提交新的SLURM作业处理剩余部分。

简化版代码示例： #

from mpi4py import MPI
import os

def main():
    comm = MPI.COMM_WORLD
    rank = comm.Get_rank()
    size = comm.Get_size()

    # 假设这是你的大任务列表，比如每个元素是一个数据块的路径
    all_tasks = [f"data_block_{i}" for i in range(100)]
    # 把任务分配给各个进程
    tasks_for_rank = all_tasks[rank::size]

    need_split = False
    completed_tasks = []

    for task in tasks_for_rank:
        # 每次处理任务前检查是否需要拆分
        if need_split or check_memory_limit():
            need_split = True
            break
        
        # 这里写你的任务处理逻辑
        process_single_task(task)
        completed_tasks.append(task)
    
    # 收集所有进程的完成情况
    all_completed = comm.gather(completed_tasks, root=0)
    
    # 主进程（rank 0）负责处理剩余任务提交
    if rank == 0:
        completed_set = set(sum(all_completed, []))
        remaining_tasks = [t for t in all_tasks if t not in completed_set]
        
        if need_split and remaining_tasks:
            # 生成新的SLURM提交脚本并提交
            submit_remaining_job(remaining_tasks)
    
    MPI.Finalize()

def submit_remaining_job(remaining_tasks):
    # 生成临时提交脚本
    script_content = f"""#!/bin/bash
#SBATCH --job-name=auto_split_job
#SBATCH --nodes=1
#SBATCH --ntasks=4
#SBATCH --mem=16G
#SBATCH --time=1:00:00
#SBATCH --signal=B:USR1@300

# 把剩余任务作为参数传给脚本
mpirun python your_script.py --tasks {' '.join(remaining_tasks)}
"""
    with open("auto_split_submit.sh", "w") as f:
        f.write(script_content)
    # 提交作业
    os.system("sbatch auto_split_submit.sh")
    print(f"已自动提交拆分作业，剩余任务数：{len(remaining_tasks)}")

def process_single_task(task):
    # 这里替换成你的实际任务处理代码
    print(f"处理任务：{task}")

额外注意事项 #

• 先通过sacct命令查看被终止作业的资源使用记录，确认是内存还是时间超限，针对性调整预警阈值。
• 自动拆分时，子任务的粒度要足够小，避免再次触发超限。可以根据历史运行数据动态调整拆分大小。
• 触发拆分时，要确保所有进程都能正常退出，避免出现僵尸进程拖慢节点。

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