feat: impl exception_auto_debugger contextmanager (#67)

sunray/__init__.py

@@ -51,4 +51,6 @@
 from ._internal.remote import remote as remote
 
 # remote pdb
+from ._internal.rpdb import post_mortem as post_mortem
 from ._internal.rpdb import set_trace as set_trace
+from ._internal.util import exception_auto_debugger as exception_auto_debugger

sunray/_internal/rpdb.py

@@ -91,10 +91,37 @@ def set_trace_by_madbg(frame: FrameType | None):
         if get_global_worker().ray_debugger_external
         else os.environ.get("REMOTE_PDB_HOST", "localhost")
     )
+    debugger: RemoteDebugger
     debugger, exit_stack = use_context(RemoteDebugger.connect_and_start(ip, port))
     debugger.set_trace(frame, done_callback=exit_stack.close)
 
 
+def post_mortem(traceback=None):
+    traceback = traceback or sys.exc_info()[2] or sys.last_traceback
+    if os.environ.get("SUNRAY_REMOTE_PDB", "yes").lower() in ["1", "yes", "true"]:
+        post_mortem_by_madbg(traceback)
+    else:
+        post_mortem_by_ray(traceback)
+
+
+def post_mortem_by_ray(traceback):
+    ray.util.rpdb._connect_ray_pdb(
+        breakpoint_uuid=None,
+        debugger_external=get_global_worker().ray_debugger_external,
+    ).post_mortem(traceback)
+
+
+def post_mortem_by_madbg(traceback):
+    port = int(os.environ.get("REMOTE_PDB_PORT", "0"))
+    ip: str = (
+        get_global_worker().node_ip_address
+        if get_global_worker().ray_debugger_external
+        else os.environ.get("REMOTE_PDB_HOST", "localhost")
+    )
+    with RemoteDebugger.connect_and_start(ip, port) as debugger:
+        debugger.post_mortem(traceback)
+
+
 class RemoteDebugger(RemoteIPythonDebugger):
     def __init__(self, stdin, stdout, context, **kwargs):
         # fix annoying `Warning: Input is not a terminal (fd=0)`

sunray/_internal/util.py

@@ -6,6 +6,7 @@
 
 
 if TYPE_CHECKING:
+    from types import TracebackType
     from typing import Callable
 
 
@@ -25,3 +26,38 @@ def get_num_returns(f: Callable) -> int:
         except AttributeError:  # pragma: no cover
             return len(ret_ast.body.slice.elts)  # type: ignore[attr-defined]
     return 1
+
+
+def exception_auto_debugger() -> ExceptionAutoDebugger:  # pragma: no cover
+    """
+    Auto post-mortem debugging for exceptions.
+
+    .. code-block:: python
+        import sunray
+
+
+        @sunray.remote
+        def f():
+            with sunray.exception_auto_debugger():
+                value1 = 1
+                value2 = 0
+                return value1 / value2
+
+
+        sunray.get(f.remote())
+    """
+
+    return ExceptionAutoDebugger()
+
+
+class ExceptionAutoDebugger:  # pragma: no cover
+    def __enter__(self) -> None: ...
+
+    def __exit__(self, exc_type, exc_value, tb: TracebackType) -> None:
+        if exc_type is not None:
+            import traceback
+
+            import sunray
+
+            traceback.print_tb(tb)
+            sunray.post_mortem(tb)