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

…ess (#3410)

Backported-from: main (25Q1)
Backported-to: 24.03

Co-authored-by: octodog <[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/common/logging.py

@@ -12,10 +12,12 @@
 import traceback
 from abc import ABCMeta, abstractmethod
 from collections import OrderedDict
+from collections.abc import Mapping, MutableMapping, Sequence
 from contextvars import ContextVar
 from datetime import datetime
 from pathlib import Path
-from typing import Any, Mapping, MutableMapping, Optional
+from types import TracebackType
+from typing import Any, Optional, TypeAlias, cast
 
 import coloredlogs
 import graypy
@@ -95,6 +97,10 @@
     }).allow_extra("*"),
 }).allow_extra("*")
 
+_SysExcInfoType: TypeAlias = (
+    tuple[type[BaseException], BaseException, TracebackType | None] | tuple[None, None, None]
+)
+
 
 class PickledException(Exception):
     """
@@ -193,10 +199,17 @@ def emit(self, record):
             self._sock.sendall(json.dumps(log).encode("utf-8"))
 
 
-def format_exception(self, ei):
-    s = "".join(ei)
-    if s[-1:] == "\n":
-        s = s[:-1]
+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