vllm-project / vllm-project/aibrix
KubernetesJob batches finalize with wrong request_counts and missing usage (metadata doesn't reflect the worker's final state)
- Dominant language
- Go
- Stars
- 5.1k
- Forks
- 694
- Avg merge
- 1d 20h
- Merged PRs (30d)
- 98
Description
### 🐛 Describe the bug
### Summary
For the `KubernetesJob` (self-hosting) path, the worker does its own
prepare → dispatch → finalize and writes the output + per-request results. To avoid the
metadata service overwriting the worker's output, the metadata side now **skips
aggregation** (`_aggregate_on_finalize = False` for self-hosting). As a side effect, the
metadata service never reflects the worker's final stats, so the persisted/reported
`BatchJob` has:
- `request_counts.completed = 0` (output is correct, e.g. 3 lines, but the counter says 0), and
- `usage = None` (token usage is computed during aggregation, which is now skipped).
The output file itself is correct; only the **reported counts/usage** are wrong.
### Where
- `python/aibrix/aibrix/batch/job_driver/base.py`
- `execute()` sets `_aggregate_on_finalize = False` when `endpoint.source is None and runtime.provisions`
(self-hosting).
- `_await_self_hosted_run()` (~L289) re-reads the job via `get_job()` but returns the metadata's
stale in-memory copy rather than the worker's authoritative metastore state.
- `finalize_job()` (~L494) therefore skips `storage.finalize_job_output_data()`, which is also
the path that sums `usage` from the assembled output (`adapter.py:~372`).
### Impact
- Cosmetic-but-misleading: a completed `KubernetesJob` batch reports `completed=0` and `usage=null`.
- Anything consuming `request_counts` / `usage` (UI, billing, quotas) sees wrong numbers.
- Contrast: the `Kubernetes` (Deployment) path aggregates in-process and reports `completed=N` + usage correctly
### Steps to Reproduce
1. Submit a 3-record batch with `provider = "KubernetesJob"` against shared S3/minio + Redis.
2. Batch reaches `completed`; download output → 3 valid results.
3. Inspect the batch object → `request_counts.completed = 0`, `usage = null`.
### Expected behavior
For self-hosting runtimes, have the metadata side **reflect** the worker's final state from
the metastore instead of recomputing/overwriting: in `_await_self_hosted_run` (after
`await_completion`), adopt the worker-written `request_counts`, `usage`, and `output_file_id`
from the metastore as the terminal `BatchJob`. Keep aggregation disabled (don't overwrite the
worker's output). One change closes both the `completed=0` and `usage=None` gaps.
## Acceptance criteria
- [ ] A completed `KubernetesJob` batch reports correct `request_counts.completed` (== input count).
- [ ] `usage` reflects the worker's accumulated token usage.
- [ ] The worker's output file is still not overwritten (no regression of the empty-output bug).
- [ ] Deployment path counts/usage unchanged.
### Environment
nightly
Contributor guide
Assessment
This issue has not been assessed yet.