add support for parallel ckpt removal (NVIDIA-NeMo#15073)

* add support for parallel ckpt removal

Signed-off-by: dimapihtar <dpihtar@gmail.com>

* Apply isort and black reformatting

Signed-off-by: dimapihtar <dimapihtar@users.noreply.github.com>

---------

Signed-off-by: dimapihtar <dpihtar@gmail.com>
Signed-off-by: dimapihtar <dimapihtar@users.noreply.github.com>
Co-authored-by: dimapihtar <dimapihtar@users.noreply.github.com>
Signed-off-by: genquan9 <genquan@google.com>

Author: dimapihtar

nemo/lightning/pytorch/callbacks/model_checkpoint.py

@@ -15,6 +15,7 @@
 import os
 import re
 import shutil
+import threading
 from datetime import timedelta
 from pathlib import Path
 from typing import Any, Dict, Iterable, List, Literal, Optional, Union
@@ -707,7 +708,16 @@ def _remove_checkpoint(self, trainer: "lightning.pytorch.Trainer", filepath: str
         # if anything goes wrong during removal, we should be able to detect that data is incomplete.
         self.set_checkpoint_unfinished_marker(filepath, barrier_after=True)
         try:
-            super()._remove_checkpoint(trainer, filepath)
+            if self.async_save:
+                threading.Thread(
+                    target=super()._remove_checkpoint,
+                    args=(
+                        trainer,
+                        filepath,
+                    ),
+                ).start()
+            else:
+                super()._remove_checkpoint(trainer, filepath)
         except Exception as e:
             logging.warning(
                 f'Error removing checkpoint, common if doing manual cleanup and restarting: {filepath}: {e}'
@@ -718,7 +728,16 @@ def _remove_checkpoint(self, trainer: "lightning.pytorch.Trainer", filepath: str
 
             filepath = self._ema_format_filepath(filepath)
             try:
-                super()._remove_checkpoint(trainer, filepath)
+                if self.async_save:
+                    threading.Thread(
+                        target=super()._remove_checkpoint,
+                        args=(
+                            trainer,
+                            filepath,
+                        ),
+                    ).start()
+                else:
+                    super()._remove_checkpoint(trainer, filepath)
             except Exception as e:
                 logging.warning(
                     f'Error removing checkpoint, common if doing manual cleanup and restarting: {filepath}: {e}'