Merge pull request #199 from NVIDIA/abasant/persistent_worker_del_after_abort

Training exit after Inprocess abort should ensure clean shutdown of PersistentAsync worker

Author: aartibasant

src/nvidia_resiliency_ext/checkpointing/async_ckpt/core.py

@@ -655,7 +655,9 @@ def close(self, abort=False):
             abort (bool, optional): Default to False. Needs to be manually set to true when
                 the checkpoint async process needs to be aborted.
         """
-        if not abort:
+        # For a clean shut down scenario with valid async processes running,
+        # finalize all pending async calls
+        if not abort and (self.persistent is False or self.persistent_caller is not None):
             self.maybe_finalize_async_calls(blocking=True)
         if self.persistent and self.persistent_caller:
             self.persistent_caller.close(abort=abort)

tests/checkpointing/unit/test_async_writer.py

@@ -308,3 +308,34 @@ def test_async_cp_with_multiple_queue_and_abort(self, tmp_path_dist_ckpt):
 
             async_queue_dist.close()
             ckpt_impl.close()
+
+    def test_async_cp_with_multiple_queue_and_abort_followed_by_delete(self, tmp_path_dist_ckpt):
+        """
+        Test that persistent async CP worker shuts down cleanly after an abort operation.
+        This test mocks the behavior of training exiting after an abort triggered by an inprocess restart.
+        """
+        Utils.initialize_distributed()
+        model = FSDP(Model((1024, 1024), 8))
+        async_queue_dist = AsyncCallsQueue(persistent=True)
+        with (
+            TempNamedDir(
+                tmp_path_dist_ckpt / 'async_checkpoint_dist', sync=True
+            ) as async_ckpt_dir_dist,
+        ):
+            state_dict = model.state_dict()
+            planner = DefaultSavePlanner()
+
+            try:
+                # Raise an exception in training process right after async CP request is submitted
+                with pytest.raises(RuntimeError) as exc_info:
+                    self.async_save_checkpoint(
+                        async_ckpt_dir_dist, state_dict, planner, async_queue_dist
+                    )
+                    raise RuntimeError("Fake exception to mock training process exception")
+                    async_queue_dist.maybe_finalize_async_calls(blocking=True, no_dist=False)
+            finally:
+                # Mock behavior of an abort operation triggered by inprocess restart when exception occurs.
+                # Abort the CP workers to mock the action of inprocess restarts
+                abort_nvrx_checkpoint()
+        # Mock training loop exit which would invoke a __del__ on async queue object
+        async_queue_dist.__del__()