Repository navigation
[Python][RunInference] Add prefix-aware batching across workers - #40905
Conversation
|
R: @damccorm |
|
Stopping reviewer notifications for this pull request: review requested by someone other than the bot, ceding control. If you'd like to restart, comment |
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.
465e446 to
07aac31
Compare
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #40905 +/- ##
============================================
+ Coverage 58.27% 58.29% +0.02%
- Complexity 2276 2277 +1
============================================
Files 1136 1137 +1
Lines 169393 169544 +151
Branches 1482 1482
============================================
+ Hits 98708 98837 +129
- Misses 68132 68160 +28
+ Partials 2553 2547 -6
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:
|
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.
| batch_key_num_shards: Optional[Union[int, | ||
| Mapping[Any, int], | ||
| Callable[[Any], int]]] = None, | ||
| large_model: bool = False, |
There was a problem hiding this comment.
Nit: ModelHandler.__init__ calls _validate_batching_args(...) below to eagerly validate max_batch_weight, element_size_fn, batch_length_fn, and batch_bucket_boundaries at handler construction time (before pipeline expansion). Could we pass batch_key_fn and batch_key_num_shards into _validate_batching_args as well so invalid combinations (e.g. batch_key_num_shards without batch_key_fn, or batch_key_fn + batch_length_fn) fail eagerly at init time?
There was a problem hiding this comment.
Done in 7acd8e9. _validate_batching_args now also covers batch_key_fn / batch_key_num_shards, using the same _validate_num_shards as BatchElements, so bad combinations (e.g. batch_key_num_shards without batch_key_fn, or batch_key_fn + batch_length_fn) fail when the handler is constructed.
| if batch_bucket_boundaries is not None: | ||
| self._batching_kwargs['bucket_boundaries'] = batch_bucket_boundaries | ||
| if batch_key_fn is not None: | ||
| self._batching_kwargs['key_fn'] = batch_key_fn |
There was a problem hiding this comment.
In _keyed_batch_elements_kwargs (around line 774), we wrap element_size_fn and length_fn in _ApplyToExample so they receive example (element[1]) instead of the (key, example) tuple when a handler is wrapped in KeyedModelHandler or MaybeKeyedModelHandler:
def _keyed_batch_elements_kwargs(batch_elements_kwargs: Mapping[str, Any]):
kwargs = dict(batch_elements_kwargs)
for name in ('element_size_fn', 'length_fn', 'key_fn'):
if name in kwargs and kwargs[name] is not None:
kwargs[name] = _ApplyToExample(kwargs[name])
return kwargsRight now 'key_fn' is missing from that loop, so wrapping a prefix-batched handler in KeyedModelHandler / MaybeKeyedModelHandler passes (key, example) into key_fn (PromptPrefixKeyFn raises TypeError: Expected str prompt for prefix key extraction, got tuple, and ChatPrefixKeyFn inspects element[0]—the key—as the first chat message). Could we add 'key_fn' there and add a unit test with KeyedModelHandler?
There was a problem hiding this comment.
Good catch, fixed in 7acd8e9. key_fn is now wrapped in _ApplyToExample like element_size_fn / length_fn. I added a unit test for keyed and maybe-keyed handlers and an end-to-end RunInference(KeyedModelHandler(...)) test.
| if self._key_fn is not None: | ||
| element_type = pcoll.element_type | ||
| # GroupIntoBatches treats a buffering duration of 0 as "no limit", which | ||
| # in a streaming global window would hold an incomplete batch forever. | ||
| buffering_secs = ( | ||
| self._max_batch_dur if self._max_batch_dur is not None else | ||
| _DEFAULT_KEYED_MAX_BUFFERING_SECS) | ||
| return ( | ||
| pcoll | ||
| | 'WithBatchKey' >> ParDo( | ||
| WithBatchKey(self._key_fn, self._num_shards)).with_output_types( | ||
| tuple[str, element_type]) | ||
| | 'GroupIntoBatchesByKey' >> GroupIntoBatches( | ||
| self._batch_size_estimator._max_batch_size, | ||
| max_buffering_duration_secs=buffering_secs, | ||
| clock=self._clock) | ||
| | 'DropBatchKey' >> Map(lambda kv: list(kv[1])).with_output_types( | ||
| list[element_type])) |
There was a problem hiding this comment.
Why use GroupIntoBatches(self._batch_size_estimator._max_batch_size, ...) here instead of plugging WithBatchKey into the existing stateful ParDo(keying_dofn) | ParDo(batching_dofn) path below (_pardo_stateful_batch_elements / _pardo_stateful_weighted_batch_elements)?
In Python Beam, GroupIntoBatches expands to _pardo_group_into_batches (a stateful DoFn with the same BagStateSpec / CombiningValueStateSpec / TimerSpec setup as _pardo_stateful_batch_elements), and any stateful DoFn already shuffles/groups elements across workers by the key emitted by keying_dofn (WithSharedKey and WithLengthBucketKey only stay worker-local because they include a per-worker UUID in the key).
Using GroupIntoBatches(self._batch_size_estimator._max_batch_size) has two downsides:
- Default
max_batch_size=10000footgun:BatchElementsdefaultsmax_batch_sizeto10000(expecting_BatchSizeEstimatorto start atmin_batch_size=1and adapt based ontarget_batch_duration_secs). If a user enablesPrefixBatchingConfig()on aModelHandlerwithout explicitly settingmax_batch_size,GroupIntoBatchesuses a fixed batch size of10,000and will buffer up to 10,000 LLM requests into a singlerun_inferencecall. - It silently ignores
_BatchSizeEstimatorparameters (min_batch_size,target_batch_overhead,target_batch_duration_secs,record_metrics) and cannot be combined withelement_size_fn/max_batch_weight(which_pardo_stateful_weighted_batch_elementsalready supports and is particularly useful for LLM token-budget batching).
There was a problem hiding this comment.
Agreed, switched in 7acd8e9. key_fn now just swaps the keying DoFn (WithBatchKey instead of WithSharedKey / WithLengthBucketKey) in front of the existing stateful batching DoFn. The estimator, element_size_fn and max_batch_weight all apply per key, and batches start at min_batch_size and grow from there instead of using the 10k max_batch_size. Setting min_batch_size == max_batch_size keeps fixed batch sizes. key_fn + length_fn is still rejected because both pick the state key.
| """Default prefix key extractor for plain-text prompts. | ||
|
|
||
| Keys a prompt string by its leading ``prefix_char_len`` characters, or by | ||
| everything before the first ``\\n\\n`` when that comes sooner, so prompts |
There was a problem hiding this comment.
Should the delimiter(s) be configurable here? I've seen some prompts like:
Foo, bar, baz. Choose from one of these options when responding:
- x
- y
- z
- ```
where I could imagine a `:` being a useful delimeter.
There was a problem hiding this comment.
Yes, done in 7acd8e9. PromptPrefixKeyFn / ChatPrefixKeyFn take delimiters, and PrefixBatchingConfig(prefix_delimiters=(':',)) sets them for the handler's default extractor. The defaults are unchanged: a blank line for prompts, none for chat. The earliest delimiter within prefix_char_len wins.
| raise TypeError( | ||
| f'Expected str prompt for prefix key extraction, got ' | ||
| f'{type(element).__name__}.') | ||
| sep_idx = element.find('\n\n') |
There was a problem hiding this comment.
Nit: Should we limit the search to element[:_prefix_char_len].find('\n\n') to shrink the search space for large prompts.
There was a problem hiding this comment.
Done in 7acd8e9. The search is bounded to the first prefix_char_len characters (plus the delimiter length), so long prompts cost the same as short ones.
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.
# Conflicts: # CHANGES.md
damccorm
left a comment
damccorm
left a comment
There was a problem hiding this comment.
Thanks - could you resolve the conflict. After that and checks run we can merge
|
Oh and please fix linting |
Yep for sure! |
# 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], ...].
…ling Wire the RunInference prefix batching (apache#40905) and the capacity-based autoscaling signal (apache#40907) 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 routed to the same worker and batched together 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.PrometheusNumLoadedModelsPublisher, which does the scraping and estimation. The metric tag is the identifier RunInference reports through ModelHandler.set_model_identifier. - The publisher starts once the vLLM server is reachable and stops with it. Its thread reaches the server only through a weak reference, 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.
Split out of #40760, following review feedback to separate the generic pieces from the vLLM-specific ones. This PR is the generic prefix-aware batching; vLLM support builds on it in a follow-up.
Why
BatchElementsbatches per worker, so an LLM model handler behind RunInference sends its server batches that mix unrelated prompt prefixes, and the server's prefix (KV) cache gets thrashed instead of reused.What
BatchElements(key_fn=..., num_shards=...)(transforms/util.py)key_fnset, a newWithBatchKeyDoFn keys each element without the per-worker UUID thatWithSharedKey/WithLengthBucketKeyadd.GroupIntoBatchesthen 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, so one key isn't capped at one worker.max_batch_duration_secs, or 1s when unset (GroupIntoBatchestreats 0 as "no limit", which would stall low-traffic keys in streaming).key_fncan't be combined withlength_fnormax_batch_weight(raisesValueError):GroupIntoBatchesonly bounds the element count, so those options would be silently ignored.RunInference plumbing (
ml/inference/base.py)ModelHandler(batch_key_fn=..., batch_key_num_shards=...)forwards toBatchElements.apache_beam.ml.inference.prefix_batching(new, generic for LLM handlers)PrefixBatchingConfig, plus default extractorsPromptPrefixKeyFn(completions) andChatPrefixKeyFn(OpenAI-style chat messages; unknown message schemas raise rather than collapsing onto one key).BatchElements/ModelHandlerdefaults stay at 1 shard.Testing
util_test.py: homogeneous keyed batches, shard specs (int / mapping / callable), validation, default buffering.prefix_batching_test.py: extractors and config.base_test.py:batch_key_fn/num_shardsreachBatchElements, and an end-to-end RunInference run produces prefix-pure batches.Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:
addresses #123orfixes #456).CHANGES.mdwith noteworthy changes.