[AMOOOMA]

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:

• 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.

[codecov]

Codecov Report

❌ Patch coverage is 89.07363% with 46 lines in your changes missing coverage. Please review.
✅ Project coverage is 79.96%. Comparing base (a8d6afd) to head (46ecede).
⚠️ Report is 70 commits behind head on master.

Files with missing lines	Patch %	Lines
...on/apache_beam/ml/inference/autoscaling_metrics.py	86.34%	34 Missing ⚠️
.../python/apache_beam/ml/inference/vllm_inference.py	74.35%	10 Missing ⚠️
sdks/python/apache_beam/transforms/util.py	96.29%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #40760      +/-   ##
==========================================
- Coverage   80.04%   79.96%   -0.09%     
==========================================
  Files         490      495       +5     
  Lines       82916    83893     +977     
==========================================
+ Hits        66373    67081     +708     
- Misses      16543    16812     +269     

Flag	Coverage Δ
python	79.96% <89.07%> (-0.09%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.