Fragment of python code for catching signals from Slurm and restarting the job

handle_slurm_signals.py

import signal

# depends on requesting SIGUSR1 in runner file: https://gist.github.com/willwhitney/e1509c86522896c6930d2fe9ea49a522

def handle_signal(signal_value, _):
    signame = signal.Signals(signal_value).name
    if signal_value == signal.SIGUSR1:
        print('Process {} got signal {}. Saving and restarting.'.format(
                os.getpid(), signame), flush=True)
        save_dynamics(epoch)
        if opt.restart_command is not None:
            os.system(opt.restart_command)
        env.close()
        if mjdynamics is not None:
            mjdynamics.close()
        sys.exit(0)
    else:
        print('Process {} got signal {}. Doing nothing.'.format(
                os.getpid(), signame), flush=True)

signal.signal(signal.SIGUSR1, handle_signal)
signal.signal(signal.SIGCONT, handle_signal)
signal.signal(signal.SIGTERM, handle_signal)

Interesting piece of code. Thanks for sharing.