From 8d1297cd6c32f7b1144c6b952540cc3431434920 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Wed, 13 Nov 2024 03:18:55 +0000 Subject: [PATCH 01/41] feat: Implement management API for controlling Harbor per-project Quota --- src/ai/backend/manager/api/schema.graphql | 9 ++ src/ai/backend/manager/models/gql.py | 3 + .../models/gql_models/container_registry.py | 136 +++++++++++++++++- 3 files changed, 147 insertions(+), 1 deletion(-) diff --git a/src/ai/backend/manager/api/schema.graphql b/src/ai/backend/manager/api/schema.graphql index 6caa04e4ea..c83e30f7af 100644 --- a/src/ai/backend/manager/api/schema.graphql +++ b/src/ai/backend/manager/api/schema.graphql @@ -1844,6 +1844,9 @@ type Mutations { """Added in 24.12.0""" disassociate_container_registry_with_group(group_id: String!, registry_id: String!): DisassociateContainerRegistryWithGroup + + """Added in 24.12.0""" + update_container_registry_quota(quota: Int!, scope_id: ScopeField!): UpdateQuota create_container_registry(hostname: String!, props: CreateContainerRegistryInput!): CreateContainerRegistry modify_container_registry(hostname: String!, props: ModifyContainerRegistryInput!): ModifyContainerRegistry delete_container_registry(hostname: String!): DeleteContainerRegistry @@ -2570,6 +2573,12 @@ type DisassociateContainerRegistryWithGroup { msg: String } +"""Added in 24.12.0.""" +type UpdateQuota { + ok: Boolean + msg: String +} + type CreateContainerRegistry { container_registry: ContainerRegistry } diff --git a/src/ai/backend/manager/models/gql.py b/src/ai/backend/manager/models/gql.py index 078c169577..912b94c02a 100644 --- a/src/ai/backend/manager/models/gql.py +++ b/src/ai/backend/manager/models/gql.py @@ -75,6 +75,7 @@ from .gql_models.container_registry import ( AssociateContainerRegistryWithGroup, DisassociateContainerRegistryWithGroup, + UpdateQuota, ) from .gql_models.domain import ( CreateDomainNode, @@ -349,6 +350,8 @@ class Mutations(graphene.ObjectType): description="Added in 24.12.0" ) + update_container_registry_quota = UpdateQuota.Field(description="Added in 24.12.0") + # Legacy mutations create_container_registry = CreateContainerRegistry.Field() modify_container_registry = ModifyContainerRegistry.Field() diff --git a/src/ai/backend/manager/models/gql_models/container_registry.py b/src/ai/backend/manager/models/gql_models/container_registry.py index ae44bef0e2..11e4c62bf6 100644 --- a/src/ai/backend/manager/models/gql_models/container_registry.py +++ b/src/ai/backend/manager/models/gql_models/container_registry.py @@ -1,17 +1,25 @@ from __future__ import annotations import logging -from typing import Self +from typing import Any, Self +import aiohttp +import aiohttp.client_exceptions import graphene import sqlalchemy as sa +import yarl +from sqlalchemy.orm import load_only from ai.backend.logging import BraceStyleAdapter +from ai.backend.manager.models.container_registry import ContainerRegistryRow, ContainerRegistryType +from ai.backend.manager.models.gql_models.fields import ScopeField +from ai.backend.manager.models.group import GroupRow from ..association_container_registries_groups import ( AssociationContainerRegistriesGroupsRow, ) from ..base import simple_db_mutate +from ..rbac import ProjectScope, ScopeType from ..user import UserRole log = BraceStyleAdapter(logging.getLogger(__spec__.name)) # type: ignore @@ -70,3 +78,129 @@ async def mutate( .where(AssociationContainerRegistriesGroupsRow.group_id == group_id) ) return await simple_db_mutate(cls, info.context, delete_query) + + +class UpdateQuota(graphene.Mutation): + """Added in 24.12.0.""" + + allowed_roles = ( + UserRole.SUPERADMIN, + UserRole.ADMIN, + UserRole.USER, + ) + + class Arguments: + scope_id = ScopeField(required=True) + quota = graphene.Int(required=True) + + ok = graphene.Boolean() + msg = graphene.String() + + @classmethod + async def mutate( + cls, + root, + info: graphene.ResolveInfo, + scope_id: ScopeType, + quota: int, + ) -> Self: + graph_ctx = info.context + + # TODO: Support other scope types + assert isinstance(scope_id, ProjectScope) + project_id = scope_id.project_id + + # user = graph_ctx.user + # client_ctx = ClientContext( + # graph_ctx.db, user["domain_name"], user["uuid"], user["role"] + # ) + + async with graph_ctx.db.begin_session() as db_sess: + group_query = ( + sa.select(GroupRow) + .where(GroupRow.id == project_id) + .options(load_only(GroupRow.container_registry)) + ) + result = await db_sess.execute(group_query) + + group = result.scalar_one_or_none() + + if ( + group is None + or group.container_registry is None + or "registry" not in group.container_registry + or "project" not in group.container_registry + ): + raise ValueError("Container registry info does not exist in the group.") + + registry_name, project = ( + group.container_registry["registry"], + group.container_registry["project"], + ) + + cr_query = sa.select(ContainerRegistryRow).where( + (ContainerRegistryRow.registry_name == registry_name) + & (ContainerRegistryRow.project == project) + ) + + result = await db_sess.execute(cr_query) + registry = result.fetchone()[0] + + if registry.type != ContainerRegistryType.HARBOR2: + raise ValueError("Only HarborV2 registry is supported for now.") + + if not registry.is_global: + get_assoc_query = sa.select( + sa.exists() + .where(AssociationContainerRegistriesGroupsRow.registry_id == registry.id) + .where(AssociationContainerRegistriesGroupsRow.group_id == project_id) + ) + assoc_exist = (await db_sess.execute(get_assoc_query)).scalar() + + if not assoc_exist: + return UpdateQuota( + ok=False, msg="The group is not associated with the container registry." + ) + + ssl_verify = registry.ssl_verify + connector = aiohttp.TCPConnector(ssl=ssl_verify) + + url = yarl.URL(registry.url) + async with aiohttp.ClientSession(connector=connector) as sess: + rqst_args: dict[str, Any] = {} + rqst_args["auth"] = aiohttp.BasicAuth( + registry.username, + registry.password, + ) + + get_project_id = url / "api" / "v2.0" / "projects" / project + + async with sess.get(get_project_id, allow_redirects=False, **rqst_args) as resp: + res = await resp.json() + harbor_project_id = res["project_id"] + + get_quota_id = (url / "api" / "v2.0" / "quotas").with_query({ + "reference": "project", + "reference_id": harbor_project_id, + }) + + async with sess.get(get_quota_id, allow_redirects=False, **rqst_args) as resp: + res = await resp.json() + # TODO: Raise error when quota is not found or multiple quotas are found. + quota_id = res[0]["id"] + + put_quota_url = url / "api" / "v2.0" / "quotas" / str(quota_id) + update_payload = {"hard": {"storage": quota}} + + async with sess.put( + put_quota_url, json=update_payload, allow_redirects=False, **rqst_args + ) as resp: + if resp.status == 200: + return UpdateQuota(ok=True, msg="Quota updated successfully.") + else: + log.error(f"Failed to update quota: {await resp.json()}") + return UpdateQuota( + ok=False, msg=f"Failed to update quota. Status code: {resp.status}" + ) + + return UpdateQuota(ok=False, msg="Unknown error!") From 9128ef781793b74a48f3e0b69417546fd625bdee Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Wed, 13 Nov 2024 03:21:17 +0000 Subject: [PATCH 02/41] chore: Add news fragment --- changes/3090.feature.md | 1 + 1 file changed, 1 insertion(+) create mode 100644 changes/3090.feature.md diff --git a/changes/3090.feature.md b/changes/3090.feature.md new file mode 100644 index 0000000000..76acfee344 --- /dev/null +++ b/changes/3090.feature.md @@ -0,0 +1 @@ +Implement management API for controlling Harbor per-project Quota. From c60bd90abd94044d27c75f23c206bdbee2803d7d Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Wed, 13 Nov 2024 03:28:39 +0000 Subject: [PATCH 03/41] fix: Disable user quota mutation --- .../models/gql_models/container_registry.py | 14 -------------- 1 file changed, 14 deletions(-) diff --git a/src/ai/backend/manager/models/gql_models/container_registry.py b/src/ai/backend/manager/models/gql_models/container_registry.py index 11e4c62bf6..2f5283ee6a 100644 --- a/src/ai/backend/manager/models/gql_models/container_registry.py +++ b/src/ai/backend/manager/models/gql_models/container_registry.py @@ -86,7 +86,6 @@ class UpdateQuota(graphene.Mutation): allowed_roles = ( UserRole.SUPERADMIN, UserRole.ADMIN, - UserRole.USER, ) class Arguments: @@ -149,19 +148,6 @@ async def mutate( if registry.type != ContainerRegistryType.HARBOR2: raise ValueError("Only HarborV2 registry is supported for now.") - if not registry.is_global: - get_assoc_query = sa.select( - sa.exists() - .where(AssociationContainerRegistriesGroupsRow.registry_id == registry.id) - .where(AssociationContainerRegistriesGroupsRow.group_id == project_id) - ) - assoc_exist = (await db_sess.execute(get_assoc_query)).scalar() - - if not assoc_exist: - return UpdateQuota( - ok=False, msg="The group is not associated with the container registry." - ) - ssl_verify = registry.ssl_verify connector = aiohttp.TCPConnector(ssl=ssl_verify) From af7596d8264415b6e10ae57a12efeba7a90da796 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Wed, 13 Nov 2024 03:34:33 +0000 Subject: [PATCH 04/41] fix: Rename variables --- .../models/gql_models/container_registry.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/src/ai/backend/manager/models/gql_models/container_registry.py b/src/ai/backend/manager/models/gql_models/container_registry.py index 2f5283ee6a..b221e7e1b8 100644 --- a/src/ai/backend/manager/models/gql_models/container_registry.py +++ b/src/ai/backend/manager/models/gql_models/container_registry.py @@ -151,7 +151,7 @@ async def mutate( ssl_verify = registry.ssl_verify connector = aiohttp.TCPConnector(ssl=ssl_verify) - url = yarl.URL(registry.url) + api_url = yarl.URL(registry.url) / "api" / "v2.0" async with aiohttp.ClientSession(connector=connector) as sess: rqst_args: dict[str, Any] = {} rqst_args["auth"] = aiohttp.BasicAuth( @@ -159,27 +159,27 @@ async def mutate( registry.password, ) - get_project_id = url / "api" / "v2.0" / "projects" / project + get_project_id_api = api_url / "projects" / project - async with sess.get(get_project_id, allow_redirects=False, **rqst_args) as resp: + async with sess.get(get_project_id_api, allow_redirects=False, **rqst_args) as resp: res = await resp.json() harbor_project_id = res["project_id"] - get_quota_id = (url / "api" / "v2.0" / "quotas").with_query({ + get_quota_id_api = (api_url / "quotas").with_query({ "reference": "project", "reference_id": harbor_project_id, }) - async with sess.get(get_quota_id, allow_redirects=False, **rqst_args) as resp: + async with sess.get(get_quota_id_api, allow_redirects=False, **rqst_args) as resp: res = await resp.json() # TODO: Raise error when quota is not found or multiple quotas are found. quota_id = res[0]["id"] - put_quota_url = url / "api" / "v2.0" / "quotas" / str(quota_id) + put_quota_api = api_url / "quotas" / str(quota_id) update_payload = {"hard": {"storage": quota}} async with sess.put( - put_quota_url, json=update_payload, allow_redirects=False, **rqst_args + put_quota_api, json=update_payload, allow_redirects=False, **rqst_args ) as resp: if resp.status == 200: return UpdateQuota(ok=True, msg="Quota updated successfully.") From ed1a4b145617e77113e11837fb620327c5daa7be Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Wed, 13 Nov 2024 04:06:27 +0000 Subject: [PATCH 05/41] feat: Add `registry_quota` to GroupNode --- src/ai/backend/manager/api/schema.graphql | 3 + .../manager/models/gql_models/group.py | 68 +++++++++++++++++++ 2 files changed, 71 insertions(+) diff --git a/src/ai/backend/manager/api/schema.graphql b/src/ai/backend/manager/api/schema.graphql index c83e30f7af..bfebd14ba9 100644 --- a/src/ai/backend/manager/api/schema.graphql +++ b/src/ai/backend/manager/api/schema.graphql @@ -698,6 +698,9 @@ type GroupNode implements Node { """Added in 24.03.7.""" container_registry: JSONString scaling_groups: [String] + + """Added in 24.12.0.""" + registry_quota: Int user_nodes(filter: String, order: String, offset: Int, before: String, after: String, first: Int, last: Int): UserConnection } diff --git a/src/ai/backend/manager/models/gql_models/group.py b/src/ai/backend/manager/models/gql_models/group.py index 4bd98da161..979cf71a9d 100644 --- a/src/ai/backend/manager/models/gql_models/group.py +++ b/src/ai/backend/manager/models/gql_models/group.py @@ -3,15 +3,20 @@ from collections.abc import Mapping from typing import ( TYPE_CHECKING, + Any, Self, Sequence, ) +import aiohttp import graphene import sqlalchemy as sa +import yarl from dateutil.parser import parse as dtparse from graphene.types.datetime import DateTime as GQLDateTime +from ai.backend.manager.models.container_registry import ContainerRegistryRow, ContainerRegistryType + from ..base import ( FilterExprArg, OrderExprArg, @@ -112,6 +117,8 @@ class Meta: lambda: graphene.String, ) + registry_quota = graphene.Int(description="Added in 24.12.0.") + user_nodes = PaginatedConnectionField( UserConnection, ) @@ -204,6 +211,67 @@ async def resolve_user_nodes( total_cnt = await db_session.scalar(cnt_query) return ConnectionResolverResult(result, cursor, pagination_order, page_size, total_cnt) + async def resolve_registry_quota(self, info: graphene.ResolveInfo) -> int: + graph_ctx = info.context + + # user = graph_ctx.user + # client_ctx = ClientContext( + # graph_ctx.db, user["domain_name"], user["uuid"], user["role"] + # ) + + async with graph_ctx.db.begin_session() as db_sess: + if ( + self.container_registry is None + or "registry" not in self.container_registry + or "project" not in self.container_registry + ): + raise ValueError("Container registry info does not exist in the group.") + + registry_name, project = ( + self.container_registry["registry"], + self.container_registry["project"], + ) + + cr_query = sa.select(ContainerRegistryRow).where( + (ContainerRegistryRow.registry_name == registry_name) + & (ContainerRegistryRow.project == project) + ) + + result = await db_sess.execute(cr_query) + registry = result.fetchone()[0] + + if registry.type != ContainerRegistryType.HARBOR2: + raise ValueError("Only HarborV2 registry is supported for now.") + + ssl_verify = registry.ssl_verify + connector = aiohttp.TCPConnector(ssl=ssl_verify) + + api_url = yarl.URL(registry.url) / "api" / "v2.0" + async with aiohttp.ClientSession(connector=connector) as sess: + rqst_args: dict[str, Any] = {} + rqst_args["auth"] = aiohttp.BasicAuth( + registry.username, + registry.password, + ) + + get_project_id_api = api_url / "projects" / project + + async with sess.get(get_project_id_api, allow_redirects=False, **rqst_args) as resp: + res = await resp.json() + harbor_project_id = res["project_id"] + + get_quota_id_api = (api_url / "quotas").with_query({ + "reference": "project", + "reference_id": harbor_project_id, + }) + + async with sess.get(get_quota_id_api, allow_redirects=False, **rqst_args) as resp: + res = await resp.json() + # TODO: Raise error when quota is not found or multiple quotas are found. + quota = res[0]["hard"]["storage"] + + return quota + @classmethod async def get_node(cls, info: graphene.ResolveInfo, id) -> Self: graph_ctx: GraphQueryContext = info.context From a1225c2f41570fc78f936267640cf533de419b5c Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Wed, 13 Nov 2024 08:34:25 +0000 Subject: [PATCH 06/41] fix: Update `UpdateQuota` mutation --- .../models/gql_models/container_registry.py | 72 +++++++++++-------- 1 file changed, 43 insertions(+), 29 deletions(-) diff --git a/src/ai/backend/manager/models/gql_models/container_registry.py b/src/ai/backend/manager/models/gql_models/container_registry.py index b221e7e1b8..afb88cafcd 100644 --- a/src/ai/backend/manager/models/gql_models/container_registry.py +++ b/src/ai/backend/manager/models/gql_models/container_registry.py @@ -103,16 +103,13 @@ async def mutate( scope_id: ScopeType, quota: int, ) -> Self: - graph_ctx = info.context + if not isinstance(scope_id, ProjectScope): + return UpdateQuota( + ok=False, msg="Quota mutation currently supports only the project scope." + ) - # TODO: Support other scope types - assert isinstance(scope_id, ProjectScope) project_id = scope_id.project_id - - # user = graph_ctx.user - # client_ctx = ClientContext( - # graph_ctx.db, user["domain_name"], user["uuid"], user["role"] - # ) + graph_ctx = info.context async with graph_ctx.db.begin_session() as db_sess: group_query = ( @@ -121,32 +118,40 @@ async def mutate( .options(load_only(GroupRow.container_registry)) ) result = await db_sess.execute(group_query) - - group = result.scalar_one_or_none() + group_row = result.scalar_one_or_none() if ( - group is None - or group.container_registry is None - or "registry" not in group.container_registry - or "project" not in group.container_registry + not group_row + or not group_row.container_registry + or "registry" not in group_row.container_registry + or "project" not in group_row.container_registry ): - raise ValueError("Container registry info does not exist in the group.") + return UpdateQuota( + ok=False, + msg=f"Container registry info does not exist in the group. (gr: {project_id})", + ) registry_name, project = ( - group.container_registry["registry"], - group.container_registry["project"], + group_row.container_registry["registry"], + group_row.container_registry["project"], ) - cr_query = sa.select(ContainerRegistryRow).where( + registry_query = sa.select(ContainerRegistryRow).where( (ContainerRegistryRow.registry_name == registry_name) & (ContainerRegistryRow.project == project) ) - result = await db_sess.execute(cr_query) - registry = result.fetchone()[0] + result = await db_sess.execute(registry_query) + registry = result.scalars().one_or_none() + + if not registry: + return UpdateQuota( + ok=False, + msg=f"Specified container registry row does not exist. (cr: {registry_name}, gr: {project})", + ) if registry.type != ContainerRegistryType.HARBOR2: - raise ValueError("Only HarborV2 registry is supported for now.") + return UpdateQuota(ok=False, msg="Only HarborV2 registry is supported for now.") ssl_verify = registry.ssl_verify connector = aiohttp.TCPConnector(ssl=ssl_verify) @@ -165,21 +170,30 @@ async def mutate( res = await resp.json() harbor_project_id = res["project_id"] - get_quota_id_api = (api_url / "quotas").with_query({ - "reference": "project", - "reference_id": harbor_project_id, - }) + get_quota_id_api = (api_url / "quotas").with_query({ + "reference": "project", + "reference_id": harbor_project_id, + }) async with sess.get(get_quota_id_api, allow_redirects=False, **rqst_args) as resp: res = await resp.json() - # TODO: Raise error when quota is not found or multiple quotas are found. + if not res: + return UpdateQuota( + ok=False, msg=f"Quota entity not found. (project_id: {harbor_project_id})" + ) + if len(res) > 1: + return UpdateQuota( + ok=False, + msg=f"Multiple quota entity found. (project_id: {harbor_project_id})", + ) + quota_id = res[0]["id"] - put_quota_api = api_url / "quotas" / str(quota_id) - update_payload = {"hard": {"storage": quota}} + put_quota_api = api_url / "quotas" / str(quota_id) + payload = {"hard": {"storage": quota}} async with sess.put( - put_quota_api, json=update_payload, allow_redirects=False, **rqst_args + put_quota_api, json=payload, allow_redirects=False, **rqst_args ) as resp: if resp.status == 200: return UpdateQuota(ok=True, msg="Quota updated successfully.") From c8fece0f81c01f11abfe9e63f00533457a932aa8 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Wed, 13 Nov 2024 08:55:50 +0000 Subject: [PATCH 07/41] fix: Update `resolve_registry_quota` --- .../models/gql_models/container_registry.py | 6 ++--- .../manager/models/gql_models/group.py | 25 +++++++++---------- 2 files changed, 15 insertions(+), 16 deletions(-) diff --git a/src/ai/backend/manager/models/gql_models/container_registry.py b/src/ai/backend/manager/models/gql_models/container_registry.py index afb88cafcd..73cc85b251 100644 --- a/src/ai/backend/manager/models/gql_models/container_registry.py +++ b/src/ai/backend/manager/models/gql_models/container_registry.py @@ -11,16 +11,16 @@ from sqlalchemy.orm import load_only from ai.backend.logging import BraceStyleAdapter -from ai.backend.manager.models.container_registry import ContainerRegistryRow, ContainerRegistryType -from ai.backend.manager.models.gql_models.fields import ScopeField -from ai.backend.manager.models.group import GroupRow from ..association_container_registries_groups import ( AssociationContainerRegistriesGroupsRow, ) from ..base import simple_db_mutate +from ..container_registry import ContainerRegistryRow, ContainerRegistryType from ..rbac import ProjectScope, ScopeType from ..user import UserRole +from .fields import ScopeField +from .group import GroupRow log = BraceStyleAdapter(logging.getLogger(__spec__.name)) # type: ignore diff --git a/src/ai/backend/manager/models/gql_models/group.py b/src/ai/backend/manager/models/gql_models/group.py index 979cf71a9d..5d2b1ba56a 100644 --- a/src/ai/backend/manager/models/gql_models/group.py +++ b/src/ai/backend/manager/models/gql_models/group.py @@ -15,14 +15,13 @@ from dateutil.parser import parse as dtparse from graphene.types.datetime import DateTime as GQLDateTime -from ai.backend.manager.models.container_registry import ContainerRegistryRow, ContainerRegistryType - from ..base import ( FilterExprArg, OrderExprArg, PaginatedConnectionField, generate_sql_info_for_gql_connection, ) +from ..container_registry import ContainerRegistryRow, ContainerRegistryType from ..gql_relay import ( AsyncNode, Connection, @@ -213,15 +212,9 @@ async def resolve_user_nodes( async def resolve_registry_quota(self, info: graphene.ResolveInfo) -> int: graph_ctx = info.context - - # user = graph_ctx.user - # client_ctx = ClientContext( - # graph_ctx.db, user["domain_name"], user["uuid"], user["role"] - # ) - async with graph_ctx.db.begin_session() as db_sess: if ( - self.container_registry is None + not self.container_registry or "registry" not in self.container_registry or "project" not in self.container_registry ): @@ -232,14 +225,16 @@ async def resolve_registry_quota(self, info: graphene.ResolveInfo) -> int: self.container_registry["project"], ) - cr_query = sa.select(ContainerRegistryRow).where( + registry_query = sa.select(ContainerRegistryRow).where( (ContainerRegistryRow.registry_name == registry_name) & (ContainerRegistryRow.project == project) ) - result = await db_sess.execute(cr_query) - registry = result.fetchone()[0] + result = await db_sess.execute(registry_query) + registry = result.scalars().one_or_none() + if not registry: + raise ValueError("Specified container registry row does not exist.") if registry.type != ContainerRegistryType.HARBOR2: raise ValueError("Only HarborV2 registry is supported for now.") @@ -267,7 +262,11 @@ async def resolve_registry_quota(self, info: graphene.ResolveInfo) -> int: async with sess.get(get_quota_id_api, allow_redirects=False, **rqst_args) as resp: res = await resp.json() - # TODO: Raise error when quota is not found or multiple quotas are found. + if not res: + raise ValueError("Quota not found.") + if len(res) > 1: + raise ValueError("Multiple quotas found.") + quota = res[0]["hard"]["storage"] return quota From 4f3448abb8630618043f1b80bcafe632060cb9bf Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Thu, 14 Nov 2024 01:23:09 +0000 Subject: [PATCH 08/41] fix: Update `resolve_registry_quota` --- src/ai/backend/manager/api/exceptions.py | 4 ++++ src/ai/backend/manager/models/gql_models/group.py | 10 +++++++--- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/src/ai/backend/manager/api/exceptions.py b/src/ai/backend/manager/api/exceptions.py index 23f381d30b..01f42af6a7 100644 --- a/src/ai/backend/manager/api/exceptions.py +++ b/src/ai/backend/manager/api/exceptions.py @@ -242,6 +242,10 @@ class EndpointTokenNotFound(ObjectNotFound): object_name = "endpoint_token" +class ContainerRegistryNotFound(ObjectNotFound): + object_name = "endpoint_token" + + class TooManySessionsMatched(BackendError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/too-many-sessions-matched" error_title = "Too many sessions matched." diff --git a/src/ai/backend/manager/models/gql_models/group.py b/src/ai/backend/manager/models/gql_models/group.py index 5d2b1ba56a..185a678b28 100644 --- a/src/ai/backend/manager/models/gql_models/group.py +++ b/src/ai/backend/manager/models/gql_models/group.py @@ -15,6 +15,8 @@ from dateutil.parser import parse as dtparse from graphene.types.datetime import DateTime as GQLDateTime +from ai.backend.manager.api.exceptions import ContainerRegistryNotFound + from ..base import ( FilterExprArg, OrderExprArg, @@ -218,7 +220,9 @@ async def resolve_registry_quota(self, info: graphene.ResolveInfo) -> int: or "registry" not in self.container_registry or "project" not in self.container_registry ): - raise ValueError("Container registry info does not exist in the group.") + raise ContainerRegistryNotFound( + "Container registry info does not exist in the group." + ) registry_name, project = ( self.container_registry["registry"], @@ -234,9 +238,9 @@ async def resolve_registry_quota(self, info: graphene.ResolveInfo) -> int: registry = result.scalars().one_or_none() if not registry: - raise ValueError("Specified container registry row does not exist.") + raise ContainerRegistryNotFound("Specified container registry row does not exist.") if registry.type != ContainerRegistryType.HARBOR2: - raise ValueError("Only HarborV2 registry is supported for now.") + raise NotImplementedError("Only HarborV2 registry is supported for now.") ssl_verify = registry.ssl_verify connector = aiohttp.TCPConnector(ssl=ssl_verify) From 780a250ac192db36683bbfff457eab0d47fd6996 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Thu, 14 Nov 2024 02:31:32 +0000 Subject: [PATCH 09/41] fix: Only authorized groups view the quota --- src/ai/backend/manager/models/gql_models/group.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/src/ai/backend/manager/models/gql_models/group.py b/src/ai/backend/manager/models/gql_models/group.py index 185a678b28..57de38c8ca 100644 --- a/src/ai/backend/manager/models/gql_models/group.py +++ b/src/ai/backend/manager/models/gql_models/group.py @@ -17,6 +17,9 @@ from ai.backend.manager.api.exceptions import ContainerRegistryNotFound +from ..association_container_registries_groups import ( + AssociationContainerRegistriesGroupsRow, +) from ..base import ( FilterExprArg, OrderExprArg, @@ -242,6 +245,17 @@ async def resolve_registry_quota(self, info: graphene.ResolveInfo) -> int: if registry.type != ContainerRegistryType.HARBOR2: raise NotImplementedError("Only HarborV2 registry is supported for now.") + if not registry.is_global: + get_assoc_query = sa.select( + sa.exists() + .where(AssociationContainerRegistriesGroupsRow.registry_id == registry.id) + .where(AssociationContainerRegistriesGroupsRow.group_id == self.row_id) + ) + assoc_exist = (await db_sess.execute(get_assoc_query)).scalar() + + if not assoc_exist: + raise ValueError("The group is not associated with the container registry.") + ssl_verify = registry.ssl_verify connector = aiohttp.TCPConnector(ssl=ssl_verify) From 1a4bf708c6cb913cde8d538f908858b7c0e39aeb Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Thu, 14 Nov 2024 02:50:56 +0000 Subject: [PATCH 10/41] feat: Add CreateQuota, DeleteQuota mutation --- src/ai/backend/manager/models/gql.py | 4 + .../models/gql_models/container_registry.py | 236 +++++++++++------- 2 files changed, 148 insertions(+), 92 deletions(-) diff --git a/src/ai/backend/manager/models/gql.py b/src/ai/backend/manager/models/gql.py index 912b94c02a..21d52c31bb 100644 --- a/src/ai/backend/manager/models/gql.py +++ b/src/ai/backend/manager/models/gql.py @@ -74,6 +74,8 @@ ) from .gql_models.container_registry import ( AssociateContainerRegistryWithGroup, + CreateQuota, + DeleteQuota, DisassociateContainerRegistryWithGroup, UpdateQuota, ) @@ -350,7 +352,9 @@ class Mutations(graphene.ObjectType): description="Added in 24.12.0" ) + create_container_registry_quota = CreateQuota.Field(description="Added in 24.12.0") update_container_registry_quota = UpdateQuota.Field(description="Added in 24.12.0") + delete_container_registry_quota = DeleteQuota.Field(description="Added in 24.12.0") # Legacy mutations create_container_registry = CreateContainerRegistry.Field() diff --git a/src/ai/backend/manager/models/gql_models/container_registry.py b/src/ai/backend/manager/models/gql_models/container_registry.py index 73cc85b251..279c15baa4 100644 --- a/src/ai/backend/manager/models/gql_models/container_registry.py +++ b/src/ai/backend/manager/models/gql_models/container_registry.py @@ -80,7 +80,109 @@ async def mutate( return await simple_db_mutate(cls, info.context, delete_query) -class UpdateQuota(graphene.Mutation): +async def update_quota( + cls: Any, info: graphene.ResolveInfo, scope_id: ScopeType, quota: int +) -> Any: + if not isinstance(scope_id, ProjectScope): + return cls(ok=False, msg="Quota mutation currently supports only the project scope.") + + project_id = scope_id.project_id + graph_ctx = info.context + + async with graph_ctx.db.begin_session() as db_sess: + group_query = ( + sa.select(GroupRow) + .where(GroupRow.id == project_id) + .options(load_only(GroupRow.container_registry)) + ) + result = await db_sess.execute(group_query) + group_row = result.scalar_one_or_none() + + if ( + not group_row + or not group_row.container_registry + or "registry" not in group_row.container_registry + or "project" not in group_row.container_registry + ): + return UpdateQuota( + ok=False, + msg=f"Container registry info does not exist in the group. (gr: {project_id})", + ) + + registry_name, project = ( + group_row.container_registry["registry"], + group_row.container_registry["project"], + ) + + registry_query = sa.select(ContainerRegistryRow).where( + (ContainerRegistryRow.registry_name == registry_name) + & (ContainerRegistryRow.project == project) + ) + + result = await db_sess.execute(registry_query) + registry = result.scalars().one_or_none() + + if not registry: + return cls( + ok=False, + msg=f"Specified container registry row does not exist. (cr: {registry_name}, gr: {project})", + ) + + if registry.type != ContainerRegistryType.HARBOR2: + return cls(ok=False, msg="Only HarborV2 registry is supported for now.") + + ssl_verify = registry.ssl_verify + connector = aiohttp.TCPConnector(ssl=ssl_verify) + + api_url = yarl.URL(registry.url) / "api" / "v2.0" + async with aiohttp.ClientSession(connector=connector) as sess: + rqst_args: dict[str, Any] = {} + rqst_args["auth"] = aiohttp.BasicAuth( + registry.username, + registry.password, + ) + + get_project_id_api = api_url / "projects" / project + + async with sess.get(get_project_id_api, allow_redirects=False, **rqst_args) as resp: + res = await resp.json() + harbor_project_id = res["project_id"] + + get_quota_id_api = (api_url / "quotas").with_query({ + "reference": "project", + "reference_id": harbor_project_id, + }) + + async with sess.get(get_quota_id_api, allow_redirects=False, **rqst_args) as resp: + res = await resp.json() + if not res: + return cls( + ok=False, msg=f"Quota entity not found. (project_id: {harbor_project_id})" + ) + if len(res) > 1: + return cls( + ok=False, + msg=f"Multiple quota entity found. (project_id: {harbor_project_id})", + ) + + quota_id = res[0]["id"] + + put_quota_api = api_url / "quotas" / str(quota_id) + payload = {"hard": {"storage": quota}} + + async with sess.put( + put_quota_api, json=payload, allow_redirects=False, **rqst_args + ) as resp: + if resp.status == 200: + return cls(ok=True, msg="Quota updated successfully.") + else: + log.error(f"Failed to update quota: {await resp.json()}") + return cls(ok=False, msg=f"Failed to update quota. Status code: {resp.status}") + + return cls(ok=False, msg="Unknown error!") + + +class CreateQuota(graphene.Mutation): """Added in 24.12.0.""" allowed_roles = ( @@ -103,104 +205,54 @@ async def mutate( scope_id: ScopeType, quota: int, ) -> Self: - if not isinstance(scope_id, ProjectScope): - return UpdateQuota( - ok=False, msg="Quota mutation currently supports only the project scope." - ) + return await update_quota(cls, info, scope_id, quota) - project_id = scope_id.project_id - graph_ctx = info.context - async with graph_ctx.db.begin_session() as db_sess: - group_query = ( - sa.select(GroupRow) - .where(GroupRow.id == project_id) - .options(load_only(GroupRow.container_registry)) - ) - result = await db_sess.execute(group_query) - group_row = result.scalar_one_or_none() - - if ( - not group_row - or not group_row.container_registry - or "registry" not in group_row.container_registry - or "project" not in group_row.container_registry - ): - return UpdateQuota( - ok=False, - msg=f"Container registry info does not exist in the group. (gr: {project_id})", - ) +class UpdateQuota(graphene.Mutation): + """Added in 24.12.0.""" - registry_name, project = ( - group_row.container_registry["registry"], - group_row.container_registry["project"], - ) + allowed_roles = ( + UserRole.SUPERADMIN, + UserRole.ADMIN, + ) - registry_query = sa.select(ContainerRegistryRow).where( - (ContainerRegistryRow.registry_name == registry_name) - & (ContainerRegistryRow.project == project) - ) + class Arguments: + scope_id = ScopeField(required=True) + quota = graphene.Int(required=True) - result = await db_sess.execute(registry_query) - registry = result.scalars().one_or_none() + ok = graphene.Boolean() + msg = graphene.String() - if not registry: - return UpdateQuota( - ok=False, - msg=f"Specified container registry row does not exist. (cr: {registry_name}, gr: {project})", - ) + @classmethod + async def mutate( + cls, + root, + info: graphene.ResolveInfo, + scope_id: ScopeType, + quota: int, + ) -> Self: + return await update_quota(cls, info, scope_id, quota) - if registry.type != ContainerRegistryType.HARBOR2: - return UpdateQuota(ok=False, msg="Only HarborV2 registry is supported for now.") - ssl_verify = registry.ssl_verify - connector = aiohttp.TCPConnector(ssl=ssl_verify) +class DeleteQuota(graphene.Mutation): + """Added in 24.12.0.""" - api_url = yarl.URL(registry.url) / "api" / "v2.0" - async with aiohttp.ClientSession(connector=connector) as sess: - rqst_args: dict[str, Any] = {} - rqst_args["auth"] = aiohttp.BasicAuth( - registry.username, - registry.password, - ) + allowed_roles = ( + UserRole.SUPERADMIN, + UserRole.ADMIN, + ) + + class Arguments: + scope_id = ScopeField(required=True) + + ok = graphene.Boolean() + msg = graphene.String() - get_project_id_api = api_url / "projects" / project - - async with sess.get(get_project_id_api, allow_redirects=False, **rqst_args) as resp: - res = await resp.json() - harbor_project_id = res["project_id"] - - get_quota_id_api = (api_url / "quotas").with_query({ - "reference": "project", - "reference_id": harbor_project_id, - }) - - async with sess.get(get_quota_id_api, allow_redirects=False, **rqst_args) as resp: - res = await resp.json() - if not res: - return UpdateQuota( - ok=False, msg=f"Quota entity not found. (project_id: {harbor_project_id})" - ) - if len(res) > 1: - return UpdateQuota( - ok=False, - msg=f"Multiple quota entity found. (project_id: {harbor_project_id})", - ) - - quota_id = res[0]["id"] - - put_quota_api = api_url / "quotas" / str(quota_id) - payload = {"hard": {"storage": quota}} - - async with sess.put( - put_quota_api, json=payload, allow_redirects=False, **rqst_args - ) as resp: - if resp.status == 200: - return UpdateQuota(ok=True, msg="Quota updated successfully.") - else: - log.error(f"Failed to update quota: {await resp.json()}") - return UpdateQuota( - ok=False, msg=f"Failed to update quota. Status code: {resp.status}" - ) - - return UpdateQuota(ok=False, msg="Unknown error!") + @classmethod + async def mutate( + cls, + root, + info: graphene.ResolveInfo, + scope_id: ScopeType, + ) -> Self: + return await update_quota(cls, info, scope_id, -1) From 63d5dbb426af938725fc47d28b33c86dd4856001 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Thu, 14 Nov 2024 02:55:51 +0000 Subject: [PATCH 11/41] chore: Rename function --- .../manager/models/gql_models/container_registry.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/ai/backend/manager/models/gql_models/container_registry.py b/src/ai/backend/manager/models/gql_models/container_registry.py index 279c15baa4..dc5c76fc7b 100644 --- a/src/ai/backend/manager/models/gql_models/container_registry.py +++ b/src/ai/backend/manager/models/gql_models/container_registry.py @@ -80,7 +80,7 @@ async def mutate( return await simple_db_mutate(cls, info.context, delete_query) -async def update_quota( +async def update_harbor_project_quota( cls: Any, info: graphene.ResolveInfo, scope_id: ScopeType, quota: int ) -> Any: if not isinstance(scope_id, ProjectScope): @@ -205,7 +205,7 @@ async def mutate( scope_id: ScopeType, quota: int, ) -> Self: - return await update_quota(cls, info, scope_id, quota) + return await update_harbor_project_quota(cls, info, scope_id, quota) class UpdateQuota(graphene.Mutation): @@ -231,7 +231,7 @@ async def mutate( scope_id: ScopeType, quota: int, ) -> Self: - return await update_quota(cls, info, scope_id, quota) + return await update_harbor_project_quota(cls, info, scope_id, quota) class DeleteQuota(graphene.Mutation): @@ -255,4 +255,4 @@ async def mutate( info: graphene.ResolveInfo, scope_id: ScopeType, ) -> Self: - return await update_quota(cls, info, scope_id, -1) + return await update_harbor_project_quota(cls, info, scope_id, -1) From 473b1254664b0d26c3a5113d7aeebc10e482abe5 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Thu, 14 Nov 2024 03:52:45 +0000 Subject: [PATCH 12/41] fix: Add exception handling for each operation --- .../models/gql_models/container_registry.py | 46 +++++++++++++++---- 1 file changed, 36 insertions(+), 10 deletions(-) diff --git a/src/ai/backend/manager/models/gql_models/container_registry.py b/src/ai/backend/manager/models/gql_models/container_registry.py index dc5c76fc7b..c5e869c52e 100644 --- a/src/ai/backend/manager/models/gql_models/container_registry.py +++ b/src/ai/backend/manager/models/gql_models/container_registry.py @@ -1,5 +1,6 @@ from __future__ import annotations +import enum import logging from typing import Any, Self @@ -80,8 +81,18 @@ async def mutate( return await simple_db_mutate(cls, info.context, delete_query) +class UpdateQuotaOperationType(enum.StrEnum): + CREATE = "create" + DELETE = "delete" + UPDATE = "update" + + async def update_harbor_project_quota( - cls: Any, info: graphene.ResolveInfo, scope_id: ScopeType, quota: int + operation_type: UpdateQuotaOperationType, + cls: Any, + info: graphene.ResolveInfo, + scope_id: ScopeType, + quota: int, ) -> Any: if not isinstance(scope_id, ProjectScope): return cls(ok=False, msg="Quota mutation currently supports only the project scope.") @@ -106,7 +117,7 @@ async def update_harbor_project_quota( ): return UpdateQuota( ok=False, - msg=f"Container registry info does not exist in the group. (gr: {project_id})", + msg=f"Container registry info does not exist or is invalid in the group. (gr: {project_id})", ) registry_name, project = ( @@ -133,8 +144,6 @@ async def update_harbor_project_quota( ssl_verify = registry.ssl_verify connector = aiohttp.TCPConnector(ssl=ssl_verify) - - api_url = yarl.URL(registry.url) / "api" / "v2.0" async with aiohttp.ClientSession(connector=connector) as sess: rqst_args: dict[str, Any] = {} rqst_args["auth"] = aiohttp.BasicAuth( @@ -142,6 +151,7 @@ async def update_harbor_project_quota( registry.password, ) + api_url = yarl.URL(registry.url) / "api" / "v2.0" get_project_id_api = api_url / "projects" / project async with sess.get(get_project_id_api, allow_redirects=False, **rqst_args) as resp: @@ -165,6 +175,14 @@ async def update_harbor_project_quota( msg=f"Multiple quota entity found. (project_id: {harbor_project_id})", ) + previous_quota = res[0]["hard"]["storage"] + if operation_type == UpdateQuotaOperationType.DELETE: + if previous_quota == -1: + return cls(ok=False, msg=f"Quota is not set. (gr: {project_id})") + elif operation_type == UpdateQuotaOperationType.CREATE: + if previous_quota > 0: + return cls(ok=False, msg=f"Quota already exists. (gr: {project_id})") + quota_id = res[0]["id"] put_quota_api = api_url / "quotas" / str(quota_id) @@ -174,10 +192,12 @@ async def update_harbor_project_quota( put_quota_api, json=payload, allow_redirects=False, **rqst_args ) as resp: if resp.status == 200: - return cls(ok=True, msg="Quota updated successfully.") + return cls(ok=True, msg="success") else: - log.error(f"Failed to update quota: {await resp.json()}") - return cls(ok=False, msg=f"Failed to update quota. Status code: {resp.status}") + log.error(f"Failed to {operation_type} quota: {await resp.json()}") + return cls( + ok=False, msg=f"Failed to {operation_type} quota. Status code: {resp.status}" + ) return cls(ok=False, msg="Unknown error!") @@ -205,7 +225,9 @@ async def mutate( scope_id: ScopeType, quota: int, ) -> Self: - return await update_harbor_project_quota(cls, info, scope_id, quota) + return await update_harbor_project_quota( + UpdateQuotaOperationType.CREATE, cls, info, scope_id, quota + ) class UpdateQuota(graphene.Mutation): @@ -231,7 +253,9 @@ async def mutate( scope_id: ScopeType, quota: int, ) -> Self: - return await update_harbor_project_quota(cls, info, scope_id, quota) + return await update_harbor_project_quota( + UpdateQuotaOperationType.UPDATE, cls, info, scope_id, quota + ) class DeleteQuota(graphene.Mutation): @@ -255,4 +279,6 @@ async def mutate( info: graphene.ResolveInfo, scope_id: ScopeType, ) -> Self: - return await update_harbor_project_quota(cls, info, scope_id, -1) + return await update_harbor_project_quota( + UpdateQuotaOperationType.DELETE, cls, info, scope_id, -1 + ) From 1e8ac1ac592eac36315e874cbed0dfb43e8a5b23 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Thu, 14 Nov 2024 03:53:59 +0000 Subject: [PATCH 13/41] chore: Update schema --- src/ai/backend/manager/api/schema.graphql | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/src/ai/backend/manager/api/schema.graphql b/src/ai/backend/manager/api/schema.graphql index bfebd14ba9..f67922d130 100644 --- a/src/ai/backend/manager/api/schema.graphql +++ b/src/ai/backend/manager/api/schema.graphql @@ -1848,8 +1848,14 @@ type Mutations { """Added in 24.12.0""" disassociate_container_registry_with_group(group_id: String!, registry_id: String!): DisassociateContainerRegistryWithGroup + """Added in 24.12.0""" + create_container_registry_quota(quota: Int!, scope_id: ScopeField!): CreateQuota + """Added in 24.12.0""" update_container_registry_quota(quota: Int!, scope_id: ScopeField!): UpdateQuota + + """Added in 24.12.0""" + delete_container_registry_quota(scope_id: ScopeField!): DeleteQuota create_container_registry(hostname: String!, props: CreateContainerRegistryInput!): CreateContainerRegistry modify_container_registry(hostname: String!, props: ModifyContainerRegistryInput!): ModifyContainerRegistry delete_container_registry(hostname: String!): DeleteContainerRegistry @@ -2576,12 +2582,24 @@ type DisassociateContainerRegistryWithGroup { msg: String } +"""Added in 24.12.0.""" +type CreateQuota { + ok: Boolean + msg: String +} + """Added in 24.12.0.""" type UpdateQuota { ok: Boolean msg: String } +"""Added in 24.12.0.""" +type DeleteQuota { + ok: Boolean + msg: String +} + type CreateContainerRegistry { container_registry: ContainerRegistry } From cfd21fce251b7461c13884213441cba514e2f74f Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Thu, 14 Nov 2024 03:57:46 +0000 Subject: [PATCH 14/41] chore: Update error msg --- .../backend/manager/models/gql_models/container_registry.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/ai/backend/manager/models/gql_models/container_registry.py b/src/ai/backend/manager/models/gql_models/container_registry.py index c5e869c52e..aa66238d51 100644 --- a/src/ai/backend/manager/models/gql_models/container_registry.py +++ b/src/ai/backend/manager/models/gql_models/container_registry.py @@ -172,16 +172,16 @@ async def update_harbor_project_quota( if len(res) > 1: return cls( ok=False, - msg=f"Multiple quota entity found. (project_id: {harbor_project_id})", + msg=f"Multiple quota entities found. (project_id: {harbor_project_id})", ) previous_quota = res[0]["hard"]["storage"] if operation_type == UpdateQuotaOperationType.DELETE: if previous_quota == -1: - return cls(ok=False, msg=f"Quota is not set. (gr: {project_id})") + return cls(ok=False, msg=f"Quota entity not found. (gr: {project_id})") elif operation_type == UpdateQuotaOperationType.CREATE: if previous_quota > 0: - return cls(ok=False, msg=f"Quota already exists. (gr: {project_id})") + return cls(ok=False, msg=f"Quota limit already exists. (gr: {project_id})") quota_id = res[0]["id"] From a49f1ff676f50ddede807d2af2fefce3cbd887d9 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Thu, 14 Nov 2024 03:59:18 +0000 Subject: [PATCH 15/41] chore: Rename variable --- .../models/gql_models/container_registry.py | 26 +++++++++++-------- 1 file changed, 15 insertions(+), 11 deletions(-) diff --git a/src/ai/backend/manager/models/gql_models/container_registry.py b/src/ai/backend/manager/models/gql_models/container_registry.py index aa66238d51..b27dc7c5f2 100644 --- a/src/ai/backend/manager/models/gql_models/container_registry.py +++ b/src/ai/backend/manager/models/gql_models/container_registry.py @@ -89,13 +89,15 @@ class UpdateQuotaOperationType(enum.StrEnum): async def update_harbor_project_quota( operation_type: UpdateQuotaOperationType, - cls: Any, + mutation_cls: Any, info: graphene.ResolveInfo, scope_id: ScopeType, quota: int, ) -> Any: if not isinstance(scope_id, ProjectScope): - return cls(ok=False, msg="Quota mutation currently supports only the project scope.") + return mutation_cls( + ok=False, msg="Quota mutation currently supports only the project scope." + ) project_id = scope_id.project_id graph_ctx = info.context @@ -134,13 +136,13 @@ async def update_harbor_project_quota( registry = result.scalars().one_or_none() if not registry: - return cls( + return mutation_cls( ok=False, msg=f"Specified container registry row does not exist. (cr: {registry_name}, gr: {project})", ) if registry.type != ContainerRegistryType.HARBOR2: - return cls(ok=False, msg="Only HarborV2 registry is supported for now.") + return mutation_cls(ok=False, msg="Only HarborV2 registry is supported for now.") ssl_verify = registry.ssl_verify connector = aiohttp.TCPConnector(ssl=ssl_verify) @@ -166,11 +168,11 @@ async def update_harbor_project_quota( async with sess.get(get_quota_id_api, allow_redirects=False, **rqst_args) as resp: res = await resp.json() if not res: - return cls( + return mutation_cls( ok=False, msg=f"Quota entity not found. (project_id: {harbor_project_id})" ) if len(res) > 1: - return cls( + return mutation_cls( ok=False, msg=f"Multiple quota entities found. (project_id: {harbor_project_id})", ) @@ -178,10 +180,12 @@ async def update_harbor_project_quota( previous_quota = res[0]["hard"]["storage"] if operation_type == UpdateQuotaOperationType.DELETE: if previous_quota == -1: - return cls(ok=False, msg=f"Quota entity not found. (gr: {project_id})") + return mutation_cls(ok=False, msg=f"Quota entity not found. (gr: {project_id})") elif operation_type == UpdateQuotaOperationType.CREATE: if previous_quota > 0: - return cls(ok=False, msg=f"Quota limit already exists. (gr: {project_id})") + return mutation_cls( + ok=False, msg=f"Quota limit already exists. (gr: {project_id})" + ) quota_id = res[0]["id"] @@ -192,14 +196,14 @@ async def update_harbor_project_quota( put_quota_api, json=payload, allow_redirects=False, **rqst_args ) as resp: if resp.status == 200: - return cls(ok=True, msg="success") + return mutation_cls(ok=True, msg="success") else: log.error(f"Failed to {operation_type} quota: {await resp.json()}") - return cls( + return mutation_cls( ok=False, msg=f"Failed to {operation_type} quota. Status code: {resp.status}" ) - return cls(ok=False, msg="Unknown error!") + return mutation_cls(ok=False, msg="Unknown error!") class CreateQuota(graphene.Mutation): From 16b45d11e5c695f6c2f19481dfd3737e258300ba Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Thu, 14 Nov 2024 04:08:21 +0000 Subject: [PATCH 16/41] fix: Remove useless strenum --- .../models/gql_models/container_registry.py | 30 +++++++------------ 1 file changed, 10 insertions(+), 20 deletions(-) diff --git a/src/ai/backend/manager/models/gql_models/container_registry.py b/src/ai/backend/manager/models/gql_models/container_registry.py index b27dc7c5f2..bf75cef10b 100644 --- a/src/ai/backend/manager/models/gql_models/container_registry.py +++ b/src/ai/backend/manager/models/gql_models/container_registry.py @@ -1,8 +1,7 @@ from __future__ import annotations -import enum import logging -from typing import Any, Self +from typing import Any, Literal, Self import aiohttp import aiohttp.client_exceptions @@ -81,19 +80,16 @@ async def mutate( return await simple_db_mutate(cls, info.context, delete_query) -class UpdateQuotaOperationType(enum.StrEnum): - CREATE = "create" - DELETE = "delete" - UPDATE = "update" - - async def update_harbor_project_quota( - operation_type: UpdateQuotaOperationType, + operation_type: Literal["create", "delete", "update"], mutation_cls: Any, info: graphene.ResolveInfo, scope_id: ScopeType, quota: int, ) -> Any: + """ + Utility function for code reuse of the HarborV2 per-project Quota CRUD API + """ if not isinstance(scope_id, ProjectScope): return mutation_cls( ok=False, msg="Quota mutation currently supports only the project scope." @@ -178,10 +174,10 @@ async def update_harbor_project_quota( ) previous_quota = res[0]["hard"]["storage"] - if operation_type == UpdateQuotaOperationType.DELETE: + if operation_type == "delete": if previous_quota == -1: return mutation_cls(ok=False, msg=f"Quota entity not found. (gr: {project_id})") - elif operation_type == UpdateQuotaOperationType.CREATE: + elif operation_type == "create": if previous_quota > 0: return mutation_cls( ok=False, msg=f"Quota limit already exists. (gr: {project_id})" @@ -229,9 +225,7 @@ async def mutate( scope_id: ScopeType, quota: int, ) -> Self: - return await update_harbor_project_quota( - UpdateQuotaOperationType.CREATE, cls, info, scope_id, quota - ) + return await update_harbor_project_quota("create", cls, info, scope_id, quota) class UpdateQuota(graphene.Mutation): @@ -257,9 +251,7 @@ async def mutate( scope_id: ScopeType, quota: int, ) -> Self: - return await update_harbor_project_quota( - UpdateQuotaOperationType.UPDATE, cls, info, scope_id, quota - ) + return await update_harbor_project_quota("update", cls, info, scope_id, quota) class DeleteQuota(graphene.Mutation): @@ -283,6 +275,4 @@ async def mutate( info: graphene.ResolveInfo, scope_id: ScopeType, ) -> Self: - return await update_harbor_project_quota( - UpdateQuotaOperationType.DELETE, cls, info, scope_id, -1 - ) + return await update_harbor_project_quota("delete", cls, info, scope_id, -1) From 3128b0a27710b4fe0b1cf0ac1c086a1f4fd164c3 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Mon, 18 Nov 2024 02:40:53 +0000 Subject: [PATCH 17/41] refactor: `mutate_harbor_project_quota` --- .../models/gql_models/container_registry.py | 136 ++++++++++-------- 1 file changed, 73 insertions(+), 63 deletions(-) diff --git a/src/ai/backend/manager/models/gql_models/container_registry.py b/src/ai/backend/manager/models/gql_models/container_registry.py index bf75cef10b..1673ed361c 100644 --- a/src/ai/backend/manager/models/gql_models/container_registry.py +++ b/src/ai/backend/manager/models/gql_models/container_registry.py @@ -8,9 +8,17 @@ import graphene import sqlalchemy as sa import yarl +from sqlalchemy.ext.asyncio import AsyncSession as SASession from sqlalchemy.orm import load_only from ai.backend.logging import BraceStyleAdapter +from ai.backend.manager.api.exceptions import ( + ContainerRegistryNotFound, + GenericBadRequest, + InternalServerError, + NotImplementedAPI, + ObjectNotFound, +) from ..association_container_registries_groups import ( AssociationContainerRegistriesGroupsRow, @@ -80,65 +88,57 @@ async def mutate( return await simple_db_mutate(cls, info.context, delete_query) -async def update_harbor_project_quota( +async def mutate_harbor_project_quota( operation_type: Literal["create", "delete", "update"], - mutation_cls: Any, - info: graphene.ResolveInfo, + db_sess: SASession, scope_id: ScopeType, quota: int, -) -> Any: +) -> None: """ Utility function for code reuse of the HarborV2 per-project Quota CRUD API """ if not isinstance(scope_id, ProjectScope): - return mutation_cls( - ok=False, msg="Quota mutation currently supports only the project scope." - ) + raise NotImplementedAPI("Quota mutation currently supports only the project scope.") project_id = scope_id.project_id - graph_ctx = info.context - - async with graph_ctx.db.begin_session() as db_sess: - group_query = ( - sa.select(GroupRow) - .where(GroupRow.id == project_id) - .options(load_only(GroupRow.container_registry)) - ) - result = await db_sess.execute(group_query) - group_row = result.scalar_one_or_none() - - if ( - not group_row - or not group_row.container_registry - or "registry" not in group_row.container_registry - or "project" not in group_row.container_registry - ): - return UpdateQuota( - ok=False, - msg=f"Container registry info does not exist or is invalid in the group. (gr: {project_id})", - ) - - registry_name, project = ( - group_row.container_registry["registry"], - group_row.container_registry["project"], + group_query = ( + sa.select(GroupRow) + .where(GroupRow.id == project_id) + .options(load_only(GroupRow.container_registry)) + ) + result = await db_sess.execute(group_query) + group_row = result.scalar_one_or_none() + + if ( + not group_row + or not group_row.container_registry + or "registry" not in group_row.container_registry + or "project" not in group_row.container_registry + ): + raise ContainerRegistryNotFound( + f"Container registry info does not exist or is invalid in the group. (gr: {project_id})" ) - registry_query = sa.select(ContainerRegistryRow).where( - (ContainerRegistryRow.registry_name == registry_name) - & (ContainerRegistryRow.project == project) - ) + registry_name, project = ( + group_row.container_registry["registry"], + group_row.container_registry["project"], + ) - result = await db_sess.execute(registry_query) - registry = result.scalars().one_or_none() + registry_query = sa.select(ContainerRegistryRow).where( + (ContainerRegistryRow.registry_name == registry_name) + & (ContainerRegistryRow.project == project) + ) - if not registry: - return mutation_cls( - ok=False, - msg=f"Specified container registry row does not exist. (cr: {registry_name}, gr: {project})", - ) + result = await db_sess.execute(registry_query) + registry = result.scalars().one_or_none() - if registry.type != ContainerRegistryType.HARBOR2: - return mutation_cls(ok=False, msg="Only HarborV2 registry is supported for now.") + if not registry: + raise ContainerRegistryNotFound( + f"Specified container registry row does not exist. (cr: {registry_name}, gr: {project})" + ) + + if registry.type != ContainerRegistryType.HARBOR2: + raise NotImplementedAPI("Only HarborV2 registry is supported for now.") ssl_verify = registry.ssl_verify connector = aiohttp.TCPConnector(ssl=ssl_verify) @@ -164,24 +164,19 @@ async def update_harbor_project_quota( async with sess.get(get_quota_id_api, allow_redirects=False, **rqst_args) as resp: res = await resp.json() if not res: - return mutation_cls( - ok=False, msg=f"Quota entity not found. (project_id: {harbor_project_id})" - ) + raise ObjectNotFound(object_name="quota entity") if len(res) > 1: - return mutation_cls( - ok=False, - msg=f"Multiple quota entities found. (project_id: {harbor_project_id})", + raise InternalServerError( + f"Multiple quota entities found. (project_id: {harbor_project_id})" ) previous_quota = res[0]["hard"]["storage"] - if operation_type == "delete": + if operation_type == "update" or operation_type == "delete": if previous_quota == -1: - return mutation_cls(ok=False, msg=f"Quota entity not found. (gr: {project_id})") + raise ObjectNotFound(object_name="quota entity") elif operation_type == "create": if previous_quota > 0: - return mutation_cls( - ok=False, msg=f"Quota limit already exists. (gr: {project_id})" - ) + raise GenericBadRequest(f"Quota limit already exists. (gr: {project_id})") quota_id = res[0]["id"] @@ -192,14 +187,14 @@ async def update_harbor_project_quota( put_quota_api, json=payload, allow_redirects=False, **rqst_args ) as resp: if resp.status == 200: - return mutation_cls(ok=True, msg="success") + return else: log.error(f"Failed to {operation_type} quota: {await resp.json()}") - return mutation_cls( - ok=False, msg=f"Failed to {operation_type} quota. Status code: {resp.status}" + raise InternalServerError( + f"Failed to {operation_type} quota. Status code: {resp.status}" ) - return mutation_cls(ok=False, msg="Unknown error!") + raise InternalServerError("Unknown error!") class CreateQuota(graphene.Mutation): @@ -225,7 +220,12 @@ async def mutate( scope_id: ScopeType, quota: int, ) -> Self: - return await update_harbor_project_quota("create", cls, info, scope_id, quota) + async with info.context.db.begin_session() as db_sess: + try: + await mutate_harbor_project_quota("create", db_sess, scope_id, quota) + return cls(ok=True, msg="success") + except Exception as e: + return cls(ok=False, msg=str(e)) class UpdateQuota(graphene.Mutation): @@ -251,7 +251,12 @@ async def mutate( scope_id: ScopeType, quota: int, ) -> Self: - return await update_harbor_project_quota("update", cls, info, scope_id, quota) + async with info.context.db.begin_session() as db_sess: + try: + await mutate_harbor_project_quota("update", db_sess, scope_id, quota) + return cls(ok=True, msg="success") + except Exception as e: + return cls(ok=False, msg=str(e)) class DeleteQuota(graphene.Mutation): @@ -275,4 +280,9 @@ async def mutate( info: graphene.ResolveInfo, scope_id: ScopeType, ) -> Self: - return await update_harbor_project_quota("delete", cls, info, scope_id, -1) + async with info.context.db.begin_session() as db_sess: + try: + await mutate_harbor_project_quota("delete", db_sess, scope_id, -1) + return cls(ok=True, msg="success") + except Exception as e: + return cls(ok=False, msg=str(e)) From 97a302221248e1440be596b283741f6f64d307ca Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Mon, 18 Nov 2024 03:16:24 +0000 Subject: [PATCH 18/41] refactor: Add read operation handling for code reuse --- .../models/gql_models/container_registry.py | 134 +-------------- .../gql_models/container_registry_utils.py | 159 ++++++++++++++++++ .../manager/models/gql_models/group.py | 87 +--------- 3 files changed, 173 insertions(+), 207 deletions(-) create mode 100644 src/ai/backend/manager/models/gql_models/container_registry_utils.py diff --git a/src/ai/backend/manager/models/gql_models/container_registry.py b/src/ai/backend/manager/models/gql_models/container_registry.py index 1673ed361c..b78939f041 100644 --- a/src/ai/backend/manager/models/gql_models/container_registry.py +++ b/src/ai/backend/manager/models/gql_models/container_registry.py @@ -1,34 +1,21 @@ from __future__ import annotations import logging -from typing import Any, Literal, Self +from typing import Self -import aiohttp -import aiohttp.client_exceptions import graphene import sqlalchemy as sa -import yarl -from sqlalchemy.ext.asyncio import AsyncSession as SASession -from sqlalchemy.orm import load_only from ai.backend.logging import BraceStyleAdapter -from ai.backend.manager.api.exceptions import ( - ContainerRegistryNotFound, - GenericBadRequest, - InternalServerError, - NotImplementedAPI, - ObjectNotFound, -) from ..association_container_registries_groups import ( AssociationContainerRegistriesGroupsRow, ) from ..base import simple_db_mutate -from ..container_registry import ContainerRegistryRow, ContainerRegistryType -from ..rbac import ProjectScope, ScopeType +from ..rbac import ScopeType from ..user import UserRole +from .container_registry_utils import handle_harbor_project_quota_operation from .fields import ScopeField -from .group import GroupRow log = BraceStyleAdapter(logging.getLogger(__spec__.name)) # type: ignore @@ -88,115 +75,6 @@ async def mutate( return await simple_db_mutate(cls, info.context, delete_query) -async def mutate_harbor_project_quota( - operation_type: Literal["create", "delete", "update"], - db_sess: SASession, - scope_id: ScopeType, - quota: int, -) -> None: - """ - Utility function for code reuse of the HarborV2 per-project Quota CRUD API - """ - if not isinstance(scope_id, ProjectScope): - raise NotImplementedAPI("Quota mutation currently supports only the project scope.") - - project_id = scope_id.project_id - group_query = ( - sa.select(GroupRow) - .where(GroupRow.id == project_id) - .options(load_only(GroupRow.container_registry)) - ) - result = await db_sess.execute(group_query) - group_row = result.scalar_one_or_none() - - if ( - not group_row - or not group_row.container_registry - or "registry" not in group_row.container_registry - or "project" not in group_row.container_registry - ): - raise ContainerRegistryNotFound( - f"Container registry info does not exist or is invalid in the group. (gr: {project_id})" - ) - - registry_name, project = ( - group_row.container_registry["registry"], - group_row.container_registry["project"], - ) - - registry_query = sa.select(ContainerRegistryRow).where( - (ContainerRegistryRow.registry_name == registry_name) - & (ContainerRegistryRow.project == project) - ) - - result = await db_sess.execute(registry_query) - registry = result.scalars().one_or_none() - - if not registry: - raise ContainerRegistryNotFound( - f"Specified container registry row does not exist. (cr: {registry_name}, gr: {project})" - ) - - if registry.type != ContainerRegistryType.HARBOR2: - raise NotImplementedAPI("Only HarborV2 registry is supported for now.") - - ssl_verify = registry.ssl_verify - connector = aiohttp.TCPConnector(ssl=ssl_verify) - async with aiohttp.ClientSession(connector=connector) as sess: - rqst_args: dict[str, Any] = {} - rqst_args["auth"] = aiohttp.BasicAuth( - registry.username, - registry.password, - ) - - api_url = yarl.URL(registry.url) / "api" / "v2.0" - get_project_id_api = api_url / "projects" / project - - async with sess.get(get_project_id_api, allow_redirects=False, **rqst_args) as resp: - res = await resp.json() - harbor_project_id = res["project_id"] - - get_quota_id_api = (api_url / "quotas").with_query({ - "reference": "project", - "reference_id": harbor_project_id, - }) - - async with sess.get(get_quota_id_api, allow_redirects=False, **rqst_args) as resp: - res = await resp.json() - if not res: - raise ObjectNotFound(object_name="quota entity") - if len(res) > 1: - raise InternalServerError( - f"Multiple quota entities found. (project_id: {harbor_project_id})" - ) - - previous_quota = res[0]["hard"]["storage"] - if operation_type == "update" or operation_type == "delete": - if previous_quota == -1: - raise ObjectNotFound(object_name="quota entity") - elif operation_type == "create": - if previous_quota > 0: - raise GenericBadRequest(f"Quota limit already exists. (gr: {project_id})") - - quota_id = res[0]["id"] - - put_quota_api = api_url / "quotas" / str(quota_id) - payload = {"hard": {"storage": quota}} - - async with sess.put( - put_quota_api, json=payload, allow_redirects=False, **rqst_args - ) as resp: - if resp.status == 200: - return - else: - log.error(f"Failed to {operation_type} quota: {await resp.json()}") - raise InternalServerError( - f"Failed to {operation_type} quota. Status code: {resp.status}" - ) - - raise InternalServerError("Unknown error!") - - class CreateQuota(graphene.Mutation): """Added in 24.12.0.""" @@ -222,7 +100,7 @@ async def mutate( ) -> Self: async with info.context.db.begin_session() as db_sess: try: - await mutate_harbor_project_quota("create", db_sess, scope_id, quota) + await handle_harbor_project_quota_operation("create", db_sess, scope_id, quota) return cls(ok=True, msg="success") except Exception as e: return cls(ok=False, msg=str(e)) @@ -253,7 +131,7 @@ async def mutate( ) -> Self: async with info.context.db.begin_session() as db_sess: try: - await mutate_harbor_project_quota("update", db_sess, scope_id, quota) + await handle_harbor_project_quota_operation("update", db_sess, scope_id, quota) return cls(ok=True, msg="success") except Exception as e: return cls(ok=False, msg=str(e)) @@ -282,7 +160,7 @@ async def mutate( ) -> Self: async with info.context.db.begin_session() as db_sess: try: - await mutate_harbor_project_quota("delete", db_sess, scope_id, -1) + await handle_harbor_project_quota_operation("delete", db_sess, scope_id, None) return cls(ok=True, msg="success") except Exception as e: return cls(ok=False, msg=str(e)) diff --git a/src/ai/backend/manager/models/gql_models/container_registry_utils.py b/src/ai/backend/manager/models/gql_models/container_registry_utils.py new file mode 100644 index 0000000000..dfd25b038b --- /dev/null +++ b/src/ai/backend/manager/models/gql_models/container_registry_utils.py @@ -0,0 +1,159 @@ +from __future__ import annotations + +import logging +from typing import Any, Literal, Optional + +import aiohttp +import aiohttp.client_exceptions +import sqlalchemy as sa +import yarl +from sqlalchemy.ext.asyncio import AsyncSession as SASession +from sqlalchemy.orm import load_only + +from ai.backend.logging import BraceStyleAdapter +from ai.backend.manager.api.exceptions import ( + ContainerRegistryNotFound, + GenericBadRequest, + InternalServerError, + NotImplementedAPI, + ObjectNotFound, +) + +from ...container_registry import ContainerRegistryRow +from ..association_container_registries_groups import ( + AssociationContainerRegistriesGroupsRow, +) +from ..group import GroupRow +from ..rbac import ProjectScope, ScopeType + +log = BraceStyleAdapter(logging.getLogger(__spec__.name)) # type: ignore + + +async def handle_harbor_project_quota_operation( + operation_type: Literal["create", "read", "update", "delete"], + db_sess: SASession, + scope_id: ScopeType, + quota: Optional[int], +) -> Optional[int]: + """ + Utility function for code reuse of the HarborV2 per-project Quota CRUD API. + + :param quota: Required for create and delete operations. For all other operations, this parameter should be set to None. + :return: The current quota value for read operations. For other operations, returns None. + """ + if not isinstance(scope_id, ProjectScope): + raise NotImplementedAPI("Quota mutation currently supports only the project scope.") + + if operation_type in ("create", "update"): + assert quota is not None, "Quota value is required for create/update operation." + else: + assert quota is None, "Quota value must be None for read/delete operation." + + project_id = scope_id.project_id + group_query = ( + sa.select(GroupRow) + .where(GroupRow.id == project_id) + .options(load_only(GroupRow.container_registry)) + ) + result = await db_sess.execute(group_query) + group_row = result.scalar_one_or_none() + + if ( + not group_row + or not group_row.container_registry + or "registry" not in group_row.container_registry + or "project" not in group_row.container_registry + ): + raise ContainerRegistryNotFound( + f"Container registry info does not exist or is invalid in the group. (gr: {project_id})" + ) + + registry_name, project = ( + group_row.container_registry["registry"], + group_row.container_registry["project"], + ) + + registry_query = sa.select(ContainerRegistryRow).where( + (ContainerRegistryRow.registry_name == registry_name) + & (ContainerRegistryRow.project == project) + ) + + result = await db_sess.execute(registry_query) + registry = result.scalars().one_or_none() + + if not registry: + raise ContainerRegistryNotFound( + f"Specified container registry row does not exist. (cr: {registry_name}, gr: {project})" + ) + + if operation_type == "read" and not registry.is_global: + get_assoc_query = sa.select( + sa.exists() + .where(AssociationContainerRegistriesGroupsRow.registry_id == registry.id) + .where(AssociationContainerRegistriesGroupsRow.group_id == group_row.row_id) + ) + assoc_exist = (await db_sess.execute(get_assoc_query)).scalar() + + if not assoc_exist: + raise ValueError("The group is not associated with the container registry.") + + ssl_verify = registry.ssl_verify + connector = aiohttp.TCPConnector(ssl=ssl_verify) + async with aiohttp.ClientSession(connector=connector) as sess: + rqst_args: dict[str, Any] = {} + rqst_args["auth"] = aiohttp.BasicAuth( + registry.username, + registry.password, + ) + + api_url = yarl.URL(registry.url) / "api" / "v2.0" + get_project_id_api = api_url / "projects" / project + + async with sess.get(get_project_id_api, allow_redirects=False, **rqst_args) as resp: + res = await resp.json() + harbor_project_id = res["project_id"] + + get_quota_id_api = (api_url / "quotas").with_query({ + "reference": "project", + "reference_id": harbor_project_id, + }) + + async with sess.get(get_quota_id_api, allow_redirects=False, **rqst_args) as resp: + res = await resp.json() + if not res: + raise ObjectNotFound(object_name="quota entity") + if len(res) > 1: + raise InternalServerError( + f"Multiple quota entities found. (project_id: {harbor_project_id})" + ) + + previous_quota = res[0]["hard"]["storage"] + + if operation_type == "create": + if previous_quota > 0: + raise GenericBadRequest(f"Quota limit already exists. (gr: {project_id})") + else: + if previous_quota == -1: + raise ObjectNotFound(object_name="quota entity") + + if operation_type == "read": + return previous_quota + + quota_id = res[0]["id"] + + put_quota_api = api_url / "quotas" / str(quota_id) + quota = quota if operation_type != "delete" else -1 + payload = {"hard": {"storage": quota}} + + async with sess.put( + put_quota_api, json=payload, allow_redirects=False, **rqst_args + ) as resp: + if resp.status == 200: + return None + else: + log.error(f"Failed to {operation_type} quota: {await resp.json()}") + raise InternalServerError( + f"Failed to {operation_type} quota. Status code: {resp.status}" + ) + + raise InternalServerError("Unknown error!") diff --git a/src/ai/backend/manager/models/gql_models/group.py b/src/ai/backend/manager/models/gql_models/group.py index 57de38c8ca..c896aae6c6 100644 --- a/src/ai/backend/manager/models/gql_models/group.py +++ b/src/ai/backend/manager/models/gql_models/group.py @@ -3,30 +3,21 @@ from collections.abc import Mapping from typing import ( TYPE_CHECKING, - Any, Self, Sequence, ) -import aiohttp import graphene import sqlalchemy as sa -import yarl from dateutil.parser import parse as dtparse from graphene.types.datetime import DateTime as GQLDateTime -from ai.backend.manager.api.exceptions import ContainerRegistryNotFound - -from ..association_container_registries_groups import ( - AssociationContainerRegistriesGroupsRow, -) from ..base import ( FilterExprArg, OrderExprArg, PaginatedConnectionField, generate_sql_info_for_gql_connection, ) -from ..container_registry import ContainerRegistryRow, ContainerRegistryType from ..gql_relay import ( AsyncNode, Connection, @@ -35,6 +26,10 @@ from ..group import AssocGroupUserRow, GroupRow, ProjectType from ..minilang.ordering import OrderSpecItem, QueryOrderParser from ..minilang.queryfilter import FieldSpecItem, QueryFilterParser +from ..rbac import ProjectScope +from .container_registry_utils import ( + handle_harbor_project_quota_operation, +) from .user import UserConnection, UserNode if TYPE_CHECKING: @@ -218,76 +213,10 @@ async def resolve_user_nodes( async def resolve_registry_quota(self, info: graphene.ResolveInfo) -> int: graph_ctx = info.context async with graph_ctx.db.begin_session() as db_sess: - if ( - not self.container_registry - or "registry" not in self.container_registry - or "project" not in self.container_registry - ): - raise ContainerRegistryNotFound( - "Container registry info does not exist in the group." - ) - - registry_name, project = ( - self.container_registry["registry"], - self.container_registry["project"], - ) - - registry_query = sa.select(ContainerRegistryRow).where( - (ContainerRegistryRow.registry_name == registry_name) - & (ContainerRegistryRow.project == project) - ) - - result = await db_sess.execute(registry_query) - registry = result.scalars().one_or_none() - - if not registry: - raise ContainerRegistryNotFound("Specified container registry row does not exist.") - if registry.type != ContainerRegistryType.HARBOR2: - raise NotImplementedError("Only HarborV2 registry is supported for now.") - - if not registry.is_global: - get_assoc_query = sa.select( - sa.exists() - .where(AssociationContainerRegistriesGroupsRow.registry_id == registry.id) - .where(AssociationContainerRegistriesGroupsRow.group_id == self.row_id) - ) - assoc_exist = (await db_sess.execute(get_assoc_query)).scalar() - - if not assoc_exist: - raise ValueError("The group is not associated with the container registry.") - - ssl_verify = registry.ssl_verify - connector = aiohttp.TCPConnector(ssl=ssl_verify) - - api_url = yarl.URL(registry.url) / "api" / "v2.0" - async with aiohttp.ClientSession(connector=connector) as sess: - rqst_args: dict[str, Any] = {} - rqst_args["auth"] = aiohttp.BasicAuth( - registry.username, - registry.password, - ) - - get_project_id_api = api_url / "projects" / project - - async with sess.get(get_project_id_api, allow_redirects=False, **rqst_args) as resp: - res = await resp.json() - harbor_project_id = res["project_id"] - - get_quota_id_api = (api_url / "quotas").with_query({ - "reference": "project", - "reference_id": harbor_project_id, - }) - - async with sess.get(get_quota_id_api, allow_redirects=False, **rqst_args) as resp: - res = await resp.json() - if not res: - raise ValueError("Quota not found.") - if len(res) > 1: - raise ValueError("Multiple quotas found.") - - quota = res[0]["hard"]["storage"] - - return quota + scope_id = ProjectScope(project_id=self.id, domain_name=None) + result = await handle_harbor_project_quota_operation("read", db_sess, scope_id, None) + assert result is not None, "Quota value must be returned for read operation." + return result @classmethod async def get_node(cls, info: graphene.ResolveInfo, id) -> Self: From 1a0378990916147bef20f9ef6b60df06047a6b97 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Mon, 18 Nov 2024 04:51:30 +0000 Subject: [PATCH 19/41] feat: Add SDK for registry quota mutations --- src/ai/backend/client/func/group.py | 100 ++++++++++++++++++++++++++++ 1 file changed, 100 insertions(+) diff --git a/src/ai/backend/client/func/group.py b/src/ai/backend/client/func/group.py index 9590ee948d..9975013ddf 100644 --- a/src/ai/backend/client/func/group.py +++ b/src/ai/backend/client/func/group.py @@ -1,6 +1,8 @@ import textwrap from typing import Any, Iterable, Optional, Sequence +from graphql_relay.utils import base64 + from ai.backend.client.output.fields import group_fields from ai.backend.client.output.types import FieldSpec @@ -311,3 +313,101 @@ async def remove_users( } data = await api_session.get().Admin._query(query, variables) return data["modify_group"] + + @api_function + @classmethod + async def get_container_registry_quota(cls, group_id: str) -> int: + """ + Delete Quota Limit for the group's container registry. + Currently, only the HarborV2 registry is supported. + + You need an admin privilege for this operation. + """ + query = textwrap.dedent( + """\ + query($id: String!) { + group_node(id: $id) { + registry_quota + } + } + """ + ) + + variables = {"id": base64(f"group_node:{group_id}")} + data = await api_session.get().Admin._query(query, variables) + return data["group_node"]["registry_quota"] + + @api_function + @classmethod + async def create_container_registry_quota(cls, group_id: str, quota: int) -> dict: + """ + Create Quota Limit for the group's container registry. + Currently, only the HarborV2 registry is supported. + + You need an admin privilege for this operation. + """ + query = textwrap.dedent( + """\ + mutation($scope_id: ScopeField!, $quota: Int!) { + create_container_registry_quota( + scope_id: $scope_id, quota: $quota) { + ok msg + } + } + """ + ) + + scope_id = f"project:{group_id}" + variables = {"scope_id": scope_id, "quota": quota} + data = await api_session.get().Admin._query(query, variables) + return data["create_container_registry_quota"] + + @api_function + @classmethod + async def update_container_registry_quota(cls, group_id: str, quota: int) -> dict: + """ + Update Quota Limit for the group's container registry. + Currently, only the HarborV2 registry is supported. + + You need an admin privilege for this operation. + """ + query = textwrap.dedent( + """\ + mutation($scope_id: ScopeField!, $quota: Int!) { + update_container_registry_quota( + scope_id: $scope_id, quota: $quota) { + ok msg + } + } + """ + ) + + scope_id = f"project:{group_id}" + variables = {"scope_id": scope_id, "quota": quota} + data = await api_session.get().Admin._query(query, variables) + return data["update_container_registry_quota"] + + @api_function + @classmethod + async def delete_container_registry_quota(cls, group_id: str) -> dict: + """ + Delete Quota Limit for the group's container registry. + Currently, only the HarborV2 registry is supported. + + You need an admin privilege for this operation. + """ + query = textwrap.dedent( + """\ + mutation($scope_id: ScopeField!) { + delete_container_registry_quota( + scope_id: $scope_id) { + ok msg + } + } + """ + ) + + scope_id = f"project:{group_id}" + variables = {"scope_id": scope_id} + data = await api_session.get().Admin._query(query, variables) + return data["delete_container_registry_quota"] From c118c7e3a58dd70345b4af99c62e24400178dae1 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Mon, 18 Nov 2024 05:15:15 +0000 Subject: [PATCH 20/41] fix: Broken CI --- src/ai/backend/client/func/group.py | 5 ++--- src/ai/backend/common/utils.py | 9 +++++++++ 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/src/ai/backend/client/func/group.py b/src/ai/backend/client/func/group.py index 9975013ddf..c2f34babd3 100644 --- a/src/ai/backend/client/func/group.py +++ b/src/ai/backend/client/func/group.py @@ -1,10 +1,9 @@ import textwrap from typing import Any, Iterable, Optional, Sequence -from graphql_relay.utils import base64 - from ai.backend.client.output.fields import group_fields from ai.backend.client.output.types import FieldSpec +from ai.backend.common.utils import b64encode from ...cli.types import Undefined, undefined from ..session import api_session @@ -333,7 +332,7 @@ async def get_container_registry_quota(cls, group_id: str) -> int: """ ) - variables = {"id": base64(f"group_node:{group_id}")} + variables = {"id": b64encode(f"group_node:{group_id}")} data = await api_session.get().Admin._query(query, variables) return data["group_node"]["registry_quota"] diff --git a/src/ai/backend/common/utils.py b/src/ai/backend/common/utils.py index 5d98b54d33..c669f58814 100644 --- a/src/ai/backend/common/utils.py +++ b/src/ai/backend/common/utils.py @@ -425,3 +425,12 @@ def join_non_empty(*args, sep): """ filtered_args = [arg for arg in args if arg] return sep.join(filtered_args) + + +def b64encode(s: str) -> str: + """ + base64 encoding method of graphql_relay. + Use it in components where the graphql_relay package is unavailable. + """ + b: bytes = s.encode("utf-8") if isinstance(s, str) else s + return base64.b64encode(b).decode("ascii") From 96fc374237cd52be8c6610a933e5991789170ac4 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Mon, 18 Nov 2024 05:16:42 +0000 Subject: [PATCH 21/41] fix: Wrong object_name in ContainerRegistryNotFound --- src/ai/backend/manager/api/exceptions.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/ai/backend/manager/api/exceptions.py b/src/ai/backend/manager/api/exceptions.py index 01f42af6a7..8db4ee7db9 100644 --- a/src/ai/backend/manager/api/exceptions.py +++ b/src/ai/backend/manager/api/exceptions.py @@ -243,7 +243,7 @@ class EndpointTokenNotFound(ObjectNotFound): class ContainerRegistryNotFound(ObjectNotFound): - object_name = "endpoint_token" + object_name = "container_registry" class TooManySessionsMatched(BackendError, web.HTTPNotFound): From a1f3c0860e2c1a7c97491594fbbce0c55c6a107d Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Mon, 18 Nov 2024 05:58:17 +0000 Subject: [PATCH 22/41] feat: Implement REST API --- src/ai/backend/manager/api/group.py | 119 ++++++++++++++++++++++++++++ src/ai/backend/manager/server.py | 1 + 2 files changed, 120 insertions(+) create mode 100644 src/ai/backend/manager/api/group.py diff --git a/src/ai/backend/manager/api/group.py b/src/ai/backend/manager/api/group.py new file mode 100644 index 0000000000..d59ecd9a63 --- /dev/null +++ b/src/ai/backend/manager/api/group.py @@ -0,0 +1,119 @@ +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Any, Iterable, Tuple + +import aiohttp_cors +import trafaret as t +from aiohttp import web + +from ai.backend.common import validators as tx +from ai.backend.logging import BraceStyleAdapter +from ai.backend.manager.models.gql_models.container_registry_utils import ( + handle_harbor_project_quota_operation, +) +from ai.backend.manager.models.rbac import ProjectScope + +if TYPE_CHECKING: + from .context import RootContext + +from .auth import superadmin_required +from .manager import READ_ALLOWED, server_status_required +from .types import CORSOptions, WebMiddleware +from .utils import check_api_params + +log = BraceStyleAdapter(logging.getLogger(__spec__.name)) + + +@server_status_required(READ_ALLOWED) +@superadmin_required +@check_api_params( + t.Dict({ + tx.AliasedKey(["group_id", "group"]): t.String, + tx.AliasedKey(["quota"]): t.Int, + }) +) +async def update_registry_quota(request: web.Request, params: Any) -> web.Response: + log.info("UPDATE_REGISTRY_QUOTA (gr:{})", params["group_id"]) + root_ctx: RootContext = request.app["_root.context"] + group_id = params["group_id"] + scope_id = ProjectScope(project_id=group_id, domain_name=None) + quota = int(params["quota"]) + + async with root_ctx.db.begin_session() as db_sess: + await handle_harbor_project_quota_operation("update", db_sess, scope_id, quota) + + return web.json_response({}) + + +@server_status_required(READ_ALLOWED) +@superadmin_required +@check_api_params( + t.Dict({ + tx.AliasedKey(["group_id", "group"]): t.String, + }) +) +async def delete_registry_quota(request: web.Request, params: Any) -> web.Response: + log.info("DELETE_REGISTRY_QUOTA (gr:{})", params["group_id"]) + root_ctx: RootContext = request.app["_root.context"] + group_id = params["group_id"] + scope_id = ProjectScope(project_id=group_id, domain_name=None) + + async with root_ctx.db.begin_session() as db_sess: + await handle_harbor_project_quota_operation("delete", db_sess, scope_id, None) + + return web.json_response({}) + + +@server_status_required(READ_ALLOWED) +@superadmin_required +@check_api_params( + t.Dict({ + tx.AliasedKey(["group_id", "group"]): t.String, + tx.AliasedKey(["quota"]): t.Int, + }) +) +async def create_registry_quota(request: web.Request, params: Any) -> web.Response: + log.info("CREATE_REGISTRY_QUOTA (gr:{})", params["group_id"]) + root_ctx: RootContext = request.app["_root.context"] + group_id = params["group_id"] + scope_id = ProjectScope(project_id=group_id, domain_name=None) + quota = int(params["quota"]) + + async with root_ctx.db.begin_session() as db_sess: + await handle_harbor_project_quota_operation("create", db_sess, scope_id, quota) + + return web.json_response({}) + + +@server_status_required(READ_ALLOWED) +@superadmin_required +@check_api_params( + t.Dict({ + tx.AliasedKey(["group_id", "group"]): t.String, + }) +) +async def read_registry_quota(request: web.Request, params: Any) -> web.Response: + log.info("READ_REGISTRY_QUOTA (gr:{})", params["group_id"]) + root_ctx: RootContext = request.app["_root.context"] + group_id = params["group_id"] + scope_id = ProjectScope(project_id=group_id, domain_name=None) + + async with root_ctx.db.begin_session() as db_sess: + quota = await handle_harbor_project_quota_operation("read", db_sess, scope_id, None) + + return web.json_response({"result": quota}) + + +def create_app( + default_cors_options: CORSOptions, +) -> Tuple[web.Application, Iterable[WebMiddleware]]: + app = web.Application() + app["api_versions"] = (1, 2, 3, 4, 5) + app["prefix"] = "group" + cors = aiohttp_cors.setup(app, defaults=default_cors_options) + cors.add(app.router.add_route("POST", "/registry-quota", create_registry_quota)) + cors.add(app.router.add_route("GET", "/registry-quota", read_registry_quota)) + cors.add(app.router.add_route("PATCH", "/registry-quota", update_registry_quota)) + cors.add(app.router.add_route("DELETE", "/registry-quota", delete_registry_quota)) + return app, [] diff --git a/src/ai/backend/manager/server.py b/src/ai/backend/manager/server.py index fc85bc9c95..635b7831dd 100644 --- a/src/ai/backend/manager/server.py +++ b/src/ai/backend/manager/server.py @@ -187,6 +187,7 @@ ".image", ".userconfig", ".domainconfig", + ".group", ".groupconfig", ".logs", ] From 217a359c358bd600f0b8d4877b9dedf26f45671f Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Mon, 18 Nov 2024 05:58:54 +0000 Subject: [PATCH 23/41] fix: Wrong exception handling --- .../manager/models/gql_models/container_registry_utils.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/src/ai/backend/manager/models/gql_models/container_registry_utils.py b/src/ai/backend/manager/models/gql_models/container_registry_utils.py index dfd25b038b..bd5ed7dd6a 100644 --- a/src/ai/backend/manager/models/gql_models/container_registry_utils.py +++ b/src/ai/backend/manager/models/gql_models/container_registry_utils.py @@ -151,9 +151,7 @@ async def handle_harbor_project_quota_operation( if resp.status == 200: return None else: - log.error(f"Failed to {operation_type} quota: {await resp.json()}") - raise InternalServerError( - f"Failed to {operation_type} quota. Status code: {resp.status}" - ) + log.error(f"Failed to {operation_type} quota! response: {resp}") + raise InternalServerError(f"Failed to {operation_type} quota! response: {resp}") raise InternalServerError("Unknown error!") From 89ab7a7683afb096d4f4e5df6f1a35543f7bb3b4 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Tue, 19 Nov 2024 01:25:01 +0000 Subject: [PATCH 24/41] chore: Update comment --- src/ai/backend/client/func/group.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/ai/backend/client/func/group.py b/src/ai/backend/client/func/group.py index c2f34babd3..03a0d6fdb9 100644 --- a/src/ai/backend/client/func/group.py +++ b/src/ai/backend/client/func/group.py @@ -318,7 +318,7 @@ async def remove_users( async def get_container_registry_quota(cls, group_id: str) -> int: """ Delete Quota Limit for the group's container registry. - Currently, only the HarborV2 registry is supported. + Currently only HarborV2 registry is supported. You need an admin privilege for this operation. """ @@ -341,7 +341,7 @@ async def get_container_registry_quota(cls, group_id: str) -> int: async def create_container_registry_quota(cls, group_id: str, quota: int) -> dict: """ Create Quota Limit for the group's container registry. - Currently, only the HarborV2 registry is supported. + Currently only HarborV2 registry is supported. You need an admin privilege for this operation. """ @@ -366,7 +366,7 @@ async def create_container_registry_quota(cls, group_id: str, quota: int) -> dic async def update_container_registry_quota(cls, group_id: str, quota: int) -> dict: """ Update Quota Limit for the group's container registry. - Currently, only the HarborV2 registry is supported. + Currently only HarborV2 registry is supported. You need an admin privilege for this operation. """ @@ -391,7 +391,7 @@ async def update_container_registry_quota(cls, group_id: str, quota: int) -> dic async def delete_container_registry_quota(cls, group_id: str) -> dict: """ Delete Quota Limit for the group's container registry. - Currently, only the HarborV2 registry is supported. + Currently only HarborV2 registry is supported. You need an admin privilege for this operation. """ From 896c44ddd5ef04565b81d3e68e116fbc4c991561 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Tue, 19 Nov 2024 01:40:08 +0000 Subject: [PATCH 25/41] chore: Rename types --- src/ai/backend/manager/models/gql.py | 18 ++++++++++++------ .../models/gql_models/container_registry.py | 6 +++--- .../gql_models/container_registry_utils.py | 2 +- 3 files changed, 16 insertions(+), 10 deletions(-) diff --git a/src/ai/backend/manager/models/gql.py b/src/ai/backend/manager/models/gql.py index 21d52c31bb..9facab1268 100644 --- a/src/ai/backend/manager/models/gql.py +++ b/src/ai/backend/manager/models/gql.py @@ -74,10 +74,10 @@ ) from .gql_models.container_registry import ( AssociateContainerRegistryWithGroup, - CreateQuota, - DeleteQuota, + CreateContainerRegistryQuota, + DeleteContainerRegistryQuota, DisassociateContainerRegistryWithGroup, - UpdateQuota, + UpdateContainerRegistryQuota, ) from .gql_models.domain import ( CreateDomainNode, @@ -352,9 +352,15 @@ class Mutations(graphene.ObjectType): description="Added in 24.12.0" ) - create_container_registry_quota = CreateQuota.Field(description="Added in 24.12.0") - update_container_registry_quota = UpdateQuota.Field(description="Added in 24.12.0") - delete_container_registry_quota = DeleteQuota.Field(description="Added in 24.12.0") + create_container_registry_quota = CreateContainerRegistryQuota.Field( + description="Added in 24.12.0" + ) + update_container_registry_quota = UpdateContainerRegistryQuota.Field( + description="Added in 24.12.0" + ) + delete_container_registry_quota = DeleteContainerRegistryQuota.Field( + description="Added in 24.12.0" + ) # Legacy mutations create_container_registry = CreateContainerRegistry.Field() diff --git a/src/ai/backend/manager/models/gql_models/container_registry.py b/src/ai/backend/manager/models/gql_models/container_registry.py index b78939f041..21589c97d5 100644 --- a/src/ai/backend/manager/models/gql_models/container_registry.py +++ b/src/ai/backend/manager/models/gql_models/container_registry.py @@ -75,7 +75,7 @@ async def mutate( return await simple_db_mutate(cls, info.context, delete_query) -class CreateQuota(graphene.Mutation): +class CreateContainerRegistryQuota(graphene.Mutation): """Added in 24.12.0.""" allowed_roles = ( @@ -106,7 +106,7 @@ async def mutate( return cls(ok=False, msg=str(e)) -class UpdateQuota(graphene.Mutation): +class UpdateContainerRegistryQuota(graphene.Mutation): """Added in 24.12.0.""" allowed_roles = ( @@ -137,7 +137,7 @@ async def mutate( return cls(ok=False, msg=str(e)) -class DeleteQuota(graphene.Mutation): +class DeleteContainerRegistryQuota(graphene.Mutation): """Added in 24.12.0.""" allowed_roles = ( diff --git a/src/ai/backend/manager/models/gql_models/container_registry_utils.py b/src/ai/backend/manager/models/gql_models/container_registry_utils.py index bd5ed7dd6a..f8c1a10d58 100644 --- a/src/ai/backend/manager/models/gql_models/container_registry_utils.py +++ b/src/ai/backend/manager/models/gql_models/container_registry_utils.py @@ -38,7 +38,7 @@ async def handle_harbor_project_quota_operation( """ Utility function for code reuse of the HarborV2 per-project Quota CRUD API. - :param quota: Required for create and delete operations. For all other operations, this parameter should be set to None. + :param quota: Required for create, delete operations. For other operations, quota should be set to None. :return: The current quota value for read operations. For other operations, returns None. """ if not isinstance(scope_id, ProjectScope): From 835d0c687d4fcfaeba8d7441a3f9f16b1c66f0f4 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Tue, 19 Nov 2024 01:42:27 +0000 Subject: [PATCH 26/41] chore: update GraphQL schema dump Co-authored-by: octodog --- src/ai/backend/manager/api/schema.graphql | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/src/ai/backend/manager/api/schema.graphql b/src/ai/backend/manager/api/schema.graphql index f67922d130..cfbce507f1 100644 --- a/src/ai/backend/manager/api/schema.graphql +++ b/src/ai/backend/manager/api/schema.graphql @@ -1849,13 +1849,13 @@ type Mutations { disassociate_container_registry_with_group(group_id: String!, registry_id: String!): DisassociateContainerRegistryWithGroup """Added in 24.12.0""" - create_container_registry_quota(quota: Int!, scope_id: ScopeField!): CreateQuota + create_container_registry_quota(quota: Int!, scope_id: ScopeField!): CreateContainerRegistryQuota """Added in 24.12.0""" - update_container_registry_quota(quota: Int!, scope_id: ScopeField!): UpdateQuota + update_container_registry_quota(quota: Int!, scope_id: ScopeField!): UpdateContainerRegistryQuota """Added in 24.12.0""" - delete_container_registry_quota(scope_id: ScopeField!): DeleteQuota + delete_container_registry_quota(scope_id: ScopeField!): DeleteContainerRegistryQuota create_container_registry(hostname: String!, props: CreateContainerRegistryInput!): CreateContainerRegistry modify_container_registry(hostname: String!, props: ModifyContainerRegistryInput!): ModifyContainerRegistry delete_container_registry(hostname: String!): DeleteContainerRegistry @@ -2583,19 +2583,19 @@ type DisassociateContainerRegistryWithGroup { } """Added in 24.12.0.""" -type CreateQuota { +type CreateContainerRegistryQuota { ok: Boolean msg: String } """Added in 24.12.0.""" -type UpdateQuota { +type UpdateContainerRegistryQuota { ok: Boolean msg: String } """Added in 24.12.0.""" -type DeleteQuota { +type DeleteContainerRegistryQuota { ok: Boolean msg: String } From 0a442c9bf853e14293e749bbd070402f452687be Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Tue, 19 Nov 2024 01:47:44 +0000 Subject: [PATCH 27/41] chore: Rename news fragment --- changes/3090.feature.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/changes/3090.feature.md b/changes/3090.feature.md index 76acfee344..e725ca5ff6 100644 --- a/changes/3090.feature.md +++ b/changes/3090.feature.md @@ -1 +1 @@ -Implement management API for controlling Harbor per-project Quota. +Implement CRUD API for managing Harbor per-project Quota. From 049dfe0185f03e95c1d6c7083820774b9bc90702 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Wed, 20 Nov 2024 05:34:27 +0000 Subject: [PATCH 28/41] fix: Use `BigInt` --- .../models/gql_models/container_registry.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/src/ai/backend/manager/models/gql_models/container_registry.py b/src/ai/backend/manager/models/gql_models/container_registry.py index 21589c97d5..5ecf68e9da 100644 --- a/src/ai/backend/manager/models/gql_models/container_registry.py +++ b/src/ai/backend/manager/models/gql_models/container_registry.py @@ -11,7 +11,7 @@ from ..association_container_registries_groups import ( AssociationContainerRegistriesGroupsRow, ) -from ..base import simple_db_mutate +from ..base import BigInt, simple_db_mutate from ..rbac import ScopeType from ..user import UserRole from .container_registry_utils import handle_harbor_project_quota_operation @@ -85,7 +85,7 @@ class CreateContainerRegistryQuota(graphene.Mutation): class Arguments: scope_id = ScopeField(required=True) - quota = graphene.Int(required=True) + quota = BigInt(required=True) ok = graphene.Boolean() msg = graphene.String() @@ -96,11 +96,11 @@ async def mutate( root, info: graphene.ResolveInfo, scope_id: ScopeType, - quota: int, + quota: int | float, ) -> Self: async with info.context.db.begin_session() as db_sess: try: - await handle_harbor_project_quota_operation("create", db_sess, scope_id, quota) + await handle_harbor_project_quota_operation("create", db_sess, scope_id, int(quota)) return cls(ok=True, msg="success") except Exception as e: return cls(ok=False, msg=str(e)) @@ -116,7 +116,7 @@ class UpdateContainerRegistryQuota(graphene.Mutation): class Arguments: scope_id = ScopeField(required=True) - quota = graphene.Int(required=True) + quota = BigInt(required=True) ok = graphene.Boolean() msg = graphene.String() @@ -127,11 +127,11 @@ async def mutate( root, info: graphene.ResolveInfo, scope_id: ScopeType, - quota: int, + quota: int | float, ) -> Self: async with info.context.db.begin_session() as db_sess: try: - await handle_harbor_project_quota_operation("update", db_sess, scope_id, quota) + await handle_harbor_project_quota_operation("update", db_sess, scope_id, int(quota)) return cls(ok=True, msg="success") except Exception as e: return cls(ok=False, msg=str(e)) From cae3feeb34634db7ba8eaf673bf98c2e726c584d Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Wed, 20 Nov 2024 05:36:56 +0000 Subject: [PATCH 29/41] chore: update GraphQL schema dump Co-authored-by: octodog --- src/ai/backend/manager/api/schema.graphql | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/ai/backend/manager/api/schema.graphql b/src/ai/backend/manager/api/schema.graphql index cfbce507f1..ff92f3b31c 100644 --- a/src/ai/backend/manager/api/schema.graphql +++ b/src/ai/backend/manager/api/schema.graphql @@ -1849,10 +1849,10 @@ type Mutations { disassociate_container_registry_with_group(group_id: String!, registry_id: String!): DisassociateContainerRegistryWithGroup """Added in 24.12.0""" - create_container_registry_quota(quota: Int!, scope_id: ScopeField!): CreateContainerRegistryQuota + create_container_registry_quota(quota: BigInt!, scope_id: ScopeField!): CreateContainerRegistryQuota """Added in 24.12.0""" - update_container_registry_quota(quota: Int!, scope_id: ScopeField!): UpdateContainerRegistryQuota + update_container_registry_quota(quota: BigInt!, scope_id: ScopeField!): UpdateContainerRegistryQuota """Added in 24.12.0""" delete_container_registry_quota(scope_id: ScopeField!): DeleteContainerRegistryQuota From 3d1c9cfbc5d05957164679287cff5f3866b143de Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Thu, 28 Nov 2024 03:16:05 +0000 Subject: [PATCH 30/41] refactor: Add `HarborQuotaManager` *(Reflect feedback) --- src/ai/backend/client/func/group.py | 2 +- src/ai/backend/manager/api/group.py | 14 +- .../models/gql_models/container_registry.py | 11 +- .../gql_models/container_registry_utils.py | 280 ++++++++++++------ .../manager/models/gql_models/group.py | 7 +- 5 files changed, 207 insertions(+), 107 deletions(-) diff --git a/src/ai/backend/client/func/group.py b/src/ai/backend/client/func/group.py index 03a0d6fdb9..c8fa49254c 100644 --- a/src/ai/backend/client/func/group.py +++ b/src/ai/backend/client/func/group.py @@ -317,7 +317,7 @@ async def remove_users( @classmethod async def get_container_registry_quota(cls, group_id: str) -> int: """ - Delete Quota Limit for the group's container registry. + Get Quota Limit for the group's container registry. Currently only HarborV2 registry is supported. You need an admin privilege for this operation. diff --git a/src/ai/backend/manager/api/group.py b/src/ai/backend/manager/api/group.py index d59ecd9a63..5434d6a88d 100644 --- a/src/ai/backend/manager/api/group.py +++ b/src/ai/backend/manager/api/group.py @@ -10,7 +10,7 @@ from ai.backend.common import validators as tx from ai.backend.logging import BraceStyleAdapter from ai.backend.manager.models.gql_models.container_registry_utils import ( - handle_harbor_project_quota_operation, + HarborQuotaManager, ) from ai.backend.manager.models.rbac import ProjectScope @@ -41,7 +41,8 @@ async def update_registry_quota(request: web.Request, params: Any) -> web.Respon quota = int(params["quota"]) async with root_ctx.db.begin_session() as db_sess: - await handle_harbor_project_quota_operation("update", db_sess, scope_id, quota) + manager = await HarborQuotaManager.new(db_sess, scope_id) + await manager.update(quota) return web.json_response({}) @@ -60,7 +61,8 @@ async def delete_registry_quota(request: web.Request, params: Any) -> web.Respon scope_id = ProjectScope(project_id=group_id, domain_name=None) async with root_ctx.db.begin_session() as db_sess: - await handle_harbor_project_quota_operation("delete", db_sess, scope_id, None) + manager = await HarborQuotaManager.new(db_sess, scope_id) + await manager.delete() return web.json_response({}) @@ -81,7 +83,8 @@ async def create_registry_quota(request: web.Request, params: Any) -> web.Respon quota = int(params["quota"]) async with root_ctx.db.begin_session() as db_sess: - await handle_harbor_project_quota_operation("create", db_sess, scope_id, quota) + manager = await HarborQuotaManager.new(db_sess, scope_id) + await manager.create(quota) return web.json_response({}) @@ -100,7 +103,8 @@ async def read_registry_quota(request: web.Request, params: Any) -> web.Response scope_id = ProjectScope(project_id=group_id, domain_name=None) async with root_ctx.db.begin_session() as db_sess: - quota = await handle_harbor_project_quota_operation("read", db_sess, scope_id, None) + manager = await HarborQuotaManager.new(db_sess, scope_id) + quota = await manager.read() return web.json_response({"result": quota}) diff --git a/src/ai/backend/manager/models/gql_models/container_registry.py b/src/ai/backend/manager/models/gql_models/container_registry.py index 5ecf68e9da..72010c565f 100644 --- a/src/ai/backend/manager/models/gql_models/container_registry.py +++ b/src/ai/backend/manager/models/gql_models/container_registry.py @@ -14,7 +14,7 @@ from ..base import BigInt, simple_db_mutate from ..rbac import ScopeType from ..user import UserRole -from .container_registry_utils import handle_harbor_project_quota_operation +from .container_registry_utils import HarborQuotaManager from .fields import ScopeField log = BraceStyleAdapter(logging.getLogger(__spec__.name)) # type: ignore @@ -100,7 +100,8 @@ async def mutate( ) -> Self: async with info.context.db.begin_session() as db_sess: try: - await handle_harbor_project_quota_operation("create", db_sess, scope_id, int(quota)) + manager = await HarborQuotaManager.new(db_sess, scope_id) + await manager.create(int(quota)) return cls(ok=True, msg="success") except Exception as e: return cls(ok=False, msg=str(e)) @@ -131,7 +132,8 @@ async def mutate( ) -> Self: async with info.context.db.begin_session() as db_sess: try: - await handle_harbor_project_quota_operation("update", db_sess, scope_id, int(quota)) + manager = await HarborQuotaManager.new(db_sess, scope_id) + await manager.update(int(quota)) return cls(ok=True, msg="success") except Exception as e: return cls(ok=False, msg=str(e)) @@ -160,7 +162,8 @@ async def mutate( ) -> Self: async with info.context.db.begin_session() as db_sess: try: - await handle_harbor_project_quota_operation("delete", db_sess, scope_id, None) + manager = await HarborQuotaManager.new(db_sess, scope_id) + await manager.delete() return cls(ok=True, msg="success") except Exception as e: return cls(ok=False, msg=str(e)) diff --git a/src/ai/backend/manager/models/gql_models/container_registry_utils.py b/src/ai/backend/manager/models/gql_models/container_registry_utils.py index f8c1a10d58..60611f5b6f 100644 --- a/src/ai/backend/manager/models/gql_models/container_registry_utils.py +++ b/src/ai/backend/manager/models/gql_models/container_registry_utils.py @@ -1,7 +1,8 @@ from __future__ import annotations import logging -from typing import Any, Literal, Optional +import uuid +from typing import Any, TypedDict import aiohttp import aiohttp.client_exceptions @@ -10,6 +11,7 @@ from sqlalchemy.ext.asyncio import AsyncSession as SASession from sqlalchemy.orm import load_only +from ai.backend.common.types import aobject from ai.backend.logging import BraceStyleAdapter from ai.backend.manager.api.exceptions import ( ContainerRegistryNotFound, @@ -29,94 +31,99 @@ log = BraceStyleAdapter(logging.getLogger(__spec__.name)) # type: ignore -async def handle_harbor_project_quota_operation( - operation_type: Literal["create", "read", "update", "delete"], - db_sess: SASession, - scope_id: ScopeType, - quota: Optional[int], -) -> Optional[int]: - """ - Utility function for code reuse of the HarborV2 per-project Quota CRUD API. +class HarborQuotaInfo(TypedDict): + previous_quota: int + quota_id: int + - :param quota: Required for create, delete operations. For other operations, quota should be set to None. - :return: The current quota value for read operations. For other operations, returns None. +class HarborQuotaManager(aobject): """ - if not isinstance(scope_id, ProjectScope): - raise NotImplementedAPI("Quota mutation currently supports only the project scope.") - - if operation_type in ("create", "update"): - assert quota is not None, "Quota value is required for create/update operation." - else: - assert quota is None, "Quota value must be None for read/delete operation." - - project_id = scope_id.project_id - group_query = ( - sa.select(GroupRow) - .where(GroupRow.id == project_id) - .options(load_only(GroupRow.container_registry)) - ) - result = await db_sess.execute(group_query) - group_row = result.scalar_one_or_none() - - if ( - not group_row - or not group_row.container_registry - or "registry" not in group_row.container_registry - or "project" not in group_row.container_registry - ): - raise ContainerRegistryNotFound( - f"Container registry info does not exist or is invalid in the group. (gr: {project_id})" - ) + Utility class for HarborV2 per-project Quota CRUD API. + """ + + db_sess: SASession + scope_id: ScopeType + group_row: GroupRow + registry: ContainerRegistryRow + project: str + project_id: uuid.UUID - registry_name, project = ( - group_row.container_registry["registry"], - group_row.container_registry["project"], - ) + def __init__(self, db_sess: SASession, scope_id: ScopeType): + if not isinstance(scope_id, ProjectScope): + raise NotImplementedAPI("Quota mutation currently supports only the project scope.") - registry_query = sa.select(ContainerRegistryRow).where( - (ContainerRegistryRow.registry_name == registry_name) - & (ContainerRegistryRow.project == project) - ) + self.db_sess = db_sess + self.scope_id = scope_id - result = await db_sess.execute(registry_query) - registry = result.scalars().one_or_none() + async def __ainit__(self) -> None: + assert isinstance(self.scope_id, ProjectScope) - if not registry: - raise ContainerRegistryNotFound( - f"Specified container registry row does not exist. (cr: {registry_name}, gr: {project})" + project_id = self.scope_id.project_id + group_query = ( + sa.select(GroupRow) + .where(GroupRow.id == project_id) + .options(load_only(GroupRow.container_registry)) ) + result = await self.db_sess.execute(group_query) + group_row = result.scalar_one_or_none() - if operation_type == "read" and not registry.is_global: - get_assoc_query = sa.select( - sa.exists() - .where(AssociationContainerRegistriesGroupsRow.registry_id == registry.id) - .where(AssociationContainerRegistriesGroupsRow.group_id == group_row.row_id) + if not HarborQuotaManager._is_valid_group_row(group_row): + raise ContainerRegistryNotFound( + f"Container registry info does not exist or is invalid in the group. (gr: {project_id})" + ) + + registry_name, project = ( + group_row.container_registry["registry"], + group_row.container_registry["project"], ) - assoc_exist = (await db_sess.execute(get_assoc_query)).scalar() - - if not assoc_exist: - raise ValueError("The group is not associated with the container registry.") - - ssl_verify = registry.ssl_verify - connector = aiohttp.TCPConnector(ssl=ssl_verify) - async with aiohttp.ClientSession(connector=connector) as sess: - rqst_args: dict[str, Any] = {} - rqst_args["auth"] = aiohttp.BasicAuth( - registry.username, - registry.password, + + registry_query = sa.select(ContainerRegistryRow).where( + (ContainerRegistryRow.registry_name == registry_name) + & (ContainerRegistryRow.project == project) ) - api_url = yarl.URL(registry.url) / "api" / "v2.0" - get_project_id_api = api_url / "projects" / project + result = await self.db_sess.execute(registry_query) + registry = result.scalars().one_or_none() + + if not registry: + raise ContainerRegistryNotFound( + f"Specified container registry row not found. (cr: {registry_name}, gr: {project})" + ) + + self.group_row = group_row + self.registry = registry + self.project = project + self.project_id = project_id + + @classmethod + def _is_valid_group_row(self, group_row: GroupRow) -> bool: + return ( + group_row + and group_row.container_registry + and "registry" in group_row.container_registry + and "project" in group_row.container_registry + ) + + async def _get_harbor_project_id( + self, sess: aiohttp.ClientSession, rqst_args: dict[str, Any] + ) -> str: + get_project_id_api = ( + yarl.URL(self.registry.url) / "api" / "v2.0" / "projects" / self.project + ) async with sess.get(get_project_id_api, allow_redirects=False, **rqst_args) as resp: res = await resp.json() harbor_project_id = res["project_id"] + return harbor_project_id - get_quota_id_api = (api_url / "quotas").with_query({ - "reference": "project", - "reference_id": harbor_project_id, - }) + async def _get_quota_info( + self, sess: aiohttp.ClientSession, rqst_args: dict[str, Any] + ) -> HarborQuotaInfo: + harbor_project_id = await self._get_harbor_project_id(sess, rqst_args) + get_quota_id_api = (yarl.URL(self.registry.url) / "api" / "v2.0" / "quotas").with_query({ + "reference": "project", + "reference_id": harbor_project_id, + }) async with sess.get(get_quota_id_api, allow_redirects=False, **rqst_args) as resp: res = await resp.json() @@ -128,30 +135,117 @@ async def handle_harbor_project_quota_operation( ) previous_quota = res[0]["hard"]["storage"] + quota_id = res[0]["id"] - if operation_type == "create": - if previous_quota > 0: - raise GenericBadRequest(f"Quota limit already exists. (gr: {project_id})") - else: - if previous_quota == -1: - raise ObjectNotFound(object_name="quota entity") + return HarborQuotaInfo(previous_quota=previous_quota, quota_id=quota_id) + + async def read(self) -> int: + if not self.registry.is_global: + get_assoc_query = sa.select( + sa.exists() + .where(AssociationContainerRegistriesGroupsRow.registry_id == self.registry.id) + .where(AssociationContainerRegistriesGroupsRow.group_id == self.group_row.row_id) + ) + assoc_exist = (await self.db_sess.execute(get_assoc_query)).scalar() + + if not assoc_exist: + raise ValueError("The group is not associated with the container registry.") + + ssl_verify = self.registry.ssl_verify + connector = aiohttp.TCPConnector(ssl=ssl_verify) + async with aiohttp.ClientSession(connector=connector) as sess: + rqst_args: dict[str, Any] = {} + rqst_args["auth"] = aiohttp.BasicAuth( + self.registry.username, + self.registry.password, + ) + + previous_quota = (await self._get_quota_info(sess, rqst_args))["previous_quota"] + if previous_quota == -1: + raise ObjectNotFound(object_name="quota entity") - if operation_type == "read": - return previous_quota + return previous_quota - quota_id = res[0]["id"] + async def create(self, quota: int) -> None: + ssl_verify = self.registry.ssl_verify + connector = aiohttp.TCPConnector(ssl=ssl_verify) + async with aiohttp.ClientSession(connector=connector) as sess: + rqst_args: dict[str, Any] = {} + rqst_args["auth"] = aiohttp.BasicAuth( + self.registry.username, + self.registry.password, + ) + + quota_info = await self._get_quota_info(sess, rqst_args) + previous_quota, quota_id = quota_info["previous_quota"], quota_info["quota_id"] + + if previous_quota > 0: + raise GenericBadRequest(f"Quota limit already exists. (gr: {self.project_id})") - put_quota_api = api_url / "quotas" / str(quota_id) - quota = quota if operation_type != "delete" else -1 + put_quota_api = yarl.URL(self.registry.url) / "api" / "v2.0" / "quotas" / str(quota_id) payload = {"hard": {"storage": quota}} - async with sess.put( - put_quota_api, json=payload, allow_redirects=False, **rqst_args - ) as resp: - if resp.status == 200: - return None - else: - log.error(f"Failed to {operation_type} quota! response: {resp}") - raise InternalServerError(f"Failed to {operation_type} quota! response: {resp}") + async with sess.put( + put_quota_api, json=payload, allow_redirects=False, **rqst_args + ) as resp: + if resp.status == 200: + return None + else: + log.error(f"Failed to create quota! response: {resp}") + raise InternalServerError(f"Failed to create quota! response: {resp}") + + async def update(self, quota: int) -> None: + ssl_verify = self.registry.ssl_verify + connector = aiohttp.TCPConnector(ssl=ssl_verify) + async with aiohttp.ClientSession(connector=connector) as sess: + rqst_args: dict[str, Any] = {} + rqst_args["auth"] = aiohttp.BasicAuth( + self.registry.username, + self.registry.password, + ) + + quota_info = await self._get_quota_info(sess, rqst_args) + previous_quota, quota_id = quota_info["previous_quota"], quota_info["quota_id"] + + if previous_quota == -1: + raise ObjectNotFound(object_name="quota entity") + + put_quota_api = yarl.URL(self.registry.url) / "api" / "v2.0" / "quotas" / str(quota_id) + payload = {"hard": {"storage": quota}} + + async with sess.put( + put_quota_api, json=payload, allow_redirects=False, **rqst_args + ) as resp: + if resp.status == 200: + return None + else: + log.error(f"Failed to update quota! response: {resp}") + raise InternalServerError(f"Failed to update quota! response: {resp}") + + async def delete(self) -> None: + ssl_verify = self.registry.ssl_verify + connector = aiohttp.TCPConnector(ssl=ssl_verify) + async with aiohttp.ClientSession(connector=connector) as sess: + rqst_args: dict[str, Any] = {} + rqst_args["auth"] = aiohttp.BasicAuth( + self.registry.username, + self.registry.password, + ) + + quota_info = await self._get_quota_info(sess, rqst_args) + previous_quota, quota_id = quota_info["previous_quota"], quota_info["quota_id"] + + if previous_quota == -1: + raise ObjectNotFound(object_name="quota entity") - raise InternalServerError("Unknown error!") + put_quota_api = yarl.URL(self.registry.url) / "api" / "v2.0" / "quotas" / str(quota_id) + payload = {"hard": {"storage": -1}} # setting quota to -1 means delete + + async with sess.put( + put_quota_api, json=payload, allow_redirects=False, **rqst_args + ) as resp: + if resp.status == 200: + return None + else: + log.error(f"Failed to delete quota! response: {resp}") + raise InternalServerError(f"Failed to delete quota! response: {resp}") diff --git a/src/ai/backend/manager/models/gql_models/group.py b/src/ai/backend/manager/models/gql_models/group.py index c896aae6c6..8b5bcb54e1 100644 --- a/src/ai/backend/manager/models/gql_models/group.py +++ b/src/ai/backend/manager/models/gql_models/group.py @@ -28,7 +28,7 @@ from ..minilang.queryfilter import FieldSpecItem, QueryFilterParser from ..rbac import ProjectScope from .container_registry_utils import ( - handle_harbor_project_quota_operation, + HarborQuotaManager, ) from .user import UserConnection, UserNode @@ -214,9 +214,8 @@ async def resolve_registry_quota(self, info: graphene.ResolveInfo) -> int: graph_ctx = info.context async with graph_ctx.db.begin_session() as db_sess: scope_id = ProjectScope(project_id=self.id, domain_name=None) - result = await handle_harbor_project_quota_operation("read", db_sess, scope_id, None) - assert result is not None, "Quota value must be returned for read operation." - return result + manager = await HarborQuotaManager.new(db_sess, scope_id) + return await manager.read() @classmethod async def get_node(cls, info: graphene.ResolveInfo, id) -> Self: From cb68069a5b1c34d5bc57bc7a5df91e7b83143908 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Thu, 28 Nov 2024 03:23:52 +0000 Subject: [PATCH 31/41] fix: Use BigInt --- src/ai/backend/manager/models/gql_models/group.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/ai/backend/manager/models/gql_models/group.py b/src/ai/backend/manager/models/gql_models/group.py index 8b5bcb54e1..afc58eb4ac 100644 --- a/src/ai/backend/manager/models/gql_models/group.py +++ b/src/ai/backend/manager/models/gql_models/group.py @@ -13,6 +13,7 @@ from graphene.types.datetime import DateTime as GQLDateTime from ..base import ( + BigInt, FilterExprArg, OrderExprArg, PaginatedConnectionField, @@ -116,7 +117,7 @@ class Meta: lambda: graphene.String, ) - registry_quota = graphene.Int(description="Added in 24.12.0.") + registry_quota = BigInt(description="Added in 24.12.0.") user_nodes = PaginatedConnectionField( UserConnection, From 1ec6c5098fdb735d4af07e1662ebf06aeea84aa3 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Thu, 28 Nov 2024 03:25:37 +0000 Subject: [PATCH 32/41] chore: self -> cls --- .../manager/models/gql_models/container_registry_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/ai/backend/manager/models/gql_models/container_registry_utils.py b/src/ai/backend/manager/models/gql_models/container_registry_utils.py index 60611f5b6f..7a241c9d8c 100644 --- a/src/ai/backend/manager/models/gql_models/container_registry_utils.py +++ b/src/ai/backend/manager/models/gql_models/container_registry_utils.py @@ -96,7 +96,7 @@ async def __ainit__(self) -> None: self.project_id = project_id @classmethod - def _is_valid_group_row(self, group_row: GroupRow) -> bool: + def _is_valid_group_row(cls, group_row: GroupRow) -> bool: return ( group_row and group_row.container_registry From 7ddcad72e68d4b94d4b4bc422022c60fc6014226 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Thu, 28 Nov 2024 03:28:00 +0000 Subject: [PATCH 33/41] chore: update GraphQL schema dump Co-authored-by: octodog --- src/ai/backend/manager/api/schema.graphql | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/ai/backend/manager/api/schema.graphql b/src/ai/backend/manager/api/schema.graphql index ff92f3b31c..dcfebc74c6 100644 --- a/src/ai/backend/manager/api/schema.graphql +++ b/src/ai/backend/manager/api/schema.graphql @@ -700,7 +700,7 @@ type GroupNode implements Node { scaling_groups: [String] """Added in 24.12.0.""" - registry_quota: Int + registry_quota: BigInt user_nodes(filter: String, order: String, offset: Int, before: String, after: String, first: Int, last: Int): UserConnection } From 90eaff7e1ce7419c65597d3527bf9b915de48301 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Thu, 28 Nov 2024 03:34:33 +0000 Subject: [PATCH 34/41] fix: Improve exception handling --- .../gql_models/container_registry_utils.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/src/ai/backend/manager/models/gql_models/container_registry_utils.py b/src/ai/backend/manager/models/gql_models/container_registry_utils.py index 7a241c9d8c..af0e9546a6 100644 --- a/src/ai/backend/manager/models/gql_models/container_registry_utils.py +++ b/src/ai/backend/manager/models/gql_models/container_registry_utils.py @@ -112,6 +112,9 @@ async def _get_harbor_project_id( ) async with sess.get(get_project_id_api, allow_redirects=False, **rqst_args) as resp: + if resp.status != 200: + raise InternalServerError(f"Failed to get harbor project_id! response: {resp}") + res = await resp.json() harbor_project_id = res["project_id"] return harbor_project_id @@ -126,6 +129,9 @@ async def _get_quota_info( }) async with sess.get(get_quota_id_api, allow_redirects=False, **rqst_args) as resp: + if resp.status != 200: + raise InternalServerError(f"Failed to get quota info! response: {resp}") + res = await resp.json() if not res: raise ObjectNotFound(object_name="quota entity") @@ -188,9 +194,7 @@ async def create(self, quota: int) -> None: async with sess.put( put_quota_api, json=payload, allow_redirects=False, **rqst_args ) as resp: - if resp.status == 200: - return None - else: + if resp.status != 200: log.error(f"Failed to create quota! response: {resp}") raise InternalServerError(f"Failed to create quota! response: {resp}") @@ -216,9 +220,7 @@ async def update(self, quota: int) -> None: async with sess.put( put_quota_api, json=payload, allow_redirects=False, **rqst_args ) as resp: - if resp.status == 200: - return None - else: + if resp.status != 200: log.error(f"Failed to update quota! response: {resp}") raise InternalServerError(f"Failed to update quota! response: {resp}") @@ -244,8 +246,6 @@ async def delete(self) -> None: async with sess.put( put_quota_api, json=payload, allow_redirects=False, **rqst_args ) as resp: - if resp.status == 200: - return None - else: + if resp.status != 200: log.error(f"Failed to delete quota! response: {resp}") raise InternalServerError(f"Failed to delete quota! response: {resp}") From 79d85e287eff0895a8a6ca0e0390862f75bfa145 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Thu, 28 Nov 2024 09:36:51 +0000 Subject: [PATCH 35/41] feat: `extra_fixtures` to `database_fixture` for effortless integration test writing --- tests/manager/conftest.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/tests/manager/conftest.py b/tests/manager/conftest.py index 92d39c7ce8..8ffd7da3c3 100644 --- a/tests/manager/conftest.py +++ b/tests/manager/conftest.py @@ -28,6 +28,7 @@ from unittest.mock import AsyncMock, MagicMock from urllib.parse import quote_plus as urlquote +import aiofiles.os import aiohttp import asyncpg import pytest @@ -419,7 +420,7 @@ async def database_engine(local_config, database): @pytest.fixture() -def database_fixture(local_config, test_db, database) -> Iterator[None]: +def database_fixture(local_config, test_db, database, extra_fixtures) -> Iterator[None]: """ Populate the example data as fixtures to the database and delete them after use. @@ -430,12 +431,20 @@ def database_fixture(local_config, test_db, database) -> Iterator[None]: db_url = f"postgresql+asyncpg://{db_user}:{urlquote(db_pass)}@{db_addr}/{test_db}" build_root = Path(os.environ["BACKEND_BUILD_ROOT"]) + + tmp_file = tempfile.NamedTemporaryFile(delete=False) + tmp_file_path = Path(tmp_file.name) + + with open(tmp_file_path, "w") as f: + json.dump(extra_fixtures, f) + fixture_paths = [ build_root / "fixtures" / "manager" / "example-users.json", build_root / "fixtures" / "manager" / "example-keypairs.json", build_root / "fixtures" / "manager" / "example-set-user-main-access-keys.json", build_root / "fixtures" / "manager" / "example-resource-presets.json", build_root / "fixtures" / "manager" / "example-container-registries-harbor.json", + tmp_file_path, ] async def init_fixture() -> None: @@ -460,6 +469,9 @@ async def init_fixture() -> None: yield async def clean_fixture() -> None: + if tmp_file_path.exists(): + await aiofiles.os.remove(tmp_file_path) + engine: SAEngine = sa.ext.asyncio.create_async_engine( db_url, connect_args=pgsql_connect_opts, From 15d396f4d42b84f69af928a0b853f0c6f7ddb932 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Thu, 28 Nov 2024 09:37:45 +0000 Subject: [PATCH 36/41] feat: Add `test_harbor_read_project_quota` --- .../models/test_container_registries.py | 95 +++++++++++++++++++ 1 file changed, 95 insertions(+) diff --git a/tests/manager/models/test_container_registries.py b/tests/manager/models/test_container_registries.py index 69abc7efcd..f5ccfdf8aa 100644 --- a/tests/manager/models/test_container_registries.py +++ b/tests/manager/models/test_container_registries.py @@ -1,10 +1,16 @@ import pytest +from aioresponses import aioresponses from graphene import Schema from graphene.test import Client +from ai.backend.common.utils import b64encode +from ai.backend.manager.api.context import RootContext from ai.backend.manager.defs import PASSWORD_PLACEHOLDER from ai.backend.manager.models.gql import GraphQueryContext, Mutations, Queries from ai.backend.manager.models.utils import ExtendedAsyncSAEngine +from ai.backend.manager.server import ( + database_ctx, +) CONTAINER_REGISTRY_FIELDS = """ hostname @@ -250,3 +256,92 @@ async def test_delete_container_registry(client: Client, database_engine: Extend response = await client.execute_async(query, variables=variables, context_value=context) assert response["data"] is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "extra_fixtures", + [ + { + "container_registries": [ + { + "id": "00000000-0000-0000-0000-000000000000", + "url": "http://mock_registry", + "registry_name": "mock_registry", + "project": "mock_project", + "username": "mock_user", + "password": "mock_password", + "ssl_verify": False, + "is_global": True, + } + ], + "groups": [ + { + "id": "00000000-0000-0000-0000-000000000000", + "name": "mock-group", + "description": "", + "is_active": True, + "domain_name": "default", + "resource_policy": "default", + "total_resource_slots": {}, + "allowed_vfolder_hosts": {}, + "container_registry": { + "registry": "mock_registry", + "project": "mock_project", + }, + "type": "general", + } + ], + }, + ], +) +async def test_harbor_read_project_quota( + client: Client, + database_fixture, + create_app_and_client, +): + test_app, _ = await create_app_and_client( + [ + database_ctx, + ], + [], + ) + + root_ctx: RootContext = test_app["_root.context"] + context = get_graphquery_context(root_ctx.db) + + # Arbitrary values for mocking Harbor API responses + HARBOR_PROJECT_ID = "123" + HARBOR_QUOTA_ID = 456 + HARBOR_QUOTA_VALUE = 1024 + + with aioresponses() as mocked: + # Mock the get project ID API call + get_project_id_url = "http://mock_registry/api/v2.0/projects/mock_project" + mocked.get(get_project_id_url, status=200, payload={"project_id": HARBOR_PROJECT_ID}) + + # Mock the get quota info API call + get_quota_url = f"http://mock_registry/api/v2.0/quotas?reference=project&reference_id={HARBOR_PROJECT_ID}" + mocked.get( + get_quota_url, + status=200, + payload=[{"id": HARBOR_QUOTA_ID, "hard": {"storage": HARBOR_QUOTA_VALUE}}], + ) + + groupnode_query = """ + query ($id: String!) { + group_node(id: $id) { + registry_quota + } + } + """ + + group_id = "00000000-0000-0000-0000-000000000000" + variables = { + "id": b64encode(f"group_node:{group_id}"), + } + + response = await client.execute_async( + groupnode_query, variables=variables, context_value=context + ) + assert response["data"]["group_node"]["registry_quota"] == HARBOR_QUOTA_VALUE From cd0d828f352a3efbbb509d6405e96e2293a837ad Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Thu, 28 Nov 2024 09:40:17 +0000 Subject: [PATCH 37/41] chore: `mock-group` -> `mock_group` --- tests/manager/models/test_container_registries.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/manager/models/test_container_registries.py b/tests/manager/models/test_container_registries.py index f5ccfdf8aa..6a29df2ed2 100644 --- a/tests/manager/models/test_container_registries.py +++ b/tests/manager/models/test_container_registries.py @@ -278,7 +278,7 @@ async def test_delete_container_registry(client: Client, database_engine: Extend "groups": [ { "id": "00000000-0000-0000-0000-000000000000", - "name": "mock-group", + "name": "mock_group", "description": "", "is_active": True, "domain_name": "default", From 515d39da0030fa9c0ca279297d17a6ed2371cf5b Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Thu, 28 Nov 2024 09:42:45 +0000 Subject: [PATCH 38/41] fix: Disjoint `FIXTURES_FOR_HARBOR_CRUD_TEST` from `test_harbor_read_project_quota` for reuse --- .../models/test_container_registries.py | 72 +++++++++---------- 1 file changed, 36 insertions(+), 36 deletions(-) diff --git a/tests/manager/models/test_container_registries.py b/tests/manager/models/test_container_registries.py index 6a29df2ed2..00bece04cc 100644 --- a/tests/manager/models/test_container_registries.py +++ b/tests/manager/models/test_container_registries.py @@ -258,43 +258,43 @@ async def test_delete_container_registry(client: Client, database_engine: Extend assert response["data"] is None -@pytest.mark.asyncio -@pytest.mark.parametrize( - "extra_fixtures", - [ - { - "container_registries": [ - { - "id": "00000000-0000-0000-0000-000000000000", - "url": "http://mock_registry", - "registry_name": "mock_registry", +FIXTURES_FOR_HARBOR_CRUD_TEST = [ + { + "container_registries": [ + { + "id": "00000000-0000-0000-0000-000000000000", + "url": "http://mock_registry", + "registry_name": "mock_registry", + "project": "mock_project", + "username": "mock_user", + "password": "mock_password", + "ssl_verify": False, + "is_global": True, + } + ], + "groups": [ + { + "id": "00000000-0000-0000-0000-000000000000", + "name": "mock_group", + "description": "", + "is_active": True, + "domain_name": "default", + "resource_policy": "default", + "total_resource_slots": {}, + "allowed_vfolder_hosts": {}, + "container_registry": { + "registry": "mock_registry", "project": "mock_project", - "username": "mock_user", - "password": "mock_password", - "ssl_verify": False, - "is_global": True, - } - ], - "groups": [ - { - "id": "00000000-0000-0000-0000-000000000000", - "name": "mock_group", - "description": "", - "is_active": True, - "domain_name": "default", - "resource_policy": "default", - "total_resource_slots": {}, - "allowed_vfolder_hosts": {}, - "container_registry": { - "registry": "mock_registry", - "project": "mock_project", - }, - "type": "general", - } - ], - }, - ], -) + }, + "type": "general", + } + ], + }, +] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("extra_fixtures", FIXTURES_FOR_HARBOR_CRUD_TEST) async def test_harbor_read_project_quota( client: Client, database_fixture, From b40f6e3b88e08d1cbe750f79f104b50d40f78a33 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Thu, 28 Nov 2024 09:44:29 +0000 Subject: [PATCH 39/41] chore: Rename variable --- tests/manager/conftest.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/tests/manager/conftest.py b/tests/manager/conftest.py index 8ffd7da3c3..257d29b9ec 100644 --- a/tests/manager/conftest.py +++ b/tests/manager/conftest.py @@ -432,10 +432,10 @@ def database_fixture(local_config, test_db, database, extra_fixtures) -> Iterato build_root = Path(os.environ["BACKEND_BUILD_ROOT"]) - tmp_file = tempfile.NamedTemporaryFile(delete=False) - tmp_file_path = Path(tmp_file.name) + extra_fixture_file = tempfile.NamedTemporaryFile(delete=False) + extra_fixture_file_path = Path(extra_fixture_file.name) - with open(tmp_file_path, "w") as f: + with open(extra_fixture_file_path, "w") as f: json.dump(extra_fixtures, f) fixture_paths = [ @@ -444,7 +444,7 @@ def database_fixture(local_config, test_db, database, extra_fixtures) -> Iterato build_root / "fixtures" / "manager" / "example-set-user-main-access-keys.json", build_root / "fixtures" / "manager" / "example-resource-presets.json", build_root / "fixtures" / "manager" / "example-container-registries-harbor.json", - tmp_file_path, + extra_fixture_file_path, ] async def init_fixture() -> None: @@ -469,8 +469,8 @@ async def init_fixture() -> None: yield async def clean_fixture() -> None: - if tmp_file_path.exists(): - await aiofiles.os.remove(tmp_file_path) + if extra_fixture_file_path.exists(): + await aiofiles.os.remove(extra_fixture_file_path) engine: SAEngine = sa.ext.asyncio.create_async_engine( db_url, From 1c04e5a58456f776d03dedbf20e786a8f8b9f15c Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Thu, 28 Nov 2024 09:55:25 +0000 Subject: [PATCH 40/41] fix: CI --- tests/manager/conftest.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/manager/conftest.py b/tests/manager/conftest.py index 257d29b9ec..b027e4b8fb 100644 --- a/tests/manager/conftest.py +++ b/tests/manager/conftest.py @@ -419,6 +419,11 @@ async def database_engine(local_config, database): yield db +@pytest.fixture() +def extra_fixtures(): + return {} + + @pytest.fixture() def database_fixture(local_config, test_db, database, extra_fixtures) -> Iterator[None]: """ From fc3060c778aeeb421c59bffcbd78cd0410422e5a Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Thu, 28 Nov 2024 10:01:10 +0000 Subject: [PATCH 41/41] chore: Add registry type --- tests/manager/models/test_container_registries.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/manager/models/test_container_registries.py b/tests/manager/models/test_container_registries.py index 00bece04cc..e4b49d7dba 100644 --- a/tests/manager/models/test_container_registries.py +++ b/tests/manager/models/test_container_registries.py @@ -263,6 +263,7 @@ async def test_delete_container_registry(client: Client, database_engine: Extend "container_registries": [ { "id": "00000000-0000-0000-0000-000000000000", + "type": "harbor2", "url": "http://mock_registry", "registry_name": "mock_registry", "project": "mock_project",