AMOOOMA opened a new pull request, #40908:
URL: https://github.com/apache/beam/pull/40908
Split out of #40760. This is the vLLM-specific part, built on the generic
pieces:
- #40905: prefix-aware batching (`BatchElements(key_fn, num_shards)`,
`PrefixBatchingConfig`)
- #40906: `ModelHandler.set_model_identifier()`
- #40907: capacity-based autoscaling signal library
> **Depends on #40905, #40906 and #40907.** Until they merge, this branch
includes their commits; **only the last commit (`vLLM: prefix-aware batching
and capacity-based autoscaling`) is new here**. I'll rebase as the dependencies
land.
## What
`VLLMCompletionsModelHandler` / `VLLMChatModelHandler`:
- **`prefix_batching_config`**: when set, requests that share a prompt
prefix are batched together on the same worker, so vLLM's prefix cache is
reused. By default the key is the leading characters of the prompt
(completions), or the role and leading content of the first message (chat);
`key_fn` overrides it.
- **Capacity-based autoscaling signal**: the handlers publish
`num_loaded_models_<model_identifier>` so Dataflow's GPU autoscaling sizes the
pool by each worker's capacity.
- On by default; `enable_autoscaling_metric` turns it off and
`autoscaling_config` tunes it.
- Reads vLLM's `vllm:request_success_total` counter and
`vllm:num_requests_running` gauge via `PrometheusCounterGaugeSampler` +
`CompletionRateCapacityEstimator`.
- The tag is the identifier RunInference reports through
`set_model_identifier`.
- The publisher starts once the vLLM server is reachable and stops with it.
The only vLLM-specific parts are the two metric names, the default prefix
extractors and the server lifecycle; everything else lives in the generic
modules.
## Testing
- `vllm_inference_test.py`:
- Publisher defaults (no double smoothing), and the multiplier and window
reaching the estimator.
- The vLLM metric names reaching the sampler.
- Tag precedence (identifier used verbatim, sanitized model-name fallback,
explicit tag wins).
- Prefix config reaching `BatchElements` (8 shards by default).
- An end-to-end RunInference run with prefix batching and
`with_preprocess_fn`.
- 4h Dataflow staging run (two 2h sine cycles, 5–40 RPS, OpenOrca, T4, up to
50 workers): GPU scaling applied only to the inference stage. The pool scaled
with load and back down at low traffic (50 → 9 and 50 → 2 VMs). An additional
single-cycle run with the default 8 shards wrote all 160,166 published requests
and scaled down to 4 workers.
------------------------
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`).
- [x] 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]
