feat: Add Associate and Disassociate Scaling Groups with Domain, User…

… Group, and KeyPair mutations (#2473) (#2724)

Co-authored-by: Sujin Kim <kimsujin@lablup.com>
Co-authored-by: Sanghun Lee <sanghun@lablup.com>

changes/2473.feature.md

@@ -0,0 +1 @@
+Allow bulk association and disassociation of scaling groups with domains, user groups, and key pairs.

src/ai/backend/manager/api/schema.graphql

@@ -1179,11 +1179,29 @@ type Mutations {
   modify_scaling_group(name: String!, props: ModifyScalingGroupInput!): ModifyScalingGroup
   delete_scaling_group(name: String!): DeleteScalingGroup
   associate_scaling_group_with_domain(domain: String!, scaling_group: String!): AssociateScalingGroupWithDomain
+
+  """Added in 24.03.9"""
+  associate_scaling_groups_with_domain(domain: String!, scaling_groups: [String]!): AssociateScalingGroupsWithDomain
   associate_scaling_group_with_user_group(scaling_group: String!, user_group: UUID!): AssociateScalingGroupWithUserGroup
+
+  """Added in 24.03.9"""
+  associate_scaling_groups_with_user_group(scaling_groups: [String]!, user_group: UUID!): AssociateScalingGroupsWithUserGroup
   associate_scaling_group_with_keypair(access_key: String!, scaling_group: String!): AssociateScalingGroupWithKeyPair
+
+  """Added in 24.03.9"""
+  associate_scaling_groups_with_keypair(access_key: String!, scaling_groups: [String]!): AssociateScalingGroupsWithKeyPair
   disassociate_scaling_group_with_domain(domain: String!, scaling_group: String!): DisassociateScalingGroupWithDomain
+
+  """Added in 24.03.9"""
+  disassociate_scaling_groups_with_domain(domain: String!, scaling_groups: [String]!): DisassociateScalingGroupsWithDomain
   disassociate_scaling_group_with_user_group(scaling_group: String!, user_group: UUID!): DisassociateScalingGroupWithUserGroup
+
+  """Added in 24.03.9"""
+  disassociate_scaling_groups_with_user_group(scaling_groups: [String]!, user_group: UUID!): DisassociateScalingGroupsWithUserGroup
   disassociate_scaling_group_with_keypair(access_key: String!, scaling_group: String!): DisassociateScalingGroupWithKeyPair
+
+  """Added in 24.03.9"""
+  disassociate_scaling_groups_with_keypair(access_key: String!, scaling_groups: [String]!): DisassociateScalingGroupsWithKeyPair
   disassociate_all_scaling_groups_with_domain(domain: String!): DisassociateAllScalingGroupsWithDomain
   disassociate_all_scaling_groups_with_group(user_group: UUID!): DisassociateAllScalingGroupsWithGroup
   set_quota_scope(props: QuotaScopeInput!, quota_scope_id: String!, storage_host_name: String!): SetQuotaScope
@@ -1747,31 +1765,67 @@ type AssociateScalingGroupWithDomain {
   msg: String
 }
 
+"""Added in 24.03.9."""
+type AssociateScalingGroupsWithDomain {
+  ok: Boolean
+  msg: String
+}
+
 type AssociateScalingGroupWithUserGroup {
   ok: Boolean
   msg: String
 }
 
+"""Added in 24.03.9."""
+type AssociateScalingGroupsWithUserGroup {
+  ok: Boolean
+  msg: String
+}
+
 type AssociateScalingGroupWithKeyPair {
   ok: Boolean
   msg: String
 }
 
+"""Added in 24.03.9."""
+type AssociateScalingGroupsWithKeyPair {
+  ok: Boolean
+  msg: String
+}
+
 type DisassociateScalingGroupWithDomain {
   ok: Boolean
   msg: String
 }
 
+"""Added in 24.03.9."""
+type DisassociateScalingGroupsWithDomain {
+  ok: Boolean
+  msg: String
+}
+
 type DisassociateScalingGroupWithUserGroup {
   ok: Boolean
   msg: String
 }
 
+"""Added in 24.03.9."""
+type DisassociateScalingGroupsWithUserGroup {
+  ok: Boolean
+  msg: String
+}
+
 type DisassociateScalingGroupWithKeyPair {
   ok: Boolean
   msg: String
 }
 
+"""Added in 24.03.9."""
+type DisassociateScalingGroupsWithKeyPair {
+  ok: Boolean
+  msg: String
+}
+
 type DisassociateAllScalingGroupsWithDomain {
   ok: Boolean
   msg: String

src/ai/backend/manager/models/gql.py

@@ -109,13 +109,19 @@
 )
 from .routing import Routing, RoutingList
 from .scaling_group import (
+    AssociateScalingGroupsWithDomain,
+    AssociateScalingGroupsWithKeyPair,
+    AssociateScalingGroupsWithUserGroup,
     AssociateScalingGroupWithDomain,
     AssociateScalingGroupWithKeyPair,
     AssociateScalingGroupWithUserGroup,
     CreateScalingGroup,
     DeleteScalingGroup,
     DisassociateAllScalingGroupsWithDomain,
     DisassociateAllScalingGroupsWithGroup,
+    DisassociateScalingGroupsWithDomain,
+    DisassociateScalingGroupsWithKeyPair,
+    DisassociateScalingGroupsWithUserGroup,
     DisassociateScalingGroupWithDomain,
     DisassociateScalingGroupWithKeyPair,
     DisassociateScalingGroupWithUserGroup,
@@ -241,11 +247,29 @@ class Mutations(graphene.ObjectType):
     modify_scaling_group = ModifyScalingGroup.Field()
     delete_scaling_group = DeleteScalingGroup.Field()
     associate_scaling_group_with_domain = AssociateScalingGroupWithDomain.Field()
+    associate_scaling_groups_with_domain = AssociateScalingGroupsWithDomain.Field(
+        description="Added in 24.03.9"
+    )
     associate_scaling_group_with_user_group = AssociateScalingGroupWithUserGroup.Field()
+    associate_scaling_groups_with_user_group = AssociateScalingGroupsWithUserGroup.Field(
+        description="Added in 24.03.9"
+    )
     associate_scaling_group_with_keypair = AssociateScalingGroupWithKeyPair.Field()
+    associate_scaling_groups_with_keypair = AssociateScalingGroupsWithKeyPair.Field(
+        description="Added in 24.03.9"
+    )
     disassociate_scaling_group_with_domain = DisassociateScalingGroupWithDomain.Field()
+    disassociate_scaling_groups_with_domain = DisassociateScalingGroupsWithDomain.Field(
+        description="Added in 24.03.9"
+    )
     disassociate_scaling_group_with_user_group = DisassociateScalingGroupWithUserGroup.Field()
+    disassociate_scaling_groups_with_user_group = DisassociateScalingGroupsWithUserGroup.Field(
+        description="Added in 24.03.9"
+    )
     disassociate_scaling_group_with_keypair = DisassociateScalingGroupWithKeyPair.Field()
+    disassociate_scaling_groups_with_keypair = DisassociateScalingGroupsWithKeyPair.Field(
+        description="Added in 24.03.9"
+    )
     disassociate_all_scaling_groups_with_domain = DisassociateAllScalingGroupsWithDomain.Field()
     disassociate_all_scaling_groups_with_group = DisassociateAllScalingGroupsWithGroup.Field()
 

src/ai/backend/manager/models/scaling_group.py

@@ -692,6 +692,32 @@ async def mutate(
         return await simple_db_mutate(cls, info.context, insert_query)
 
 
+class AssociateScalingGroupsWithDomain(graphene.Mutation):
+    """Added in 24.03.9."""
+
+    allowed_roles = (UserRole.SUPERADMIN,)
+
+    class Arguments:
+        scaling_groups = graphene.List(graphene.String, required=True)
+        domain = graphene.String(required=True)
+
+    ok = graphene.Boolean()
+    msg = graphene.String()
+
+    @classmethod
+    async def mutate(
+        cls,
+        root,
+        info: graphene.ResolveInfo,
+        scaling_groups: Sequence[str],
+        domain: str,
+    ) -> AssociateScalingGroupsWithDomain:
+        insert_query = sa.insert(sgroups_for_domains).values([
+            {"scaling_group": scaling_group, "domain": domain} for scaling_group in scaling_groups
+        ])
+        return await simple_db_mutate(cls, info.context, insert_query)
+
+
 class DisassociateScalingGroupWithDomain(graphene.Mutation):
     allowed_roles = (UserRole.SUPERADMIN,)
 
@@ -717,6 +743,33 @@ async def mutate(
         return await simple_db_mutate(cls, info.context, delete_query)
 
 
+class DisassociateScalingGroupsWithDomain(graphene.Mutation):
+    """Added in 24.03.9."""
+
+    allowed_roles = (UserRole.SUPERADMIN,)
+
+    class Arguments:
+        scaling_groups = graphene.List(graphene.String, required=True)
+        domain = graphene.String(required=True)
+
+    ok = graphene.Boolean()
+    msg = graphene.String()
+
+    @classmethod
+    async def mutate(
+        cls,
+        root,
+        info: graphene.ResolveInfo,
+        scaling_groups: Sequence[str],
+        domain: str,
+    ) -> DisassociateScalingGroupsWithDomain:
+        delete_query = sa.delete(sgroups_for_domains).where(
+            (sgroups_for_domains.c.scaling_group.in_(scaling_groups))
+            & (sgroups_for_domains.c.domain == domain),
+        )
+        return await simple_db_mutate(cls, info.context, delete_query)
+
+
 class DisassociateAllScalingGroupsWithDomain(graphene.Mutation):
     allowed_roles = (UserRole.SUPERADMIN,)
 
@@ -762,6 +815,33 @@ async def mutate(
         return await simple_db_mutate(cls, info.context, insert_query)
 
 
+class AssociateScalingGroupsWithUserGroup(graphene.Mutation):
+    """Added in 24.03.9."""
+
+    allowed_roles = (UserRole.SUPERADMIN,)
+
+    class Arguments:
+        scaling_groups = graphene.List(graphene.String, required=True)
+        user_group = graphene.UUID(required=True)
+
+    ok = graphene.Boolean()
+    msg = graphene.String()
+
+    @classmethod
+    async def mutate(
+        cls,
+        root,
+        info: graphene.ResolveInfo,
+        scaling_groups: Sequence[str],
+        user_group: uuid.UUID,
+    ) -> AssociateScalingGroupsWithUserGroup:
+        insert_query = sa.insert(sgroups_for_groups).values([
+            {"scaling_group": scaling_group, "group": user_group}
+            for scaling_group in scaling_groups
+        ])
+        return await simple_db_mutate(cls, info.context, insert_query)
+
+
 class DisassociateScalingGroupWithUserGroup(graphene.Mutation):
     allowed_roles = (UserRole.SUPERADMIN,)
 
@@ -787,6 +867,33 @@ async def mutate(
         return await simple_db_mutate(cls, info.context, delete_query)
 
 
+class DisassociateScalingGroupsWithUserGroup(graphene.Mutation):
+    """Added in 24.03.9."""
+
+    allowed_roles = (UserRole.SUPERADMIN,)
+
+    class Arguments:
+        scaling_groups = graphene.List(graphene.String, required=True)
+        user_group = graphene.UUID(required=True)
+
+    ok = graphene.Boolean()
+    msg = graphene.String()
+
+    @classmethod
+    async def mutate(
+        cls,
+        root,
+        info: graphene.ResolveInfo,
+        scaling_groups: Sequence[str],
+        user_group: uuid.UUID,
+    ) -> DisassociateScalingGroupsWithUserGroup:
+        delete_query = sa.delete(sgroups_for_groups).where(
+            (sgroups_for_groups.c.scaling_group.in_(scaling_groups))
+            & (sgroups_for_groups.c.group == user_group),
+        )
+        return await simple_db_mutate(cls, info.context, delete_query)
+
+
 class DisassociateAllScalingGroupsWithGroup(graphene.Mutation):
     allowed_roles = (UserRole.SUPERADMIN,)
 
@@ -832,6 +939,33 @@ async def mutate(
         return await simple_db_mutate(cls, info.context, insert_query)
 
 
+class AssociateScalingGroupsWithKeyPair(graphene.Mutation):
+    """Added in 24.03.9."""
+
+    allowed_roles = (UserRole.SUPERADMIN,)
+
+    class Arguments:
+        scaling_groups = graphene.List(graphene.String, required=True)
+        access_key = graphene.String(required=True)
+
+    ok = graphene.Boolean()
+    msg = graphene.String()
+
+    @classmethod
+    async def mutate(
+        cls,
+        root,
+        info: graphene.ResolveInfo,
+        scaling_groups: Sequence[str],
+        access_key: str,
+    ) -> AssociateScalingGroupsWithKeyPair:
+        insert_query = sa.insert(sgroups_for_keypairs).values([
+            {"scaling_group": scaling_group, "access_key": access_key}
+            for scaling_group in scaling_groups
+        ])
+        return await simple_db_mutate(cls, info.context, insert_query)
+
+
 class DisassociateScalingGroupWithKeyPair(graphene.Mutation):
     allowed_roles = (UserRole.SUPERADMIN,)
 
@@ -855,3 +989,30 @@ async def mutate(
             & (sgroups_for_keypairs.c.access_key == access_key),
         )
         return await simple_db_mutate(cls, info.context, delete_query)
+
+
+class DisassociateScalingGroupsWithKeyPair(graphene.Mutation):
+    """Added in 24.03.9."""
+
+    allowed_roles = (UserRole.SUPERADMIN,)
+
+    class Arguments:
+        scaling_groups = graphene.List(graphene.String, required=True)
+        access_key = graphene.String(required=True)
+
+    ok = graphene.Boolean()
+    msg = graphene.String()
+
+    @classmethod
+    async def mutate(
+        cls,
+        root,
+        info: graphene.ResolveInfo,
+        scaling_groups: Sequence[str],
+        access_key: str,
+    ) -> DisassociateScalingGroupsWithKeyPair:
+        delete_query = sa.delete(sgroups_for_keypairs).where(
+            (sgroups_for_keypairs.c.scaling_group.in_(scaling_groups))
+            & (sgroups_for_keypairs.c.access_key == access_key),
+        )
+        return await simple_db_mutate(cls, info.context, delete_query)