AMOOOMA opened a new pull request, #40907:
URL: https://github.com/apache/beam/pull/40907
Split out of #40760, following review feedback to separate the generic
pieces from the vLLM-specific ones and to split autoscaling into
collection/publishing vs. vLLM consumption. This PR is the generic library; the
vLLM handler uses it in a follow-up, and a Triton handler is planned to reuse
it.
## Why
Dataflow's GPU autoscaling reads `num_loaded_models_<tag>` in the
`BeamML_ModelManager` namespace as *the number of messages each worker can
hold*, and sizes the pool roughly as `ceil(backlog / reported)`. A handler that
serves inference through an out-of-process model server always has exactly one
loaded copy per worker, so it needs to publish an estimate of the worker's
capacity under that metric instead.
## What (`apache_beam.ml.inference.autoscaling_metrics`, new)
Nothing in this module depends on a particular model server.
- **`NumLoadedModelsPublisher`** (publishing): polls a capacity signal on a
background thread and publishes it under `num_loaded_models_<tag>`.
- Resets the process-wide distribution before each update, so the runner
sees the latest value rather than a lifetime mean.
- Optional aggregation on top (`last` by default; also `lifetime_max`,
`max`, `pNN`, `mean`).
- Never publishes 0 on start/stop, and a reloaded server's publisher
supersedes the stale one.
- **`PrometheusCounterGaugeSampler`** (collection): reads a server's
completed-requests counter and in-flight gauge from its Prometheus `/metrics`
endpoint. The two metric names are the only server-specific input.
- **`CompletionRateCapacityEstimator`** (collection): turns those samples
into `capacity = per-slot throughput × p90 concurrency × multiplier` over a
trailing window, with warmup and idle-tick handling.
- Does no I/O: it's fed `(completed_total, running)` samples.
- Normalizing by concurrency keeps the estimate steady when a worker is
briefly starved after a rebalance, instead of cratering and triggering a
downscale spiral.
- **`AutoscalingMetricConfig`**: user-facing knobs, metric-tag resolution
(`resolve_tag`: an explicit tag, then the RunInference model identifier
verbatim, then a sanitized model name), and factories for the estimator and
publisher.
Pairs with the `ModelHandler.set_model_identifier()` hook (separate PR),
which supplies the RunInference model identifier for `resolve_tag`. No
dependency between the two PRs.
## Testing
`autoscaling_metrics_test.py`:
- Prometheus parsing and sampler.
- Estimator math, counter reset, clock skew, warmup, idle ticks.
- Publisher aggregation, cell reset, supersession.
- Config tag precedence.
The estimator gives identical results to the version validated in #40760's
staging runs (checked on 2,000 randomized traces).
------------------------
Thank you for your contribution! Follow this checklist to help us
incorporate your contribution quickly and easily:
- [x] Replace these instructions with a description of your change and why
you have made it.
- [x] Mention the appropriate issue in your description (for example:
`addresses #123` or `fixes #456`).
- [ ] Update `CHANGES.md` with noteworthy changes.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
