Preserve the original exception type when converting to RayTaskError (#5799)

python/ray/exceptions.py

@@ -28,18 +28,49 @@ class RayTaskError(RayError):
         traceback_str (str): The traceback from the exception.
     """
 
-    def __init__(self, function_name, traceback_str):
+    def __init__(self,
+                 function_name,
+                 traceback_str,
+                 cause_cls,
+                 pid=None,
+                 host=None):
         """Initialize a RayTaskError."""
         if setproctitle:
             self.proctitle = setproctitle.getproctitle()
         else:
             self.proctitle = "ray_worker"
-        self.pid = os.getpid()
-        self.host = os.uname()[1]
+        self.pid = pid or os.getpid()
+        self.host = host or os.uname()[1]
         self.function_name = function_name
         self.traceback_str = traceback_str
+        self.cause_cls = cause_cls
         assert traceback_str is not None
 
+    def as_instanceof_cause(self):
+        """Returns copy that is an instance of the cause's Python class.
+
+        The returned exception will inherit from both RayTaskError and the
+        cause class.
+        """
+
+        if issubclass(RayTaskError, self.cause_cls):
+            return self  # already satisfied
+
+        class cls(self.cause_cls, RayTaskError):
+            def __init__(self, function_name, traceback_str, cause_cls, pid,
+                         host):
+                RayTaskError.__init__(self, function_name, traceback_str,
+                                      cause_cls, pid, host)
+
+        name = "RayTaskError({})".format(self.cause_cls.__name__)
+        cls.__name__ = name
+        cls.__qualname__ = name
+
+        return cls(self.function_name, self.traceback_str, self.cause_cls,
+                   self.pid, self.host)
+        cls.original = self
+        return cls
+
     def __str__(self):
         """Format a RayTaskError as a string."""
         lines = self.traceback_str.split("\n")