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

Posted 2023-03-27

【中文标题】当程序在 SLURM 集群上学习时收到 TIME LIMIT 时保存模型权重

【英文标题】：Save a model weights when a program receives TIME LIMIT while learning on a SLURM cluster

【发布时间】：2021-11-25 05:00:30

【问题描述】：

我使用用 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 以在可用时间结束时保存模型？

我在每个纪元之后都知道automatic saving，但我只有一个持续时间超过 10 小时的长纪元，所以这种情况对我不起作用。

【参考方案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，我更喜欢编写独立的提交脚本:)）

编辑：This link 对信号传播和 Slurm 所涉及的问题进行了很好的总结。

【讨论】：

srun 会以交互模式还是批处理模式启动我的程序？

通过批处理脚本，它将是批处理模式。我不确定信号处理在交互模式下的工作方式是否相同。