From 18b0e9426051e8430ab57087c23ed0d6321ae4a6 Mon Sep 17 00:00:00 2001 From: Valentyn Tymofieiev Date: Wed, 14 Oct 2026 19:37:15 -0700 Subject: [PATCH 1/2] [Python] Support max_batch_weight in BatchElements and honor element_size_fn in stateful batching * Add max_batch_weight to BatchElements. When set, the batch size limits count elements and element_size_fn gives the weight of an element. This fixes RunInference failing with a TypeError when a ModelHandler sets max_batch_weight. * Honor element_size_fn when max_batch_duration_secs is set. This uses a separate stateful DoFn that only adds a 'weight' state and keeps the transform label, so pipelines that don't set element_size_fn are unaffected. * Never emit an empty batch when an element exceeds the limits on its own. * Warn when length_fn is set without max_batch_duration_secs, since it is ignored. --- CHANGES.md | 2 + sdks/python/apache_beam/transforms/util.py | 272 ++++++++++++++++-- .../apache_beam/transforms/util_test.py | 174 ++++++++++- 3 files changed, 429 insertions(+), 19 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index f98790bc2e19..c260a16cbadc 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -109,6 +109,8 @@ * (Python) `ReadFromText` with `escapechar` no longer skips a delimiter that directly follows an escaped delimiter, which merged two records into one ([#40459](https://github.com/apache/beam/issues/40459)). * (Python) DataFrame API: `agg` with a list that includes `'count'` or `'size'` now labels those rows `count`/`size` instead of `sum` ([#40504](https://github.com/apache/beam/issues/40504)). * (Python) DataFrame API: `Series.value_counts(bins=...)` no longer always fails with a `TypeError` ([#40503](https://github.com/apache/beam/issues/40503)). +* (Python) Setting `max_batch_weight` on a RunInference `ModelHandler` no longer fails with a `TypeError`: `BatchElements` now supports `max_batch_weight`, which limits the total `element_size_fn` of a batch in addition to its number of elements ([#ISSUE](https://github.com/apache/beam/issues/ISSUE)). +* (Python) `BatchElements` now honors `element_size_fn` when `max_batch_duration_secs` is set, and no longer emits an empty batch before an element that exceeds the batch size on its own ([#ISSUE](https://github.com/apache/beam/issues/ISSUE)). * Fixed X (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). ## Security Fixes diff --git a/sdks/python/apache_beam/transforms/util.py b/sdks/python/apache_beam/transforms/util.py index c8c7cc86d718..2bce3ee61b74 100644 --- a/sdks/python/apache_beam/transforms/util.py +++ b/sdks/python/apache_beam/transforms/util.py @@ -759,28 +759,57 @@ def stats(self): self._data) +def _unit_element_size(unused_element): + return 1 + + +def _element_size_and_weight(element, element_size_fn, max_batch_weight): + """Returns the contribution of an element to the size and weight of a batch. + + Without max_batch_weight, element_size_fn determines the size of a batch. + With max_batch_weight, the size of a batch is its number of elements, and + element_size_fn determines its weight. + """ + if max_batch_weight is None: + return element_size_fn(element), 0 + return 1, element_size_fn(element) + + class _GlobalWindowsBatchingDoFn(DoFn): - def __init__(self, batch_size_estimator, element_size_fn): + def __init__( + self, batch_size_estimator, element_size_fn, max_batch_weight=None): self._batch_size_estimator = batch_size_estimator self._element_size_fn = element_size_fn + self._max_batch_weight = max_batch_weight def start_bundle(self): self._batch = [] self._running_batch_size = 0 + self._running_batch_weight = 0 self._target_batch_size = self._batch_size_estimator.next_batch_size() # The first emit often involves non-trivial setup. self._batch_size_estimator.ignore_next_timing() def process(self, element): - element_size = self._element_size_fn(element) - if self._running_batch_size + element_size > self._target_batch_size: + element_size, element_weight = _element_size_and_weight( + element, self._element_size_fn, self._max_batch_weight) + would_exceed_size = ( + self._running_batch_size + element_size > self._target_batch_size) + would_exceed_weight = ( + self._max_batch_weight is not None and + self._running_batch_weight + element_weight > self._max_batch_weight) + # Never emit an empty batch: an element that exceeds the limits on its own + # is emitted in a batch by itself. + if self._batch and (would_exceed_size or would_exceed_weight): with self._batch_size_estimator.record_time(self._running_batch_size): yield window.GlobalWindows.windowed_value_at_end_of_window(self._batch) self._batch = [] self._running_batch_size = 0 + self._running_batch_weight = 0 self._target_batch_size = self._batch_size_estimator.next_batch_size() self._batch.append(element) self._running_batch_size += element_size + self._running_batch_weight += element_weight def finish_bundle(self): if self._batch: @@ -797,15 +826,18 @@ class _SizedBatch(): def __init__(self): self.elements = [] self.size = 0 + self.weight = 0 class _WindowAwareBatchingDoFn(DoFn): _MAX_LIVE_WINDOWS = 10 - def __init__(self, batch_size_estimator, element_size_fn): + def __init__( + self, batch_size_estimator, element_size_fn, max_batch_weight=None): self._batch_size_estimator = batch_size_estimator self._element_size_fn = element_size_fn + self._max_batch_weight = max_batch_weight def start_bundle(self): self._batches = collections.defaultdict(_SizedBatch) @@ -815,8 +847,15 @@ def start_bundle(self): def process(self, element, window=DoFn.WindowParam): batch = self._batches[window] - element_size = self._element_size_fn(element) - if batch.size + element_size > self._target_batch_size: + element_size, element_weight = _element_size_and_weight( + element, self._element_size_fn, self._max_batch_weight) + would_exceed_size = batch.size + element_size > self._target_batch_size + would_exceed_weight = ( + self._max_batch_weight is not None and + batch.weight + element_weight > self._max_batch_weight) + # Never emit an empty batch: an element that exceeds the limits on its own + # is emitted in a batch by itself. + if batch.elements and (would_exceed_size or would_exceed_weight): with self._batch_size_estimator.record_time(batch.size): yield windowed_value.WindowedValue( batch.elements, window.max_timestamp(), (window, )) @@ -825,6 +864,7 @@ def process(self, element, window=DoFn.WindowParam): self._batches[window].elements.append(element) self._batches[window].size += element_size + self._batches[window].weight += element_weight if len(self._batches) > self._MAX_LIVE_WINDOWS: window, batch = max( @@ -947,6 +987,173 @@ def flush_batch( return _StatefulBatchElementsDoFn() +def _pardo_stateful_weighted_batch_elements( + input_coder: coders.Coder, + batch_size_estimator: _BatchSizeEstimator, + max_buffering_duration_secs: int, + element_size_fn: Callable[[Any], Any], + max_batch_weight: Optional[int] = None, + clock=time.time): + """Like _pardo_stateful_batch_elements, but honors element_size_fn. + + Without max_batch_weight, element_size_fn determines the contribution of each + element to the batch size, as in the non-stateful implementation. With + max_batch_weight, the batch size is the number of elements, and + max_batch_weight limits the total element_size_fn of a batch. + + This DoFn only adds a 'weight' state to the state and timers of the DoFn of + _pardo_stateful_batch_elements, and keeps its transform label, to minimize the + impact on the update compatibility of streaming pipelines. Pipelines that do + not set element_size_fn keep using _pardo_stateful_batch_elements. + """ + ELEMENT_STATE = BagStateSpec('values', input_coder) + COUNT_STATE = CombiningValueStateSpec('count', input_coder, CountCombineFn()) + # The sum of element_size_fn over the buffered elements. + WEIGHT_STATE = ReadModifyWriteStateSpec( + 'weight', coders.FastPrimitivesCoder()) + BATCH_SIZE_STATE = ReadModifyWriteStateSpec('batch_size', input_coder) + WINDOW_TIMER = TimerSpec('window_end', TimeDomain.WATERMARK) + BUFFERING_TIMER = TimerSpec('buffering_end', TimeDomain.REAL_TIME) + BATCH_ESTIMATOR_STATE = ReadModifyWriteStateSpec( + 'batch_estimator', coders.PickleCoder()) + + def batch_size(count, weight): + # The size of a batch, in the unit of the batch sizes of the estimator. + return weight if max_batch_weight is None else count + + def exceeds_limits(count, weight, target_size): + return batch_size(count, weight) > target_size or ( + max_batch_weight is not None and weight > max_batch_weight) + + def reaches_limits(count, weight, target_size): + return batch_size(count, weight) >= target_size or ( + max_batch_weight is not None and weight >= max_batch_weight) + + class _StatefulWeightedBatchElementsDoFn(DoFn): + def default_label(self): + # Matches the label of the DoFn of _pardo_stateful_batch_elements, so that + # setting element_size_fn doesn't change the transform names of a + # pipeline, which streaming pipeline updates rely on. + return '_StatefulBatchElementsDoFn' + + def process( + self, + element, + window=DoFn.WindowParam, + element_state=DoFn.StateParam(ELEMENT_STATE), + count_state=DoFn.StateParam(COUNT_STATE), + weight_state=DoFn.StateParam(WEIGHT_STATE), + batch_size_state=DoFn.StateParam(BATCH_SIZE_STATE), + batch_estimator_state=DoFn.StateParam(BATCH_ESTIMATOR_STATE), + window_timer=DoFn.TimerParam(WINDOW_TIMER), + buffering_timer=DoFn.TimerParam(BUFFERING_TIMER)): + window_timer.set(window.end) + # Drop the fixed key since we don't care about it + value = element[1] + element_weight = element_size_fn(value) + target_size = batch_size_state.read() + # Should only happen on the first element + if target_size is None: + batch_estimator = batch_size_estimator + target_size = batch_estimator.next_batch_size() + batch_size_state.write(target_size) + batch_estimator_state.write(batch_estimator) + + count = count_state.read() + weight = weight_state.read() or 0 + if count and exceeds_limits( + count + 1, weight + element_weight, target_size): + # Emit the buffered elements first, so that adding this element + # doesn't make the batch exceed the limits. An element that exceeds + # the limits on its own is emitted in a batch by itself below. + yield from self.flush_batch( + element_state, + count_state, + weight_state, + batch_size_state, + batch_estimator_state, + buffering_timer) + count, weight = 0, 0 + target_size = batch_size_state.read() + + element_state.add(value) + count_state.add(1) + count += 1 + weight += element_weight + weight_state.write(weight) + + if count == 1 and max_buffering_duration_secs > 0: + # First element in batch, start buffering timer + buffering_timer.set(clock() + max_buffering_duration_secs) + + if reaches_limits(count, weight, target_size): + yield from self.flush_batch( + element_state, + count_state, + weight_state, + batch_size_state, + batch_estimator_state, + buffering_timer) + + @on_timer(WINDOW_TIMER) + def on_window_timer( + self, + element_state=DoFn.StateParam(ELEMENT_STATE), + count_state=DoFn.StateParam(COUNT_STATE), + weight_state=DoFn.StateParam(WEIGHT_STATE), + batch_size_state=DoFn.StateParam(BATCH_SIZE_STATE), + batch_estimator_state=DoFn.StateParam(BATCH_ESTIMATOR_STATE), + buffering_timer=DoFn.TimerParam(BUFFERING_TIMER)): + return self.flush_batch( + element_state, + count_state, + weight_state, + batch_size_state, + batch_estimator_state, + buffering_timer) + + @on_timer(BUFFERING_TIMER) + def on_buffering_timer( + self, + element_state=DoFn.StateParam(ELEMENT_STATE), + count_state=DoFn.StateParam(COUNT_STATE), + weight_state=DoFn.StateParam(WEIGHT_STATE), + batch_size_state=DoFn.StateParam(BATCH_SIZE_STATE), + batch_estimator_state=DoFn.StateParam(BATCH_ESTIMATOR_STATE), + buffering_timer=DoFn.TimerParam(BUFFERING_TIMER)): + return self.flush_batch( + element_state, + count_state, + weight_state, + batch_size_state, + batch_estimator_state, + buffering_timer) + + def flush_batch( + self, + element_state, + count_state, + weight_state, + batch_size_state, + batch_estimator_state, + buffering_timer): + batch = [element for element in element_state.read()] + if not batch: + return + size = batch_size(len(batch), weight_state.read() or 0) + element_state.clear() + count_state.clear() + weight_state.clear() + batch_estimator = batch_estimator_state.read() + with batch_estimator.record_time(size): + yield batch + batch_size_state.write(batch_estimator.next_batch_size()) + batch_estimator_state.write(batch_estimator) + buffering_timer.clear() + + return _StatefulWeightedBatchElementsDoFn() + + class SharedKey(): """A class that holds a per-process UUID used to key elements for streaming BatchElements. @@ -1048,7 +1255,15 @@ class BatchElements(PTransform): element_size_fn: (optional) A mapping of an element to its contribution to batch size, defaulting to every element having size 1. When provided, attempts to provide batches of optimal total size which may consist of - a varying number of elements. + a varying number of elements. If max_batch_weight is set, the size of a + batch is its number of elements instead, and element_size_fn gives the + weight of an element. + max_batch_weight: (optional) the largest total weight of a batch, where the + weight of an element is given by element_size_fn, which must be set. + This bounds batches by both their number of elements, between + min_batch_size and max_batch_size, and their total weight, e.g. their + number of tokens. An element that exceeds max_batch_weight on its own + is emitted in a batch by itself. variance: (optional) the permitted (relative) amount of deviation from the (estimated) ideal batch size used to produce a wider base for linear interpolation @@ -1059,7 +1274,8 @@ class BatchElements(PTransform): length_fn: (optional) a callable mapping an element to its length (int). When set together with bucket_boundaries, enables length-aware bucketed keying on the stateful path so that elements of similar length are - routed to the same batch, reducing padding waste. + routed to the same batch, reducing padding waste. Ignored, with a + warning, if max_batch_duration_secs is not set. bucket_boundaries: (optional) a sorted list of positive boundary values for length bucketing. Boundaries are lower-inclusive (bisect_right semantics): e.g., for boundaries [10, 50], buckets are (-inf, 10), @@ -1077,12 +1293,19 @@ def __init__( target_batch_duration_secs_including_fixed_cost=None, max_batch_duration_secs=None, *, - element_size_fn=lambda x: 1, + element_size_fn=None, + max_batch_weight=None, variance=0.25, clock=time.time, record_metrics=True, length_fn=None, bucket_boundaries=None): + 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.') + if max_batch_weight < 1: + raise ValueError( + f'max_batch_weight must be >= 1, got {max_batch_weight}') if bucket_boundaries is not None and length_fn is None: raise ValueError('bucket_boundaries requires length_fn to be set.') if bucket_boundaries is not None: @@ -1091,6 +1314,10 @@ def __init__( raise ValueError( 'bucket_boundaries must be a non-empty sorted list of ' 'positive values.') + if length_fn is not None and max_batch_duration_secs is None: + _LOGGER.warning( + 'BatchElements ignores length_fn and bucket_boundaries unless ' + 'max_batch_duration_secs is set.') self._batch_size_estimator = _BatchSizeEstimator( min_batch_size=min_batch_size, max_batch_size=max_batch_size, @@ -1102,6 +1329,7 @@ def __init__( clock=clock, record_metrics=record_metrics) self._element_size_fn = element_size_fn + self._max_batch_weight = max_batch_weight self._max_batch_dur = max_batch_duration_secs self._clock = clock self._length_fn = length_fn @@ -1120,22 +1348,32 @@ def expand(self, pcoll): self._length_fn, self._bucket_boundaries) else: keying_dofn = WithSharedKey() - return pcoll | ParDo(keying_dofn) | ParDo( - _pardo_stateful_batch_elements( - coder, - self._batch_size_estimator, - self._max_batch_dur, - self._clock)) + if self._element_size_fn is None: + batching_dofn = _pardo_stateful_batch_elements( + coder, self._batch_size_estimator, self._max_batch_dur, self._clock) + else: + batching_dofn = _pardo_stateful_weighted_batch_elements( + coder, + self._batch_size_estimator, + self._max_batch_dur, + self._element_size_fn, + self._max_batch_weight, + self._clock) + return pcoll | ParDo(keying_dofn) | ParDo(batching_dofn) elif pcoll.windowing.is_default(): # This is the same logic as _GlobalWindowsBatchingDoFn, but optimized # for that simpler case. return pcoll | ParDo( _GlobalWindowsBatchingDoFn( - self._batch_size_estimator, self._element_size_fn)) + self._batch_size_estimator, + self._element_size_fn or _unit_element_size, + self._max_batch_weight)) else: return pcoll | ParDo( _WindowAwareBatchingDoFn( - self._batch_size_estimator, self._element_size_fn)) + self._batch_size_estimator, + self._element_size_fn or _unit_element_size, + self._max_batch_weight)) class _SortAndBatchElementsDoFn(DoFn): diff --git a/sdks/python/apache_beam/transforms/util_test.py b/sdks/python/apache_beam/transforms/util_test.py index 1e3db695e074..4456fea18536 100644 --- a/sdks/python/apache_beam/transforms/util_test.py +++ b/sdks/python/apache_beam/transforms/util_test.py @@ -567,6 +567,86 @@ def test_sized_windowed_batches(self): 'a' * 7, ])) + def test_sized_batches_oversized_element(self): + # An element that exceeds the batch size on its own is emitted in a batch + # by itself, instead of after an empty batch. + with TestPipeline('FnApiRunner') as p: + res = ( + p + | beam.Create(['a' * 20, 'a', 'a'], reshuffle=False) + | util.BatchElements( + min_batch_size=10, max_batch_size=10, element_size_fn=len) + | beam.Map(lambda batch: tuple(len(e) for e in batch))) + assert_that(res, equal_to([(20, ), (1, 1)])) + + def test_sized_windowed_batches_oversized_element(self): + with TestPipeline('FnApiRunner') as p: + res = ( + p + | beam.Create(['a' * 20, 'a', 'a'], reshuffle=False) + | beam.Map(lambda e: window.TimestampedValue(e, 1)) + | beam.WindowInto(window.FixedWindows(10)) + | util.BatchElements( + min_batch_size=10, + max_batch_size=10, + element_size_fn=len, + clock=FakeClock()) + | beam.Map(lambda batch: tuple(len(e) for e in batch))) + assert_that(res, equal_to([(20, ), (1, 1)])) + + def test_max_batch_weight(self): + # Batches have at most 3 elements, with a total length of at most 10. + # The first batch of a bundle is limited by its number of elements, as the + # estimator replays its size as the next target. + with TestPipeline('FnApiRunner') as p: + res = ( + p + | beam.Create([1, 1, 1, 4, 4, 4, 1, 20, 1], reshuffle=False) + | beam.Map(lambda length: 'a' * length) + | util.BatchElements( + min_batch_size=3, + max_batch_size=3, + max_batch_weight=10, + element_size_fn=len) + | beam.Map(lambda batch: tuple(len(e) for e in batch))) + assert_that(res, equal_to([(1, 1, 1), (4, 4), (4, 1), (20, ), (1, )])) + + def test_max_batch_weight_windowed(self): + # As above, the first batch of the bundle is limited by its number of + # elements. + with TestPipeline('FnApiRunner') as p: + res = ( + p + | beam.Create([(1, 11), (1, 11), (1, 11), (1, 11), (4, 1), (4, 1), + (4, 1)], + reshuffle=False) + | beam.MapTuple( + lambda length, ts: window.TimestampedValue('a' * length, ts)) + | beam.WindowInto(window.FixedWindows(10)) + | util.BatchElements( + min_batch_size=3, + max_batch_size=3, + max_batch_weight=10, + element_size_fn=len, + clock=FakeClock()) + | beam.Map(lambda batch: tuple(len(e) for e in batch))) + assert_that( + res, + equal_to([ + (4, 4), + (4, ), # Elements in [0, 10) + (1, 1, 1), + (1, ), # Elements in [10, 20) + ])) + + def test_max_batch_weight_requires_element_size_fn(self): + with self.assertRaisesRegex(ValueError, 'requires element_size_fn'): + util.BatchElements(max_batch_weight=10) + + def test_max_batch_weight_must_be_positive(self): + with self.assertRaisesRegex(ValueError, 'must be >= 1'): + util.BatchElements(max_batch_weight=0, element_size_fn=len) + def test_target_duration(self): clock = FakeClock() batch_estimator = util._BatchSizeEstimator( @@ -879,6 +959,93 @@ def test_stateful_grows_to_max_batch(self): | beam.Map(len)) assert_that(res, equal_to([1, 1, 2, 4, 8, 16, 32, 50, 50])) + def test_stateful_sized_batches(self): + # The stateful implementation used to ignore element_size_fn. + with TestPipeline('FnApiRunner') as p: + res = ( + p + | beam.Create(['aaaaa'] * 6) + | util.BatchElements( + min_batch_size=10, + max_batch_size=10, + max_batch_duration_secs=100, + element_size_fn=len) + | beam.Map(len)) + assert_that(res, equal_to([2, 2, 2])) + + def test_stateful_max_batch_weight(self): + def batch_lengths(label, elements): + return ( + p + | f'Create{label}' >> beam.Create(elements) + | f'Batch{label}' >> util.BatchElements( + min_batch_size=3, + max_batch_size=3, + max_batch_duration_secs=100, + max_batch_weight=10, + element_size_fn=len) + | f'Len{label}' >> beam.Map(len)) + + with TestPipeline('FnApiRunner') as p: + # At most 2 elements of length 4 fit within max_batch_weight. + assert_that( + batch_lengths('Heavy', ['aaaa'] * 6), + equal_to([2, 2, 2]), + label='CheckHeavy') + # At most 3 elements fit within max_batch_size. + assert_that( + batch_lengths('Light', ['a'] * 7), + equal_to([3, 3, 1]), + label='CheckLight') + + def test_stateful_max_batch_weight_oversized_element(self): + with TestPipeline('FnApiRunner') as p: + res = ( + p + | beam.Create(['a' * 20, 'a', 'a']) + | util.BatchElements( + min_batch_size=3, + max_batch_size=3, + max_batch_duration_secs=100, + max_batch_weight=10, + element_size_fn=len) + | beam.Map(lambda batch: sorted(len(e) for e in batch))) + + def check(batches): + # Whatever the order of the elements, the oversized element is emitted + # in a batch by itself, and no batch is empty. + assert [20] in batches, batches + assert all(batches), batches + assert sorted(sum(batches, [])) == [1, 1, 20], batches + + assert_that(res, check) + + def test_stateful_element_size_fn_keeps_state_and_label(self): + # Setting element_size_fn only adds a 'weight' state, and keeps the + # transform names, to minimize the impact on the update compatibility of + # streaming pipelines. + def stateful_transforms(**kwargs): + p = beam.Pipeline() + batch_elements = util.BatchElements(max_batch_duration_secs=100, **kwargs) + _ = p | beam.Create(['a']) | 'Batch' >> batch_elements + transforms = {} + for transform in p.to_runner_api().components.transforms.values(): + if transform.spec.urn == common_urns.primitives.PAR_DO.urn: + payload = proto_utils.parse_Bytes( + transform.spec.payload, beam_runner_api_pb2.ParDoPayload) + if payload.state_specs: + transforms[transform.unique_name] = set(payload.state_specs) + return transforms + + state = {'values', 'count', 'batch_size', 'batch_estimator'} + label = 'Batch/ParDo(_StatefulBatchElementsDoFn)' + self.assertEqual(stateful_transforms(), {label: state}) + self.assertEqual( + stateful_transforms(element_size_fn=len), {label: state | {'weight'}}) + self.assertEqual( + stateful_transforms(element_size_fn=len, max_batch_weight=10), + {label: state | {'weight'}}) + def test_length_bucket_assignment(self): """WithLengthBucketKey assigns correct bucket indices.""" boundaries = [10, 50, 100] @@ -956,12 +1123,15 @@ def test_bucket_boundaries_must_be_positive(self): def test_length_fn_without_stateful_is_ignored(self): """length_fn without max_batch_duration_secs uses non-stateful path.""" + with self.assertLogs(util.__name__, level='WARNING') as logs: + batch_elements = util.BatchElements( + min_batch_size=3, max_batch_size=3, length_fn=len) + self.assertIn('ignores length_fn', logs.output[0]) with TestPipeline() as p: res = ( p | beam.Create(['a', 'bb', 'ccc']) - | util.BatchElements( - min_batch_size=3, max_batch_size=3, length_fn=len) + | batch_elements | beam.Map(len)) assert_that(res, equal_to([3])) From 80e18e04b3e4b3a929d6bf1ebe63902ad55b49d5 Mon Sep 17 00:00:00 2001 From: Valentyn Tymofieiev Date: Wed, 14 Oct 2026 19:48:44 -0700 Subject: [PATCH 2/2] Apply element_size_fn and batch_length_fn to examples in keyed RunInference handlers. --- CHANGES.md | 1 + sdks/python/apache_beam/ml/inference/base.py | 41 +++++++- .../apache_beam/ml/inference/base_test.py | 99 +++++++++++++++++++ 3 files changed, 138 insertions(+), 3 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index c260a16cbadc..256865905758 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -111,6 +111,7 @@ * (Python) DataFrame API: `Series.value_counts(bins=...)` no longer always fails with a `TypeError` ([#40503](https://github.com/apache/beam/issues/40503)). * (Python) Setting `max_batch_weight` on a RunInference `ModelHandler` no longer fails with a `TypeError`: `BatchElements` now supports `max_batch_weight`, which limits the total `element_size_fn` of a batch in addition to its number of elements ([#ISSUE](https://github.com/apache/beam/issues/ISSUE)). * (Python) `BatchElements` now honors `element_size_fn` when `max_batch_duration_secs` is set, and no longer emits an empty batch before an element that exceeds the batch size on its own ([#ISSUE](https://github.com/apache/beam/issues/ISSUE)). +* (Python) `KeyedModelHandler` and `MaybeKeyedModelHandler` now apply the `element_size_fn` and `batch_length_fn` of the wrapped `ModelHandler` to examples, instead of to `(key, example)` pairs ([#ISSUE](https://github.com/apache/beam/issues/ISSUE)). * Fixed X (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). ## Security Fixes diff --git a/sdks/python/apache_beam/ml/inference/base.py b/sdks/python/apache_beam/ml/inference/base.py index 7a19eba35117..361d05e0e90f 100644 --- a/sdks/python/apache_beam/ml/inference/base.py +++ b/sdks/python/apache_beam/ml/inference/base.py @@ -190,8 +190,13 @@ def __init__( max_batch_size: the maximum batch size to use when batching inputs. max_batch_duration_secs: the maximum amount of time to buffer a batch before emitting; used in streaming contexts. - max_batch_weight: the maximum weight of a batch. Requires element_size_fn. + max_batch_weight: the maximum total weight of a batch, where the weight + of an element is given by element_size_fn, which must be set. The + batch size limits (min_batch_size and max_batch_size) then count + elements. element_size_fn: a function that returns the size (weight) of an element. + Without max_batch_weight, the batch size limits are expressed in this + size instead of in number of elements. batch_length_fn: a callable mapping an element to its length (int). When set together with max_batch_duration_secs, enables length-aware bucketed keying so that elements of similar length are batched together, reducing @@ -674,6 +679,35 @@ def __init__( self.mh = mh +class _ApplyToExample: + """Applies a function of an example to a (key, example) pair. + + The element_size_fn and batch_length_fn of a ModelHandler take examples, but + BatchElements sees (key, example) pairs when the ModelHandler is wrapped in a + KeyedModelHandler or a MaybeKeyedModelHandler. + """ + def __init__(self, fn: Callable[[Any], Any], maybe_keyed: bool = False): + self._fn = fn + self._maybe_keyed = maybe_keyed + + def __call__(self, element): + # Like MaybeKeyedModelHandler.run_inference, consider tuples to be keyed. + if self._maybe_keyed and not isinstance(element, tuple): + return self._fn(element) + return self._fn(element[1]) + + +def _keyed_batch_elements_kwargs( + batch_elements_kwargs: Mapping[str, Any], + maybe_keyed: bool = False) -> dict[str, Any]: + """Adapts the BatchElements kwargs of a ModelHandler to keyed examples.""" + kwargs = dict(batch_elements_kwargs) + for name in ('element_size_fn', 'length_fn'): + if kwargs.get(name) is not None: + kwargs[name] = _ApplyToExample(kwargs[name], maybe_keyed) + return kwargs + + class KeyedModelHandler(Generic[KeyT, ExampleT, PredictionT, ModelT], ModelHandler[tuple[KeyT, ExampleT], tuple[KeyT, PredictionT], @@ -929,7 +963,7 @@ def get_resource_hints(self): def batch_elements_kwargs(self): if self._single_model: - return self._unkeyed.batch_elements_kwargs() + return _keyed_batch_elements_kwargs(self._unkeyed.batch_elements_kwargs()) return {} def validate_inference_args(self, inference_args: Optional[dict[str, Any]]): @@ -1132,7 +1166,8 @@ def get_resource_hints(self): return self._unkeyed.get_resource_hints() def batch_elements_kwargs(self): - return self._unkeyed.batch_elements_kwargs() + return _keyed_batch_elements_kwargs( + self._unkeyed.batch_elements_kwargs(), maybe_keyed=True) def validate_inference_args(self, inference_args: Optional[dict[str, Any]]): return self._unkeyed.validate_inference_args(inference_args) diff --git a/sdks/python/apache_beam/ml/inference/base_test.py b/sdks/python/apache_beam/ml/inference/base_test.py index a8ffc81e5acf..1085c35882df 100644 --- a/sdks/python/apache_beam/ml/inference/base_test.py +++ b/sdks/python/apache_beam/ml/inference/base_test.py @@ -2395,6 +2395,105 @@ def test_run_inference_with_length_aware_batch_elements(self): assert_that(results, equal_to(['a:2', 'bb:2', 'cccccc:7', 'ddddddd:7'])) +class BatchLengthsModelHandler(base.ModelHandler[str, tuple, FakeModel]): + """Predicts the lengths of the examples in the batch of each example.""" + def load_model(self): + return FakeModel() + + def run_inference(self, batch, model, inference_args=None): + lengths = tuple(len(example) for example in batch) + return [lengths] * len(batch) + + +class RunInferenceWeightedBatchingTest(unittest.TestCase): + """End-to-end tests for max_batch_weight and element_size_fn.""" + def _handler(self, **kwargs): + # At most 2 examples of length 4 fit within max_batch_weight. + return BatchLengthsModelHandler( + min_batch_size=3, + max_batch_size=3, + max_batch_weight=10, + element_size_fn=len, + **kwargs) + + def test_max_batch_weight(self): + with TestPipeline('FnApiRunner') as p: + results = ( + p + | beam.Create(['aaaa'] * 6, reshuffle=False) + | base.RunInference(self._handler())) + assert_that(results, equal_to([(4, 4)] * 6)) + + def test_max_batch_weight_with_max_batch_duration_secs(self): + with TestPipeline('FnApiRunner') as p: + results = ( + p + | beam.Create(['aaaa'] * 6) + | base.RunInference(self._handler(max_batch_duration_secs=60))) + assert_that(results, equal_to([(4, 4)] * 6)) + + def test_keyed_model_handler_sizes_examples(self): + # element_size_fn applies to examples, not to (key, example) pairs. + handler = base.KeyedModelHandler(self._handler()) + with TestPipeline('FnApiRunner') as p: + results = ( + p + | beam.Create([(i, 'aaaa') for i in range(6)], reshuffle=False) + | base.RunInference(handler) + | beam.Values()) + assert_that(results, equal_to([(4, 4)] * 6)) + + def test_maybe_keyed_model_handler_sizes_examples(self): + handler = base.MaybeKeyedModelHandler(self._handler()) + with TestPipeline('FnApiRunner') as p: + keyed = ( + p + | 'CreateKeyed' >> beam.Create([(i, 'aaaa') for i in range(6)], + reshuffle=False) + | 'RunInferenceKeyed' >> base.RunInference(handler) + | beam.Values()) + assert_that(keyed, equal_to([(4, 4)] * 6), label='CheckKeyed') + unkeyed = ( + p + | 'CreateUnkeyed' >> beam.Create(['aaaa'] * 6, reshuffle=False) + | 'RunInferenceUnkeyed' >> base.RunInference(handler)) + assert_that(unkeyed, equal_to([(4, 4)] * 6), label='CheckUnkeyed') + + def test_keyed_model_handler_buckets_examples(self): + # batch_length_fn applies to examples, not to (key, example) pairs. + handler = base.KeyedModelHandler( + PaddingReportingStringModelHandler( + min_batch_size=2, + max_batch_size=2, + max_batch_duration_secs=60, + batch_length_fn=len, + batch_bucket_boundaries=[5])) + examples = [('k1', 'a'), ('k2', 'cccccc'), ('k3', 'bb'), ('k4', 'ddddddd')] + with TestPipeline('FnApiRunner') as p: + results = ( + p + | beam.Create(examples, reshuffle=False) + | base.RunInference(handler)) + assert_that( + results, + equal_to([('k1', 'a:2'), ('k2', 'cccccc:7'), ('k3', 'bb:2'), + ('k4', 'ddddddd:7')])) + + def test_keyed_batch_elements_kwargs(self): + unkeyed = FakeModelHandlerForBatching( + min_batch_size=2, element_size_fn=len, batch_length_fn=len) + keyed_kwargs = base.KeyedModelHandler(unkeyed).batch_elements_kwargs() + self.assertEqual(keyed_kwargs['min_batch_size'], 2) + self.assertEqual(keyed_kwargs['element_size_fn'](('key', 'abc')), 3) + self.assertEqual(keyed_kwargs['length_fn'](('key', 'abc')), 3) + maybe_keyed_kwargs = base.MaybeKeyedModelHandler( + unkeyed).batch_elements_kwargs() + self.assertEqual(maybe_keyed_kwargs['element_size_fn'](('key', 'abc')), 3) + self.assertEqual(maybe_keyed_kwargs['element_size_fn']('abc'), 3) + # The kwargs of the unkeyed ModelHandler are unchanged. + self.assertIs(unkeyed.batch_elements_kwargs()['element_size_fn'], len) + + class SimpleFakeModelHandler(base.ModelHandler[int, int, FakeModel]): def load_model(self): return FakeModel()