sync entire dirpath prior to checkpointing (#867)

Summary: Pull Request resolved: #867

Reviewed By: diego-urgell

Differential Revision: D59929876

fbshipit-source-id: 542eee83b2fd69f55e14234657fdb6a8e61dc2d7

Author: JKSenthil

tests/framework/callbacks/test_base_checkpointer.py

@@ -1045,6 +1045,48 @@ def test_multi_phase_e2e(self) -> None:
                 min_optim, [str(x) for x in metric_cb._checkpoint_manager._ckpt_paths]
             )
 
+    @skip_if_not_distributed
+    def test_directory_path_synced(self) -> None:
+        spawn_multi_process(
+            2,
+            "gloo",
+            self._test_directory_path_synced,
+        )
+
+    @staticmethod
+    def _test_directory_path_synced() -> None:
+        init_from_env()
+        tc = unittest.TestCase()
+
+        temp_dir = tempfile.mkdtemp() if get_global_rank() == 0 else ""
+        bcs = BaseCheckpointSaver(
+            temp_dir,
+            save_every_n_epochs=1,
+        )
+
+        try:
+            state = get_dummy_train_state()
+            my_train_unit = MyTrainLossUnit()
+
+            if dist.get_rank() == 0:
+                my_train_unit.train_progress._num_epochs_completed = 10
+            else:
+                my_train_unit.train_progress._num_epochs_completed = 3
+
+            bcs.on_train_epoch_end(state, my_train_unit)
+            tc.assertEqual(len(bcs._checkpoint_manager._ckpt_paths), 1)
+            tc.assertEqual(
+                str(bcs._checkpoint_manager._ckpt_paths[0]),
+                os.path.join(bcs.dirpath, "epoch_10_train_step_0"),
+            )
+            tc.assertEqual(
+                os.listdir(bcs.dirpath),
+                ["epoch_10_train_step_0"],
+            )
+        finally:
+            if get_global_rank() == 0:
+                shutil.rmtree(temp_dir)  # delete temp directory
+
 
 class MyValLossUnit(TrainUnit[Batch]):
     def __init__(self) -> None:

tests/utils/test_checkpoint.py

@@ -608,6 +608,32 @@ def test_generate_checkpoint_path(self) -> None:
         ):
             ckpt_manager.generate_checkpoint_path(1, 2, MetricData("val_loss", 3.5))
 
+    @skip_if_not_distributed
+    def test_generate_checkpoint_path_distributed(self) -> None:
+        spawn_multi_process(
+            world_size=2,
+            backend="gloo",
+            method=self._test_generate_checkpoint_path_distributed,
+        )
+
+    @staticmethod
+    def _test_generate_checkpoint_path_distributed() -> None:
+        tc = unittest.TestCase()
+
+        init_from_env()
+
+        ckpt_manager = CheckpointManager("foo")
+
+        if dist.get_rank() == 0:
+            path = ckpt_manager.generate_checkpoint_path(1, 1).path
+        else:
+            path = ckpt_manager.generate_checkpoint_path(3, 41).path
+
+        tc.assertEqual(
+            path,
+            "foo/epoch_1_step_1",
+        )
+
     def test_append_checkpoint_by_recency(self) -> None:
         ckpt_manager = CheckpointManager("foo", keep_last_n_checkpoints=3)
         ckpt_manager._ckpt_paths = [CheckpointPath("foo", 0, 0)]

torchtnt/framework/callbacks/base_checkpointer.py

@@ -171,7 +171,10 @@ def _generate_checkpoint_and_upkeep(
             )
 
         checkpoint_path = self._checkpoint_manager.generate_checkpoint_path(
-            epoch, step_mapping, metric_data
+            epoch,
+            step_mapping,
+            metric_data,
+            process_group=self._process_group,
         )
 
         # 2) Determine if we should save checkpoint

torchtnt/utils/checkpoint.py

@@ -400,11 +400,13 @@ def prune_surplus_checkpoints(self) -> None:
             for _ in range(len(self._ckpt_paths) - keep_last_n_checkpoints):
                 self.remove_checkpoint()
 
+    @rank_zero_read_and_broadcast
     def generate_checkpoint_path(
         self,
         epoch: int,
         step: Union[int, Dict[Phase, int]],
         metric_data: Optional[MetricData] = None,
+        process_group: Optional[dist.ProcessGroup] = None,
     ) -> CheckpointPath:
         """
         Given the current epoch, step, and possibly a metric_data value, determine the path