Update signal_connector.py

Author: KAVYANSHTYAGI

src/lightning/pytorch/trainer/connectors/signal_connector.py

@@ -6,6 +6,8 @@
 from subprocess import call
 from types import FrameType
 from typing import Any, Callable, Union
+import torch
+import torch.distributed as dist
 
 import lightning.pytorch as pl
 from lightning.fabric.plugins.environments import SLURMEnvironment
@@ -104,14 +106,19 @@ def _slurm_sigusr_handler_fn(self, signum: _SIGNUM, _: FrameType) -> None:
 
     def _sigterm_notifier_fn(self, signum: _SIGNUM, _: FrameType) -> None:
         log.info(rank_prefixed_message(f"Received SIGTERM: {signum}", self.trainer.local_rank))
-        # subprocesses killing the parent process is not supported, only the parent (rank 0) does it
         if not self.received_sigterm:
-            # send the same signal to the subprocesses
             launcher = self.trainer.strategy.launcher
             if launcher is not None:
                 launcher.kill(signum)
+
+        # New broadcast logic
+        if dist.is_available() and dist.is_initialized() and self.trainer.world_size > 1:
+            sigterm_tensor = torch.tensor([1], device=self.trainer.strategy.root_device)
+            dist.broadcast(sigterm_tensor, src=0)
+
         self.received_sigterm = True
 
+
     def _sigterm_handler_fn(self, signum: _SIGNUM, _: FrameType) -> None:
         log.info(f"Bypassing SIGTERM: {signum}")