From 06d1dc8813c3011facc2edc7f7ce897e9feca0b9 Mon Sep 17 00:00:00 2001 From: Tanmay Singh Date: Sun, 1 Nov 2026 00:31:40 +0530 Subject: [PATCH] [RunInference] Raise when a model handler returns fewer predictions than inputs utils._convert_to_result paired inputs and predictions with zip(), which stops at the shorter of the two, so inputs without a prediction were dropped without an error. It now raises a ValueError that gives the counts and names any dict output without one entry per input. Extra predictions are still ignored, so pipelines in which every input already gets a result are unaffected. Also fixes three tests that only passed because of the truncation. Fixes #40928 --- CHANGES.md | 1 + .../ml/inference/sklearn_inference_test.py | 3 +- .../ml/inference/tensorflow_inference_test.py | 4 +- sdks/python/apache_beam/ml/inference/utils.py | 62 ++++++++++++- .../apache_beam/ml/inference/utils_test.py | 89 +++++++++++++++++++ 5 files changed, 155 insertions(+), 4 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index 350a20c95955..e69b92c65675 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -83,6 +83,7 @@ * (Python) RunInference now measures model memory (the `model_byte_size` metric) when `torch` is not installed. It was always 0, because a failed `ModelManager` import also disabled the `resource` module ([#40854](https://github.com/apache/beam/issues/40854)). * (Python) Fixed `GeminiModelHandler` returning only one result per batch: each prompt is now sent as its own request (up to 32 at a time), so every input gets its own result. A custom `request_fn` must now return one response per input ([#40913](https://github.com/apache/beam/issues/40913)). +* (Python) RunInference now raises a `ValueError` when a model handler returns fewer predictions than inputs, instead of silently dropping the inputs that have no prediction. For example, this happened with PyTorch and Keras models that return a tuple or list of outputs, and with Hugging Face causal language models that return `past_key_values` ([#40928](https://github.com/apache/beam/issues/40928)). * Fixed X (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). ## Security Fixes diff --git a/sdks/python/apache_beam/ml/inference/sklearn_inference_test.py b/sdks/python/apache_beam/ml/inference/sklearn_inference_test.py index 76d6bc65729c..f91037df0325 100644 --- a/sdks/python/apache_beam/ml/inference/sklearn_inference_test.py +++ b/sdks/python/apache_beam/ml/inference/sklearn_inference_test.py @@ -160,7 +160,7 @@ def alternate_numpy_inference_fn( model: BaseEstimator, batch: Sequence[numpy.ndarray], inference_args: Optional[dict[str, Any]] = None) -> Any: - return [0] + return [0] * len(batch) def alternate_pandas_inference_fn( @@ -212,6 +212,7 @@ def test_custom_inference_fn(self): PredictionResult(numpy.array([7, 8, 9]), 0) ] inferences = inference_runner.run_inference(batched_examples, fake_model) + self.assertEqual(len(inferences), len(expected_predictions)) for actual, expected in zip(inferences, expected_predictions): self.assertTrue(_compare_prediction_result(actual, expected)) diff --git a/sdks/python/apache_beam/ml/inference/tensorflow_inference_test.py b/sdks/python/apache_beam/ml/inference/tensorflow_inference_test.py index 3a2e58e378eb..8f367603fe83 100644 --- a/sdks/python/apache_beam/ml/inference/tensorflow_inference_test.py +++ b/sdks/python/apache_beam/ml/inference/tensorflow_inference_test.py @@ -382,6 +382,9 @@ class TFRunInferenceTestWithMocks(unittest.TestCase): def setUp(self): self._load_model = tensorflow_inference._load_model tensorflow_inference._load_model = unittest.mock.MagicMock() + # The mocked model returns one prediction per input. + tensorflow_inference._load_model.return_value.predict.side_effect = ( + lambda batch, **kwargs: batch) def tearDown(self): tensorflow_inference._load_model = self._load_model @@ -433,7 +436,6 @@ def test_env_vars_set_correctly_numpy(self): os.environ.pop('FOO', None) self.assertFalse('FOO' in os.environ) batched_examples = [numpy.array([1]), numpy.array([10]), numpy.array([100])] - tensorflow_inference._load_model = unittest.mock.MagicMock() with TestPipeline() as pipeline: _ = ( pipeline diff --git a/sdks/python/apache_beam/ml/inference/utils.py b/sdks/python/apache_beam/ml/inference/utils.py index 8cbe836e22dd..dbf7417aa0ce 100644 --- a/sdks/python/apache_beam/ml/inference/utils.py +++ b/sdks/python/apache_beam/ml/inference/utils.py @@ -21,6 +21,7 @@ """ import os from collections.abc import Iterable +from collections.abc import Sized from functools import partial from typing import Any from typing import Optional @@ -45,6 +46,17 @@ def _convert_to_result( predictions: Union[Iterable, dict[Any, Iterable]], model_id: Optional[str] = None, ) -> Iterable[PredictionResult]: + """Pairs each input in the batch with its prediction. + + Predictions beyond the number of inputs are ignored. + + Raises: + ValueError: if there are fewer predictions than inputs. Pairing stops at + the shorter of the two, so without this check the remaining inputs + would silently get no result. + """ + if not isinstance(batch, Sized): + batch = list(batch) if isinstance(predictions, dict): # Go from one dictionary of type: {key_type1: Iterable, # key_type2: Iterable, ...} where each Iterable is of @@ -53,11 +65,57 @@ def _convert_to_result( predictions_per_tensor = [ dict(zip(predictions.keys(), v)) for v in zip(*predictions.values()) ] - return [ + results = [ PredictionResult(x, y, model_id) for x, y in zip(batch, predictions_per_tensor) ] - return [PredictionResult(x, y, model_id) for x, y in zip(batch, predictions)] + else: + results = [ + PredictionResult(x, y, model_id) for x, y in zip(batch, predictions) + ] + if len(results) < len(batch): + raise ValueError( + _missing_predictions_message( + len(batch), len(results), predictions, model_id)) + return results + + +def _count(number: int, singular: str, plural: str) -> str: + return f'{number} {singular if number == 1 else plural}' + + +def _missing_predictions_message( + num_inputs: int, + num_predictions: int, + predictions: Union[Iterable, dict[Any, Iterable]], + model_id: Optional[str]) -> str: + expected = 'Expected one prediction per input' + if model_id: + expected += f' from model {model_id}' + inputs = _count(num_inputs, 'input', 'inputs') + if isinstance(predictions, dict): + short_outputs = [] + for key, value in predictions.items(): + try: + length = len(value) + except TypeError: + continue + if length < num_inputs: + short_outputs.append( + f'{key!r} has {_count(length, "entry", "entries")}') + if short_outputs: + outputs = 'output' if len(short_outputs) == 1 else 'outputs' + return ( + f'{expected}, but {outputs} {", ".join(short_outputs)} for a batch ' + f'of {inputs}. Every output must have one entry per input, so remove ' + 'any that do not (for example, pass inference_args={"use_cache": ' + 'False} so that Hugging Face models do not return past_key_values), ' + 'or use an inference function that returns one prediction per input.') + predictions_count = _count(num_predictions, 'prediction', 'predictions') + return ( + f'{expected}, but got {predictions_count} for a batch of {inputs}. If ' + 'the model returns several outputs, such as a tuple, use an inference ' + 'function that returns one prediction per input.') class _ConvertIterToSingleton(beam.DoFn): diff --git a/sdks/python/apache_beam/ml/inference/utils_test.py b/sdks/python/apache_beam/ml/inference/utils_test.py index 66499a5a6f48..e8dd0162b143 100644 --- a/sdks/python/apache_beam/ml/inference/utils_test.py +++ b/sdks/python/apache_beam/ml/inference/utils_test.py @@ -18,9 +18,14 @@ import unittest +import numpy + import apache_beam as beam from apache_beam.io.filesystem import FileMetadata from apache_beam.ml.inference import utils +from apache_beam.ml.inference.base import ModelHandler +from apache_beam.ml.inference.base import PredictionResult +from apache_beam.ml.inference.base import RunInference from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to @@ -99,5 +104,89 @@ def test_emitting_singleton_output(self): assert_that(files_pc, equal_to(['', 'path3.py', 'path4.py'])) +class _ModelHandlerWithMissingPredictions(ModelHandler[int, + PredictionResult, + None]): + """Returns a prediction for every input except the last one in a batch.""" + def load_model(self): + return None + + def run_inference(self, batch, model, inference_args=None): + return utils._convert_to_result(batch, [x * 10 for x in batch[:-1]]) + + +class ConvertToResultTest(unittest.TestCase): + def test_pairs_each_input_with_its_prediction(self): + self.assertEqual( + utils._convert_to_result(['a', 'b'], [1, 2], model_id='m'), + [PredictionResult('a', 1, 'm'), PredictionResult('b', 2, 'm')]) + + def test_splits_batched_array_into_one_prediction_per_input(self): + results = utils._convert_to_result(['a', 'b', 'c'], + numpy.array([[1, 2], [3, 4], [5, 6]])) + self.assertEqual([r.inference.tolist() for r in results], + [[1, 2], [3, 4], [5, 6]]) + + def test_splits_dict_predictions_into_one_dict_per_input(self): + results = utils._convert_to_result(['a', 'b'], { + 'label': [0, 1], 'score': [0.9, 0.8] + }) + self.assertEqual([r.inference for r in results], [{ + 'label': 0, 'score': 0.9 + }, { + 'label': 1, 'score': 0.8 + }]) + + def test_accepts_predictions_from_a_generator(self): + results = utils._convert_to_result(['a', 'b'], (x for x in [1, 2])) + self.assertEqual([r.inference for r in results], [1, 2]) + + def test_raises_when_there_are_fewer_predictions_than_inputs(self): + # For example, a model that returns a tuple of two outputs for a batch of + # five inputs: zip() would pair the inputs with the outputs and drop the + # last three inputs. + with self.assertRaisesRegex(ValueError, + 'got 2 predictions for a batch of 5 inputs'): + utils._convert_to_result( + list(range(5)), ('logits', 'scores'), model_id='m') + + def test_raises_when_a_generator_yields_too_few_predictions(self): + with self.assertRaisesRegex(ValueError, + 'got 1 prediction for a batch of 3 inputs'): + utils._convert_to_result(['a', 'b', 'c'], (x for x in [1])) + + def test_raises_naming_dict_outputs_without_one_entry_per_input(self): + # For example, a Hugging Face causal LM returns past_key_values with one + # entry per layer rather than one per input. + predictions = { + 'logits': [[0.1], [0.2], [0.3]], + 'past_key_values': [('k1', 'v1'), ('k2', 'v2')] + } + with self.assertRaisesRegex( + ValueError, + "output 'past_key_values' has 2 entries for a batch of 3 inputs"): + utils._convert_to_result(['a', 'b', 'c'], predictions) + + def test_extra_predictions_are_ignored(self): + # Only missing predictions raise, so pipelines in which every input + # already gets a result keep working. + results = utils._convert_to_result(['a'], ('logits', 'scores')) + self.assertEqual([r.inference for r in results], ['logits']) + + def test_run_inference_sends_inputs_without_predictions_to_failed_output( + self): + with TestPipeline() as p: + main, other = ( + p + | beam.Create([1, 2, 3, 4]) + | RunInference(_ModelHandlerWithMissingPredictions() + ).with_exception_handling()) + assert_that(main, equal_to([]), label='CheckNoResults') + failed_inputs = ( + other.failed_inferences | beam.FlatMap(lambda failure: failure[0])) + assert_that( + failed_inputs, equal_to([1, 2, 3, 4]), label='CheckFailedInputs') + + if __name__ == '__main__': unittest.main()