DCP checkpoint restore backward compatibility in TorchTNT (#898)

Summary:
Pull Request resolved: #898

 DCP checkpoint restore backward compatibility. Current implementation assumes the checkpoint to be a ModelStore Checkpoint. There are customers who are using DCP saver with path based checkpoint id.

Reviewed By: JKSenthil

Differential Revision: D62542227

fbshipit-source-id: 6981cad0e4195308c0fc2e8a5da621e2fd9c024a

tests/utils/test_checkpoint.py

@@ -1339,6 +1339,26 @@ def test_metadata_exists(self) -> None:
             os.remove(os.path.join(dirpath, SNAPSHOT_METADATA_FNAME))
             self.assertFalse(_metadata_exists(fs, dirpath, SNAPSHOT_METADATA_FNAME))
 
+    def test_does_checkpoint_metadata_exist(self) -> None:
+        app_state = {"module": nn.Linear(2, 2)}
+
+        with tempfile.TemporaryDirectory() as temp_dir:
+            dirpath = os.path.join(temp_dir, "checkpoint")
+            Snapshot.take(dirpath, app_state=app_state)
+
+            self.assertTrue(
+                CheckpointManager.does_checkpoint_metadata_exist(
+                    dirpath, SNAPSHOT_METADATA_FNAME
+                )
+            )
+
+            os.remove(os.path.join(dirpath, SNAPSHOT_METADATA_FNAME))
+            self.assertFalse(
+                CheckpointManager.does_checkpoint_metadata_exist(
+                    dirpath, SNAPSHOT_METADATA_FNAME
+                )
+            )
+
 
 class MyValLossUnit(TrainUnit[Batch]):
     def __init__(self) -> None:

torchtnt/utils/checkpoint.py

@@ -539,6 +539,18 @@ def does_checkpoint_exist(
             _metadata_exists(fs, ckpt.path, fname) for fname in self._metadata_fnames
         )
 
+    @staticmethod
+    def does_checkpoint_metadata_exist(
+        checkpoint_path: str,
+        metadata_fname: str,
+    ) -> bool:
+        """
+        Checking whether a checkpoint metadata file exists in the directory.
+        If the checkpointer has that metadata file, this function will returns True. Returns False otherwise.
+        """
+        fs, _ = url_to_fs(checkpoint_path)
+        return _metadata_exists(fs, checkpoint_path, metadata_fname)
+
     @staticmethod
     @rank_zero_read_and_broadcast
     def _sync_dirpath_to_all_ranks(