Add exception classification to torch.multiprocessing.spawn (pytorch#45174)

Summary:
Pull Request resolved: pytorch#45174

Introduce different types of exceptions that map to different failures
of torch.multiprocessing.spawn. The change introduces three different exception types:
ProcessRaisedException - occurs when the process initiated by spawn raises an exception
ProcessExitedException - occurs when the process initiated by spawn exits
The following logic will allow frameworks that use mp.spawn to categorize failures.
This can be helpful for tracking metrics and enhancing logs.

Test Plan: Imported from OSS

Reviewed By: taohe

Differential Revision: D23889400

Pulled By: tierex

fbshipit-source-id: 8849624c616230a6a81158c52ce0c18beb437330

Author: aivanou

test/test_multiprocessing_spawn.py

@@ -54,6 +54,15 @@ def test_nested_child_body(i, ready_queue, nested_child_sleep):
     time.sleep(nested_child_sleep)
 
 
+def test_infinite_task(i):
+    while True:
+        time.sleep(1)
+
+
+def test_process_exit(idx):
+    sys.exit(12)
+
+
 def test_nested(i, pids_queue, nested_child_sleep, start_method):
     context = mp.get_context(start_method)
     nested_child_ready_queue = context.Queue()
@@ -184,6 +193,23 @@ def test_nested(self):
 class SpawnTest(TestCase, _TestMultiProcessing):
     start_method = 'spawn'
 
+    def test_exception_raises(self):
+        with self.assertRaises(mp.ProcessRaisedException):
+            mp.spawn(test_success_first_then_exception_func, args=(), nprocs=1)
+
+    def test_signal_raises(self):
+        context = mp.spawn(test_infinite_task, args=(), nprocs=1, join=False)
+        for pid in context.pids():
+            os.kill(pid, signal.SIGTERM)
+        with self.assertRaises(mp.ProcessExitedException):
+            context.join()
+
+    def test_process_exited(self):
+        with self.assertRaises(mp.ProcessExitedException) as e:
+            mp.spawn(test_process_exit, args=(), nprocs=1)
+            self.assertEqual(12, e.exit_code)
+
+
 @unittest.skipIf(
     IS_WINDOWS,
     "Fork is only available on Unix",

torch/multiprocessing/__init__.py

@@ -42,7 +42,8 @@
 
 """Add helper function to spawn N processes and wait for completion of any of
 them. This depends `mp.get_context` which was added in Python 3.4."""
-from .spawn import spawn, SpawnContext, _supports_context, start_processes, ProcessContext
+from .spawn import spawn, SpawnContext, _supports_context, start_processes, ProcessContext, \
+    ProcessRaisedException, ProcessExitedException
 
 
 if sys.platform == 'darwin' or sys.platform == 'win32':

torch/multiprocessing/spawn.py

@@ -1,4 +1,5 @@
 
+from typing import Optional
 import multiprocessing
 import multiprocessing.connection
 import signal
@@ -8,6 +9,45 @@
 from . import _prctl_pr_set_pdeathsig
 
 
+class ProcessException(Exception):
+    __slots__ = ["error_index", "error_pid"]
+
+    def __init__(self, msg: str, error_index: int, pid: int):
+        super().__init__(msg)
+        self.error_index = error_index
+        self.pid = pid
+
+
+class ProcessRaisedException(ProcessException):
+    """
+    Exception is thrown when the process failed due to exception
+    raised by the code.
+    """
+    def __init__(
+        self,
+        msg: str,
+        error_index: int,
+        error_pid: int,
+    ):
+        super().__init__(msg, error_index, error_pid)
+
+
+class ProcessExitedException(ProcessException):
+    """
+    Exception is thrown when the process failed due to signal
+    or exited with a specific code.
+    """
+    __slots__ = ["exit_code"]
+
+    def __init__(
+            self, msg: str, error_index: int, error_pid: int,
+            exit_code: int, signal_name: Optional[str] = None
+    ):
+        super().__init__(msg, error_index, error_pid)
+        self.exit_code = exit_code
+        self.signal_name = signal_name
+
+
 def _wrap(fn, i, args, error_queue):
     # prctl(2) is a Linux specific system call.
     # On other systems the following function call has no effect.
@@ -98,24 +138,32 @@ def join(self, timeout=None):
             process.join()
 
         # There won't be an error on the queue if the process crashed.
+        failed_process = self.processes[error_index]
         if self.error_queues[error_index].empty():
             exitcode = self.processes[error_index].exitcode
             if exitcode < 0:
                 name = signal.Signals(-exitcode).name
-                raise Exception(
+                raise ProcessExitedException(
                     "process %d terminated with signal %s" %
-                    (error_index, name)
+                    (error_index, name),
+                    error_index=error_index,
+                    error_pid=failed_process.pid,
+                    exit_code=exitcode,
+                    signal_name=name
                 )
             else:
-                raise Exception(
+                raise ProcessExitedException(
                     "process %d terminated with exit code %d" %
-                    (error_index, exitcode)
+                    (error_index, exitcode),
+                    error_index=error_index,
+                    error_pid=failed_process.pid,
+                    exit_code=exitcode
                 )
 
         original_trace = self.error_queues[error_index].get()
         msg = "\n\n-- Process %d terminated with the following error:\n" % error_index
         msg += original_trace
-        raise Exception(msg)
+        raise ProcessRaisedException(msg, error_index, failed_process.pid)
 
 
 class SpawnContext(ProcessContext):