Repository navigation
Conversation
…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.
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.
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:
|
- Docstrings: Google-style Args/Attributes/Raises throughout; module docstring shows how a handler wires the pieces together; the publisher docstring now reflects that 'last' is the default policy; the estimator docstring states which servers it suits (latencies long enough to be seen in flight) and shows the formula in one block. - Parse Prometheus sample and label lines with regexes instead of a hand-rolled index parser. Quoted label values may now contain commas, braces and escapes. - Validate the publisher's aggregation policy up front (unknown policies raise ValueError instead of silently behaving like 'max') and split ingestion into _record(), which also stops 'last' from retaining samples. - Store estimator samples as a NamedTuple instead of positional tuples. Outputs are unchanged for the estimator, every valid aggregation policy, and valid Prometheus exposition (checked against the previous commit on randomized inputs).
AMOOOMA
force-pushed
the
pr-vllm-integration
branch
from
November 2, 2026 23:58
24ee97a to
60398e4
Compare
Each component now has one job: the sampler reads, the estimator computes
(the only window and the only multiplier), the publisher publishes, and the
config wires them together.
Remove the publisher's aggregation policies, its aggregation window and its
multiplier, and the matching AutoscalingMetricConfig fields. The capacity
estimate is already smoothed and its multiplier already applied, so for the
estimator these were a no-op ('last', 1.0) that read like a second smoothing
layer. A handler whose signal needs extra smoothing can apply it in its
signal_fn.
The estimator, sampler and Prometheus parsing are unchanged.
AMOOOMA
force-pushed
the
pr-vllm-integration
branch
from
November 3, 2026 00:22
60398e4 to
53a456a
Compare
Add AutoscalingMetricConfig.make_model_server_publisher(metrics_port_fn, completed_metric, running_metric). A handler for a model server that exposes a completed-requests counter and an in-flight gauge now only supplies those two metric names and how to find the server's port; the library builds the sampler and estimator, wires them together, honors load_signal_fn, skips ticks while the server is down, and returns a ready publisher. The sampler, estimator and make_publisher stay public for handlers that need a different signal.
AMOOOMA
force-pushed
the
pr-vllm-integration
branch
from
November 3, 2026 00:35
53a456a to
19344fd
Compare
Producing one capacity value crossed five classes (publisher -> signal -> sampler -> estimator -> publisher), two of which were only glue. Write the poll as a single function inside make_model_server_publisher (port, scrape, parse the two metrics, estimate) and remove PrometheusCounterGaugeSampler, _ModelServerCapacitySignal and AutoscalingMetricConfig.make_capacity_estimator. What remains is three classes with one job each (estimator computes, publisher publishes, config configures) plus the scrape/parse functions. Published values are unchanged.
AMOOOMA
force-pushed
the
pr-vllm-integration
branch
from
November 3, 2026 00:43
19344fd to
fa42fda
Compare
The config class mixed knobs with construction: resolve_tag() returned a copy with the tag filled in, and make_model_server_publisher() held the polling logic. Move construction to module-level functions: - make_model_server_publisher(...): the one call a model-server handler makes (metric names, port, model identifier, optional config). - make_publisher(...): for a handler's own capacity signal. Both resolve the metric tag internally, so handlers no longer call a separate resolve_tag step, and both take keyword-only arguments so the identifier, model name and metric names can't be swapped. AutoscalingMetricConfig is now just the knobs. Published values are unchanged.
AMOOOMA
force-pushed
the
pr-vllm-integration
branch
from
November 3, 2026 00:51
fa42fda to
6c32c92
Compare
Remove the make_publisher and make_model_server_publisher wrappers. A handler now constructs NumLoadedModelsPublisher directly with its server's metric names, a metrics_port_fn, the model identifier and an optional AutoscalingMetricConfig; the publisher resolves the metric tag, owns its CompletionRateCapacityEstimator, and on each poll reads the port, scrapes /metrics (or calls config.load_signal_fn), updates the estimate and publishes it. Behavior is unchanged.
AMOOOMA
force-pushed
the
pr-vllm-integration
branch
from
November 3, 2026 01:03
6c32c92 to
916f3a4
Compare
AMOOOMA
force-pushed
the
pr-vllm-integration
branch
from
November 3, 2026 01:09
916f3a4 to
70ff600
Compare
CompletionRateCapacityEstimator and AutoscalingMetricConfig take an aggregation: 'concurrency' (the default and the existing estimate), 'window_mean', 'window_max' or 'lifetime_max'. The last three use the raw completion rate and ignore concurrency, so they also work for servers whose in-flight gauge is almost always 0. The default is unchanged.
AMOOOMA
force-pushed
the
pr-vllm-integration
branch
from
November 3, 2026 01:20
70ff600 to
e30e639
Compare
Group the estimator docstring by mode, and make start() a no-op once a publisher is stopped: previously a superseded publisher that was started again would take its metric back from the newer one and then publish nothing. Add tests for scrape failures, non-numeric samples, negative gauge readings, running_percentile, window_secs=0, errors in the polling loop, and repeated start() calls; the module now has full line and branch coverage.
AMOOOMA
force-pushed
the
pr-vllm-integration
branch
from
November 3, 2026 01:30
e30e639 to
af8d757
Compare
BatchElements(key_fn=...) groups with GroupIntoBatches, which only bounds the element count, so length_fn and max_batch_weight were silently ignored; reject the combination instead. WithBatchKey now matches an unhashable key against a num_shards mapping by its str() form instead of raising, and rejects a num_shards callable that returns a bool.
The SDK harness reads process-wide metrics from another thread, and reset() followed by update() briefly left the distribution empty (count=0), which a concurrent report could pick up. Swap the cell's data in one step under its lock instead. Time estimator samples with time.monotonic() so wall-clock adjustments can't distort intervals.
AMOOOMA
force-pushed
the
pr-vllm-integration
branch
from
November 3, 2026 01:48
af8d757 to
1a5bc09
Compare
The previous atomic swap took DistributionCell._lock and assigned cell.data, neither of which is reachable from Python when apache_beam.metrics is compiled with Cython (as in CI and on workers), so _publish raised AttributeError. Install a fresh DistributionCell holding (sum=value, count=1) in the process-wide container with a single dict store instead; that is still atomic under the GIL, so a concurrent reader never sees an empty cell. Tests re-fetch the cell for every read since it is now replaced rather than mutated.
AMOOOMA
force-pushed
the
pr-vllm-integration
branch
from
November 5, 2026 18:50
1a5bc09 to
c713837
Compare
Address review feedback on keyed (prefix-aware) batching: - BatchElements(key_fn=...) now keys with WithBatchKey and reuses the existing stateful batching DoFns instead of GroupIntoBatches, so batch sizes adapt via _BatchSizeEstimator between min_batch_size and max_batch_size, and element_size_fn / max_batch_weight apply per key. key_fn + length_fn is still rejected since both choose the state key. - ModelHandler validates batch_key_fn / batch_key_num_shards eagerly in _validate_batching_args. - Keyed and maybe-keyed handlers pass the example, not the (key, example) pair, to batch_key_fn. - PromptPrefixKeyFn / ChatPrefixKeyFn take configurable delimiters (PrefixBatchingConfig.prefix_delimiters), and only search the first prefix_char_len characters for them.
- Add DistributionCell.reset_to(value), which replaces the distribution with a single sample under the cell's lock. It works when apache_beam.metrics is compiled with Cython, where _lock and data are not reachable from outside the class. NumLoadedModelsPublisher uses it instead of swapping the cell in the container. - CompletionRateCapacityEstimator returns None when requests are in flight but none completed within the window (slow first requests, first poll after an idle period), keeping the previously published capacity instead of reporting the floor. lifetime_max keeps reporting its peak.
AMOOOMA
force-pushed
the
pr-vllm-integration
branch
from
November 5, 2026 19:23
c713837 to
587a990
Compare
# Conflicts: # CHANGES.md
# Conflicts: # sdks/python/apache_beam/ml/inference/base_test.py
resolve() also calls the default extractor with delimiters=, so type it as Callable[..., Callable[[Any], Any]] instead of Callable[[int], ...].
Address review feedback: - Rename NumLoadedModelsPublisher to PrometheusNumLoadedModelsPublisher and document what the model server must provide: a local Prometheus text-format /metrics endpoint with a completed-requests counter and an in-flight gauge (or AutoscalingMetricConfig.load_signal_fn otherwise). - Export only AutoscalingMetricConfig and PrometheusNumLoadedModelsPublisher. The capacity estimator and the Prometheus scrape/parse helpers are now private; the aggregation modes are documented on AutoscalingMetricConfig.
…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 server passes vLLM's vllm:request_success_total counter and vllm:num_requests_running gauge to an autoscaling_metrics.NumLoadedModelsPublisher, which handles the scraping and estimation. 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
November 6, 2026 00:36
2eb9415 to
5687806
Compare
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 to anautoscaling_metrics.PrometheusNumLoadedModelsPublisher, which does the rest.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.