From 60633eaa3214e88fa66379f21d49d60e072b5429 Mon Sep 17 00:00:00 2001 From: William Falcon Date: Fri, 6 Sep 2019 11:54:51 -0400 Subject: [PATCH] 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 --- docs/Trainer/SLURM Managed Cluster.md | 42 +++--- pytorch_lightning/models/trainer.py | 14 +- pytorch_lightning/root_module/model_saving.py | 127 +++++++++++++++--- tests/test_models.py | 23 ++-- 4 files changed, 139 insertions(+), 67 deletions(-) diff --git a/docs/Trainer/SLURM Managed Cluster.md b/docs/Trainer/SLURM Managed Cluster.md index 4ceb38bc..dc85195f 100644 --- a/docs/Trainer/SLURM Managed Cluster.md +++ b/docs/Trainer/SLURM Managed Cluster.md @@ -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 diff --git a/pytorch_lightning/models/trainer.py b/pytorch_lightning/models/trainer.py index 14bec3aa..e0e4333a 100644 --- a/pytorch_lightning/models/trainer.py +++ b/pytorch_lightning/models/trainer.py @@ -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: diff --git a/pytorch_lightning/root_module/model_saving.py b/pytorch_lightning/root_module/model_saving.py index 7d6d2c8c..a9889318 100644 --- a/pytorch_lightning/root_module/model_saving.py +++ b/pytorch_lightning/root_module/model_saving.py @@ -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): """ diff --git a/tests/test_models.py b/tests/test_models.py index 049ecf1e..e7084949 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -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) )