diff --git a/changes/2165.feature.md b/changes/2165.feature.md new file mode 100644 index 0000000000..f756564995 --- /dev/null +++ b/changes/2165.feature.md @@ -0,0 +1 @@ +Add relay-aware `VirtualFolderNode` GQL Query diff --git a/src/ai/backend/manager/api/schema.graphql b/src/ai/backend/manager/api/schema.graphql index 9b28813d78..d724ef3813 100644 --- a/src/ai/backend/manager/api/schema.graphql +++ b/src/ai/backend/manager/api/schema.graphql @@ -76,6 +76,12 @@ type Queries { storage_volume(id: String): StorageVolume storage_volume_list(limit: Int!, offset: Int!, filter: String, order: String): StorageVolumeList vfolder(id: String): VirtualFolder + + """Added in 24.03.4.""" + vfolder_node(id: String!): VirtualFolderNode + + """Added in 24.03.4.""" + vfolder_nodes(filter: String, order: String, offset: Int, before: String, after: String, first: Int, last: Int): VirtualFolderConnection vfolder_list(limit: Int!, offset: Int!, filter: String, order: String, domain_name: String, group_id: UUID, access_key: String): VirtualFolderList vfolder_permission_list(limit: Int!, offset: Int!, filter: String, order: String): VirtualFolderPermissionList vfolder_own_list(limit: Int!, offset: Int!, filter: String, order: String, domain_name: String, access_key: String): VirtualFolderList @@ -691,6 +697,52 @@ type StorageVolumeList implements PaginatedList { total_count: Int! } +type VirtualFolderNode implements Node { + """The ID of the object""" + id: ID! + host: String + quota_scope_id: String + name: String + user: UUID + user_email: String + group: UUID + group_name: String + creator: String + unmanaged_path: String + usage_mode: String + permission: String + ownership_type: String + max_files: Int + max_size: BigInt + created_at: DateTime + modified_at: DateTime + last_used: DateTime + num_files: Int + cur_size: BigInt + cloneable: Boolean + status: String +} + +type VirtualFolderConnection { + """Pagination data for this connection.""" + pageInfo: PageInfo! + + """Contains the nodes in this connection.""" + edges: [VirtualFolderEdge]! + + """Total count of the GQL nodes of the query.""" + count: Int +} + +"""A Relay edge containing a `VirtualFolder` and its cursor.""" +type VirtualFolderEdge { + """The item at the end of the edge""" + node: VirtualFolderNode + + """A cursor for use in pagination""" + cursor: String! +} + type VirtualFolderList implements PaginatedList { items: [VirtualFolder]! total_count: Int! diff --git a/src/ai/backend/manager/models/endpoint.py b/src/ai/backend/manager/models/endpoint.py index f9eec62db8..9d0162ae2d 100644 --- a/src/ai/backend/manager/models/endpoint.py +++ b/src/ai/backend/manager/models/endpoint.py @@ -19,7 +19,7 @@ from ai.backend.common.types import ClusterMode, ResourceSlot from ai.backend.manager.defs import SERVICE_MAX_RETRIES -from ..api.exceptions import EndpointNotFound, EndpointTokenNotFound +from ..api.exceptions import EndpointNotFound, EndpointTokenNotFound, ObjectNotFound from .base import ( GUID, Base, @@ -38,6 +38,7 @@ from .image import ImageNode, ImageRefType, ImageRow from .routing import RouteStatus, Routing from .user import UserRole +from .vfolder import VFolderRow, VirtualFolderNode if TYPE_CHECKING: from .gql import GraphQueryContext @@ -155,6 +156,7 @@ class EndpointRow(Base): routings = relationship("RoutingRow", back_populates="endpoint_row") tokens = relationship("EndpointTokenRow", back_populates="endpoint_row") image_row = relationship("ImageRow", back_populates="endpoints") + model_row = relationship("VFolderRow", back_populates="endpoints") created_user_row = relationship( "UserRow", back_populates="created_endpoints", foreign_keys="EndpointRow.created_user" ) @@ -437,6 +439,7 @@ class Meta: resource_slots = graphene.JSONString() url = graphene.String() model = graphene.UUID() + model_vfolder = VirtualFolderNode() model_mount_destiation = graphene.String( deprecation_reason="Deprecated since 24.03.4; use `model_mount_destination` instead" ) @@ -663,6 +666,18 @@ async def resolve_status(self, info: graphene.ResolveInfo) -> str: return "DEGRADED" return "PROVISIONING" + async def resolve_model_vfolder(self, info: graphene.ResolveInfo) -> VirtualFolderNode: + if not self.model: + raise ObjectNotFound(object_name="VFolder") + + ctx: GraphQueryContext = info.context + + async with ctx.db.begin_readonly_session() as sess: + vfolder_row = await VFolderRow.get( + sess, uuid.UUID(self.model), load_user=True, load_group=True + ) + return VirtualFolderNode.from_row(info, vfolder_row) + async def resolve_errors(self, info: graphene.ResolveInfo) -> Any: error_routes = [r for r in self.routings if r.status == RouteStatus.FAILED_TO_START.name] errors = [] diff --git a/src/ai/backend/manager/models/gql.py b/src/ai/backend/manager/models/gql.py index 0e1c4442f1..2d6a904608 100644 --- a/src/ai/backend/manager/models/gql.py +++ b/src/ai/backend/manager/models/gql.py @@ -141,7 +141,9 @@ SetQuotaScope, UnsetQuotaScope, VirtualFolder, + VirtualFolderConnection, VirtualFolderList, + VirtualFolderNode, VirtualFolderPermission, VirtualFolderPermissionList, ensure_quota_scope_accessible_by_user, @@ -512,6 +514,13 @@ class Queries(graphene.ObjectType): id=graphene.String(), ) + vfolder_node = graphene.Field( + VirtualFolderNode, id=graphene.String(required=True), description="Added in 24.03.4." + ) + vfolder_nodes = PaginatedConnectionField( + VirtualFolderConnection, description="Added in 24.03.4." + ) + vfolder_list = graphene.Field( # legacy non-paginated list VirtualFolderList, limit=graphene.Int(required=True), @@ -846,6 +855,7 @@ async def resolve_domains( ) -> Sequence[Domain]: return await Domain.load_all(info.context, is_active=is_active) + @staticmethod async def resolve_group_node( root: Any, info: graphene.ResolveInfo, @@ -853,6 +863,7 @@ async def resolve_group_node( ): return await GroupNode.get_node(info, id) + @staticmethod async def resolve_group_nodes( root: Any, info: graphene.ResolveInfo, @@ -876,6 +887,38 @@ async def resolve_group_nodes( last, ) + @staticmethod + async def resolve_vfolder_node( + root: Any, + info: graphene.ResolveInfo, + id: str, + ): + return await VirtualFolderNode.get_node(info, id) + + @staticmethod + async def resolve_vfolder_nodes( + root: Any, + info: graphene.ResolveInfo, + *, + filter: str | None = None, + order: str | None = None, + offset: int | None = None, + after: str | None = None, + first: int | None = None, + before: str | None = None, + last: int | None = None, + ) -> ConnectionResolverResult: + return await VirtualFolderNode.get_connection( + info, + filter, + order, + offset, + after, + first, + before, + last, + ) + @staticmethod async def resolve_group( root: Any, diff --git a/src/ai/backend/manager/models/group.py b/src/ai/backend/manager/models/group.py index 39fe85145d..6a4709ea8b 100644 --- a/src/ai/backend/manager/models/group.py +++ b/src/ai/backend/manager/models/group.py @@ -197,6 +197,7 @@ class GroupRow(Base): users = relationship("AssocGroupUserRow", back_populates="group") resource_policy_row = relationship("ProjectResourcePolicyRow", back_populates="projects") kernels = relationship("KernelRow", back_populates="group_row") + vfolder_row = relationship("VFolderRow", back_populates="group_row") def _build_group_query(cond: sa.sql.BinaryExpression, domain_name: str) -> sa.sql.Select: diff --git a/src/ai/backend/manager/models/user.py b/src/ai/backend/manager/models/user.py index 3685b6e6cb..a3f847b51a 100644 --- a/src/ai/backend/manager/models/user.py +++ b/src/ai/backend/manager/models/user.py @@ -186,6 +186,8 @@ class UserRow(Base): main_keypair = relationship("KeyPairRow", foreign_keys=users.c.main_access_key) + vfolder_row = relationship("VFolderRow", back_populates="user_row") + class UserGroup(graphene.ObjectType): id = graphene.UUID() diff --git a/src/ai/backend/manager/models/vfolder.py b/src/ai/backend/manager/models/vfolder.py index c593853e6b..72c2c1cf40 100644 --- a/src/ai/backend/manager/models/vfolder.py +++ b/src/ai/backend/manager/models/vfolder.py @@ -23,7 +23,7 @@ from sqlalchemy.engine.row import Row from sqlalchemy.ext.asyncio import AsyncConnection as SAConnection from sqlalchemy.ext.asyncio import AsyncSession as SASession -from sqlalchemy.orm import load_only, selectinload +from sqlalchemy.orm import load_only, relationship, selectinload from ai.backend.common.bgtask import ProgressReporter from ai.backend.common.config import model_definition_iv @@ -42,6 +42,7 @@ from ..api.exceptions import ( InvalidAPIParameters, + ObjectNotFound, VFolderNotFound, VFolderOperationFailed, VFolderPermissionError, @@ -417,6 +418,29 @@ class VFolderCloneInfo(NamedTuple): class VFolderRow(Base): __table__ = vfolders + endpoints = relationship("EndpointRow", back_populates="model_row") + user_row = relationship("UserRow", back_populates="vfolder_row") + group_row = relationship("GroupRow", back_populates="vfolder_row") + + @classmethod + async def get( + cls, + session: SASession, + id: uuid.UUID, + load_user=False, + load_group=False, + ) -> "VFolderRow": + query = sa.select(VFolderRow).where(VFolderRow.id == id) + if load_user: + query = query.options(selectinload(VFolderRow.user_row)) + if load_group: + query = query.options(selectinload(VFolderRow.group_row)) + + result = await session.scalar(query) + if not result: + raise ObjectNotFound(object_name="VFolder") + return result + def __contains__(self, key): return key in self.__dir__() @@ -426,6 +450,10 @@ def __getitem__(self, item): except AttributeError: raise KeyError(item) + @property + def vfid(self) -> VFolderID: + return VFolderID(self.quota_scope_id, self.id) + def verify_vfolder_name(folder: str) -> bool: if folder in RESERVED_VFOLDERS: @@ -1695,6 +1723,197 @@ class Meta: items = graphene.List(VirtualFolder, required=True) +class VirtualFolderNode(graphene.ObjectType): + class Meta: + interfaces = (AsyncNode,) + + host = graphene.String() + quota_scope_id = graphene.String() + name = graphene.String() + user = graphene.UUID() # User.id (current owner, null in project vfolders) + user_email = graphene.String() # User.email (current owner, null in project vfolders) + group = graphene.UUID() # Group.id (current owner, null in user vfolders) + group_name = graphene.String() # Group.name (current owenr, null in user vfolders) + creator = graphene.String() # User.email (always set) + unmanaged_path = graphene.String() + usage_mode = graphene.String() + permission = graphene.String() + ownership_type = graphene.String() + max_files = graphene.Int() + max_size = BigInt() # in MiB + created_at = GQLDateTime() + modified_at = GQLDateTime() + last_used = GQLDateTime() + + num_files = graphene.Int() + cur_size = BigInt() + # num_attached = graphene.Int() + cloneable = graphene.Boolean() + status = graphene.String() + + _queryfilter_fieldspec: Mapping[str, FieldSpecItem] = { + "id": ("vfolders_id", uuid.UUID), + "host": ("vfolders_host", None), + "quota_scope_id": ("vfolders_quota_scope_id", None), + "name": ("vfolders_name", None), + "group": ("vfolders_group", uuid.UUID), + "group_name": ("groups_name", None), + "user": ("vfolders_user", uuid.UUID), + "user_email": ("users_email", None), + "creator": ("vfolders_creator", None), + "unmanaged_path": ("vfolders_unmanaged_path", None), + "usage_mode": ( + "vfolders_usage_mode", + enum_field_getter(VFolderUsageMode), + ), + "permission": ( + "vfolders_permission", + enum_field_getter(VFolderPermission), + ), + "ownership_type": ( + "vfolders_ownership_type", + enum_field_getter(VFolderOwnershipType), + ), + "max_files": ("vfolders_max_files", None), + "max_size": ("vfolders_max_size", None), + "created_at": ("vfolders_created_at", dtparse), + "last_used": ("vfolders_last_used", dtparse), + "cloneable": ("vfolders_cloneable", None), + "status": ( + "vfolders_status", + enum_field_getter(VFolderOperationStatus), + ), + } + + _queryorder_colmap: Mapping[str, OrderSpecItem] = { + "id": ("vfolders_id", None), + "host": ("vfolders_host", None), + "quota_scope_id": ("vfolders_quota_scope_id", None), + "name": ("vfolders_name", None), + "group": ("vfolders_group", None), + "group_name": ("groups_name", None), + "user": ("vfolders_user", None), + "user_email": ("users_email", None), + "creator": ("vfolders_creator", None), + "usage_mode": ("vfolders_usage_mode", None), + "permission": ("vfolders_permission", None), + "ownership_type": ("vfolders_ownership_type", None), + "max_files": ("vfolders_max_files", None), + "max_size": ("vfolders_max_size", None), + "created_at": ("vfolders_created_at", None), + "last_used": ("vfolders_last_used", None), + "cloneable": ("vfolders_cloneable", None), + "status": ("vfolders_status", None), + "cur_size": ("vfolders_cur_size", None), + } + + def resolve_created_at( + self, + info: graphene.ResolveInfo, + ) -> datetime: + try: + return dtparse(self.created_at) + except ParserError: + return self.created_at + + @classmethod + def from_row(cls, info: graphene.ResolveInfo, row: VFolderRow) -> VirtualFolderNode: + return cls( + id=row.id, + host=row.host, + quota_scope_id=row.quota_scope_id, + name=row.name, + user=row.user_row, + user_email=row.user_row.email if row.user_row else None, + group=row.group_row, + group_name=row.group_row.name if row.group_row else None, + creator=row.creator, + unmanaged_path=row.unmanaged_path, + usage_mode=row.usage_mode, + permission=row.permission, + ownership_type=row.ownership_type, + max_files=row.max_files, + max_size=row.max_size, # in B + created_at=row.created_at, + last_used=row.last_used, + cloneable=row.cloneable, + status=row.status, + cur_size=row.cur_size, + ) + + @classmethod + async def get_node(cls, info: graphene.ResolveInfo, id: str) -> VirtualFolderNode: + graph_ctx: GraphQueryContext = info.context + + _, vfolder_row_id = AsyncNode.resolve_global_id(info, id) + async with graph_ctx.db.begin_readonly_session() as db_session: + vfolder_row = await VFolderRow.get( + db_session, uuid.UUID(vfolder_row_id), load_user=True, load_group=True + ) + if vfolder_row.status in DEAD_VFOLDER_STATUSES: + raise ValueError( + f"The vfolder is deleted. (id: {vfolder_row_id}, status: {vfolder_row.status})" + ) + return cls.from_row(info, vfolder_row) + + @classmethod + async def get_connection( + cls, + info: graphene.ResolveInfo, + filter_expr: str | None = None, + order_expr: str | None = None, + offset: int | None = None, + after: str | None = None, + first: int | None = None, + before: str | None = None, + last: int | None = None, + ) -> ConnectionResolverResult: + graph_ctx: GraphQueryContext = info.context + _filter_arg = ( + FilterExprArg(filter_expr, QueryFilterParser(cls._queryfilter_fieldspec)) + if filter_expr is not None + else None + ) + _order_expr = ( + OrderExprArg(order_expr, QueryOrderParser(cls._queryorder_colmap)) + if order_expr is not None + else None + ) + ( + query, + conditions, + cursor, + pagination_order, + page_size, + ) = generate_sql_info_for_gql_connection( + info, + VFolderRow, + VFolderRow.id, + _filter_arg, + _order_expr, + offset, + after=after, + first=first, + before=before, + last=last, + ) + cnt_query = sa.select(sa.func.count()).select_from(VFolderRow) + for cond in conditions: + cnt_query = cnt_query.where(cond) + + async with graph_ctx.db.begin_readonly_session() as db_session: + vfolder_rows = (await db_session.scalars(query)).all() + result = [(cls.from_row(info, vf)) for vf in vfolder_rows] + + total_cnt = await db_session.scalar(cnt_query) + return ConnectionResolverResult(result, cursor, pagination_order, page_size, total_cnt) + + +class VirtualFolderConnection(Connection): + class Meta: + node = VirtualFolderNode + + class VirtualFolderPermission(graphene.ObjectType): class Meta: interfaces = (Item,) @@ -2220,9 +2439,8 @@ async def get_node(cls, info: graphene.ResolveInfo, id: str) -> ModelCard: graph_ctx: GraphQueryContext = info.context _, vfolder_row_id = AsyncNode.resolve_global_id(info, id) - query = sa.select(VFolderRow).where(VFolderRow.id == vfolder_row_id) async with graph_ctx.db.begin_readonly_session() as db_session: - vfolder_row = (await db_session.scalars(query)).first() + vfolder_row = await VFolderRow.get(db_session, uuid.UUID(vfolder_row_id)) if vfolder_row.usage_mode != VFolderUsageMode.MODEL: raise ValueError( f"The vfolder is not model. expect: {VFolderUsageMode.MODEL.value}, got:"