[RunInference] Add content-aware dynamic batching via element_size_fn (Issue #37414)

Author: Eliaaazzz

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

@@ -301,6 +301,22 @@ def with_postprocess_fn(
     inference result in order from first applied to last applied."""
     return _PostProcessingModelHandler(self, fn)
 
+  def with_element_size_fn(
+      self, fn: Callable[[Union[ExampleT, tuple[KeyT, ExampleT]]], int]
+  ) -> 'ModelHandler[ExampleT, PredictionT, ModelT]':
+    """Returns a new `ModelHandler` that uses `fn` for element sizing.
+
+    The provided sizing function is passed through to `beam.BatchElements`
+    via `batch_elements_kwargs` as `element_size_fn`.
+
+    Args:
+      fn: A callable that returns the size (as an `int`) for each element.
+
+    Returns:
+      A `ModelHandler` wrapping this handler, with `fn` used for batching.
+    """
+    return _SizingModelHandler(self, fn)
+
   def with_no_batching(
       self
   ) -> """ModelHandler[Union[
@@ -1287,6 +1303,85 @@ def get_postprocess_fns(self) -> Iterable[Callable[[Any], Any]]:
     return self._base.get_postprocess_fns() + [self._postprocess_fn]
 
 
+class _SizingModelHandler(Generic[ExampleT, PredictionT, ModelT],
+                          ModelHandler[ExampleT, PredictionT, ModelT]):
+  def __init__(
+      self,
+      base: ModelHandler[ExampleT, PredictionT, ModelT],
+      element_size_fn: Callable[[Union[ExampleT, tuple[KeyT, ExampleT]]], int]):
+    """A ModelHandler that has an element_size_fn associated with it.
+
+    Args:
+      base: An implementation of the underlying model handler.
+      element_size_fn: the element sizing function to use for batching.
+    """
+    self._base = base
+    self._env_vars = getattr(base, '_env_vars', {})
+    self._element_size_fn = element_size_fn
+
+  def set_environment_vars(self):
+    return self._base.set_environment_vars()
+
+  def load_model(self) -> ModelT:
+    return self._base.load_model()
+
+  def run_inference(
+      self,
+      batch: Sequence[Union[ExampleT, tuple[KeyT, ExampleT]]],
+      model: ModelT,
+      inference_args: Optional[dict[str, Any]] = None
+  ) -> Union[Iterable[PredictionT], Iterable[tuple[KeyT, PredictionT]]]:
+    return self._base.run_inference(batch, model, inference_args)
+
+  def get_num_bytes(
+      self, batch: Sequence[Union[ExampleT, tuple[KeyT, ExampleT]]]) -> int:
+    return self._base.get_num_bytes(batch)
+
+  def get_metrics_namespace(self) -> str:
+    return self._base.get_metrics_namespace()
+
+  def get_resource_hints(self):
+    return self._base.get_resource_hints()
+
+  def batch_elements_kwargs(self):
+    kwargs = dict(self._base.batch_elements_kwargs())
+    kwargs["element_size_fn"] = self._element_size_fn
+    return kwargs
+
+  def validate_inference_args(self, inference_args: Optional[dict[str, Any]]):
+    return self._base.validate_inference_args(inference_args)
+
+  def update_model_path(self, model_path: Optional[str] = None):
+    return self._base.update_model_path(model_path=model_path)
+
+  def update_model_paths(
+      self,
+      model: ModelT,
+      model_paths: Optional[Union[str, list[KeyModelPathMapping]]] = None):
+    return self._base.update_model_paths(model, model_paths)
+
+  def get_preprocess_fns(self) -> Iterable[Callable[[Any], Any]]:
+    return self._base.get_preprocess_fns()
+
+  def should_skip_batching(self) -> bool:
+    return self._base.should_skip_batching()
+
+  def share_model_across_processes(self) -> bool:
+    return self._base.share_model_across_processes()
+
+  def model_copies(self) -> int:
+    return self._base.model_copies()
+
+  def override_metrics(self, metrics_namespace: str = '') -> bool:
+    return self._base.override_metrics(metrics_namespace=metrics_namespace)
+
+  def should_garbage_collect_on_timeout(self) -> bool:
+    return self._base.should_garbage_collect_on_timeout()
+
+  def get_postprocess_fns(self) -> Iterable[Callable[[Any], Any]]:
+    return self._base.get_postprocess_fns()
+
+
 class RunInference(beam.PTransform[beam.PCollection[Union[ExampleT,
                                                           Iterable[ExampleT]]],
                                    beam.PCollection[PredictionT]]):

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

@@ -2133,5 +2133,80 @@ def request(self, batch, model, inference_args=None):
       model_handler.run_inference([1], FakeModel())
 
 
+class FakeModelHandlerForSizing(base.ModelHandler[int, int, FakeModel]):
+  """A ModelHandler used to test element sizing behavior."""
+  def __init__(self, max_batch_size: int = 10):
+    self._max_batch_size = max_batch_size
+
+  def load_model(self) -> FakeModel:
+    return FakeModel()
+
+  def run_inference(self, batch, model, inference_args=None):
+    return [model.predict(x) for x in batch]
+
+  def batch_elements_kwargs(self):
+    return {'max_batch_size': self._max_batch_size}
+
+
+class RunInferenceSizeTest(unittest.TestCase):
+  """Tests for ModelHandler.with_element_size_fn."""
+  def test_kwargs_are_passed_correctly(self):
+    """Adds element_size_fn without clobbering existing kwargs."""
+    def size_fn(x):
+      return 10
+
+    base_handler = FakeModelHandlerForSizing(max_batch_size=20)
+    sized_handler = base_handler.with_element_size_fn(size_fn)
+
+    kwargs = sized_handler.batch_elements_kwargs()
+
+    self.assertEqual(kwargs['max_batch_size'], 20)
+    self.assertIn('element_size_fn', kwargs)
+    self.assertEqual(kwargs['element_size_fn'](1), 10)
+
+  def test_element_size_fn_wrapper_delegates_correctly(self):
+    """_SizingModelHandler delegates methods to the base handler."""
+    base_handler = FakeModelHandlerForSizing()
+    size_fn = lambda x: x * 2
+    sized_handler = base_handler.with_element_size_fn(size_fn)
+
+    model = sized_handler.load_model()
+    self.assertIsInstance(model, FakeModel)
+
+    result = list(sized_handler.run_inference([1, 2], model))
+    expected = [2, 3]  # FakeModel.predict(x) = x + 1
+    self.assertEqual(result, expected)
+
+    self.assertEqual(sized_handler.get_metrics_namespace(), 'RunInference')
+
+  def test_multiple_wrappers_can_be_chained(self):
+    """Sizing can be chained with other ModelHandler wrappers."""
+    base_handler = FakeModelHandlerForSizing()
+    preprocess_fn = lambda x: x * 10
+    size_fn = lambda x: 5
+
+    chained_handler = (
+        base_handler.with_preprocess_fn(preprocess_fn).with_element_size_fn(
+            size_fn))
+
+    kwargs = chained_handler.batch_elements_kwargs()
+    self.assertIn('element_size_fn', kwargs)
+    self.assertEqual(kwargs['element_size_fn'](1), 5)
+
+  def test_sizing_with_edge_cases(self):
+    """Allows extreme values from element_size_fn."""
+    base_handler = FakeModelHandlerForSizing(max_batch_size=1)
+
+    zero_size_fn = lambda x: 0
+    sized_handler = base_handler.with_element_size_fn(zero_size_fn)
+    kwargs = sized_handler.batch_elements_kwargs()
+    self.assertEqual(kwargs['element_size_fn'](999), 0)
+
+    large_size_fn = lambda x: 1000000
+    sized_handler = base_handler.with_element_size_fn(large_size_fn)
+    kwargs = sized_handler.batch_elements_kwargs()
+    self.assertEqual(kwargs['element_size_fn'](1), 1000000)
+
+
 if __name__ == '__main__':
   unittest.main()