程序在 SLURM 集群上学习时收到 TIME LIMIT 时保存模型权重

Question

我使用用 pytorch_lightning (pytorch) 编写的深度学习模型，并在 slurm 集群上训练它们。我这样提交作业：

sbatch --gpus=1 -t 100 python train.py

当要求

GPU

时间结束时，slurm 会杀死我的程序并显示这样的消息：

Epoch 0: : 339it [01:10,  4.84it/s, loss=-34]  slurmstepd: error: *** JOB 375083 ON cn-007 CANCELLED AT 2021-10-04T22:20:54 DUE TO TIME LIMIT *** 

如何配置 Trainer 在可用时间结束时保存模型？

我知道每个纪元后都会有自动保存，但我只有一个持续> 10小时的长纪元，所以这种情况不适合我。

Answer 1

您可以使用 Slurm 的信号机制在时间限制的特定秒数内将信号传递给您的应用程序（请参阅

man sbatch

）。在您的提交脚本中，使用

--signal=USR1@30

在达到时间限制前 30 秒发送

USR1

。您的提交脚本将包含以下行：

#SBATCH -t 100
#SBATCH --signal=USR1@30
srun python train.py

然后，在您的代码中，您可以像这样处理该信号：

import signal

def handler(signum, frame):
    print('Signal handler got signal ', signum)
    # e.g. exit(0), or call your pytorch save routines

# enable the handler
signal.signal(signal.SIGUSR1, handler)

# your code here

您需要通过

srun

调用您的 Python 应用程序，以便 Slurm 能够将信号传播到 Python 进程。 （您可能可以在命令行上使用

--signal

到

sbatch

，我更喜欢编写独立的提交脚本:)）

编辑：此链接很好地总结了与信号传播和 Slurm 相关的问题。

Answer 2

如果您知道抢占不会成为问题，您也可以在初始化训练器时设置

max_time

。

您可以将其设置为静态值或从 SLURM 环境变量获取它。