Adding Actor Id to Exceptions in Monarch Tracing Scuba (#1263)

Summary:
Pull Request resolved: #1263

The Actor Id column used to be null for Exceptions:
{F1982094709}
https://fburl.com/scuba/monarch_tracing/13hx70aj

Actor Id would show up for spans:
{F1982094744}
https://fburl.com/scuba/monarch_tracing/s2f6glbu

This change allows Exceptions to have Actor Id column populated as well:
{F1982094856}
https://fburl.com/scuba/monarch_tracing/757k16xp

Reviewed By: eliothedeman

Differential Revision: D82748854

fbshipit-source-id: af08f202e46928f83736c826c8df3420a4ff9ed5

Author: Hiren Patel

monarch_hyperactor/src/telemetry.rs

@@ -33,6 +33,13 @@ pub fn forward_to_tracing(py: Python, record: PyObject) -> PyResult<()> {
     let file = record.getattr(py, "filename")?;
     let file: &str = file.extract(py)?;
     let level: i32 = record.getattr(py, "levelno")?.extract(py)?;
+
+    // Extract actor_id from the Python record object if available
+    let actor_id = record
+        .getattr(py, "actor_id")
+        .ok()
+        .and_then(|attr| attr.extract::<String>(py).ok());
+
     // Map level number to level name
     match level {
         40 | 50 => {
@@ -53,21 +60,36 @@ pub fn forward_to_tracing(py: Python, record: PyObject) -> PyResult<()> {
                         file = file,
                         lineno = lineno,
                         stacktrace = traceback,
+                        actor_id = actor_id.as_deref(),
                         message
                     );
                 }
                 None => {
-                    tracing::error!(file = file, lineno = lineno, message);
+                    tracing::error!(
+                        file = file,
+                        lineno = lineno,
+                        actor_id = actor_id.as_deref(),
+                        message
+                    );
                 }
             }
         }
-        30 => tracing::warn!(target:"log_events", file = file, lineno = lineno, message),
-        20 => tracing::info!(target:"log_events", file = file, lineno = lineno, message),
-        10 => tracing::debug!(target:"log_events", file = file, lineno = lineno, message),
-        _ => tracing::info!(target:"log_events", file = file, lineno = lineno, message),
+        30 => {
+            tracing::warn!(target:"log_events", file = file, lineno = lineno, actor_id = actor_id.as_deref(), message)
+        }
+        20 => {
+            tracing::info!(target:"log_events", file = file, lineno = lineno, actor_id = actor_id.as_deref(), message)
+        }
+        10 => {
+            tracing::debug!(target:"log_events", file = file, lineno = lineno, actor_id = actor_id.as_deref(), message)
+        }
+        _ => {
+            tracing::info!(target:"log_events", file = file, lineno = lineno, actor_id = actor_id.as_deref(), message)
+        }
     }
     Ok(())
 }
+
 #[pyfunction]
 pub fn use_real_clock() -> PyResult<()> {
     swap_telemetry_clock(ClockKind::Real(RealClock));

python/monarch/_src/actor/actor_mesh.py

@@ -83,7 +83,7 @@
 from monarch._src.actor.python_extension_methods import rust_struct
 from monarch._src.actor.shape import MeshTrait, NDSlice
 from monarch._src.actor.sync_state import fake_sync_state
-from monarch._src.actor.telemetry import METER
+from monarch._src.actor.telemetry import METER, TracingForwarder
 from monarch._src.actor.tensor_engine_shim import actor_rref, actor_send
 from typing_extensions import Self
 
@@ -97,6 +97,7 @@
 CallMethod = PythonMessageKind.CallMethod
 
 logger: logging.Logger = logging.getLogger(__name__)
+logging.root.addHandler(TracingForwarder(level=logging.DEBUG))
 
 TRACER = get_monarch_tracer()
 

python/monarch/_src/actor/telemetry/__init__.py

@@ -28,6 +28,19 @@
 
 class TracingForwarder(logging.Handler):
     def emit(self, record: logging.LogRecord) -> None:
+        # Try to add actor_id from the current context to the logging record
+        try:
+            from monarch._src.actor.actor_mesh import context
+
+            ctx = context()
+            if ctx and ctx.actor_instance and ctx.actor_instance.actor_id:
+                # Add actor_id as an attribute to the logging record
+                setattr(record, "actor_id", str(ctx.actor_instance.actor_id))
+        except Exception:
+            # If we can't get the context or actor_id for any reason, just continue
+            # without adding the actor_id field
+            pass
+
         forward_to_tracing(record)