Fix python logs in asyncio code not showing up in scuba

Differential Revision: D75957462

Author: eliothedeman

hyperactor_extension/src/telemetry.rs

@@ -8,9 +8,47 @@
 
 #![allow(unsafe_op_in_unsafe_fn)]
 
+use std::cell::Cell;
+
 use pyo3::prelude::*;
+use tracing::span::EnteredSpan;
+// Thread local to store the current span
+thread_local! {
+    static ACTIVE_ACTOR_ID: Cell<Option<String>> = const { Cell::new(None) };
+    static ACTIVE_ACTOR_SPAN: Cell<Option<EnteredSpan>> = const { Cell::new(None) };
+}
+
+pub fn set_py_actor(name: String) {
+    ACTIVE_ACTOR_ID.replace(Some(name));
+}
+
+/// Enter the span stored in the thread local
+#[pyfunction]
+pub fn enter_span(module_name: String, method_name: String, actor_id: String) -> PyResult<()> {
+    let mut maybe_span = ACTIVE_ACTOR_SPAN.take();
+    if maybe_span.is_none() {
+        maybe_span = Some(
+            tracing::info_span!(
+                "py_actor_method",
+                name = method_name,
+                target = module_name,
+                actor_id = actor_id
+            )
+            .entered(),
+        );
+    }
+    ACTIVE_ACTOR_SPAN.set(maybe_span);
+    Ok(())
+}
 
-/// Log a message with the given metadata
+/// Exit the span stored in the thread local
+#[pyfunction]
+pub fn exit_span() -> PyResult<()> {
+    ACTIVE_ACTOR_SPAN.replace(None);
+    Ok(())
+}
+
+/// Log a message with the given metaata
 #[pyfunction]
 pub fn forward_to_tracing(message: &str, file: &str, lineno: i64, level: i32) {
     // Map level number to level name
@@ -23,15 +61,29 @@ pub fn forward_to_tracing(message: &str, file: &str, lineno: i64, level: i32) {
     }
 }
 
-use pyo3::Bound;
-use pyo3::types::PyModule;
-
 pub fn register_python_bindings(module: &Bound<'_, PyModule>) -> PyResult<()> {
+    // Register the forward_to_tracing function
     let f = wrap_pyfunction!(forward_to_tracing, module)?;
     f.setattr(
         "__module__",
         "monarch._rust_bindings.hyperactor_extension.telemetry",
     )?;
     module.add_function(f)?;
+
+    // Register the span-related functions
+    let enter_span_fn = wrap_pyfunction!(enter_span, module)?;
+    enter_span_fn.setattr(
+        "__module__",
+        "monarch._rust_bindings.hyperactor_extension.telemetry",
+    )?;
+    module.add_function(enter_span_fn)?;
+
+    let exit_span_fn = wrap_pyfunction!(exit_span, module)?;
+    exit_span_fn.setattr(
+        "__module__",
+        "monarch._rust_bindings.hyperactor_extension.telemetry",
+    )?;
+    module.add_function(exit_span_fn)?;
+
     Ok(())
 }

monarch_hyperactor/src/actor.rs

@@ -33,6 +33,7 @@ use serde::Deserialize;
 use serde::Serialize;
 use serde_bytes::ByteBuf;
 use tokio::sync::Mutex;
+use tracing::span::Id;
 
 use crate::mailbox::PyMailbox;
 use crate::proc::InstanceWrapper;
@@ -292,6 +293,7 @@ impl Handler<PythonMessage> for PythonActor {
         this: &Instance<Self>,
         message: PythonMessage,
     ) -> anyhow::Result<()> {
+        this.self_id();
         let future = Python::with_gil(|py| -> PyResult<_> {
             let mailbox = PyMailbox {
                 inner: this.mailbox_for_py().clone(),

python/monarch/_rust_bindings/hyperactor_extension/telemetry.pyi

@@ -15,3 +15,20 @@ def forward_to_tracing(message: str, file: str, lineno: int, level: int) -> None
     - level (int): The log level (10 for debug, 20 for info, 30 for warn, 40 for error).
     """
     ...
+
+def enter_span(module_name: str, method_name: str, actor_id: str) -> None:
+    """
+    Enter a tracing span for a Python actor method.
+
+    Args:
+    - module_name (str): The name of the module containing the actor.
+    - method_name (str): The name of the method being called.
+    - actor_id (str): The ID of the actor instance.
+    """
+    ...
+
+def exit_span() -> None:
+    """
+    Exit the current tracing span for a Python actor method.
+    """
+    ...

python/monarch/actor_mesh.py

@@ -7,6 +7,7 @@
 import asyncio
 import collections
 import contextvars
+import functools
 import inspect
 
 import itertools
@@ -38,6 +39,7 @@
 
 import monarch
 from monarch import ActorFuture as Future
+from monarch._rust_bindings.hyperactor_extension.telemetry import enter_span, exit_span
 
 from monarch._rust_bindings.monarch_hyperactor.actor import PythonMessage
 from monarch._rust_bindings.monarch_hyperactor.actor_mesh import PythonActorMesh
@@ -49,6 +51,7 @@
 )
 from monarch._rust_bindings.monarch_hyperactor.proc import ActorId
 from monarch._rust_bindings.monarch_hyperactor.shape import Point as HyPoint, Shape
+
 from monarch.common.pickle_flatten import flatten, unflatten
 from monarch.common.shape import MeshTrait, NDSlice, Shape
 
@@ -487,13 +490,31 @@ def handle_cast(
                 return None
             else:
                 the_method = getattr(self.instance, message.method)._method
-                result = the_method(self.instance, *args, **kwargs)
+
                 if not inspect.iscoroutinefunction(the_method):
+                    enter_span(
+                        the_method.__module__, message.method, str(ctx.mailbox.actor_id)
+                    )
+                    result = the_method(self.instance, *args, **kwargs)
+                    exit_span()
                     if port is not None:
                         port.send("result", result)
                     return None
 
-                return self.run_async(ctx, self.run_task(port, result))
+                async def instrumented():
+                    enter_span(
+                        the_method.__module__, message.method, str(ctx.mailbox.actor_id)
+                    )
+                    result = await the_method(self.instance, *args, **kwargs)
+                    exit_span()
+                    return result
+
+                return self.run_async(
+                    self.instance.__module__,
+                    message.method,
+                    ctx,
+                    self.run_task(port, instrumented()),
+                )
         except Exception as e:
             traceback.print_exc()
             s = ActorMeshRefCallFailedException(e)
@@ -505,7 +526,13 @@ def handle_cast(
             else:
                 raise s from None
 
-    async def run_async(self, ctx, coroutine):
+    async def run_async(
+        self,
+        module_name: str,
+        method_name: str,
+        ctx: MonarchContext,
+        coroutine: Coroutine[Any, None, Any],
+    ) -> None:
         _context.set(ctx)
         if self.complete_task is None:
             asyncio.create_task(self._complete())
@@ -545,6 +572,12 @@ def _unpickle(data: bytes, mailbox: Mailbox) -> Any:
 
 
 class Actor(MeshTrait):
+    @functools.cached_property
+    def logger(cls) -> logging.Logger:
+        lgr = logging.getLogger(cls.__class__.__name__)
+        lgr.setLevel(logging.DEBUG)
+        return lgr
+
     @property
     def _ndslice(self) -> NDSlice:
         raise NotImplementedError(

python/monarch/bootstrap_main.py

@@ -55,7 +55,7 @@ def emit(self, record: logging.LogRecord) -> None:
 
     # forward logs to rust tracing. Defaults to on.
     if os.environ.get("MONARCH_PYTHON_LOG_TRACING", "1") == "1":
-        logging.root.addHandler(TracingForwarder())
+        logging.root.addHandler(TracingForwarder(level=logging.DEBUG))
 
     with (
         importlib.resources.path("monarch", "py-spy") as pyspy,