transit_session_status() returns SessionStatus

lablup · Jun 18, 2024 · 5be5355 · 5be5355
1 parent 3753e07
commit 5be5355
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/src/ai/backend/manager/registry.py b/src/ai/backend/manager/registry.py
@@ -3226,7 +3226,7 @@ async def transit_session_status(
         db_conn: SAConnection,
         session_id: SessionId,
         status_changed_at: datetime | None = None,
-    ) -> None:
+    ) -> SessionStatus:
         now = status_changed_at or datetime.now(tzutc())
 
         async def _get_and_transit(
@@ -3241,7 +3241,7 @@ async def _get_and_transit(
         )
 
         if not transited:
-            return
+            return session_row.status
         match session_row.status:
             case SessionStatus.RUNNING:
                 log.debug(
@@ -3269,6 +3269,7 @@ async def _get_and_transit(
                 )
             case _:
                 pass
+        return session_row.status
 
     async def set_status_updatable_session(self, session_id: SessionId) -> None:
         await redis_helper.execute(