damccorm commented on code in PR #40905:
URL: https://github.com/apache/beam/pull/40905#discussion_r4483928738
##########
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,
Review 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?
##########
sdks/python/apache_beam/transforms/util.py:
##########
@@ -1339,9 +1461,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]))
Review 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:
1. **Default `max_batch_size=10000` footgun:** `BatchElements` defaults
`max_batch_size` to `10000` (expecting `_BatchSizeEstimator` to start at
`min_batch_size=1` and adapt based on `target_batch_duration_secs`). If a user
enables `PrefixBatchingConfig()` on a `ModelHandler` without explicitly setting
`max_batch_size`, `GroupIntoBatches` uses a fixed batch size of `10,000` and
will buffer up to 10,000 LLM requests into a single `run_inference` call.
2. It silently ignores `_BatchSizeEstimator` parameters (`min_batch_size`,
`target_batch_overhead`, `target_batch_duration_secs`, `record_metrics`) and
cannot be combined with `element_size_fn` / `max_batch_weight` (which
`_pardo_stateful_weighted_batch_elements` already supports and is particularly
useful for LLM token-budget batching).
##########
sdks/python/apache_beam/ml/inference/base.py:
##########
@@ -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
Review 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`:
```python
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 kwargs
```
Right 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`?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
