[Eliaaazzz]

Finishes the max_batch_weight implementation for RunInference, starting from the commits in #40603 (thanks, @tvalentyn; they are kept as-is with their authorship), plus the issue number in CHANGES.md and two RunInference-level regression tests.

Addresses the problems reported in #40644:

• BatchElements now accepts max_batch_weight. When it is set, min_batch_size/max_batch_size count elements and element_size_fn gives the weight of an element; a batch is flushed before it would exceed either limit. Setting max_batch_weight on a ModelHandler no longer fails at pipeline construction with TypeError: BatchElements.__init__() got an unexpected keyword argument 'max_batch_weight'.
• The stateful path (max_batch_duration_secs) now honors element_size_fn and max_batch_weight. A separate DoFn adds only a weight state and keeps the _StatefulBatchElementsDoFn label, so pipelines that do not set element_size_fn keep using the existing DoFn and transform names stay the same (covered by test_stateful_element_size_fn_keeps_state_and_label).
• No more empty batches: an element that exceeds the limits on its own is emitted in a batch by itself, so run_inference never receives [].
• KeyedModelHandler and MaybeKeyedModelHandler apply the wrapped handler's element_size_fn and batch_length_fn to the example instead of the (key, example) pair.
• length_fn without max_batch_duration_secs logs a warning instead of being silently ignored.

Argument validation at the ModelHandler level (requiring the sizing functions when batch limits are set, rejecting mixed weight-based and length-based options) is being handled separately by @alisha-1000, as discussed in #40644. BatchElements itself raises a ValueError when max_batch_weight is set without element_size_fn or is smaller than 1.

Tests:

• transforms/util_test.py: weight limits on the non-stateful, windowed and stateful paths, oversized elements, the length_fn warning, and the state/label stability check.
• ml/inference/base_test.py: end-to-end RunInference with max_batch_weight (with and without max_batch_duration_secs), keyed and maybe-keyed handlers sizing and bucketing examples, the pipeline construction repro from the issue, and an oversized example reaching run_inference alone.
• Ran pytest apache_beam/transforms/util_test.py -k "Batch or Stateful" and pytest apache_beam/ml/inference/base_test.py -k "WeightedBatching or LengthAware or Keyed or batch" locally; yapf is clean.

Part of #40644.

R: @tvalentyn

[github-actions]

📋 Review Summary

This PR adds support for max_batch_weight to BatchElements across global, windowed, and stateful batching paths, honors element_size_fn on the stateful path while preserving DoFn labels and existing state names for streaming update compatibility, and fixes KeyedModelHandler / MaybeKeyedModelHandler to apply element_size_fn and length_fn to examples rather than (key, example) tuples. It also prevents empty batches from being emitted before oversized elements and adds comprehensive unit and integration tests.

🔍 General Feedback

• The separation of _StatefulWeightedBatchElementsDoFn with an overridden default_label() returning '_StatefulBatchElementsDoFn' is a clean design that preserves transform hierarchy and state specs for existing pipelines that do not set element_size_fn.
• Guarding batch flushes with if self._batch: / if batch.elements: / if count: cleanly prevents empty batches when an individual element exceeds the target size or weight limits.
• The _ApplyToExample wrapper in ml/inference/base.py is picklable and properly adapts both element_size_fn and length_fn without mutating the underlying unkeyed handler's kwargs.

---

💡 Tip: Comment /gemini review on this pull request to trigger a new Gemini code review.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 80.04%. Comparing base (76f6e6d) to head (a8d6afd).
⚠️ Report is 242 commits behind head on master.

Additional details and impacted files

@@              Coverage Diff              @@
##             master   #40682       +/-   ##
=============================================
+ Coverage     56.16%   80.04%   +23.87%     
=============================================
  Files          1124      490      -634     
  Lines        178336    82916    -95420     
  Branches       1489        0     -1489     
=============================================
- Hits         100170    66373    -33797     
+ Misses        75645    16543    -59102     
+ Partials       2521        0     -2521     

Flag	Coverage Δ
go	?
java	?
python	80.04% <100.00%> (+0.49%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[github-actions]

Assigning reviewers:

R: @damccorm for label python.
Backup reviewer: tvalentyn (if the main reviewer is unresponsive, you can escalate to them by tagging their username)

This pull request likely touches a core component ("core" label). Please review with scrutiny.

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

[tvalentyn]

Thanks, @Eliaaazzz !