[Python] Support max_batch_weight in BatchElements and honor element_size_fn in stateful batching

CHANGES.md

@@ -109,6 +109,9 @@
 * (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)).
+* (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

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)

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()