Update vfolder_mounts type from VFolderMount to UUID

lablup · Jan 24, 2025 · 59c8381 · 59c8381
1 parent 3b1e12e
commit 59c8381
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/src/ai/backend/manager/models/gql_models/session.py b/src/ai/backend/manager/models/gql_models/session.py
@@ -22,7 +22,7 @@
 from sqlalchemy.orm import selectinload
 
 from ai.backend.common import validators as tx
-from ai.backend.common.types import ClusterMode, SessionId, SessionResult, VFolderMount
+from ai.backend.common.types import ClusterMode, SessionId, SessionResult
 from ai.backend.manager.idle import ReportInfo
 
 from ..base import (
@@ -291,8 +291,7 @@ async def resolve_vfolder_nodes(
         info: graphene.ResolveInfo,
     ) -> ConnectionResolverResult[VirtualFolderNode]:
         ctx: GraphQueryContext = info.context
-        vfolder_mounts = cast(list[VFolderMount], self.vfolder_mounts)
-        _folder_ids = [vf_mount.vfid.folder_id for vf_mount in vfolder_mounts]
+        _folder_ids = cast(list[uuid.UUID], self.vfolder_mounts)
         loader = ctx.dataloader_manager.get_loader_by_func(ctx, VirtualFolderNode.batch_load_by_id)
         result = cast(list[list[VirtualFolderNode]], await loader.load_many(_folder_ids))