[Python][RunInference] Add prefix-aware batching across workers

CHANGES.md

@@ -70,6 +70,7 @@
 
 * X feature added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)).
 * (Python) Added support for Protobuf 7.x ([#40580](https://github.com/apache/beam/issues/40580)).
+* (Python) `BatchElements` can group elements by key across workers (`key_fn`, `num_shards`), and RunInference model handlers can use it through `batch_key_fn` / `batch_key_num_shards`. `apache_beam.ml.inference.prefix_batching` adds a `PrefixBatchingConfig` and prompt/chat prefix extractors so LLM handlers can batch requests that share a prompt prefix and hit the server's prefix cache ([#40905](https://github.com/apache/beam/pull/40905)).
 
 ## Breaking Changes
 

sdks/python/apache_beam/ml/inference/base.py

@@ -213,6 +213,10 @@ def __init__(
       element_size_fn: Optional[Callable[[Any], int]] = None,
       batch_length_fn: Optional[Callable[[Any], int]] = None,
       batch_bucket_boundaries: Optional[list[int]] = None,
+      batch_key_fn: Optional[Callable[[Any], Any]] = None,
+      batch_key_num_shards: Optional[Union[int,
+                                           Mapping[Any, int],
+                                           Callable[[Any], int]]] = None,
       large_model: bool = False,
       model_copies: Optional[int] = None,
       **kwargs):
@@ -243,6 +247,21 @@ def __init__(
         semantics): bucket i covers lengths in [boundaries[i-1], boundaries[i]).
         Requires batch_length_fn. Defaults to [16, 32, 64, 128, 256, 512] when
         batch_length_fn is set.
+      batch_key_fn: a callable mapping an element to a grouping key. It
+        receives the element as it enters batching, i.e. *after* any
+        ``with_preprocess_fn`` functions have been applied. When set,
+        BatchElements keys elements via WithBatchKey (without a per-worker UUID)
+        and groups them across workers using GroupIntoBatches so that elements
+        sharing the same key (e.g., prompt prefix or category) are routed to
+        the same worker and batched together.
+      batch_key_num_shards: bounds hot-key fanout when batch_key_fn is set.
+        Either an int applied to every key, a mapping from batch key to shard
+        count for a known skewed distribution (keys not in the mapping get 1
+        shard), or a callable from batch key to shard count. Each key is
+        spread across at most that many random sub-keys / workers. Defaults
+        to 1, which serves each key on one worker at a time; skewed key
+        distributions need more shards for their hot keys (see
+        ``apache_beam.ml.inference.prefix_batching.PrefixBatchingConfig``).
       large_model: set to true if your model is large enough to run into
         memory pressure if you load multiple copies.
       model_copies: The exact number of models that you would like loaded
@@ -273,6 +292,10 @@ def __init__(
       self._batching_kwargs['length_fn'] = batch_length_fn
     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
+    if batch_key_num_shards is not None:
+      self._batching_kwargs['num_shards'] = batch_key_num_shards
     self._large_model = large_model
     self._model_copies = model_copies
     self._share_across_processes = large_model or (model_copies is not None)

sdks/python/apache_beam/ml/inference/base_test.py

@@ -2817,6 +2817,20 @@ def test_batching_kwargs_none_values_omitted(self):
     self.assertNotIn('bucket_boundaries', kwargs)
     self.assertEqual(kwargs['min_batch_size'], 5)
 
+  def test_batch_key_fn_and_num_shards_passed_to_kwargs(self):
+    key_fn = lambda s: s[:4]
+    handler = FakeModelHandlerForBatching(
+        batch_key_fn=key_fn, batch_key_num_shards=4)
+    kwargs = handler.batch_elements_kwargs()
+    self.assertIs(kwargs['key_fn'], key_fn)
+    self.assertEqual(kwargs['num_shards'], 4)
+
+  def test_batch_key_num_shards_mapping_passed_to_kwargs(self):
+    shard_map = {'hot_prefix': 8}
+    handler = FakeModelHandlerForBatching(
+        batch_key_fn=lambda s: s, batch_key_num_shards=shard_map)
+    self.assertEqual(handler.batch_elements_kwargs()['num_shards'], shard_map)
+
 
 class PaddingReportingStringModelHandler(base.ModelHandler[str, str,
                                                            FakeModel]):
@@ -2829,6 +2843,18 @@ def run_inference(self, batch, model, inference_args=None):
     return [f'{s}:{max_len}' for s in batch]
 
 
+class PrefixPeerReportingStringModelHandler(base.ModelHandler[str,
+                                                              str,
+                                                              FakeModel]):
+  """Reports each element with the sorted prefix set of the batch it ran in."""
+  def load_model(self):
+    return FakeModel()
+
+  def run_inference(self, batch, model, inference_args=None):
+    prefixes = ','.join(sorted({s.split(':', 1)[0] for s in batch}))
+    return [f'{s}|batch_prefixes={prefixes}|size={len(batch)}' for s in batch]
+
+
 class RunInferenceLengthAwareBatchingTest(unittest.TestCase):
   """End-to-end tests for PR2 length-aware batching in RunInference."""
   def test_run_inference_with_length_aware_batch_elements(self):
@@ -2847,6 +2873,44 @@ def test_run_inference_with_length_aware_batch_elements(self):
           | base.RunInference(handler))
       assert_that(results, equal_to(['a:2', 'bb:2', 'cccccc:7', 'ddddddd:7']))
 
+  def test_run_inference_with_batch_key_fn_and_preprocess_fn(self):
+    # batch_key_fn receives the element *after* preprocessing (the prompt
+    # string here), not the raw input record.
+    handler = PrefixPeerReportingStringModelHandler(
+        min_batch_size=2,
+        max_batch_size=2,
+        max_batch_duration_secs=1.0,
+        batch_key_fn=lambda prompt: prompt.split(':', 1)[0]).with_preprocess_fn(
+            lambda rec: f"{rec['sector']}:{rec['id']}")
+
+    records = [
+        {
+            'sector': 'fin', 'id': '1'
+        },
+        {
+            'sector': 'med', 'id': '1'
+        },
+        {
+            'sector': 'fin', 'id': '2'
+        },
+        {
+            'sector': 'med', 'id': '2'
+        },
+    ]
+    with TestPipeline('FnApiRunner') as p:
+      results = (
+          p
+          | beam.Create(records, reshuffle=False)
+          | base.RunInference(handler))
+      assert_that(
+          results,
+          equal_to([
+              'fin:1|batch_prefixes=fin|size=2',
+              'fin:2|batch_prefixes=fin|size=2',
+              'med:1|batch_prefixes=med|size=2',
+              'med:2|batch_prefixes=med|size=2',
+          ]))
+
 
 class BatchLengthsModelHandler(base.ModelHandler[str, tuple, FakeModel]):
   """Predicts the lengths of the examples in the batch of each example."""

sdks/python/apache_beam/ml/inference/prefix_batching.py

@@ -0,0 +1,229 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Prefix-aware batching for LLM model handlers.
+
+Generic across LLM servers that keep a prefix (KV) cache, such as vLLM, SGLang
+or TensorRT-LLM; nothing here depends on a particular server.
+
+``ModelHandler(batch_key_fn=...)`` makes ``BatchElements`` route elements that
+share a key to the same worker via ``GroupIntoBatches``, so each batch the
+server sees is homogeneous and its prefix cache gets hit instead of thrashed.
+This module packages the knobs for that into one :class:`PrefixBatchingConfig`
+plus the default key extractors a handler installs when the user does not
+supply their own:
+
+* :class:`PromptPrefixKeyFn` -- for plain-text prompts (completions API).
+* :class:`ChatPrefixKeyFn` -- for chat message lists (chat API).
+
+Both see the element *after* any ``with_preprocess_fn`` has run, i.e. the
+prompt / message list the handler will actually send. A handler exposes a
+``PrefixBatchingConfig`` argument, picks which extractor is its default, and
+passes :func:`resolve_batching_kwargs` to ``ModelHandler.__init__``.
+"""
+
+# pytype: skip-file
+
+import dataclasses
+from collections.abc import Callable
+from collections.abc import Mapping
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import Any
+from typing import Optional
+from typing import Union
+
+__all__ = [
+    'ChatPrefixKeyFn',
+    'PrefixBatchingConfig',
+    'PromptPrefixKeyFn',
+]
+
+# Characters of prompt (or leading message content) the default extractors key
+# on. Long enough to cover a typical system prompt, short enough that keys are
+# cheap to hash and shuffle.
+DEFAULT_PREFIX_CHAR_LEN = 128
+# GroupIntoBatches buffering applied when the handler's own
+# ``max_batch_duration_secs`` is unset, so low-traffic keys still flush in
+# streaming instead of waiting for a full batch.
+DEFAULT_KEYED_BATCH_DURATION_SECS = 1.0
+# Shards per prefix key. A single shard serves a hot prefix (e.g. a common
+# system prompt) on one worker at a time; a few shards keep that prefix's
+# throughput scaling with the pool while every shard still batches a single
+# prefix.
+DEFAULT_NUM_SHARDS = 8
+
+# Accepted forms of ``num_shards``: one count for every key, a per-key table,
+# or a callable computed per key (see ``WithBatchKey`` in transforms/util.py).
+NumShardsSpec = Union[int, Mapping[Any, int], Callable[[Any], int]]
+
+
+class PromptPrefixKeyFn:
+  """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
+  sharing the same system prompt or instruction header are routed to the same
+  worker and batched together.
+  """
+  def __init__(self, prefix_char_len: int = DEFAULT_PREFIX_CHAR_LEN):
+    if prefix_char_len <= 0:
+      raise ValueError('prefix_char_len must be positive.')
+    self._prefix_char_len = prefix_char_len
+
+  def __call__(self, element: str) -> str:
+    if not isinstance(element, str):
+      raise TypeError(
+          f'Expected str prompt for prefix key extraction, got '
+          f'{type(element).__name__}.')
+    sep_idx = element.find('\n\n')
+    if 0 < sep_idx <= self._prefix_char_len:
+      return element[:sep_idx]
+    return element[:self._prefix_char_len]
+
+
+class ChatPrefixKeyFn:
+  """Default prefix key extractor for chat message lists.
+
+  Keys a conversation by the role and leading ``prefix_char_len`` characters of
+  its first (typically system) message, so conversations sharing the same
+  system prompt are routed to the same worker and batched together. Messages
+  may be objects with ``role``/``content`` attributes (e.g.
+  ``OpenAIChatMessage``) or dicts with those keys. Any other message schema
+  raises ``TypeError`` rather than silently collapsing every conversation onto
+  one key; pass ``PrefixBatchingConfig(key_fn=...)`` for custom schemas.
+  """
+  def __init__(self, prefix_char_len: int = DEFAULT_PREFIX_CHAR_LEN):
+    if prefix_char_len <= 0:
+      raise ValueError('prefix_char_len must be positive.')
+    self._prefix_char_len = prefix_char_len
+
+  def __call__(self, element: Sequence[Any]) -> str:
+    if not isinstance(element, (list, tuple)):
+      raise TypeError(
+          f'Expected a sequence of chat messages for prefix key extraction, '
+          f'got {type(element).__name__}.')
+    if not element:
+      return ''
+    first_msg = element[0]
+    if hasattr(first_msg, 'role') and hasattr(first_msg, 'content'):
+      role, content = first_msg.role, first_msg.content
+    elif (isinstance(first_msg, Mapping) and 'role' in first_msg and
+          'content' in first_msg):
+      role, content = first_msg['role'], first_msg['content']
+    else:
+      raise TypeError(
+          'Chat message has no role/content attributes or keys (got '
+          f'{type(first_msg).__name__}); pass '
+          'PrefixBatchingConfig(key_fn=...) to key a custom message schema.')
+    content_str = str(content or '')
+    return f'{role or ""}:{content_str[:self._prefix_char_len]}'
+
+
+@dataclass(frozen=True)
+class PrefixBatchingConfig:
+  """User-facing knobs for prefix-aware batching.
+
+  Pass an instance to an LLM model handler that accepts one (the handler
+  forwards it to ``ModelHandler`` via :func:`resolve_batching_kwargs`) to turn
+  keyed batching on; leave it unset to keep the default unkeyed batching. Each
+  field:
+
+  * ``key_fn`` -- callable mapping an element (as it enters batching, i.e.
+    after any ``with_preprocess_fn``) to a grouping key. Leave unset (the
+    default) to use the handler's default prefix extractor
+    (:class:`PromptPrefixKeyFn` or :class:`ChatPrefixKeyFn`).
+  * ``prefix_char_len`` -- how many leading characters the default extractor
+    keys on (default 128). Ignored when ``key_fn`` is set.
+  * ``num_shards`` -- bounds hot-key fanout: an ``int`` applied to every key,
+    a ``Mapping[key, int]`` with per-key shard counts for a known skewed
+    prefix distribution (e.g. ``{'<the common system prompt>': 32}``;
+    unlisted keys get 1), or a ``Callable[[key], int]``. Each key is spread
+    across at most that many workers (default 8).
+
+    **Sharding is what lets skewed workloads scale.** ``GroupIntoBatches``
+    keeps per-key state, so all batches for one shard of a key are processed
+    by a single worker at a time. With a single shard, a prefix shared by most
+    traffic (e.g. a common system prompt) is served at the throughput of one
+    worker no matter how many workers the pipeline has, and its requests
+    queue up behind that key. Runners that size the pool from the number of
+    keys with pending work, such as Dataflow's GPU autoscaling, also do not
+    see that queue and will not scale up to drain it. The default of 8 covers
+    moderate skew; raise it, or map the hot prefixes to a larger count, when
+    one prefix carries most of the traffic and the pool scales well past 8
+    workers. The trade-off is prefix-cache locality: each shard of a prefix
+    warms its own copy of that prefix's KV cache, so set ``num_shards=1`` only
+    when prefixes are evenly spread and cache reuse matters most.
+
+  The handler calls :meth:`resolve` with its default extractor and then
+  :meth:`batching_kwargs` to obtain the ``ModelHandler`` arguments.
+  """
+  key_fn: Optional[Callable[[Any], Any]] = None
+  prefix_char_len: int = DEFAULT_PREFIX_CHAR_LEN
+  num_shards: Optional[NumShardsSpec] = DEFAULT_NUM_SHARDS
+
+  def resolve(
+      self, default_key_fn: Callable[[int], Callable[[Any], Any]]
+  ) -> 'PrefixBatchingConfig':
+    """Returns a copy whose ``key_fn`` is set.
+
+    Args:
+      default_key_fn: Factory (e.g. :class:`PromptPrefixKeyFn`) invoked with
+        ``prefix_char_len`` when ``key_fn`` is unset.
+    """
+    if self.key_fn is not None:
+      return self
+    return dataclasses.replace(
+        self, key_fn=default_key_fn(self.prefix_char_len))
+
+  def batching_kwargs(self,
+                      max_batch_duration_secs: Optional[float] = None
+                      ) -> dict[str, Any]:
+    """Returns the ``ModelHandler`` batching kwargs for a resolved config.
+
+    Args:
+      max_batch_duration_secs: The handler's own buffering duration; when
+        unset, ``DEFAULT_KEYED_BATCH_DURATION_SECS`` is used so low-traffic
+        keys still flush in streaming.
+    """
+    if self.key_fn is None:
+      raise ValueError(
+          'PrefixBatchingConfig.key_fn is unset; call resolve(default_key_fn) '
+          'first.')
+    return {
+        'batch_key_fn': self.key_fn,
+        'batch_key_num_shards': self.num_shards,
+        'max_batch_duration_secs': (
+            DEFAULT_KEYED_BATCH_DURATION_SECS
+            if max_batch_duration_secs is None else max_batch_duration_secs),
+    }
+
+
+def resolve_batching_kwargs(
+    config: Optional[PrefixBatchingConfig],
+    default_key_fn: Callable[[int], Callable[[Any], Any]],
+    max_batch_duration_secs: Optional[float] = None,
+) -> dict[str, Any]:
+  """Batching kwargs for an optional config; unkeyed when ``config`` is None.
+
+  Convenience for handlers: always yields ``max_batch_duration_secs`` and, when
+  prefix batching is on, ``batch_key_fn`` / ``batch_key_num_shards`` too, so
+  the result can be splatted straight into ``ModelHandler.__init__``.
+  """
+  if config is None:
+    return {'max_batch_duration_secs': max_batch_duration_secs}
+  return config.resolve(default_key_fn).batching_kwargs(max_batch_duration_secs)