[RunInference] Raise when a model handler returns fewer predictions than inputs

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

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))
 

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

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<val_type1>,
     # key_type2: Iterable<val_type2>, ...} 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):

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()