feat(BA-6827): add agent resource-group update service

changes/12917.feature.md

@@ -0,0 +1 @@
+Add a service to change an agent's resource group that also handles the sessions still running on it: it fails unless forced (so an admin can drain first), and on force it terminates those sessions before applying the change.

src/ai/backend/manager/errors/agent.py

@@ -36,3 +36,48 @@ def error_code(self) -> ErrorCode:
             operation=ErrorOperation.ACCESS,
             error_detail=ErrorDetail.UNAVAILABLE,
         )
+
+
+class AgentHasConflictingSessions(BackendAIError, web.HTTPConflict):
+    """
+    Raised when an agent has sessions conflicting with its resource group and
+    the caller did not request forced cleanup (the admin must drain first).
+    """
+
+    error_type = "https://api.backend.ai/probs/agent-has-conflicting-sessions"
+    error_title = "Agent has sessions conflicting with its resource group."
+
+    def __init__(self, agent_id: AgentId, conflicting_count: int) -> None:
+        self.agent_id = agent_id
+        self.conflicting_count = conflicting_count
+        super().__init__(
+            f"Agent {agent_id} has {conflicting_count} session(s) conflicting with its "
+            "resource group; drain them or retry with force."
+        )
+
+    @override
+    def error_code(self) -> ErrorCode:
+        return ErrorCode(
+            domain=ErrorDomain.AGENT,
+            operation=ErrorOperation.UPDATE,
+            error_detail=ErrorDetail.CONFLICT,
+        )
+
+
+class ConflictingSessionRescheduleNotSupported(BackendAIError, web.HTTPNotImplemented):
+    """
+    Raised when the RESCHEDULE cleanup policy is requested. Re-enqueueing
+    conflicting sessions depends on the RESCHEDULING state from the preemption
+    work, which is not yet available; only TERMINATE is implemented.
+    """
+
+    error_type = "https://api.backend.ai/probs/conflicting-session-reschedule-not-supported"
+    error_title = "Rescheduling conflicting sessions is not supported yet."
+
+    @override
+    def error_code(self) -> ErrorCode:
+        return ErrorCode(
+            domain=ErrorDomain.AGENT,
+            operation=ErrorOperation.UPDATE,
+            error_detail=ErrorDetail.NOT_IMPLEMENTED,
+        )

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

@@ -8,6 +8,7 @@
 from sqlalchemy.orm import selectinload
 
 from ai.backend.common.exception import AgentNotFound
+from ai.backend.common.identifier.resource_group import ResourceGroupID
 from ai.backend.common.types import AgentId, ImageID
 from ai.backend.logging.utils import BraceStyleAdapter
 from ai.backend.manager.data.agent.types import (
@@ -19,10 +20,13 @@
     UpsertResult,
 )
 from ai.backend.manager.data.image.types import ImageDataWithDetails, ImageIdentifier
+from ai.backend.manager.data.kernel.types import KernelInfo, KernelStatus
+from ai.backend.manager.errors.agent import AgentHasConflictingSessions
 from ai.backend.manager.errors.resource import UnresolvableResourceGroup
 from ai.backend.manager.models.agent import ADMIN_PERMISSIONS as ADMIN_AGENT_PERMISSIONS
 from ai.backend.manager.models.agent import AgentRow, agents
 from ai.backend.manager.models.image import ImageRow
+from ai.backend.manager.models.kernel import KernelRow
 from ai.backend.manager.models.resource_slot import AgentResourceRow
 from ai.backend.manager.models.scaling_group import ScalingGroupRow
 from ai.backend.manager.models.utils import ExtendedAsyncSAEngine
@@ -181,6 +185,54 @@ async def update_agent_status(self, updater: Updater[AgentRow]) -> None:
         async with self._db.begin_session() as session:
             await execute_updater(session, updater)
 
+    async def update_resource_group(
+        self,
+        agent_id: AgentId,
+        resource_group_id: ResourceGroupID,
+        *,
+        force: bool,
+    ) -> list[KernelInfo]:
+        """
+        Change the agent's resource group, gating on the kernels running on it.
+
+        Finds the active kernels on the agent. If any exist and ``force`` is not
+        set, raises without changing anything. Otherwise updates the agent's group
+        (name + id columns) and returns those kernels so the caller can transition
+        their sessions. The check and the update run in one transaction.
+        """
+        active_statuses = (
+            KernelStatus.resource_occupied_statuses() | KernelStatus.resource_requested_statuses()
+        )
+        async with self._db.begin_session_read_committed() as session:
+            rows = (
+                (
+                    await session.execute(
+                        sa.select(KernelRow).where(
+                            KernelRow.agent == agent_id,
+                            KernelRow.status.in_(active_statuses),
+                        )
+                    )
+                )
+                .scalars()
+                .all()
+            )
+            kernels = [row.to_kernel_info() for row in rows]
+            if kernels and not force:
+                distinct_sessions = len({kernel.session.session_id for kernel in kernels})
+                raise AgentHasConflictingSessions(agent_id, distinct_sessions)
+
+            await session.execute(
+                sa.update(agents)
+                .where(agents.c.id == agent_id)
+                .values(
+                    resource_group_id=resource_group_id,
+                    scaling_group=sa.select(ScalingGroupRow.name)
+                    .where(ScalingGroupRow.id == resource_group_id)
+                    .scalar_subquery(),
+                )
+            )
+        return kernels
+
     async def search_agents(
         self,
         querier: BatchQuerier,

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

@@ -9,6 +9,7 @@
 from ai.backend.common.clients.valkey_client.valkey_stat.client import ValkeyStatClient
 from ai.backend.common.data.image.types import ScannedImage
 from ai.backend.common.exception import BackendAIError
+from ai.backend.common.identifier.resource_group import ResourceGroupID
 from ai.backend.common.metrics.metric import DomainType, LayerType
 from ai.backend.common.resilience.policies.metrics import MetricArgs, MetricPolicy
 from ai.backend.common.resilience.policies.retry import BackoffStrategy, RetryArgs, RetryPolicy
@@ -23,6 +24,7 @@
     UpsertResult,
 )
 from ai.backend.manager.data.image.types import ImageDataWithDetails, ImageIdentifier
+from ai.backend.manager.data.kernel.types import KernelInfo
 from ai.backend.manager.models.agent import AgentRow
 from ai.backend.manager.models.clauses import QueryCondition, QueryOrder
 from ai.backend.manager.models.resource_slot import AgentResourceRow
@@ -159,6 +161,21 @@ async def update_agent_status(self, agent_id: AgentId, spec: AgentStatusUpdaterS
         updater = Updater[AgentRow](spec=spec, pk_value=agent_id)
         await self._db_source.update_agent_status(updater)
 
+    @agent_repository_resilience.apply()
+    async def update_resource_group(
+        self,
+        agent_id: AgentId,
+        resource_group_id: ResourceGroupID,
+        *,
+        force: bool,
+    ) -> list[KernelInfo]:
+        """Change the agent's resource group, returning the kernels still on it.
+
+        Raises AgentHasConflictingSessions if the agent has active kernels and
+        ``force`` is not set; otherwise updates the group and returns those kernels.
+        """
+        return await self._db_source.update_resource_group(agent_id, resource_group_id, force=force)
+
     # For compatibility with redis key made with image canonical strings
     # Use remove_agent_from_images instead of this if possible
     @agent_repository_resilience.apply()

src/ai/backend/manager/services/agent/actions/update_resource_group.py

@@ -0,0 +1,40 @@
+from dataclasses import dataclass
+from typing import override
+
+from ai.backend.common.identifier.resource_group import ResourceGroupID
+from ai.backend.common.types import AgentId, SessionId
+from ai.backend.manager.actions.action import BaseActionResult
+from ai.backend.manager.actions.types import ActionOperationType
+from ai.backend.manager.services.agent.actions.base import AgentAction
+from ai.backend.manager.services.agent.types import ConflictingSessionCleanupPolicy
+
+
+@dataclass
+class UpdateAgentResourceGroupAction(AgentAction):
+    agent_id: AgentId
+    # Target resource group id (already resolved by the caller).
+    resource_group_id: ResourceGroupID
+    # How to handle sessions still running on the agent under the old group.
+    policy: ConflictingSessionCleanupPolicy
+    force: bool
+
+    @override
+    def entity_id(self) -> str | None:
+        return str(self.agent_id)
+
+    @override
+    @classmethod
+    def operation_type(cls) -> ActionOperationType:
+        return ActionOperationType.UPDATE
+
+
+@dataclass
+class UpdateAgentResourceGroupActionResult(BaseActionResult):
+    agent_id: AgentId
+    resource_group_id: ResourceGroupID
+    conflicting_session_ids: list[SessionId]
+    terminating_session_ids: list[SessionId]
+
+    @override
+    def entity_id(self) -> str | None:
+        return str(self.agent_id)

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

@@ -49,6 +49,10 @@
     SyncAgentRegistryAction,
     SyncAgentRegistryActionResult,
 )
+from ai.backend.manager.services.agent.actions.update_resource_group import (
+    UpdateAgentResourceGroupAction,
+    UpdateAgentResourceGroupActionResult,
+)
 from ai.backend.manager.services.agent.actions.watcher_agent_restart import (
     WatcherAgentRestartAction,
     WatcherAgentRestartActionResult,
@@ -87,6 +91,9 @@ class AgentProcessors(AbstractProcessorPackage):
     load_container_counts: ActionProcessor[
         LoadContainerCountsAction, LoadContainerCountsActionResult
     ]
+    update_resource_group: ActionProcessor[
+        UpdateAgentResourceGroupAction, UpdateAgentResourceGroupActionResult
+    ]
 
     def __init__(
         self,
@@ -118,6 +125,7 @@ def __init__(
         )
         self.search_agents = ActionProcessor(service.search_agents, action_monitors)
         self.load_container_counts = ActionProcessor(service.load_container_counts, action_monitors)
+        self.update_resource_group = ActionProcessor(service.update_resource_group, action_monitors)
 
     @override
     def supported_actions(self) -> list[ActionSpec]:
@@ -134,4 +142,5 @@ def supported_actions(self) -> list[ActionSpec]:
             RemoveAgentFromImagesByCanonicalsAction.spec(),
             SearchAgentsAction.spec(),
             LoadContainerCountsAction.spec(),
+            UpdateAgentResourceGroupAction.spec(),
         ]

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

@@ -1,6 +1,7 @@
 import logging
 from datetime import datetime
 from typing import Any, Literal, cast
+from uuid import UUID
 
 import aiohttp
 import yarl
@@ -20,6 +21,7 @@
 from ai.backend.common.plugin.hook import HookPluginContext
 from ai.backend.common.types import (
     AgentId,
+    SessionId,
 )
 from ai.backend.logging.utils import BraceStyleAdapter
 from ai.backend.manager.agent_cache import AgentRPCCache
@@ -28,6 +30,7 @@
     AgentHeartbeatUpsert,
     UpsertResult,
 )
+from ai.backend.manager.errors.agent import ConflictingSessionRescheduleNotSupported
 from ai.backend.manager.registry import AgentRegistry
 from ai.backend.manager.repositories.agent.repository import AgentRepository
 from ai.backend.manager.repositories.agent.updaters import AgentStatusUpdaterSpec
@@ -76,6 +79,10 @@
     SyncAgentRegistryAction,
     SyncAgentRegistryActionResult,
 )
+from ai.backend.manager.services.agent.actions.update_resource_group import (
+    UpdateAgentResourceGroupAction,
+    UpdateAgentResourceGroupActionResult,
+)
 from ai.backend.manager.services.agent.actions.watcher_agent_restart import (
     WatcherAgentRestartAction,
     WatcherAgentRestartActionResult,
@@ -88,17 +95,22 @@
     WatcherAgentStopAction,
     WatcherAgentStopActionResult,
 )
+from ai.backend.manager.services.agent.types import ConflictingSessionCleanupPolicy
+from ai.backend.manager.sokovan.scheduling_controller import SchedulingController
 from ai.backend.manager.types import OptionalState
 
 log = BraceStyleAdapter(logging.getLogger(__spec__.name))
 
+_RESOURCE_GROUP_CHANGED_REASON = "AGENT_RESOURCE_GROUP_CHANGED"
+
 
 class AgentService:
     _etcd: AsyncEtcd
     _config_provider: ManagerConfigProvider
     _agent_registry: AgentRegistry
     _agent_repository: AgentRepository
     _scheduler_repository: SchedulerRepository
+    _scheduling_controller: SchedulingController
     _hook_plugin_ctx: HookPluginContext
     _event_producer: EventProducer
     _agent_cache: AgentRPCCache
@@ -110,6 +122,7 @@ def __init__(
         config_provider: ManagerConfigProvider,
         agent_repository: AgentRepository,
         scheduler_repository: SchedulerRepository,
+        scheduling_controller: SchedulingController,
         hook_plugin_ctx: HookPluginContext,
         event_producer: EventProducer,
         agent_cache: AgentRPCCache,
@@ -119,6 +132,7 @@ def __init__(
         self._config_provider = config_provider
         self._agent_repository = agent_repository
         self._scheduler_repository = scheduler_repository
+        self._scheduling_controller = scheduling_controller
         self._hook_plugin_ctx = hook_plugin_ctx
         self._event_producer = event_producer
         self._agent_cache = agent_cache
@@ -153,6 +167,48 @@ async def sync_agent_registry(
 
         return SyncAgentRegistryActionResult(result=None, agent_data=agent_data)
 
+    async def update_resource_group(
+        self, action: UpdateAgentResourceGroupAction
+    ) -> UpdateAgentResourceGroupActionResult:
+        """
+        Change an agent's resource group, handling the sessions running on it.
+
+        The repository atomically gates on the agent's active sessions (raising
+        AgentHasConflictingSessions when they exist and ``force`` is unset),
+        commits the group change, and returns those sessions. The returned
+        sessions are then transitioned per ``policy``; their cleanup proceeds
+        asynchronously. ``RESCHEDULE`` is not implemented yet.
+        """
+        agent_id = action.agent_id
+        if action.policy is ConflictingSessionCleanupPolicy.RESCHEDULE:
+            # Rejected up front so no group change or termination happens.
+            raise ConflictingSessionRescheduleNotSupported()
+
+        kernels = await self._agent_repository.update_resource_group(
+            agent_id, action.resource_group_id, force=action.force
+        )
+        conflicting_session_ids = list({
+            SessionId(UUID(kernel.session.session_id)) for kernel in kernels
+        })
+
+        terminating_session_ids: list[SessionId] = []
+        if conflicting_session_ids:
+            # Graceful termination: sessions transition to TERMINATING and the
+            # container cleanup proceeds asynchronously in the next schedule cycle.
+            mark_result = await self._scheduling_controller.mark_sessions_for_termination(
+                conflicting_session_ids,
+                reason=_RESOURCE_GROUP_CHANGED_REASON,
+                forced=False,
+            )
+            terminating_session_ids = mark_result.terminating_sessions
+
+        return UpdateAgentResourceGroupActionResult(
+            agent_id=agent_id,
+            resource_group_id=action.resource_group_id,
+            conflicting_session_ids=conflicting_session_ids,
+            terminating_session_ids=terminating_session_ids,
+        )
+
     async def _request_watcher(
         self,
         agent_id: AgentId,

src/ai/backend/manager/services/agent/types.py

@@ -0,0 +1,16 @@
+import enum
+
+
+class ConflictingSessionCleanupPolicy(enum.StrEnum):
+    """
+    How to clean up sessions that conflict with an agent's resource group.
+
+    A session conflicts when its kernels were scheduled on the agent under a
+    resource group that no longer matches the agent's current (authoritative)
+    resource group.
+    """
+
+    TERMINATE = "terminate"
+    # Re-enqueue conflicting sessions back to PENDING. Not yet implemented: it
+    # depends on the RESCHEDULING state introduced by the preemption work.
+    RESCHEDULE = "reschedule"

src/ai/backend/manager/services/factory.py

@@ -170,6 +170,7 @@ def create_services(args: ServiceArgs) -> Services:
             args.config_provider,
             repositories.agent.repository,
             repositories.scheduler.repository,
+            args.scheduling_controller,
             args.hook_plugin_ctx,
             args.event_producer,
             args.agent_cache,