Repository navigation
[Bug]: [RunInference] Pipelines using ModelHandlers with max_batch_weight fail at construction time #40644
Description
What happened?
Repro (LLM generated):
"""Repro: max_batch_weight on a RunInference ModelHandler raises TypeError.
ModelHandler.__init__ accepts max_batch_weight and returns it from
batch_elements_kwargs(), but RunInference.expand passes those kwargs to
beam.BatchElements, which has no max_batch_weight parameter. The pipeline
fails at construction time, before any model is loaded or any data is
processed.
Run from sdks/python on a checkout without the fix:
PYTHONPATH=. python3 repro_max_batch_weight.py
"""
import apache_beam as beam
from apache_beam.ml.inference import base
from apache_beam.ml.inference.sklearn_inference import SklearnModelHandlerNumpy
class MinimalHandler(base.ModelHandler):
def load_model(self):
return None
def run_inference(self, batch, model, inference_args=None):
return batch
def build(name, handler):
print(f'--- {name} ---')
print('batch_elements_kwargs():', handler.batch_elements_kwargs())
try:
p = beam.Pipeline()
# Only constructs the pipeline: RunInference.expand runs here, the model
# is never loaded and the pipeline is never run.
_ = p | beam.Create(['a', 'bb']) | base.RunInference(handler)
print('OK: pipeline constructed')
except TypeError as e:
print(f'TypeError: {e}')
print('apache_beam', beam.__version__, 'from', beam.__file__)
# 1. Minimal custom handler: the base ModelHandler.__init__ argument.
build(
'custom ModelHandler',
MinimalHandler(max_batch_weight=10, element_size_fn=len))
# 2. Built-in handler: max_batch_weight is a documented constructor argument
# of every built-in handler. The model URI is never read.
build(
'SklearnModelHandlerNumpy',
SklearnModelHandlerNumpy(
model_uri='gs://unused/model.pkl',
max_batch_weight=10,
element_size_fn=len))
# 3. Control: the same handler without max_batch_weight constructs fine.
build('control (no max_batch_weight)', MinimalHandler(element_size_fn=len))
Output:
python repro_max_batch_weight.py
apache_beam 2.78.0.dev from /usr/local/google/home/valentyn/projects/beam/beam/beam/sdks/python/apache_beam/__init__.py
--- custom ModelHandler ---
batch_elements_kwargs(): {'max_batch_weight': 10, 'element_size_fn': <built-in function len>}
TypeError: BatchElements.__init__() got an unexpected keyword argument 'max_batch_weight'
--- SklearnModelHandlerNumpy ---
batch_elements_kwargs(): {'max_batch_weight': 10, 'element_size_fn': <built-in function len>}
TypeError: BatchElements.__init__() got an unexpected keyword argument 'max_batch_weight'
--- control (no max_batch_weight) ---
batch_elements_kwargs(): {'element_size_fn': <built-in function len>}
OK: pipeline constructed
Issue Priority
Priority: 2 (default / most bugs should be filed as P2)
Issue Components
- Component: Python SDK
- Component: Java SDK
- Component: Go SDK
- Component: Typescript SDK
- Component: IO connector
- Component: Beam YAML
- Component: Beam examples
- Component: Beam playground
- Component: Beam katas
- Component: Website
- Component: Infrastructure
- Component: Spark Runner
- Component: Flink Runner
- Component: Prism Runner
- Component: Twister2 Runner
- Component: Hazelcast Jet Runner
- Component: Google Cloud Dataflow Runner
Activity
[-][Bug]: [RunInference] Pipelines using ModelHandlers with max_batch_by_weight fail at construction time[/-][+][Bug]: [RunInference] Pipelines using ModelHandlers with max_batch_weight fail at construction time[/+]tvalentyn commentedon Oct 16, 2026
Hi @Eliaaazzz , do you know if we finished implementation of this functionality?
tvalentyn commentedon Oct 16, 2026
it seems publicly documented so we should either fix it up + add a test or remove, hopefully it's a regression that we somehow didn't catch.
tvalentyn commentedon Oct 16, 2026
Couple other issues identified by an LLM (I haven't looked in detail but seems convincing):
BatchElements emits empty batches
Symptom: when a single element's size exceeds the current target batch size, BatchElements first emits an empty batch [], then the element.
Cause: both non-stateful DoFns flush when the running size plus the next element's size exceeds the target, without checking whether the batch is empty.
Impact: RunInference calls run_inference([]) on the model handler. Handlers that index batch[0], such as MaybeKeyedModelHandler, or stack tensors can crash. This gets more likely once weight limits work (bug 1), because one long input is enough to trigger it.
Keyed handlers pass (key, example) pairs to element_size_fn and batch_length_fn
Symptom: with KeyedModelHandler or MaybeKeyedModelHandler, the wrapped handler's functions receive a (key, example) tuple instead of the example. For example, element_size_fn=len returns 2 for every element.
Cause: both wrappers forward the inner handler's batch_elements_kwargs() unchanged, but BatchElements runs on the keyed elements.
Impact: size limits and length bucketing are silently wrong for keyed inference, a common pattern for joining predictions back to inputs. Two 60-character examples were batched together despite a weight cap of 100, and every element lands in length bucket 0.
tvalentyn commentedon Oct 16, 2026
Also, we could do better argument validation so that required sizing functions are supplied when batch limits are set, and also to prevent mixing weight-based and length-based params.
Eliaaazzz commentedon Oct 17, 2026
Thanks for flagging this, and sorry for the trouble. No, this was never finished; I left that gap in #37428. Since 2.72, max_batch_weight has been passed to BatchElements, which does not accept it. My tests only checked the kwargs dict, so they missed the pipeline construction failure. #37532 added support in SortAndBatchElements, which RunInference never uses.
I reproduced both other issues on master. A first element heavier than the current target triggers an empty batch in both non-stateful DoFns. The keyed wrappers pass the full tuple to element_size_fn, so len returns 2 for every keyed example.
I lean toward fixing it, since BatchElements has no separate count and weight limits today, and your #40603 already covers all three issues. If that works for you, I can open a PR on top of #40603 with the validation you mentioned and a test that builds and runs a pipeline. If you plan to land #40603 yourself, I am glad to review.
tvalentyn commentedon Oct 19, 2026
Thanks, I think our goal should be reach a state where API features either work as advertised or removed. If it is almost ready, and the risk of breaking something else is low, let's finish that. If that opens up a Pandora box for which we don't have time right now, we can choose to remove the APIs.
If you are open to finish the implementation, you are most welcome and that would be my preference since you have more context, feel free to use commits from #40603 if they are helpful or start from scratch as you see fit. Thanks a lot!
alisha-1000 commentedon Oct 20, 2026
Hi @tvalentyn , I checked PR #40603 and noticed it addresses the main
max_batch_weightissue. I'd like to contribute to any remaining work on #40644, particularly regression tests or validation for the related batching edge cases. Is there a specific part that would be useful for me to take on?Eliaaazzz commentedon Oct 20, 2026
Thanks, @tvalentyn! That makes sense, and I appreciate the guidance. I’m happy to take this forward, starting from #40603 where useful and keeping the changes focused. If it turns out to need a broader redesign, I’ll flag that here before expanding the scope.
@alisha-1000, thanks for offering to help! Would you be interested in taking on the argument validation and its unit tests: checking that the required sizing functions are provided and rejecting incompatible weight-based and length-based options? I can handle the core fixes and pipeline-level regression coverage, and I’d be happy to review your PR and help with any questions along the way. @tvalentyn, does that split sound reasonable to you?
alisha-1000 commentedon Oct 20, 2026
Thanks, Eliaaazzz! Yes, I'd be happy to take on the argument validation and unit tests.
I'll first review the existing
ModelHandlerand batching configuration APIs, along with PR #40603, so the validation stays consistent with the current implementation. I'll keep the changes focused on validating required sizing functions and rejecting incompatible weight-based and length-based options.I'll share my proposed validation rules before opening the PR if anything is ambiguous.
tvalentyn commentedon Oct 21, 2026
sure, @Eliaaazzz I am deferring to you as the expert on this corner on RunInference functionality :)
Finish max_batch_weight support for RunInference: issue number in CHA…
Merge pull request #40682 from Eliaaazzz/fix/runinference-max-batch-w…