Add validation to prevent scaling group update with active kernels
- Dominant language
- Python
- Stars
- 670
- Forks
- 183
- Avg merge
- 17h 7m
- Merged PRs (30d)
- 358
Description
## Objective
Add validation logic to the `ModifyAgent` mutation to prevent updating an agent's scaling group when it has active kernels running.
## Background
When `update_scaling_group()` is called on an agent with running kernels, only the agent's `scaling_group` field is updated while kernel records retain their old `scaling_group` values. This causes the Sokovan scheduler to incorrectly calculate available resources because it filters kernels by `scaling_group` when computing occupancy.
## Implementation Details
### Location
- File: `src/ai/backend/manager/models/gql_models/agent.py`
- Method: `ModifyAgent.mutate()` (around line 888-891)
### Changes Required
Add a check before calling `registry.update_scaling_group()`:
```python
if (scaling_group := data.get("scaling_group")) is not None:
# Check if agent has active kernels
kernel_count_query = sa.select(sa.func.count()).select_from(kernels).where(
sa.and_(
kernels.c.agent == id,
kernels.c.status.in_(AGENT_RESOURCE_OCCUPYING_KERNEL_STATUSES)
)
)
kernel_count = await graph_ctx.db.scalar(kernel_count_query)
if kernel_count > 0:
raise InvalidAPIParameters(
f"Cannot change scaling group while agent has {kernel_count} active kernels"
)
await graph_ctx.registry.update_scaling_group(id, scaling_group)
```
## Acceptance Criteria
- [ ] Validation logic added to `ModifyAgent.mutate()`
- [ ] Error raised with clear message when active kernels exist
- [ ] Error message includes count of active kernels
- [ ] Unit test added to verify validation works
- [ ] Integration test added to verify end-to-end behavior
- [ ] Update succeeds when no active kernels exist (regression test)
JIRA Issue: BA-3570
Contributor guide
Assessment
This issue has not been assessed yet.