[AMOOOMA]

Split out of #40760. This is the vLLM-specific part, built on the generic pieces that have all merged:

• [Python][RunInference] Add prefix-aware batching across workers #40905: prefix-aware batching (BatchElements(key_fn, num_shards), PrefixBatchingConfig)
• [Python][RunInference] Report the model identifier to model handlers #40906: ModelHandler.set_model_identifier()
• [Python][RunInference] Add a capacity-based autoscaling signal for model-server handlers #40907: capacity-based autoscaling signal (autoscaling_metrics.PrometheusNumLoadedModelsPublisher)

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.
  • The server passes vLLM's vllm:request_success_total counter and vllm:num_requests_running gauge to an autoscaling_metrics.PrometheusNumLoadedModelsPublisher, which does the rest.
  • The tag is the identifier RunInference reports through set_model_identifier.
  • In Dynamo mode (use_dynamo=True) the frontend doesn't serve vLLM's metrics, so the dynamo.vllm worker is started with DYN_SYSTEM_PORT and the publisher reads that port instead.
  • The publisher starts once the vLLM server is reachable and stops with it. It only holds a weak reference to the server, so a dropped server is still garbage-collected (which stops the publisher and the vLLM processes).

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:
  • Capacity published from vLLM's metrics end to end (server down skips the tick; scraped counter and gauge become the estimate, with the configured multiplier and poll interval).
  • Publisher lifecycle: none when disabled, started once the server is up, stopped and unregistered when the server is torn down.
  • A dropped server is garbage-collected and its publisher thread exits.
  • Dynamo mode starts the worker with a DYN_SYSTEM_PORT below 32768 (Dynamo parses it as a signed 16-bit integer) and scrapes it; native vLLM still reads the API server port.
  • 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.
• Dynamo staging run (L4, ai-dynamo 1.6.2 / vLLM 0.31, use_dynamo=True): each worker scraped its dynamo.vllm system port, found the vllm:* metrics and published num_loaded_models continuously; the autoscaler's per-worker target matched the published capacity and the job scaled to 2 workers.

---

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.83834% with 44 lines in your changes missing coverage. Please review.
✅ Project coverage is 79.88%. Comparing base (12bbb17) to head (24ee97a).
⚠️ Report is 11 commits behind head on master.

Files with missing lines	Patch %	Lines
...on/apache_beam/ml/inference/autoscaling_metrics.py	88.67%	29 Missing ⚠️
.../python/apache_beam/ml/inference/vllm_inference.py	69.04%	13 Missing ⚠️
sdks/python/apache_beam/transforms/util.py	96.36%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #40908      +/-   ##
==========================================
+ Coverage   79.83%   79.88%   +0.04%     
==========================================
  Files         494      496       +2     
  Lines       83575    84013     +438     
==========================================
+ Hits        66722    67113     +391     
- Misses      16853    16900      +47     

Flag	Coverage Δ
python	79.88% <89.83%> (+0.04%)	⬆️

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.