From ca28e1e7a82b4467012e27934c2cddda39ea2c6c Mon Sep 17 00:00:00 2001
From: Mert <101130780+mertalev@users.noreply.github.com>
Date: Thu, 25 Jan 2024 19:26:27 -0500
Subject: [PATCH] fix(ml): error logging (#6646)

* fix ml error logging

* exclude certain libraries from traceback
---
 machine-learning/app/config.py | 24 +++++++++++++++++++++---
 machine-learning/log_conf.json | 11 +++++------
 2 files changed, 26 insertions(+), 9 deletions(-)

diff --git a/machine-learning/app/config.py b/machine-learning/app/config.py
index 8c588ea548..08eeae7d5f 100644
--- a/machine-learning/app/config.py
+++ b/machine-learning/app/config.py
@@ -1,10 +1,10 @@
+import concurrent.futures
 import logging
 import os
 import sys
 from pathlib import Path
 from socket import socket
 
-import starlette
 from gunicorn.arbiter import Arbiter
 from pydantic import BaseSettings
 from rich.console import Console
@@ -74,10 +74,28 @@ log_settings = LogSettings()
 class CustomRichHandler(RichHandler):
     def __init__(self) -> None:
         console = Console(color_system="standard", no_color=log_settings.no_color)
-        super().__init__(show_path=False, omit_repeated_times=False, console=console, tracebacks_suppress=[starlette])
+        self.excluded = ["uvicorn", "starlette", "fastapi"]
+        super().__init__(
+            show_path=False,
+            omit_repeated_times=False,
+            console=console,
+            rich_tracebacks=True,
+            tracebacks_suppress=[*self.excluded, concurrent.futures],
+        )
+
+    # hack to exclude certain modules from rich tracebacks
+    def emit(self, record: logging.LogRecord) -> None:
+        if record.exc_info is not None:
+            tb = record.exc_info[2]
+            while tb is not None:
+                if any(excluded in tb.tb_frame.f_code.co_filename for excluded in self.excluded):
+                    tb.tb_frame.f_locals["_rich_traceback_omit"] = True
+                tb = tb.tb_next
+
+        return super().emit(record)
 
 
-log = logging.getLogger("gunicorn.access")
+log = logging.getLogger("ml.log")
 log.setLevel(LOG_LEVELS.get(log_settings.log_level.lower(), logging.INFO))
 
 
diff --git a/machine-learning/log_conf.json b/machine-learning/log_conf.json
index ae44f3a032..8cb09fc666 100644
--- a/machine-learning/log_conf.json
+++ b/machine-learning/log_conf.json
@@ -1,16 +1,15 @@
 {
   "version": 1,
-  "disable_existing_loggers": true,
-  "formatters": { "rich": { "show_path": false, "omit_repeated_times": false } },
+  "disable_existing_loggers": false,
   "handlers": {
     "console": {
-      "class": "app.config.CustomRichHandler",
-      "formatter": "rich"
+      "class": "app.config.CustomRichHandler"
     }
   },
   "loggers": {
-    "gunicorn.access": { "propagate": true },
-    "gunicorn.error": { "propagate": true }
+    "gunicorn.error": {
+      "handlers": ["console"]
+    }
   },
   "root": { "handlers": ["console"] }
 }