Update training_epoch_loop.py

Author: KAVYANSHTYAGI

src/lightning/pytorch/loops/training_epoch_loop.py

@@ -276,20 +276,29 @@ def advance(self, data_fetcher: _DataFetcher) -> None:
         self.val_loop.restarting = False
 
         # =====================================================================
-        # FINAL: Check for SIGTERM broadcast and exit synchronously across ranks
         from lightning.pytorch.utilities.exceptions import SIGTERMException
-
+        
         if dist.is_available() and dist.is_initialized() and self.trainer.world_size > 1:
-            with suppress(Exception):  # never crash CI
+            try:
+                # Prepare the SIGTERM signal tensor (1 if signal received, else 0)
                 sigterm_tensor = torch.tensor(
                     [1 if getattr(self.trainer, "received_sigterm", False) else 0],
                     device=self.trainer.strategy.root_device,
                 )
+                # Broadcast the SIGTERM flag from rank 0 to all other ranks
                 dist.broadcast(sigterm_tensor, src=0)
+            except Exception:
+                # In case broadcast fails (e.g., CPU-only or non-DDP), fallback to no SIGTERM
+                sigterm_tensor = torch.tensor([0], device=self.trainer.strategy.root_device)
+        
+            # If SIGTERM flag is set, synchronize all ranks and raise exception to exit
+            if sigterm_tensor.item() == 1:
+                try:
+                    dist.barrier()  # prevent deadlocks by syncing all ranks before exit
+                except Exception:
+                    pass  # Don't fail if barrier fails in fallback mode
+                raise SIGTERMException()
 
-                if sigterm_tensor.item() == 1:
-                    dist.barrier()
-                    raise SIGTERMException()
         # =====================================================================
 
         if using_dataloader_iter := isinstance(data_fetcher, _DataLoaderIterDataFetcher):