Eliaaazzz opened a new pull request, #40682: URL: https://github.com/apache/beam/pull/40682
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 -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
