[Train] Add a barrier in RayTrainReportCallback to ensure synchronous…

… reporting. (#40875) Signed-off-by: woshiyyya <xiaoyunxuan1998@gmail.com>
ray-project · Nov 3, 2023 · c1e387f · c1e387f
1 parent f08498e
commit c1e387f
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/python/ray/train/lightning/_lightning_utils.py b/python/ray/train/lightning/_lightning_utils.py
@@ -270,6 +270,9 @@ def on_train_epoch_end(self, trainer, pl_module) -> None:
         checkpoint = Checkpoint.from_directory(tmpdir)
         train.report(metrics=metrics, checkpoint=checkpoint)
 
+        # Add a barrier to ensure all workers finished reporting here
+        torch.distributed.barrier()
+
         if self.local_rank == 0:
             shutil.rmtree(tmpdir)