caffe2: preserve python exception type from PythonOp (#36267)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/36267

This makes PythonOp throw the original python exception instead of wrapping it in a c10::Error type. This allows throwing exceptions from Python and preserving the type when they're caught again in Python. This is important for structured logging and handling non-retryable error types.

Test Plan: buck test caffe2/caffe2/python:python_op_test

Reviewed By: wenqicaofb

Differential Revision: D20928098

fbshipit-source-id: 001747f022c657b420f8450b84d64f4d57f6cdf6
diff --git a/caffe2/python/pybind_state.h b/caffe2/python/pybind_state.h
index 5a3628d..53dbbdc 100644
--- a/caffe2/python/pybind_state.h
+++ b/caffe2/python/pybind_state.h
@@ -364,11 +364,10 @@
                      OperatorBase::template GetSingleArgument<bool>(
                          "pass_workspace", false)});
       } catch (const py::error_already_set& e) {
-        std::stringstream error;
-        error << "Python exception encountered while creating PythonOp: "
-              << e.what();
-        LOG(ERROR) << error.str();
-        CAFFE_THROW(error.str());
+        LOG(ERROR) << "Python exception encountered while creating PythonOp: "
+                   << e.what();
+        // Rethrow exception to preserve python exception type.
+        throw;
       }
     }
   }
@@ -439,11 +438,10 @@
           pyFunc->py_func(inputs, outputs);
         }
       } catch (const py::error_already_set& e) {
-        std::stringstream error;
-        error << "Exception encountered running PythonOp function: "
-              << e.what();
-        LOG(ERROR) << error.str();
-        CAFFE_THROW(error.str());
+        LOG(ERROR) << "Exception encountered running PythonOp function: "
+                   << e.what();
+        // Rethrow exception to preserve python exception type.
+        throw;
       }
     }
     return true;
diff --git a/caffe2/python/python_op_test.py b/caffe2/python/python_op_test.py
index 3a83899..20a24e3 100644
--- a/caffe2/python/python_op_test.py
+++ b/caffe2/python/python_op_test.py
@@ -11,14 +11,20 @@
 import six
 
 
-def SubFunctionThatThrowsRuntimeError():
-    raise RuntimeError("This is an intentional exception.")
+class CustomError(Exception):
+    pass
 
 
-def MainOpFunctionThatThrowsRuntimeError(inputs, _):
-    return SubFunctionThatThrowsRuntimeError()
+def SubFunctionThatThrowsCustomError():
+    raise CustomError("This is an intentional exception.")
 
 
+def MainOpFunctionThatThrowsCustomError(inputs, _):
+    return SubFunctionThatThrowsCustomError()
+
+def MainOpFunctionThatThrowsCustomErrorInBuilder(inputs, _):
+    raise CustomError("This is an intentional exception in builder.")
+
 def op_builder(name, index, extra):
     iterations = [0]
     assert name == 'name'
@@ -48,8 +54,13 @@
         workspace.RunOperatorOnce(op)
 
     def test_exception(self):
-        op = CreatePythonOperator(MainOpFunctionThatThrowsRuntimeError, [], [])
-        with six.assertRaisesRegex(self, RuntimeError, "This is an intentional exception."):
+        op = CreatePythonOperator(MainOpFunctionThatThrowsCustomError, [], [])
+        with six.assertRaisesRegex(self, CustomError, "This is an intentional exception."):
+            workspace.RunOperatorOnce(op)
+
+    def test_exception_builder(self):
+        op = CreatePythonOperator(MainOpFunctionThatThrowsCustomErrorInBuilder, [], [])
+        with six.assertRaisesRegex(self, CustomError, "This is an intentional exception in builder."):
             workspace.RunOperatorOnce(op)
 
     @given(x=hu.tensor())