[monarch] Add configurable stacktrace on unawaited pytokio PythonTask

Pull Request resolved: #1679

This diff adds the environment variable `MONARCH_HYPERACTOR_UNAWAITED_PYTOKIO_TRACEBACK`, which is 0 by default. When it is set to 1, errors in unawaited pytokio `PythonTask` instances will print a traceback indicating where the `PythonTask` was created in addition to the error itself. This should only be used for development/debugging, because it adds overhead to compute the traceback for each new `PythonTask`.

This diff also makes it so that the panic trace isn't printed anymore, and the new warning makes it clear that the unawaited task is *not* the cause of a crash. This should ideally make it less scary and less likely to cause people to go down a fruitless rabbit hole.
ghstack-source-id: 319337699
@exported-using-ghexport

Differential Revision: [D85614449](https://our.internmc.facebook.com/intern/diff/D85614449/)

**NOTE FOR REVIEWERS**: This PR has internal Meta-specific changes or comments, please review them on [Phabricator](https://our.internmc.facebook.com/intern/diff/D85614449/)!

Author: samlurye

monarch_hyperactor/src/actor.rs

@@ -813,7 +813,7 @@ impl Handler<PythonMessage> for PythonActor {
         tokio::spawn(
             handle_async_endpoint_panic(
                 self.panic_sender.clone(),
-                PythonTask::new(future),
+                PythonTask::new(future)?,
                 receiver,
             )
             .instrument(

monarch_hyperactor/src/mailbox.rs

@@ -336,9 +336,9 @@ async fn recv_async(
 
 #[pymethods]
 impl PythonPortReceiver {
-    fn recv_task(&mut self) -> PyPythonTask {
+    fn recv_task(&mut self) -> PyResult<PyPythonTask> {
         let receiver = self.inner.clone();
-        PythonTask::new(recv_async(receiver)).into()
+        Ok(PythonTask::new(recv_async(receiver))?.into())
     }
 }
 
@@ -516,7 +516,7 @@ impl PythonOncePortReceiver {
                 .map_err(|err| PyErr::new::<PyEOFError, _>(format!("Port closed: {}", err)))
                 .and_then(|message| Python::with_gil(|py| message.into_py_any(py)))
         };
-        Ok(PythonTask::new(fut).into())
+        Ok(PythonTask::new(fut)?.into())
     }
 }
 

monarch_hyperactor/src/proc_mesh.rs

@@ -429,7 +429,7 @@ impl PyProcMesh {
         Ok(PythonTask::new(async move {
             Self::stop_mesh(inner, proc_events).await?;
             Python::with_gil(|py| Ok(py.None()))
-        })
+        })?
         .into())
     }
 }

monarch_hyperactor/src/pytokio.rs

@@ -10,8 +10,12 @@ use std::error::Error;
 use std::future::Future;
 use std::pin::Pin;
 
+use hyperactor::attrs::declare_attrs;
 use hyperactor::clock::Clock;
 use hyperactor::clock::RealClock;
+use hyperactor::config;
+use hyperactor::config::CONFIG;
+use hyperactor::config::ConfigAttr;
 use monarch_types::SerializablePyErr;
 use pyo3::IntoPyObjectExt;
 use pyo3::exceptions::PyRuntimeError;
@@ -20,6 +24,7 @@ use pyo3::exceptions::PyTimeoutError;
 use pyo3::exceptions::PyValueError;
 use pyo3::prelude::*;
 use pyo3::types::PyNone;
+use pyo3::types::PyString;
 use pyo3::types::PyType;
 use tokio::sync::Mutex;
 use tokio::sync::watch;
@@ -28,20 +33,68 @@ use tokio::task::JoinHandle;
 use crate::runtime::get_tokio_runtime;
 use crate::runtime::signal_safe_block_on;
 
+declare_attrs! {
+    /// If true, when a pytokio PythonTask fails, the traceback of the original callsite
+    /// will be logged.
+    @meta(CONFIG = ConfigAttr {
+        env_name: Some("MONARCH_HYPERACTOR_UNAWAITED_PYTOKIO_TRACEBACK".to_string()),
+        py_name: Some("unawaited_pytokio_traceback".to_string()),
+    })
+    pub attr UNAWAITED_PYTOKIO_TRACEBACK: u8 = 0;
+}
+
+fn current_traceback() -> PyResult<Option<PyObject>> {
+    if config::global::get(UNAWAITED_PYTOKIO_TRACEBACK) != 0 {
+        Python::with_gil(|py| {
+            Ok(Some(
+                py.import("traceback")?
+                    .call_method0("extract_stack")?
+                    .unbind(),
+            ))
+        })
+    } else {
+        Ok(None)
+    }
+}
+
+fn format_traceback(py: Python<'_>, traceback: &PyObject) -> PyResult<String> {
+    let tb = py
+        .import("traceback")?
+        .call_method1("format_list", (traceback,))?;
+    PyString::new(py, "")
+        .call_method1("join", (tb,))?
+        .extract::<String>()
+}
+
 /// Helper struct to make a Python future passable in an actor message.
 ///
 /// Also so that we don't have to write this massive type signature everywhere
 pub(crate) struct PythonTask {
     future: Mutex<Pin<Box<dyn Future<Output = PyResult<PyObject>> + Send + 'static>>>,
+    traceback: Option<PyObject>,
 }
 
 impl PythonTask {
-    pub(crate) fn new(fut: impl Future<Output = PyResult<PyObject>> + Send + 'static) -> Self {
+    fn new_with_traceback(
+        fut: impl Future<Output = PyResult<PyObject>> + Send + 'static,
+        traceback: Option<PyObject>,
+    ) -> Self {
         Self {
             future: Mutex::new(Box::pin(fut)),
+            traceback,
         }
     }
 
+    pub(crate) fn new(
+        fut: impl Future<Output = PyResult<PyObject>> + Send + 'static,
+    ) -> PyResult<Self> {
+        Ok(Self::new_with_traceback(fut, current_traceback()?))
+    }
+
+    fn traceback(&self) -> &Option<PyObject> {
+        &self.traceback
+    }
+
     pub(crate) fn take(self) -> Pin<Box<dyn Future<Output = PyResult<PyObject>> + Send + 'static>> {
         self.future.into_inner()
     }
@@ -101,17 +154,28 @@ impl PythonTaskAwaitIterator {
 }
 
 impl PyPythonTask {
-    pub fn new<F, T>(fut: F) -> PyResult<Self>
+    fn new_with_traceback<F, T>(fut: F, traceback: Option<PyObject>) -> PyResult<Self>
     where
         F: Future<Output = PyResult<T>> + Send + 'static,
         T: for<'py> IntoPyObject<'py>,
     {
-        Ok(PythonTask::new(async {
-            fut.await
-                .and_then(|t| Python::with_gil(|py| t.into_py_any(py)))
-        })
+        Ok(PythonTask::new_with_traceback(
+            async {
+                fut.await
+                    .and_then(|t| Python::with_gil(|py| t.into_py_any(py)))
+            },
+            traceback,
+        )
         .into())
     }
+
+    pub fn new<F, T>(fut: F) -> PyResult<Self>
+    where
+        F: Future<Output = PyResult<T>> + Send + 'static,
+        T: for<'py> IntoPyObject<'py>,
+    {
+        Self::new_with_traceback(fut, current_traceback()?)
+    }
 }
 
 fn to_py_error<T>(e: T) -> PyErr
@@ -131,6 +195,16 @@ impl PyPythonTask {
             .ok_or_else(|| PyValueError::new_err("PythonTask already consumed"))
     }
 
+    fn traceback(&self) -> PyResult<Option<PyObject>> {
+        if let Some(task) = &self.inner {
+            Ok(Python::with_gil(|py| {
+                task.traceback().as_ref().map(|t| t.clone_ref(py))
+            }))
+        } else {
+            Err(PyValueError::new_err("PythonTask already consumed"))
+        }
+    }
+
     /// Prefer spawn_abortable over spawn if the future can be safely cancelled
     /// when it is dropped.
     /// This way any resources it is using will be freed up. This is especially
@@ -141,30 +215,41 @@ impl PyPythonTask {
     /// PyShared is dropped.
     pub(crate) fn spawn_abortable(&mut self) -> PyResult<PyShared> {
         let (tx, rx) = watch::channel(None);
+        let traceback = self.traceback()?;
+        let traceback1 = self.traceback()?;
         let task = self.take_task()?;
         let handle = get_tokio_runtime().spawn(async move {
-            send_result(tx, task.await);
+            send_result(tx, task.await, traceback1);
         });
         Ok(PyShared {
             rx,
             handle,
             abort: true,
+            traceback,
         })
     }
 }
 
 fn send_result(
     tx: tokio::sync::watch::Sender<Option<PyResult<PyObject>>>,
     result: PyResult<PyObject>,
+    traceback: Option<PyObject>,
 ) {
     // a SendErr just means that there are no consumers of the value left.
     match tx.send(Some(result)) {
         Err(tokio::sync::watch::error::SendError(Some(Err(pyerr)))) => {
             Python::with_gil(|py| {
-                panic!(
-                    "PythonTask errored but is not being awaited: {}",
-                    SerializablePyErr::from(py, &pyerr)
-                )
+                let tb = if let Some(tb) = traceback {
+                    format_traceback(py, &tb).unwrap()
+                } else {
+                    "None (run with `MONARCH_HYPERACTOR_UNAWAITED_PYTOKIO_TRACEBACK=1` to see a traceback here)\n".into()
+                };
+                tracing::error!(
+                    "PythonTask errored but is not being awaited; this will not crash your program, but indicates that \
+                    something went wrong.\n{}\nTraceback where the task was created (most recent call last):\n{}",
+                    SerializablePyErr::from(py, &pyerr),
+                    tb
+                );
             });
         }
         _ => {}
@@ -175,6 +260,7 @@ fn send_result(
 impl PyPythonTask {
     fn block_on(mut slf: PyRefMut<PyPythonTask>, py: Python<'_>) -> PyResult<PyObject> {
         let task = slf.take_task()?;
+
         // mutable references to python objects must be dropped before calling
         // signal_safe_block_on. It will release the GIL, and any other thread
         // trying to access slf will throw.
@@ -184,14 +270,17 @@ impl PyPythonTask {
 
     pub(crate) fn spawn(&mut self) -> PyResult<PyShared> {
         let (tx, rx) = watch::channel(None);
+        let traceback = self.traceback()?;
+        let traceback1 = self.traceback()?;
         let task = self.take_task()?;
         let handle = get_tokio_runtime().spawn(async move {
-            send_result(tx, task.await);
+            send_result(tx, task.await, traceback1);
         });
         Ok(PyShared {
             rx,
             handle,
             abort: false,
+            traceback,
         })
     }
 
@@ -278,26 +367,35 @@ impl PyPythonTask {
     }
 
     fn with_timeout(&mut self, seconds: f64) -> PyResult<PyPythonTask> {
+        let tb = self.traceback()?;
         let task = self.take_task()?;
-        PyPythonTask::new(async move {
-            RealClock
-                .timeout(std::time::Duration::from_secs_f64(seconds), task)
-                .await
-                .map_err(|_| PyTimeoutError::new_err(()))?
-        })
+        PyPythonTask::new_with_traceback(
+            async move {
+                RealClock
+                    .timeout(std::time::Duration::from_secs_f64(seconds), task)
+                    .await
+                    .map_err(|_| PyTimeoutError::new_err(()))?
+            },
+            tb,
+        )
     }
 
     #[staticmethod]
     fn spawn_blocking(f: PyObject) -> PyResult<PyShared> {
         let (tx, rx) = watch::channel(None);
+        let traceback = current_traceback()?;
+        let traceback1 = traceback
+            .as_ref()
+            .map_or_else(|| None, |t| Python::with_gil(|py| Some(t.clone_ref(py))));
         let handle = get_tokio_runtime().spawn_blocking(move || {
             let result = Python::with_gil(|py| f.call0(py));
-            send_result(tx, result);
+            send_result(tx, result, traceback1);
         });
         Ok(PyShared {
             rx,
             handle,
             abort: false,
+            traceback,
         })
     }
 
@@ -342,6 +440,7 @@ pub struct PyShared {
     rx: watch::Receiver<Option<PyResult<PyObject>>>,
     handle: JoinHandle<()>,
     abort: bool,
+    traceback: Option<PyObject>,
 }
 
 impl Drop for PyShared {
@@ -363,20 +462,27 @@ impl PyShared {
         // we can have multiple awaiters get triggered by the same change.
         // self.rx will always be in the state where it hasn't see the change yet.
         let mut rx = self.rx.clone();
-        PyPythonTask::new(async move {
-            rx.changed().await.map_err(to_py_error)?;
-            let b = rx.borrow();
-            let r = b.as_ref().unwrap();
-            Python::with_gil(|py| match r {
-                Ok(v) => Ok(v.bind(py).clone().unbind()),
-                Err(err) => Err(err.clone_ref(py)),
-            })
-        })
+        PyPythonTask::new_with_traceback(
+            async move {
+                rx.changed().await.map_err(to_py_error)?;
+                let b = rx.borrow();
+                let r = b.as_ref().unwrap();
+                Python::with_gil(|py| match r {
+                    Ok(v) => Ok(v.bind(py).clone().unbind()),
+                    Err(err) => Err(err.clone_ref(py)),
+                })
+            },
+            self.traceback
+                .as_ref()
+                .map_or_else(|| None, |t| Python::with_gil(|py| Some(t.clone_ref(py)))),
+        )
     }
+
     fn __await__(&mut self, py: Python<'_>) -> PyResult<PythonTaskAwaitIterator> {
         let task = self.task()?;
         Ok(PythonTaskAwaitIterator::new(task.into_py_any(py)?))
     }
+
     pub fn block_on(mut slf: PyRefMut<PyShared>, py: Python<'_>) -> PyResult<PyObject> {
         let task = slf.task()?.take_task()?;
         // mutable references to python objects must be dropped before calling