Change ActorError print to TracebackException to capture context (#765)

Summary:
Pull Request resolved: #765

The current ActorError printout didn't include nested exceptions, which can be critical
for debugging if there are large wrapper try/except blocks.

Use the `TracebackException.format` method to print the exception instead to include
the nested exceptions.

Reviewed By: zdevito

Differential Revision: D79664075

fbshipit-source-id: 7c4ca55131567e753a527aae74930ba7873734d6

Author: dulinriley

python/monarch/_src/actor/actor_mesh.py

@@ -16,10 +16,9 @@
 import traceback
 
 from dataclasses import dataclass
-from traceback import extract_tb, StackSummary
+from traceback import TracebackException
 from typing import (
     Any,
-    AsyncGenerator,
     Awaitable,
     Callable,
     cast,
@@ -32,7 +31,6 @@
     Iterator,
     List,
     Literal,
-    NamedTuple,
     Optional,
     overload,
     ParamSpec,
@@ -902,16 +900,25 @@ def __init__(
         message: str = "A remote actor call has failed.",
     ) -> None:
         self.exception = exception
-        self.actor_mesh_ref_frames: StackSummary = extract_tb(exception.__traceback__)
+        # Need to stringify the exception early, because the PyPI package
+        # exceptiongroup may monkeypatch the "TracebackException" class for python
+        # versions < 3.11. If it gets unpickled in a different scope without
+        # using that monkeypatch, it'll have an exception in "format()".
+        # Store the traceback string instead which shouldn't change between machines.
+        actor_mesh_ref_tb = TracebackException.from_exception(exception).format()
+        # Replace any traceback lines to indicate it's a remote call traceback.
+        actor_mesh_ref_tb = (
+            s.replace(
+                "Traceback (most recent call last):",
+                "Traceback of where the remote call failed (most recent call last):",
+            )
+            for s in actor_mesh_ref_tb
+        )
+        self.exception_formatted = "".join(actor_mesh_ref_tb)
         self.message = message
 
     def __str__(self) -> str:
-        exe = str(self.exception)
-        actor_mesh_ref_tb = "".join(traceback.format_list(self.actor_mesh_ref_frames))
-        return (
-            f"{self.message}\n"
-            f"Traceback of where the remote call failed (most recent call last):\n{actor_mesh_ref_tb}{type(self.exception).__name__}: {exe}"
-        )
+        return f"{self.message}\n {self.exception_formatted}"
 
 
 def current_actor_name() -> str:

python/tests/test_actor_error.py

@@ -34,6 +34,24 @@ def raise_exception(self) -> None:
         raise Exception("This is a test exception")
 
 
+class NestedExceptionActor(Actor):
+    @endpoint
+    async def raise_exception_with_context(self) -> None:
+        try:
+            raise Exception("Inner exception")
+        except Exception:
+            # Don't use from here to set __context__ instead of __cause__
+            raise Exception("Outer exception")
+
+    @endpoint
+    async def raise_exception_with_cause(self) -> None:
+        try:
+            raise Exception("Inner exception")
+        except Exception as e:
+            # Use from here to set __cause__ instead of __context__
+            raise Exception("Outer exception") from e
+
+
 class BrokenPickleClass:
     """A class that can be configured to raise exceptions during pickling/unpickling."""
 
@@ -116,6 +134,41 @@ def test_actor_exception_sync(mesh, actor_class, num_procs):
             exception_actor.raise_exception.call().get()
 
 
+@pytest.mark.parametrize(
+    "mesh",
+    [local_proc_mesh, proc_mesh],
+    ids=["local_proc_mesh", "distributed_proc_mesh"],
+)
+async def test_actor_error_message(mesh):
+    """
+    Test that exceptions raised in actor endpoints capture nested exceptions.
+    """
+    proc = mesh(gpus=2)
+    exception_actor = await proc.spawn("exception_actor", NestedExceptionActor)
+
+    with pytest.raises(ActorError) as exc_info:
+        await exception_actor.raise_exception_with_cause.call()
+
+    # Make sure both exception messages are present in the message.
+    assert "Inner exception" in str(exc_info.value)
+    assert "Outer exception" in str(exc_info.value)
+    # Make sure the "cause" is set.
+    assert "The above exception was the direct cause of the following exception" in str(
+        exc_info.value
+    )
+
+    with pytest.raises(ActorError) as exc_info:
+        await exception_actor.raise_exception_with_context.call()
+
+    # Make sure both exception messages are present in the message.
+    assert "Inner exception" in str(exc_info.value)
+    assert "Outer exception" in str(exc_info.value)
+    # Make sure the "cause" is set.
+    assert "During handling of the above exception, another exception occurred" in str(
+        exc_info.value
+    )
+
+
 '''
 # oss_skip: importlib not pulling resource correctly in git CI, needs to be revisited
 @pytest.mark.oss_skip