AMOOOMA opened a new pull request, #40760: URL: https://github.com/apache/beam/pull/40760
## Summary Productionizes prefix-aware batching and a capacity-based GPU autoscaling signal for vLLM `RunInference` pipelines (the approach evaluated in https://arxiv.org/abs/2610.22731). Two problems today: 1. Dataflow's GPU autoscaler reads `num_loaded_models_<model_identifier>` as "messages each worker can hold", but a model server such as vLLM always has exactly one loaded copy per worker, so the pool cannot scale on actual serving capacity. 2. `BatchElements` assembles batches per worker regardless of prompt prefix, so vLLM's prefix (KV) cache is thrashed by heterogeneous batches. ### Changes **Keyed / prefix-aware batching (`transforms/util.py`, `ml/inference/base.py`)** - `BatchElements(key_fn=..., num_shards=...)`: a new `WithBatchKey` DoFn keys elements *without* the per-worker UUID that `WithSharedKey` / `WithLengthBucketKey` add, then `GroupIntoBatches` shuffles elements sharing a key to the same worker so each batch is homogeneous. `num_shards` (int, per-key mapping, or callable) bounds hot-key fanout. - `ModelHandler(batch_key_fn=..., batch_key_num_shards=...)` forwards these to `BatchElements`. - New `ModelHandler.set_model_identifier()` hook: `_RunInferenceDoFn.setup()` reports the `model_identifier` (explicit or generated) to the handler before the first `load_model()`, so handlers that publish `num_loaded_models_<tag>` themselves use the exact tag the runner annotates the step with (stable across side-input model updates). **Autoscaling signal (`ml/inference/autoscaling_metrics.py`, new)** - `CompletionRateSignal` scrapes a server's Prometheus `/metrics` and estimates per-worker capacity as per-slot throughput × p90 in-flight concurrency over a trailing window (load-independent, so a briefly starved worker after a rebalance does not report a collapsed rate). - `NumLoadedModelsPublisher` publishes it on a background thread into the process-wide `BeamML_ModelManager` distribution (resetting the cell so the latest value, not a lifetime mean, is reported). Optional `'lifetime_max'` / windowed `'max'` / `'pNN'` / `'mean'` policies on top; default `'last'` since the signal is already smoothed. - `AutoscalingMetricConfig` holds the user-facing knobs. **Prefix batching config (`ml/inference/prefix_batching.py`, new)** - `PrefixBatchingConfig` plus default key extractors `PromptPrefixKeyFn` (completions) and `ChatPrefixKeyFn` (chat). **vLLM handlers (`ml/inference/vllm_inference.py`)** - `VLLMCompletionsModelHandler` / `VLLMChatModelHandler` gain `prefix_batching_config`, `enable_autoscaling_metric` (default `True`), and `autoscaling_config`. `_VLLMModelServer` starts the publisher once the server is reachable and stops it on teardown. Depends on #40671 (already merged), which scopes the `model_identifier` annotation to the `BeamML_RunInference` step; without it the keying stage added by `GroupIntoBatches` would also be sized by the GPU autoscaler. ## Testing - New/updated unit tests in `autoscaling_metrics_test.py`, `prefix_batching_test.py`, `vllm_inference_test.py`, `base_test.py` (identifier hook, end-to-end prefix batching through `RunInference`), and `util_test.py` (`WithBatchKey`, sharding specs, `BatchElements(key_fn=...)`). - 4h Dataflow staging run (2 × 2h sine load, 5–40 RPS, OpenOrca prompts, prefix batching with 4 shards): the `RunInference` stage scaled 1 → 32–46 → 1 GPU workers in both cycles and the worker pool followed it down (50 → 13 and 50 → 6 VMs), with ~302k results written to BigQuery. --- 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. - [ ] 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]
