Improve Python Stack retention (#213)

Summary:
Pull Request resolved: #213

This makes SerializedPyError match the exact formatting of python exceptions by using TracebackException to format it. It also ensures that it prepends any current members of the python stack to the Exceptions traceback before formatting. Normally exceptions are only reported after being completely unwound to the root of the python stack. But in most cases when we are capturing a Python exception it is because we are going to move it across the network. We might not have competely unwound the stack at the point we do the serialization, but this unwound stack is still potentially useful. So we also include these frames by adding them to the traceback object.

This also updates two places in PythonActor where exceptions would lose stack traces by going from PyErr -> anyhow::Error. By changing the result type to Result<_, SerialiablePyErr> we force it to capture the stack trace.

Reviewed By: suo, mariusae

Differential Revision: D76316820

fbshipit-source-id: 8093ea29af4061fa90d484d31088f40be4a8cc2a

Author: zdevito

monarch_hyperactor/src/actor.rs

@@ -25,6 +25,7 @@ use hyperactor::message::IndexedErasedUnbound;
 use hyperactor::message::Unbind;
 use hyperactor_mesh::actor_mesh::Cast;
 use monarch_types::PickledPyObject;
+use monarch_types::SerializablePyErr;
 use pyo3::exceptions::PyBaseException;
 use pyo3::exceptions::PyRuntimeError;
 use pyo3::prelude::*;
@@ -276,7 +277,7 @@ impl Actor for PythonActor {
     type Params = PickledPyObject;
 
     async fn new(actor_type: PickledPyObject) -> Result<Self, anyhow::Error> {
-        Ok(Python::with_gil(|py| -> PyResult<Self> {
+        Ok(Python::with_gil(|py| -> Result<Self, SerializablePyErr> {
             let unpickled = actor_type.unpickle(py)?;
             let class_type: &Bound<'_, PyType> = unpickled.downcast()?;
             let actor: PyObject = class_type.call0()?.to_object(py);
@@ -411,7 +412,7 @@ impl Handler<PythonMessage> for PythonActor {
         // See [Panics in async endpoints].
         let (sender, receiver) = oneshot::channel();
 
-        let future = Python::with_gil(|py| -> PyResult<_> {
+        let future = Python::with_gil(|py| -> Result<_, SerializablePyErr> {
             let mailbox = PyMailbox {
                 inner: this.mailbox_for_py().clone(),
             };
@@ -438,6 +439,7 @@ impl Handler<PythonMessage> for PythonActor {
                 awaitable.into_bound(py),
             )
             .map(Some)
+            .map_err(|err| err.into())
         })?;
 
         if let Some(future) = future {

monarch_types/src/python.rs

@@ -16,6 +16,7 @@ use pyo3::ToPyObject;
 use pyo3::prelude::*;
 use pyo3::types::PyDict;
 use pyo3::types::PyList;
+use pyo3::types::PyNone;
 use pyo3::types::PyTuple;
 use serde::Deserialize;
 use serde::Serialize;
@@ -110,39 +111,68 @@ where
 /// A wrapper around `PyErr` that contains a serialized traceback.
 #[derive(Debug, Clone, Serialize, Deserialize, derive_more::Error)]
 pub struct SerializablePyErr {
-    pub etype: String,
-    pub value: String,
-    pub traceback: Option<Result<String, String>>,
+    pub message: String,
 }
 
 impl SerializablePyErr {
     pub fn from(py: Python, err: &PyErr) -> Self {
-        let etype = format!("{}", err.get_type_bound(py));
-        let value = format!("{}", err.value_bound(py));
-        let traceback = err
-            .traceback_bound(py)
-            .map(|tb| tb.format().map_err(|e| format!("{}", e)));
-        Self {
-            etype,
-            value,
-            traceback,
+        // first construct the full traceback including any python frames that were used
+        // to invoke where we currently are. This is pre-pended to the traceback of the
+        // currently unwinded frames (err.traceback_bound())
+        let inspect = py.import_bound("inspect").unwrap();
+        let types = py.import_bound("types").unwrap();
+        let traceback_type = types.getattr("TracebackType").unwrap();
+        let traceback = py.import_bound("traceback").unwrap();
+
+        let mut f = inspect
+            .call_method0("currentframe")
+            .unwrap_or(PyNone::get_bound(py).to_owned().into_any());
+        let mut tb: Bound<'_, PyAny> = err.traceback_bound(py).unwrap().as_any().clone();
+        while !f.is_none() {
+            let lasti = f.getattr("f_lasti").unwrap();
+            let lineno = f.getattr("f_lineno").unwrap();
+            let back = f.getattr("f_back").unwrap();
+            tb = traceback_type.call1((tb, f, lasti, lineno)).unwrap();
+            f = back;
         }
+
+        let traceback_exception = traceback.getattr("TracebackException").unwrap();
+
+        let tb = traceback_exception
+            .call1((err.get_type_bound(py), err.value_bound(py), tb))
+            .unwrap();
+
+        let message: String = tb
+            .getattr("format")
+            .unwrap()
+            .call0()
+            .unwrap()
+            .iter()
+            .unwrap()
+            .map(|x| -> String { x.unwrap().extract().unwrap() })
+            .collect::<Vec<String>>()
+            .join("");
+
+        Self { message }
     }
 
-    pub fn from_fn<'py>(py: Python<'py>) -> impl Fn(PyErr) -> Self + use<'py> {
+    pub fn from_fn<'py>(py: Python<'py>) -> impl Fn(PyErr) -> Self + 'py {
         move |err| Self::from(py, &err)
     }
 }
 
 impl std::fmt::Display for SerializablePyErr {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        if let Some(tb_res) = &self.traceback {
-            match tb_res {
-                Ok(tb) => write!(f, "{}", tb)?,
-                Err(err) => write!(f, "Failed to extract traceback: {}", err)?,
-            }
-        }
-        write!(f, "{}: {}", self.etype, self.value)
+        write!(f, "{}", self.message)
+    }
+}
+
+impl<T> From<T> for SerializablePyErr
+where
+    T: Into<PyErr>,
+{
+    fn from(value: T) -> Self {
+        Python::with_gil(|py| SerializablePyErr::from(py, &value.into()))
     }
 }
 
@@ -177,12 +207,13 @@ mod tests {
 
             let err = SerializablePyErr::from(py, &module.call_method0("func3").unwrap_err());
             assert_eq!(
-                err.traceback.unwrap().unwrap().as_str(),
+                err.message.as_str(),
                 indoc! {r#"
                     Traceback (most recent call last):
                       File "test_helpers.py", line 8, in func3
                       File "test_helpers.py", line 5, in func2
                       File "test_helpers.py", line 2, in func1
+                    Exception: test
                 "#}
             );