Moves hpc auto-resubmit to trainer from test-tube (#207)

* added slurm signal handler * added restore weight functions * set slurm signal handling inside process * added resubmit docs * added resubmit docs * fixed missing param * Update trainer.py * fixed missing param * fixed missing param * debugging tests * debugging tests * debugging tests * debugging tests * debugging tests * debugging tests * debugging tests
2026-09-10 12:21:57 +08:00 · 2019-09-06 11:54:51 -04:00
parent 7ed928dfac
commit 60633eaa32
4 changed files with 139 additions and 67 deletions
@@ -1,8 +1,10 @@
 Lightning supports model training on a cluster managed by SLURM in the following cases:    

-1. Training on single or multi-cpus only.
-2. Training on single or multi-gpus on the same node.
-3. Coming SOON: Training across multiple nodes.
+1. Training on a single cpu or single GPU.
+2. Train on multiple GPUs on the same node using DataParallel or DistributedDataParallel
+3. Training across multiple GPUs on multiple different nodes via DistributedDataParallel.
+
+**Note: A node means a machine with multiple GPUs**

 ---
 #### Running grid search on a cluster
@@ -55,8 +57,8 @@ cluster.memory_mb_per_node = 10000
 cluster.job_time = '10:00'
 ```

-(3). Give trainer the cluster_manager in your main function:    
-
+(3). Make a main function with your model and trainer. Each job will call this function with a particular
+hparams configuration.    
 ```{.python}
 from pytorch_lightning import Trainer

@@ -66,12 +68,12 @@ def train_fx(trial_hparams, cluster_manager, _):
    my_model = MyLightningModel()
    
    # give the trainer the cluster object
-    trainer = Trainer(cluster=cluster_manager)
+    trainer = Trainer()
    trainer.fit(my_model)

 ```

-(4). Start the grid search     
+(3). Start the grid/random search     
 ```{.python}
 # run the models on the cluster
 cluster.optimize_parallel_cluster_gpu(
@@ -81,24 +83,22 @@ cluster.optimize_parallel_cluster_gpu(
    job_display_name='my_exp')
 ```

-That's it! The SlurmCluster object will automatically checkpoint the lightning model and resubmit if it runs into the walltime!
-
-
 ---
 #### Walltime auto-resubmit
-Lightning automatically resubmits jobs when they reach the walltime. You get this behavior for free if you give lightning
-a slurm cluster object.
+Lightning automatically resubmits jobs when they reach the walltime. Make sure to set the SIGUSR1 signal in 
+your SLURM script.   

-```{.python}
-def my_main_fx(hparams, slurm_manager, _):
-    trainer = Trainer(cluster=slurm_manager)
+```bash
+# 90 seconds before training ends
+#SBATCH --signal=SIGUSR1@90
 ``` 

-(See the grid search example above for cluster configuration).
-With this feature lightning will:    
+When lightning receives the SIGUSR1 signal it will:
+1. save a checkpoint with 'hpc_ckpt' in the name.
+2. resubmit the job using the SLURM_JOB_ID  
+
+When the script starts again, Lightning will:
+1. search for a 'hpc_ckpt' checkpoint. 
+2. restore the model, optimizers, schedulers, epoch, etc...   

-1. automatically checkpoint the model
-2. checkpoint the trainer session
-3. resubmit a continuation job.
-4. load the checkpoint and trainer session in the new model

@@ -58,7 +58,6 @@ class Trainer(TrainerIO):
                 early_stop_callback=None,
                 checkpoint_callback=None,
                 gradient_clip=0,
-                 cluster=None,
                 process_position=0,
                 current_gpu_name=0,
                 nb_gpu_nodes=1,
@@ -90,7 +89,6 @@ class Trainer(TrainerIO):
        :param early_stop_callback: from pytorch_lightning import EarlyStopping
        :param checkpoint_callback: from pytorch_lightning import Checkpoint
        :param gradient_clip:
-        :param cluster:
        :param process_position:
        :param current_gpu_name:
        :param nb_gpu_nodes:
@@ -181,7 +179,6 @@ class Trainer(TrainerIO):
        self.use_ddp = False
        self.use_dp = False
        self.single_gpu = False
-        self.cluster = cluster
        self.__set_distributed_mode(distributed_backend, nb_gpu_nodes)

        # init flags for SLURM+ddp to work
@@ -816,6 +813,9 @@ class Trainer(TrainerIO):
        ref_model.use_amp = self.use_amp
        ref_model.testing = self.testing

+        # register auto-resubmit when on SLURM
+        self.register_slurm_signal_handlers()
+
        # transfer data loaders from model
        self.get_dataloaders(ref_model)

@@ -839,13 +839,7 @@ class Trainer(TrainerIO):
        self.model = model

        # restore training and model before hpc call
-        self.restore_state_if_existing_checkpoint()
-
-        # enable cluster checkpointing
-        # also restores training state
-        # hpc checkpoint overrides any other checkpoints loaded before
-        if self.cluster is not None:  # pragma: no cover
-            self.enable_auto_hpc_walltime_manager()
+        self.restore_weights(model)

        # progress bar init
        if self.show_progress_bar:
@@ -1,5 +1,7 @@
 import os
 import re
+import signal
+from subprocess import call

 import torch

@@ -51,6 +53,96 @@ class TrainerIO(object):
        model = self.model.module if is_dp_module else self.model
        return model

+    # --------------------
+    # CHECK-POINTING
+    # --------------------
+    def restore_weights(self, model):
+        """
+        To restore weights we have two cases.
+        First, if we use the same experiment version, then restore the latest ckpt.
+        AFTER that, if we find weights from hpc checkpoint, then restore that.
+        :param model:
+        :return:
+        """
+        # do nothing if there's not dir or callback
+        no_ckpt_callback = self.checkpoint_callback is None
+        if no_ckpt_callback or not os.path.exists(self.checkpoint_callback.filepath):
+            return
+
+        # restore weights if same exp version
+        self.restore_state_if_checkpoint_exists(model)
+
+        # if script called from hpc resubmit, load weights
+        self.restore_hpc_weights_if_needed(model)
+
+    def restore_state_if_checkpoint_exists(self, model):
+        # restore trainer state and model if there is a weight for this experiment
+        last_epoch = -1
+        last_ckpt_name = None
+
+        # find last epoch
+        checkpoints = os.listdir(self.checkpoint_callback.filepath)
+        for name in checkpoints:
+            # ignore hpc ckpts
+            if 'hpc_' in name:
+                continue
+
+            if '.ckpt' in name:
+                epoch = name.split('epoch_')[1]
+                epoch = int(re.sub('[^0-9]', '', epoch))
+
+                if epoch > last_epoch:
+                    last_epoch = epoch
+                    last_ckpt_name = name
+
+        # restore last checkpoint
+        if last_ckpt_name is not None:
+            last_ckpt_path = os.path.join(self.checkpoint_callback.filepath, last_ckpt_name)
+            self.restore(last_ckpt_path, self.on_gpu)
+            print(f'model and trainer restored from checkpoint: {last_ckpt_path}')
+
+    # --------------------
+    # HPC SIGNAL HANDLING
+    # --------------------
+    def register_slurm_signal_handlers(self):
+        # see if we're using slurm (not interactive)
+        on_slurm = False
+        try:
+            job_name = os.environ['SLURM_JOB_NAME']
+            if job_name != 'bash':
+                on_slurm = True
+        except Exception as e:
+            pass
+
+        if on_slurm and self.proc_rank == 0:
+            print('set slurm handle signals')
+            signal.signal(signal.SIGUSR1, self.sig_handler)
+            signal.signal(signal.SIGTERM, self.term_handler)
+
+    def sig_handler(self, signum, frame):
+        if self.proc_rank == 0:
+            # save weights
+            print('handling SIGUSR1')
+            self.hpc_save(self.checkpoint_callback.filepath, self.experiment)
+
+            # find job id
+            job_id = os.environ['SLURM_JOB_ID']
+            cmd = 'scontrol requeue {}'.format(job_id)
+
+            # requeue job
+            print('\nrequeing job {}...'.format(job_id))
+            result = call(cmd, shell=True)
+
+            # print result text
+            if result == 0:
+                print('requeued exp ', job_id)
+            else:
+                print('requeue failed...')
+
+    def term_handler(self, signum, frame):
+        # save
+        print("bypassing sigterm")
+
    # --------------------
    # MODEL SAVE CHECKPOINT
    # --------------------
@@ -116,28 +208,21 @@ class TrainerIO(object):
    # --------------------
    # HPC IO
    # --------------------
-    def enable_auto_hpc_walltime_manager(self):
-        if self.cluster is None:
-            return
+    def restore_hpc_weights_if_needed(self, model):
+        """
+        If there is a set of hpc weights, use as signal to restore model
+        :param model:
+        :return:
+        """
+        # look for hpc weights
+        folderpath = self.checkpoint_callback.filepath
+        if os.path.exists(folderpath):
+            files = os.listdir(folderpath)
+            hpc_weight_paths = [x for x in files if 'hpc_ckpt' in x]

-        # allow test tube to handle model check pointing automatically
-        # only if proc 0 so we don't trigger world_size resubmits
-        if self.proc_rank == 0:
-            self.cluster.set_checkpoint_save_function(
-                self.hpc_save,
-                kwargs={
-                    'folderpath': self.checkpoint_callback.filepath,
-                    'experiment': self.experiment
-                }
-            )
-
-        self.cluster.set_checkpoint_load_function(
-            self.hpc_load,
-            kwargs={
-                'folderpath': self.checkpoint_callback.filepath,
-                'on_gpu': self.on_gpu
-            }
-        )
+            # if hpc weights exist restore model
+            if len(hpc_weight_paths) > 0:
+                self.hpc_load(folderpath, self.on_gpu)

    def restore_training_state(self, checkpoint):
        """
@@ -6,7 +6,7 @@ from argparse import Namespace
 import pytest
 import numpy as np
 import torch
-from test_tube import Experiment, SlurmCluster
+from test_tube import Experiment

 # sys.path += [os.path.abspath('..'), os.path.abspath('../..')]
 from pytorch_lightning import Trainer
@@ -465,7 +465,6 @@ def test_no_val_module():

    trainer_options = dict(
        max_nb_epochs=1,
-        cluster=SlurmCluster(),
        experiment=exp,
        checkpoint_callback=ModelCheckpoint(save_dir)
    )
@@ -512,7 +511,6 @@ def test_no_val_end_module():

    trainer_options = dict(
        max_nb_epochs=1,
-        cluster=SlurmCluster(),
        experiment=exp,
        checkpoint_callback=ModelCheckpoint(save_dir)
    )
@@ -699,10 +697,10 @@ def test_cpu_slurm_save_load():
    exp.argparse(hparams)
    exp.save()

-    cluster_a = SlurmCluster()
+    version = exp.version
+
    trainer_options = dict(
        max_nb_epochs=1,
-        cluster=cluster_a,
        experiment=exp,
        checkpoint_callback=ModelCheckpoint(save_dir)
    )
@@ -726,22 +724,18 @@ def test_cpu_slurm_save_load():
    model.eval()
    pred_before_saving = model(x)

-    # test registering a save function
-    trainer.enable_auto_hpc_walltime_manager()
-
    # test HPC saving
    # simulate snapshot on slurm
    saved_filepath = trainer.hpc_save(save_dir, exp)
    assert os.path.exists(saved_filepath)

-    # wipe-out trainer and model
-    # retrain with not much data... this simulates picking training back up after slurm
-    # we want to see if the weights come back correctly
-    continue_tng_hparams = get_hparams(continue_training=True,
-                                       hpc_exp_number=cluster_a.hpc_exp_number)
+    # new exp file to get meta
+    exp = get_exp(False, version=version)
+    exp.argparse(hparams)
+    exp.save()
+
    trainer_options = dict(
        max_nb_epochs=1,
-        cluster=SlurmCluster(continue_tng_hparams),
        experiment=exp,
        checkpoint_callback=ModelCheckpoint(save_dir),
    )
@@ -822,7 +816,6 @@ def test_model_saving_loading():

    trainer_options = dict(
        max_nb_epochs=1,
-        cluster=SlurmCluster(),
        experiment=exp,
        checkpoint_callback=ModelCheckpoint(save_dir)
    )