capturing exit codes after sigterm/sigkill from torch elastic. (pytorch#160908)

Summary:
**Background**
Torch Elastic sends SIGKILL/SIGTERM signals if any process fails while others are still running. However, processes terminated by these signals do not generate termination logs, causing confusion.

**Solution**
Capture exit codes after SIGTERM signals to ensure complete and accurate termination logging.

Test Plan:
unit tests

https://www.internalfb.com/mlhub/pipelines/runs/mast/f773486907-TrainingApplication__13_D79584569?job_attempt=1&version=0&tab=summary&env=PRODUCTION

Rollback Plan:

Differential Revision: D79584569

Pull Request resolved: pytorch#160908
Approved by: https://github.com/d4l3k

Author: aschhabra

test/distributed/elastic/multiprocessing/api_test.py

@@ -568,9 +568,8 @@ def test_binary_exit(self):
             )
 
             results = pc.wait(period=0.1)
-
             self.assertTrue(results.is_failed())
-            self.assertEqual(1, len(results.failures))
+            self.assertEqual(2, len(results.failures))
 
             failure = results.failures[0]
             self.assertEqual(138, failure.exitcode)
@@ -583,6 +582,13 @@ def test_binary_exit(self):
             self.assertTrue(pc._stderr_tail.stopped())
             self.assertTrue(pc._stdout_tail.stopped())
 
+            failure = results.failures[1]
+            self.assertEqual(-15, failure.exitcode)
+            self.assertEqual("SIGTERM", failure.signal_name())
+            self.assertEqual("<NONE>", failure.error_file_data["message"])
+            # Assert that the failure message contains expected substrings
+            self.assertIn("Signal 15 (SIGTERM) received by PID", failure.message)
+
         def test_binary_raises(self):
             pc = start_processes(
                 name="echo",

test/distributed/elastic/multiprocessing/bin/echo1.py

@@ -9,6 +9,7 @@
 import argparse
 import os
 import sys
+import time
 
 
 if __name__ == "__main__":
@@ -23,5 +24,6 @@
         print(f"exit {exitcode} from {rank}", file=sys.stderr)
         sys.exit(exitcode)
     else:
+        time.sleep(1000)
         print(f"{args.msg} stdout from {rank}")
         print(f"{args.msg} stderr from {rank}", file=sys.stderr)

torch/distributed/elastic/multiprocessing/api.py

@@ -875,8 +875,7 @@ def _start(self):
             for local_rank in range(self.nprocs)
         }
 
-    def _poll(self) -> Optional[RunProcsResult]:
-        done_local_ranks = set()
+    def _capture_process_failures(self, done_local_ranks: set[int]):
         for local_rank in self._running_local_ranks:
             handler = self.subprocess_handlers[local_rank]
             exitcode = handler.proc.poll()
@@ -891,11 +890,19 @@ def _poll(self) -> Optional[RunProcsResult]:
                     )
                 # else: --> succeeded; nothing to do
 
+    def _poll(self) -> Optional[RunProcsResult]:
+        done_local_ranks: set[int] = set()
+        self._capture_process_failures(done_local_ranks)
+
         self._running_local_ranks.difference_update(done_local_ranks)
 
         # if ALL procs are finished or ANY have failed
         if not self._running_local_ranks or self._failures:
             self.close()  # terminate all running procs
+            self._capture_process_failures(
+                done_local_ranks
+            )  # log sigterms and sigkill exit codes in the self._failures for bookkeeping purposes
+
             result = RunProcsResult(
                 failures=self._failures,
                 stdouts=self.stdouts,