From 7feeb066da3972c4fe39ab039c9ac86deafbc992 Mon Sep 17 00:00:00 2001
From: Sanghun Lee <sanghun@lablup.com>
Date: Mon, 2 Sep 2024 07:59:22 +0000
Subject: [PATCH] chore: update GraphQL schema dump

---
 src/ai/backend/manager/api/schema.graphql           | 4 ++--
 src/ai/backend/manager/api/session.py               | 7 ++++++-
 src/ai/backend/manager/models/gql_models/session.py | 8 +++++++-
 3 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/src/ai/backend/manager/api/schema.graphql b/src/ai/backend/manager/api/schema.graphql
index b52d1e3686..8c657c425c 100644
--- a/src/ai/backend/manager/api/schema.graphql
+++ b/src/ai/backend/manager/api/schema.graphql
@@ -1406,7 +1406,7 @@ type Mutations {
   modify_container_registry(hostname: String!, props: ModifyContainerRegistryInput!): ModifyContainerRegistry
   delete_container_registry(hostname: String!): DeleteContainerRegistry
   modify_endpoint(endpoint_id: UUID!, props: ModifyEndpointInput!): ModifyEndpoint
-  check_and_transit_session_status(session_id: UUID!): CheckAndTransitStatus
+  check_and_transit_session_status(session_ids: [UUID]!): CheckAndTransitStatus
 }
 
 type ModifyAgent {
@@ -2134,5 +2134,5 @@ input ExtraMountInput {
 }
 
 type CheckAndTransitStatus {
-  session: ComputeSessionNode
+  sessions: [ComputeSessionNode]
 }
\ No newline at end of file
diff --git a/src/ai/backend/manager/api/session.py b/src/ai/backend/manager/api/session.py
index 5d7e875e25..56510c1fdb 100644
--- a/src/ai/backend/manager/api/session.py
+++ b/src/ai/backend/manager/api/session.py
@@ -1009,7 +1009,12 @@ async def check_and_transit_status(
                 )
     if accessible_session_ids:
         now = datetime.now(tzutc())
-        session_rows = await root_ctx.registry.transit_session_status(accessible_session_ids, now)
+        session_rows = await root_ctx.registry.session_lifecycle_mgr.transit_session_status(
+            accessible_session_ids, now
+        )
+        await root_ctx.registry.session_lifecycle_mgr.deregister_status_updatable_session([
+            row.id for row in session_rows
+        ])
         result = {row.id: row.status.name for row in session_rows}
     else:
         result = {}
diff --git a/src/ai/backend/manager/models/gql_models/session.py b/src/ai/backend/manager/models/gql_models/session.py
index 014f83e374..889c8f8f78 100644
--- a/src/ai/backend/manager/models/gql_models/session.py
+++ b/src/ai/backend/manager/models/gql_models/session.py
@@ -417,10 +417,16 @@ async def mutate(
                     accessible_session_ids.append(sid)
 
         if accessible_session_ids:
+            await graph_ctx.registry.session_lifecycle_mgr.deregister_status_updatable_session(
+                accessible_session_ids
+            )
             now = datetime.now(tzutc())
-            session_rows = await graph_ctx.registry.transit_session_status(
+            session_rows = await graph_ctx.registry.session_lifecycle_mgr.transit_session_status(
                 accessible_session_ids, now
             )
+            await graph_ctx.registry.session_lifecycle_mgr.deregister_status_updatable_session([
+                row.id for row in session_rows
+            ])
             result = [
                 ComputeSessionNode.parse_with_status_field_only(info, row) for row in session_rows
             ]