feat: add trace id to monitoring module (#4732)

* feat: add trace id to monitoring module * fix: consider else case
bentoml · May 16, 2024 · 2cf64ea · 2cf64ea
1 parent 50a053c
commit 2cf64ea
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 20 deletions.
diff --git a/src/bentoml/_internal/monitoring/default.py b/src/bentoml/_internal/monitoring/default.py
@@ -56,7 +56,11 @@ class DefaultMonitor(MonitorBase["JSONSerializable"]):
     data is logged as a JSON array.
     """
 
-    PRESERVED_COLUMNS = (COLUMN_TIME, COLUMN_RID) = ("timestamp", "request_id")
+    PRESERVED_COLUMNS = (COLUMN_TIME, COLUMN_RID, COLUMN_TID) = (
+        "timestamp",
+        "request_id",
+        "trace_id",
+    )
 
     def __init__(
         self,
@@ -136,10 +140,11 @@ def export_data(
             self._init_logger()
             assert self.data_logger is not None
 
-        extra_columns = dict(
-            timestamp=datetime.datetime.now().isoformat(),
-            request_id=str(trace_context.request_id),
-        )
+        extra_columns = {
+            self.COLUMN_TIME: datetime.datetime.now().isoformat(),
+            self.COLUMN_RID: str(trace_context.request_id),
+            self.COLUMN_TID: str(trace_context.trace_id),
+        }
         while True:
             try:
                 record = {k: v.popleft() for k, v in datas.items()}

diff --git a/src/bentoml/_internal/monitoring/otlp.py b/src/bentoml/_internal/monitoring/otlp.py
@@ -89,9 +89,10 @@ class OTLPMonitor(MonitorBase["JSONSerializable"]):
 
     """
 
-    PRESERVED_COLUMNS = (COLUMN_TIME, COLUMN_RID, COLUMN_META) = (
+    PRESERVED_COLUMNS = (COLUMN_TIME, COLUMN_RID, COLUMN_TID, COLUMN_META) = (
         "timestamp",
         "request_id",
+        "trace_id",
         "bento_meta",
     )
 
@@ -211,23 +212,22 @@ def export_data(
             self._init_logger()
             assert self.data_logger is not None
 
+        extra_columns: dict[str, t.Any] = {
+            self.COLUMN_TIME: datetime.datetime.now().timestamp(),
+            self.COLUMN_RID: str(trace_context.request_id),
+            self.COLUMN_TID: str(trace_context.trace_id),
+        }
+
         if self._will_export_schema or random.random() < self.meta_sample_rate:
-            extra_columns = {
-                self.COLUMN_TIME: datetime.datetime.now().timestamp(),
-                self.COLUMN_RID: str(trace_context.request_id),
-                self.COLUMN_META: {
-                    "bento_name": server_context.bento_name,
-                    "bento_version": server_context.bento_version,
-                    "monitor_name": self.name,
-                    "schema": self._schema,
-                },
+            extra_columns[self.COLUMN_META] = {
+                "bento_name": server_context.bento_name,
+                "bento_version": server_context.bento_version,
+                "monitor_name": self.name,
+                "schema": self._schema,
             }
+
             self._will_export_schema = False
-        else:
-            extra_columns = {
-                self.COLUMN_TIME: datetime.datetime.now().timestamp(),
-                self.COLUMN_RID: str(trace_context.request_id),
-            }
+
         while True:
             try:
                 record = {k: v.popleft() for k, v in datas.items()}