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 Aug 20, 2024
1 parent ee7792f commit c0d9a00
Show file tree
Hide file tree
Showing 3 changed files with 17 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 @@ -27,6 +27,7 @@
from redis.asyncio.client import Pipeline

from . import redis_helper
from .defs import BackgroundTaskLogLevel as LogLevel
from .events import (
BgtaskCancelledEvent,
BgtaskDoneEvent,
Expand Down Expand Up @@ -70,6 +71,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 @@ -87,6 +89,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 @@ -99,6 +102,7 @@ async def _pipe_builder(r: Redis) -> Pipeline:
message=message,
current_progress=current,
total_progress=total,
log_level=log_level,
),
)

Expand Down Expand Up @@ -157,6 +161,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()
4 changes: 4 additions & 0 deletions src/ai/backend/common/events.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@
from redis.asyncio import ConnectionPool

from . import msgpack, redis_helper
from .defs import BackgroundTaskLogLevel
from .logging import BraceStyleAdapter
from .types import (
AgentId,
Expand Down Expand Up @@ -559,13 +560,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 +578,7 @@ def deserialize(cls, value: tuple):
value[1],
value[2],
value[3],
BackgroundTaskLogLevel(value[4]),
)


Expand Down

0 comments on commit c0d9a00

Please sign in to comment.