Replace exit with `sys.exit (#21052)

* replace exit with sys.exit
* fix in upgrade checkpoint

(cherry picked from commit 64d3112)

Author: SkafteNicki

src/lightning/fabric/utilities/consolidate_checkpoint.py

@@ -1,4 +1,5 @@
 import logging
+import sys
 from argparse import ArgumentParser, Namespace
 from pathlib import Path
 
@@ -40,23 +41,23 @@ def _parse_cli_args() -> Namespace:
 def _process_cli_args(args: Namespace) -> Namespace:
     if not _TORCH_GREATER_EQUAL_2_3:
         _log.error("Processing distributed checkpoints requires PyTorch >= 2.3.")
-        exit(1)
+        sys.exit(1)
 
     checkpoint_folder = Path(args.checkpoint_folder)
     if not checkpoint_folder.exists():
         _log.error(f"The provided checkpoint folder does not exist: {checkpoint_folder}")
-        exit(1)
+        sys.exit(1)
     if not checkpoint_folder.is_dir():
         _log.error(
             f"The provided checkpoint path must be a folder, containing the checkpoint shards: {checkpoint_folder}"
         )
-        exit(1)
+        sys.exit(1)
     if not (checkpoint_folder / _METADATA_FILENAME).is_file():
         _log.error(
             "Only FSDP-sharded checkpoints saved with Lightning are supported for consolidation. The provided folder"
             f" is not in that format: {checkpoint_folder}"
         )
-        exit(1)
+        sys.exit(1)
 
     if args.output_file is None:
         output_file = checkpoint_folder.with_suffix(checkpoint_folder.suffix + ".consolidated")
@@ -67,7 +68,7 @@ def _process_cli_args(args: Namespace) -> Namespace:
             "The path for the converted checkpoint already exists. Choose a different path by providing"
             f" `--output_file` or move/delete the file first: {output_file}"
         )
-        exit(1)
+        sys.exit(1)
 
     return Namespace(checkpoint_folder=checkpoint_folder, output_file=output_file)
 

src/lightning/pytorch/trainer/call.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 import logging
 import signal
+import sys
 from copy import deepcopy
 from typing import Any, Callable, Optional, Union
 
@@ -62,7 +63,7 @@ def _call_and_handle_interrupt(trainer: "pl.Trainer", trainer_fn: Callable, *arg
         launcher = trainer.strategy.launcher
         if isinstance(launcher, _SubprocessScriptLauncher):
             launcher.kill(_get_sigkill_signal())
-        exit(1)
+        sys.exit(1)
 
     except BaseException as exception:
         _interrupt(trainer, exception)

src/lightning/pytorch/utilities/upgrade_checkpoint.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 import glob
 import logging
+import sys
 from argparse import ArgumentParser, Namespace
 from pathlib import Path
 from shutil import copyfile
@@ -35,7 +36,7 @@ def _upgrade(args: Namespace) -> None:
             f"The path {path} does not exist. Please provide a valid path to a checkpoint file or a directory"
             f" containing checkpoints ending in {extension}."
         )
-        exit(1)
+        sys.exit(1)
 
     if path.is_file():
         files = [path]
@@ -46,7 +47,7 @@ def _upgrade(args: Namespace) -> None:
             f"No checkpoint files with extension {extension} were found in {path}."
             f" HINT: Try setting the `--extension` option to specify the right file extension to look for."
         )
-        exit(1)
+        sys.exit(1)
 
     _log.info("Creating a backup of the existing checkpoint files before overwriting in the upgrade process.")
     for file in files: