Repository: horovod/horovod
Pull request: Fix lightning and pytorch estimator loading from checkpoint
## Checklist before submitting

- [x] Did you read the [contributor guide](https://github.com/horovod/horovod/blob/master/CONTRIBUTING.md)?
- [ ] Did you update the docs?
- [ ] Did you write any tests to validate this change?  
- [ ] Did you update the [CHANGELOG](https://github.com/horovod/horovod/blob/master/CHANGELOG.md), if this change affects users?

## Description

Fixes #3268. Creating named temporary output directories which are later copied to run path, then fetching the checkpoints from this named directory in run path to ensure latest checkpoints are fetched in multiple iterations.


Problem statement:
Unable to load most recent checkpoint for Pytorch and Pytorch lightning Estimator
**Environment:**
1. Framework: PyTorch
2. Framework version: 1.8.1
3. Horovod version: 0.23.0
4. MPI version:
5. CUDA version:
6. NCCL version:
7. Python version: 3.8
8. Spark / PySpark version: 3.1.2
9. Ray version:
10. OS and version:
11. GCC version:
12. CMake version:

**Bug report:**
In case of pytorch lightning estimator, the _read_checkpoint() API does not return the latest checkpoint stored in the run path. 
Reason: Pytorch lightning estimator calls store.get_checkpoints() which looks for a folder named 'checkpoint' in run path while there is no folder named checkpoint, instead there is a temp folder generated via tempfile.TemporaryDirectory()

In case of pytorch estimator, the checkpoint stored in run path is not overwritten if multiple iterations are done using the same run path, which leads to _load_checkpoint() API returning the stale checkpoint.

Hints from issue discussion:

Commit message: Enabled checkpoint test for lightning

Signed-off-by: Kamal Sharma <kamalbhardwaj020@gmail.com>

Patch to review:
diff --git a/horovod/spark/common/store.py b/horovod/spark/common/store.py
index e39ac92992..7ba4315912 100644
--- a/horovod/spark/common/store.py
+++ b/horovod/spark/common/store.py
@@ -150,6 +150,7 @@ def _remote_attrs(self, run_id, dataset_idx):
             'checkpoint_filename': self.get_checkpoint_filename(),
             'logs_subdir': self.get_logs_subdir(),
             'get_local_output_dir': self.get_local_output_dir_fn(run_id),
+            'get_local_run_dir': self.get_local_run_dir_fn(run_id),
             'sync': self.sync_fn(run_id)
         }
 
@@ -247,12 +248,22 @@ def get_runs_path(self):
     def get_run_path(self, run_id):
         return os.path.join(self.get_runs_path(), run_id)
 
-    def get_checkpoint_path(self, run_id):
-        return os.path.join(self.get_run_path(run_id), self.get_checkpoint_filename()) \
+    def get_checkpoint_dir(self, run_id):
+        return os.path.join(self.get_run_path(run_id), self.get_checkpoint_dirname()) \
             if self._save_runs else None
 
+    def get_checkpoint_path(self, run_id):
+        if self._save_runs:
+            dirpath = self.get_checkpoint_dir(run_id)
+            if self.fs.exists(dirpath) and self.fs.isdir(dirpath):
+                return os.path.join(dirpath, self.get_checkpoint_filename())
+            else:
+                return os.path.join(self.get_run_path(run_id), self.get_checkpoint_filename())
+        else:
+            return None
+
     def get_checkpoints(self, run_id, suffix='.ckpt'):
-        checkpoint_dir = self.get_localized_path(self.get_checkpoint_path(run_id))
+        checkpoint_dir = self.get_localized_path(self.get_checkpoint_dir(run_id))
         filenames = self.fs.ls(checkpoint_dir)
         return sorted([name for name in filenames if name.endswith(suffix)])
 
@@ -262,6 +273,9 @@ def get_logs_path(self, run_id):
 
     def get_checkpoint_filename(self):
         return 'checkpoint'
+    
+    def get_checkpoint_dirname(self):
+        return 'checkpoints'
 
     def get_logs_subdir(self):
         return 'logs'
@@ -281,6 +295,17 @@ def local_run_path():
                 yield tmpdir
         return local_run_path
 
+    def get_local_run_dir_fn(self, run_id):
+        @contextlib.contextmanager
+        def local_run_path():
+            dirpath = os.path.join(tempfile._get_default_tempdir(),self.get_checkpoint_dir)
+            os.makedirs(dirpath, exist_ok=True)
+            try:
+                yield dirpath
+            finally:
+                shutil.rmtree(dirpath)
+        return local_run_path
+
     def get_localized_path(self, path):
         raise NotImplementedError()
 
diff --git a/horovod/spark/lightning/remote.py b/horovod/spark/lightning/remote.py
index dce397907e..a1f932aabe 100644
--- a/horovod/spark/lightning/remote.py
+++ b/horovod/spark/lightning/remote.py
@@ -111,7 +111,7 @@ def train(serialized_model):
         _checkpoint_callback = None
         require_checkpoint = False
 
-        with remote_store.get_local_output_dir() as run_output_dir:
+        with remote_store.get_local_run_dir() as run_output_dir:
             logs_path = os.path.join(run_output_dir, remote_store.logs_subdir)
             os.makedirs(logs_path, exist_ok=True)
             print(f"Made directory {logs_path} for horovod rank {hvd.rank()}")
@@ -262,7 +262,7 @@ def on_epoch_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -
             if hvd.rank() == 0:
                 if remote_store.saving_runs and trainer.profiler:
                     # One more file sync to push profiler result.
-                    remote_store.sync(logs_path)
+                    remote_store.sync(run_output_dir)
 
                 # rank 0 overwrites model with best checkpoint and returns.
                 if require_checkpoint:
diff --git a/horovod/spark/torch/remote.py b/horovod/spark/torch/remote.py
index 6b2f325caa..e155f2d9ed 100644
--- a/horovod/spark/torch/remote.py
+++ b/horovod/spark/torch/remote.py
@@ -198,7 +198,7 @@ def train(serialized_model, optimizer_cls, model_opt_state_serialized,
         else:
             steps_per_epoch = train_steps_per_epoch
 
-        with remote_store.get_local_output_dir() as run_output_dir:
+        with remote_store.get_local_run_dir() as run_output_dir:
             logs_dir = os.path.join(run_output_dir, remote_store.logs_subdir)
             log_writer = SummaryWriter(logs_dir) if hvd.rank() == 0 else None
             ckpt_file = os.path.join(run_output_dir, remote_store.checkpoint_filename)
diff --git a/test/integration/test_spark_lightning.py b/test/integration/test_spark_lightning.py
index 6c784b58ae..c3278fb96d 100644
--- a/test/integration/test_spark_lightning.py
+++ b/test/integration/test_spark_lightning.py
@@ -198,8 +198,6 @@ def test_legacy_fit_model(self):
 
     # TODO: Add this test back after checkpoint call back is supported
     def test_restore_from_checkpoint(self):
-        self.skipTest('There is a deadlock bug for checkpoint call back. ' +
-                      'Will add this test back when it is solved.')
 
         model = create_xor_model()
 
@@ -236,8 +234,6 @@ def test_restore_from_checkpoint(self):
 
     # TODO: Add this test back after checkpoint call back is supported
     def test_legacy_restore_from_checkpoint(self):
-        self.skipTest('There is a deadlock bug for checkpoint call back. ' +
-                      'Will add this test back when it is solved.')
 
         model = create_legacy_xor_model()
         optimizer = torch.optim.SGD(model.parameters(), lr=0.1)


Code review comments:
diff_hunk: @@ -262,7 +262,7 @@ def on_epoch_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -
             if hvd.rank() == 0:
                 if remote_store.saving_runs and trainer.profiler:
                     # One more file sync to push profiler result.
-                    remote_store.sync(logs_path)
+                    remote_store.sync(run_output_dir)
path: horovod/spark/lightning/remote.py
text: @user1:
Checkpoint is synced at every epoch.
No need to sync checkpoint files again here.
This is only for uploading profiler results, which is available at the end of training.

@author:
Got it, will revert this change.