Repository: horovod/horovod

Issue:
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.



Discussion:

Resolving patch:
diff --git a/horovod/spark/common/store.py b/horovod/spark/common/store.py
--- 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
--- 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()}")
diff --git a/horovod/spark/torch/remote.py b/horovod/spark/torch/remote.py
--- 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)


Test patch:
diff --git a/test/integration/test_spark_lightning.py b/test/integration/test_spark_lightning.py
--- 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)