Migrate endpoint autoscaling KERNEL/INFERENCE_FRAMEWORK rules to Prometheus and remove legacy KernelStatistics/EndpointStatistics
- Dominant language
- Python
- Stars
- 670
- Forks
- 183
- Avg merge
- 15h 13m
- Merged PRs (30d)
- 368
Description
## Background
Endpoint auto-scaling evaluates rules through three AutoScalingMetricSource variants. PROMETHEUS already runs through executor._fetch_prometheus_metrics (executor.py:745), but KERNEL and INFERENCE_FRAMEWORK still pull from Valkey via the legacy GQL dataloaders KernelStatistics.batch_load_by_kernel_impl and EndpointStatistics.batch_load_by_endpoint_impl. These are the last call sites that keep KernelStatistics / EndpointStatistics (in src/ai/backend/manager/api/gql_legacy/statistics.py) alive for kernel/endpoint live stats; once they move to Prometheus, those legacy classes can be deleted along with the GraphQL kernel.live_stat dataloader path that still depends on them.
## Target call sites
- src/ai/backend/manager/repositories/deployment/repository.py:847 — KernelStatistics.batch_load_by_kernel_impl (KERNEL branch around lines 820 and 906)
- src/ai/backend/manager/repositories/deployment/repository.py:859 — EndpointStatistics.batch_load_by_endpoint_impl (INFERENCE_FRAMEWORK branch around lines 825/936)
- src/ai/backend/manager/api/gql_legacy/kernel.py:202 — get_session_statistics_batch dataloader (only consumer of the kernel side after autoscaling moves)
- src/ai/backend/manager/api/gql_legacy/statistics.py — KernelStatistics / EndpointStatistics class definitions
## Scope
Phase A — autoscaling cutover:
- Replace KernelStatistics.batch_load_by_kernel_impl call in the KERNEL branch with a Prometheus query that returns the same shape used downstream (cpu_util / mem / etc.). Reuse the existing PROMETHEUS branch helpers where possible.
- Replace EndpointStatistics.batch_load_by_endpoint_impl call in the INFERENCE_FRAMEWORK branch with a Prometheus query for inference-app metrics. If a suitable helper does not exist yet, add a thin wrapper in MetricRepository / ContainerMetricService rather than introducing one-off code in the autoscaling path.
- Keep the AutoScalingMetricSource enum and rule schema unchanged — only the backing data source changes.
Phase B — legacy removal:
- Remove KernelStatistics and EndpointStatistics classes from api/gql_legacy/statistics.py.
- Remove the inference_metrics_by_kernel dataloader registration in api/gql_legacy/session.py if it relied on these classes (verify it does not reference valkey_live separately; current note: the dataloader name string is reused but the data path goes through valkey_live.get_session_statistics_batch, so confirm before deleting).
- Update imports in deployment/repository.py and any other consumers.
## Success Criteria
### Autoscaling
- [ ] AutoScalingMetricSource.KERNEL rule evaluation reads from Prometheus (verified by toggling Valkey stat writes off and confirming the rule still fires)
- [ ] AutoScalingMetricSource.INFERENCE_FRAMEWORK rule evaluation reads from Prometheus (same verification)
- [ ] AutoScalingMetricSource.PROMETHEUS rule evaluation is unchanged (no regression)
### Legacy removal
- [ ] api/gql_legacy/statistics.py no longer defines KernelStatistics or EndpointStatistics
- [ ] grep -rn "KernelStatistics\|EndpointStatistics" src/ai/backend/manager/ returns no matches outside graphene type registrations that are also removed
- [ ] kernel.live_stat / endpoint.live_stat GraphQL fields still resolve through the new Prometheus path (BA-5744 already migrated the resolvers — verify nothing broke)
### Common
- [ ] pants check passes for manager package
- [ ] pants test passes for affected packages
JIRA Issue: BA-6115
Contributor guide
Research direction
Start with the KERNEL and INFERENCE_FRAMEWORK branches in src/ai/backend/manager/repositories/deployment/repository.py and compare them with the existing PROMETHEUS path and helpers. Then inspect api/gql_legacy/kernel.py, session.py, and statistics.py to trace the remaining dataloader and class consumers. Verify the autoscaling and live-stat success criteria, then run pants check and pants test for the affected manager packages.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- graphql, prometheus, python
- Domain
- api, backend, observability
- Issue type
- Refactor
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100