如何处理 Slurm 中的作业取消？

Question

我在 HPC 集群上使用 Slurm 作业管理器。有时会有这样的情况，当一个工作因为时间限制而被取消时，我想优雅地完成我的程序。

据我了解，取消过程分两个阶段进行，正是为了让软件开发人员能够优雅地完成程序：

srun: Job step aborted: Waiting up to 62 seconds for job step to finish.                                                                                                                           
slurmstepd: error: *** JOB 18522559 ON ncm0317 CANCELLED AT 2020-12-14T19:42:43 DUE TO TIME LIMIT ***

你可以看到我有 62 秒的时间按照我希望的方式完成工作（通过保存一些文件等）。

问题：如何做到这一点？我知道首先一些 Unix 信号被发送到我的工作，我需要正确地响应它。但是，我无法在 Slurm 文档中找到有关此信号是什么的任何信息。此外，我在 Python 中没有具体如何处理它，可能是通过异常处理。

Answer 1

在大多数编程语言中，Unix 信号是使用回调来捕获的。 Python也不例外。要使用 Python 捕获 Unix 信号，只需使用 signal package.

例如，要优雅地退出：

import signal, sys
def terminate_signal(signalnum, handler):
    print ('Terminate the process')
    # save results, whatever...
    sys.exit()

# initialize signal with a callback
signal.signal(signal.SIGTERM, terminate_signal)

while True:
    pass  # work

possible signals 列表。 SIGTERM 是用来“礼貌地要求程序终止”的那个。

Answer 2

在 Slurm 中，您可以决定在作业达到时间限制之前的哪个时刻发送哪个信号。

来自sbatch man page：

--signal=[[R][B]:]<sig_num>[@<sig_time>] When a job is within sig_time seconds of its end time, send it the signal sig_num.

如此设置

#SBATCH --signal=B:TERM@05:00

让 Slurm 在分配结束前 5 分钟用 SIGTERM 发出作业信号。请注意，根据您开始工作的方式，您可能需要删除 B: 部分。

在您的 Python 脚本中，使用 signal 包。您需要定义一个“信号处理程序”，一个在接收到信号时将调用的函数，并为特定信号“注册”该函数。由于该函数在调用时会扰乱正常流程，因此您需要使其保持简短以避免不必要的副作用，尤其是对于多线程代码。

Slurm 环境中的一个典型方案是具有如下脚本框架：

#! /bin/env python

import signal, os, sys

# Global Boolean variable that indicates that a signal has been received
interrupted = False

# Global Boolean variable that indicates then natural end of the computations
converged = False

# Definition of the signal handler. All it does is flip the 'interrupted' variable
def signal_handler(signum, frame):
    global interrupted
    interrupted = True

# Register the signal handler
signal.signal(signal.SIGTERM, signal_handler)

try:
    # Try to recover a state file with the relevant variables stored
    # from previous stop if any
    with open('state', 'r') as file: 
        vars = file.read()
except:
    # Otherwise bootstrap (start from scratch)
    vars = init_computation()

while not interrupted and not converged:
    do_computation_iteration()    

# Save current state 
if interrupted:
    with open('state', 'w') as file: 
        file.write(vars)
    sys.exit(99)
sys.exit(0)

这首先尝试重新启动作业的前一个运行留下的计算，否则将引导它。如果它被中断，它会让当前循环迭代正确完成，然后将需要的变量保存到磁盘。然后它以 99 return 代码退出。如果为它配置了 Slurm，这允许自动重新排队作业以进行进一步迭代。

如果没有为它配置 slurm，您可以在提交脚本中手动配置，如下所示：

python myscript.py || scontrol requeue $SLURM_JOB_ID

如何处理 Slurm 中的作业取消？

How to handle job cancelation in Slurm?

python

signals

slurm