fix: Missing notification of bgtask cancellation or failure when shut…

…ting down the server (#2579) (#2960) Co-authored-by: Gyubong Lee <jopemachine@naver.com> Co-authored-by: Joongi Kim <joongi@lablup.com>
lablup · Oct 24, 2024 · f68d98f · f68d98f
1 parent d75f5c4
commit f68d98f
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/changes/2579.fix.md b/changes/2579.fix.md
@@ -0,0 +1 @@
+Fix missing notification of cancellation or failure of background tasks when shutting down the server
diff --git a/src/ai/backend/common/bgtask.py b/src/ai/backend/common/bgtask.py
@@ -167,7 +167,11 @@ async def push_bgtask_events(
                             else:
                                 await resp.send("{}", event="bgtask_done")
                             await resp.send("{}", event="server_close")
-                        case BgtaskCancelledEvent() | BgtaskFailedEvent():
+                        case BgtaskCancelledEvent():
+                            await resp.send(json.dumps(body), event="bgtask_cancelled")
+                            await resp.send("{}", event="server_close")
+                        case BgtaskFailedEvent():
+                            await resp.send(json.dumps(body), event="bgtask_failed")
                             await resp.send("{}", event="server_close")
             except:
                 log.exception("")