feat: Add Associate and Disassociate Scaling Groups with Domain, User Group, and KeyPair mutations (#2473)

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

@@ -1163,27 +1163,45 @@
   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
   unset_quota_scope(quota_scope_id: String!, storage_host_name: String!): UnsetQuotaScope
   create_container_registry(hostname: String!, props: CreateContainerRegistryInput!): CreateContainerRegistry
   modify_container_registry(hostname: String!, props: ModifyContainerRegistryInput!): ModifyContainerRegistry
   delete_container_registry(hostname: String!): DeleteContainerRegistry
   modify_endpoint(endpoint_id: UUID!, props: ModifyEndpointInput!): ModifyEndpoint
 }
 
 type ModifyAgent {
   ok: Boolean
   msg: String
 }
 
 input ModifyAgentInput {
   schedulable: Boolean
   scaling_group: String
 }
@@ -1731,32 +1749,68 @@
   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
 }
@@ -1767,7 +1821,7 @@
 }
 
 type SetQuotaScope {
   quota_scope: QuotaScope
 }
 
 input QuotaScopeInput {

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

@@ -627,6 +627,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,)
 
@@ -652,6 +678,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,)
 
@@ -697,6 +750,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,)
 
@@ -722,6 +802,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,)
 
@@ -767,6 +874,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,)
 
@@ -790,3 +924,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)