tf.errors.OpError and its subclasses are pickleable

tensorflow/python/framework/errors_impl.py

@@ -18,6 +18,7 @@
 from __future__ import division
 from __future__ import print_function
 
+import inspect
 import traceback
 import warnings
 
@@ -47,11 +48,17 @@ def __init__(self, node_def, op, message, error_code):
       error_code: The `error_codes_pb2.Code` describing the error.
     """
     super(OpError, self).__init__()
-    self._message = message
     self._node_def = node_def
     self._op = op
+    self._message = message
     self._error_code = error_code
 
+  def __reduce__(self):
+    # Allow the subclasses to accept less arguments in their __init__.
+    init_argspec = inspect.getargspec(self.__class__.__init__)
+    args = tuple(getattr(self, arg) for arg in init_argspec.args[1:])
+    return self.__class__, args
+
   @property
   def message(self):
     """The error message that describes the error."""

tensorflow/python/framework/errors_test.py

@@ -19,6 +19,7 @@
 from __future__ import print_function
 
 import gc
+import pickle
 import warnings
 
 from tensorflow.core.lib.core import error_codes_pb2
@@ -108,5 +109,34 @@ def testStatusDoesNotLeak(self):
     self.assertEqual(0, self._CountReferences(c_api_util.ScopedTFStatus))
 
 
+  def testPickleable(self):
+    for error_code in [
+        errors.CANCELLED,
+        errors.UNKNOWN,
+        errors.INVALID_ARGUMENT,
+        errors.DEADLINE_EXCEEDED,
+        errors.NOT_FOUND,
+        errors.ALREADY_EXISTS,
+        errors.PERMISSION_DENIED,
+        errors.UNAUTHENTICATED,
+        errors.RESOURCE_EXHAUSTED,
+        errors.FAILED_PRECONDITION,
+        errors.ABORTED,
+        errors.OUT_OF_RANGE,
+        errors.UNIMPLEMENTED,
+        errors.INTERNAL,
+        errors.UNAVAILABLE,
+        errors.DATA_LOSS,
+    ]:
+      # pylint: disable=protected-access
+      exc = errors_impl._make_specific_exception(None, None, None, error_code)
+      # pylint: enable=protected-access
+      unpickled = pickle.loads(pickle.dumps(exc))
+      self.assertEqual(exc.node_def, unpickled.node_def)
+      self.assertEqual(exc.op, unpickled.op)
+      self.assertEqual(exc.message, unpickled.message)
+      self.assertEqual(exc.error_code, unpickled.error_code)
+
+
 if __name__ == "__main__":
   test.main()