From 7762a25b9043ca0727048aa4b8e8bbbb79e1c728 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Wed, 13 Nov 2024 03:18:55 +0000 Subject: [PATCH 01/75] feat: Implement management API for controlling Harbor per-project Quota --- docs/manager/graphql-reference/schema.graphql | 9 ++ src/ai/backend/manager/models/gql.py | 4 + .../models/gql_models/container_registry.py | 135 +++++++++++++++++- 3 files changed, 147 insertions(+), 1 deletion(-) diff --git a/docs/manager/graphql-reference/schema.graphql b/docs/manager/graphql-reference/schema.graphql index 5dae5abf0a..e685383803 100644 --- a/docs/manager/graphql-reference/schema.graphql +++ b/docs/manager/graphql-reference/schema.graphql @@ -2053,6 +2053,9 @@ type Mutations { """Added in 25.1.0.""" delete_endpoint_auto_scaling_rule_node(id: String!): DeleteEndpointAutoScalingRuleNode + """Added in 24.12.0""" + update_container_registry_quota(quota: Int!, scope_id: ScopeField!): UpdateQuota + """Deprecated since 24.09.0. use `CreateContainerRegistryNode` instead""" create_container_registry(hostname: String!, props: CreateContainerRegistryInput!): CreateContainerRegistry @deprecated(reason: "Deprecated since 24.09.0. use `create_container_registry_node_v2` instead.") @@ -2955,6 +2958,12 @@ type DeleteEndpointAutoScalingRuleNode { msg: String } +"""Added in 24.12.0.""" +type UpdateQuota { + ok: Boolean + msg: String +} + """Deprecated since 24.09.0. use `CreateContainerRegistryNode` instead""" type CreateContainerRegistry { container_registry: ContainerRegistry diff --git a/src/ai/backend/manager/models/gql.py b/src/ai/backend/manager/models/gql.py index dbc768d15a..e8a10fd311 100644 --- a/src/ai/backend/manager/models/gql.py +++ b/src/ai/backend/manager/models/gql.py @@ -84,6 +84,9 @@ AgentSummaryList, ModifyAgent, ) +from .gql_models.container_registry import ( + UpdateQuota, +) from .gql_models.domain import ( CreateDomainNode, DomainConnection, @@ -382,6 +385,7 @@ class Mutations(graphene.ObjectType): delete_endpoint_auto_scaling_rule_node = DeleteEndpointAutoScalingRuleNode.Field( description="Added in 25.1.0." ) + update_container_registry_quota = UpdateQuota.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 2e595a8fa6..4793ecf409 100644 --- a/src/ai/backend/manager/models/gql_models/container_registry.py +++ b/src/ai/backend/manager/models/gql_models/container_registry.py @@ -3,12 +3,15 @@ import logging import uuid from collections.abc import Sequence -from typing import TYPE_CHECKING, Any, Optional, cast +from typing import TYPE_CHECKING, Any, Optional, Self, cast +import aiohttp import graphene import graphql import sqlalchemy as sa +import yarl from graphql import Undefined, UndefinedType +from sqlalchemy.orm import load_only from ai.backend.common.container_registry import AllowedGroupsModel, ContainerRegistryType from ai.backend.common.logging_utils import BraceStyleAdapter @@ -17,8 +20,12 @@ AssociationContainerRegistriesGroupsRow, ) from ai.backend.manager.models.container_registry import ContainerRegistryRow +from ai.backend.manager.models.gql_models.fields import ScopeField +from ai.backend.manager.models.group import GroupRow from ai.backend.manager.models.rbac import ( ContainerRegistryScope, + ProjectScope, + ScopeType, SystemScope, ) from ai.backend.manager.models.user import UserRole @@ -479,3 +486,129 @@ async def mutate( ) return cls(container_registry=container_registry) + + +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 49f60c988dece9d34253f554247f72bd6eb0f229 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Wed, 13 Nov 2024 03:21:17 +0000 Subject: [PATCH 02/75] 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 683f9275ec3addbe3b8f1eed36bb534b71841c53 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Wed, 13 Nov 2024 03:28:39 +0000 Subject: [PATCH 03/75] 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 4793ecf409..7382e3920c 100644 --- a/src/ai/backend/manager/models/gql_models/container_registry.py +++ b/src/ai/backend/manager/models/gql_models/container_registry.py @@ -494,7 +494,6 @@ class UpdateQuota(graphene.Mutation): allowed_roles = ( UserRole.SUPERADMIN, UserRole.ADMIN, - UserRole.USER, ) class Arguments: @@ -557,19 +556,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 b832cf9a9e57b7f6820f9a3af2530c9822f6e8a7 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Wed, 13 Nov 2024 03:34:33 +0000 Subject: [PATCH 04/75] 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 7382e3920c..bdc5e00a55 100644 --- a/src/ai/backend/manager/models/gql_models/container_registry.py +++ b/src/ai/backend/manager/models/gql_models/container_registry.py @@ -559,7 +559,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( @@ -567,27 +567,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 71d2aaa8a4a282bcc03f30918caba03d9275a5c0 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Wed, 13 Nov 2024 04:06:27 +0000 Subject: [PATCH 05/75] feat: Add `registry_quota` to GroupNode --- docs/manager/graphql-reference/schema.graphql | 3 + .../manager/models/gql_models/group.py | 67 +++++++++++++++++++ 2 files changed, 70 insertions(+) diff --git a/docs/manager/graphql-reference/schema.graphql b/docs/manager/graphql-reference/schema.graphql index e685383803..01c74d72b9 100644 --- a/docs/manager/graphql-reference/schema.graphql +++ b/docs/manager/graphql-reference/schema.graphql @@ -719,6 +719,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 47ee961717..f518bc9c71 100644 --- a/src/ai/backend/manager/models/gql_models/group.py +++ b/src/ai/backend/manager/models/gql_models/group.py @@ -9,12 +9,16 @@ Sequence, ) +import aiohttp import graphene import graphql 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, @@ -118,6 +122,8 @@ class Meta: lambda: graphene.String, ) + registry_quota = graphene.Int(description="Added in 24.12.0.") + user_nodes = PaginatedConnectionField( UserConnection, ) @@ -210,6 +216,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 7231fc3ce5eea90f18fe86ea2d28e56f34ce26da Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Wed, 13 Nov 2024 08:34:25 +0000 Subject: [PATCH 06/75] 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 bdc5e00a55..f4ca557c3d 100644 --- a/src/ai/backend/manager/models/gql_models/container_registry.py +++ b/src/ai/backend/manager/models/gql_models/container_registry.py @@ -511,16 +511,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 = ( @@ -529,32 +526,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) @@ -573,21 +578,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 954306e69e324a58ba8f52c932b7da2555427597 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Wed, 13 Nov 2024 08:55:50 +0000 Subject: [PATCH 07/75] fix: Update `resolve_registry_quota` --- .../manager/models/gql_models/group.py | 25 +++++++++---------- 1 file changed, 12 insertions(+), 13 deletions(-) diff --git a/src/ai/backend/manager/models/gql_models/group.py b/src/ai/backend/manager/models/gql_models/group.py index f518bc9c71..4e8218b918 100644 --- a/src/ai/backend/manager/models/gql_models/group.py +++ b/src/ai/backend/manager/models/gql_models/group.py @@ -17,14 +17,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, @@ -218,15 +217,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 ): @@ -237,14 +230,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.") @@ -272,7 +267,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 9ac12226ef7cef06f102aa9045736899b9c568b3 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Thu, 14 Nov 2024 01:23:09 +0000 Subject: [PATCH 08/75] fix: Update `resolve_registry_quota` --- src/ai/backend/manager/models/gql_models/group.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/src/ai/backend/manager/models/gql_models/group.py b/src/ai/backend/manager/models/gql_models/group.py index 4e8218b918..cad41045d9 100644 --- a/src/ai/backend/manager/models/gql_models/group.py +++ b/src/ai/backend/manager/models/gql_models/group.py @@ -17,6 +17,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, @@ -223,7 +225,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"], @@ -239,9 +243,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 eaf0949ec38a7b6eac5c9274898896bf8703e8b6 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Thu, 14 Nov 2024 02:31:32 +0000 Subject: [PATCH 09/75] 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 cad41045d9..2d44449df8 100644 --- a/src/ai/backend/manager/models/gql_models/group.py +++ b/src/ai/backend/manager/models/gql_models/group.py @@ -19,6 +19,9 @@ from ai.backend.manager.api.exceptions import ContainerRegistryNotFound +from ..association_container_registries_groups import ( + AssociationContainerRegistriesGroupsRow, +) from ..base import ( FilterExprArg, OrderExprArg, @@ -247,6 +250,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 b5d220b76dc65ed4e2ae291d6cefbcbc9bc758f8 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Thu, 14 Nov 2024 02:50:56 +0000 Subject: [PATCH 10/75] feat: Add CreateQuota, DeleteQuota mutation --- src/ai/backend/manager/models/gql.py | 4 + .../models/gql_models/container_registry.py | 232 +++++++++++------- 2 files changed, 146 insertions(+), 90 deletions(-) diff --git a/src/ai/backend/manager/models/gql.py b/src/ai/backend/manager/models/gql.py index e8a10fd311..d9907953d1 100644 --- a/src/ai/backend/manager/models/gql.py +++ b/src/ai/backend/manager/models/gql.py @@ -85,6 +85,8 @@ ModifyAgent, ) from .gql_models.container_registry import ( + CreateQuota, + DeleteQuota, UpdateQuota, ) from .gql_models.domain import ( @@ -385,7 +387,9 @@ class Mutations(graphene.ObjectType): delete_endpoint_auto_scaling_rule_node = DeleteEndpointAutoScalingRuleNode.Field( description="Added in 25.1.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 f4ca557c3d..95f908dcc8 100644 --- a/src/ai/backend/manager/models/gql_models/container_registry.py +++ b/src/ai/backend/manager/models/gql_models/container_registry.py @@ -488,7 +488,109 @@ async def mutate( return cls(container_registry=container_registry) -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 = ( @@ -511,104 +613,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." - ) - - 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() + return await update_quota(cls, info, scope_id, quota) - 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: - return UpdateQuota(ok=False, msg="Only HarborV2 registry is supported for now.") +class UpdateQuota(graphene.Mutation): + """Added in 24.12.0.""" - ssl_verify = registry.ssl_verify - connector = aiohttp.TCPConnector(ssl=ssl_verify) + allowed_roles = ( + UserRole.SUPERADMIN, + UserRole.ADMIN, + ) - 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, - ) + class Arguments: + scope_id = ScopeField(required=True) + quota = graphene.Int(required=True) - get_project_id_api = api_url / "projects" / project + ok = graphene.Boolean() + msg = graphene.String() - 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"] + @classmethod + async def mutate( + cls, + root, + info: graphene.ResolveInfo, + scope_id: ScopeType, + quota: int, + ) -> Self: + return await update_quota(cls, info, scope_id, quota) - 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})", - ) +class DeleteQuota(graphene.Mutation): + """Added in 24.12.0.""" - quota_id = res[0]["id"] + allowed_roles = ( + UserRole.SUPERADMIN, + UserRole.ADMIN, + ) - put_quota_api = api_url / "quotas" / str(quota_id) - payload = {"hard": {"storage": quota}} + class Arguments: + scope_id = ScopeField(required=True) - 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}" - ) + ok = graphene.Boolean() + msg = graphene.String() - 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 32eba43be5c5730ae4a98438f314e0707307c6e6 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Thu, 14 Nov 2024 02:55:51 +0000 Subject: [PATCH 11/75] 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 95f908dcc8..6ac3e4c9c6 100644 --- a/src/ai/backend/manager/models/gql_models/container_registry.py +++ b/src/ai/backend/manager/models/gql_models/container_registry.py @@ -488,7 +488,7 @@ async def mutate( return cls(container_registry=container_registry) -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): @@ -613,7 +613,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): @@ -639,7 +639,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): @@ -663,4 +663,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 89b79e2f8887bcf8586d7b6946911190ad8cc35b Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Thu, 14 Nov 2024 03:52:45 +0000 Subject: [PATCH 12/75] 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 6ac3e4c9c6..d5d65fe717 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 import uuid from collections.abc import Sequence @@ -488,8 +489,18 @@ async def mutate( return cls(container_registry=container_registry) +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.") @@ -514,7 +525,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 = ( @@ -541,8 +552,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( @@ -550,6 +559,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: @@ -573,6 +583,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) @@ -582,10 +600,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!") @@ -613,7 +633,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): @@ -639,7 +661,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): @@ -663,4 +687,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 5ae4e947e1fcb7f29bda9e4b77b217962f47decb Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Thu, 14 Nov 2024 03:53:59 +0000 Subject: [PATCH 13/75] chore: Update schema --- docs/manager/graphql-reference/schema.graphql | 18 ++++++++++++++++++ .../backend/manager/models/gql_models/group.py | 4 +++- 2 files changed, 21 insertions(+), 1 deletion(-) diff --git a/docs/manager/graphql-reference/schema.graphql b/docs/manager/graphql-reference/schema.graphql index 01c74d72b9..a9b0adb920 100644 --- a/docs/manager/graphql-reference/schema.graphql +++ b/docs/manager/graphql-reference/schema.graphql @@ -2056,9 +2056,15 @@ type Mutations { """Added in 25.1.0.""" delete_endpoint_auto_scaling_rule_node(id: String!): DeleteEndpointAutoScalingRuleNode + """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 + """Deprecated since 24.09.0. use `CreateContainerRegistryNode` instead""" create_container_registry(hostname: String!, props: CreateContainerRegistryInput!): CreateContainerRegistry @deprecated(reason: "Deprecated since 24.09.0. use `create_container_registry_node_v2` instead.") @@ -2961,12 +2967,24 @@ type DeleteEndpointAutoScalingRuleNode { 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 +} + """Deprecated since 24.09.0. use `CreateContainerRegistryNode` instead""" type CreateContainerRegistry { container_registry: ContainerRegistry diff --git a/src/ai/backend/manager/models/gql_models/group.py b/src/ai/backend/manager/models/gql_models/group.py index 2d44449df8..26053ed689 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 @@ PaginatedConnectionField, generate_sql_info_for_gql_connection, ) -from ..container_registry import ContainerRegistryRow, ContainerRegistryType +from ..container_registry import ContainerRegistryType from ..gql_relay import ( AsyncNode, Connection, @@ -221,6 +221,8 @@ async def resolve_user_nodes( return ConnectionResolverResult(result, cursor, pagination_order, page_size, total_cnt) async def resolve_registry_quota(self, info: graphene.ResolveInfo) -> int: + from ..container_registry import ContainerRegistryRow + graph_ctx = info.context async with graph_ctx.db.begin_session() as db_sess: if ( From 23aceb614ae4baf7cd4a746c2ec93c28141a97c2 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Thu, 14 Nov 2024 03:57:46 +0000 Subject: [PATCH 14/75] 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 d5d65fe717..8e20214c39 100644 --- a/src/ai/backend/manager/models/gql_models/container_registry.py +++ b/src/ai/backend/manager/models/gql_models/container_registry.py @@ -580,16 +580,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 321cf13fe8b5a0dc0ff927f175f64c5059fe3085 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Thu, 14 Nov 2024 03:59:18 +0000 Subject: [PATCH 15/75] 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 8e20214c39..aee079c4fc 100644 --- a/src/ai/backend/manager/models/gql_models/container_registry.py +++ b/src/ai/backend/manager/models/gql_models/container_registry.py @@ -497,13 +497,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 @@ -542,13 +544,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) @@ -574,11 +576,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})", ) @@ -586,10 +588,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"] @@ -600,14 +604,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 db2389191f93dc3fda885a22d882c704481842cb Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Thu, 14 Nov 2024 04:08:21 +0000 Subject: [PATCH 16/75] fix: Remove useless strenum --- .../models/gql_models/container_registry.py | 28 ++++++------------- 1 file changed, 9 insertions(+), 19 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 aee079c4fc..48200a642f 100644 --- a/src/ai/backend/manager/models/gql_models/container_registry.py +++ b/src/ai/backend/manager/models/gql_models/container_registry.py @@ -1,6 +1,5 @@ from __future__ import annotations -import enum import logging import uuid from collections.abc import Sequence @@ -489,19 +488,16 @@ async def mutate( return cls(container_registry=container_registry) -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." @@ -586,10 +582,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})" @@ -637,9 +633,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): @@ -665,9 +659,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): @@ -691,6 +683,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 6edd92166cde6d5ab5ce3add6d5b193b85a45607 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Mon, 18 Nov 2024 02:40:53 +0000 Subject: [PATCH 17/75] refactor: `mutate_harbor_project_quota` --- .../models/gql_models/container_registry.py | 142 ++++++++++-------- 1 file changed, 76 insertions(+), 66 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 48200a642f..00b67e11fe 100644 --- a/src/ai/backend/manager/models/gql_models/container_registry.py +++ b/src/ai/backend/manager/models/gql_models/container_registry.py @@ -3,7 +3,7 @@ import logging import uuid from collections.abc import Sequence -from typing import TYPE_CHECKING, Any, Optional, Self, cast +from typing import TYPE_CHECKING, Any, Literal, Optional, Self, cast import aiohttp import graphene @@ -11,13 +11,18 @@ import sqlalchemy as sa import yarl from graphql import Undefined, UndefinedType +from sqlalchemy.ext.asyncio import AsyncSession as SASession from sqlalchemy.orm import load_only from ai.backend.common.container_registry import AllowedGroupsModel, ContainerRegistryType from ai.backend.common.logging_utils import BraceStyleAdapter -from ai.backend.manager.api.exceptions import ContainerRegistryNotFound -from ai.backend.manager.models.association_container_registries_groups import ( - AssociationContainerRegistriesGroupsRow, +from ai.backend.logging import BraceStyleAdapter +from ai.backend.manager.api.exceptions import ( + ContainerRegistryNotFound, + GenericBadRequest, + InternalServerError, + NotImplementedAPI, + ObjectNotFound, ) from ai.backend.manager.models.container_registry import ContainerRegistryRow from ai.backend.manager.models.gql_models.fields import ScopeField @@ -32,6 +37,9 @@ from ai.backend.manager.models.utils import ExtendedAsyncSAEngine from ...defs import PASSWORD_PLACEHOLDER +from ..association_container_registries_groups import ( + AssociationContainerRegistriesGroupsRow, +) from ..base import ( FilterExprArg, OrderExprArg, @@ -488,65 +496,57 @@ async def mutate( return cls(container_registry=container_registry) -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)) + 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})" ) - 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"], - ) + 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) - ) + 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() + result = await db_sess.execute(registry_query) + registry = result.scalars().one_or_none() - if not registry: - return mutation_cls( - ok=False, - msg=f"Specified container registry row does not exist. (cr: {registry_name}, gr: {project})", - ) + if not registry: + raise ContainerRegistryNotFound( + f"Specified container registry row does not exist. (cr: {registry_name}, gr: {project})" + ) - if registry.type != ContainerRegistryType.HARBOR2: - return mutation_cls(ok=False, msg="Only HarborV2 registry is supported for now.") + 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) @@ -572,24 +572,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"] @@ -600,14 +595,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): @@ -633,7 +628,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): @@ -659,7 +659,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): @@ -683,4 +688,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 b38a8a65a5b3c50f10a6adabe3a1a65aa4298a3c Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Mon, 18 Nov 2024 03:16:24 +0000 Subject: [PATCH 18/75] refactor: Add read operation handling for code reuse --- .../models/gql_models/container_registry.py | 131 +-------------- .../gql_models/container_registry_utils.py | 159 ++++++++++++++++++ .../manager/models/gql_models/group.py | 89 +--------- 3 files changed, 175 insertions(+), 204 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 00b67e11fe..48c1a7d61e 100644 --- a/src/ai/backend/manager/models/gql_models/container_registry.py +++ b/src/ai/backend/manager/models/gql_models/container_registry.py @@ -3,33 +3,23 @@ import logging import uuid from collections.abc import Sequence -from typing import TYPE_CHECKING, Any, Literal, Optional, Self, cast +from typing import TYPE_CHECKING, Any, Optional, Self, cast -import aiohttp import graphene import graphql import sqlalchemy as sa -import yarl from graphql import Undefined, UndefinedType -from sqlalchemy.ext.asyncio import AsyncSession as SASession -from sqlalchemy.orm import load_only from ai.backend.common.container_registry import AllowedGroupsModel, ContainerRegistryType from ai.backend.common.logging_utils import BraceStyleAdapter from ai.backend.logging import BraceStyleAdapter from ai.backend.manager.api.exceptions import ( ContainerRegistryNotFound, - GenericBadRequest, - InternalServerError, - NotImplementedAPI, - ObjectNotFound, ) from ai.backend.manager.models.container_registry import ContainerRegistryRow from ai.backend.manager.models.gql_models.fields import ScopeField -from ai.backend.manager.models.group import GroupRow from ai.backend.manager.models.rbac import ( ContainerRegistryScope, - ProjectScope, ScopeType, SystemScope, ) @@ -53,6 +43,10 @@ if TYPE_CHECKING: from ..gql import GraphQueryContext +from ..rbac import ScopeType +from ..user import UserRole +from .container_registry_utils import handle_harbor_project_quota_operation +from .fields import ScopeField log = BraceStyleAdapter(logging.getLogger(__spec__.name)) # type: ignore @@ -496,115 +490,6 @@ async def mutate( return cls(container_registry=container_registry) -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.""" @@ -630,7 +515,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)) @@ -661,7 +546,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)) @@ -690,7 +575,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 26053ed689..c6071c126b 100644 --- a/src/ai/backend/manager/models/gql_models/group.py +++ b/src/ai/backend/manager/models/gql_models/group.py @@ -3,32 +3,23 @@ from collections.abc import Mapping from typing import ( TYPE_CHECKING, - Any, Optional, Self, Sequence, ) -import aiohttp import graphene import graphql 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 ContainerRegistryType from ..gql_relay import ( AsyncNode, Connection, @@ -37,8 +28,12 @@ from ..group import AssocGroupUserRow, GroupRow, ProjectType, get_permission_ctx from ..minilang.ordering import OrderSpecItem, QueryOrderParser from ..minilang.queryfilter import FieldSpecItem, QueryFilterParser +from ..rbac import ProjectScope from ..rbac.context import ClientContext from ..rbac.permission_defs import ProjectPermission +from .container_registry_utils import ( + handle_harbor_project_quota_operation, +) from .user import UserConnection, UserNode if TYPE_CHECKING: @@ -221,80 +216,12 @@ async def resolve_user_nodes( return ConnectionResolverResult(result, cursor, pagination_order, page_size, total_cnt) async def resolve_registry_quota(self, info: graphene.ResolveInfo) -> int: - from ..container_registry import ContainerRegistryRow - 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 ff568abc259ef9384d2dbe4ec1d21a138ab3f1dd Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Mon, 18 Nov 2024 04:51:30 +0000 Subject: [PATCH 19/75] 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 434460b9cf..e413d4b165 100644 --- a/src/ai/backend/client/func/group.py +++ b/src/ai/backend/client/func/group.py @@ -1,5 +1,7 @@ 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 @@ -293,3 +295,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 e465c3d9420bd33e0c29050ae4020a5ccd8f35d4 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Mon, 18 Nov 2024 05:15:15 +0000 Subject: [PATCH 20/75] 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 e413d4b165..2d80b17962 100644 --- a/src/ai/backend/client/func/group.py +++ b/src/ai/backend/client/func/group.py @@ -1,9 +1,8 @@ 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 @@ -315,7 +314,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 4eed2a0366..f8d31b7366 100644 --- a/src/ai/backend/common/utils.py +++ b/src/ai/backend/common/utils.py @@ -425,3 +425,12 @@ def join_non_empty(*args: Optional[str], sep: str) -> str: """ 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 f022d4cb934b1cf70b088bf4ee4a56add28684db Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Mon, 18 Nov 2024 05:58:17 +0000 Subject: [PATCH 21/75] 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 d78382386e..bb8c76f7ab 100644 --- a/src/ai/backend/manager/server.py +++ b/src/ai/backend/manager/server.py @@ -194,6 +194,7 @@ ".image", ".userconfig", ".domainconfig", + ".group", ".groupconfig", ".logs", ] From 3c2ea70b780c076c932243c5e0c0a96e44967494 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Mon, 18 Nov 2024 05:58:54 +0000 Subject: [PATCH 22/75] 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 ee79ed4ff8c28b2352e3f8bc3a30c337b571bf8b Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Tue, 19 Nov 2024 01:25:01 +0000 Subject: [PATCH 23/75] 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 2d80b17962..1135791dc6 100644 --- a/src/ai/backend/client/func/group.py +++ b/src/ai/backend/client/func/group.py @@ -300,7 +300,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. """ @@ -323,7 +323,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. """ @@ -348,7 +348,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. """ @@ -373,7 +373,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 f3f7691089a99222b8f78d613919347e09a6fe28 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Tue, 19 Nov 2024 01:40:08 +0000 Subject: [PATCH 24/75] 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 d9907953d1..ac4716dc05 100644 --- a/src/ai/backend/manager/models/gql.py +++ b/src/ai/backend/manager/models/gql.py @@ -85,9 +85,9 @@ ModifyAgent, ) from .gql_models.container_registry import ( - CreateQuota, - DeleteQuota, - UpdateQuota, + CreateContainerRegistryQuota, + DeleteContainerRegistryQuota, + UpdateContainerRegistryQuota, ) from .gql_models.domain import ( CreateDomainNode, @@ -387,9 +387,15 @@ class Mutations(graphene.ObjectType): delete_endpoint_auto_scaling_rule_node = DeleteEndpointAutoScalingRuleNode.Field( description="Added in 25.1.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 48c1a7d61e..e2f672389c 100644 --- a/src/ai/backend/manager/models/gql_models/container_registry.py +++ b/src/ai/backend/manager/models/gql_models/container_registry.py @@ -490,7 +490,7 @@ async def mutate( return cls(container_registry=container_registry) -class CreateQuota(graphene.Mutation): +class CreateContainerRegistryQuota(graphene.Mutation): """Added in 24.12.0.""" allowed_roles = ( @@ -521,7 +521,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 = ( @@ -552,7 +552,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 720075e7c1966994120c2f0b95c795681b1c41b5 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Tue, 19 Nov 2024 01:42:27 +0000 Subject: [PATCH 25/75] chore: update GraphQL schema dump Co-authored-by: octodog --- docs/manager/graphql-reference/schema.graphql | 12 ++++++------ .../backend/manager/container_registry/__init__.py | 1 + .../models/gql_models/container_registry_utils.py | 3 ++- 3 files changed, 9 insertions(+), 7 deletions(-) diff --git a/docs/manager/graphql-reference/schema.graphql b/docs/manager/graphql-reference/schema.graphql index a9b0adb920..4e8858520b 100644 --- a/docs/manager/graphql-reference/schema.graphql +++ b/docs/manager/graphql-reference/schema.graphql @@ -2057,13 +2057,13 @@ type Mutations { delete_endpoint_auto_scaling_rule_node(id: String!): DeleteEndpointAutoScalingRuleNode """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 """Deprecated since 24.09.0. use `CreateContainerRegistryNode` instead""" create_container_registry(hostname: String!, props: CreateContainerRegistryInput!): CreateContainerRegistry @deprecated(reason: "Deprecated since 24.09.0. use `create_container_registry_node_v2` instead.") @@ -2968,19 +2968,19 @@ type DeleteEndpointAutoScalingRuleNode { } """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 } diff --git a/src/ai/backend/manager/container_registry/__init__.py b/src/ai/backend/manager/container_registry/__init__.py index 9c58a9a7cf..bb16e6190f 100644 --- a/src/ai/backend/manager/container_registry/__init__.py +++ b/src/ai/backend/manager/container_registry/__init__.py @@ -9,6 +9,7 @@ if TYPE_CHECKING: from ..models.container_registry import ContainerRegistryRow from .base import BaseContainerRegistry + from ..container_registry import ContainerRegistryRow def get_container_registry_cls(registry_info: ContainerRegistryRow) -> Type[BaseContainerRegistry]: 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..20912020f7 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 @@ -19,7 +19,6 @@ ObjectNotFound, ) -from ...container_registry import ContainerRegistryRow from ..association_container_registries_groups import ( AssociationContainerRegistriesGroupsRow, ) @@ -41,6 +40,8 @@ async def handle_harbor_project_quota_operation( :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. """ + from ..container_registry import ContainerRegistryRow + if not isinstance(scope_id, ProjectScope): raise NotImplementedAPI("Quota mutation currently supports only the project scope.") From 370a11e4f1aef58c2565460385c43267b2130ccb Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Tue, 19 Nov 2024 01:47:44 +0000 Subject: [PATCH 26/75] 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 4afa7b102c4db29227d9c487a2ea3fb20260bb0d Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Wed, 20 Nov 2024 05:34:27 +0000 Subject: [PATCH 27/75] fix: Use `BigInt` --- .../backend/manager/container_registry/__init__.py | 2 +- .../manager/models/gql_models/container_registry.py | 13 +++++++------ 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/src/ai/backend/manager/container_registry/__init__.py b/src/ai/backend/manager/container_registry/__init__.py index bb16e6190f..7de7a7f4ae 100644 --- a/src/ai/backend/manager/container_registry/__init__.py +++ b/src/ai/backend/manager/container_registry/__init__.py @@ -7,9 +7,9 @@ from ai.backend.common.container_registry import ContainerRegistryType if TYPE_CHECKING: + from ..container_registry import ContainerRegistryRow from ..models.container_registry import ContainerRegistryRow from .base import BaseContainerRegistry - from ..container_registry import ContainerRegistryRow def get_container_registry_cls(registry_info: ContainerRegistryRow) -> Type[BaseContainerRegistry]: 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 e2f672389c..52316e76ac 100644 --- a/src/ai/backend/manager/models/gql_models/container_registry.py +++ b/src/ai/backend/manager/models/gql_models/container_registry.py @@ -31,6 +31,7 @@ AssociationContainerRegistriesGroupsRow, ) from ..base import ( + BigInt, FilterExprArg, OrderExprArg, PaginatedConnectionField, @@ -500,7 +501,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() @@ -511,11 +512,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)) @@ -531,7 +532,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() @@ -542,11 +543,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 4dcb980542b4c8c7b9b129b33b25439d3e5a6179 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Wed, 20 Nov 2024 05:36:56 +0000 Subject: [PATCH 28/75] chore: update GraphQL schema dump Co-authored-by: octodog --- docs/manager/graphql-reference/schema.graphql | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/manager/graphql-reference/schema.graphql b/docs/manager/graphql-reference/schema.graphql index 4e8858520b..cddf78dfd4 100644 --- a/docs/manager/graphql-reference/schema.graphql +++ b/docs/manager/graphql-reference/schema.graphql @@ -2057,10 +2057,10 @@ type Mutations { delete_endpoint_auto_scaling_rule_node(id: String!): DeleteEndpointAutoScalingRuleNode """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 47fa15d27abd2ddce359693d1c57ddcd41bed66f Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Thu, 28 Nov 2024 03:16:05 +0000 Subject: [PATCH 29/75] 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 | 279 ++++++++++++------ .../manager/models/gql_models/group.py | 7 +- 5 files changed, 206 insertions(+), 107 deletions(-) diff --git a/src/ai/backend/client/func/group.py b/src/ai/backend/client/func/group.py index 1135791dc6..7596aeec8f 100644 --- a/src/ai/backend/client/func/group.py +++ b/src/ai/backend/client/func/group.py @@ -299,7 +299,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 52316e76ac..2323c3177a 100644 --- a/src/ai/backend/manager/models/gql_models/container_registry.py +++ b/src/ai/backend/manager/models/gql_models/container_registry.py @@ -46,7 +46,7 @@ from ..gql import GraphQueryContext 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 @@ -516,7 +516,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)) @@ -547,7 +548,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)) @@ -576,7 +578,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 20912020f7..a6fe494698 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, @@ -28,96 +30,100 @@ 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): + """ + Utility class for HarborV2 per-project Quota CRUD API. """ from ..container_registry import ContainerRegistryRow - 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})" - ) + db_sess: SASession + scope_id: ScopeType + group_row: GroupRow + registry: ContainerRegistryRow + project: str + project_id: uuid.UUID + + 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_name, project = ( - group_row.container_registry["registry"], - group_row.container_registry["project"], - ) + self.db_sess = db_sess + self.scope_id = scope_id - registry_query = sa.select(ContainerRegistryRow).where( - (ContainerRegistryRow.registry_name == registry_name) - & (ContainerRegistryRow.project == project) - ) + async def __ainit__(self) -> None: + assert isinstance(self.scope_id, ProjectScope) - result = await db_sess.execute(registry_query) - registry = result.scalars().one_or_none() + 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 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})" + ) - if not registry: - raise ContainerRegistryNotFound( - f"Specified container registry row does not exist. (cr: {registry_name}, gr: {project})" + registry_name, project = ( + group_row.container_registry["registry"], + group_row.container_registry["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) + registry_query = sa.select(ContainerRegistryRow).where( + (ContainerRegistryRow.registry_name == registry_name) + & (ContainerRegistryRow.project == 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, + + 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 ) - api_url = yarl.URL(registry.url) / "api" / "v2.0" - get_project_id_api = api_url / "projects" / project + 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() @@ -129,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"] - put_quota_api = api_url / "quotas" / str(quota_id) - quota = quota if operation_type != "delete" else -1 + if previous_quota > 0: + raise GenericBadRequest(f"Quota limit already exists. (gr: {self.project_id})") + + 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 c6071c126b..6ac05cfe5c 100644 --- a/src/ai/backend/manager/models/gql_models/group.py +++ b/src/ai/backend/manager/models/gql_models/group.py @@ -32,7 +32,7 @@ from ..rbac.context import ClientContext from ..rbac.permission_defs import ProjectPermission from .container_registry_utils import ( - handle_harbor_project_quota_operation, + HarborQuotaManager, ) from .user import UserConnection, UserNode @@ -219,9 +219,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 0ef1b359f3fe20a53b0f8f7fa0db7e7bc0d3294b Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Thu, 28 Nov 2024 03:23:52 +0000 Subject: [PATCH 30/75] 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 6ac05cfe5c..6345e8bf44 100644 --- a/src/ai/backend/manager/models/gql_models/group.py +++ b/src/ai/backend/manager/models/gql_models/group.py @@ -15,6 +15,7 @@ from graphene.types.datetime import DateTime as GQLDateTime from ..base import ( + BigInt, FilterExprArg, OrderExprArg, PaginatedConnectionField, @@ -121,7 +122,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 bf3e6aa1d90fc1fc29deedb88c48a11b19b68f2c Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Thu, 28 Nov 2024 03:25:37 +0000 Subject: [PATCH 31/75] 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 a6fe494698..00130976a4 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 4c34c01fa2610323e8724113dd6e55b1b6b389f3 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Thu, 28 Nov 2024 03:28:00 +0000 Subject: [PATCH 32/75] chore: update GraphQL schema dump Co-authored-by: octodog --- docs/manager/graphql-reference/schema.graphql | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/manager/graphql-reference/schema.graphql b/docs/manager/graphql-reference/schema.graphql index cddf78dfd4..acef847f73 100644 --- a/docs/manager/graphql-reference/schema.graphql +++ b/docs/manager/graphql-reference/schema.graphql @@ -721,7 +721,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 34a6a54f0175990c1c05e66805ca92114718359f Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Thu, 28 Nov 2024 03:34:33 +0000 Subject: [PATCH 33/75] 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 00130976a4..5b5a0c2b93 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 db2b25dcb87cc99d0948f11e936db0bf949c9891 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Thu, 28 Nov 2024 09:37:45 +0000 Subject: [PATCH 34/75] 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 0da01fab80..e80363c627 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 @@ -251,3 +257,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 00b9463dada14dd2f64dc82ae2c30ef1598b645d Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Thu, 28 Nov 2024 09:40:17 +0000 Subject: [PATCH 35/75] 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 e80363c627..8e3ab03b1a 100644 --- a/tests/manager/models/test_container_registries.py +++ b/tests/manager/models/test_container_registries.py @@ -279,7 +279,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 77a719d8464507e411ded2e2764a73deb7a51a5c Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Thu, 28 Nov 2024 09:42:45 +0000 Subject: [PATCH 36/75] 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 8e3ab03b1a..fa77f54558 100644 --- a/tests/manager/models/test_container_registries.py +++ b/tests/manager/models/test_container_registries.py @@ -259,43 +259,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 9814c4cc5de0d94058b5da3d51ca6f17f3c19142 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Thu, 28 Nov 2024 10:01:10 +0000 Subject: [PATCH 37/75] 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 fa77f54558..1a615651d9 100644 --- a/tests/manager/models/test_container_registries.py +++ b/tests/manager/models/test_container_registries.py @@ -264,6 +264,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", From a74fb708f2fb5a27c71c3d2c52d149f3956eb27f Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Sun, 1 Dec 2024 07:08:21 +0000 Subject: [PATCH 38/75] feat: Add create GQL mutation test case --- .../models/test_container_registries.py | 78 ++++++++++++++++++- 1 file changed, 76 insertions(+), 2 deletions(-) diff --git a/tests/manager/models/test_container_registries.py b/tests/manager/models/test_container_registries.py index 1a615651d9..c501442fb2 100644 --- a/tests/manager/models/test_container_registries.py +++ b/tests/manager/models/test_container_registries.py @@ -318,11 +318,9 @@ async def test_harbor_read_project_quota( 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, @@ -347,3 +345,79 @@ async def test_harbor_read_project_quota( groupnode_query, variables=variables, context_value=context ) assert response["data"]["group_node"]["registry_quota"] == HARBOR_QUOTA_VALUE + + +@pytest.mark.asyncio +@pytest.mark.parametrize("extra_fixtures", FIXTURES_FOR_HARBOR_CRUD_TEST) +async def test_harbor_create_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 + + creation_query = """ + mutation ($scopie_id: ScopeField!, $quota: BigInt!) { + create_container_registry_quota(scope_id: $scopie_id, quota: $quota) { + ok + msg + } + } + """ + variables = { + "scopie_id": "project:00000000-0000-0000-0000-000000000000", + "quota": 100, + } + + # Normal case: create a new quota + with aioresponses() as mocked: + 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}) + + 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": -1}}], + ) + + put_quota_url = f"http://mock_registry/api/v2.0/quotas/{HARBOR_QUOTA_ID}" + mocked.put( + put_quota_url, + status=200, + ) + + response = await client.execute_async( + creation_query, variables=variables, context_value=context + ) + assert response["data"]["create_container_registry_quota"]["ok"] + assert response["data"]["create_container_registry_quota"]["msg"] == "success" + + # If the quota already exists, the mutation should fail + with aioresponses() as mocked: + 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}) + + 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": 100}}], + ) + + response = await client.execute_async( + creation_query, variables=variables, context_value=context + ) + assert not response["data"]["create_container_registry_quota"]["ok"] From 234d0c99782e55031f743424e596b0847f47615a Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Mon, 2 Dec 2024 01:41:23 +0000 Subject: [PATCH 39/75] feat: Add update, delete GQL mutation test cases --- .../models/test_container_registries.py | 163 +++++++++++++++++- 1 file changed, 157 insertions(+), 6 deletions(-) diff --git a/tests/manager/models/test_container_registries.py b/tests/manager/models/test_container_registries.py index c501442fb2..2644e2dcf7 100644 --- a/tests/manager/models/test_container_registries.py +++ b/tests/manager/models/test_container_registries.py @@ -368,16 +368,16 @@ async def test_harbor_create_project_quota( HARBOR_PROJECT_ID = "123" HARBOR_QUOTA_ID = 456 - creation_query = """ - mutation ($scopie_id: ScopeField!, $quota: BigInt!) { - create_container_registry_quota(scope_id: $scopie_id, quota: $quota) { + create_query = """ + mutation ($scope_id: ScopeField!, $quota: BigInt!) { + create_container_registry_quota(scope_id: $scope_id, quota: $quota) { ok msg } } """ variables = { - "scopie_id": "project:00000000-0000-0000-0000-000000000000", + "scope_id": "project:00000000-0000-0000-0000-000000000000", "quota": 100, } @@ -400,7 +400,7 @@ async def test_harbor_create_project_quota( ) response = await client.execute_async( - creation_query, variables=variables, context_value=context + create_query, variables=variables, context_value=context ) assert response["data"]["create_container_registry_quota"]["ok"] assert response["data"]["create_container_registry_quota"]["msg"] == "success" @@ -418,6 +418,157 @@ async def test_harbor_create_project_quota( ) response = await client.execute_async( - creation_query, variables=variables, context_value=context + create_query, variables=variables, context_value=context ) assert not response["data"]["create_container_registry_quota"]["ok"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("extra_fixtures", FIXTURES_FOR_HARBOR_CRUD_TEST) +async def test_harbor_update_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 + + update_query = """ + mutation ($scope_id: ScopeField!, $quota: BigInt!) { + update_container_registry_quota(scope_id: $scope_id, quota: $quota) { + ok + msg + } + } + """ + variables = { + "scope_id": "project:00000000-0000-0000-0000-000000000000", + "quota": 200, + } + + # Normal case: update quota + with aioresponses() as mocked: + 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}) + + 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": 100}}], + ) + + put_quota_url = f"http://mock_registry/api/v2.0/quotas/{HARBOR_QUOTA_ID}" + mocked.put( + put_quota_url, + status=200, + ) + + response = await client.execute_async( + update_query, variables=variables, context_value=context + ) + assert response["data"]["update_container_registry_quota"]["ok"] + assert response["data"]["update_container_registry_quota"]["msg"] == "success" + + # If the quota doesn't exist, the mutation should fail + with aioresponses() as mocked: + 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}) + + 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": -1}}], + ) + + response = await client.execute_async( + update_query, variables=variables, context_value=context + ) + assert not response["data"]["update_container_registry_quota"]["ok"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("extra_fixtures", FIXTURES_FOR_HARBOR_CRUD_TEST) +async def test_harbor_delete_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 + + delete_query = """ + mutation ($scope_id: ScopeField!) { + delete_container_registry_quota(scope_id: $scope_id) { + ok + msg + } + } + """ + variables = { + "scope_id": "project:00000000-0000-0000-0000-000000000000", + } + + # Normal case: update quota + with aioresponses() as mocked: + 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}) + + 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": 100}}], + ) + + put_quota_url = f"http://mock_registry/api/v2.0/quotas/{HARBOR_QUOTA_ID}" + mocked.put( + put_quota_url, + status=200, + ) + + response = await client.execute_async( + delete_query, variables=variables, context_value=context + ) + assert response["data"]["delete_container_registry_quota"]["ok"] + assert response["data"]["delete_container_registry_quota"]["msg"] == "success" + + # If the quota doesn't exist, the mutation should fail + with aioresponses() as mocked: + 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}) + + 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": -1}}], + ) + + response = await client.execute_async( + delete_query, variables=variables, context_value=context + ) + assert not response["data"]["delete_container_registry_quota"]["ok"] From d83b193bc090d446e41964e8314b37043760ac1c Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Mon, 2 Dec 2024 01:42:42 +0000 Subject: [PATCH 40/75] chore: fix typo --- 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 2644e2dcf7..b74a7e43c6 100644 --- a/tests/manager/models/test_container_registries.py +++ b/tests/manager/models/test_container_registries.py @@ -532,7 +532,7 @@ async def test_harbor_delete_project_quota( "scope_id": "project:00000000-0000-0000-0000-000000000000", } - # Normal case: update quota + # Normal case: delete quota with aioresponses() as mocked: 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}) From 0e52491bd07badd0fdc83c2a9dd1881d6e65449d Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Tue, 3 Dec 2024 05:46:58 +0000 Subject: [PATCH 41/75] feat: Add REST API `test_harbor_read_project_quota` test --- tests/manager/api/test_group.py | 91 +++++++++++++++++++++++++++++++++ 1 file changed, 91 insertions(+) create mode 100644 tests/manager/api/test_group.py diff --git a/tests/manager/api/test_group.py b/tests/manager/api/test_group.py new file mode 100644 index 0000000000..3577e6d2d7 --- /dev/null +++ b/tests/manager/api/test_group.py @@ -0,0 +1,91 @@ +from urllib.parse import urlencode + +import pytest +from aioresponses import aioresponses + +from ai.backend.manager.server import ( + database_ctx, + hook_plugin_ctx, + monitoring_ctx, + redis_ctx, + shared_config_ctx, +) + +FIXTURES_FOR_HARBOR_CRUD_TEST = [ + { + "container_registries": [ + { + "id": "00000000-0000-0000-0000-000000000000", + "type": "harbor2", + "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", + } + ], + }, +] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("extra_fixtures", FIXTURES_FOR_HARBOR_CRUD_TEST) +async def test_harbor_read_project_quota( + etcd_fixture, + database_fixture, + create_app_and_client, + get_headers, +): + app, client = await create_app_and_client( + [ + shared_config_ctx, + database_ctx, + monitoring_ctx, + hook_plugin_ctx, + redis_ctx, + ], + [".group", ".auth"], + ) + + HARBOR_PROJECT_ID = "123" + HARBOR_QUOTA_ID = 456 + HARBOR_QUOTA_VALUE = 1024 + + with aioresponses(passthrough=["http://127.0.0.1"]) as mocked: + 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}) + + 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}}], + ) + + url_path = "/group/registry-quota" + params = {"group_id": "00000000-0000-0000-0000-000000000000"} + query_string = urlencode(params) + full_url = f"{url_path}?{query_string}" + headers = get_headers("GET", full_url, b"") + + resp = await client.get(url_path, params=params, headers=headers) + assert resp.status == 200 From 0d7b6c7214438eaa89605821efc2f2b152869832 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Tue, 3 Dec 2024 06:11:36 +0000 Subject: [PATCH 42/75] chore: Rename variables --- tests/manager/api/test_group.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/tests/manager/api/test_group.py b/tests/manager/api/test_group.py index 3577e6d2d7..1df2d60c19 100644 --- a/tests/manager/api/test_group.py +++ b/tests/manager/api/test_group.py @@ -81,11 +81,10 @@ async def test_harbor_read_project_quota( payload=[{"id": HARBOR_QUOTA_ID, "hard": {"storage": HARBOR_QUOTA_VALUE}}], ) - url_path = "/group/registry-quota" + url = "/group/registry-quota" params = {"group_id": "00000000-0000-0000-0000-000000000000"} - query_string = urlencode(params) - full_url = f"{url_path}?{query_string}" + full_url = f"{url}?{urlencode(params)}" headers = get_headers("GET", full_url, b"") - resp = await client.get(url_path, params=params, headers=headers) + resp = await client.get(url, params=params, headers=headers) assert resp.status == 200 From d6625dc2824d69a8273a3f19122f5c9a93d3b607 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Tue, 3 Dec 2024 06:20:10 +0000 Subject: [PATCH 43/75] fix: Add `test_harbor_update_project_quota` test --- tests/manager/api/test_group.py | 70 +++++++++++++++++++++++++++++++++ 1 file changed, 70 insertions(+) diff --git a/tests/manager/api/test_group.py b/tests/manager/api/test_group.py index 1df2d60c19..a0bc0522a3 100644 --- a/tests/manager/api/test_group.py +++ b/tests/manager/api/test_group.py @@ -1,3 +1,4 @@ +import json from urllib.parse import urlencode import pytest @@ -88,3 +89,72 @@ async def test_harbor_read_project_quota( resp = await client.get(url, params=params, headers=headers) assert resp.status == 200 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("extra_fixtures", FIXTURES_FOR_HARBOR_CRUD_TEST) +async def test_harbor_update_project_quota( + etcd_fixture, + database_fixture, + create_app_and_client, + get_headers, +): + app, client = await create_app_and_client( + [ + shared_config_ctx, + database_ctx, + monitoring_ctx, + hook_plugin_ctx, + redis_ctx, + ], + [".group", ".auth"], + ) + + HARBOR_PROJECT_ID = "123" + HARBOR_QUOTA_ID = 456 + HARBOR_QUOTA_VALUE = 1024 + + # Normal case: update quota + with aioresponses(passthrough=["http://127.0.0.1"]) as mocked: + 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}) + + 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": 100}}], + ) + + put_quota_url = f"http://mock_registry/api/v2.0/quotas/{HARBOR_QUOTA_ID}" + mocked.put( + put_quota_url, + status=200, + ) + + url = "/group/registry-quota" + params = {"group_id": "00000000-0000-0000-0000-000000000000", "quota": HARBOR_QUOTA_VALUE} + req_bytes = json.dumps(params).encode() + headers = get_headers("PATCH", url, req_bytes) + + resp = await client.patch(url, data=req_bytes, headers=headers) + assert resp.status == 200 + + # If the quota doesn't exist, the mutation should fail + with aioresponses(passthrough=["http://127.0.0.1"]) as mocked: + 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}) + + 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": -1}}], + ) + url = "/group/registry-quota" + params = {"group_id": "00000000-0000-0000-0000-000000000000", "quota": HARBOR_QUOTA_VALUE} + req_bytes = json.dumps(params).encode() + headers = get_headers("PATCH", url, req_bytes) + + resp = await client.patch(url, data=req_bytes, headers=headers) + assert resp.status == 404 From 11969a56807c3454b744fd997ef0c99e417db711 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Tue, 3 Dec 2024 06:21:46 +0000 Subject: [PATCH 44/75] chore: Hoist the variable --- tests/manager/api/test_group.py | 14 +++++--------- 1 file changed, 5 insertions(+), 9 deletions(-) diff --git a/tests/manager/api/test_group.py b/tests/manager/api/test_group.py index a0bc0522a3..61dcb447de 100644 --- a/tests/manager/api/test_group.py +++ b/tests/manager/api/test_group.py @@ -114,6 +114,11 @@ async def test_harbor_update_project_quota( HARBOR_QUOTA_ID = 456 HARBOR_QUOTA_VALUE = 1024 + url = "/group/registry-quota" + params = {"group_id": "00000000-0000-0000-0000-000000000000", "quota": HARBOR_QUOTA_VALUE} + req_bytes = json.dumps(params).encode() + headers = get_headers("PATCH", url, req_bytes) + # Normal case: update quota with aioresponses(passthrough=["http://127.0.0.1"]) as mocked: get_project_id_url = "http://mock_registry/api/v2.0/projects/mock_project" @@ -132,11 +137,6 @@ async def test_harbor_update_project_quota( status=200, ) - url = "/group/registry-quota" - params = {"group_id": "00000000-0000-0000-0000-000000000000", "quota": HARBOR_QUOTA_VALUE} - req_bytes = json.dumps(params).encode() - headers = get_headers("PATCH", url, req_bytes) - resp = await client.patch(url, data=req_bytes, headers=headers) assert resp.status == 200 @@ -151,10 +151,6 @@ async def test_harbor_update_project_quota( status=200, payload=[{"id": HARBOR_QUOTA_ID, "hard": {"storage": -1}}], ) - url = "/group/registry-quota" - params = {"group_id": "00000000-0000-0000-0000-000000000000", "quota": HARBOR_QUOTA_VALUE} - req_bytes = json.dumps(params).encode() - headers = get_headers("PATCH", url, req_bytes) resp = await client.patch(url, data=req_bytes, headers=headers) assert resp.status == 404 From 9c374a71d40e88a26f9eb61ba32be3768243656a Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Tue, 3 Dec 2024 06:27:26 +0000 Subject: [PATCH 45/75] feat: Add `test_harbor_delete_project_quota` test --- tests/manager/api/test_group.py | 64 +++++++++++++++++++++++++++++++++ 1 file changed, 64 insertions(+) diff --git a/tests/manager/api/test_group.py b/tests/manager/api/test_group.py index 61dcb447de..ade352d760 100644 --- a/tests/manager/api/test_group.py +++ b/tests/manager/api/test_group.py @@ -154,3 +154,67 @@ async def test_harbor_update_project_quota( resp = await client.patch(url, data=req_bytes, headers=headers) assert resp.status == 404 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("extra_fixtures", FIXTURES_FOR_HARBOR_CRUD_TEST) +async def test_harbor_delete_project_quota( + etcd_fixture, + database_fixture, + create_app_and_client, + get_headers, +): + app, client = await create_app_and_client( + [ + shared_config_ctx, + database_ctx, + monitoring_ctx, + hook_plugin_ctx, + redis_ctx, + ], + [".group", ".auth"], + ) + + HARBOR_PROJECT_ID = "123" + HARBOR_QUOTA_ID = 456 + + url = "/group/registry-quota" + params = {"group_id": "00000000-0000-0000-0000-000000000000"} + req_bytes = json.dumps(params).encode() + headers = get_headers("DELETE", url, req_bytes) + + # Normal case: delete quota + with aioresponses(passthrough=["http://127.0.0.1"]) as mocked: + 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}) + + 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": 100}}], + ) + + put_quota_url = f"http://mock_registry/api/v2.0/quotas/{HARBOR_QUOTA_ID}" + mocked.put( + put_quota_url, + status=200, + ) + + resp = await client.delete(url, data=req_bytes, headers=headers) + assert resp.status == 200 + + # If the quota doesn't exist, the mutation should fail + with aioresponses(passthrough=["http://127.0.0.1"]) as mocked: + 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}) + + 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": -1}}], + ) + + resp = await client.delete(url, data=req_bytes, headers=headers) + assert resp.status == 404 From 771ceaf394b6a1c7f29b657549aefd1a17c7a750 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Tue, 3 Dec 2024 06:30:42 +0000 Subject: [PATCH 46/75] feat: Add `test_harbor_create_project_quota` test --- tests/manager/api/test_group.py | 65 +++++++++++++++++++++++++++++++++ 1 file changed, 65 insertions(+) diff --git a/tests/manager/api/test_group.py b/tests/manager/api/test_group.py index ade352d760..8033ae4e07 100644 --- a/tests/manager/api/test_group.py +++ b/tests/manager/api/test_group.py @@ -48,6 +48,71 @@ ] +@pytest.mark.asyncio +@pytest.mark.parametrize("extra_fixtures", FIXTURES_FOR_HARBOR_CRUD_TEST) +async def test_harbor_create_project_quota( + etcd_fixture, + database_fixture, + create_app_and_client, + get_headers, +): + app, client = await create_app_and_client( + [ + shared_config_ctx, + database_ctx, + monitoring_ctx, + hook_plugin_ctx, + redis_ctx, + ], + [".group", ".auth"], + ) + + HARBOR_PROJECT_ID = "123" + HARBOR_QUOTA_ID = 456 + HARBOR_QUOTA_VALUE = 1024 + + url = "/group/registry-quota" + params = {"group_id": "00000000-0000-0000-0000-000000000000", "quota": HARBOR_QUOTA_VALUE} + req_bytes = json.dumps(params).encode() + headers = get_headers("POST", url, req_bytes) + + # Normal case: create a new quota + with aioresponses(passthrough=["http://127.0.0.1"]) as mocked: + 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}) + + 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": -1}}], + ) + + put_quota_url = f"http://mock_registry/api/v2.0/quotas/{HARBOR_QUOTA_ID}" + mocked.put( + put_quota_url, + status=200, + ) + + resp = await client.post(url, data=req_bytes, headers=headers) + assert resp.status == 200 + + # If the quota already exists, the mutation should fail + with aioresponses(passthrough=["http://127.0.0.1"]) as mocked: + 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}) + + 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": 100}}], + ) + + resp = await client.post(url, data=req_bytes, headers=headers) + assert resp.status == 400 + + @pytest.mark.asyncio @pytest.mark.parametrize("extra_fixtures", FIXTURES_FOR_HARBOR_CRUD_TEST) async def test_harbor_read_project_quota( From 1fcb3f8d96798d1a317fe55791fe74ee6e17c398 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Tue, 3 Dec 2024 06:42:31 +0000 Subject: [PATCH 47/75] fix: Change the `test_harbor_read_project_quota` test location --- tests/manager/models/gql_models/test_group.py | 127 ++++++++++++++++++ .../models/test_container_registries.py | 53 -------- 2 files changed, 127 insertions(+), 53 deletions(-) create mode 100644 tests/manager/models/gql_models/test_group.py diff --git a/tests/manager/models/gql_models/test_group.py b/tests/manager/models/gql_models/test_group.py new file mode 100644 index 0000000000..9f9679e052 --- /dev/null +++ b/tests/manager/models/gql_models/test_group.py @@ -0,0 +1,127 @@ +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.models.gql import GraphQueryContext, Mutations, Queries +from ai.backend.manager.models.utils import ExtendedAsyncSAEngine +from ai.backend.manager.server import ( + database_ctx, +) + + +@pytest.fixture(scope="module") +def client() -> Client: + return Client(Schema(query=Queries, mutation=Mutations, auto_camelcase=False)) + + +def get_graphquery_context(database_engine: ExtendedAsyncSAEngine) -> GraphQueryContext: + return GraphQueryContext( + schema=None, # type: ignore + dataloader_manager=None, # type: ignore + local_config=None, # type: ignore + shared_config=None, # type: ignore + etcd=None, # type: ignore + user={"domain": "default", "role": "superadmin"}, + access_key="AKIAIOSFODNN7EXAMPLE", + db=database_engine, # type: ignore + redis_stat=None, # type: ignore + redis_image=None, # type: ignore + redis_live=None, # type: ignore + manager_status=None, # type: ignore + known_slot_types=None, # type: ignore + background_task_manager=None, # type: ignore + storage_manager=None, # type: ignore + registry=None, # type: ignore + idle_checker_host=None, # type: ignore + ) + + +FIXTURES_FOR_HARBOR_CRUD_TEST = [ + { + "container_registries": [ + { + "id": "00000000-0000-0000-0000-000000000000", + "type": "harbor2", + "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", + } + ], + }, +] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("extra_fixtures", FIXTURES_FOR_HARBOR_CRUD_TEST) +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: + 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}) + + 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 diff --git a/tests/manager/models/test_container_registries.py b/tests/manager/models/test_container_registries.py index b74a7e43c6..417ac48410 100644 --- a/tests/manager/models/test_container_registries.py +++ b/tests/manager/models/test_container_registries.py @@ -3,7 +3,6 @@ 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 @@ -295,58 +294,6 @@ async def test_delete_container_registry(client: Client, database_engine: Extend ] -@pytest.mark.asyncio -@pytest.mark.parametrize("extra_fixtures", FIXTURES_FOR_HARBOR_CRUD_TEST) -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: - 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}) - - 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 - - @pytest.mark.asyncio @pytest.mark.parametrize("extra_fixtures", FIXTURES_FOR_HARBOR_CRUD_TEST) async def test_harbor_create_project_quota( From e21327be45ac700ec750be064f289e8e7187d5c3 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Tue, 3 Dec 2024 06:46:38 +0000 Subject: [PATCH 48/75] fix: Change test code location --- .../gql_models/test_container_registries.py | 260 +++++++++++++++++ .../models/test_container_registries.py | 268 ------------------ 2 files changed, 260 insertions(+), 268 deletions(-) diff --git a/tests/manager/models/gql_models/test_container_registries.py b/tests/manager/models/gql_models/test_container_registries.py index fd61b6a936..e721298728 100644 --- a/tests/manager/models/gql_models/test_container_registries.py +++ b/tests/manager/models/gql_models/test_container_registries.py @@ -1,9 +1,15 @@ import pytest +from aioresponses import aioresponses from graphene import Schema from graphene.test import Client +from tests.manager.models.gql_models.test_group import FIXTURES_FOR_HARBOR_CRUD_TEST +from ai.backend.manager.api.context import RootContext 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, +) @pytest.fixture(scope="module") @@ -32,3 +38,257 @@ def get_graphquery_context(database_engine: ExtendedAsyncSAEngine) -> GraphQuery idle_checker_host=None, # type: ignore network_plugin_ctx=None, # type: ignore ) + + +@pytest.fixture(scope="module") +def client() -> Client: + return Client(Schema(query=Queries, mutation=Mutations, auto_camelcase=False)) + + +def get_graphquery_context(database_engine: ExtendedAsyncSAEngine) -> GraphQueryContext: + return GraphQueryContext( + schema=None, # type: ignore + dataloader_manager=None, # type: ignore + local_config=None, # type: ignore + shared_config=None, # type: ignore + etcd=None, # type: ignore + user={"domain": "default", "role": "superadmin"}, + access_key="AKIAIOSFODNN7EXAMPLE", + db=database_engine, # type: ignore + redis_stat=None, # type: ignore + redis_image=None, # type: ignore + redis_live=None, # type: ignore + manager_status=None, # type: ignore + known_slot_types=None, # type: ignore + background_task_manager=None, # type: ignore + storage_manager=None, # type: ignore + registry=None, # type: ignore + idle_checker_host=None, # type: ignore + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("extra_fixtures", FIXTURES_FOR_HARBOR_CRUD_TEST) +async def test_harbor_create_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 + + create_query = """ + mutation ($scope_id: ScopeField!, $quota: BigInt!) { + create_container_registry_quota(scope_id: $scope_id, quota: $quota) { + ok + msg + } + } + """ + variables = { + "scope_id": "project:00000000-0000-0000-0000-000000000000", + "quota": 100, + } + + # Normal case: create a new quota + with aioresponses() as mocked: + 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}) + + 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": -1}}], + ) + + put_quota_url = f"http://mock_registry/api/v2.0/quotas/{HARBOR_QUOTA_ID}" + mocked.put( + put_quota_url, + status=200, + ) + + response = await client.execute_async( + create_query, variables=variables, context_value=context + ) + assert response["data"]["create_container_registry_quota"]["ok"] + assert response["data"]["create_container_registry_quota"]["msg"] == "success" + + # If the quota already exists, the mutation should fail + with aioresponses() as mocked: + 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}) + + 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": 100}}], + ) + + response = await client.execute_async( + create_query, variables=variables, context_value=context + ) + assert not response["data"]["create_container_registry_quota"]["ok"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("extra_fixtures", FIXTURES_FOR_HARBOR_CRUD_TEST) +async def test_harbor_update_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 + + update_query = """ + mutation ($scope_id: ScopeField!, $quota: BigInt!) { + update_container_registry_quota(scope_id: $scope_id, quota: $quota) { + ok + msg + } + } + """ + variables = { + "scope_id": "project:00000000-0000-0000-0000-000000000000", + "quota": 200, + } + + # Normal case: update quota + with aioresponses() as mocked: + 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}) + + 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": 100}}], + ) + + put_quota_url = f"http://mock_registry/api/v2.0/quotas/{HARBOR_QUOTA_ID}" + mocked.put( + put_quota_url, + status=200, + ) + + response = await client.execute_async( + update_query, variables=variables, context_value=context + ) + assert response["data"]["update_container_registry_quota"]["ok"] + assert response["data"]["update_container_registry_quota"]["msg"] == "success" + + # If the quota doesn't exist, the mutation should fail + with aioresponses() as mocked: + 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}) + + 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": -1}}], + ) + + response = await client.execute_async( + update_query, variables=variables, context_value=context + ) + assert not response["data"]["update_container_registry_quota"]["ok"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("extra_fixtures", FIXTURES_FOR_HARBOR_CRUD_TEST) +async def test_harbor_delete_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 + + delete_query = """ + mutation ($scope_id: ScopeField!) { + delete_container_registry_quota(scope_id: $scope_id) { + ok + msg + } + } + """ + variables = { + "scope_id": "project:00000000-0000-0000-0000-000000000000", + } + + # Normal case: delete quota + with aioresponses() as mocked: + 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}) + + 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": 100}}], + ) + + put_quota_url = f"http://mock_registry/api/v2.0/quotas/{HARBOR_QUOTA_ID}" + mocked.put( + put_quota_url, + status=200, + ) + + response = await client.execute_async( + delete_query, variables=variables, context_value=context + ) + assert response["data"]["delete_container_registry_quota"]["ok"] + assert response["data"]["delete_container_registry_quota"]["msg"] == "success" + + # If the quota doesn't exist, the mutation should fail + with aioresponses() as mocked: + 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}) + + 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": -1}}], + ) + + response = await client.execute_async( + delete_query, variables=variables, context_value=context + ) + assert not response["data"]["delete_container_registry_quota"]["ok"] diff --git a/tests/manager/models/test_container_registries.py b/tests/manager/models/test_container_registries.py index 417ac48410..0da01fab80 100644 --- a/tests/manager/models/test_container_registries.py +++ b/tests/manager/models/test_container_registries.py @@ -1,15 +1,10 @@ import pytest -from aioresponses import aioresponses from graphene import Schema from graphene.test import Client -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 @@ -256,266 +251,3 @@ 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 - - -FIXTURES_FOR_HARBOR_CRUD_TEST = [ - { - "container_registries": [ - { - "id": "00000000-0000-0000-0000-000000000000", - "type": "harbor2", - "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", - } - ], - }, -] - - -@pytest.mark.asyncio -@pytest.mark.parametrize("extra_fixtures", FIXTURES_FOR_HARBOR_CRUD_TEST) -async def test_harbor_create_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 - - create_query = """ - mutation ($scope_id: ScopeField!, $quota: BigInt!) { - create_container_registry_quota(scope_id: $scope_id, quota: $quota) { - ok - msg - } - } - """ - variables = { - "scope_id": "project:00000000-0000-0000-0000-000000000000", - "quota": 100, - } - - # Normal case: create a new quota - with aioresponses() as mocked: - 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}) - - 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": -1}}], - ) - - put_quota_url = f"http://mock_registry/api/v2.0/quotas/{HARBOR_QUOTA_ID}" - mocked.put( - put_quota_url, - status=200, - ) - - response = await client.execute_async( - create_query, variables=variables, context_value=context - ) - assert response["data"]["create_container_registry_quota"]["ok"] - assert response["data"]["create_container_registry_quota"]["msg"] == "success" - - # If the quota already exists, the mutation should fail - with aioresponses() as mocked: - 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}) - - 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": 100}}], - ) - - response = await client.execute_async( - create_query, variables=variables, context_value=context - ) - assert not response["data"]["create_container_registry_quota"]["ok"] - - -@pytest.mark.asyncio -@pytest.mark.parametrize("extra_fixtures", FIXTURES_FOR_HARBOR_CRUD_TEST) -async def test_harbor_update_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 - - update_query = """ - mutation ($scope_id: ScopeField!, $quota: BigInt!) { - update_container_registry_quota(scope_id: $scope_id, quota: $quota) { - ok - msg - } - } - """ - variables = { - "scope_id": "project:00000000-0000-0000-0000-000000000000", - "quota": 200, - } - - # Normal case: update quota - with aioresponses() as mocked: - 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}) - - 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": 100}}], - ) - - put_quota_url = f"http://mock_registry/api/v2.0/quotas/{HARBOR_QUOTA_ID}" - mocked.put( - put_quota_url, - status=200, - ) - - response = await client.execute_async( - update_query, variables=variables, context_value=context - ) - assert response["data"]["update_container_registry_quota"]["ok"] - assert response["data"]["update_container_registry_quota"]["msg"] == "success" - - # If the quota doesn't exist, the mutation should fail - with aioresponses() as mocked: - 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}) - - 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": -1}}], - ) - - response = await client.execute_async( - update_query, variables=variables, context_value=context - ) - assert not response["data"]["update_container_registry_quota"]["ok"] - - -@pytest.mark.asyncio -@pytest.mark.parametrize("extra_fixtures", FIXTURES_FOR_HARBOR_CRUD_TEST) -async def test_harbor_delete_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 - - delete_query = """ - mutation ($scope_id: ScopeField!) { - delete_container_registry_quota(scope_id: $scope_id) { - ok - msg - } - } - """ - variables = { - "scope_id": "project:00000000-0000-0000-0000-000000000000", - } - - # Normal case: delete quota - with aioresponses() as mocked: - 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}) - - 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": 100}}], - ) - - put_quota_url = f"http://mock_registry/api/v2.0/quotas/{HARBOR_QUOTA_ID}" - mocked.put( - put_quota_url, - status=200, - ) - - response = await client.execute_async( - delete_query, variables=variables, context_value=context - ) - assert response["data"]["delete_container_registry_quota"]["ok"] - assert response["data"]["delete_container_registry_quota"]["msg"] == "success" - - # If the quota doesn't exist, the mutation should fail - with aioresponses() as mocked: - 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}) - - 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": -1}}], - ) - - response = await client.execute_async( - delete_query, variables=variables, context_value=context - ) - assert not response["data"]["delete_container_registry_quota"]["ok"] From 11a726300b990415b88ff177bf9d9c3609eb87af Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Tue, 3 Dec 2024 06:51:47 +0000 Subject: [PATCH 49/75] fix: Reuse `FIXTURES_FOR_HARBOR_CRUD_TEST` --- src/ai/backend/testutils/extra_fixtures.py | 34 +++++++++++++++++ tests/manager/api/test_group.py | 36 +----------------- tests/manager/models/gql_models/test_group.py | 37 +------------------ 3 files changed, 36 insertions(+), 71 deletions(-) create mode 100644 src/ai/backend/testutils/extra_fixtures.py diff --git a/src/ai/backend/testutils/extra_fixtures.py b/src/ai/backend/testutils/extra_fixtures.py new file mode 100644 index 0000000000..a7e1810877 --- /dev/null +++ b/src/ai/backend/testutils/extra_fixtures.py @@ -0,0 +1,34 @@ +FIXTURES_FOR_HARBOR_CRUD_TEST = [ + { + "container_registries": [ + { + "id": "00000000-0000-0000-0000-000000000000", + "type": "harbor2", + "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", + } + ], + }, +] diff --git a/tests/manager/api/test_group.py b/tests/manager/api/test_group.py index 8033ae4e07..20cb32c6a6 100644 --- a/tests/manager/api/test_group.py +++ b/tests/manager/api/test_group.py @@ -11,41 +11,7 @@ redis_ctx, shared_config_ctx, ) - -FIXTURES_FOR_HARBOR_CRUD_TEST = [ - { - "container_registries": [ - { - "id": "00000000-0000-0000-0000-000000000000", - "type": "harbor2", - "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", - } - ], - }, -] +from ai.backend.testutils.extra_fixtures import FIXTURES_FOR_HARBOR_CRUD_TEST @pytest.mark.asyncio diff --git a/tests/manager/models/gql_models/test_group.py b/tests/manager/models/gql_models/test_group.py index 9f9679e052..dab3a17302 100644 --- a/tests/manager/models/gql_models/test_group.py +++ b/tests/manager/models/gql_models/test_group.py @@ -10,6 +10,7 @@ from ai.backend.manager.server import ( database_ctx, ) +from ai.backend.testutils.extra_fixtures import FIXTURES_FOR_HARBOR_CRUD_TEST @pytest.fixture(scope="module") @@ -39,42 +40,6 @@ def get_graphquery_context(database_engine: ExtendedAsyncSAEngine) -> GraphQuery ) -FIXTURES_FOR_HARBOR_CRUD_TEST = [ - { - "container_registries": [ - { - "id": "00000000-0000-0000-0000-000000000000", - "type": "harbor2", - "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", - } - ], - }, -] - - @pytest.mark.asyncio @pytest.mark.parametrize("extra_fixtures", FIXTURES_FOR_HARBOR_CRUD_TEST) async def test_harbor_read_project_quota( From aa4bc878c6bdc4f187ac39c4f5054ccbf76a5693 Mon Sep 17 00:00:00 2001 From: Gyubong Date: Thu, 5 Dec 2024 10:22:44 +0900 Subject: [PATCH 50/75] fix: Broken CI --- .../gql_models/test_container_registries.py | 27 ------------------- 1 file changed, 27 deletions(-) diff --git a/tests/manager/models/gql_models/test_container_registries.py b/tests/manager/models/gql_models/test_container_registries.py index e721298728..df3a5013c5 100644 --- a/tests/manager/models/gql_models/test_container_registries.py +++ b/tests/manager/models/gql_models/test_container_registries.py @@ -40,33 +40,6 @@ def get_graphquery_context(database_engine: ExtendedAsyncSAEngine) -> GraphQuery ) -@pytest.fixture(scope="module") -def client() -> Client: - return Client(Schema(query=Queries, mutation=Mutations, auto_camelcase=False)) - - -def get_graphquery_context(database_engine: ExtendedAsyncSAEngine) -> GraphQueryContext: - return GraphQueryContext( - schema=None, # type: ignore - dataloader_manager=None, # type: ignore - local_config=None, # type: ignore - shared_config=None, # type: ignore - etcd=None, # type: ignore - user={"domain": "default", "role": "superadmin"}, - access_key="AKIAIOSFODNN7EXAMPLE", - db=database_engine, # type: ignore - redis_stat=None, # type: ignore - redis_image=None, # type: ignore - redis_live=None, # type: ignore - manager_status=None, # type: ignore - known_slot_types=None, # type: ignore - background_task_manager=None, # type: ignore - storage_manager=None, # type: ignore - registry=None, # type: ignore - idle_checker_host=None, # type: ignore - ) - - @pytest.mark.asyncio @pytest.mark.parametrize("extra_fixtures", FIXTURES_FOR_HARBOR_CRUD_TEST) async def test_harbor_create_project_quota( From 83332c842ca0367ec0fb9c026e3478388c1f6789 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Fri, 6 Dec 2024 07:48:23 +0000 Subject: [PATCH 51/75] refactor: Add `test_case` parametrize annotation --- .../gql_models/test_container_registries.py | 211 +++++++++++------- 1 file changed, 125 insertions(+), 86 deletions(-) diff --git a/tests/manager/models/gql_models/test_container_registries.py b/tests/manager/models/gql_models/test_container_registries.py index df3a5013c5..b8c7cf0a3c 100644 --- a/tests/manager/models/gql_models/test_container_registries.py +++ b/tests/manager/models/gql_models/test_container_registries.py @@ -42,8 +42,39 @@ def get_graphquery_context(database_engine: ExtendedAsyncSAEngine) -> GraphQuery @pytest.mark.asyncio @pytest.mark.parametrize("extra_fixtures", FIXTURES_FOR_HARBOR_CRUD_TEST) +@pytest.mark.parametrize( + "test_case", + [ + { + "mock_harbor_responses": { + "get_project_id": {"project_id": "1"}, + "get_quotas": [ + { + "id": 1, + "hard": {"storage": -1}, + } + ], + }, + "expected": True, + }, + { + "mock_harbor_responses": { + "get_project_id": {"project_id": "1"}, + "get_quotas": [ + { + "id": 1, + "hard": {"storage": 100}, + } + ], + }, + "expected": False, + }, + ], + ids=["Normal case", "Project Quota already exist"], +) async def test_harbor_create_project_quota( client: Client, + test_case, database_fixture, create_app_and_client, ): @@ -57,10 +88,6 @@ async def test_harbor_create_project_quota( 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 - create_query = """ mutation ($scope_id: ScopeField!, $quota: BigInt!) { create_container_registry_quota(scope_id: $scope_id, quota: $quota) { @@ -74,19 +101,22 @@ async def test_harbor_create_project_quota( "quota": 100, } - # Normal case: create a new quota + mock_harbor_responses = test_case["mock_harbor_responses"] + with aioresponses() as mocked: 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}) + mocked.get(get_project_id_url, status=200, payload=mock_harbor_responses["get_project_id"]) - get_quota_url = f"http://mock_registry/api/v2.0/quotas?reference=project&reference_id={HARBOR_PROJECT_ID}" + harbor_project_id = mock_harbor_responses["get_project_id"]["project_id"] + get_quotas_url = f"http://mock_registry/api/v2.0/quotas?reference=project&reference_id={harbor_project_id}" mocked.get( - get_quota_url, + get_quotas_url, status=200, - payload=[{"id": HARBOR_QUOTA_ID, "hard": {"storage": -1}}], + payload=mock_harbor_responses["get_quotas"], ) - put_quota_url = f"http://mock_registry/api/v2.0/quotas/{HARBOR_QUOTA_ID}" + harbor_quota_id = mock_harbor_responses["get_quotas"][0]["id"] + put_quota_url = f"http://mock_registry/api/v2.0/quotas/{harbor_quota_id}" mocked.put( put_quota_url, status=200, @@ -95,31 +125,45 @@ async def test_harbor_create_project_quota( response = await client.execute_async( create_query, variables=variables, context_value=context ) - assert response["data"]["create_container_registry_quota"]["ok"] - assert response["data"]["create_container_registry_quota"]["msg"] == "success" - - # If the quota already exists, the mutation should fail - with aioresponses() as mocked: - 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}) - 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": 100}}], - ) - - response = await client.execute_async( - create_query, variables=variables, context_value=context - ) - assert not response["data"]["create_container_registry_quota"]["ok"] + assert response["data"]["create_container_registry_quota"]["ok"] == test_case["expected"] @pytest.mark.asyncio @pytest.mark.parametrize("extra_fixtures", FIXTURES_FOR_HARBOR_CRUD_TEST) +@pytest.mark.parametrize( + "test_case", + [ + { + "mock_harbor_responses": { + "get_project_id": {"project_id": "1"}, + "get_quotas": [ + { + "id": 1, + "hard": {"storage": 100}, + } + ], + }, + "expected": True, + }, + { + "mock_harbor_responses": { + "get_project_id": {"project_id": "1"}, + "get_quotas": [ + { + "id": 1, + "hard": {"storage": -1}, + } + ], + }, + "expected": False, + }, + ], + ids=["Normal case", "Project Quota not found"], +) async def test_harbor_update_project_quota( client: Client, + test_case, database_fixture, create_app_and_client, ): @@ -133,10 +177,6 @@ async def test_harbor_update_project_quota( 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 - update_query = """ mutation ($scope_id: ScopeField!, $quota: BigInt!) { update_container_registry_quota(scope_id: $scope_id, quota: $quota) { @@ -150,19 +190,23 @@ async def test_harbor_update_project_quota( "quota": 200, } - # Normal case: update quota + mock_harbor_responses = test_case["mock_harbor_responses"] + with aioresponses() as mocked: 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}) + mocked.get(get_project_id_url, status=200, payload=mock_harbor_responses["get_project_id"]) + + harbor_project_id = mock_harbor_responses["get_project_id"]["project_id"] - get_quota_url = f"http://mock_registry/api/v2.0/quotas?reference=project&reference_id={HARBOR_PROJECT_ID}" + get_quotas_url = f"http://mock_registry/api/v2.0/quotas?reference=project&reference_id={harbor_project_id}" mocked.get( - get_quota_url, + get_quotas_url, status=200, - payload=[{"id": HARBOR_QUOTA_ID, "hard": {"storage": 100}}], + payload=mock_harbor_responses["get_quotas"], ) - put_quota_url = f"http://mock_registry/api/v2.0/quotas/{HARBOR_QUOTA_ID}" + harbor_quota_id = mock_harbor_responses["get_quotas"][0]["id"] + put_quota_url = f"http://mock_registry/api/v2.0/quotas/{harbor_quota_id}" mocked.put( put_quota_url, status=200, @@ -171,31 +215,44 @@ async def test_harbor_update_project_quota( response = await client.execute_async( update_query, variables=variables, context_value=context ) - assert response["data"]["update_container_registry_quota"]["ok"] - assert response["data"]["update_container_registry_quota"]["msg"] == "success" - - # If the quota doesn't exist, the mutation should fail - with aioresponses() as mocked: - 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}) - - 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": -1}}], - ) - - response = await client.execute_async( - update_query, variables=variables, context_value=context - ) - assert not response["data"]["update_container_registry_quota"]["ok"] + assert response["data"]["update_container_registry_quota"]["ok"] == test_case["expected"] @pytest.mark.asyncio @pytest.mark.parametrize("extra_fixtures", FIXTURES_FOR_HARBOR_CRUD_TEST) +@pytest.mark.parametrize( + "test_case", + [ + { + "mock_harbor_responses": { + "get_project_id": {"project_id": "1"}, + "get_quotas": [ + { + "id": 1, + "hard": {"storage": 100}, + } + ], + }, + "expected": True, + }, + { + "mock_harbor_responses": { + "get_project_id": {"project_id": "1"}, + "get_quotas": [ + { + "id": 1, + "hard": {"storage": -1}, + } + ], + }, + "expected": False, + }, + ], + ids=["Normal case", "Project Quota not found"], +) async def test_harbor_delete_project_quota( client: Client, + test_case, database_fixture, create_app_and_client, ): @@ -209,10 +266,6 @@ async def test_harbor_delete_project_quota( 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 - delete_query = """ mutation ($scope_id: ScopeField!) { delete_container_registry_quota(scope_id: $scope_id) { @@ -225,19 +278,23 @@ async def test_harbor_delete_project_quota( "scope_id": "project:00000000-0000-0000-0000-000000000000", } - # Normal case: delete quota + mock_harbor_responses = test_case["mock_harbor_responses"] + with aioresponses() as mocked: 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}) + mocked.get(get_project_id_url, status=200, payload=mock_harbor_responses["get_project_id"]) - get_quota_url = f"http://mock_registry/api/v2.0/quotas?reference=project&reference_id={HARBOR_PROJECT_ID}" + harbor_project_id = mock_harbor_responses["get_project_id"]["project_id"] + + get_quotas_url = f"http://mock_registry/api/v2.0/quotas?reference=project&reference_id={harbor_project_id}" mocked.get( - get_quota_url, + get_quotas_url, status=200, - payload=[{"id": HARBOR_QUOTA_ID, "hard": {"storage": 100}}], + payload=mock_harbor_responses["get_quotas"], ) - put_quota_url = f"http://mock_registry/api/v2.0/quotas/{HARBOR_QUOTA_ID}" + harbor_quota_id = mock_harbor_responses["get_quotas"][0]["id"] + put_quota_url = f"http://mock_registry/api/v2.0/quotas/{harbor_quota_id}" mocked.put( put_quota_url, status=200, @@ -246,22 +303,4 @@ async def test_harbor_delete_project_quota( response = await client.execute_async( delete_query, variables=variables, context_value=context ) - assert response["data"]["delete_container_registry_quota"]["ok"] - assert response["data"]["delete_container_registry_quota"]["msg"] == "success" - - # If the quota doesn't exist, the mutation should fail - with aioresponses() as mocked: - 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}) - - 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": -1}}], - ) - - response = await client.execute_async( - delete_query, variables=variables, context_value=context - ) - assert not response["data"]["delete_container_registry_quota"]["ok"] + assert response["data"]["delete_container_registry_quota"]["ok"] == test_case["expected"] From d46bd31edf07bffdea79f7c72d6008a8a7b9b038 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Sun, 8 Dec 2024 04:25:31 +0000 Subject: [PATCH 52/75] refactor: `test_group` test cases using `test_case` --- tests/manager/api/test_group.py | 240 +++++++++++++++++++++----------- 1 file changed, 160 insertions(+), 80 deletions(-) diff --git a/tests/manager/api/test_group.py b/tests/manager/api/test_group.py index 20cb32c6a6..2a18c69a99 100644 --- a/tests/manager/api/test_group.py +++ b/tests/manager/api/test_group.py @@ -16,7 +16,38 @@ @pytest.mark.asyncio @pytest.mark.parametrize("extra_fixtures", FIXTURES_FOR_HARBOR_CRUD_TEST) +@pytest.mark.parametrize( + "test_case", + [ + { + "mock_harbor_responses": { + "get_project_id": {"project_id": "1"}, + "get_quotas": [ + { + "id": 1, + "hard": {"storage": -1}, + } + ], + }, + "expected_code": 200, + }, + { + "mock_harbor_responses": { + "get_project_id": {"project_id": "1"}, + "get_quotas": [ + { + "id": 1, + "hard": {"storage": 100}, + } + ], + }, + "expected_code": 400, + }, + ], + ids=["Normal case", "Project Quota already exist"], +) async def test_harbor_create_project_quota( + test_case, etcd_fixture, database_fixture, create_app_and_client, @@ -33,55 +64,74 @@ async def test_harbor_create_project_quota( [".group", ".auth"], ) - HARBOR_PROJECT_ID = "123" - HARBOR_QUOTA_ID = 456 - HARBOR_QUOTA_VALUE = 1024 + mock_harbor_responses = test_case["mock_harbor_responses"] url = "/group/registry-quota" - params = {"group_id": "00000000-0000-0000-0000-000000000000", "quota": HARBOR_QUOTA_VALUE} + params = {"group_id": "00000000-0000-0000-0000-000000000000", "quota": 100} req_bytes = json.dumps(params).encode() headers = get_headers("POST", url, req_bytes) - # Normal case: create a new quota with aioresponses(passthrough=["http://127.0.0.1"]) as mocked: 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}) + mocked.get( + get_project_id_url, + status=200, + payload=mock_harbor_responses["get_project_id"], + ) - get_quota_url = f"http://mock_registry/api/v2.0/quotas?reference=project&reference_id={HARBOR_PROJECT_ID}" + harbor_project_id = mock_harbor_responses["get_project_id"]["project_id"] + 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": -1}}], + payload=mock_harbor_responses["get_quotas"], ) - put_quota_url = f"http://mock_registry/api/v2.0/quotas/{HARBOR_QUOTA_ID}" + harbor_quota_id = mock_harbor_responses["get_quotas"][0]["id"] + put_quota_url = f"http://mock_registry/api/v2.0/quotas/{harbor_quota_id}" mocked.put( put_quota_url, status=200, ) resp = await client.post(url, data=req_bytes, headers=headers) - assert resp.status == 200 - - # If the quota already exists, the mutation should fail - with aioresponses(passthrough=["http://127.0.0.1"]) as mocked: - 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}) - - 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": 100}}], - ) - - resp = await client.post(url, data=req_bytes, headers=headers) - assert resp.status == 400 + assert resp.status == test_case["expected_code"] @pytest.mark.asyncio @pytest.mark.parametrize("extra_fixtures", FIXTURES_FOR_HARBOR_CRUD_TEST) +@pytest.mark.parametrize( + "test_case", + [ + { + "mock_harbor_responses": { + "get_project_id": {"project_id": "1"}, + "get_quotas": [ + { + "id": 1, + "hard": {"storage": 100}, + } + ], + }, + "expected_code": 200, + }, + { + "mock_harbor_responses": { + "get_project_id": {"project_id": "1"}, + "get_quotas": [ + { + "id": 1, + "hard": {"storage": -1}, + } + ], + }, + "expected_code": 404, + }, + ], + ids=["Normal case", "Project Quota doesn't exist"], +) async def test_harbor_read_project_quota( + test_case, etcd_fixture, database_fixture, create_app_and_client, @@ -98,19 +148,18 @@ async def test_harbor_read_project_quota( [".group", ".auth"], ) - HARBOR_PROJECT_ID = "123" - HARBOR_QUOTA_ID = 456 - HARBOR_QUOTA_VALUE = 1024 + mock_harbor_responses = test_case["mock_harbor_responses"] with aioresponses(passthrough=["http://127.0.0.1"]) as mocked: 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}) + mocked.get(get_project_id_url, status=200, payload=mock_harbor_responses["get_project_id"]) + harbor_project_id = mock_harbor_responses["get_project_id"]["project_id"] - get_quota_url = f"http://mock_registry/api/v2.0/quotas?reference=project&reference_id={HARBOR_PROJECT_ID}" + 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}}], + payload=mock_harbor_responses["get_quotas"], ) url = "/group/registry-quota" @@ -119,12 +168,43 @@ async def test_harbor_read_project_quota( headers = get_headers("GET", full_url, b"") resp = await client.get(url, params=params, headers=headers) - assert resp.status == 200 + assert resp.status == test_case["expected_code"] @pytest.mark.asyncio @pytest.mark.parametrize("extra_fixtures", FIXTURES_FOR_HARBOR_CRUD_TEST) +@pytest.mark.parametrize( + "test_case", + [ + { + "mock_harbor_responses": { + "get_project_id": {"project_id": "1"}, + "get_quotas": [ + { + "id": 1, + "hard": {"storage": 100}, + } + ], + }, + "expected_code": 200, + }, + { + "mock_harbor_responses": { + "get_project_id": {"project_id": "1"}, + "get_quotas": [ + { + "id": 1, + "hard": {"storage": -1}, + } + ], + }, + "expected_code": 404, + }, + ], + ids=["Normal case", "Project Quota not found"], +) async def test_harbor_update_project_quota( + test_case, etcd_fixture, database_fixture, create_app_and_client, @@ -141,55 +221,70 @@ async def test_harbor_update_project_quota( [".group", ".auth"], ) - HARBOR_PROJECT_ID = "123" - HARBOR_QUOTA_ID = 456 - HARBOR_QUOTA_VALUE = 1024 + mock_harbor_responses = test_case["mock_harbor_responses"] url = "/group/registry-quota" - params = {"group_id": "00000000-0000-0000-0000-000000000000", "quota": HARBOR_QUOTA_VALUE} + params = {"group_id": "00000000-0000-0000-0000-000000000000", "quota": 200} req_bytes = json.dumps(params).encode() headers = get_headers("PATCH", url, req_bytes) - # Normal case: update quota with aioresponses(passthrough=["http://127.0.0.1"]) as mocked: 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}) + mocked.get(get_project_id_url, status=200, payload=mock_harbor_responses["get_project_id"]) + harbor_project_id = mock_harbor_responses["get_project_id"]["project_id"] - get_quota_url = f"http://mock_registry/api/v2.0/quotas?reference=project&reference_id={HARBOR_PROJECT_ID}" + 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": 100}}], + payload=mock_harbor_responses["get_quotas"], ) + harbor_quota_id = mock_harbor_responses["get_quotas"][0]["id"] - put_quota_url = f"http://mock_registry/api/v2.0/quotas/{HARBOR_QUOTA_ID}" + put_quota_url = f"http://mock_registry/api/v2.0/quotas/{harbor_quota_id}" mocked.put( put_quota_url, status=200, ) resp = await client.patch(url, data=req_bytes, headers=headers) - assert resp.status == 200 - - # If the quota doesn't exist, the mutation should fail - with aioresponses(passthrough=["http://127.0.0.1"]) as mocked: - 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}) - - 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": -1}}], - ) - - resp = await client.patch(url, data=req_bytes, headers=headers) - assert resp.status == 404 + assert resp.status == test_case["expected_code"] @pytest.mark.asyncio @pytest.mark.parametrize("extra_fixtures", FIXTURES_FOR_HARBOR_CRUD_TEST) +@pytest.mark.parametrize( + "test_case", + [ + { + "mock_harbor_responses": { + "get_project_id": {"project_id": "1"}, + "get_quotas": [ + { + "id": 1, + "hard": {"storage": 100}, + } + ], + }, + "expected_code": 200, + }, + { + "mock_harbor_responses": { + "get_project_id": {"project_id": "1"}, + "get_quotas": [ + { + "id": 1, + "hard": {"storage": -1}, + } + ], + }, + "expected_code": 404, + }, + ], + ids=["Normal case", "Project Quota not found"], +) async def test_harbor_delete_project_quota( + test_case, etcd_fixture, database_fixture, create_app_and_client, @@ -206,46 +301,31 @@ async def test_harbor_delete_project_quota( [".group", ".auth"], ) - HARBOR_PROJECT_ID = "123" - HARBOR_QUOTA_ID = 456 + mock_harbor_responses = test_case["mock_harbor_responses"] url = "/group/registry-quota" params = {"group_id": "00000000-0000-0000-0000-000000000000"} req_bytes = json.dumps(params).encode() headers = get_headers("DELETE", url, req_bytes) - # Normal case: delete quota with aioresponses(passthrough=["http://127.0.0.1"]) as mocked: 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}) + mocked.get(get_project_id_url, status=200, payload=mock_harbor_responses["get_project_id"]) + harbor_project_id = mock_harbor_responses["get_project_id"]["project_id"] - get_quota_url = f"http://mock_registry/api/v2.0/quotas?reference=project&reference_id={HARBOR_PROJECT_ID}" + 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": 100}}], + payload=mock_harbor_responses["get_quotas"], ) + harbor_quota_id = mock_harbor_responses["get_quotas"][0]["id"] - put_quota_url = f"http://mock_registry/api/v2.0/quotas/{HARBOR_QUOTA_ID}" + put_quota_url = f"http://mock_registry/api/v2.0/quotas/{harbor_quota_id}" mocked.put( put_quota_url, status=200, ) resp = await client.delete(url, data=req_bytes, headers=headers) - assert resp.status == 200 - - # If the quota doesn't exist, the mutation should fail - with aioresponses(passthrough=["http://127.0.0.1"]) as mocked: - 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}) - - 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": -1}}], - ) - - resp = await client.delete(url, data=req_bytes, headers=headers) - assert resp.status == 404 + assert resp.status == test_case["expected_code"] From c501ffb9a071affda7b399334e0ed8842b49de3f Mon Sep 17 00:00:00 2001 From: jopemachine Date: Mon, 6 Jan 2025 02:33:59 +0000 Subject: [PATCH 53/75] fix: Broken CI --- src/ai/backend/client/func/group.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/ai/backend/client/func/group.py b/src/ai/backend/client/func/group.py index 7596aeec8f..ba48a95952 100644 --- a/src/ai/backend/client/func/group.py +++ b/src/ai/backend/client/func/group.py @@ -1,3 +1,4 @@ +import textwrap from typing import Any, Iterable, Optional, Sequence from ai.backend.client.output.fields import group_fields From da0d6e03811bcf91a40b42fcd818a58de4f14017 Mon Sep 17 00:00:00 2001 From: jopemachine Date: Mon, 6 Jan 2025 02:55:59 +0000 Subject: [PATCH 54/75] fix: Broken CI --- tests/manager/models/gql_models/test_group.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/manager/models/gql_models/test_group.py b/tests/manager/models/gql_models/test_group.py index dab3a17302..9aa5088ae8 100644 --- a/tests/manager/models/gql_models/test_group.py +++ b/tests/manager/models/gql_models/test_group.py @@ -37,6 +37,7 @@ def get_graphquery_context(database_engine: ExtendedAsyncSAEngine) -> GraphQuery storage_manager=None, # type: ignore registry=None, # type: ignore idle_checker_host=None, # type: ignore + network_plugin_ctx=None, # type: ignore ) From 3d80dc645fd93df0c2e2e683b129d670805f5373 Mon Sep 17 00:00:00 2001 From: jopemachine Date: Wed, 8 Jan 2025 03:42:17 +0000 Subject: [PATCH 55/75] fix: Update milestone --- docs/manager/graphql-reference/schema.graphql | 10 +++++----- src/ai/backend/manager/models/gql.py | 6 +++--- .../manager/models/gql_models/container_registry.py | 6 +++--- src/ai/backend/manager/models/gql_models/group.py | 2 +- 4 files changed, 12 insertions(+), 12 deletions(-) diff --git a/docs/manager/graphql-reference/schema.graphql b/docs/manager/graphql-reference/schema.graphql index acef847f73..3cc002ffe7 100644 --- a/docs/manager/graphql-reference/schema.graphql +++ b/docs/manager/graphql-reference/schema.graphql @@ -720,7 +720,7 @@ type GroupNode implements Node { container_registry: JSONString scaling_groups: [String] - """Added in 24.12.0.""" + """Added in 25.01.0.""" registry_quota: BigInt user_nodes(filter: String, order: String, offset: Int, before: String, after: String, first: Int, last: Int): UserConnection } @@ -2059,10 +2059,10 @@ type Mutations { """Added in 24.12.0""" create_container_registry_quota(quota: BigInt!, scope_id: ScopeField!): CreateContainerRegistryQuota - """Added in 24.12.0""" + """Added in 25.01.0.""" update_container_registry_quota(quota: BigInt!, scope_id: ScopeField!): UpdateContainerRegistryQuota - """Added in 24.12.0""" + """Added in 25.01.0.""" delete_container_registry_quota(scope_id: ScopeField!): DeleteContainerRegistryQuota """Deprecated since 24.09.0. use `CreateContainerRegistryNode` instead""" @@ -2973,13 +2973,13 @@ type CreateContainerRegistryQuota { msg: String } -"""Added in 24.12.0.""" +"""Added in 25.01.0.""" type UpdateContainerRegistryQuota { ok: Boolean msg: String } -"""Added in 24.12.0.""" +"""Added in 25.01.0.""" type DeleteContainerRegistryQuota { ok: Boolean msg: String diff --git a/src/ai/backend/manager/models/gql.py b/src/ai/backend/manager/models/gql.py index ac4716dc05..69d4393c1f 100644 --- a/src/ai/backend/manager/models/gql.py +++ b/src/ai/backend/manager/models/gql.py @@ -388,13 +388,13 @@ class Mutations(graphene.ObjectType): description="Added in 25.1.0." ) create_container_registry_quota = CreateContainerRegistryQuota.Field( - description="Added in 24.12.0" + description="Added in 25.01.0." ) update_container_registry_quota = UpdateContainerRegistryQuota.Field( - description="Added in 24.12.0" + description="Added in 25.01.0." ) delete_container_registry_quota = DeleteContainerRegistryQuota.Field( - description="Added in 24.12.0" + description="Added in 25.01.0." ) # Legacy mutations 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 2323c3177a..a66cd989b6 100644 --- a/src/ai/backend/manager/models/gql_models/container_registry.py +++ b/src/ai/backend/manager/models/gql_models/container_registry.py @@ -492,7 +492,7 @@ async def mutate( class CreateContainerRegistryQuota(graphene.Mutation): - """Added in 24.12.0.""" + """Added in 25.01.0.""" allowed_roles = ( UserRole.SUPERADMIN, @@ -524,7 +524,7 @@ async def mutate( class UpdateContainerRegistryQuota(graphene.Mutation): - """Added in 24.12.0.""" + """Added in 25.01.0.""" allowed_roles = ( UserRole.SUPERADMIN, @@ -556,7 +556,7 @@ async def mutate( class DeleteContainerRegistryQuota(graphene.Mutation): - """Added in 24.12.0.""" + """Added in 25.01.0.""" allowed_roles = ( UserRole.SUPERADMIN, diff --git a/src/ai/backend/manager/models/gql_models/group.py b/src/ai/backend/manager/models/gql_models/group.py index 6345e8bf44..4cbe0aff8a 100644 --- a/src/ai/backend/manager/models/gql_models/group.py +++ b/src/ai/backend/manager/models/gql_models/group.py @@ -122,7 +122,7 @@ class Meta: lambda: graphene.String, ) - registry_quota = BigInt(description="Added in 24.12.0.") + registry_quota = BigInt(description="Added in 25.01.0.") user_nodes = PaginatedConnectionField( UserConnection, From 305778d4bbfc309de450f35746e285622bf9949e Mon Sep 17 00:00:00 2001 From: jopemachine Date: Wed, 8 Jan 2025 05:43:17 +0000 Subject: [PATCH 56/75] fix: Update milestone --- docs/manager/graphql-reference/schema.graphql | 10 +++++----- src/ai/backend/manager/models/gql.py | 6 +++--- .../manager/models/gql_models/container_registry.py | 6 +++--- src/ai/backend/manager/models/gql_models/group.py | 2 +- 4 files changed, 12 insertions(+), 12 deletions(-) diff --git a/docs/manager/graphql-reference/schema.graphql b/docs/manager/graphql-reference/schema.graphql index 3cc002ffe7..6da63f2966 100644 --- a/docs/manager/graphql-reference/schema.graphql +++ b/docs/manager/graphql-reference/schema.graphql @@ -720,7 +720,7 @@ type GroupNode implements Node { container_registry: JSONString scaling_groups: [String] - """Added in 25.01.0.""" + """Added in 25.1.0.""" registry_quota: BigInt user_nodes(filter: String, order: String, offset: Int, before: String, after: String, first: Int, last: Int): UserConnection } @@ -2059,10 +2059,10 @@ type Mutations { """Added in 24.12.0""" create_container_registry_quota(quota: BigInt!, scope_id: ScopeField!): CreateContainerRegistryQuota - """Added in 25.01.0.""" + """Added in 25.1.0.""" update_container_registry_quota(quota: BigInt!, scope_id: ScopeField!): UpdateContainerRegistryQuota - """Added in 25.01.0.""" + """Added in 25.1.0.""" delete_container_registry_quota(scope_id: ScopeField!): DeleteContainerRegistryQuota """Deprecated since 24.09.0. use `CreateContainerRegistryNode` instead""" @@ -2973,13 +2973,13 @@ type CreateContainerRegistryQuota { msg: String } -"""Added in 25.01.0.""" +"""Added in 25.1.0.""" type UpdateContainerRegistryQuota { ok: Boolean msg: String } -"""Added in 25.01.0.""" +"""Added in 25.1.0.""" type DeleteContainerRegistryQuota { ok: Boolean msg: String diff --git a/src/ai/backend/manager/models/gql.py b/src/ai/backend/manager/models/gql.py index 69d4393c1f..7920e34f67 100644 --- a/src/ai/backend/manager/models/gql.py +++ b/src/ai/backend/manager/models/gql.py @@ -388,13 +388,13 @@ class Mutations(graphene.ObjectType): description="Added in 25.1.0." ) create_container_registry_quota = CreateContainerRegistryQuota.Field( - description="Added in 25.01.0." + description="Added in 25.1.0." ) update_container_registry_quota = UpdateContainerRegistryQuota.Field( - description="Added in 25.01.0." + description="Added in 25.1.0." ) delete_container_registry_quota = DeleteContainerRegistryQuota.Field( - description="Added in 25.01.0." + description="Added in 25.1.0." ) # Legacy mutations 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 a66cd989b6..4f151b64e2 100644 --- a/src/ai/backend/manager/models/gql_models/container_registry.py +++ b/src/ai/backend/manager/models/gql_models/container_registry.py @@ -492,7 +492,7 @@ async def mutate( class CreateContainerRegistryQuota(graphene.Mutation): - """Added in 25.01.0.""" + """Added in 25.1.0.""" allowed_roles = ( UserRole.SUPERADMIN, @@ -524,7 +524,7 @@ async def mutate( class UpdateContainerRegistryQuota(graphene.Mutation): - """Added in 25.01.0.""" + """Added in 25.1.0.""" allowed_roles = ( UserRole.SUPERADMIN, @@ -556,7 +556,7 @@ async def mutate( class DeleteContainerRegistryQuota(graphene.Mutation): - """Added in 25.01.0.""" + """Added in 25.1.0.""" allowed_roles = ( UserRole.SUPERADMIN, diff --git a/src/ai/backend/manager/models/gql_models/group.py b/src/ai/backend/manager/models/gql_models/group.py index 4cbe0aff8a..57176240ff 100644 --- a/src/ai/backend/manager/models/gql_models/group.py +++ b/src/ai/backend/manager/models/gql_models/group.py @@ -122,7 +122,7 @@ class Meta: lambda: graphene.String, ) - registry_quota = BigInt(description="Added in 25.01.0.") + registry_quota = BigInt(description="Added in 25.1.0.") user_nodes = PaginatedConnectionField( UserConnection, From 4860a60c1ecfe5ba889a26501ee865ed6649481f Mon Sep 17 00:00:00 2001 From: jopemachine Date: Fri, 10 Jan 2025 05:44:15 +0000 Subject: [PATCH 57/75] fix: Update milestone --- docs/manager/graphql-reference/schema.graphql | 4 ++-- src/ai/backend/manager/container_registry/__init__.py | 1 - .../manager/models/gql_models/container_registry_utils.py | 8 ++++++-- 3 files changed, 8 insertions(+), 5 deletions(-) diff --git a/docs/manager/graphql-reference/schema.graphql b/docs/manager/graphql-reference/schema.graphql index 6da63f2966..25167796c6 100644 --- a/docs/manager/graphql-reference/schema.graphql +++ b/docs/manager/graphql-reference/schema.graphql @@ -2056,7 +2056,7 @@ type Mutations { """Added in 25.1.0.""" delete_endpoint_auto_scaling_rule_node(id: String!): DeleteEndpointAutoScalingRuleNode - """Added in 24.12.0""" + """Added in 25.1.0.""" create_container_registry_quota(quota: BigInt!, scope_id: ScopeField!): CreateContainerRegistryQuota """Added in 25.1.0.""" @@ -2967,7 +2967,7 @@ type DeleteEndpointAutoScalingRuleNode { msg: String } -"""Added in 24.12.0.""" +"""Added in 25.1.0.""" type CreateContainerRegistryQuota { ok: Boolean msg: String diff --git a/src/ai/backend/manager/container_registry/__init__.py b/src/ai/backend/manager/container_registry/__init__.py index 7de7a7f4ae..9c58a9a7cf 100644 --- a/src/ai/backend/manager/container_registry/__init__.py +++ b/src/ai/backend/manager/container_registry/__init__.py @@ -7,7 +7,6 @@ from ai.backend.common.container_registry import ContainerRegistryType if TYPE_CHECKING: - from ..container_registry import ContainerRegistryRow from ..models.container_registry import ContainerRegistryRow from .base import BaseContainerRegistry 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 5b5a0c2b93..cca7bb943e 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 @@ -2,7 +2,7 @@ import logging import uuid -from typing import Any, TypedDict +from typing import TYPE_CHECKING, Any, TypedDict import aiohttp import aiohttp.client_exceptions @@ -29,6 +29,9 @@ log = BraceStyleAdapter(logging.getLogger(__spec__.name)) # type: ignore +if TYPE_CHECKING: + from ..container_registry import ContainerRegistryRow + class HarborQuotaInfo(TypedDict): previous_quota: int @@ -39,7 +42,6 @@ class HarborQuotaManager(aobject): """ Utility class for HarborV2 per-project Quota CRUD API. """ - from ..container_registry import ContainerRegistryRow db_sess: SASession scope_id: ScopeType @@ -56,6 +58,8 @@ def __init__(self, db_sess: SASession, scope_id: ScopeType): self.scope_id = scope_id async def __ainit__(self) -> None: + from ..container_registry import ContainerRegistryRow + assert isinstance(self.scope_id, ProjectScope) project_id = self.scope_id.project_id From 0dba50b7763c9aa4abbfe34204511fdac96fe02a Mon Sep 17 00:00:00 2001 From: jopemachine Date: Tue, 14 Jan 2025 05:55:06 +0000 Subject: [PATCH 58/75] fix: Remove cyclic import --- .../manager/models/gql_models/container_registry_utils.py | 3 +++ 1 file changed, 3 insertions(+) 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 cca7bb943e..1851172a0e 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 @@ -27,6 +27,9 @@ from ..group import GroupRow from ..rbac import ProjectScope, ScopeType +if TYPE_CHECKING: + from ...container_registry import ContainerRegistryRow + log = BraceStyleAdapter(logging.getLogger(__spec__.name)) # type: ignore if TYPE_CHECKING: From 2185cc9378aef61164731be2b12f0fb97f52775e Mon Sep 17 00:00:00 2001 From: jopemachine Date: Tue, 14 Jan 2025 06:23:22 +0000 Subject: [PATCH 59/75] docs: Update milestone --- docs/manager/graphql-reference/schema.graphql | 14 +++++++------- src/ai/backend/manager/models/gql.py | 6 +++--- .../models/gql_models/container_registry.py | 6 +++--- src/ai/backend/manager/models/gql_models/group.py | 2 +- 4 files changed, 14 insertions(+), 14 deletions(-) diff --git a/docs/manager/graphql-reference/schema.graphql b/docs/manager/graphql-reference/schema.graphql index 25167796c6..2db3ae19c4 100644 --- a/docs/manager/graphql-reference/schema.graphql +++ b/docs/manager/graphql-reference/schema.graphql @@ -720,7 +720,7 @@ type GroupNode implements Node { container_registry: JSONString scaling_groups: [String] - """Added in 25.1.0.""" + """Added in 25.2.0.""" registry_quota: BigInt user_nodes(filter: String, order: String, offset: Int, before: String, after: String, first: Int, last: Int): UserConnection } @@ -2056,13 +2056,13 @@ type Mutations { """Added in 25.1.0.""" delete_endpoint_auto_scaling_rule_node(id: String!): DeleteEndpointAutoScalingRuleNode - """Added in 25.1.0.""" + """Added in 25.2.0.""" create_container_registry_quota(quota: BigInt!, scope_id: ScopeField!): CreateContainerRegistryQuota - """Added in 25.1.0.""" + """Added in 25.2.0.""" update_container_registry_quota(quota: BigInt!, scope_id: ScopeField!): UpdateContainerRegistryQuota - """Added in 25.1.0.""" + """Added in 25.2.0.""" delete_container_registry_quota(scope_id: ScopeField!): DeleteContainerRegistryQuota """Deprecated since 24.09.0. use `CreateContainerRegistryNode` instead""" @@ -2967,19 +2967,19 @@ type DeleteEndpointAutoScalingRuleNode { msg: String } -"""Added in 25.1.0.""" +"""Added in 25.2.0.""" type CreateContainerRegistryQuota { ok: Boolean msg: String } -"""Added in 25.1.0.""" +"""Added in 25.2.0.""" type UpdateContainerRegistryQuota { ok: Boolean msg: String } -"""Added in 25.1.0.""" +"""Added in 25.2.0.""" type DeleteContainerRegistryQuota { ok: Boolean msg: String diff --git a/src/ai/backend/manager/models/gql.py b/src/ai/backend/manager/models/gql.py index 7920e34f67..74c7a0e4a6 100644 --- a/src/ai/backend/manager/models/gql.py +++ b/src/ai/backend/manager/models/gql.py @@ -388,13 +388,13 @@ class Mutations(graphene.ObjectType): description="Added in 25.1.0." ) create_container_registry_quota = CreateContainerRegistryQuota.Field( - description="Added in 25.1.0." + description="Added in 25.2.0." ) update_container_registry_quota = UpdateContainerRegistryQuota.Field( - description="Added in 25.1.0." + description="Added in 25.2.0." ) delete_container_registry_quota = DeleteContainerRegistryQuota.Field( - description="Added in 25.1.0." + description="Added in 25.2.0." ) # Legacy mutations 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 4f151b64e2..d1a2469434 100644 --- a/src/ai/backend/manager/models/gql_models/container_registry.py +++ b/src/ai/backend/manager/models/gql_models/container_registry.py @@ -492,7 +492,7 @@ async def mutate( class CreateContainerRegistryQuota(graphene.Mutation): - """Added in 25.1.0.""" + """Added in 25.2.0.""" allowed_roles = ( UserRole.SUPERADMIN, @@ -524,7 +524,7 @@ async def mutate( class UpdateContainerRegistryQuota(graphene.Mutation): - """Added in 25.1.0.""" + """Added in 25.2.0.""" allowed_roles = ( UserRole.SUPERADMIN, @@ -556,7 +556,7 @@ async def mutate( class DeleteContainerRegistryQuota(graphene.Mutation): - """Added in 25.1.0.""" + """Added in 25.2.0.""" allowed_roles = ( UserRole.SUPERADMIN, diff --git a/src/ai/backend/manager/models/gql_models/group.py b/src/ai/backend/manager/models/gql_models/group.py index 57176240ff..a35b3d1419 100644 --- a/src/ai/backend/manager/models/gql_models/group.py +++ b/src/ai/backend/manager/models/gql_models/group.py @@ -122,7 +122,7 @@ class Meta: lambda: graphene.String, ) - registry_quota = BigInt(description="Added in 25.1.0.") + registry_quota = BigInt(description="Added in 25.2.0.") user_nodes = PaginatedConnectionField( UserConnection, From 0800ad73dcf9edbd26345e015deaab33b6dc9a88 Mon Sep 17 00:00:00 2001 From: jopemachine Date: Fri, 24 Jan 2025 01:03:42 +0000 Subject: [PATCH 60/75] fix: Broken import --- tests/manager/models/gql_models/test_container_registries.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/manager/models/gql_models/test_container_registries.py b/tests/manager/models/gql_models/test_container_registries.py index b8c7cf0a3c..b0c32be06f 100644 --- a/tests/manager/models/gql_models/test_container_registries.py +++ b/tests/manager/models/gql_models/test_container_registries.py @@ -2,7 +2,6 @@ from aioresponses import aioresponses from graphene import Schema from graphene.test import Client -from tests.manager.models.gql_models.test_group import FIXTURES_FOR_HARBOR_CRUD_TEST from ai.backend.manager.api.context import RootContext from ai.backend.manager.models.gql import GraphQueryContext, Mutations, Queries @@ -10,6 +9,7 @@ from ai.backend.manager.server import ( database_ctx, ) +from ai.backend.testutils.extra_fixtures import FIXTURES_FOR_HARBOR_CRUD_TEST @pytest.fixture(scope="module") From f8da8e4a1d7e8e16331071e6608250a262d08228 Mon Sep 17 00:00:00 2001 From: jopemachine Date: Thu, 30 Jan 2025 10:52:57 +0000 Subject: [PATCH 61/75] chore: update api schema dump Co-authored-by: octodog --- docs/manager/rest-reference/openapi.json | 136 +++++++++++++++++++++++ 1 file changed, 136 insertions(+) diff --git a/docs/manager/rest-reference/openapi.json b/docs/manager/rest-reference/openapi.json index 248d84cf02..020c65e4a8 100644 --- a/docs/manager/rest-reference/openapi.json +++ b/docs/manager/rest-reference/openapi.json @@ -8147,6 +8147,142 @@ "description": "\n**Preconditions:**\n* Admin privilege required.\n* Manager status required: one of FROZEN, RUNNING\n" } }, + "/group/registry-quota": { + "post": { + "operationId": "group.create_registry_quota", + "tags": [ + "group" + ], + "responses": { + "200": { + "description": "Successful response" + } + }, + "security": [ + { + "TokenAuth": [] + } + ], + "requestBody": { + "content": { + "application/json": { + "schema": { + "type": "object", + "properties": { + "group_id": { + "type": "string" + }, + "quota": { + "type": "integer" + } + }, + "required": [ + "group_id", + "quota" + ] + }, + "examples": {} + } + } + }, + "parameters": [], + "description": "\n**Preconditions:**\n* Superadmin privilege required.\n* Manager status required: one of FROZEN, RUNNING\n" + }, + "get": { + "operationId": "group.read_registry_quota", + "tags": [ + "group" + ], + "responses": { + "200": { + "description": "Successful response" + } + }, + "security": [ + { + "TokenAuth": [] + } + ], + "parameters": [ + { + "name": "group_id", + "schema": { + "type": "string" + }, + "required": true, + "in": "query" + } + ], + "description": "\n**Preconditions:**\n* Superadmin privilege required.\n* Manager status required: one of FROZEN, RUNNING\n" + }, + "patch": { + "operationId": "group.update_registry_quota", + "tags": [ + "group" + ], + "responses": { + "200": { + "description": "Successful response" + } + }, + "security": [ + { + "TokenAuth": [] + } + ], + "requestBody": { + "content": { + "application/json": { + "schema": { + "type": "object", + "properties": { + "group_id": { + "type": "string" + }, + "quota": { + "type": "integer" + } + }, + "required": [ + "group_id", + "quota" + ] + }, + "examples": {} + } + } + }, + "parameters": [], + "description": "\n**Preconditions:**\n* Superadmin privilege required.\n* Manager status required: one of FROZEN, RUNNING\n" + }, + "delete": { + "operationId": "group.delete_registry_quota", + "tags": [ + "group" + ], + "responses": { + "200": { + "description": "Successful response" + } + }, + "security": [ + { + "TokenAuth": [] + } + ], + "parameters": [ + { + "name": "group_id", + "schema": { + "type": "string" + }, + "required": true, + "in": "query" + } + ], + "description": "\n**Preconditions:**\n* Superadmin privilege required.\n* Manager status required: one of FROZEN, RUNNING\n" + } + }, "/group-config/dotfiles": { "post": { "operationId": "group-config.create", From fd0ee6f965436a45fb5a5a57f31130fc5d73ead3 Mon Sep 17 00:00:00 2001 From: jopemachine Date: Fri, 31 Jan 2025 05:45:08 +0000 Subject: [PATCH 62/75] refactor: Implement `services_ctx` and refactoring using this --- src/ai/backend/manager/api/admin.py | 1 + src/ai/backend/manager/api/context.py | 2 + src/ai/backend/manager/api/group.py | 28 +- src/ai/backend/manager/api/services/BUILD | 1 + .../backend/manager/api/services/__init__.py | 0 src/ai/backend/manager/api/services/base.py | 25 ++ .../api/services/container_registries/BUILD | 1 + .../services/container_registries/__init__.py | 0 .../api/services/container_registries/base.py | 73 +++++ .../services/container_registries/harbor.py | 210 ++++++++++++++ src/ai/backend/manager/models/gql.py | 2 + .../models/gql_models/container_registry.py | 61 +++-- .../gql_models/container_registry_utils.py | 258 ------------------ .../manager/models/gql_models/group.py | 14 +- src/ai/backend/manager/server.py | 11 + .../gql_models/test_container_registries.py | 1 + .../test_container_registry_nodes_v2.py | 1 + tests/manager/models/gql_models/test_group.py | 1 + .../models/test_container_registries.py | 1 + 19 files changed, 381 insertions(+), 310 deletions(-) create mode 100644 src/ai/backend/manager/api/services/BUILD create mode 100644 src/ai/backend/manager/api/services/__init__.py create mode 100644 src/ai/backend/manager/api/services/base.py create mode 100644 src/ai/backend/manager/api/services/container_registries/BUILD create mode 100644 src/ai/backend/manager/api/services/container_registries/__init__.py create mode 100644 src/ai/backend/manager/api/services/container_registries/base.py create mode 100644 src/ai/backend/manager/api/services/container_registries/harbor.py delete mode 100644 src/ai/backend/manager/models/gql_models/container_registry_utils.py diff --git a/src/ai/backend/manager/api/admin.py b/src/ai/backend/manager/api/admin.py index de97f51e0c..5e59c349e8 100644 --- a/src/ai/backend/manager/api/admin.py +++ b/src/ai/backend/manager/api/admin.py @@ -85,6 +85,7 @@ async def _handle_gql_common(request: web.Request, params: Any) -> ExecutionResu manager_status=manager_status, known_slot_types=known_slot_types, background_task_manager=root_ctx.background_task_manager, + services_ctx=root_ctx.services_ctx, storage_manager=root_ctx.storage_manager, registry=root_ctx.registry, idle_checker_host=root_ctx.idle_checker_host, diff --git a/src/ai/backend/manager/api/context.py b/src/ai/backend/manager/api/context.py index 5d7cf4bb5d..cd9c925999 100644 --- a/src/ai/backend/manager/api/context.py +++ b/src/ai/backend/manager/api/context.py @@ -5,6 +5,7 @@ import attrs from ai.backend.common.metrics.metric import CommonMetricRegistry +from ai.backend.manager.api.services.base import ServicesContext from ai.backend.manager.plugin.network import NetworkPluginContext if TYPE_CHECKING: @@ -50,6 +51,7 @@ class RootContext(BaseContext): storage_manager: StorageSessionManager hook_plugin_ctx: HookPluginContext network_plugin_ctx: NetworkPluginContext + services_ctx: ServicesContext registry: AgentRegistry agent_cache: AgentRPCCache diff --git a/src/ai/backend/manager/api/group.py b/src/ai/backend/manager/api/group.py index 5434d6a88d..77ca2d864f 100644 --- a/src/ai/backend/manager/api/group.py +++ b/src/ai/backend/manager/api/group.py @@ -9,9 +9,6 @@ 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 ( - HarborQuotaManager, -) from ai.backend.manager.models.rbac import ProjectScope if TYPE_CHECKING: @@ -40,11 +37,8 @@ async def update_registry_quota(request: web.Request, params: Any) -> web.Respon scope_id = ProjectScope(project_id=group_id, domain_name=None) quota = int(params["quota"]) - async with root_ctx.db.begin_session() as db_sess: - manager = await HarborQuotaManager.new(db_sess, scope_id) - await manager.update(quota) - - return web.json_response({}) + await root_ctx.services_ctx.per_project_container_registries_quota.update(scope_id, quota) + return web.Response(status=204) @server_status_required(READ_ALLOWED) @@ -60,11 +54,8 @@ async def delete_registry_quota(request: web.Request, params: Any) -> web.Respon 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: - manager = await HarborQuotaManager.new(db_sess, scope_id) - await manager.delete() - - return web.json_response({}) + await root_ctx.services_ctx.per_project_container_registries_quota.delete(scope_id) + return web.Response(status=204) @server_status_required(READ_ALLOWED) @@ -82,11 +73,8 @@ async def create_registry_quota(request: web.Request, params: Any) -> web.Respon scope_id = ProjectScope(project_id=group_id, domain_name=None) quota = int(params["quota"]) - async with root_ctx.db.begin_session() as db_sess: - manager = await HarborQuotaManager.new(db_sess, scope_id) - await manager.create(quota) - - return web.json_response({}) + await root_ctx.services_ctx.per_project_container_registries_quota.create(scope_id, quota) + return web.Response(status=204) @server_status_required(READ_ALLOWED) @@ -102,9 +90,7 @@ async def read_registry_quota(request: web.Request, params: Any) -> web.Response 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: - manager = await HarborQuotaManager.new(db_sess, scope_id) - quota = await manager.read() + quota = await root_ctx.services_ctx.per_project_container_registries_quota.read(scope_id) return web.json_response({"result": quota}) diff --git a/src/ai/backend/manager/api/services/BUILD b/src/ai/backend/manager/api/services/BUILD new file mode 100644 index 0000000000..7357442404 --- /dev/null +++ b/src/ai/backend/manager/api/services/BUILD @@ -0,0 +1 @@ +python_sources(name="src") diff --git a/src/ai/backend/manager/api/services/__init__.py b/src/ai/backend/manager/api/services/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/src/ai/backend/manager/api/services/base.py b/src/ai/backend/manager/api/services/base.py new file mode 100644 index 0000000000..4b42e2c0b7 --- /dev/null +++ b/src/ai/backend/manager/api/services/base.py @@ -0,0 +1,25 @@ +from ai.backend.manager.models.utils import ExtendedAsyncSAEngine + +from .container_registries.base import PerProjectRegistryQuotaRepository +from .container_registries.harbor import ( + PerProjectContainerRegistryQuotaProtocol, + PerProjectContainerRegistryQuotaService, +) + + +class ServicesContext: + """ + In the API layer, requests are processed through the ServicesContext and + its subordinate layers, including the DB, Client, and Repository layers. + Each layer separates the responsibilities specific to its respective level. + """ + + db: ExtendedAsyncSAEngine + + def __init__(self, db: ExtendedAsyncSAEngine) -> None: + self.db = db + + @property + def per_project_container_registries_quota(self) -> PerProjectContainerRegistryQuotaProtocol: + repository = PerProjectRegistryQuotaRepository(db=self.db) + return PerProjectContainerRegistryQuotaService(repository=repository) diff --git a/src/ai/backend/manager/api/services/container_registries/BUILD b/src/ai/backend/manager/api/services/container_registries/BUILD new file mode 100644 index 0000000000..7357442404 --- /dev/null +++ b/src/ai/backend/manager/api/services/container_registries/BUILD @@ -0,0 +1 @@ +python_sources(name="src") diff --git a/src/ai/backend/manager/api/services/container_registries/__init__.py b/src/ai/backend/manager/api/services/container_registries/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/src/ai/backend/manager/api/services/container_registries/base.py b/src/ai/backend/manager/api/services/container_registries/base.py new file mode 100644 index 0000000000..ce0abd4dde --- /dev/null +++ b/src/ai/backend/manager/api/services/container_registries/base.py @@ -0,0 +1,73 @@ +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING + +import sqlalchemy as sa +from sqlalchemy.orm import load_only + +from ai.backend.logging import BraceStyleAdapter +from ai.backend.manager.api.exceptions import ( + ContainerRegistryNotFound, +) +from ai.backend.manager.models.container_registry import ContainerRegistryRow +from ai.backend.manager.models.group import GroupRow +from ai.backend.manager.models.rbac import ProjectScope +from ai.backend.manager.models.utils import ExtendedAsyncSAEngine + +if TYPE_CHECKING: + pass + +log = BraceStyleAdapter(logging.getLogger(__spec__.name)) + + +class PerProjectRegistryQuotaRepository: + """ """ + + def __init__(self, db: ExtendedAsyncSAEngine): + self.db = db + + @classmethod + def _is_valid_group_row(cls, 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 fetch_container_registry_row(self, scope_id: ProjectScope) -> ContainerRegistryRow: + async with self.db.begin_readonly_session() as db_sess: + 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 PerProjectRegistryQuotaRepository._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"], + ) + + 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 not found. (cr: {registry_name}, gr: {project})" + ) + + return registry diff --git a/src/ai/backend/manager/api/services/container_registries/harbor.py b/src/ai/backend/manager/api/services/container_registries/harbor.py new file mode 100644 index 0000000000..9f3b6b0300 --- /dev/null +++ b/src/ai/backend/manager/api/services/container_registries/harbor.py @@ -0,0 +1,210 @@ +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Any, Protocol, TypedDict + +import aiohttp +import yarl + +from ai.backend.common.container_registry import ContainerRegistryType +from ai.backend.logging import BraceStyleAdapter +from ai.backend.manager.api.exceptions import GenericBadRequest, InternalServerError, ObjectNotFound +from ai.backend.manager.api.services.container_registries.base import ( + ContainerRegistryRow, + PerProjectRegistryQuotaRepository, +) +from ai.backend.manager.models.rbac import ProjectScope + +if TYPE_CHECKING: + pass + +log = BraceStyleAdapter(logging.getLogger(__spec__.name)) + + +__all__ = ( + "AbstractPerProjectRegistryQuotaClient", + "PerProjectHarborQuotaClient", + "PerProjectRegistryQuotaRepository", + "PerProjectContainerRegistryQuotaProtocol", + "PerProjectContainerRegistryQuotaService", +) + + +class HarborProjectQuotaInfo(TypedDict): + previous_quota: int + quota_id: int + + +class AbstractPerProjectRegistryQuotaClient(Protocol): + async def create(self, registry_row: ContainerRegistryRow, quota: int) -> None: ... + async def update(self, registry_row: ContainerRegistryRow, quota: int) -> None: ... + async def delete(self, registry_row: ContainerRegistryRow) -> None: ... + async def read(self, registry_row: ContainerRegistryRow) -> int: ... + + +class PerProjectHarborQuotaClient(AbstractPerProjectRegistryQuotaClient): + async def _get_harbor_project_id( + self, + sess: aiohttp.ClientSession, + registry_row: ContainerRegistryRow, + rqst_args: dict[str, Any], + ) -> str: + get_project_id_api = ( + yarl.URL(registry_row.url) / "api" / "v2.0" / "projects" / registry_row.project + ) + + 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 str(harbor_project_id) + + async def _get_quota_info( + self, + sess: aiohttp.ClientSession, + registry_row: ContainerRegistryRow, + rqst_args: dict[str, Any], + ) -> HarborProjectQuotaInfo: + harbor_project_id = await self._get_harbor_project_id(sess, registry_row, rqst_args) + get_quota_id_api = (yarl.URL(registry_row.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: + 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") + if len(res) > 1: + raise InternalServerError( + f"Multiple quota entities found. (project_id: {harbor_project_id})" + ) + + previous_quota = res[0]["hard"]["storage"] + quota_id = res[0]["id"] + return HarborProjectQuotaInfo(previous_quota=previous_quota, quota_id=quota_id) + + async def read(self, registry_row: ContainerRegistryRow) -> int: + ssl_verify = True + connector = aiohttp.TCPConnector(ssl=ssl_verify) + async with aiohttp.ClientSession(connector=connector) as sess: + rqst_args: dict[str, Any] = {} + + quota_info = await self._get_quota_info(sess, registry_row, rqst_args) + previous_quota = quota_info["previous_quota"] + if previous_quota == -1: + raise ObjectNotFound(object_name="quota entity") + return previous_quota + + async def create(self, registry_row: ContainerRegistryRow, quota: int) -> None: + ssl_verify = True + 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_row.username, registry_row.password) + + quota_info = await self._get_quota_info(sess, registry_row, rqst_args) + previous_quota, quota_id = quota_info["previous_quota"], quota_info["quota_id"] + + if previous_quota > 0: + raise GenericBadRequest("Quota limit already exists!") + + put_quota_api = yarl.URL(registry_row.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: + log.error(f"Failed to create quota! response: {resp}") + raise InternalServerError(f"Failed to create quota! response: {resp}") + + async def update(self, registry_row: ContainerRegistryRow, quota: int) -> None: + ssl_verify = True + 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_row.username, registry_row.password) + + quota_info = await self._get_quota_info(sess, registry_row, 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(registry_row.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: + log.error(f"Failed to update quota! response: {resp}") + raise InternalServerError(f"Failed to update quota! response: {resp}") + + async def delete(self, registry_row: ContainerRegistryRow) -> None: + ssl_verify = True + 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_row.username, registry_row.password) + + quota_info = await self._get_quota_info(sess, registry_row, 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(registry_row.url) / "api" / "v2.0" / "quotas" / str(quota_id) + payload = {"hard": {"storage": -1}} + + async with sess.put( + put_quota_api, json=payload, allow_redirects=False, **rqst_args + ) as resp: + if resp.status != 200: + log.error(f"Failed to delete quota! response: {resp}") + raise InternalServerError(f"Failed to delete quota! response: {resp}") + + +class PerProjectContainerRegistryQuotaProtocol(Protocol): + async def create(self, scope_id: ProjectScope, quota: int) -> None: ... + async def update(self, scope_id: ProjectScope, quota: int) -> None: ... + async def delete(self, scope_id: ProjectScope) -> None: ... + async def read(self, scope_id: ProjectScope) -> int: ... + + +class PerProjectContainerRegistryQuotaService(PerProjectContainerRegistryQuotaProtocol): + repository: PerProjectRegistryQuotaRepository + + def __init__(self, repository: PerProjectRegistryQuotaRepository): + self.repository = repository + + def make_client(self, type_: ContainerRegistryType) -> AbstractPerProjectRegistryQuotaClient: + match type_: + case ContainerRegistryType.HARBOR2: + return PerProjectHarborQuotaClient() + case _: + raise GenericBadRequest( + f"{type_} does not support registry quota per project management." + ) + + async def create(self, scope_id: ProjectScope, quota: int) -> None: + registry_row = await self.repository.fetch_container_registry_row(scope_id) + await self.make_client(registry_row.type).create(registry_row, quota) + + async def update(self, scope_id: ProjectScope, quota: int) -> None: + registry_row = await self.repository.fetch_container_registry_row(scope_id) + await self.make_client(registry_row.type).update(registry_row, quota) + + async def delete(self, scope_id: ProjectScope) -> None: + registry_row = await self.repository.fetch_container_registry_row(scope_id) + await self.make_client(registry_row.type).delete(registry_row) + + async def read(self, scope_id: ProjectScope) -> int: + registry_row = await self.repository.fetch_container_registry_row(scope_id) + return await self.make_client(registry_row.type).read(registry_row) diff --git a/src/ai/backend/manager/models/gql.py b/src/ai/backend/manager/models/gql.py index 74c7a0e4a6..5b8c6ba3a7 100644 --- a/src/ai/backend/manager/models/gql.py +++ b/src/ai/backend/manager/models/gql.py @@ -11,6 +11,7 @@ from graphql import OperationType, Undefined from graphql.type import GraphQLField +from ai.backend.manager.api.services.base import ServicesContext from ai.backend.manager.plugin.network import NetworkPluginContext from .gql_models.container_registry import ( @@ -241,6 +242,7 @@ class GraphQueryContext: access_key: str db: ExtendedAsyncSAEngine network_plugin_ctx: NetworkPluginContext + services_ctx: ServicesContext redis_stat: RedisConnectionInfo redis_live: RedisConnectionInfo redis_image: RedisConnectionInfo 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 d1a2469434..0182d634d2 100644 --- a/src/ai/backend/manager/models/gql_models/container_registry.py +++ b/src/ai/backend/manager/models/gql_models/container_registry.py @@ -46,7 +46,6 @@ from ..gql import GraphQueryContext from ..rbac import ScopeType from ..user import UserRole -from .container_registry_utils import HarborQuotaManager from .fields import ScopeField log = BraceStyleAdapter(logging.getLogger(__spec__.name)) # type: ignore @@ -514,13 +513,19 @@ async def mutate( scope_id: ScopeType, quota: int | float, ) -> Self: - async with info.context.db.begin_session() as db_sess: - try: - 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)) + graph_ctx: GraphQueryContext = info.context + try: + match scope_id: + case ProjectScope(_): + await graph_ctx.services_ctx.per_project_container_registries_quota.create( + scope_id, int(quota) + ) + case _: + raise NotImplementedError("Only project scope is supported for now.") + + return cls(ok=True, msg="success") + except Exception as e: + return cls(ok=False, msg=str(e)) class UpdateContainerRegistryQuota(graphene.Mutation): @@ -546,13 +551,19 @@ async def mutate( scope_id: ScopeType, quota: int | float, ) -> Self: - async with info.context.db.begin_session() as db_sess: - try: - 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)) + graph_ctx: GraphQueryContext = info.context + try: + match scope_id: + case ProjectScope(_): + await graph_ctx.services_ctx.per_project_container_registries_quota.update( + scope_id, int(quota) + ) + case _: + raise NotImplementedError("Only project scope is supported for now.") + + return cls(ok=True, msg="success") + except Exception as e: + return cls(ok=False, msg=str(e)) class DeleteContainerRegistryQuota(graphene.Mutation): @@ -576,10 +587,16 @@ async def mutate( info: graphene.ResolveInfo, scope_id: ScopeType, ) -> Self: - async with info.context.db.begin_session() as db_sess: - try: - 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)) + graph_ctx: GraphQueryContext = info.context + try: + match scope_id: + case ProjectScope(_): + await graph_ctx.services_ctx.per_project_container_registries_quota.delete( + scope_id + ) + case _: + raise NotImplementedError("Only project scope is supported for now.") + + 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 deleted file mode 100644 index 1851172a0e..0000000000 --- a/src/ai/backend/manager/models/gql_models/container_registry_utils.py +++ /dev/null @@ -1,258 +0,0 @@ -from __future__ import annotations - -import logging -import uuid -from typing import TYPE_CHECKING, Any, TypedDict - -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.common.types import aobject -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 ..group import GroupRow -from ..rbac import ProjectScope, ScopeType - -if TYPE_CHECKING: - from ...container_registry import ContainerRegistryRow - -log = BraceStyleAdapter(logging.getLogger(__spec__.name)) # type: ignore - -if TYPE_CHECKING: - from ..container_registry import ContainerRegistryRow - - -class HarborQuotaInfo(TypedDict): - previous_quota: int - quota_id: int - - -class HarborQuotaManager(aobject): - """ - 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 - - 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.") - - self.db_sess = db_sess - self.scope_id = scope_id - - async def __ainit__(self) -> None: - from ..container_registry import ContainerRegistryRow - - assert isinstance(self.scope_id, ProjectScope) - - 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 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"], - ) - - registry_query = sa.select(ContainerRegistryRow).where( - (ContainerRegistryRow.registry_name == registry_name) - & (ContainerRegistryRow.project == 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(cls, 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: - 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 - - 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: - 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") - if len(res) > 1: - raise InternalServerError( - f"Multiple quota entities found. (project_id: {harbor_project_id})" - ) - - previous_quota = res[0]["hard"]["storage"] - quota_id = res[0]["id"] - - 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") - - return previous_quota - - 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 = 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: - 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: - 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") - - 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: - 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 a35b3d1419..117ab6c8fe 100644 --- a/src/ai/backend/manager/models/gql_models/group.py +++ b/src/ai/backend/manager/models/gql_models/group.py @@ -14,6 +14,8 @@ from dateutil.parser import parse as dtparse from graphene.types.datetime import DateTime as GQLDateTime +from ai.backend.manager.models.rbac import ProjectScope + from ..base import ( BigInt, FilterExprArg, @@ -29,12 +31,8 @@ from ..group import AssocGroupUserRow, GroupRow, ProjectType, get_permission_ctx from ..minilang.ordering import OrderSpecItem, QueryOrderParser from ..minilang.queryfilter import FieldSpecItem, QueryFilterParser -from ..rbac import ProjectScope from ..rbac.context import ClientContext from ..rbac.permission_defs import ProjectPermission -from .container_registry_utils import ( - HarborQuotaManager, -) from .user import UserConnection, UserNode if TYPE_CHECKING: @@ -217,11 +215,9 @@ async def resolve_user_nodes( return ConnectionResolverResult(result, cursor, pagination_order, page_size, total_cnt) 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) - manager = await HarborQuotaManager.new(db_sess, scope_id) - return await manager.read() + graph_ctx: GraphQueryContext = info.context + scope_id = ProjectScope(project_id=self.id, domain_name=None) + return await graph_ctx.services_ctx.per_project_container_registries_quota.read(scope_id) @classmethod async def get_node(cls, info: graphene.ResolveInfo, id) -> Self: diff --git a/src/ai/backend/manager/server.py b/src/ai/backend/manager/server.py index bb8c76f7ab..95ec2466f9 100644 --- a/src/ai/backend/manager/server.py +++ b/src/ai/backend/manager/server.py @@ -62,6 +62,7 @@ from ai.backend.common.types import AgentSelectionStrategy, HostPortPair from ai.backend.common.utils import env_info from ai.backend.logging import BraceStyleAdapter, Logger, LogLevel +from ai.backend.manager.api.services.base import ServicesContext from ai.backend.manager.plugin.network import NetworkPluginContext from . import __version__ @@ -692,6 +693,12 @@ async def _force_terminate_hanging_sessions( await task +@actxmgr +async def services_ctx(root_ctx: RootContext) -> AsyncIterator[None]: + root_ctx.services_ctx = ServicesContext(root_ctx.db) + yield None + + class background_task_ctx: def __init__(self, root_ctx: RootContext) -> None: self.root_ctx = root_ctx @@ -859,6 +866,7 @@ def build_root_app( manager_status_ctx, redis_ctx, database_ctx, + services_ctx, distributed_lock_ctx, event_dispatcher_ctx, idle_checker_ctx, @@ -910,6 +918,9 @@ async def _call_cleanup_context_shutdown_handlers(app: web.Application) -> None: if pidx == 0: log.info("Loading module: {0}", pkg_name[1:]) subapp_mod = importlib.import_module(pkg_name, "ai.backend.manager.api") + + if pkg_name == ".service": + continue init_subapp(pkg_name, app, getattr(subapp_mod, "create_app")) vendor_path = importlib.resources.files("ai.backend.manager.vendor") diff --git a/tests/manager/models/gql_models/test_container_registries.py b/tests/manager/models/gql_models/test_container_registries.py index b0c32be06f..d1189f4697 100644 --- a/tests/manager/models/gql_models/test_container_registries.py +++ b/tests/manager/models/gql_models/test_container_registries.py @@ -37,6 +37,7 @@ def get_graphquery_context(database_engine: ExtendedAsyncSAEngine) -> GraphQuery registry=None, # type: ignore idle_checker_host=None, # type: ignore network_plugin_ctx=None, # type: ignore + services_ctx=None, # type: ignore ) diff --git a/tests/manager/models/gql_models/test_container_registry_nodes_v2.py b/tests/manager/models/gql_models/test_container_registry_nodes_v2.py index fae23df7f7..e99e966254 100644 --- a/tests/manager/models/gql_models/test_container_registry_nodes_v2.py +++ b/tests/manager/models/gql_models/test_container_registry_nodes_v2.py @@ -108,6 +108,7 @@ def mock_shared_config_api_getitem(key): registry=None, # type: ignore idle_checker_host=None, # type: ignore network_plugin_ctx=None, # type: ignore + services_ctx=None, # type: ignore ) diff --git a/tests/manager/models/gql_models/test_group.py b/tests/manager/models/gql_models/test_group.py index 9aa5088ae8..37880b2f34 100644 --- a/tests/manager/models/gql_models/test_group.py +++ b/tests/manager/models/gql_models/test_group.py @@ -38,6 +38,7 @@ def get_graphquery_context(database_engine: ExtendedAsyncSAEngine) -> GraphQuery registry=None, # type: ignore idle_checker_host=None, # type: ignore network_plugin_ctx=None, # type: ignore + services_ctx=None, # type: ignore ) diff --git a/tests/manager/models/test_container_registries.py b/tests/manager/models/test_container_registries.py index 0da01fab80..33e2f10990 100644 --- a/tests/manager/models/test_container_registries.py +++ b/tests/manager/models/test_container_registries.py @@ -45,6 +45,7 @@ def get_graphquery_context(database_engine: ExtendedAsyncSAEngine) -> GraphQuery storage_manager=None, # type: ignore registry=None, # type: ignore idle_checker_host=None, # type: ignore + services_ctx=None, # type: ignore ) From 937bc8dddb224a2dd5755bc7689162082155c0d4 Mon Sep 17 00:00:00 2001 From: jopemachine Date: Mon, 3 Feb 2025 08:17:06 +0000 Subject: [PATCH 63/75] refactor: Reflect feedbacks --- .../backend/manager/api/container_registry.py | 2 +- src/ai/backend/manager/api/context.py | 2 +- src/ai/backend/manager/api/group.py | 16 +- src/ai/backend/manager/api/services/base.py | 25 --- .../services/container_registries/harbor.py | 210 ------------------ .../manager/{api/services => client}/BUILD | 0 .../{api/services => client}/__init__.py | 0 .../container_registry}/BUILD | 0 .../container_registry}/__init__.py | 0 .../client/container_registry/harbor.py | 174 +++++++++++++++ src/ai/backend/manager/models/gql.py | 2 +- .../models/gql_models/container_registry.py | 18 +- .../manager/models/gql_models/group.py | 4 +- src/ai/backend/manager/server.py | 19 +- src/ai/backend/manager/service/BUILD | 1 + src/ai/backend/manager/service/__init__.py | 0 src/ai/backend/manager/service/base.py | 18 ++ .../manager/service/container_registry/BUILD | 1 + .../service/container_registry/__init__.py | 0 .../container_registry}/base.py | 54 ++++- .../service/container_registry/harbor.py | 110 +++++++++ 21 files changed, 388 insertions(+), 268 deletions(-) delete mode 100644 src/ai/backend/manager/api/services/base.py delete mode 100644 src/ai/backend/manager/api/services/container_registries/harbor.py rename src/ai/backend/manager/{api/services => client}/BUILD (100%) rename src/ai/backend/manager/{api/services => client}/__init__.py (100%) rename src/ai/backend/manager/{api/services/container_registries => client/container_registry}/BUILD (100%) rename src/ai/backend/manager/{api/services/container_registries => client/container_registry}/__init__.py (100%) create mode 100644 src/ai/backend/manager/client/container_registry/harbor.py create mode 100644 src/ai/backend/manager/service/BUILD create mode 100644 src/ai/backend/manager/service/__init__.py create mode 100644 src/ai/backend/manager/service/base.py create mode 100644 src/ai/backend/manager/service/container_registry/BUILD create mode 100644 src/ai/backend/manager/service/container_registry/__init__.py rename src/ai/backend/manager/{api/services/container_registries => service/container_registry}/base.py (59%) create mode 100644 src/ai/backend/manager/service/container_registry/harbor.py diff --git a/src/ai/backend/manager/api/container_registry.py b/src/ai/backend/manager/api/container_registry.py index 399ae8fee2..111d581d95 100644 --- a/src/ai/backend/manager/api/container_registry.py +++ b/src/ai/backend/manager/api/container_registry.py @@ -38,7 +38,7 @@ async def patch_container_registry( from ..models.container_registry import ContainerRegistryRow registry_id = uuid.UUID(request.match_info["registry_id"]) - log.info("PATCH_CONTAINER_REGISTRY (cr:{})", registry_id) + log.info("PATCH_CONTAINER_REGISTRY (registry:{})", registry_id) root_ctx: RootContext = request.app["_root.context"] registry_row_updates = params.model_dump(exclude={"allowed_groups"}, exclude_none=True) diff --git a/src/ai/backend/manager/api/context.py b/src/ai/backend/manager/api/context.py index cd9c925999..dcd2654076 100644 --- a/src/ai/backend/manager/api/context.py +++ b/src/ai/backend/manager/api/context.py @@ -5,8 +5,8 @@ import attrs from ai.backend.common.metrics.metric import CommonMetricRegistry -from ai.backend.manager.api.services.base import ServicesContext from ai.backend.manager.plugin.network import NetworkPluginContext +from ai.backend.manager.service.base import ServicesContext if TYPE_CHECKING: from ai.backend.common.bgtask import BackgroundTaskManager diff --git a/src/ai/backend/manager/api/group.py b/src/ai/backend/manager/api/group.py index 77ca2d864f..2c68457933 100644 --- a/src/ai/backend/manager/api/group.py +++ b/src/ai/backend/manager/api/group.py @@ -31,13 +31,13 @@ }) ) async def update_registry_quota(request: web.Request, params: Any) -> web.Response: - log.info("UPDATE_REGISTRY_QUOTA (gr:{})", params["group_id"]) + log.info("UPDATE_REGISTRY_QUOTA (group:{})", 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"]) - await root_ctx.services_ctx.per_project_container_registries_quota.update(scope_id, quota) + await root_ctx.services_ctx.per_project_container_registries_quota.update_quota(scope_id, quota) return web.Response(status=204) @@ -49,12 +49,12 @@ async def update_registry_quota(request: web.Request, params: Any) -> web.Respon }) ) async def delete_registry_quota(request: web.Request, params: Any) -> web.Response: - log.info("DELETE_REGISTRY_QUOTA (gr:{})", params["group_id"]) + log.info("DELETE_REGISTRY_QUOTA (group:{})", 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) - await root_ctx.services_ctx.per_project_container_registries_quota.delete(scope_id) + await root_ctx.services_ctx.per_project_container_registries_quota.delete_quota(scope_id) return web.Response(status=204) @@ -67,13 +67,13 @@ async def delete_registry_quota(request: web.Request, params: Any) -> web.Respon }) ) async def create_registry_quota(request: web.Request, params: Any) -> web.Response: - log.info("CREATE_REGISTRY_QUOTA (gr:{})", params["group_id"]) + log.info("CREATE_REGISTRY_QUOTA (group:{})", 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"]) - await root_ctx.services_ctx.per_project_container_registries_quota.create(scope_id, quota) + await root_ctx.services_ctx.per_project_container_registries_quota.create_quota(scope_id, quota) return web.Response(status=204) @@ -85,12 +85,12 @@ async def create_registry_quota(request: web.Request, params: Any) -> web.Respon }) ) async def read_registry_quota(request: web.Request, params: Any) -> web.Response: - log.info("READ_REGISTRY_QUOTA (gr:{})", params["group_id"]) + log.info("READ_REGISTRY_QUOTA (group:{})", 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 = await root_ctx.services_ctx.per_project_container_registries_quota.read(scope_id) + quota = await root_ctx.services_ctx.per_project_container_registries_quota.read_quota(scope_id) return web.json_response({"result": quota}) diff --git a/src/ai/backend/manager/api/services/base.py b/src/ai/backend/manager/api/services/base.py deleted file mode 100644 index 4b42e2c0b7..0000000000 --- a/src/ai/backend/manager/api/services/base.py +++ /dev/null @@ -1,25 +0,0 @@ -from ai.backend.manager.models.utils import ExtendedAsyncSAEngine - -from .container_registries.base import PerProjectRegistryQuotaRepository -from .container_registries.harbor import ( - PerProjectContainerRegistryQuotaProtocol, - PerProjectContainerRegistryQuotaService, -) - - -class ServicesContext: - """ - In the API layer, requests are processed through the ServicesContext and - its subordinate layers, including the DB, Client, and Repository layers. - Each layer separates the responsibilities specific to its respective level. - """ - - db: ExtendedAsyncSAEngine - - def __init__(self, db: ExtendedAsyncSAEngine) -> None: - self.db = db - - @property - def per_project_container_registries_quota(self) -> PerProjectContainerRegistryQuotaProtocol: - repository = PerProjectRegistryQuotaRepository(db=self.db) - return PerProjectContainerRegistryQuotaService(repository=repository) diff --git a/src/ai/backend/manager/api/services/container_registries/harbor.py b/src/ai/backend/manager/api/services/container_registries/harbor.py deleted file mode 100644 index 9f3b6b0300..0000000000 --- a/src/ai/backend/manager/api/services/container_registries/harbor.py +++ /dev/null @@ -1,210 +0,0 @@ -from __future__ import annotations - -import logging -from typing import TYPE_CHECKING, Any, Protocol, TypedDict - -import aiohttp -import yarl - -from ai.backend.common.container_registry import ContainerRegistryType -from ai.backend.logging import BraceStyleAdapter -from ai.backend.manager.api.exceptions import GenericBadRequest, InternalServerError, ObjectNotFound -from ai.backend.manager.api.services.container_registries.base import ( - ContainerRegistryRow, - PerProjectRegistryQuotaRepository, -) -from ai.backend.manager.models.rbac import ProjectScope - -if TYPE_CHECKING: - pass - -log = BraceStyleAdapter(logging.getLogger(__spec__.name)) - - -__all__ = ( - "AbstractPerProjectRegistryQuotaClient", - "PerProjectHarborQuotaClient", - "PerProjectRegistryQuotaRepository", - "PerProjectContainerRegistryQuotaProtocol", - "PerProjectContainerRegistryQuotaService", -) - - -class HarborProjectQuotaInfo(TypedDict): - previous_quota: int - quota_id: int - - -class AbstractPerProjectRegistryQuotaClient(Protocol): - async def create(self, registry_row: ContainerRegistryRow, quota: int) -> None: ... - async def update(self, registry_row: ContainerRegistryRow, quota: int) -> None: ... - async def delete(self, registry_row: ContainerRegistryRow) -> None: ... - async def read(self, registry_row: ContainerRegistryRow) -> int: ... - - -class PerProjectHarborQuotaClient(AbstractPerProjectRegistryQuotaClient): - async def _get_harbor_project_id( - self, - sess: aiohttp.ClientSession, - registry_row: ContainerRegistryRow, - rqst_args: dict[str, Any], - ) -> str: - get_project_id_api = ( - yarl.URL(registry_row.url) / "api" / "v2.0" / "projects" / registry_row.project - ) - - 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 str(harbor_project_id) - - async def _get_quota_info( - self, - sess: aiohttp.ClientSession, - registry_row: ContainerRegistryRow, - rqst_args: dict[str, Any], - ) -> HarborProjectQuotaInfo: - harbor_project_id = await self._get_harbor_project_id(sess, registry_row, rqst_args) - get_quota_id_api = (yarl.URL(registry_row.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: - 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") - if len(res) > 1: - raise InternalServerError( - f"Multiple quota entities found. (project_id: {harbor_project_id})" - ) - - previous_quota = res[0]["hard"]["storage"] - quota_id = res[0]["id"] - return HarborProjectQuotaInfo(previous_quota=previous_quota, quota_id=quota_id) - - async def read(self, registry_row: ContainerRegistryRow) -> int: - ssl_verify = True - connector = aiohttp.TCPConnector(ssl=ssl_verify) - async with aiohttp.ClientSession(connector=connector) as sess: - rqst_args: dict[str, Any] = {} - - quota_info = await self._get_quota_info(sess, registry_row, rqst_args) - previous_quota = quota_info["previous_quota"] - if previous_quota == -1: - raise ObjectNotFound(object_name="quota entity") - return previous_quota - - async def create(self, registry_row: ContainerRegistryRow, quota: int) -> None: - ssl_verify = True - 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_row.username, registry_row.password) - - quota_info = await self._get_quota_info(sess, registry_row, rqst_args) - previous_quota, quota_id = quota_info["previous_quota"], quota_info["quota_id"] - - if previous_quota > 0: - raise GenericBadRequest("Quota limit already exists!") - - put_quota_api = yarl.URL(registry_row.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: - log.error(f"Failed to create quota! response: {resp}") - raise InternalServerError(f"Failed to create quota! response: {resp}") - - async def update(self, registry_row: ContainerRegistryRow, quota: int) -> None: - ssl_verify = True - 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_row.username, registry_row.password) - - quota_info = await self._get_quota_info(sess, registry_row, 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(registry_row.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: - log.error(f"Failed to update quota! response: {resp}") - raise InternalServerError(f"Failed to update quota! response: {resp}") - - async def delete(self, registry_row: ContainerRegistryRow) -> None: - ssl_verify = True - 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_row.username, registry_row.password) - - quota_info = await self._get_quota_info(sess, registry_row, 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(registry_row.url) / "api" / "v2.0" / "quotas" / str(quota_id) - payload = {"hard": {"storage": -1}} - - async with sess.put( - put_quota_api, json=payload, allow_redirects=False, **rqst_args - ) as resp: - if resp.status != 200: - log.error(f"Failed to delete quota! response: {resp}") - raise InternalServerError(f"Failed to delete quota! response: {resp}") - - -class PerProjectContainerRegistryQuotaProtocol(Protocol): - async def create(self, scope_id: ProjectScope, quota: int) -> None: ... - async def update(self, scope_id: ProjectScope, quota: int) -> None: ... - async def delete(self, scope_id: ProjectScope) -> None: ... - async def read(self, scope_id: ProjectScope) -> int: ... - - -class PerProjectContainerRegistryQuotaService(PerProjectContainerRegistryQuotaProtocol): - repository: PerProjectRegistryQuotaRepository - - def __init__(self, repository: PerProjectRegistryQuotaRepository): - self.repository = repository - - def make_client(self, type_: ContainerRegistryType) -> AbstractPerProjectRegistryQuotaClient: - match type_: - case ContainerRegistryType.HARBOR2: - return PerProjectHarborQuotaClient() - case _: - raise GenericBadRequest( - f"{type_} does not support registry quota per project management." - ) - - async def create(self, scope_id: ProjectScope, quota: int) -> None: - registry_row = await self.repository.fetch_container_registry_row(scope_id) - await self.make_client(registry_row.type).create(registry_row, quota) - - async def update(self, scope_id: ProjectScope, quota: int) -> None: - registry_row = await self.repository.fetch_container_registry_row(scope_id) - await self.make_client(registry_row.type).update(registry_row, quota) - - async def delete(self, scope_id: ProjectScope) -> None: - registry_row = await self.repository.fetch_container_registry_row(scope_id) - await self.make_client(registry_row.type).delete(registry_row) - - async def read(self, scope_id: ProjectScope) -> int: - registry_row = await self.repository.fetch_container_registry_row(scope_id) - return await self.make_client(registry_row.type).read(registry_row) diff --git a/src/ai/backend/manager/api/services/BUILD b/src/ai/backend/manager/client/BUILD similarity index 100% rename from src/ai/backend/manager/api/services/BUILD rename to src/ai/backend/manager/client/BUILD diff --git a/src/ai/backend/manager/api/services/__init__.py b/src/ai/backend/manager/client/__init__.py similarity index 100% rename from src/ai/backend/manager/api/services/__init__.py rename to src/ai/backend/manager/client/__init__.py diff --git a/src/ai/backend/manager/api/services/container_registries/BUILD b/src/ai/backend/manager/client/container_registry/BUILD similarity index 100% rename from src/ai/backend/manager/api/services/container_registries/BUILD rename to src/ai/backend/manager/client/container_registry/BUILD diff --git a/src/ai/backend/manager/api/services/container_registries/__init__.py b/src/ai/backend/manager/client/container_registry/__init__.py similarity index 100% rename from src/ai/backend/manager/api/services/container_registries/__init__.py rename to src/ai/backend/manager/client/container_registry/__init__.py diff --git a/src/ai/backend/manager/client/container_registry/harbor.py b/src/ai/backend/manager/client/container_registry/harbor.py new file mode 100644 index 0000000000..5fc520a958 --- /dev/null +++ b/src/ai/backend/manager/client/container_registry/harbor.py @@ -0,0 +1,174 @@ +from __future__ import annotations + +import abc +import logging +from typing import TYPE_CHECKING, Any, override + +import aiohttp +import yarl + +from ai.backend.logging import BraceStyleAdapter +from ai.backend.manager.api.exceptions import GenericBadRequest, InternalServerError, ObjectNotFound + +if TYPE_CHECKING: + from ai.backend.manager.service.container_registry.harbor import ( + HarborAuthArgs, + HarborProjectInfo, + HarborProjectQuotaInfo, + ) + +log = BraceStyleAdapter(logging.getLogger(__spec__.name)) + + +def _get_harbor_auth_args(auth_args: HarborAuthArgs) -> dict[str, Any]: + return {"auth": aiohttp.BasicAuth(auth_args["username"], auth_args["password"])} + + +class AbstractPerProjectRegistryQuotaClient(abc.ABC): + async def create_quota( + self, project_info: HarborProjectInfo, quota: int, auth_args: HarborAuthArgs + ) -> None: + raise NotImplementedError + + async def update_quota( + self, project_info: HarborProjectInfo, quota: int, auth_args: HarborAuthArgs + ) -> None: + raise NotImplementedError + + async def delete_quota( + self, project_info: HarborProjectInfo, auth_args: HarborAuthArgs + ) -> None: + raise NotImplementedError + + async def read_quota(self, project_info: HarborProjectInfo) -> int: + raise NotImplementedError + + +class PerProjectHarborQuotaClient(AbstractPerProjectRegistryQuotaClient): + async def _get_harbor_project_id( + self, + sess: aiohttp.ClientSession, + project_info: HarborProjectInfo, + rqst_args: dict[str, Any], + ) -> str: + get_project_id_api = ( + yarl.URL(project_info.url) / "api" / "v2.0" / "projects" / project_info.project + ) + + 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 str(harbor_project_id) + + async def _get_quota_info( + self, + sess: aiohttp.ClientSession, + project_info: HarborProjectInfo, + rqst_args: dict[str, Any], + ) -> HarborProjectQuotaInfo: + from ...service.container_registry.harbor import HarborProjectQuotaInfo + + harbor_project_id = await self._get_harbor_project_id(sess, project_info, rqst_args) + get_quota_id_api = (yarl.URL(project_info.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: + 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") + if len(res) > 1: + raise InternalServerError( + f"Multiple quota entities found. (project_id: {harbor_project_id})" + ) + + previous_quota = res[0]["hard"]["storage"] + quota_id = res[0]["id"] + return HarborProjectQuotaInfo(previous_quota=previous_quota, quota_id=quota_id) + + @override + async def read_quota(self, project_info: HarborProjectInfo) -> int: + connector = aiohttp.TCPConnector(ssl=project_info.ssl_verify) + async with aiohttp.ClientSession(connector=connector) as sess: + rqst_args: dict[str, Any] = {} + quota_info = await self._get_quota_info(sess, project_info, rqst_args) + previous_quota = quota_info["previous_quota"] + if previous_quota == -1: + raise ObjectNotFound(object_name="quota entity") + return previous_quota + + @override + async def create_quota( + self, project_info: HarborProjectInfo, quota: int, auth_args: HarborAuthArgs + ) -> None: + connector = aiohttp.TCPConnector(ssl=project_info.ssl_verify) + async with aiohttp.ClientSession(connector=connector) as sess: + rqst_args = _get_harbor_auth_args(auth_args) + quota_info = await self._get_quota_info(sess, project_info, rqst_args) + previous_quota, quota_id = quota_info["previous_quota"], quota_info["quota_id"] + + if previous_quota > 0: + raise GenericBadRequest("Quota limit already exists!") + + put_quota_api = yarl.URL(project_info.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: + log.error(f"Failed to create quota! response: {resp}") + raise InternalServerError(f"Failed to create quota! response: {resp}") + + @override + async def update_quota( + self, project_info: HarborProjectInfo, quota: int, auth_args: HarborAuthArgs + ) -> None: + connector = aiohttp.TCPConnector(ssl=project_info.ssl_verify) + async with aiohttp.ClientSession(connector=connector) as sess: + rqst_args = _get_harbor_auth_args(auth_args) + quota_info = await self._get_quota_info(sess, project_info, 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(project_info.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: + log.error(f"Failed to update quota! response: {resp}") + raise InternalServerError(f"Failed to update quota! response: {resp}") + + @override + async def delete_quota( + self, project_info: HarborProjectInfo, auth_args: HarborAuthArgs + ) -> None: + connector = aiohttp.TCPConnector(ssl=project_info.ssl_verify) + async with aiohttp.ClientSession(connector=connector) as sess: + rqst_args = _get_harbor_auth_args(auth_args) + quota_info = await self._get_quota_info(sess, project_info, 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(project_info.url) / "api" / "v2.0" / "quotas" / str(quota_id) + payload = {"hard": {"storage": -1}} + + async with sess.put( + put_quota_api, json=payload, allow_redirects=False, **rqst_args + ) as resp: + if resp.status != 200: + 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.py b/src/ai/backend/manager/models/gql.py index 5b8c6ba3a7..df03b09bc7 100644 --- a/src/ai/backend/manager/models/gql.py +++ b/src/ai/backend/manager/models/gql.py @@ -11,8 +11,8 @@ from graphql import OperationType, Undefined from graphql.type import GraphQLField -from ai.backend.manager.api.services.base import ServicesContext from ai.backend.manager.plugin.network import NetworkPluginContext +from ai.backend.manager.service.base import ServicesContext from .gql_models.container_registry import ( ContainerRegistryConnection, 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 0182d634d2..5f7693132f 100644 --- a/src/ai/backend/manager/models/gql_models/container_registry.py +++ b/src/ai/backend/manager/models/gql_models/container_registry.py @@ -517,8 +517,10 @@ async def mutate( try: match scope_id: case ProjectScope(_): - await graph_ctx.services_ctx.per_project_container_registries_quota.create( - scope_id, int(quota) + await ( + graph_ctx.services_ctx.per_project_container_registries_quota.create_quota( + scope_id, int(quota) + ) ) case _: raise NotImplementedError("Only project scope is supported for now.") @@ -555,8 +557,10 @@ async def mutate( try: match scope_id: case ProjectScope(_): - await graph_ctx.services_ctx.per_project_container_registries_quota.update( - scope_id, int(quota) + await ( + graph_ctx.services_ctx.per_project_container_registries_quota.update_quota( + scope_id, int(quota) + ) ) case _: raise NotImplementedError("Only project scope is supported for now.") @@ -591,8 +595,10 @@ async def mutate( try: match scope_id: case ProjectScope(_): - await graph_ctx.services_ctx.per_project_container_registries_quota.delete( - scope_id + await ( + graph_ctx.services_ctx.per_project_container_registries_quota.delete_quota( + scope_id + ) ) case _: raise NotImplementedError("Only project scope is supported for now.") diff --git a/src/ai/backend/manager/models/gql_models/group.py b/src/ai/backend/manager/models/gql_models/group.py index 117ab6c8fe..47fa77c300 100644 --- a/src/ai/backend/manager/models/gql_models/group.py +++ b/src/ai/backend/manager/models/gql_models/group.py @@ -217,7 +217,9 @@ async def resolve_user_nodes( async def resolve_registry_quota(self, info: graphene.ResolveInfo) -> int: graph_ctx: GraphQueryContext = info.context scope_id = ProjectScope(project_id=self.id, domain_name=None) - return await graph_ctx.services_ctx.per_project_container_registries_quota.read(scope_id) + return await graph_ctx.services_ctx.per_project_container_registries_quota.read_quota( + scope_id + ) @classmethod async def get_node(cls, info: graphene.ResolveInfo, id) -> Self: diff --git a/src/ai/backend/manager/server.py b/src/ai/backend/manager/server.py index 95ec2466f9..c93f31d443 100644 --- a/src/ai/backend/manager/server.py +++ b/src/ai/backend/manager/server.py @@ -62,8 +62,12 @@ from ai.backend.common.types import AgentSelectionStrategy, HostPortPair from ai.backend.common.utils import env_info from ai.backend.logging import BraceStyleAdapter, Logger, LogLevel -from ai.backend.manager.api.services.base import ServicesContext from ai.backend.manager.plugin.network import NetworkPluginContext +from ai.backend.manager.service.base import ServicesContext +from ai.backend.manager.service.container_registry.base import PerProjectRegistryQuotaRepository +from ai.backend.manager.service.container_registry.harbor import ( + PerProjectContainerRegistryQuotaService, +) from . import __version__ from .agent_cache import AgentRPCCache @@ -695,7 +699,15 @@ async def _force_terminate_hanging_sessions( @actxmgr async def services_ctx(root_ctx: RootContext) -> AsyncIterator[None]: - root_ctx.services_ctx = ServicesContext(root_ctx.db) + db = root_ctx.db + + per_project_container_registries_quota = PerProjectContainerRegistryQuotaService( + repository=PerProjectRegistryQuotaRepository(db) + ) + + root_ctx.services_ctx = ServicesContext( + per_project_container_registries_quota, + ) yield None @@ -918,9 +930,6 @@ async def _call_cleanup_context_shutdown_handlers(app: web.Application) -> None: if pidx == 0: log.info("Loading module: {0}", pkg_name[1:]) subapp_mod = importlib.import_module(pkg_name, "ai.backend.manager.api") - - if pkg_name == ".service": - continue init_subapp(pkg_name, app, getattr(subapp_mod, "create_app")) vendor_path = importlib.resources.files("ai.backend.manager.vendor") diff --git a/src/ai/backend/manager/service/BUILD b/src/ai/backend/manager/service/BUILD new file mode 100644 index 0000000000..7357442404 --- /dev/null +++ b/src/ai/backend/manager/service/BUILD @@ -0,0 +1 @@ +python_sources(name="src") diff --git a/src/ai/backend/manager/service/__init__.py b/src/ai/backend/manager/service/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/src/ai/backend/manager/service/base.py b/src/ai/backend/manager/service/base.py new file mode 100644 index 0000000000..bd4a22aa21 --- /dev/null +++ b/src/ai/backend/manager/service/base.py @@ -0,0 +1,18 @@ +from .container_registry.harbor import ( + PerProjectContainerRegistryQuota, +) + + +class ServicesContext: + """ + In the API layer, requests are processed through the ServicesContext and + its subordinate layers, including the DB, Client, and Repository layers. + Each layer separates the responsibilities specific to its respective level. + """ + + per_project_container_registries_quota: PerProjectContainerRegistryQuota + + def __init__( + self, per_project_container_registries_quota: PerProjectContainerRegistryQuota + ) -> None: + self.per_project_container_registries_quota = per_project_container_registries_quota diff --git a/src/ai/backend/manager/service/container_registry/BUILD b/src/ai/backend/manager/service/container_registry/BUILD new file mode 100644 index 0000000000..7357442404 --- /dev/null +++ b/src/ai/backend/manager/service/container_registry/BUILD @@ -0,0 +1 @@ +python_sources(name="src") diff --git a/src/ai/backend/manager/service/container_registry/__init__.py b/src/ai/backend/manager/service/container_registry/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/src/ai/backend/manager/api/services/container_registries/base.py b/src/ai/backend/manager/service/container_registry/base.py similarity index 59% rename from src/ai/backend/manager/api/services/container_registries/base.py rename to src/ai/backend/manager/service/container_registry/base.py index ce0abd4dde..f9567e2cdf 100644 --- a/src/ai/backend/manager/api/services/container_registries/base.py +++ b/src/ai/backend/manager/service/container_registry/base.py @@ -1,11 +1,15 @@ from __future__ import annotations +import abc import logging -from typing import TYPE_CHECKING +import uuid +from dataclasses import dataclass +from typing import Any, override import sqlalchemy as sa from sqlalchemy.orm import load_only +from ai.backend.common.container_registry import ContainerRegistryType from ai.backend.logging import BraceStyleAdapter from ai.backend.manager.api.exceptions import ( ContainerRegistryNotFound, @@ -15,15 +19,31 @@ from ai.backend.manager.models.rbac import ProjectScope from ai.backend.manager.models.utils import ExtendedAsyncSAEngine -if TYPE_CHECKING: - pass - log = BraceStyleAdapter(logging.getLogger(__spec__.name)) -class PerProjectRegistryQuotaRepository: - """ """ +@dataclass +class ContainerRegistryRowInfo: + id: uuid.UUID + url: str + registry_name: str + type: ContainerRegistryType + project: str + username: str + password: str + ssl_verify: bool + is_global: bool + extra: dict[str, Any] + + +class AbstractPerProjectRegistryQuotaRepository(abc.ABC): + async def fetch_container_registry_row( + self, scope_id: ProjectScope + ) -> ContainerRegistryRowInfo: + raise NotImplementedError + +class PerProjectRegistryQuotaRepository(AbstractPerProjectRegistryQuotaRepository): def __init__(self, db: ExtendedAsyncSAEngine): self.db = db @@ -36,7 +56,10 @@ def _is_valid_group_row(cls, group_row: GroupRow) -> bool: and "project" in group_row.container_registry ) - async def fetch_container_registry_row(self, scope_id: ProjectScope) -> ContainerRegistryRow: + @override + async def fetch_container_registry_row( + self, scope_id: ProjectScope + ) -> ContainerRegistryRowInfo: async with self.db.begin_readonly_session() as db_sess: project_id = scope_id.project_id group_query = ( @@ -49,7 +72,7 @@ async def fetch_container_registry_row(self, scope_id: ProjectScope) -> Containe if not PerProjectRegistryQuotaRepository._is_valid_group_row(group_row): raise ContainerRegistryNotFound( - f"Container registry info does not exist or is invalid in the group. (gr: {project_id})" + f"Container registry info does not exist or is invalid in the group. (group: {project_id})" ) registry_name, project = ( @@ -67,7 +90,18 @@ async def fetch_container_registry_row(self, scope_id: ProjectScope) -> Containe if not registry: raise ContainerRegistryNotFound( - f"Specified container registry row not found. (cr: {registry_name}, gr: {project})" + f"Container registry row not found. (registry: {registry_name}, group: {project})" ) - return registry + return ContainerRegistryRowInfo( + id=registry.id, + url=registry.url, + registry_name=registry.registry_name, + type=registry.type, + project=registry.project, + username=registry.username, + password=registry.password, + ssl_verify=registry.ssl_verify, + is_global=registry.is_global, + extra=registry.extra, + ) diff --git a/src/ai/backend/manager/service/container_registry/harbor.py b/src/ai/backend/manager/service/container_registry/harbor.py new file mode 100644 index 0000000000..f74ec4d9c4 --- /dev/null +++ b/src/ai/backend/manager/service/container_registry/harbor.py @@ -0,0 +1,110 @@ +from __future__ import annotations + +import abc +import logging +from dataclasses import dataclass +from typing import TypedDict, override + +from ai.backend.common.container_registry import ContainerRegistryType +from ai.backend.logging import BraceStyleAdapter +from ai.backend.manager.api.exceptions import GenericBadRequest +from ai.backend.manager.client.container_registry.harbor import ( + AbstractPerProjectRegistryQuotaClient, + PerProjectHarborQuotaClient, +) +from ai.backend.manager.models.rbac import ProjectScope +from ai.backend.manager.service.container_registry.base import ( + ContainerRegistryRowInfo, + PerProjectRegistryQuotaRepository, +) + +log = BraceStyleAdapter(logging.getLogger(__spec__.name)) + + +@dataclass +class HarborProjectInfo: + url: str + project: str + ssl_verify: bool + + +class HarborAuthArgs(TypedDict): + username: str + password: str + + +class HarborProjectQuotaInfo(TypedDict): + previous_quota: int + quota_id: int + + +class PerProjectContainerRegistryQuota(abc.ABC): + async def create_quota(self, scope_id: ProjectScope, quota: int) -> None: + raise NotImplementedError + + async def update_quota(self, scope_id: ProjectScope, quota: int) -> None: + raise NotImplementedError + + async def delete_quota(self, scope_id: ProjectScope) -> None: + raise NotImplementedError + + async def read_quota(self, scope_id: ProjectScope) -> int: + raise NotImplementedError + + +class PerProjectContainerRegistryQuotaService(PerProjectContainerRegistryQuota): + repository: PerProjectRegistryQuotaRepository + + def __init__(self, repository: PerProjectRegistryQuotaRepository): + self.repository = repository + + def _registry_row_to_harbor_project_info( + self, registry_info: ContainerRegistryRowInfo + ) -> HarborProjectInfo: + return HarborProjectInfo( + url=registry_info.url, + project=registry_info.project, + ssl_verify=registry_info.ssl_verify, + ) + + def _make_client(self, type_: ContainerRegistryType) -> AbstractPerProjectRegistryQuotaClient: + match type_: + case ContainerRegistryType.HARBOR2: + return PerProjectHarborQuotaClient() + case _: + raise GenericBadRequest( + f"{type_} does not support registry quota per project management." + ) + + @override + async def create_quota(self, scope_id: ProjectScope, quota: int) -> None: + registry_info = await self.repository.fetch_container_registry_row(scope_id) + project_info = self._registry_row_to_harbor_project_info(registry_info) + credential = HarborAuthArgs( + username=registry_info.username, password=registry_info.password + ) + await self._make_client(registry_info.type).create_quota(project_info, quota, credential) + + @override + async def update_quota(self, scope_id: ProjectScope, quota: int) -> None: + registry_info = await self.repository.fetch_container_registry_row(scope_id) + project_info = self._registry_row_to_harbor_project_info(registry_info) + credential = HarborAuthArgs( + username=registry_info.username, password=registry_info.password + ) + await self._make_client(registry_info.type).update_quota(project_info, quota, credential) + + @override + async def delete_quota(self, scope_id: ProjectScope) -> None: + registry_info = await self.repository.fetch_container_registry_row(scope_id) + project_info = self._registry_row_to_harbor_project_info(registry_info) + credential = HarborAuthArgs( + username=registry_info.username, password=registry_info.password + ) + await self._make_client(registry_info.type).delete_quota(project_info, credential) + + @override + async def read_quota(self, scope_id: ProjectScope) -> int: + registry_info = await self.repository.fetch_container_registry_row(scope_id) + project_info = self._registry_row_to_harbor_project_info(registry_info) + return await self._make_client(registry_info.type).read_quota(project_info) From d498374ad0ffd6ff7d9847670768b2308aefaaa5 Mon Sep 17 00:00:00 2001 From: jopemachine Date: Wed, 5 Feb 2025 04:52:33 +0000 Subject: [PATCH 64/75] fix: try to fix CI --- tests/manager/models/test_image.py | 1 + tests/manager/models/test_vfolder.py | 1 + 2 files changed, 2 insertions(+) diff --git a/tests/manager/models/test_image.py b/tests/manager/models/test_image.py index b956ce5638..b5b13cd690 100644 --- a/tests/manager/models/test_image.py +++ b/tests/manager/models/test_image.py @@ -60,6 +60,7 @@ def get_graphquery_context( registry=None, # type: ignore idle_checker_host=None, # type: ignore network_plugin_ctx=None, # type: ignore + services_ctx=None, # type: ignore ) diff --git a/tests/manager/models/test_vfolder.py b/tests/manager/models/test_vfolder.py index 89e9118eb7..9f0b2332ab 100644 --- a/tests/manager/models/test_vfolder.py +++ b/tests/manager/models/test_vfolder.py @@ -30,6 +30,7 @@ def get_graphquery_context(database_engine: ExtendedAsyncSAEngine) -> GraphQuery registry=None, # type: ignore idle_checker_host=None, # type: ignore network_plugin_ctx=None, # type: ignore + services_ctx=None, # type: ignore ) From ff7382ca52b06071ff0d58e50df6f91de96c6c4f Mon Sep 17 00:00:00 2001 From: jopemachine Date: Wed, 5 Feb 2025 05:12:17 +0000 Subject: [PATCH 65/75] fix: Fix Broken CI --- .../gql_models/test_container_registries.py | 16 +++++++++++----- tests/manager/models/gql_models/test_group.py | 10 +++++++--- tests/manager/models/test_image.py | 12 +++++++++--- 3 files changed, 27 insertions(+), 11 deletions(-) diff --git a/tests/manager/models/gql_models/test_container_registries.py b/tests/manager/models/gql_models/test_container_registries.py index d1189f4697..763d5867a6 100644 --- a/tests/manager/models/gql_models/test_container_registries.py +++ b/tests/manager/models/gql_models/test_container_registries.py @@ -8,6 +8,7 @@ from ai.backend.manager.models.utils import ExtendedAsyncSAEngine from ai.backend.manager.server import ( database_ctx, + services_ctx, ) from ai.backend.testutils.extra_fixtures import FIXTURES_FOR_HARBOR_CRUD_TEST @@ -17,7 +18,9 @@ def client() -> Client: return Client(Schema(query=Queries, mutation=Mutations, auto_camelcase=False)) -def get_graphquery_context(database_engine: ExtendedAsyncSAEngine) -> GraphQueryContext: +def get_graphquery_context( + database_engine: ExtendedAsyncSAEngine, services_ctx +) -> GraphQueryContext: return GraphQueryContext( schema=None, # type: ignore dataloader_manager=None, # type: ignore @@ -37,7 +40,7 @@ def get_graphquery_context(database_engine: ExtendedAsyncSAEngine) -> GraphQuery registry=None, # type: ignore idle_checker_host=None, # type: ignore network_plugin_ctx=None, # type: ignore - services_ctx=None, # type: ignore + services_ctx=services_ctx, # type: ignore ) @@ -82,12 +85,13 @@ async def test_harbor_create_project_quota( test_app, _ = await create_app_and_client( [ database_ctx, + services_ctx, ], [], ) root_ctx: RootContext = test_app["_root.context"] - context = get_graphquery_context(root_ctx.db) + context = get_graphquery_context(root_ctx.db, root_ctx.services_ctx) create_query = """ mutation ($scope_id: ScopeField!, $quota: BigInt!) { @@ -171,12 +175,13 @@ async def test_harbor_update_project_quota( test_app, _ = await create_app_and_client( [ database_ctx, + services_ctx, ], [], ) root_ctx: RootContext = test_app["_root.context"] - context = get_graphquery_context(root_ctx.db) + context = get_graphquery_context(root_ctx.db, root_ctx.services_ctx) update_query = """ mutation ($scope_id: ScopeField!, $quota: BigInt!) { @@ -260,12 +265,13 @@ async def test_harbor_delete_project_quota( test_app, _ = await create_app_and_client( [ database_ctx, + services_ctx, ], [], ) root_ctx: RootContext = test_app["_root.context"] - context = get_graphquery_context(root_ctx.db) + context = get_graphquery_context(root_ctx.db, root_ctx.services_ctx) delete_query = """ mutation ($scope_id: ScopeField!) { diff --git a/tests/manager/models/gql_models/test_group.py b/tests/manager/models/gql_models/test_group.py index 37880b2f34..4f64fa2578 100644 --- a/tests/manager/models/gql_models/test_group.py +++ b/tests/manager/models/gql_models/test_group.py @@ -9,6 +9,7 @@ from ai.backend.manager.models.utils import ExtendedAsyncSAEngine from ai.backend.manager.server import ( database_ctx, + services_ctx, ) from ai.backend.testutils.extra_fixtures import FIXTURES_FOR_HARBOR_CRUD_TEST @@ -18,7 +19,9 @@ def client() -> Client: return Client(Schema(query=Queries, mutation=Mutations, auto_camelcase=False)) -def get_graphquery_context(database_engine: ExtendedAsyncSAEngine) -> GraphQueryContext: +def get_graphquery_context( + database_engine: ExtendedAsyncSAEngine, services_ctx +) -> GraphQueryContext: return GraphQueryContext( schema=None, # type: ignore dataloader_manager=None, # type: ignore @@ -38,7 +41,7 @@ def get_graphquery_context(database_engine: ExtendedAsyncSAEngine) -> GraphQuery registry=None, # type: ignore idle_checker_host=None, # type: ignore network_plugin_ctx=None, # type: ignore - services_ctx=None, # type: ignore + services_ctx=services_ctx, # type: ignore ) @@ -52,12 +55,13 @@ async def test_harbor_read_project_quota( test_app, _ = await create_app_and_client( [ database_ctx, + services_ctx, ], [], ) root_ctx: RootContext = test_app["_root.context"] - context = get_graphquery_context(root_ctx.db) + context = get_graphquery_context(root_ctx.db, root_ctx.services_ctx) # Arbitrary values for mocking Harbor API responses HARBOR_PROJECT_ID = "123" diff --git a/tests/manager/models/test_image.py b/tests/manager/models/test_image.py index b5b13cd690..785de1751e 100644 --- a/tests/manager/models/test_image.py +++ b/tests/manager/models/test_image.py @@ -28,6 +28,7 @@ hook_plugin_ctx, monitoring_ctx, redis_ctx, + services_ctx, shared_config_ctx, ) from ai.backend.testutils.mock import mock_aioresponses_sequential_payloads @@ -39,7 +40,9 @@ def client() -> Client: def get_graphquery_context( - background_task_manager, database_engine: ExtendedAsyncSAEngine + background_task_manager, + services_ctx, + database_engine: ExtendedAsyncSAEngine, ) -> GraphQueryContext: return GraphQueryContext( schema=None, # type: ignore @@ -60,7 +63,7 @@ def get_graphquery_context( registry=None, # type: ignore idle_checker_host=None, # type: ignore network_plugin_ctx=None, # type: ignore - services_ctx=None, # type: ignore + services_ctx=services_ctx, # type: ignore ) @@ -176,6 +179,7 @@ async def test_image_rescan_on_docker_registry( hook_plugin_ctx, redis_ctx, event_dispatcher_ctx, + services_ctx, background_task_ctx, ], [".events", ".auth"], @@ -272,7 +276,9 @@ def setup_dockerhub_mocking(mocked): with aioresponses() as mocked: setup_dockerhub_mocking(mocked) - context = get_graphquery_context(root_ctx.background_task_manager, root_ctx.db) + context = get_graphquery_context( + root_ctx.background_task_manager, root_ctx.services_ctx, root_ctx.db + ) image_rescan_query = """ mutation ($registry: String, $project: String) { rescan_images(registry: $registry, project: $project) { From 946565dfd28f730030bf669eb01be0fa6c51e1a3 Mon Sep 17 00:00:00 2001 From: jopemachine Date: Wed, 5 Feb 2025 08:06:34 +0000 Subject: [PATCH 66/75] fix: Broken lint --- src/ai/backend/manager/models/gql_models/group.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/ai/backend/manager/models/gql_models/group.py b/src/ai/backend/manager/models/gql_models/group.py index 47fa77c300..aacf1b25cf 100644 --- a/src/ai/backend/manager/models/gql_models/group.py +++ b/src/ai/backend/manager/models/gql_models/group.py @@ -3,6 +3,7 @@ from collections.abc import Mapping from typing import ( TYPE_CHECKING, + Any, Optional, Self, Sequence, From abbbd7a3acdbe25d4c869ca0e42fa6934f2b5937 Mon Sep 17 00:00:00 2001 From: jopemachine Date: Wed, 5 Feb 2025 08:21:15 +0000 Subject: [PATCH 67/75] fix: Broken test --- tests/manager/api/test_group.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/tests/manager/api/test_group.py b/tests/manager/api/test_group.py index 2a18c69a99..27a024af99 100644 --- a/tests/manager/api/test_group.py +++ b/tests/manager/api/test_group.py @@ -9,6 +9,7 @@ hook_plugin_ctx, monitoring_ctx, redis_ctx, + services_ctx, shared_config_ctx, ) from ai.backend.testutils.extra_fixtures import FIXTURES_FOR_HARBOR_CRUD_TEST @@ -29,7 +30,7 @@ } ], }, - "expected_code": 200, + "expected_code": 204, }, { "mock_harbor_responses": { @@ -60,6 +61,7 @@ async def test_harbor_create_project_quota( monitoring_ctx, hook_plugin_ctx, redis_ctx, + services_ctx, ], [".group", ".auth"], ) @@ -144,6 +146,7 @@ async def test_harbor_read_project_quota( monitoring_ctx, hook_plugin_ctx, redis_ctx, + services_ctx, ], [".group", ".auth"], ) @@ -186,7 +189,7 @@ async def test_harbor_read_project_quota( } ], }, - "expected_code": 200, + "expected_code": 204, }, { "mock_harbor_responses": { @@ -217,6 +220,7 @@ async def test_harbor_update_project_quota( monitoring_ctx, hook_plugin_ctx, redis_ctx, + services_ctx, ], [".group", ".auth"], ) @@ -266,7 +270,7 @@ async def test_harbor_update_project_quota( } ], }, - "expected_code": 200, + "expected_code": 204, }, { "mock_harbor_responses": { @@ -297,6 +301,7 @@ async def test_harbor_delete_project_quota( monitoring_ctx, hook_plugin_ctx, redis_ctx, + services_ctx, ], [".group", ".auth"], ) From 51af272fa9f82b143865d867670d2f9b21340392 Mon Sep 17 00:00:00 2001 From: jopemachine Date: Thu, 6 Feb 2025 09:23:00 +0000 Subject: [PATCH 68/75] refactor: improve abstraction --- src/ai/backend/manager/api/group.py | 35 ++++++- .../models/gql_models/container_registry.py | 14 ++- src/ai/backend/manager/service/base.py | 7 +- .../service/container_registry/harbor.py | 93 +++++++++++++------ 4 files changed, 110 insertions(+), 39 deletions(-) diff --git a/src/ai/backend/manager/api/group.py b/src/ai/backend/manager/api/group.py index 2c68457933..cb0c4ad22d 100644 --- a/src/ai/backend/manager/api/group.py +++ b/src/ai/backend/manager/api/group.py @@ -10,6 +10,9 @@ from ai.backend.common import validators as tx from ai.backend.logging import BraceStyleAdapter from ai.backend.manager.models.rbac import ProjectScope +from ai.backend.manager.service.container_registry.harbor import ( + PerProjectContainerRegistryQuotaClientPool, +) if TYPE_CHECKING: from .context import RootContext @@ -37,7 +40,13 @@ async def update_registry_quota(request: web.Request, params: Any) -> web.Respon scope_id = ProjectScope(project_id=group_id, domain_name=None) quota = int(params["quota"]) - await root_ctx.services_ctx.per_project_container_registries_quota.update_quota(scope_id, quota) + registry_info = await root_ctx.services_ctx.per_project_container_registries_quota.repository.fetch_container_registry_row( + scope_id + ) + client = PerProjectContainerRegistryQuotaClientPool.make_client(registry_info.type) + await root_ctx.services_ctx.per_project_container_registries_quota.update_quota( + client, registry_info, quota + ) return web.Response(status=204) @@ -54,7 +63,13 @@ async def delete_registry_quota(request: web.Request, params: Any) -> web.Respon group_id = params["group_id"] scope_id = ProjectScope(project_id=group_id, domain_name=None) - await root_ctx.services_ctx.per_project_container_registries_quota.delete_quota(scope_id) + registry_info = await root_ctx.services_ctx.per_project_container_registries_quota.repository.fetch_container_registry_row( + scope_id + ) + client = PerProjectContainerRegistryQuotaClientPool.make_client(scope_id) + await root_ctx.services_ctx.per_project_container_registries_quota.delete_quota( + client, registry_info + ) return web.Response(status=204) @@ -73,7 +88,13 @@ async def create_registry_quota(request: web.Request, params: Any) -> web.Respon scope_id = ProjectScope(project_id=group_id, domain_name=None) quota = int(params["quota"]) - await root_ctx.services_ctx.per_project_container_registries_quota.create_quota(scope_id, quota) + registry_info = await root_ctx.services_ctx.per_project_container_registries_quota.repository.fetch_container_registry_row( + scope_id + ) + client = PerProjectContainerRegistryQuotaClientPool.make_client(registry_info.type) + await root_ctx.services_ctx.per_project_container_registries_quota.create_quota( + client, registry_info, quota + ) return web.Response(status=204) @@ -90,7 +111,13 @@ async def read_registry_quota(request: web.Request, params: Any) -> web.Response group_id = params["group_id"] scope_id = ProjectScope(project_id=group_id, domain_name=None) - quota = await root_ctx.services_ctx.per_project_container_registries_quota.read_quota(scope_id) + registry_info = await root_ctx.services_ctx.per_project_container_registries_quota.repository.fetch_container_registry_row( + scope_id + ) + client = PerProjectContainerRegistryQuotaClientPool.make_client(registry_info.type) + quota = await root_ctx.services_ctx.per_project_container_registries_quota.read_quota( + client, registry_info + ) 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 5f7693132f..cc47e562a3 100644 --- a/src/ai/backend/manager/models/gql_models/container_registry.py +++ b/src/ai/backend/manager/models/gql_models/container_registry.py @@ -25,6 +25,9 @@ ) from ai.backend.manager.models.user import UserRole from ai.backend.manager.models.utils import ExtendedAsyncSAEngine +from ai.backend.manager.service.container_registry.harbor import ( + PerProjectContainerRegistryQuotaClientPool, +) from ...defs import PASSWORD_PLACEHOLDER from ..association_container_registries_groups import ( @@ -44,7 +47,7 @@ if TYPE_CHECKING: from ..gql import GraphQueryContext -from ..rbac import ScopeType +from ..rbac import ProjectScope, ScopeType from ..user import UserRole from .fields import ScopeField @@ -514,12 +517,17 @@ async def mutate( quota: int | float, ) -> Self: graph_ctx: GraphQueryContext = info.context + + registry_info = await graph_ctx.services_ctx.per_project_container_registries_quota.repository.fetch_container_registry_row( + scope_id + ) + client = PerProjectContainerRegistryQuotaClientPool.make_client(registry_info.type) try: match scope_id: - case ProjectScope(_): + case ProjectScope(): await ( graph_ctx.services_ctx.per_project_container_registries_quota.create_quota( - scope_id, int(quota) + client, registry_info, int(quota) ) ) case _: diff --git a/src/ai/backend/manager/service/base.py b/src/ai/backend/manager/service/base.py index bd4a22aa21..3ab8cc508b 100644 --- a/src/ai/backend/manager/service/base.py +++ b/src/ai/backend/manager/service/base.py @@ -1,5 +1,5 @@ from .container_registry.harbor import ( - PerProjectContainerRegistryQuota, + AbstractPerProjectContainerRegistryQuotaService, ) @@ -10,9 +10,10 @@ class ServicesContext: Each layer separates the responsibilities specific to its respective level. """ - per_project_container_registries_quota: PerProjectContainerRegistryQuota + per_project_container_registries_quota: AbstractPerProjectContainerRegistryQuotaService def __init__( - self, per_project_container_registries_quota: PerProjectContainerRegistryQuota + self, + per_project_container_registries_quota: AbstractPerProjectContainerRegistryQuotaService, ) -> None: self.per_project_container_registries_quota = per_project_container_registries_quota diff --git a/src/ai/backend/manager/service/container_registry/harbor.py b/src/ai/backend/manager/service/container_registry/harbor.py index f74ec4d9c4..fce41bc8b8 100644 --- a/src/ai/backend/manager/service/container_registry/harbor.py +++ b/src/ai/backend/manager/service/container_registry/harbor.py @@ -12,7 +12,6 @@ AbstractPerProjectRegistryQuotaClient, PerProjectHarborQuotaClient, ) -from ai.backend.manager.models.rbac import ProjectScope from ai.backend.manager.service.container_registry.base import ( ContainerRegistryRowInfo, PerProjectRegistryQuotaRepository, @@ -38,24 +37,57 @@ class HarborProjectQuotaInfo(TypedDict): quota_id: int -class PerProjectContainerRegistryQuota(abc.ABC): - async def create_quota(self, scope_id: ProjectScope, quota: int) -> None: +class AbstractPerProjectContainerRegistryQuotaService(abc.ABC): + async def create_quota( + self, + client: AbstractPerProjectRegistryQuotaClient, + registry_info: ContainerRegistryRowInfo, + quota: int, + ) -> None: raise NotImplementedError - async def update_quota(self, scope_id: ProjectScope, quota: int) -> None: + async def update_quota( + self, + client: AbstractPerProjectRegistryQuotaClient, + registry_info: ContainerRegistryRowInfo, + quota: int, + ) -> None: raise NotImplementedError - async def delete_quota(self, scope_id: ProjectScope) -> None: + async def delete_quota( + self, + client: AbstractPerProjectRegistryQuotaClient, + registry_info: ContainerRegistryRowInfo, + ) -> None: raise NotImplementedError - async def read_quota(self, scope_id: ProjectScope) -> int: + async def read_quota( + self, + client: AbstractPerProjectRegistryQuotaClient, + registry_info: ContainerRegistryRowInfo, + ) -> int: raise NotImplementedError -class PerProjectContainerRegistryQuotaService(PerProjectContainerRegistryQuota): +class PerProjectContainerRegistryQuotaClientPool(abc.ABC): + @staticmethod + def make_client(type_: ContainerRegistryType) -> AbstractPerProjectRegistryQuotaClient: + match type_: + case ContainerRegistryType.HARBOR2: + return PerProjectHarborQuotaClient() + case _: + raise GenericBadRequest( + f"{type_} does not support registry quota per project management." + ) + + +class PerProjectContainerRegistryQuotaService(AbstractPerProjectContainerRegistryQuotaService): repository: PerProjectRegistryQuotaRepository - def __init__(self, repository: PerProjectRegistryQuotaRepository): + def __init__( + self, + repository: PerProjectRegistryQuotaRepository, + ): self.repository = repository def _registry_row_to_harbor_project_info( @@ -67,44 +99,47 @@ def _registry_row_to_harbor_project_info( ssl_verify=registry_info.ssl_verify, ) - def _make_client(self, type_: ContainerRegistryType) -> AbstractPerProjectRegistryQuotaClient: - match type_: - case ContainerRegistryType.HARBOR2: - return PerProjectHarborQuotaClient() - case _: - raise GenericBadRequest( - f"{type_} does not support registry quota per project management." - ) - @override - async def create_quota(self, scope_id: ProjectScope, quota: int) -> None: - registry_info = await self.repository.fetch_container_registry_row(scope_id) + async def create_quota( + self, + client: AbstractPerProjectRegistryQuotaClient, + registry_info: ContainerRegistryRowInfo, + quota: int, + ) -> None: project_info = self._registry_row_to_harbor_project_info(registry_info) credential = HarborAuthArgs( username=registry_info.username, password=registry_info.password ) - await self._make_client(registry_info.type).create_quota(project_info, quota, credential) + await client.create_quota(project_info, quota, credential) @override - async def update_quota(self, scope_id: ProjectScope, quota: int) -> None: - registry_info = await self.repository.fetch_container_registry_row(scope_id) + async def update_quota( + self, + client: AbstractPerProjectRegistryQuotaClient, + registry_info: ContainerRegistryRowInfo, + quota: int, + ) -> None: project_info = self._registry_row_to_harbor_project_info(registry_info) credential = HarborAuthArgs( username=registry_info.username, password=registry_info.password ) - await self._make_client(registry_info.type).update_quota(project_info, quota, credential) + await client.update_quota(project_info, quota, credential) @override - async def delete_quota(self, scope_id: ProjectScope) -> None: - registry_info = await self.repository.fetch_container_registry_row(scope_id) + async def delete_quota( + self, + client: AbstractPerProjectRegistryQuotaClient, + registry_info: ContainerRegistryRowInfo, + ) -> None: project_info = self._registry_row_to_harbor_project_info(registry_info) credential = HarborAuthArgs( username=registry_info.username, password=registry_info.password ) - await self._make_client(registry_info.type).delete_quota(project_info, credential) + await client.delete_quota(project_info, credential) @override - async def read_quota(self, scope_id: ProjectScope) -> int: - registry_info = await self.repository.fetch_container_registry_row(scope_id) + async def read_quota( + self, client: AbstractPerProjectRegistryQuotaClient, registry_info: ContainerRegistryRowInfo + ) -> int: project_info = self._registry_row_to_harbor_project_info(registry_info) - return await self._make_client(registry_info.type).read_quota(project_info) + return await client.read_quota(project_info) From 1b2bd519c6e396e18d26efb2cb7642a0a702f4bb Mon Sep 17 00:00:00 2001 From: jopemachine Date: Thu, 6 Feb 2025 09:30:22 +0000 Subject: [PATCH 69/75] refactor: improve abstraction --- src/ai/backend/manager/api/group.py | 10 +++---- .../models/gql_models/container_registry.py | 27 ++++++++++++++----- .../manager/models/gql_models/group.py | 11 +++++++- .../service/container_registry/harbor.py | 11 ++++++++ 4 files changed, 46 insertions(+), 13 deletions(-) diff --git a/src/ai/backend/manager/api/group.py b/src/ai/backend/manager/api/group.py index cb0c4ad22d..f45e20dbc0 100644 --- a/src/ai/backend/manager/api/group.py +++ b/src/ai/backend/manager/api/group.py @@ -40,7 +40,7 @@ async def update_registry_quota(request: web.Request, params: Any) -> web.Respon scope_id = ProjectScope(project_id=group_id, domain_name=None) quota = int(params["quota"]) - registry_info = await root_ctx.services_ctx.per_project_container_registries_quota.repository.fetch_container_registry_row( + registry_info = await root_ctx.services_ctx.per_project_container_registries_quota.fetch_container_registry_row( scope_id ) client = PerProjectContainerRegistryQuotaClientPool.make_client(registry_info.type) @@ -63,10 +63,10 @@ async def delete_registry_quota(request: web.Request, params: Any) -> web.Respon group_id = params["group_id"] scope_id = ProjectScope(project_id=group_id, domain_name=None) - registry_info = await root_ctx.services_ctx.per_project_container_registries_quota.repository.fetch_container_registry_row( + registry_info = await root_ctx.services_ctx.per_project_container_registries_quota.fetch_container_registry_row( scope_id ) - client = PerProjectContainerRegistryQuotaClientPool.make_client(scope_id) + client = PerProjectContainerRegistryQuotaClientPool.make_client(registry_info.type) await root_ctx.services_ctx.per_project_container_registries_quota.delete_quota( client, registry_info ) @@ -88,7 +88,7 @@ async def create_registry_quota(request: web.Request, params: Any) -> web.Respon scope_id = ProjectScope(project_id=group_id, domain_name=None) quota = int(params["quota"]) - registry_info = await root_ctx.services_ctx.per_project_container_registries_quota.repository.fetch_container_registry_row( + registry_info = await root_ctx.services_ctx.per_project_container_registries_quota.fetch_container_registry_row( scope_id ) client = PerProjectContainerRegistryQuotaClientPool.make_client(registry_info.type) @@ -111,7 +111,7 @@ async def read_registry_quota(request: web.Request, params: Any) -> web.Response group_id = params["group_id"] scope_id = ProjectScope(project_id=group_id, domain_name=None) - registry_info = await root_ctx.services_ctx.per_project_container_registries_quota.repository.fetch_container_registry_row( + registry_info = await root_ctx.services_ctx.per_project_container_registries_quota.fetch_container_registry_row( scope_id ) client = PerProjectContainerRegistryQuotaClientPool.make_client(registry_info.type) 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 cc47e562a3..0ae74b17d1 100644 --- a/src/ai/backend/manager/models/gql_models/container_registry.py +++ b/src/ai/backend/manager/models/gql_models/container_registry.py @@ -517,14 +517,15 @@ async def mutate( quota: int | float, ) -> Self: graph_ctx: GraphQueryContext = info.context - - registry_info = await graph_ctx.services_ctx.per_project_container_registries_quota.repository.fetch_container_registry_row( - scope_id - ) - client = PerProjectContainerRegistryQuotaClientPool.make_client(registry_info.type) try: match scope_id: case ProjectScope(): + registry_info = await graph_ctx.services_ctx.per_project_container_registries_quota.fetch_container_registry_row( + scope_id + ) + client = PerProjectContainerRegistryQuotaClientPool.make_client( + registry_info.type + ) await ( graph_ctx.services_ctx.per_project_container_registries_quota.create_quota( client, registry_info, int(quota) @@ -565,9 +566,15 @@ async def mutate( try: match scope_id: case ProjectScope(_): + registry_info = await graph_ctx.services_ctx.per_project_container_registries_quota.fetch_container_registry_row( + scope_id + ) + client = PerProjectContainerRegistryQuotaClientPool.make_client( + registry_info.type + ) await ( graph_ctx.services_ctx.per_project_container_registries_quota.update_quota( - scope_id, int(quota) + client, registry_info, int(quota) ) ) case _: @@ -603,9 +610,15 @@ async def mutate( try: match scope_id: case ProjectScope(_): + registry_info = await graph_ctx.services_ctx.per_project_container_registries_quota.fetch_container_registry_row( + scope_id + ) + client = PerProjectContainerRegistryQuotaClientPool.make_client( + registry_info.type + ) await ( graph_ctx.services_ctx.per_project_container_registries_quota.delete_quota( - scope_id + client, registry_info ) ) case _: diff --git a/src/ai/backend/manager/models/gql_models/group.py b/src/ai/backend/manager/models/gql_models/group.py index aacf1b25cf..1a934e3564 100644 --- a/src/ai/backend/manager/models/gql_models/group.py +++ b/src/ai/backend/manager/models/gql_models/group.py @@ -16,6 +16,9 @@ from graphene.types.datetime import DateTime as GQLDateTime from ai.backend.manager.models.rbac import ProjectScope +from ai.backend.manager.service.container_registry.harbor import ( + PerProjectContainerRegistryQuotaClientPool, +) from ..base import ( BigInt, @@ -218,9 +221,15 @@ async def resolve_user_nodes( async def resolve_registry_quota(self, info: graphene.ResolveInfo) -> int: graph_ctx: GraphQueryContext = info.context scope_id = ProjectScope(project_id=self.id, domain_name=None) - return await graph_ctx.services_ctx.per_project_container_registries_quota.read_quota( + + registry_info = await graph_ctx.services_ctx.per_project_container_registries_quota.fetch_container_registry_row( scope_id ) + client = PerProjectContainerRegistryQuotaClientPool.make_client(registry_info.type) + + return await graph_ctx.services_ctx.per_project_container_registries_quota.read_quota( + client, registry_info + ) @classmethod async def get_node(cls, info: graphene.ResolveInfo, id) -> Self: diff --git a/src/ai/backend/manager/service/container_registry/harbor.py b/src/ai/backend/manager/service/container_registry/harbor.py index fce41bc8b8..6618c3c40d 100644 --- a/src/ai/backend/manager/service/container_registry/harbor.py +++ b/src/ai/backend/manager/service/container_registry/harbor.py @@ -12,6 +12,7 @@ AbstractPerProjectRegistryQuotaClient, PerProjectHarborQuotaClient, ) +from ai.backend.manager.models.rbac import ProjectScope from ai.backend.manager.service.container_registry.base import ( ContainerRegistryRowInfo, PerProjectRegistryQuotaRepository, @@ -68,6 +69,11 @@ async def read_quota( ) -> int: raise NotImplementedError + async def fetch_container_registry_row( + self, scope_id: ProjectScope + ) -> ContainerRegistryRowInfo: + raise NotImplementedError + class PerProjectContainerRegistryQuotaClientPool(abc.ABC): @staticmethod @@ -143,3 +149,8 @@ async def read_quota( ) -> int: project_info = self._registry_row_to_harbor_project_info(registry_info) return await client.read_quota(project_info) + + async def fetch_container_registry_row( + self, scope_id: ProjectScope + ) -> ContainerRegistryRowInfo: + return await self.repository.fetch_container_registry_row(scope_id) From 2bcb1cdd4afd4fe2d1def4de6c3f2756ea036cda Mon Sep 17 00:00:00 2001 From: jopemachine Date: Thu, 6 Feb 2025 09:34:07 +0000 Subject: [PATCH 70/75] fix: Broken import --- src/ai/backend/manager/models/gql_models/group.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/ai/backend/manager/models/gql_models/group.py b/src/ai/backend/manager/models/gql_models/group.py index 1a934e3564..49f773023a 100644 --- a/src/ai/backend/manager/models/gql_models/group.py +++ b/src/ai/backend/manager/models/gql_models/group.py @@ -16,9 +16,6 @@ from graphene.types.datetime import DateTime as GQLDateTime from ai.backend.manager.models.rbac import ProjectScope -from ai.backend.manager.service.container_registry.harbor import ( - PerProjectContainerRegistryQuotaClientPool, -) from ..base import ( BigInt, @@ -219,6 +216,10 @@ async def resolve_user_nodes( return ConnectionResolverResult(result, cursor, pagination_order, page_size, total_cnt) async def resolve_registry_quota(self, info: graphene.ResolveInfo) -> int: + from ai.backend.manager.service.container_registry.harbor import ( + PerProjectContainerRegistryQuotaClientPool, + ) + graph_ctx: GraphQueryContext = info.context scope_id = ProjectScope(project_id=self.id, domain_name=None) From a72a3fee0e2831d1290fe552109bcd1576e64d7e Mon Sep 17 00:00:00 2001 From: jopemachine Date: Thu, 6 Feb 2025 09:48:30 +0000 Subject: [PATCH 71/75] fix: Reflect feedback --- src/ai/backend/manager/api/group.py | 19 ++++++++++------- .../models/gql_models/container_registry.py | 21 +++++++++++-------- .../manager/models/gql_models/group.py | 8 +++---- src/ai/backend/manager/server.py | 4 +++- .../service/container_registry/harbor.py | 20 +++++++++++++----- 5 files changed, 45 insertions(+), 27 deletions(-) diff --git a/src/ai/backend/manager/api/group.py b/src/ai/backend/manager/api/group.py index f45e20dbc0..a1493f87ac 100644 --- a/src/ai/backend/manager/api/group.py +++ b/src/ai/backend/manager/api/group.py @@ -10,9 +10,6 @@ from ai.backend.common import validators as tx from ai.backend.logging import BraceStyleAdapter from ai.backend.manager.models.rbac import ProjectScope -from ai.backend.manager.service.container_registry.harbor import ( - PerProjectContainerRegistryQuotaClientPool, -) if TYPE_CHECKING: from .context import RootContext @@ -43,7 +40,9 @@ async def update_registry_quota(request: web.Request, params: Any) -> web.Respon registry_info = await root_ctx.services_ctx.per_project_container_registries_quota.fetch_container_registry_row( scope_id ) - client = PerProjectContainerRegistryQuotaClientPool.make_client(registry_info.type) + client = root_ctx.services_ctx.per_project_container_registries_quota.make_client( + registry_info.type + ) await root_ctx.services_ctx.per_project_container_registries_quota.update_quota( client, registry_info, quota ) @@ -66,7 +65,9 @@ async def delete_registry_quota(request: web.Request, params: Any) -> web.Respon registry_info = await root_ctx.services_ctx.per_project_container_registries_quota.fetch_container_registry_row( scope_id ) - client = PerProjectContainerRegistryQuotaClientPool.make_client(registry_info.type) + client = root_ctx.services_ctx.per_project_container_registries_quota.make_client( + registry_info.type + ) await root_ctx.services_ctx.per_project_container_registries_quota.delete_quota( client, registry_info ) @@ -91,7 +92,9 @@ async def create_registry_quota(request: web.Request, params: Any) -> web.Respon registry_info = await root_ctx.services_ctx.per_project_container_registries_quota.fetch_container_registry_row( scope_id ) - client = PerProjectContainerRegistryQuotaClientPool.make_client(registry_info.type) + client = root_ctx.services_ctx.per_project_container_registries_quota.make_client( + registry_info.type + ) await root_ctx.services_ctx.per_project_container_registries_quota.create_quota( client, registry_info, quota ) @@ -114,7 +117,9 @@ async def read_registry_quota(request: web.Request, params: Any) -> web.Response registry_info = await root_ctx.services_ctx.per_project_container_registries_quota.fetch_container_registry_row( scope_id ) - client = PerProjectContainerRegistryQuotaClientPool.make_client(registry_info.type) + client = root_ctx.services_ctx.per_project_container_registries_quota.make_client( + registry_info.type + ) quota = await root_ctx.services_ctx.per_project_container_registries_quota.read_quota( client, registry_info ) 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 0ae74b17d1..d376963354 100644 --- a/src/ai/backend/manager/models/gql_models/container_registry.py +++ b/src/ai/backend/manager/models/gql_models/container_registry.py @@ -25,9 +25,6 @@ ) from ai.backend.manager.models.user import UserRole from ai.backend.manager.models.utils import ExtendedAsyncSAEngine -from ai.backend.manager.service.container_registry.harbor import ( - PerProjectContainerRegistryQuotaClientPool, -) from ...defs import PASSWORD_PLACEHOLDER from ..association_container_registries_groups import ( @@ -523,8 +520,10 @@ async def mutate( registry_info = await graph_ctx.services_ctx.per_project_container_registries_quota.fetch_container_registry_row( scope_id ) - client = PerProjectContainerRegistryQuotaClientPool.make_client( - registry_info.type + client = ( + graph_ctx.services_ctx.per_project_container_registries_quota.make_client( + registry_info.type + ) ) await ( graph_ctx.services_ctx.per_project_container_registries_quota.create_quota( @@ -569,8 +568,10 @@ async def mutate( registry_info = await graph_ctx.services_ctx.per_project_container_registries_quota.fetch_container_registry_row( scope_id ) - client = PerProjectContainerRegistryQuotaClientPool.make_client( - registry_info.type + client = ( + graph_ctx.services_ctx.per_project_container_registries_quota.make_client( + registry_info.type + ) ) await ( graph_ctx.services_ctx.per_project_container_registries_quota.update_quota( @@ -613,8 +614,10 @@ async def mutate( registry_info = await graph_ctx.services_ctx.per_project_container_registries_quota.fetch_container_registry_row( scope_id ) - client = PerProjectContainerRegistryQuotaClientPool.make_client( - registry_info.type + client = ( + graph_ctx.services_ctx.per_project_container_registries_quota.make_client( + registry_info.type + ) ) await ( graph_ctx.services_ctx.per_project_container_registries_quota.delete_quota( diff --git a/src/ai/backend/manager/models/gql_models/group.py b/src/ai/backend/manager/models/gql_models/group.py index 49f773023a..726e4dc0c8 100644 --- a/src/ai/backend/manager/models/gql_models/group.py +++ b/src/ai/backend/manager/models/gql_models/group.py @@ -216,17 +216,15 @@ async def resolve_user_nodes( return ConnectionResolverResult(result, cursor, pagination_order, page_size, total_cnt) async def resolve_registry_quota(self, info: graphene.ResolveInfo) -> int: - from ai.backend.manager.service.container_registry.harbor import ( - PerProjectContainerRegistryQuotaClientPool, - ) - graph_ctx: GraphQueryContext = info.context scope_id = ProjectScope(project_id=self.id, domain_name=None) registry_info = await graph_ctx.services_ctx.per_project_container_registries_quota.fetch_container_registry_row( scope_id ) - client = PerProjectContainerRegistryQuotaClientPool.make_client(registry_info.type) + client = graph_ctx.services_ctx.per_project_container_registries_quota.make_client( + registry_info.type + ) return await graph_ctx.services_ctx.per_project_container_registries_quota.read_quota( client, registry_info diff --git a/src/ai/backend/manager/server.py b/src/ai/backend/manager/server.py index c93f31d443..3a7ba95d10 100644 --- a/src/ai/backend/manager/server.py +++ b/src/ai/backend/manager/server.py @@ -66,6 +66,7 @@ from ai.backend.manager.service.base import ServicesContext from ai.backend.manager.service.container_registry.base import PerProjectRegistryQuotaRepository from ai.backend.manager.service.container_registry.harbor import ( + PerProjectContainerRegistryQuotaClientPool, PerProjectContainerRegistryQuotaService, ) @@ -702,7 +703,8 @@ async def services_ctx(root_ctx: RootContext) -> AsyncIterator[None]: db = root_ctx.db per_project_container_registries_quota = PerProjectContainerRegistryQuotaService( - repository=PerProjectRegistryQuotaRepository(db) + repository=PerProjectRegistryQuotaRepository(db), + client_pool=PerProjectContainerRegistryQuotaClientPool(), ) root_ctx.services_ctx = ServicesContext( diff --git a/src/ai/backend/manager/service/container_registry/harbor.py b/src/ai/backend/manager/service/container_registry/harbor.py index 6618c3c40d..e19536749a 100644 --- a/src/ai/backend/manager/service/container_registry/harbor.py +++ b/src/ai/backend/manager/service/container_registry/harbor.py @@ -74,10 +74,12 @@ async def fetch_container_registry_row( ) -> ContainerRegistryRowInfo: raise NotImplementedError + def make_client(self, type_: ContainerRegistryType) -> AbstractPerProjectRegistryQuotaClient: + raise NotImplementedError + class PerProjectContainerRegistryQuotaClientPool(abc.ABC): - @staticmethod - def make_client(type_: ContainerRegistryType) -> AbstractPerProjectRegistryQuotaClient: + def make_client(self, type_: ContainerRegistryType) -> AbstractPerProjectRegistryQuotaClient: match type_: case ContainerRegistryType.HARBOR2: return PerProjectHarborQuotaClient() @@ -88,13 +90,16 @@ def make_client(type_: ContainerRegistryType) -> AbstractPerProjectRegistryQuota class PerProjectContainerRegistryQuotaService(AbstractPerProjectContainerRegistryQuotaService): - repository: PerProjectRegistryQuotaRepository + _repository: PerProjectRegistryQuotaRepository + _client_pool: PerProjectContainerRegistryQuotaClientPool def __init__( self, repository: PerProjectRegistryQuotaRepository, + client_pool: PerProjectContainerRegistryQuotaClientPool, ): - self.repository = repository + self._repository = repository + self._client_pool = client_pool def _registry_row_to_harbor_project_info( self, registry_info: ContainerRegistryRowInfo @@ -150,7 +155,12 @@ async def read_quota( project_info = self._registry_row_to_harbor_project_info(registry_info) return await client.read_quota(project_info) + @override async def fetch_container_registry_row( self, scope_id: ProjectScope ) -> ContainerRegistryRowInfo: - return await self.repository.fetch_container_registry_row(scope_id) + return await self._repository.fetch_container_registry_row(scope_id) + + @override + def make_client(self, type_: ContainerRegistryType) -> AbstractPerProjectRegistryQuotaClient: + return self._client_pool.make_client(type_) From f477f33ee2db722d246cec2b09091968c2382f38 Mon Sep 17 00:00:00 2001 From: jopemachine Date: Thu, 6 Feb 2025 13:41:04 +0000 Subject: [PATCH 72/75] fix: Reflect feedback --- src/ai/backend/manager/api/group.py | 38 ++------------ .../models/gql_models/container_registry.py | 30 ++--------- .../manager/models/gql_models/group.py | 9 +--- .../service/container_registry/harbor.py | 51 ++++++------------- 4 files changed, 24 insertions(+), 104 deletions(-) diff --git a/src/ai/backend/manager/api/group.py b/src/ai/backend/manager/api/group.py index a1493f87ac..e0b170ad8e 100644 --- a/src/ai/backend/manager/api/group.py +++ b/src/ai/backend/manager/api/group.py @@ -37,15 +37,7 @@ async def update_registry_quota(request: web.Request, params: Any) -> web.Respon scope_id = ProjectScope(project_id=group_id, domain_name=None) quota = int(params["quota"]) - registry_info = await root_ctx.services_ctx.per_project_container_registries_quota.fetch_container_registry_row( - scope_id - ) - client = root_ctx.services_ctx.per_project_container_registries_quota.make_client( - registry_info.type - ) - await root_ctx.services_ctx.per_project_container_registries_quota.update_quota( - client, registry_info, quota - ) + await root_ctx.services_ctx.per_project_container_registries_quota.update_quota(scope_id, quota) return web.Response(status=204) @@ -62,15 +54,7 @@ async def delete_registry_quota(request: web.Request, params: Any) -> web.Respon group_id = params["group_id"] scope_id = ProjectScope(project_id=group_id, domain_name=None) - registry_info = await root_ctx.services_ctx.per_project_container_registries_quota.fetch_container_registry_row( - scope_id - ) - client = root_ctx.services_ctx.per_project_container_registries_quota.make_client( - registry_info.type - ) - await root_ctx.services_ctx.per_project_container_registries_quota.delete_quota( - client, registry_info - ) + await root_ctx.services_ctx.per_project_container_registries_quota.delete_quota(scope_id) return web.Response(status=204) @@ -89,15 +73,7 @@ async def create_registry_quota(request: web.Request, params: Any) -> web.Respon scope_id = ProjectScope(project_id=group_id, domain_name=None) quota = int(params["quota"]) - registry_info = await root_ctx.services_ctx.per_project_container_registries_quota.fetch_container_registry_row( - scope_id - ) - client = root_ctx.services_ctx.per_project_container_registries_quota.make_client( - registry_info.type - ) - await root_ctx.services_ctx.per_project_container_registries_quota.create_quota( - client, registry_info, quota - ) + await root_ctx.services_ctx.per_project_container_registries_quota.create_quota(scope_id, quota) return web.Response(status=204) @@ -114,14 +90,8 @@ async def read_registry_quota(request: web.Request, params: Any) -> web.Response group_id = params["group_id"] scope_id = ProjectScope(project_id=group_id, domain_name=None) - registry_info = await root_ctx.services_ctx.per_project_container_registries_quota.fetch_container_registry_row( - scope_id - ) - client = root_ctx.services_ctx.per_project_container_registries_quota.make_client( - registry_info.type - ) quota = await root_ctx.services_ctx.per_project_container_registries_quota.read_quota( - client, registry_info + scope_id, ) 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 d376963354..6c395294d8 100644 --- a/src/ai/backend/manager/models/gql_models/container_registry.py +++ b/src/ai/backend/manager/models/gql_models/container_registry.py @@ -517,17 +517,9 @@ async def mutate( try: match scope_id: case ProjectScope(): - registry_info = await graph_ctx.services_ctx.per_project_container_registries_quota.fetch_container_registry_row( - scope_id - ) - client = ( - graph_ctx.services_ctx.per_project_container_registries_quota.make_client( - registry_info.type - ) - ) await ( graph_ctx.services_ctx.per_project_container_registries_quota.create_quota( - client, registry_info, int(quota) + scope_id, int(quota) ) ) case _: @@ -565,17 +557,9 @@ async def mutate( try: match scope_id: case ProjectScope(_): - registry_info = await graph_ctx.services_ctx.per_project_container_registries_quota.fetch_container_registry_row( - scope_id - ) - client = ( - graph_ctx.services_ctx.per_project_container_registries_quota.make_client( - registry_info.type - ) - ) await ( graph_ctx.services_ctx.per_project_container_registries_quota.update_quota( - client, registry_info, int(quota) + scope_id, int(quota) ) ) case _: @@ -611,17 +595,9 @@ async def mutate( try: match scope_id: case ProjectScope(_): - registry_info = await graph_ctx.services_ctx.per_project_container_registries_quota.fetch_container_registry_row( - scope_id - ) - client = ( - graph_ctx.services_ctx.per_project_container_registries_quota.make_client( - registry_info.type - ) - ) await ( graph_ctx.services_ctx.per_project_container_registries_quota.delete_quota( - client, registry_info + scope_id ) ) case _: diff --git a/src/ai/backend/manager/models/gql_models/group.py b/src/ai/backend/manager/models/gql_models/group.py index 726e4dc0c8..26cb9244ad 100644 --- a/src/ai/backend/manager/models/gql_models/group.py +++ b/src/ai/backend/manager/models/gql_models/group.py @@ -219,15 +219,8 @@ async def resolve_registry_quota(self, info: graphene.ResolveInfo) -> int: graph_ctx: GraphQueryContext = info.context scope_id = ProjectScope(project_id=self.id, domain_name=None) - registry_info = await graph_ctx.services_ctx.per_project_container_registries_quota.fetch_container_registry_row( - scope_id - ) - client = graph_ctx.services_ctx.per_project_container_registries_quota.make_client( - registry_info.type - ) - return await graph_ctx.services_ctx.per_project_container_registries_quota.read_quota( - client, registry_info + scope_id, ) @classmethod diff --git a/src/ai/backend/manager/service/container_registry/harbor.py b/src/ai/backend/manager/service/container_registry/harbor.py index e19536749a..c666ff2e9c 100644 --- a/src/ai/backend/manager/service/container_registry/harbor.py +++ b/src/ai/backend/manager/service/container_registry/harbor.py @@ -41,42 +41,30 @@ class HarborProjectQuotaInfo(TypedDict): class AbstractPerProjectContainerRegistryQuotaService(abc.ABC): async def create_quota( self, - client: AbstractPerProjectRegistryQuotaClient, - registry_info: ContainerRegistryRowInfo, + scope_id: ProjectScope, quota: int, ) -> None: raise NotImplementedError async def update_quota( self, - client: AbstractPerProjectRegistryQuotaClient, - registry_info: ContainerRegistryRowInfo, + scope_id: ProjectScope, quota: int, ) -> None: raise NotImplementedError async def delete_quota( self, - client: AbstractPerProjectRegistryQuotaClient, - registry_info: ContainerRegistryRowInfo, + scope_id: ProjectScope, ) -> None: raise NotImplementedError async def read_quota( self, - client: AbstractPerProjectRegistryQuotaClient, - registry_info: ContainerRegistryRowInfo, + scope_id: ProjectScope, ) -> int: raise NotImplementedError - async def fetch_container_registry_row( - self, scope_id: ProjectScope - ) -> ContainerRegistryRowInfo: - raise NotImplementedError - - def make_client(self, type_: ContainerRegistryType) -> AbstractPerProjectRegistryQuotaClient: - raise NotImplementedError - class PerProjectContainerRegistryQuotaClientPool(abc.ABC): def make_client(self, type_: ContainerRegistryType) -> AbstractPerProjectRegistryQuotaClient: @@ -113,10 +101,11 @@ def _registry_row_to_harbor_project_info( @override async def create_quota( self, - client: AbstractPerProjectRegistryQuotaClient, - registry_info: ContainerRegistryRowInfo, + scope_id: ProjectScope, quota: int, ) -> None: + registry_info = await self._repository.fetch_container_registry_row(scope_id) + client = self._client_pool.make_client(registry_info.type) project_info = self._registry_row_to_harbor_project_info(registry_info) credential = HarborAuthArgs( username=registry_info.username, password=registry_info.password @@ -126,10 +115,11 @@ async def create_quota( @override async def update_quota( self, - client: AbstractPerProjectRegistryQuotaClient, - registry_info: ContainerRegistryRowInfo, + scope_id: ProjectScope, quota: int, ) -> None: + registry_info = await self._repository.fetch_container_registry_row(scope_id) + client = self._client_pool.make_client(registry_info.type) project_info = self._registry_row_to_harbor_project_info(registry_info) credential = HarborAuthArgs( username=registry_info.username, password=registry_info.password @@ -139,9 +129,10 @@ async def update_quota( @override async def delete_quota( self, - client: AbstractPerProjectRegistryQuotaClient, - registry_info: ContainerRegistryRowInfo, + scope_id: ProjectScope, ) -> None: + registry_info = await self._repository.fetch_container_registry_row(scope_id) + client = self._client_pool.make_client(registry_info.type) project_info = self._registry_row_to_harbor_project_info(registry_info) credential = HarborAuthArgs( username=registry_info.username, password=registry_info.password @@ -149,18 +140,8 @@ async def delete_quota( await client.delete_quota(project_info, credential) @override - async def read_quota( - self, client: AbstractPerProjectRegistryQuotaClient, registry_info: ContainerRegistryRowInfo - ) -> int: + async def read_quota(self, scope_id: ProjectScope) -> int: + registry_info = await self._repository.fetch_container_registry_row(scope_id) + client = self._client_pool.make_client(registry_info.type) project_info = self._registry_row_to_harbor_project_info(registry_info) return await client.read_quota(project_info) - - @override - async def fetch_container_registry_row( - self, scope_id: ProjectScope - ) -> ContainerRegistryRowInfo: - return await self._repository.fetch_container_registry_row(scope_id) - - @override - def make_client(self, type_: ContainerRegistryType) -> AbstractPerProjectRegistryQuotaClient: - return self._client_pool.make_client(type_) From c381c71cef381ddfe11e1eecf87e3e44d6d08a30 Mon Sep 17 00:00:00 2001 From: jopemachine Date: Mon, 10 Feb 2025 08:48:28 +0000 Subject: [PATCH 73/75] fix: Update version --- src/ai/backend/manager/models/gql.py | 6 +++--- .../manager/models/gql_models/container_registry.py | 11 ++++------- src/ai/backend/manager/models/gql_models/group.py | 2 +- 3 files changed, 8 insertions(+), 11 deletions(-) diff --git a/src/ai/backend/manager/models/gql.py b/src/ai/backend/manager/models/gql.py index df03b09bc7..b22f5b8a71 100644 --- a/src/ai/backend/manager/models/gql.py +++ b/src/ai/backend/manager/models/gql.py @@ -390,13 +390,13 @@ class Mutations(graphene.ObjectType): description="Added in 25.1.0." ) create_container_registry_quota = CreateContainerRegistryQuota.Field( - description="Added in 25.2.0." + description="Added in 25.3.0." ) update_container_registry_quota = UpdateContainerRegistryQuota.Field( - description="Added in 25.2.0." + description="Added in 25.3.0." ) delete_container_registry_quota = DeleteContainerRegistryQuota.Field( - description="Added in 25.2.0." + description="Added in 25.3.0." ) # Legacy mutations 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 6c395294d8..fabda6098d 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,6 @@ from graphql import Undefined, UndefinedType from ai.backend.common.container_registry import AllowedGroupsModel, ContainerRegistryType -from ai.backend.common.logging_utils import BraceStyleAdapter from ai.backend.logging import BraceStyleAdapter from ai.backend.manager.api.exceptions import ( ContainerRegistryNotFound, @@ -20,10 +19,10 @@ from ai.backend.manager.models.gql_models.fields import ScopeField from ai.backend.manager.models.rbac import ( ContainerRegistryScope, + ProjectScope, ScopeType, SystemScope, ) -from ai.backend.manager.models.user import UserRole from ai.backend.manager.models.utils import ExtendedAsyncSAEngine from ...defs import PASSWORD_PLACEHOLDER @@ -44,9 +43,7 @@ if TYPE_CHECKING: from ..gql import GraphQueryContext -from ..rbac import ProjectScope, ScopeType from ..user import UserRole -from .fields import ScopeField log = BraceStyleAdapter(logging.getLogger(__spec__.name)) # type: ignore @@ -491,7 +488,7 @@ async def mutate( class CreateContainerRegistryQuota(graphene.Mutation): - """Added in 25.2.0.""" + """Added in 25.3.0.""" allowed_roles = ( UserRole.SUPERADMIN, @@ -531,7 +528,7 @@ async def mutate( class UpdateContainerRegistryQuota(graphene.Mutation): - """Added in 25.2.0.""" + """Added in 25.3.0.""" allowed_roles = ( UserRole.SUPERADMIN, @@ -571,7 +568,7 @@ async def mutate( class DeleteContainerRegistryQuota(graphene.Mutation): - """Added in 25.2.0.""" + """Added in 25.3.0.""" allowed_roles = ( UserRole.SUPERADMIN, diff --git a/src/ai/backend/manager/models/gql_models/group.py b/src/ai/backend/manager/models/gql_models/group.py index 26cb9244ad..aab50bb303 100644 --- a/src/ai/backend/manager/models/gql_models/group.py +++ b/src/ai/backend/manager/models/gql_models/group.py @@ -121,7 +121,7 @@ class Meta: lambda: graphene.String, ) - registry_quota = BigInt(description="Added in 25.2.0.") + registry_quota = BigInt(description="Added in 25.3.0.") user_nodes = PaginatedConnectionField( UserConnection, From e552f0a041e6728f57ad04f24583e169d538a9a3 Mon Sep 17 00:00:00 2001 From: jopemachine Date: Mon, 10 Feb 2025 08:49:53 +0000 Subject: [PATCH 74/75] fix: Update version --- docs/manager/graphql-reference/schema.graphql | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/docs/manager/graphql-reference/schema.graphql b/docs/manager/graphql-reference/schema.graphql index 2db3ae19c4..74d56366fc 100644 --- a/docs/manager/graphql-reference/schema.graphql +++ b/docs/manager/graphql-reference/schema.graphql @@ -720,7 +720,7 @@ type GroupNode implements Node { container_registry: JSONString scaling_groups: [String] - """Added in 25.2.0.""" + """Added in 25.3.0.""" registry_quota: BigInt user_nodes(filter: String, order: String, offset: Int, before: String, after: String, first: Int, last: Int): UserConnection } @@ -2056,13 +2056,13 @@ type Mutations { """Added in 25.1.0.""" delete_endpoint_auto_scaling_rule_node(id: String!): DeleteEndpointAutoScalingRuleNode - """Added in 25.2.0.""" + """Added in 25.3.0.""" create_container_registry_quota(quota: BigInt!, scope_id: ScopeField!): CreateContainerRegistryQuota - """Added in 25.2.0.""" + """Added in 25.3.0.""" update_container_registry_quota(quota: BigInt!, scope_id: ScopeField!): UpdateContainerRegistryQuota - """Added in 25.2.0.""" + """Added in 25.3.0.""" delete_container_registry_quota(scope_id: ScopeField!): DeleteContainerRegistryQuota """Deprecated since 24.09.0. use `CreateContainerRegistryNode` instead""" @@ -2967,19 +2967,19 @@ type DeleteEndpointAutoScalingRuleNode { msg: String } -"""Added in 25.2.0.""" +"""Added in 25.3.0.""" type CreateContainerRegistryQuota { ok: Boolean msg: String } -"""Added in 25.2.0.""" +"""Added in 25.3.0.""" type UpdateContainerRegistryQuota { ok: Boolean msg: String } -"""Added in 25.2.0.""" +"""Added in 25.3.0.""" type DeleteContainerRegistryQuota { ok: Boolean msg: String From b695731074e37aac26e605061035153c489eccec Mon Sep 17 00:00:00 2001 From: jopemachine Date: Mon, 10 Feb 2025 11:00:07 +0000 Subject: [PATCH 75/75] fix: Broken CI --- tests/manager/models/gql_models/test_container_registry_nodes.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/manager/models/gql_models/test_container_registry_nodes.py b/tests/manager/models/gql_models/test_container_registry_nodes.py index a99d9ae8fd..98df3c8e56 100644 --- a/tests/manager/models/gql_models/test_container_registry_nodes.py +++ b/tests/manager/models/gql_models/test_container_registry_nodes.py @@ -65,6 +65,7 @@ def mock_shared_config_api_getitem(key): registry=None, # type: ignore idle_checker_host=None, # type: ignore network_plugin_ctx=None, # type: ignore + services_ctx=None, # type: ignore )