feat(BA-3490): Implement ScalingGroup Keypair Association Actions (#7655)

Author: jopemachine

changes/7655.feature.md

@@ -0,0 +1 @@
+Implement ScalingGroup Keypair Association Actions

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

@@ -4,8 +4,10 @@
 from dataclasses import dataclass, field
 from typing import Any, Optional, override
 
+from ai.backend.common.types import AccessKey
 from ai.backend.manager.models.scaling_group import (
     ScalingGroupForDomainRow,
+    ScalingGroupForKeypairsRow,
     ScalingGroupOpts,
     ScalingGroupRow,
 )
@@ -58,3 +60,18 @@ def build_row(self) -> ScalingGroupForDomainRow:
             scaling_group=self.scaling_group,
             domain=self.domain,
         )
+
+
+@dataclass
+class ScalingGroupForKeypairsCreatorSpec(CreatorSpec[ScalingGroupForKeypairsRow]):
+    """CreatorSpec for associating a scaling group with a keypair."""
+
+    scaling_group: str
+    access_key: AccessKey
+
+    @override
+    def build_row(self) -> ScalingGroupForKeypairsRow:
+        return ScalingGroupForKeypairsRow(
+            scaling_group=self.scaling_group,
+            access_key=self.access_key,
+        )

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

@@ -11,7 +11,11 @@
 from ai.backend.manager.errors.resource import ScalingGroupNotFound
 from ai.backend.manager.models.endpoint import EndpointRow
 from ai.backend.manager.models.routing import RoutingRow
-from ai.backend.manager.models.scaling_group import ScalingGroupForDomainRow, ScalingGroupRow
+from ai.backend.manager.models.scaling_group import (
+    ScalingGroupForDomainRow,
+    ScalingGroupForKeypairsRow,
+    ScalingGroupRow,
+)
 from ai.backend.manager.models.session import SessionRow
 from ai.backend.manager.repositories.base import BatchQuerier, execute_batch_querier
 from ai.backend.manager.repositories.base.creator import (
@@ -189,3 +193,37 @@ async def check_scaling_group_domain_association_exists(
             )
             result = await session.scalar(query)
             return (result or 0) > 0
+
+    async def associate_scaling_group_with_keypairs(
+        self,
+        bulk_creator: BulkCreator[ScalingGroupForKeypairsRow],
+    ) -> None:
+        """Associates a scaling group with multiple keypairs."""
+        async with self._db.begin_session() as session:
+            await execute_bulk_creator(session, bulk_creator)
+
+    async def disassociate_scaling_group_with_keypairs(
+        self,
+        purger: BatchPurger[ScalingGroupForKeypairsRow],
+    ) -> None:
+        """Disassociates a scaling group from multiple keypairs."""
+        async with self._db.begin_session() as session:
+            await execute_batch_purger(session, purger)
+
+    async def check_scaling_group_keypair_association_exists(
+        self,
+        scaling_group_name: str,
+        access_key: str,
+    ) -> bool:
+        """Checks if a scaling group is associated with a keypair."""
+        async with self._db.begin_readonly_session() as session:
+            query = sa.select(
+                sa.exists().where(
+                    sa.and_(
+                        ScalingGroupForKeypairsRow.scaling_group == scaling_group_name,
+                        ScalingGroupForKeypairsRow.access_key == access_key,
+                    )
+                )
+            )
+            result = await session.execute(query)
+            return result.scalar() or False

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

@@ -5,7 +5,11 @@
 
 import sqlalchemy as sa
 
-from ai.backend.manager.models.scaling_group import ScalingGroupForDomainRow
+from ai.backend.common.types import AccessKey
+from ai.backend.manager.models.scaling_group import (
+    ScalingGroupForDomainRow,
+    ScalingGroupForKeypairsRow,
+)
 from ai.backend.manager.repositories.base.purger import BatchPurger, BatchPurgerSpec
 
 
@@ -26,6 +30,23 @@ def build_subquery(self) -> sa.sql.Select[tuple[ScalingGroupForDomainRow]]:
         )
 
 
+@dataclass
+class ScalingGroupForKeypairsPurgerSpec(BatchPurgerSpec[ScalingGroupForKeypairsRow]):
+    """PurgerSpec for disassociating a scaling group from a keypair."""
+
+    scaling_group: str
+    access_key: AccessKey
+
+    @override
+    def build_subquery(self) -> sa.sql.Select[tuple[ScalingGroupForKeypairsRow]]:
+        return sa.select(ScalingGroupForKeypairsRow).where(
+            sa.and_(
+                ScalingGroupForKeypairsRow.scaling_group == self.scaling_group,
+                ScalingGroupForKeypairsRow.access_key == self.access_key,
+            )
+        )
+
+
 def create_scaling_group_for_domain_purger(
     scaling_group: str,
     domain: str,
@@ -38,3 +59,17 @@ def create_scaling_group_for_domain_purger(
         ),
         batch_size=1,  # We expect only one row to be deleted
     )
+
+
+def create_scaling_group_for_keypairs_purger(
+    scaling_group: str,
+    access_key: AccessKey,
+) -> BatchPurger[ScalingGroupForKeypairsRow]:
+    """Create a BatchPurger for disassociating a scaling group from a keypair."""
+    return BatchPurger(
+        spec=ScalingGroupForKeypairsPurgerSpec(
+            scaling_group=scaling_group,
+            access_key=access_key,
+        ),
+        batch_size=1,  # We expect only one row to be deleted
+    )

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

@@ -12,7 +12,11 @@
 )
 from ai.backend.common.resilience.policies.retry import BackoffStrategy
 from ai.backend.manager.data.scaling_group.types import ScalingGroupData, ScalingGroupListResult
-from ai.backend.manager.models.scaling_group import ScalingGroupForDomainRow, ScalingGroupRow
+from ai.backend.manager.models.scaling_group import (
+    ScalingGroupForDomainRow,
+    ScalingGroupForKeypairsRow,
+    ScalingGroupRow,
+)
 from ai.backend.manager.repositories.base import BatchQuerier
 from ai.backend.manager.repositories.base.creator import BulkCreator, Creator
 from ai.backend.manager.repositories.base.purger import BatchPurger, Purger
@@ -114,3 +118,27 @@ async def check_scaling_group_domain_association_exists(
             scaling_group=scaling_group,
             domain=domain,
         )
+
+    async def associate_scaling_group_with_keypairs(
+        self,
+        bulk_creator: BulkCreator[ScalingGroupForKeypairsRow],
+    ) -> None:
+        """Associates a scaling group with multiple keypairs."""
+        await self._db_source.associate_scaling_group_with_keypairs(bulk_creator)
+
+    async def disassociate_scaling_group_with_keypairs(
+        self,
+        purger: BatchPurger[ScalingGroupForKeypairsRow],
+    ) -> None:
+        """Disassociates a scaling group from multiple keypairs."""
+        await self._db_source.disassociate_scaling_group_with_keypairs(purger)
+
+    async def check_scaling_group_keypair_association_exists(
+        self,
+        scaling_group_name: str,
+        access_key: str,
+    ) -> bool:
+        """Checks if a scaling group is associated with a keypair."""
+        return await self._db_source.check_scaling_group_keypair_association_exists(
+            scaling_group_name, access_key
+        )

src/ai/backend/manager/services/scaling_group/actions/associate_with_keypair.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 ScalingGroupForKeypairsRow
+from ai.backend.manager.repositories.base.creator import BulkCreator
+
+from .base import ScalingGroupAction
+
+
+@dataclass
+class AssociateScalingGroupWithKeypairsAction(ScalingGroupAction):
+    """Action to associate a scaling group with multiple keypairs."""
+
+    bulk_creator: BulkCreator[ScalingGroupForKeypairsRow]
+
+    @override
+    @classmethod
+    def operation_type(cls) -> str:
+        return "associate_with_keypairs"
+
+    @override
+    def entity_id(self) -> Optional[str]:
+        return None
+
+
+@dataclass
+class AssociateScalingGroupWithKeypairsActionResult(BaseActionResult):
+    """Result of associating a scaling group with keypairs."""
+
+    @override
+    def entity_id(self) -> Optional[str]:
+        return None

src/ai/backend/manager/services/scaling_group/actions/disassociate_with_keypair.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 ScalingGroupForKeypairsRow
+from ai.backend.manager.repositories.base.purger import BatchPurger
+
+from .base import ScalingGroupAction
+
+
+@dataclass
+class DisassociateScalingGroupWithKeypairsAction(ScalingGroupAction):
+    """Action to disassociate a scaling group from multiple keypairs."""
+
+    purger: BatchPurger[ScalingGroupForKeypairsRow]
+
+    @override
+    @classmethod
+    def operation_type(cls) -> str:
+        return "disassociate_with_keypairs"
+
+    @override
+    def entity_id(self) -> Optional[str]:
+        return None
+
+
+@dataclass
+class DisassociateScalingGroupWithKeypairsActionResult(BaseActionResult):
+    """Result of disassociating a scaling group from keypairs."""
+
+    @override
+    def entity_id(self) -> Optional[str]:
+        return None

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

@@ -7,6 +7,10 @@
     AssociateScalingGroupWithDomainsAction,
     AssociateScalingGroupWithDomainsActionResult,
 )
+from ai.backend.manager.services.scaling_group.actions.associate_with_keypair import (
+    AssociateScalingGroupWithKeypairsAction,
+    AssociateScalingGroupWithKeypairsActionResult,
+)
 from ai.backend.manager.services.scaling_group.actions.create import (
     CreateScalingGroupAction,
     CreateScalingGroupActionResult,
@@ -15,6 +19,10 @@
     DisassociateScalingGroupWithDomainsAction,
     DisassociateScalingGroupWithDomainsActionResult,
 )
+from ai.backend.manager.services.scaling_group.actions.disassociate_with_keypair import (
+    DisassociateScalingGroupWithKeypairsAction,
+    DisassociateScalingGroupWithKeypairsActionResult,
+)
 from ai.backend.manager.services.scaling_group.actions.list_scaling_groups import (
     SearchScalingGroupsAction,
     SearchScalingGroupsActionResult,
@@ -43,6 +51,12 @@ class ScalingGroupProcessors(AbstractProcessorPackage):
     disassociate_scaling_group_with_domains: ActionProcessor[
         DisassociateScalingGroupWithDomainsAction, DisassociateScalingGroupWithDomainsActionResult
     ]
+    associate_scaling_group_with_keypairs: ActionProcessor[
+        AssociateScalingGroupWithKeypairsAction, AssociateScalingGroupWithKeypairsActionResult
+    ]
+    disassociate_scaling_group_with_keypairs: ActionProcessor[
+        DisassociateScalingGroupWithKeypairsAction, DisassociateScalingGroupWithKeypairsActionResult
+    ]
 
     def __init__(self, service: ScalingGroupService, action_monitors: list[ActionMonitor]) -> None:
         self.create_scaling_group = ActionProcessor(service.create_scaling_group, action_monitors)
@@ -55,6 +69,12 @@ def __init__(self, service: ScalingGroupService, action_monitors: list[ActionMon
         self.disassociate_scaling_group_with_domains = ActionProcessor(
             service.disassociate_scaling_group_with_domains, action_monitors
         )
+        self.associate_scaling_group_with_keypairs = ActionProcessor(
+            service.associate_scaling_group_with_keypairs, action_monitors
+        )
+        self.disassociate_scaling_group_with_keypairs = ActionProcessor(
+            service.disassociate_scaling_group_with_keypairs, action_monitors
+        )
 
     @override
     def supported_actions(self) -> list[ActionSpec]:
@@ -65,4 +85,6 @@ def supported_actions(self) -> list[ActionSpec]:
             SearchScalingGroupsAction.spec(),
             AssociateScalingGroupWithDomainsAction.spec(),
             DisassociateScalingGroupWithDomainsAction.spec(),
+            AssociateScalingGroupWithKeypairsAction.spec(),
+            DisassociateScalingGroupWithKeypairsAction.spec(),
         ]

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

@@ -6,6 +6,10 @@
     AssociateScalingGroupWithDomainsAction,
     AssociateScalingGroupWithDomainsActionResult,
 )
+from ai.backend.manager.services.scaling_group.actions.associate_with_keypair import (
+    AssociateScalingGroupWithKeypairsAction,
+    AssociateScalingGroupWithKeypairsActionResult,
+)
 from ai.backend.manager.services.scaling_group.actions.create import (
     CreateScalingGroupAction,
     CreateScalingGroupActionResult,
@@ -14,6 +18,10 @@
     DisassociateScalingGroupWithDomainsAction,
     DisassociateScalingGroupWithDomainsActionResult,
 )
+from ai.backend.manager.services.scaling_group.actions.disassociate_with_keypair import (
+    DisassociateScalingGroupWithKeypairsAction,
+    DisassociateScalingGroupWithKeypairsActionResult,
+)
 from ai.backend.manager.services.scaling_group.actions.list_scaling_groups import (
     SearchScalingGroupsAction,
     SearchScalingGroupsActionResult,
@@ -85,3 +93,17 @@ async def disassociate_scaling_group_with_domains(
         """Disassociates a scaling group from multiple domains."""
         await self._repository.disassociate_scaling_group_with_domains(action.purger)
         return DisassociateScalingGroupWithDomainsActionResult()
+
+    async def associate_scaling_group_with_keypairs(
+        self, action: AssociateScalingGroupWithKeypairsAction
+    ) -> AssociateScalingGroupWithKeypairsActionResult:
+        """Associates a scaling group with multiple keypairs."""
+        await self._repository.associate_scaling_group_with_keypairs(action.bulk_creator)
+        return AssociateScalingGroupWithKeypairsActionResult()
+
+    async def disassociate_scaling_group_with_keypairs(
+        self, action: DisassociateScalingGroupWithKeypairsAction
+    ) -> DisassociateScalingGroupWithKeypairsActionResult:
+        """Disassociates a scaling group from multiple keypairs."""
+        await self._repository.disassociate_scaling_group_with_keypairs(action.purger)
+        return DisassociateScalingGroupWithKeypairsActionResult()