Move critical log of exceptions on endpoints to monarch file (#1934)

Summary:
Pull Request resolved: #1934

Move the log of endpoint exceptions from the user logger (logging in python)
to use monarch's tracing logger which goes to a separate file.
Log as info because it's only an error to the user's perspective, not from monarch's
perspective.

Reviewed By: zdevito

Differential Revision: D87459440

fbshipit-source-id: 586a38c2ebeb19efa8596cdd362e6b84b0e5af4b

Author: dulinriley

monarch_hyperactor/src/v1/logging.rs

@@ -24,6 +24,7 @@ use ndslice::View;
 use pyo3::Bound;
 use pyo3::prelude::*;
 use pyo3::types::PyModule;
+use pyo3::types::PyString;
 
 use crate::context::PyInstance;
 use crate::instance_dispatch;
@@ -416,13 +417,53 @@ impl Drop for LoggingMeshClient {
     }
 }
 
+/// Turns a python exception into a string with a traceback. If the traceback doesn't
+/// exist or can't be formatted, returns just the exception message.
+fn format_traceback<'py>(py: Python<'py>, err: PyErr) -> String {
+    let traceback = err.traceback(py);
+    if traceback.is_some() {
+        let inner = || -> PyResult<String> {
+            let formatted = py
+                .import("traceback")?
+                .call_method1("format_exception", (err.clone_ref(py),))?;
+            Ok(PyString::new(py, "")
+                .call_method1("join", (formatted,))?
+                .to_string())
+        };
+        match inner() {
+            Ok(s) => s,
+            Err(e) => format!("{}: no traceback {}", err, e),
+        }
+    } else {
+        err.to_string()
+    }
+}
+
+#[pyfunction]
+fn log_endpoint_exception<'py>(py: Python<'py>, e: PyObject, endpoint: PyObject) {
+    let pyerr = PyErr::from_value(e.into_bound(py));
+    let exception_str = format_traceback(py, pyerr);
+    let endpoint = endpoint.into_bound(py).to_string();
+    tracing::info!(
+        %endpoint,
+        "exception occurred in endpoint: {}",
+        exception_str,
+    );
+}
+
 /// Register the Python-facing types for this module.
 ///
 /// `pyo3` calls this when building `monarch._rust_bindings...`. We
 /// expose `LoggingMeshClient` so that Python can construct it and
 /// call its methods (`spawn`, `set_mode`, `flush`, ...).
 pub fn register_python_bindings(module: &Bound<'_, PyModule>) -> PyResult<()> {
     module.add_class::<LoggingMeshClient>()?;
+    let log_endpoint_exception = wrap_pyfunction!(log_endpoint_exception, module.py())?;
+    log_endpoint_exception.setattr(
+        "__module__",
+        "monarch._rust_bindings.monarch_hyperactor.v1.logging",
+    )?;
+    module.add_function(log_endpoint_exception)?;
     Ok(())
 }
 

python/monarch/_rust_bindings/monarch_hyperactor/v1/logging.pyi

@@ -32,3 +32,5 @@ class LoggingMeshClient:
         level: int,
     ) -> None: ...
     def flush(self, instance: Instance) -> PythonTask[None]: ...
+
+def log_endpoint_exception(e: Exception, endpoint: str | None) -> None: ...

python/monarch/_src/actor/actor_mesh.py

@@ -72,6 +72,7 @@
     Region,
     Shape,
 )
+from monarch._rust_bindings.monarch_hyperactor.v1.logging import log_endpoint_exception
 from monarch._rust_bindings.monarch_hyperactor.value_mesh import (
     ValueMesh as HyValueMesh,
 )
@@ -970,6 +971,7 @@ async def handle(
         local_state: Iterable[Any],
         response_port: "PortProtocol[Any]",
     ) -> None:
+        method_name = None
         MESSAGES_HANDLED.add(1)
         # response_port can be None. If so, then sending to port will drop the response,
         # and raise any exceptions to the caller.
@@ -1000,14 +1002,15 @@ async def handle(
                         self._saved_error = ActorError(
                             e, f"Remote actor {Class}.__init__ call failed."
                         )
-                        raise e
+                        raise
                     response_port.send(None)
                     return None
                 case MethodSpecifier.ReturnsResponse(name=method_name):
                     pass
                 case MethodSpecifier.ExplicitPort(name=method_name):
                     args = (response_port, *args)
                     response_port = DroppingPort()
+            assert isinstance(method_name, str)
 
             if self.instance is None:
                 # This could happen because of the following reasons. Both
@@ -1036,22 +1039,15 @@ async def handle(
                 the_method = functools.partial(the_method._method, self.instance)
 
             if inspect.iscoroutinefunction(the_method):
-                try:
-                    if should_instrument:
-                        with TRACER.start_as_current_span(
-                            method_name,
-                            attributes={"actor_id": str(ctx.actor_instance.actor_id)},
-                        ):
-                            result = await the_method(*args, **kwargs)
-                    else:
+                if should_instrument:
+                    with TRACER.start_as_current_span(
+                        method_name,
+                        attributes={"actor_id": str(ctx.actor_instance.actor_id)},
+                    ):
                         result = await the_method(*args, **kwargs)
-                    self._maybe_exit_debugger()
-                except Exception as e:
-                    logging.critical(
-                        "Unhandled exception in actor endpoint",
-                        exc_info=e,
-                    )
-                    raise e
+                else:
+                    result = await the_method(*args, **kwargs)
+                self._maybe_exit_debugger()
             else:
                 with fake_sync_state():
                     if should_instrument:
@@ -1066,6 +1062,7 @@ async def handle(
 
             response_port.send(result)
         except Exception as e:
+            log_endpoint_exception(e, method_name)
             self._post_mortem_debug(e.__traceback__)
             response_port.exception(ActorError(e))
         except BaseException as e: