From e563a3c09864f1cd0f2d1815745f0712618195ae Mon Sep 17 00:00:00 2001
From: Sanghun Lee <sanghun@lablup.com>
Date: Tue, 20 Aug 2024 15:01:18 +0900
Subject: [PATCH] feat: Add log-level to background task message

---
 src/ai/backend/common/bgtask.py | 5 +++++
 src/ai/backend/common/defs.py   | 8 ++++++++
 src/ai/backend/common/events.py | 5 +++++
 3 files changed, 18 insertions(+)

diff --git a/src/ai/backend/common/bgtask.py b/src/ai/backend/common/bgtask.py
index 243b3c2d2e..30449882fd 100644
--- a/src/ai/backend/common/bgtask.py
+++ b/src/ai/backend/common/bgtask.py
@@ -29,6 +29,7 @@
 from ai.backend.logging import BraceStyleAdapter
 
 from . import redis_helper
+from .defs import BackgroundTaskLogLevel as LogLevel
 from .events import (
     BgtaskCancelledEvent,
     BgtaskDoneEvent,
@@ -71,6 +72,7 @@ async def update(
         self,
         increment: Union[int, float] = 0,
         message: str | None = None,
+        log_level: LogLevel = LogLevel.INFO,
     ) -> None:
         self.current_progress += increment
         # keep the state as local variables because they might be changed
@@ -88,6 +90,7 @@ async def _pipe_builder(r: Redis) -> Pipeline:
                     "total": str(total),
                     "msg": message or "",
                     "last_update": str(time.time()),
+                    "log_level": str(log_level),
                 },
             )
             await pipe.expire(tracker_key, MAX_BGTASK_ARCHIVE_PERIOD)
@@ -100,6 +103,7 @@ async def _pipe_builder(r: Redis) -> Pipeline:
                 message=message,
                 current_progress=current,
                 total_progress=total,
+                log_level=log_level,
             ),
         )
 
@@ -158,6 +162,7 @@ async def push_bgtask_events(
                         case BgtaskUpdatedEvent():
                             body["current_progress"] = event.current_progress
                             body["total_progress"] = event.total_progress
+                            body["log_level"] = event.log_level
                             await resp.send(json.dumps(body), event=event.name, retry=5)
                         case BgtaskDoneEvent():
                             if extra_data:
diff --git a/src/ai/backend/common/defs.py b/src/ai/backend/common/defs.py
index 9dadda4c90..0881d88c7b 100644
--- a/src/ai/backend/common/defs.py
+++ b/src/ai/backend/common/defs.py
@@ -1,3 +1,4 @@
+import enum
 from typing import Final
 
 # Redis database IDs depending on purposes
@@ -10,3 +11,10 @@
 
 
 DEFAULT_FILE_IO_TIMEOUT: Final = 10
+
+
+class BackgroundTaskLogLevel(enum.StrEnum):
+    INFO = enum.auto()
+    WARNING = enum.auto()
+    ERROR = enum.auto()
+    DEBUG = enum.auto()
diff --git a/src/ai/backend/common/events.py b/src/ai/backend/common/events.py
index 7869f07f39..3e7ef9780a 100644
--- a/src/ai/backend/common/events.py
+++ b/src/ai/backend/common/events.py
@@ -36,6 +36,8 @@
 from ai.backend.logging import BraceStyleAdapter, LogLevel
 
 from . import msgpack, redis_helper
+from .defs import BackgroundTaskLogLevel
+from .logging import BraceStyleAdapter
 from .types import (
     AgentId,
     EtcdRedisConfig,
@@ -559,6 +561,7 @@ class BgtaskUpdatedEvent(AbstractEvent):
     current_progress: float = attrs.field()
     total_progress: float = attrs.field()
     message: Optional[str] = attrs.field(default=None)
+    log_level: BackgroundTaskLogLevel = attrs.field(default=BackgroundTaskLogLevel.INFO)
 
     def serialize(self) -> tuple:
         return (
@@ -566,6 +569,7 @@ def serialize(self) -> tuple:
             self.current_progress,
             self.total_progress,
             self.message,
+            str(self.log_level),
         )
 
     @classmethod
@@ -575,6 +579,7 @@ def deserialize(cls, value: tuple):
             value[1],
             value[2],
             value[3],
+            BackgroundTaskLogLevel(value[4]),
         )