Improve trace collection for python actors (pytorch-labs#242)

Summary:
Pull Request resolved: pytorch-labs#242

stacktraces are now correctly collected from python exceptions.

Reviewed By: kaiyuan-li

Differential Revision: D76468993

Author: eliothedeman

hyperactor/src/proc.rs

@@ -876,7 +876,6 @@ impl<A: Actor> Instance<A> {
         Ok(actor_handle)
     }
 
-    #[crate::instrument_infallible(fields(actor_id=tracing::field::display(self.self_id())))]
     async fn serve(mut self, mut actor: A) {
         let result = self.run_actor_tree(&mut actor).await;
 

hyperactor_extension/src/telemetry.rs

@@ -11,6 +11,7 @@
 use std::cell::Cell;
 
 use pyo3::prelude::*;
+use pyo3::types::PyTraceback;
 use tracing::span::EnteredSpan;
 // Thread local to store the current span
 thread_local! {
@@ -45,15 +46,47 @@ pub fn exit_span() -> PyResult<()> {
 
 /// Log a message with the given metaata
 #[pyfunction]
-pub fn forward_to_tracing(message: &str, file: &str, lineno: i64, level: i32) {
+pub fn forward_to_tracing(py: Python, record: PyObject) -> PyResult<()> {
+    let message = record.call_method0(py, "getMessage")?;
+    let message: &str = message.extract(py)?;
+    let lineno: i64 = record.getattr(py, "lineno")?.extract(py)?;
+    let file = record.getattr(py, "filename")?;
+    let file: &str = file.extract(py)?;
+    let level: i32 = record.getattr(py, "levelno")?.extract(py)?;
+
     // Map level number to level name
     match level {
-        40 => tracing::error!(file = file, lineno = lineno, message),
+        40 | 50 => {
+            let exc = record.getattr(py, "exc_info").ok();
+            let traceback = exc
+                .and_then(|exc| {
+                    if exc.is_none(py) {
+                        return None;
+                    }
+                    exc.extract::<(PyObject, PyObject, Bound<'_, PyTraceback>)>(py)
+                        .ok()
+                })
+                .map(|(_, _, tb)| tb.format().unwrap_or_default());
+            match traceback {
+                Some(traceback) => {
+                    tracing::error!(
+                        file = file,
+                        lineno = lineno,
+                        stacktrace = traceback,
+                        message
+                    );
+                }
+                None => {
+                    tracing::error!(file = file, lineno = lineno, message);
+                }
+            }
+        }
         30 => tracing::warn!(file = file, lineno = lineno, message),
         20 => tracing::info!(file = file, lineno = lineno, message),
         10 => tracing::debug!(file = file, lineno = lineno, message),
         _ => tracing::info!(file = file, lineno = lineno, message),
     }
+    Ok(())
 }
 
 pub fn register_python_bindings(module: &Bound<'_, PyModule>) -> PyResult<()> {

python/monarch/_rust_bindings/hyperactor_extension/telemetry.pyi

@@ -4,23 +4,27 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-def forward_to_tracing(message: str, file: str, lineno: int, level: int) -> None:
+import logging
+
+def forward_to_tracing(record: logging.LogRecord) -> None:
     """
     Log a message with the given metadata using the tracing system.
 
     This function forwards Python log messages to the Rust tracing system,
     preserving the original source location and log level.
 
     Args:
-    - message (str): The log message content.
-    - file (str): The file where the log message originated.
-    - lineno (int): The line number where the log message originated.
-    - level (int): The log level:
-        - 10: DEBUG
-        - 20: INFO
-        - 30: WARN
-        - 40: ERROR
-        - other values default to INFO
+    - record (logging.LogRecord): The log record containing message, file, lineno, and level information.
+      The function extracts:
+        - message: The log message content via record.getMessage()
+        - file: The filename via record.filename
+        - lineno: The line number via record.lineno
+        - level: The log level via record.levelno:
+            - 10: DEBUG
+            - 20: INFO
+            - 30: WARN
+            - 40: ERROR
+            - other values default to INFO
     """
     ...
 

python/monarch/actor_mesh.py

@@ -525,7 +525,14 @@ async def instrumented():
                     enter_span(
                         the_method.__module__, message.method, str(ctx.mailbox.actor_id)
                     )
-                    result = await the_method(self.instance, *args, **kwargs)
+                    try:
+                        result = await the_method(self.instance, *args, **kwargs)
+                    except Exception as e:
+                        logging.critical(
+                            "Unahndled exception in actor endpoint",
+                            exc_info=e,
+                        )
+                        raise e
                     exit_span()
                     return result
 

python/monarch/bootstrap_main.py

@@ -30,28 +30,9 @@ def invoke_main():
     # behavior of std out as if it were a terminal.
     sys.stdout.reconfigure(line_buffering=True)
     global bootstrap_main
-    from monarch._rust_bindings.hyperactor_extension.telemetry import (  # @manual=//monarch/monarch_extension:monarch_extension  # @manual=//monarch/monarch_extension:monarch_extension
-        forward_to_tracing,
-    )
 
     # TODO: figure out what from worker_main.py we should reproduce here.
-
-    class TracingForwarder(logging.Handler):
-        def emit(self, record: logging.LogRecord) -> None:
-            try:
-                forward_to_tracing(
-                    record.getMessage(),
-                    record.filename or "",
-                    record.lineno or 0,
-                    record.levelno,
-                )
-            except AttributeError:
-                forward_to_tracing(
-                    record.__str__(),
-                    record.filename or "",
-                    record.lineno or 0,
-                    record.levelno,
-                )
+    from monarch.telemetry import TracingForwarder
 
     if os.environ.get("MONARCH_ERROR_DURING_BOOTSTRAP_FOR_TESTING") == "1":
         raise RuntimeError("Error during bootstrap for testing")

python/monarch/telemetry.py

@@ -0,0 +1,19 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-strict
+
+
+import logging
+
+from monarch._rust_bindings.hyperactor_extension.telemetry import (  # @manual=//monarch/monarch_extension:monarch_extension
+    forward_to_tracing,
+)
+
+
+class TracingForwarder(logging.Handler):
+    def emit(self, record: logging.LogRecord) -> None:
+        forward_to_tracing(record)