tannnmayy opened a new issue, #40928:
URL: https://github.com/apache/beam/issues/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 (5d82bbc693), Python 3.12, torch
2.11 (CPU), TensorFlow 2.21 and transformers 4.55, sending 20 inputs in batches
of 5:
| Setup | Results | `num_inferences` |
|---|---|---|
| `PytorchModelHandlerTensor`, model returns a tuple `(logits, score)` | 8 |
20 |
| `PytorchModelHandlerTensor` with `torch.nn.LSTM` (returns `(output, (h,
c))`) | 8 | 20 |
| `TFModelHandlerNumpy`, Keras model with two outputs | 8 | 20 |
| `HuggingFaceModelHandlerKeyedTensor` + `GPT2LMHeadModel` (2 layers),
default arguments | 8 | 20 |
| Same, with `inference_args={'use_cache': False, 'output_hidden_states':
True}` | 12 | 20 |
| `OpenAIModelHandler` / `AnthropicModelHandler` with a `request_fn` that
returns one response per batch | 4 | 20 |
| The tuple model with `with_exception_handling()` | 8 results, 0 failures |
20 |
| Controls: single-output PyTorch and Keras models, BERT, scikit-learn, a
per-prompt `request_fn` | 20 | 20 |
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):
```python
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 #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 Priority
Priority: 2 (default / most bugs should be filed as P2)
### Issue Components
- [x] 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
--
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]
