From 1a835969a6d76e31aed12e4abdbf0b9267b8f01c Mon Sep 17 00:00:00 2001
From: William Falcon <falc@fb.com>
Date: Fri, 26 Jul 2019 12:14:58 -0400
Subject: [PATCH] added saving tests to cpu

---
 pytorch_lightning/root_module/model_saving.py |  4 +-
 tests/test_models.py                          | 39 ++++++++++++-------
 2 files changed, 27 insertions(+), 16 deletions(-)

diff --git a/pytorch_lightning/root_module/model_saving.py b/pytorch_lightning/root_module/model_saving.py
index 3d242ada..25377851 100644
--- a/pytorch_lightning/root_module/model_saving.py
+++ b/pytorch_lightning/root_module/model_saving.py
@@ -86,7 +86,7 @@ class TrainerIO(object):
     # --------------------
     # HPC IO
     # --------------------
-    def enable_auto_hpc_walltime_manager(self):  # pragma: no cover
+    def enable_auto_hpc_walltime_manager(self):
         if self.cluster is None:
             return
 
@@ -157,6 +157,8 @@ class TrainerIO(object):
         # do the actual save
         torch.save(checkpoint_dict, filepath)
 
+        return filepath
+
     def hpc_load(self, folderpath, on_gpu):
         filepath = '{}/hpc_ckpt_{}.ckpt'.format(folderpath, self.max_ckpt_in_folder(folderpath))
 
diff --git a/tests/test_models.py b/tests/test_models.py
index fd214009..455111d4 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -3,7 +3,7 @@ from pytorch_lightning import Trainer
 from pytorch_lightning.examples.new_project_templates.lightning_module_template import LightningTemplateModel
 from pytorch_lightning.testing_models.lm_test_module import LightningTestModel
 from argparse import Namespace
-from test_tube import Experiment
+from test_tube import Experiment, SlurmCluster
 from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping
 from pytorch_lightning.utils.debugging import MisconfigurationException
 from pytorch_lightning.root_module import memory
@@ -43,7 +43,7 @@ def test_dp_output_reduce():
     assert reduced['b']['c'] == out['b']['c']
 
 
-def test_cpu_slurm_managed():
+def test_cpu_slurm_saving_loading():
     """
     Verify model save/load/checkpoint on CPU
     :return:
@@ -51,10 +51,6 @@ def test_cpu_slurm_managed():
     hparams = get_hparams()
     model = LightningTestModel(hparams)
 
-    trainer_options = dict(
-        max_nb_epochs=1,
-    )
-
     save_dir = init_save_dir()
 
     # exp file to get meta
@@ -62,20 +58,28 @@ def test_cpu_slurm_managed():
     exp.argparse(hparams)
     exp.save()
 
-    # exp file to get weights
-    checkpoint = ModelCheckpoint(save_dir)
-
-    # add these to the trainer options
-    trainer_options['checkpoint_callback'] = checkpoint
-    trainer_options['experiment'] = exp
+    trainer_options = dict(
+        max_nb_epochs=1,
+        cluster=SlurmCluster(),
+        experiment=exp,
+        checkpoint_callback=ModelCheckpoint(save_dir)
+    )
 
     # fit model
     trainer = Trainer(**trainer_options)
     result = trainer.fit(model)
+    real_global_step = trainer.global_step
 
-    # correct result and ok accuracy
+    # traning complete
     assert result == 1, 'amp + ddp model failed to complete'
 
+    # test saving checkpoint
+    ckpt_test = os.path.join(save_dir, 'test.ckpt')
+    trainer.save_checkpoint(ckpt_test)
+
+    # test registering a save function
+    trainer.enable_auto_hpc_walltime_manager()
+
     # test model loading with a map_location
     pretrained_model = load_model(exp, save_dir, True)
 
@@ -85,9 +89,14 @@ def test_cpu_slurm_managed():
     trainer.model = pretrained_model
     trainer.optimizers = pretrained_model.configure_optimizers()
 
-    # test HPC loading / saving
-    trainer.hpc_save(save_dir, exp)
+    # test HPC saving
+    saved_filepath = trainer.hpc_save(save_dir, exp)
+    assert os.path.exists(saved_filepath)
+
+    # test HPC loading
+    trainer.global_step = 20000000
     trainer.hpc_load(save_dir, on_gpu=False)
+    assert trainer.global_step == real_global_step and trainer.global_step != 20000000
 
     # test freeze on gpu
     model.freeze()