feat(BA-3489): Implement ScalingGroup User Group Association Actions (#7654)

Author: jopemachine

changes/7654.feature.md

@@ -0,0 +1 @@
+Implement ScalingGroup User Group Association, Disassociation Actions

src/ai/backend/manager/repositories/scaling_group/creators.py

@@ -3,11 +3,13 @@
 from collections.abc import Mapping
 from dataclasses import dataclass, field
 from typing import Any, Optional, override
+from uuid import UUID
 
 from ai.backend.common.types import AccessKey
 from ai.backend.manager.models.scaling_group import (
     ScalingGroupForDomainRow,
     ScalingGroupForKeypairsRow,
+    ScalingGroupForProjectRow,
     ScalingGroupOpts,
     ScalingGroupRow,
 )
@@ -75,3 +77,18 @@ def build_row(self) -> ScalingGroupForKeypairsRow:
             scaling_group=self.scaling_group,
             access_key=self.access_key,
         )
+
+
+@dataclass
+class ScalingGroupForProjectCreatorSpec(CreatorSpec[ScalingGroupForProjectRow]):
+    """CreatorSpec for associating a scaling group with a project (user group)."""
+
+    scaling_group: str
+    project: UUID
+
+    @override
+    def build_row(self) -> ScalingGroupForProjectRow:
+        return ScalingGroupForProjectRow(
+            scaling_group=self.scaling_group,
+            group=self.project,
+        )

src/ai/backend/manager/repositories/scaling_group/db_source/db_source.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import uuid
 from typing import TYPE_CHECKING, cast
 
 import sqlalchemy as sa
@@ -14,6 +15,7 @@
 from ai.backend.manager.models.scaling_group import (
     ScalingGroupForDomainRow,
     ScalingGroupForKeypairsRow,
+    ScalingGroupForProjectRow,
     ScalingGroupRow,
 )
 from ai.backend.manager.models.session import SessionRow
@@ -227,3 +229,39 @@ async def check_scaling_group_keypair_association_exists(
             )
             result = await session.execute(query)
             return result.scalar() or False
+
+    async def associate_scaling_group_with_user_groups(
+        self,
+        bulk_creator: BulkCreator[ScalingGroupForProjectRow],
+    ) -> None:
+        """Associates a scaling group with multiple user groups (projects)."""
+        async with self._db.begin_session() as session:
+            await execute_bulk_creator(session, bulk_creator)
+
+    async def disassociate_scaling_group_with_user_groups(
+        self,
+        purger: BatchPurger[ScalingGroupForProjectRow],
+    ) -> None:
+        """Disassociates a single scaling group from a user group (project)."""
+        async with self._db.begin_session() as session:
+            await execute_batch_purger(session, purger)
+
+    async def check_scaling_group_user_group_association_exists(
+        self,
+        scaling_group: str,
+        user_group: uuid.UUID,
+    ) -> bool:
+        """Checks if a scaling group is associated with a user group (project)."""
+        async with self._db.begin_readonly_session() as session:
+            query = (
+                sa.select(sa.func.count())
+                .select_from(ScalingGroupForProjectRow)
+                .where(
+                    sa.and_(
+                        ScalingGroupForProjectRow.scaling_group == scaling_group,
+                        ScalingGroupForProjectRow.group == user_group,
+                    )
+                )
+            )
+            result = await session.scalar(query)
+            return (result or 0) > 0

src/ai/backend/manager/repositories/scaling_group/purgers.py

@@ -2,13 +2,15 @@
 
 from dataclasses import dataclass
 from typing import override
+from uuid import UUID
 
 import sqlalchemy as sa
 
 from ai.backend.common.types import AccessKey
 from ai.backend.manager.models.scaling_group import (
     ScalingGroupForDomainRow,
     ScalingGroupForKeypairsRow,
+    ScalingGroupForProjectRow,
 )
 from ai.backend.manager.repositories.base.purger import BatchPurger, BatchPurgerSpec
 
@@ -73,3 +75,34 @@ def create_scaling_group_for_keypairs_purger(
         ),
         batch_size=1,  # We expect only one row to be deleted
     )
+
+
+@dataclass
+class ScalingGroupForProjectPurgerSpec(BatchPurgerSpec[ScalingGroupForProjectRow]):
+    """PurgerSpec for disassociating a scaling group from a project (user group)."""
+
+    scaling_group: str
+    project: UUID
+
+    @override
+    def build_subquery(self) -> sa.sql.Select[tuple[ScalingGroupForProjectRow]]:
+        return sa.select(ScalingGroupForProjectRow).where(
+            sa.and_(
+                ScalingGroupForProjectRow.scaling_group == self.scaling_group,
+                ScalingGroupForProjectRow.group == self.project,
+            )
+        )
+
+
+def create_scaling_group_for_project_purger(
+    scaling_group: str,
+    project: UUID,
+) -> BatchPurger[ScalingGroupForProjectRow]:
+    """Create a BatchPurger for disassociating a scaling group from a project."""
+    return BatchPurger(
+        spec=ScalingGroupForProjectPurgerSpec(
+            scaling_group=scaling_group,
+            project=project,
+        ),
+        batch_size=1,  # We expect only one row to be deleted
+    )

src/ai/backend/manager/repositories/scaling_group/repository.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from typing import TYPE_CHECKING
+from uuid import UUID
 
 from ai.backend.common.metrics.metric import DomainType, LayerType
 from ai.backend.common.resilience import (
@@ -15,6 +16,7 @@
 from ai.backend.manager.models.scaling_group import (
     ScalingGroupForDomainRow,
     ScalingGroupForKeypairsRow,
+    ScalingGroupForProjectRow,
     ScalingGroupRow,
 )
 from ai.backend.manager.repositories.base import BatchQuerier
@@ -143,3 +145,28 @@ async def check_scaling_group_keypair_association_exists(
         return await self._db_source.check_scaling_group_keypair_association_exists(
             scaling_group_name, access_key
         )
+
+    async def associate_scaling_group_with_user_groups(
+        self,
+        bulk_creator: BulkCreator[ScalingGroupForProjectRow],
+    ) -> None:
+        """Associates a scaling group with multiple user groups (projects)."""
+        await self._db_source.associate_scaling_group_with_user_groups(bulk_creator)
+
+    async def disassociate_scaling_group_with_user_groups(
+        self,
+        purger: BatchPurger[ScalingGroupForProjectRow],
+    ) -> None:
+        """Disassociates a single scaling group from a user group (project)."""
+        await self._db_source.disassociate_scaling_group_with_user_groups(purger)
+
+    async def check_scaling_group_user_group_association_exists(
+        self,
+        scaling_group: str,
+        user_group: UUID,
+    ) -> bool:
+        """Checks if a scaling group is associated with a user group (project)."""
+        return await self._db_source.check_scaling_group_user_group_association_exists(
+            scaling_group=scaling_group,
+            user_group=user_group,
+        )

src/ai/backend/manager/services/scaling_group/actions/associate_with_user_group.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional, override
+
+from ai.backend.manager.actions.action import BaseActionResult
+from ai.backend.manager.models.scaling_group import ScalingGroupForProjectRow
+from ai.backend.manager.repositories.base.creator import BulkCreator
+
+from .base import ScalingGroupAction
+
+
+@dataclass
+class AssociateScalingGroupWithUserGroupsAction(ScalingGroupAction):
+    """Action to associate a scaling group with multiple user groups (projects)."""
+
+    bulk_creator: BulkCreator[ScalingGroupForProjectRow]
+
+    @override
+    @classmethod
+    def operation_type(cls) -> str:
+        return "associate_with_user_groups"
+
+    @override
+    def entity_id(self) -> Optional[str]:
+        return None
+
+
+@dataclass
+class AssociateScalingGroupWithUserGroupsActionResult(BaseActionResult):
+    """Result of associating a scaling group with user groups."""
+
+    @override
+    def entity_id(self) -> Optional[str]:
+        return None

src/ai/backend/manager/services/scaling_group/actions/disassociate_with_user_group.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional, override
+
+from ai.backend.manager.actions.action import BaseActionResult
+from ai.backend.manager.models.scaling_group import ScalingGroupForProjectRow
+from ai.backend.manager.repositories.base.purger import BatchPurger
+
+from .base import ScalingGroupAction
+
+
+@dataclass
+class DisassociateScalingGroupWithUserGroupsAction(ScalingGroupAction):
+    """Action to disassociate a single scaling group from a user group (project)."""
+
+    purger: BatchPurger[ScalingGroupForProjectRow]
+
+    @override
+    @classmethod
+    def operation_type(cls) -> str:
+        return "disassociate_with_user_groups"
+
+    @override
+    def entity_id(self) -> Optional[str]:
+        return None
+
+
+@dataclass
+class DisassociateScalingGroupWithUserGroupsActionResult(BaseActionResult):
+    """Result of disassociating a scaling group from a user group."""
+
+    @override
+    def entity_id(self) -> Optional[str]:
+        return None

src/ai/backend/manager/services/scaling_group/processors.py

@@ -11,6 +11,10 @@
     AssociateScalingGroupWithKeypairsAction,
     AssociateScalingGroupWithKeypairsActionResult,
 )
+from ai.backend.manager.services.scaling_group.actions.associate_with_user_group import (
+    AssociateScalingGroupWithUserGroupsAction,
+    AssociateScalingGroupWithUserGroupsActionResult,
+)
 from ai.backend.manager.services.scaling_group.actions.create import (
     CreateScalingGroupAction,
     CreateScalingGroupActionResult,
@@ -23,6 +27,10 @@
     DisassociateScalingGroupWithKeypairsAction,
     DisassociateScalingGroupWithKeypairsActionResult,
 )
+from ai.backend.manager.services.scaling_group.actions.disassociate_with_user_group import (
+    DisassociateScalingGroupWithUserGroupsAction,
+    DisassociateScalingGroupWithUserGroupsActionResult,
+)
 from ai.backend.manager.services.scaling_group.actions.list_scaling_groups import (
     SearchScalingGroupsAction,
     SearchScalingGroupsActionResult,
@@ -57,6 +65,13 @@ class ScalingGroupProcessors(AbstractProcessorPackage):
     disassociate_scaling_group_with_keypairs: ActionProcessor[
         DisassociateScalingGroupWithKeypairsAction, DisassociateScalingGroupWithKeypairsActionResult
     ]
+    associate_scaling_group_with_user_groups: ActionProcessor[
+        AssociateScalingGroupWithUserGroupsAction, AssociateScalingGroupWithUserGroupsActionResult
+    ]
+    disassociate_scaling_group_with_user_groups: ActionProcessor[
+        DisassociateScalingGroupWithUserGroupsAction,
+        DisassociateScalingGroupWithUserGroupsActionResult,
+    ]
 
     def __init__(self, service: ScalingGroupService, action_monitors: list[ActionMonitor]) -> None:
         self.create_scaling_group = ActionProcessor(service.create_scaling_group, action_monitors)
@@ -75,6 +90,12 @@ def __init__(self, service: ScalingGroupService, action_monitors: list[ActionMon
         self.disassociate_scaling_group_with_keypairs = ActionProcessor(
             service.disassociate_scaling_group_with_keypairs, action_monitors
         )
+        self.associate_scaling_group_with_user_groups = ActionProcessor(
+            service.associate_scaling_group_with_user_groups, action_monitors
+        )
+        self.disassociate_scaling_group_with_user_groups = ActionProcessor(
+            service.disassociate_scaling_group_with_user_groups, action_monitors
+        )
 
     @override
     def supported_actions(self) -> list[ActionSpec]:
@@ -87,4 +108,6 @@ def supported_actions(self) -> list[ActionSpec]:
             DisassociateScalingGroupWithDomainsAction.spec(),
             AssociateScalingGroupWithKeypairsAction.spec(),
             DisassociateScalingGroupWithKeypairsAction.spec(),
+            AssociateScalingGroupWithUserGroupsAction.spec(),
+            DisassociateScalingGroupWithUserGroupsAction.spec(),
         ]

src/ai/backend/manager/services/scaling_group/service.py

@@ -10,6 +10,10 @@
     AssociateScalingGroupWithKeypairsAction,
     AssociateScalingGroupWithKeypairsActionResult,
 )
+from ai.backend.manager.services.scaling_group.actions.associate_with_user_group import (
+    AssociateScalingGroupWithUserGroupsAction,
+    AssociateScalingGroupWithUserGroupsActionResult,
+)
 from ai.backend.manager.services.scaling_group.actions.create import (
     CreateScalingGroupAction,
     CreateScalingGroupActionResult,
@@ -22,6 +26,10 @@
     DisassociateScalingGroupWithKeypairsAction,
     DisassociateScalingGroupWithKeypairsActionResult,
 )
+from ai.backend.manager.services.scaling_group.actions.disassociate_with_user_group import (
+    DisassociateScalingGroupWithUserGroupsAction,
+    DisassociateScalingGroupWithUserGroupsActionResult,
+)
 from ai.backend.manager.services.scaling_group.actions.list_scaling_groups import (
     SearchScalingGroupsAction,
     SearchScalingGroupsActionResult,
@@ -107,3 +115,17 @@ async def disassociate_scaling_group_with_keypairs(
         """Disassociates a scaling group from multiple keypairs."""
         await self._repository.disassociate_scaling_group_with_keypairs(action.purger)
         return DisassociateScalingGroupWithKeypairsActionResult()
+
+    async def associate_scaling_group_with_user_groups(
+        self, action: AssociateScalingGroupWithUserGroupsAction
+    ) -> AssociateScalingGroupWithUserGroupsActionResult:
+        """Associates a scaling group with multiple user groups (projects)."""
+        await self._repository.associate_scaling_group_with_user_groups(action.bulk_creator)
+        return AssociateScalingGroupWithUserGroupsActionResult()
+
+    async def disassociate_scaling_group_with_user_groups(
+        self, action: DisassociateScalingGroupWithUserGroupsAction
+    ) -> DisassociateScalingGroupWithUserGroupsActionResult:
+        """Disassociates a single scaling group from a user group (project)."""
+        await self._repository.disassociate_scaling_group_with_user_groups(action.purger)
+        return DisassociateScalingGroupWithUserGroupsActionResult()