diff --git a/CHANGES.md b/CHANGES.md index cd87a9609674..0c65e09baf8e 100644 --- a/CHANGES.md +++ b/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 diff --git a/sdks/python/apache_beam/ml/inference/base.py b/sdks/python/apache_beam/ml/inference/base.py index 4d5eb4f74c0a..ae4590692ec8 100644 --- a/sdks/python/apache_beam/ml/inference/base.py +++ b/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) diff --git a/sdks/python/apache_beam/ml/inference/base_test.py b/sdks/python/apache_beam/ml/inference/base_test.py index d4cfce8ec4f4..0c02829fb11e 100644 --- a/sdks/python/apache_beam/ml/inference/base_test.py +++ b/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.""" diff --git a/sdks/python/apache_beam/ml/inference/prefix_batching.py b/sdks/python/apache_beam/ml/inference/prefix_batching.py new file mode 100644 index 000000000000..21a6da0d8929 --- /dev/null +++ b/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. ``{'': 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) diff --git a/sdks/python/apache_beam/ml/inference/prefix_batching_test.py b/sdks/python/apache_beam/ml/inference/prefix_batching_test.py new file mode 100644 index 000000000000..d9eba644c2c5 --- /dev/null +++ b/sdks/python/apache_beam/ml/inference/prefix_batching_test.py @@ -0,0 +1,117 @@ +# +# 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. +# +import unittest + +from apache_beam.ml.inference import prefix_batching +from apache_beam.ml.inference.prefix_batching import ChatPrefixKeyFn +from apache_beam.ml.inference.prefix_batching import PrefixBatchingConfig +from apache_beam.ml.inference.prefix_batching import PromptPrefixKeyFn + + +class _Msg: + def __init__(self, role, content): + self.role = role + self.content = content + + +class PrefixBatchingTest(unittest.TestCase): + def test_prompt_prefix_key_fn(self): + key_fn = PromptPrefixKeyFn(prefix_char_len=5) + self.assertEqual(key_fn('abcdefghij'), 'abcde') + # Stops at the first blank line when it comes before prefix_char_len. + self.assertEqual(PromptPrefixKeyFn(prefix_char_len=20)('ab\n\ncdef'), 'ab') + self.assertEqual(PromptPrefixKeyFn(prefix_char_len=1)('ab\n\ncdef'), 'a') + with self.assertRaises(TypeError): + key_fn({'prompt': 'abcdefghij'}) + with self.assertRaises(ValueError): + PromptPrefixKeyFn(prefix_char_len=0) + + def test_chat_prefix_key_fn(self): + key_fn = ChatPrefixKeyFn(prefix_char_len=3) + msgs = [_Msg('system', 'You are A'), _Msg('user', 'Hello')] + self.assertEqual(key_fn(msgs), 'system:You') + # Dict-style messages and empty conversations are handled too. + self.assertEqual( + key_fn([{ + 'role': 'system', 'content': 'You are B' + }]), 'system:You') + self.assertEqual(key_fn([]), '') + with self.assertRaises(TypeError): + key_fn({'messages': msgs}) + with self.assertRaises(ValueError): + ChatPrefixKeyFn(prefix_char_len=0) + + def test_chat_prefix_key_fn_rejects_unknown_message_schema(self): + """Unknown schemas must fail loudly instead of all sharing one key.""" + key_fn = ChatPrefixKeyFn() + with self.assertRaisesRegex(TypeError, 'key_fn'): + key_fn([{'speaker': 'system', 'text': 'You are A'}]) + with self.assertRaisesRegex(TypeError, 'key_fn'): + key_fn(['You are A']) + # A message with a None content (e.g. an assistant tool call) is still a + # valid chat message. + self.assertEqual( + key_fn([{ + 'role': 'assistant', 'content': None + }]), 'assistant:') + + def test_config_resolve_and_batching_kwargs(self): + # Default: the handler's extractor is built with prefix_char_len. + cfg = PrefixBatchingConfig(prefix_char_len=32, num_shards=4) + resolved = cfg.resolve(PromptPrefixKeyFn) + self.assertIsInstance(resolved.key_fn, PromptPrefixKeyFn) + self.assertEqual(resolved.key_fn._prefix_char_len, 32) + self.assertEqual( + resolved.batching_kwargs(), + { + 'batch_key_fn': resolved.key_fn, + 'batch_key_num_shards': 4, + 'max_batch_duration_secs': prefix_batching. + DEFAULT_KEYED_BATCH_DURATION_SECS, + }) + # An explicit handler duration wins over the keyed default. + self.assertEqual( + resolved.batching_kwargs( + max_batch_duration_secs=2.5)['max_batch_duration_secs'], + 2.5) + # A custom key_fn is kept verbatim. + custom = PrefixBatchingConfig(key_fn=len).resolve(PromptPrefixKeyFn) + self.assertIs(custom.key_fn, len) + with self.assertRaises(ValueError): + PrefixBatchingConfig().batching_kwargs() + + def test_resolve_batching_kwargs_without_config(self): + # No config -> unkeyed batching with the handler's duration passed through. + self.assertEqual( + prefix_batching.resolve_batching_kwargs(None, PromptPrefixKeyFn, 3.0), + {'max_batch_duration_secs': 3.0}) + self.assertEqual( + prefix_batching.resolve_batching_kwargs(None, PromptPrefixKeyFn), + {'max_batch_duration_secs': None}) + kwargs = prefix_batching.resolve_batching_kwargs( + PrefixBatchingConfig(), ChatPrefixKeyFn) + self.assertIsInstance(kwargs['batch_key_fn'], ChatPrefixKeyFn) + # Prefix keys are sharded by default so a hot prefix is not served by a + # single worker. + self.assertEqual(kwargs['batch_key_num_shards'], 8) + self.assertEqual( + kwargs['batch_key_num_shards'], prefix_batching.DEFAULT_NUM_SHARDS) + self.assertEqual(kwargs['max_batch_duration_secs'], 1.0) + + +if __name__ == '__main__': + unittest.main() diff --git a/sdks/python/apache_beam/transforms/util.py b/sdks/python/apache_beam/transforms/util.py index df9607977950..448d479d28d8 100644 --- a/sdks/python/apache_beam/transforms/util.py +++ b/sdks/python/apache_beam/transforms/util.py @@ -33,6 +33,7 @@ import uuid from collections.abc import Callable from collections.abc import Iterable +from collections.abc import Mapping from typing import TYPE_CHECKING from typing import Any from typing import Optional @@ -1206,6 +1207,96 @@ def process(self, element): yield ((self.key, bucket), element) +def _validate_num_shards(num_shards: Any) -> Any: + """Validates a ``num_shards`` spec for ``WithBatchKey`` / ``BatchElements``. + + Accepts an int (uniform shard count), a mapping from batch key to shard count + (keys absent from the mapping get 1 shard), or a callable from batch key to + shard count. Returns the normalized spec (``None`` becomes ``1``). + """ + if num_shards is None: + return 1 + if isinstance(num_shards, bool): + raise ValueError('num_shards must be an int, mapping, or callable.') + if isinstance(num_shards, int): + if num_shards < 1: + raise ValueError(f'num_shards ({num_shards}) must be >= 1.') + return num_shards + if isinstance(num_shards, Mapping): + bad = { + k: v + for k, v in num_shards.items() + if not (isinstance(v, int) and not isinstance(v, bool) and v >= 1) + } + if bad: + raise ValueError( + f'num_shards mapping values must be ints >= 1; got {bad}.') + return dict(num_shards) + if callable(num_shards): + return num_shards + raise ValueError( + 'num_shards must be an int, a mapping from key to int, or a callable ' + f'from key to int; got {type(num_shards).__name__}.') + + +class WithBatchKey(DoFn): + """Keys elements with a deterministic string key for cross-worker + GroupIntoBatches batching. + + Unlike ``WithSharedKey`` and ``WithLengthBucketKey`` (which include a + per-worker UUID so elements remain worker-local), ``WithBatchKey`` omits any + worker UUID so that elements sharing the same key across all upstream workers + shuffle to the same downstream worker and are grouped into homogeneous + batches (e.g., for LLM prefix caching). + + ``num_shards`` bounds hot-key fanout: when a key resolves to ``n > 1`` + shards, a random suffix ``:`` in ``[0, n)`` is appended so that key is + spread across at most ``n`` workers. It may be + + * an ``int`` -- the same shard count for every key; + * a ``Mapping[key, int]`` -- per-key shard counts for a known skewed + distribution (e.g. ``{'default_system_prompt': 8}``); keys not in the + mapping get 1 shard; + * a ``Callable[[key], int]`` -- computed per key (e.g. from a frequency + table loaded at ``setup``). + + The mapping / callable is looked up with the raw ``key_fn`` result first and + its ``str()`` form second, so either may be used as the mapping key. + """ + def __init__(self, key_fn: Callable[[Any], Any], num_shards: Any = 1): + self._key_fn = key_fn + self._num_shards = _validate_num_shards(num_shards) + + def _shards_for(self, raw_key: Any, key: str) -> int: + spec = self._num_shards + if isinstance(spec, int): + return spec + if isinstance(spec, Mapping): + n = spec.get(raw_key) + if n is None: + n = spec.get(key, 1) + return n + n = spec(raw_key) + if not isinstance(n, int) or n < 1: + raise ValueError( + f'num_shards callable returned {n!r} for key {key!r}; expected an ' + 'int >= 1.') + return n + + def process(self, element): + raw_key = self._key_fn(element) + key = str(raw_key) if raw_key is not None else '' + num_shards = self._shards_for(raw_key, key) + if num_shards > 1: + key = f'{key}:{random.randrange(num_shards)}' + yield (key, element) + + +# Buffering applied by BatchElements(key_fn=...) when max_batch_duration_secs is +# unset, so incomplete per-key batches still flush in streaming pipelines. +_DEFAULT_KEYED_MAX_BUFFERING_SECS = 1.0 + + @typehints.with_input_types(T) @typehints.with_output_types(list[T]) class BatchElements(PTransform): @@ -1283,6 +1374,21 @@ class BatchElements(PTransform): semantics): e.g., for boundaries [10, 50], buckets are (-inf, 10), [10, 50), [50, inf). Defaults to [16, 32, 64, 128, 256, 512] when length_fn is set. Requires length_fn. + key_fn: (optional) a callable mapping an element to a grouping key. When + set, elements are keyed via ``WithBatchKey`` (without a per-worker UUID) + and grouped across workers using ``GroupIntoBatches`` up to + ``max_batch_size``, producing homogeneous batches per key (e.g., for LLM + prefix caching). Incomplete batches are flushed after + ``max_batch_duration_secs``, or after 1 second when it is unset, so + low-traffic keys still make progress in streaming pipelines. + num_shards: (optional) bounds hot-key fanout when ``key_fn`` is set: an + ``int`` applied to every key, a ``Mapping[key, int]`` giving per-key + shard counts for a known skewed distribution (unlisted keys get 1), or + a ``Callable[[key], int]``. 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, so a hot key's throughput does not grow with + the number of workers; give hot keys more shards when the key + distribution is skewed. See ``WithBatchKey``. """ _DEFAULT_BUCKET_BOUNDARIES = [16, 32, 64, 128, 256, 512] @@ -1301,7 +1407,9 @@ def __init__( clock=time.time, record_metrics=True, length_fn=None, - bucket_boundaries=None): + bucket_boundaries=None, + key_fn=None, + num_shards=1): if max_batch_weight is not None: if element_size_fn is None: raise ValueError('max_batch_weight requires element_size_fn to be set.') @@ -1320,6 +1428,11 @@ def __init__( _LOGGER.warning( 'BatchElements ignores length_fn and bucket_boundaries unless ' 'max_batch_duration_secs is set.') + if key_fn is not None and not callable(key_fn): + raise ValueError('key_fn must be callable.') + num_shards = _validate_num_shards(num_shards) + if num_shards != 1 and key_fn is None: + raise ValueError('num_shards requires key_fn to be set.') self._batch_size_estimator = _BatchSizeEstimator( min_batch_size=min_batch_size, max_batch_size=max_batch_size, @@ -1339,9 +1452,29 @@ def __init__( self._bucket_boundaries = self._DEFAULT_BUCKET_BOUNDARIES else: self._bucket_boundaries = bucket_boundaries + self._key_fn = key_fn + self._num_shards = num_shards def expand(self, pcoll): - if getattr(pcoll.pipeline.runner, 'is_streaming', False): + 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])) + elif getattr(pcoll.pipeline.runner, 'is_streaming', False): raise NotImplementedError("Requires stateful processing (BEAM-2687)") elif self._max_batch_dur is not None: coder = coders.registry.get_coder(pcoll) diff --git a/sdks/python/apache_beam/transforms/util_test.py b/sdks/python/apache_beam/transforms/util_test.py index 4456fea18536..d6d60b9896ff 100644 --- a/sdks/python/apache_beam/transforms/util_test.py +++ b/sdks/python/apache_beam/transforms/util_test.py @@ -1281,6 +1281,109 @@ def check_batch_homogeneity(batch): checks = batches | beam.Map(check_batch_homogeneity) assert_that(checks, is_not_empty()) + def test_batch_elements_key_fn_groups_homogeneous_batches(self): + """BatchElements with key_fn groups elements by key via GroupIntoBatches.""" + elements = ['sysA:q1', 'sysB:q1', 'sysA:q2', 'sysB:q2', 'sysA:q3'] + with TestPipeline('FnApiRunner') as p: + batches = ( + p + | beam.Create(elements, reshuffle=False) + | util.BatchElements( + min_batch_size=1, + max_batch_size=2, + max_batch_duration_secs=1.0, + key_fn=lambda s: s.split(':', 1)[0])) + + def check_homogeneous(batch_list): + flattened = [] + for b in batch_list: + prefixes = {s.split(':', 1)[0] for s in b} + assert len(prefixes) == 1, f'Mixed prefixes in batch: {b}' + assert len(b) <= 2, f'Batch exceeded max_batch_size: {b}' + flattened.extend(b) + assert sorted(flattened) == sorted(elements) + + assert_that(batches, check_homogeneous) + + def test_with_batch_key_num_shards(self): + """WithBatchKey appends a bounded random shard suffix to the key.""" + dofn = util.WithBatchKey(key_fn=lambda d: d['cat'], num_shards=4) + element = {'cat': 'math', 'q': '1+1'} + out = list(dofn.process(element)) + self.assertEqual(len(out), 1) + key, val = out[0] + self.assertIs(val, element) + self.assertTrue(key.startswith('math:')) + shard = int(key.split(':', 1)[1]) + self.assertIn(shard, range(4)) + + def test_with_batch_key_num_shards_mapping_handles_skew(self): + """A per-key shard mapping fans out only the listed hot keys.""" + dofn = util.WithBatchKey( + key_fn=lambda d: d['cat'], num_shards={ + 'hot': 4, 'warm': 2 + }) + hot_keys = {list(dofn.process({'cat': 'hot'}))[0][0] for _ in range(200)} + warm_keys = {list(dofn.process({'cat': 'warm'}))[0][0] for _ in range(200)} + cold_keys = {list(dofn.process({'cat': 'cold'}))[0][0] for _ in range(200)} + self.assertEqual(hot_keys, {f'hot:{i}' for i in range(4)}) + self.assertEqual(warm_keys, {'warm:0', 'warm:1'}) + # Keys absent from the mapping are not sharded at all. + self.assertEqual(cold_keys, {'cold'}) + + def test_with_batch_key_num_shards_mapping_matches_str_key(self): + """Mapping lookup falls back to the str() form of the key.""" + dofn = util.WithBatchKey(key_fn=lambda d: d['tenant'], num_shards={'7': 3}) + keys = {list(dofn.process({'tenant': 7}))[0][0] for _ in range(100)} + self.assertEqual(keys, {'7:0', '7:1', '7:2'}) + + def test_with_batch_key_num_shards_callable(self): + """A callable computes the shard count per key.""" + dofn = util.WithBatchKey( + key_fn=lambda s: s, num_shards=lambda k: 3 if k == 'hot' else 1) + hot_keys = {list(dofn.process('hot'))[0][0] for _ in range(100)} + self.assertEqual(hot_keys, {'hot:0', 'hot:1', 'hot:2'}) + self.assertEqual(list(dofn.process('cold')), [('cold', 'cold')]) + bad = util.WithBatchKey(key_fn=lambda s: s, num_shards=lambda k: 0) + with self.assertRaises(ValueError): + list(bad.process('x')) + + def test_batch_elements_key_fn_validation(self): + with self.assertRaises(ValueError): + util.BatchElements(key_fn='not_callable') + with self.assertRaises(ValueError): + util.BatchElements(key_fn=len, num_shards=0) + with self.assertRaises(ValueError): + util.BatchElements(num_shards=2) + with self.assertRaises(ValueError): + util.BatchElements(num_shards={'a': 2}) + with self.assertRaises(ValueError): + util.BatchElements(key_fn=len, num_shards={'a': 0}) + with self.assertRaises(ValueError): + util.BatchElements(key_fn=len, num_shards='4') + # Valid specs. + util.BatchElements(key_fn=len, num_shards={'a': 2}) + util.BatchElements(key_fn=len, num_shards=lambda k: 2) + + def test_batch_elements_key_fn_default_buffering(self): + """Keyed batching never runs GroupIntoBatches without a flush timer.""" + def buffering_secs(**kwargs): + p = beam.Pipeline() + _ = p | beam.Create(['a']) | util.BatchElements(key_fn=len, **kwargs) + found = [] + + class Visitor(beam.pipeline.PipelineVisitor): + def enter_composite_transform(self, node): + if isinstance(node.transform, util.GroupIntoBatches): + found.append(node.transform.params.max_buffering_duration_secs) + + p.visit(Visitor()) + self.assertEqual(len(found), 1) + return found[0] + + self.assertEqual(buffering_secs(), 1.0) + self.assertEqual(buffering_secs(max_batch_duration_secs=2.5), 2.5) + def test_global_batching_dofn_single_vs_multiple_bundles(self): # This test directly verifies how bundling affects the batch sizes produced by # the internal _GlobalWindowsBatchingDoFn of BatchElements.