From 5be5355a44257d9e4d4a196a09bf017f149aa5b7 Mon Sep 17 00:00:00 2001
From: Sanghun Lee <sanghun@lablup.com>
Date: Tue, 18 Jun 2024 23:29:56 +0900
Subject: [PATCH] transit_session_status() returns SessionStatus

---
 src/ai/backend/manager/registry.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/src/ai/backend/manager/registry.py b/src/ai/backend/manager/registry.py
index e9277482b6b..45f21e97356 100644
--- 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(