Skip to content

Commit

Permalink
feat: Add log-level to background task message
Browse files Browse the repository at this point in the history
  • Loading branch information
fregataa committed Nov 9, 2024
1 parent da3b607 commit e563a3c
Show file tree
Hide file tree
Showing 3 changed files with 18 additions and 0 deletions.
5 changes: 5 additions & 0 deletions src/ai/backend/common/bgtask.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand All @@ -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)
Expand All @@ -100,6 +103,7 @@ async def _pipe_builder(r: Redis) -> Pipeline:
message=message,
current_progress=current,
total_progress=total,
log_level=log_level,
),
)

Expand Down Expand Up @@ -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:
Expand Down
8 changes: 8 additions & 0 deletions src/ai/backend/common/defs.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import enum
from typing import Final

# Redis database IDs depending on purposes
Expand All @@ -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()
5 changes: 5 additions & 0 deletions src/ai/backend/common/events.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -559,13 +561,15 @@ 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 (
str(self.task_id),
self.current_progress,
self.total_progress,
self.message,
str(self.log_level),
)

@classmethod
Expand All @@ -575,6 +579,7 @@ def deserialize(cls, value: tuple):
value[1],
value[2],
value[3],
BackgroundTaskLogLevel(value[4]),
)


Expand Down

0 comments on commit e563a3c

Please sign in to comment.