Make the Slurm Signal Configurable

[Queuecumber]

🚀 Feature

Lightning currently assumes that the slurm preemption signal is set to USR1, it would be nice to add an option to SLURMEnvironment to change this.

Motivation

I use submitit to handle slurm submissions and they recently moved away from USR1 because of an NCCL conflict.

This breaks the requeuing in pytorch lightning.

See also facebookincubator/submitit#1709

Pitch

SLURMEnvironment can be trivially extended to account for this with something like

def __init__(self, auto_requeue: bool = True, signal: str = "USR1") -> None:
        super().__init__()
        self.auto_requeue = auto_requeue
        self.signal = signal

then SignalConnector ~ln 60 can be changed to something like:

sig = getattr(signal, "SIG" + environment.signal, None)
self._register_signal(sig, HandlersCompose(sigusr1_handlers))

cc @Borda @awaelchli

[Queuecumber]

Happy to contribute this if it looks alright BTW

[awaelchli]

Hi @Queuecumber

I think what you propose is very reasonable. If you want to contribute it, that would be great and I will be happy to take a look at your PR.

[krshrimali]

Hi, @Queuecumber - Thanks for showing the interest. I assume you want to contribute, so I've taken the liberty to assign you to this issue. If you are stuck, or have any questions, please feel free to ask them there. :)