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

CHANGES.md

@@ -86,6 +86,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)).
 * (Java) BigQuery Storage Write API successful-row output now propagates fused downstream exceptions instead of allowing the bundle to succeed with missing downstream results ([#40004](https://github.com/apache/beam/pull/40004)).
+* (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

@@ -53,11 +53,19 @@ 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(
+        f'Expected one prediction per input, but got {len(results)} for a '
+        f'batch of size {len(batch)}.')
+  return results
 
 
 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,85 @@ 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 for a batch of size 5'):
+      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 for a batch of size 3'):
+      utils._convert_to_result(['a', 'b', 'c'], (x for x in [1]))
+
+  def test_raises_when_a_dict_output_has_too_few_entries(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, 'got 2 for a batch of size 3'):
+      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()