[Bug]: RunInference silently drops inputs when a model returns fewer predictions than inputs #40928

What happened?

RunInference pairs each batch of inputs with the model's predictions in utils._convert_to_result, using zip(). zip() stops at the shorter of the two, so when a model handler produces fewer predictions than inputs, the leftover inputs get no result and no error. The num_inferences metric still counts every input, and with_exception_handling() reports no failures.

This has already caused two bugs that were only found after release: #36587 (Hugging Face pipelines returned one result per batch, fixed in #36593) and #40913 (Gemini, fixed in #40914). It can still happen with built-in handlers.

I reproduced it with real models on master (5d82bbc), Python 3.12, torch 2.11 (CPU), TensorFlow 2.21 and transformers 4.55, sending 20 inputs in batches of 5:

It's easy to miss even in Beam's own tests. Three existing tests only pass because of the truncation: sklearn_inference_test.py::test_custom_inference_fn uses an inference function that returns one prediction for three inputs (and its own zip() compares only the first result), and the two env-var tests in tensorflow_inference_test.py mock the model, so their pipelines produce no results at all.

For a tuple output, zip() pairs the 5 inputs with the 2 outputs, so the first input gets the whole batch's logits, the second gets the whole batch's scores, and the other three get nothing. For dict outputs, zip(*predictions.values()) stops at the shortest value. GPT-2's past_key_values has one entry per layer, so a batch larger than the layer count is cut down to the layer count; with the real 12-layer GPT-2, that's any batch above 12. Keyed inputs are affected the same way.

Minimal repro (no download; a tiny model with random weights):

import os, tempfile, torch
import apache_beam as beam
from apache_beam.ml.inference.base import RunInference
from apache_beam.ml.inference.pytorch_inference import PytorchModelHandlerTensor

class TwoOutputs(torch.nn.Module):
  def __init__(self):
    super().__init__()
    self.linear = torch.nn.Linear(3, 2)
  def forward(self, x):
    logits = self.linear(x)
    return logits, logits.sum(dim=1)

path = os.path.join(tempfile.mkdtemp(), 'model.pt')
torch.save(TwoOutputs().state_dict(), path)
handler = PytorchModelHandlerTensor(
    state_dict_path=path, model_class=TwoOutputs, model_params={},
    min_batch_size=5, max_batch_size=5)

with beam.Pipeline() as p:
  (p
   | beam.Create([torch.randn(3) for _ in range(20)])
   | RunInference(handler)
   | beam.combiners.Count.Globally()
   | beam.Map(print))  # prints 8, not 20

Proposal: make _convert_to_result raise a ValueError when there are fewer predictions than inputs. The message would give the counts and, for dict outputs, name the output whose length doesn't match, with a hint such as inference_args={'use_cache': False}. Only missing predictions would raise, so pipelines in which every input already gets a result keep working (for example, batch size 1, or GPT-2 with use_cache=False). I have this ready locally. With it, every row above that loses data fails with the message instead, the controls are unchanged, and the whole apache_beam/ml test suite gives the same results as master once the three tests mentioned above are fixed to return one prediction per input.

Questions:

1. Raise (as proposed) or log a warning?
2. Should RunInference also check the count itself, in _RunInferenceDoFn, to cover custom handlers that pair results with their own zip() (as the Hugging Face handler did before use utils._convert_to_result for huggingface_inference #36593)? I tried a prototype of that check locally. It catches that case, and it changes no test outcome in apache_beam/ml. But it would also apply to every custom handler, so I'd rather ask first.
3. Should extra predictions (more than inputs) raise too? That would also catch tuple outputs at batch size 1, but it would break pipelines that work today.

cc @jrmccluskey @liferoad

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: Hazelcast Jet Runner
• Component: Google Cloud Dataflow Runner