fix(BA-477): Logging and formatting exceptions raised from local proc…

…ess (#3410) (#3414)

Co-authored-by: Joongi Kim <[email protected]>

changes/3410.fix.md

@@ -0,0 +1 @@
+Fix formatting errors when logging exceptions raised from the current local process that did not pass our custom serialization step

src/ai/backend/logging/formatter.py

@@ -3,17 +3,31 @@
 import logging
 import pprint
 import time
+import traceback
+from collections.abc import Sequence
 from datetime import datetime
-from typing import Any
+from types import TracebackType
+from typing import Any, TypeAlias, cast
 
 import coloredlogs
 from pythonjsonlogger.json import JsonFormatter
 
-
-def format_exception(self, ei) -> str:
-    s = "".join(ei)
-    if s[-1:] == "\n":
-        s = s[:-1]
+_SysExcInfoType: TypeAlias = (
+    tuple[type[BaseException], BaseException, TracebackType | None] | tuple[None, None, None]
+)
+
+
+def format_exception(self, ei: Sequence[str] | _SysExcInfoType) -> str:
+    match ei:
+        case (str(), *_):
+            # Already foramtted from the source process for ease of serialization
+            s = "".join(cast(Sequence[str], ei))  # cast is required for mypy
+        case (type(), BaseException(), _):
+            # A live exc_info object from the current process
+            s = "".join(traceback.format_exception(*ei))
+        case _:
+            s = "<exception-info-unavailable>"
+    s = s.rstrip("\n")
     return s