[RunInference] Add content-aware dynamic batching via element_size_fn…

[Eliaaazzz]

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

Rationale

This PR addresses #37414 by introducing content-aware dynamic batching to RunInference.

Currently, RunInference relies on BatchElements with a strict count-based limit (max_batch_size). However, for workloads like NLP and LLMs, variable-length inputs (tokens) lead to significant variance in computational cost and memory usage. A batch of 10 short sentences is vastly different from a batch of 10 long documents.

This change allows users to provide a custom element_size_fn to ModelHandler, which is then passed down to the underlying BatchElements transform. This enables batching based on total "weight" (e.g., token count) rather than just element count, improving GPU utilization and preventing OOM errors.

Design Principles

This implementation prioritizes modularity and type safety through the following design choices:

• Decorator Pattern (Composition over Inheritance):
  Implemented _SizingModelHandler as a wrapper to dynamically attach sizing behavior to any ModelHandler implementation. This avoids the combinatorial explosion of subclasses (e.g., TFModelHandlerWithSizing, PyTorchModelHandlerWithSizing) and keeps the codebase DRY.

• Open-Closed Principle (OCP):
  The change is strictly additive. The base ModelHandler remains closed for modification, ensuring zero regression risk for existing pipelines. Functionality is extended purely by overriding batch_elements_kwargs in the wrapper and safely delegating all other methods to the base instance.

• Architectural Consistency:
  The implementation mirrors the existing _PreProcessingModelHandler pattern in Apache Beam. This ensures API consistency and reduces cognitive load for maintainers.

Changes

• apache_beam/ml/inference/base.py:
  • Added with_element_size_fn method to the ModelHandler interface.
  • Implemented _SizingModelHandler wrapper class.
  • Overrode batch_elements_kwargs to inject element_size_fn while preserving existing configuration (using safe dictionary copy).
  • Implemented full delegation for all ModelHandler methods (e.g., update_model_paths, get_metrics_namespace) to ensure transparency.

Usage Example

def token_counter(text: str) -> int:
    return len(text.split())

# Configure the handler to batch based on token count (e.g., max 1000 tokens per batch)
# The max_batch_size in BatchElements will now act as a limit on the sum of element sizes.
model_handler = MyModelHandler().with_element_size_fn(token_counter)

Testing

Comprehensive tests were added in sdks/python/apache_beam/ml/inference/base_test.py:

1. test_kwargs_are_passed_correctly

   • Verifies that element_size_fn is correctly injected into batch_elements_kwargs.

2. test_batch_elements_integration_with_beam_pipeline

   • Verifies run-time behavior.
   • Scenario: Input elements with weight 5, max_batch_size set to 10.
   • Result: BatchElements correctly creates batches of 2 elements (5+5=10), confirming the dynamic sizing logic.

3. test_element_size_fn_wrapper_delegates_correctly

   • Ensures all ModelHandler methods are properly delegated (Critical for features like model updates).

4. test_multiple_wrappers_can_be_chained

   • Verifies compatibility when chained with with_preprocess_fn.

---

Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:

[x] fixes #37414

[gemini-code-assist]

Summary of Changes

Hello @Eliaaazzz, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request enhances the RunInference component by introducing a flexible content-aware dynamic batching mechanism. Previously, batching was strictly count-based, which was inefficient for variable-length inputs common in NLP and LLM workloads. The new feature allows users to define a custom function (element_size_fn) to determine the "size" or "weight" of each element (e.g., token count), enabling RunInference to create batches that optimize resource utilization and prevent out-of-memory issues. The implementation follows a modular design using a wrapper pattern, ensuring that existing pipelines remain unaffected while providing powerful new capabilities.

Highlights

• Content-Aware Dynamic Batching: Introduced the ability to perform dynamic batching based on element size rather than just element count, addressing limitations with variable-length inputs in ML inference workloads like NLP and LLMs.
• element_size_fn Integration: Users can now provide a custom element_size_fn to ModelHandler, which is then passed to beam.BatchElements to define how elements are sized for batching, enabling optimization based on 'weight' (e.g., token count).
• Modular Design: Implemented using a _SizingModelHandler wrapper (decorator pattern) to extend ModelHandler functionality without modifying existing classes, ensuring backward compatibility, type safety, and adherence to the Open-Closed Principle.
• Comprehensive Testing: New unit tests were added to verify correct element_size_fn injection into batching arguments, proper delegation of ModelHandler methods by the wrapper, and compatibility when chaining multiple ModelHandler wrappers.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[github-actions]

Assigning reviewers:

R: @shunping for label python.

Note: If you would like to opt out of this review, comment assign to next reviewer.

Available commands:

• stop reviewer notifications - opt out of the automated review tooling
• remind me after tests pass - tag the comment author after tests pass
• waiting on author - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)

The PR bot will only process comments in the main thread (not review comments).

[Eliaaazzz]

Hi @damccorm, just a heads-up that the PR is ready and all CI checks have passed.

I've implemented the weighted BatchElements approach as discussed with Robert. Ready for review whenever you have a chance. Thanks!