lablup / lablup/backend.ai

Migrate endpoint autoscaling KERNEL/INFERENCE_FRAMEWORK rules to Prometheus and remove legacy KernelStatistics/EndpointStatistics

Open
#11,704 0 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.