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())