Better graceful shutdown for KeyboardInterrupt (#19976)

Author: awaelchli

src/lightning/fabric/CHANGELOG.md

@@ -5,6 +5,40 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 
+## [unreleased] - YYYY-MM-DD
+
+### Added
+
+-
+
+-
+
+### Changed
+
+-
+
+-
+
+### Deprecated
+
+-
+
+-
+
+### Removed
+
+-
+
+-
+
+### Fixed
+
+-
+
+-
+
+
+
 ## [2.3.0] - 2024-06-13
 
 ### Added

src/lightning/fabric/utilities/distributed.py

@@ -2,6 +2,7 @@
 import contextlib
 import logging
 import os
+import signal
 import time
 from contextlib import nullcontext
 from datetime import timedelta
@@ -306,8 +307,11 @@ def _init_dist_connection(
 
 
 def _destroy_dist_connection() -> None:
+    # Don't allow Ctrl+C to interrupt this handler
+    signal.signal(signal.SIGINT, signal.SIG_IGN)
     if _distributed_is_initialized():
         torch.distributed.destroy_process_group()
+    signal.signal(signal.SIGINT, signal.SIG_DFL)
 
 
 def _get_default_process_group_backend_for_device(device: torch.device) -> str:

src/lightning/pytorch/CHANGELOG.md

@@ -4,6 +4,41 @@ All notable changes to this project will be documented in this file.
 
 The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
+
+## [unreleased] - YYYY-MM-DD
+
+### Added
+
+-
+
+-
+
+### Changed
+
+- Triggering KeyboardInterrupt (Ctrl+C) during `.fit()`, `.evaluate()`, `.test()` or `.predict()` now terminates all processes launched by the Trainer and exits the program ([#19976](https://github.com/Lightning-AI/pytorch-lightning/pull/19976))
+
+-
+
+### Deprecated
+
+-
+
+-
+
+### Removed
+
+-
+
+-
+
+### Fixed
+
+-
+
+-
+
+
+
 ## [2.3.0] - 2024-06-13
 
 ### Added

src/lightning/pytorch/strategies/launchers/multiprocessing.py

@@ -259,7 +259,7 @@ def update_main_process_results(self, trainer: "pl.Trainer", extra: Dict[str, An
     def kill(self, signum: _SIGNUM) -> None:
         for proc in self.procs:
             if proc.is_alive() and proc.pid is not None:
-                log.info(f"pid {os.getpid()} killing {proc.pid} with {signum}")
+                log.debug(f"Process {os.getpid()} is terminating {proc.pid} with {signum}")
                 with suppress(ProcessLookupError):
                     os.kill(proc.pid, signum)
 

src/lightning/pytorch/strategies/launchers/subprocess_script.py

@@ -107,7 +107,7 @@ def launch(self, function: Callable, *args: Any, trainer: Optional["pl.Trainer"]
     @override
     def kill(self, signum: _SIGNUM) -> None:
         for proc in self.procs:
-            log.info(f"pid {os.getpid()} killing {proc.pid} with {signum}")
+            log.debug(f"Process {os.getpid()} is terminating {proc.pid} with {signum}")
             # this skips subprocesses already terminated
             proc.send_signal(signum)
 

src/lightning/pytorch/trainer/call.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import logging
+import signal
 from copy import deepcopy
 from typing import Any, Callable, Dict, Optional, Type, Union
 
@@ -20,10 +21,12 @@
 import lightning.pytorch as pl
 from lightning.fabric.utilities.device_dtype_mixin import _DeviceDtypeModuleMixin
 from lightning.pytorch.callbacks import Checkpoint, EarlyStopping
+from lightning.pytorch.strategies.launchers import _SubprocessScriptLauncher
+from lightning.pytorch.trainer.connectors.signal_connector import _get_sigkill_signal
 from lightning.pytorch.trainer.states import TrainerStatus
 from lightning.pytorch.utilities.exceptions import _TunerExitException
 from lightning.pytorch.utilities.model_helpers import is_overridden
-from lightning.pytorch.utilities.rank_zero import rank_zero_warn
+from lightning.pytorch.utilities.rank_zero import rank_zero_info, rank_zero_warn
 
 log = logging.getLogger(__name__)
 
@@ -49,12 +52,17 @@ def _call_and_handle_interrupt(trainer: "pl.Trainer", trainer_fn: Callable, *arg
         trainer.state.status = TrainerStatus.FINISHED
         trainer.state.stage = None
 
-    # TODO: Unify both exceptions below, where `KeyboardError` doesn't re-raise
     except KeyboardInterrupt as exception:
-        rank_zero_warn("Detected KeyboardInterrupt, attempting graceful shutdown...")
-        # user could press Ctrl+c many times... only shutdown once
-        if not trainer.interrupted:
-            _interrupt(trainer, exception)
+        rank_zero_info("\nDetected KeyboardInterrupt, attempting graceful shutdown ...")
+        # user could press Ctrl+C many times, disable KeyboardInterrupt for shutdown
+        signal.signal(signal.SIGINT, signal.SIG_IGN)
+        _interrupt(trainer, exception)
+        trainer._teardown()
+        launcher = trainer.strategy.launcher
+        if isinstance(launcher, _SubprocessScriptLauncher):
+            launcher.kill(_get_sigkill_signal())
+        exit(1)
+
     except BaseException as exception:
         _interrupt(trainer, exception)
         trainer._teardown()

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

@@ -2,7 +2,6 @@
 import os
 import re
 import signal
-import sys
 import threading
 from subprocess import call
 from types import FrameType
@@ -54,7 +53,7 @@ def register_signal_handlers(self) -> None:
             sigterm_handlers.append(self._sigterm_handler_fn)
 
         # Windows seems to have signal incompatibilities
-        if not self._is_on_windows():
+        if not _IS_WINDOWS:
             sigusr = environment.requeue_signal if isinstance(environment, SLURMEnvironment) else signal.SIGUSR1
             assert sigusr is not None
             if sigusr_handlers and not self._has_already_handler(sigusr):
@@ -155,10 +154,6 @@ def _valid_signals() -> Set[signal.Signals]:
             }
         return set(signal.Signals)
 
-    @staticmethod
-    def _is_on_windows() -> bool:
-        return sys.platform == "win32"
-
     @staticmethod
     def _has_already_handler(signum: _SIGNUM) -> bool:
         return signal.getsignal(signum) not in (None, signal.SIG_DFL)
@@ -172,3 +167,7 @@ def __getstate__(self) -> Dict:
         state = self.__dict__.copy()
         state["_original_handlers"] = {}
         return state
+
+
+def _get_sigkill_signal() -> _SIGNUM:
+    return signal.SIGTERM if _IS_WINDOWS else signal.SIGKILL

tests/tests_pytorch/callbacks/progress/test_rich_progress_bar.py

@@ -143,7 +143,7 @@ def on_train_start(self) -> None:
 
     with mock.patch(
         "lightning.pytorch.callbacks.progress.rich_progress.Progress.stop", autospec=True
-    ) as mock_progress_stop:
+    ) as mock_progress_stop, pytest.raises(SystemExit):
         progress_bar = RichProgressBar()
         trainer = Trainer(
             default_root_dir=tmp_path,

tests/tests_pytorch/callbacks/test_lambda_function.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 from functools import partial
 
+import pytest
 from lightning.pytorch import Trainer, seed_everything
 from lightning.pytorch.callbacks import Callback, LambdaCallback
 from lightning.pytorch.demos.boring_classes import BoringModel
@@ -23,10 +24,13 @@
 def test_lambda_call(tmp_path):
     seed_everything(42)
 
+    class CustomException(Exception):
+        pass
+
     class CustomModel(BoringModel):
         def on_train_epoch_start(self):
             if self.current_epoch > 1:
-                raise KeyboardInterrupt
+                raise CustomException("Custom exception to trigger `on_exception` hooks")
 
     checker = set()
 
@@ -59,7 +63,8 @@ def call(hook, *_, **__):
         limit_predict_batches=1,
         callbacks=[LambdaCallback(**hooks_args)],
     )
-    trainer.fit(model, ckpt_path=ckpt_path)
+    with pytest.raises(CustomException):
+        trainer.fit(model, ckpt_path=ckpt_path)
     trainer.test(model)
     trainer.predict(model)
 

tests/tests_pytorch/trainer/test_states.py

@@ -84,5 +84,6 @@ def on_train_batch_start(self, trainer, pl_module, batch, batch_idx):
 
     trainer = Trainer(callbacks=[InterruptCallback()], default_root_dir=tmp_path, **extra_params)
 
-    trainer.fit(model)
+    with pytest.raises(SystemExit):
+        trainer.fit(model)
     assert trainer.interrupted