Repository navigation
Conversation
BatchElements batches per worker, so an LLM server behind RunInference sees batches that mix unrelated prompt prefixes and its prefix (KV) cache is thrashed. Add key_fn and num_shards to BatchElements. With key_fn set, a new WithBatchKey DoFn keys each element without the per-worker UUID that WithSharedKey and WithLengthBucketKey add, and GroupIntoBatches routes elements that share a key to the same worker, so every batch holds a single key. num_shards (an int, a per-key mapping, or a callable) spreads a hot key over a bounded number of random sub-keys. Incomplete keyed batches flush after max_batch_duration_secs, or 1s when it is unset, so low-traffic keys still make progress in streaming. ModelHandler forwards batch_key_fn / batch_key_num_shards to BatchElements, and the new apache_beam.ml.inference.prefix_batching module gives LLM handlers a PrefixBatchingConfig plus default prompt and chat prefix extractors. The config defaults to 8 shards per prefix: GroupIntoBatches serves each shard on one worker at a time, so a single shard caps a prefix shared by most traffic at one worker's throughput.
…er handlers Dataflow's GPU autoscaling reads num_loaded_models_<tag> in the BeamML_ModelManager namespace as the number of messages each worker can hold. A model handler that serves inference through an out-of-process server (vLLM, Triton, ...) always has exactly one loaded copy per worker, so it needs to publish a capacity estimate under that metric instead. Add a generic library for that in apache_beam.ml.inference.autoscaling_metrics: - NumLoadedModelsPublisher publishes a capacity signal on a background thread, resetting the process-wide distribution so the latest value (not a lifetime mean) is reported, with optional aggregation policies on top. - PrometheusCounterGaugeSampler reads a server's completed-requests counter and in-flight gauge from its Prometheus /metrics endpoint; the metric names are the only server-specific input. - CompletionRateCapacityEstimator turns those samples into a per-worker capacity estimate: per-slot throughput times p90 concurrency over a trailing window. It does no I/O. - AutoscalingMetricConfig holds the user-facing knobs, resolves the metric tag (explicit tag, then the RunInference model identifier verbatim, then a sanitized model name), and builds the estimator and publisher.
Add ModelHandler.set_model_identifier(), which RunInference calls on each worker before the first load_model() with the step's model_identifier (the value passed to RunInference, or the uuid it generates). Handlers that publish num_loaded_models_<identifier> themselves, such as handlers that front an out-of-process model server, can then use the exact tag the step is annotated with. The value is stable across side-input model updates, matching the metric tag ModelManager uses since apache#40714. The hook is a no-op by default and is forwarded by the keyed, maybe-keyed, prebatched, preprocessing and postprocessing wrappers.
…ling Wire the generic pieces into the vLLM handlers: - VLLMCompletionsModelHandler / VLLMChatModelHandler accept a prefix_batching_config. When set, requests sharing a prompt prefix (the leading characters of the prompt, or the role and leading content of the first chat message, unless a custom key_fn is given) are batched together on the same worker so vLLM's prefix cache is reused. - The handlers publish num_loaded_models_<model_identifier> (on by default, enable_autoscaling_metric / autoscaling_config to tune or disable) so Dataflow's GPU autoscaling sizes the pool by each worker's capacity. The signal reads vLLM's vllm:request_success_total counter and vllm:num_requests_running gauge through PrometheusCounterGaugeSampler and CompletionRateCapacityEstimator. The metric tag is the identifier RunInference reports through set_model_identifier, and the publisher starts once the 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.
AMOOOMA
force-pushed
the
pr-vllm-integration
branch
from
October 30, 2026 17:21
ff5bb1b to
24ee97a
Compare
Codecov Report❌ Patch coverage is 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
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Split out of #40760. This is the vLLM-specific part, built on the generic pieces:
BatchElements(key_fn, num_shards),PrefixBatchingConfig)ModelHandler.set_model_identifier()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_fnoverrides it.num_loaded_models_<model_identifier>so Dataflow's GPU autoscaling sizes the pool by each worker's capacity.enable_autoscaling_metricturns it off andautoscaling_configtunes it.vllm:request_success_totalcounter andvllm:num_requests_runninggauge viaPrometheusCounterGaugeSampler+CompletionRateCapacityEstimator.set_model_identifier.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:BatchElements(8 shards by default).with_preprocess_fn.Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:
addresses #123orfixes #456).CHANGES.mdwith noteworthy changes.