chore: update GraphQL schema dump

lablup · Sep 10, 2024 · 0ab732d · 0ab732d
1 parent 7f7367b
commit 0ab732d
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 4 deletions.
diff --git 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]
 }
diff --git 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
@@ -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
             ]