fix invalid event subscription pattern

lablup · Nov 30, 2023 · eaaefdf · eaaefdf
1 parent 0e810ba
commit eaaefdf
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 6 deletions.
diff --git a/src/ai/backend/manager/api/service.py b/src/ai/backend/manager/api/service.py
@@ -28,6 +28,7 @@
     KernelLifecycleEventReason,
     ModelServiceStatusEvent,
     SessionCancelledEvent,
+    SessionPreparingEvent,
     SessionStartedEvent,
     SessionTerminatedEvent,
 )
@@ -549,10 +550,16 @@ async def _handle_event(
                     forced=True,
                 )
 
-        session_event_matcher = lambda args: args[0] == result["sessionId"]
-        model_service_event_matcher = lambda args: args[1] == result["sessionId"]
+        session_event_matcher = lambda args: args[0] == str(result["sessionId"])
+        model_service_event_matcher = lambda args: args[1] == str(result["sessionId"])
 
         handlers: list[EventHandler] = [
+            root_ctx.event_dispatcher.subscribe(
+                SessionPreparingEvent,
+                None,
+                _handle_event,
+                args_matcher=session_event_matcher,
+            ),
             root_ctx.event_dispatcher.subscribe(
                 SessionStartedEvent,
                 None,
@@ -565,13 +572,13 @@ async def _handle_event(
                 _handle_event,
                 args_matcher=session_event_matcher,
             ),
-            root_ctx.event_dispatcher.consume(
+            root_ctx.event_dispatcher.subscribe(
                 SessionTerminatedEvent,
                 None,
                 _handle_event,
                 args_matcher=session_event_matcher,
             ),
-            root_ctx.event_dispatcher.consume(
+            root_ctx.event_dispatcher.subscribe(
                 ModelServiceStatusEvent,
                 None,
                 _handle_event,
@@ -586,7 +593,7 @@ async def _handle_event(
                 root_ctx.event_dispatcher.unsubscribe(handler)
 
     task_id = await background_task_manager.start(_task)
-    return web.json_response({"task_id": task_id})
+    return web.json_response({"task_id": str(task_id)})
 
 
 @auth_required

diff --git a/src/ai/backend/manager/registry.py b/src/ai/backend/manager/registry.py
@@ -3407,7 +3407,9 @@ async def handle_model_service_status_update(
                 kernel_loading_strategy=KernelLoadingStrategy.MAIN_KERNEL_ONLY,
             )
             route = await RoutingRow.get_by_session(db_sess, session.id, load_endpoint=True)
-    except SessionNotFound | NoResultFound:
+    except SessionNotFound:
+        return
+    except NoResultFound:
         return
 
     async def _update():