[RunInference] Send one Gemini request per input so every element gets a result

CHANGES.md

@@ -82,6 +82,7 @@
 ## Bugfixes
 
 * (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)).
 * Fixed X (Java/Python) ([#X](https://github.com/apache/beam/issues/X)).
 
 ## Security Fixes

sdks/python/apache_beam/ml/inference/gemini_inference.py

@@ -19,8 +19,10 @@
 from collections.abc import Callable
 from collections.abc import Iterable
 from collections.abc import Sequence
+from concurrent.futures import ThreadPoolExecutor
 from typing import Any
 from typing import Optional
+from typing import TypeVar
 from typing import Union
 from typing import cast
 
@@ -37,6 +39,16 @@
 
 LOGGER = logging.getLogger("GeminiModelHandler")
 
+# Upper bound on the number of in-flight requests issued concurrently for a
+# single batch by the request functions in this module, which send one
+# generate_content request per input. Bounds thread usage when adaptive
+# batching produces large batches; use max_batch_size to further limit
+# per-batch concurrency.
+_MAX_CONCURRENT_REQUESTS_PER_BATCH = 32
+
+_InputT = TypeVar('_InputT')
+_OutputT = TypeVar('_OutputT')
+
 
 def _retry_on_appropriate_service_error(exception: Exception) -> bool:
   """
@@ -57,14 +69,37 @@ def _retry_on_appropriate_service_error(exception: Exception) -> bool:
   return exception.code == 429 or exception.code >= 500
 
 
+def _fan_out(fn: Callable[[_InputT], _OutputT],
+             items: Sequence[_InputT]) -> list[_OutputT]:
+  """Applies fn to each item concurrently and returns results in input order.
+
+  The google-genai client can be shared across threads (its HTTP connection
+  pool is thread-safe and credential refreshes are guarded by a lock), so
+  per-item requests are dispatched on a bounded thread pool. If any request
+  raises, pending requests are cancelled and the first exception (in input
+  order) is propagated so Beam's retry logic can handle the batch.
+  """
+  if len(items) <= 1:
+    return [fn(item) for item in items]
+  executor = ThreadPoolExecutor(
+      max_workers=min(len(items), _MAX_CONCURRENT_REQUESTS_PER_BATCH),
+      thread_name_prefix='GeminiModelHandler')
+  try:
+    return list(executor.map(fn, items))
+  finally:
+    executor.shutdown(wait=True, cancel_futures=True)
+
+
 def generate_from_string(
     model_name: str,
     batch: Sequence[str],
     model: genai.Client,
     inference_args: dict[str, Any]):
   """ Request function that expects inputs to be composed of strings, then
   sends requests to Gemini to generate text responses based on the text
-  prompts.
+  prompts. Each prompt is sent as its own generate_content request (up to
+  _MAX_CONCURRENT_REQUESTS_PER_BATCH at a time), because Gemini treats a list
+  of strings in a single request as one message and returns one response.
 
   Args:
     model_name: the Gemini model to use for the request. This model should be
@@ -73,9 +108,14 @@ def generate_from_string(
     model: the genai Client
     inference_args: any additional arguments passed to the generate_content
       call.
+
+  Returns:
+    A list with one response per prompt, in the same order as the batch.
   """
-  return model.models.generate_content(
-      model=model_name, contents=cast(Any, batch), **inference_args)
+  return _fan_out(
+      lambda prompt: model.models.generate_content(
+          model=model_name, contents=prompt, **inference_args),
+      batch)
 
 
 def generate_image_from_strings_and_images(
@@ -86,7 +126,10 @@ def generate_image_from_strings_and_images(
   """ Request function that expects inputs to be composed of lists of strings
   and PIL Image instances, then sends requests to Gemini to generate images
   based on the text prompts and contextual images. This is currently intended
-  to be used with the gemini-2.5-flash-image model (AKA Nano Banana.)
+  to be used with the gemini-2.5-flash-image model (AKA Nano Banana.) Each
+  element of the batch is sent as its own generate_content request (up to
+  _MAX_CONCURRENT_REQUESTS_PER_BATCH at a time), so every element gets its own
+  response.
 
   Args:
     model_name: the Gemini model to use for the request. This model should be
@@ -96,9 +139,14 @@ def generate_image_from_strings_and_images(
     model: the genai Client
     inference_args: any additional arguments passed to the generate_content
       call.
+
+  Returns:
+    A list with one response per element, in the same order as the batch.
   """
-  return model.models.generate_content(
-      model=model_name, contents=cast(Any, batch), **inference_args)
+  return _fan_out(
+      lambda element: model.models.generate_content(
+          model=model_name, contents=cast(Any, element), **inference_args),
+      batch)
 
 
 class GeminiModelHandler(RemoteModelHandler[Any, PredictionResult,
@@ -132,9 +180,9 @@ def __init__(
     Args:
       model_name: the Gemini model to send the request to
       request_fn: the function to use to send the request. Should take the
-        model name and the parameters from request() and return the responses
-        from Gemini. The class will handle bundling the inputs and responses
-        together.
+        model name and the parameters from request() and return a list with
+        one response from Gemini per input, in the same order as the inputs.
+        The class will handle bundling the inputs and responses together.
       api_key: the Gemini Developer API key to use for the requests. Setting
         this parameter sends requests for this job to the Gemini Developer API.
         If this paramter is provided, do not set the project or location
@@ -244,9 +292,9 @@ def request(
       model: genai.Client,
       inference_args: Optional[dict[str, Any]] = None
   ) -> Iterable[PredictionResult]:
-    """ Sends a prediction request to a Gemini service containing a batch
-    of inputs and matches that input with the prediction response from
-    the endpoint as an iterable of PredictionResults.
+    """ Sends a batch of inputs to a Gemini service using the request function
+    and matches each input with its response as an iterable of
+    PredictionResults.
 
     Args:
       batch: a sequence of any values to be passed to the Gemini service.
@@ -256,11 +304,22 @@ def request(
         prediction request.
 
     Returns:
-      An iterable of Predictions.
+      An iterable of Predictions, one per input.
     """
     if inference_args is None:
       inference_args = {}
-    # Wrap the responses in a list to prevent zip() call from treating the
-    # response itself as an iterable of individual responses.
-    responses = [self.request_fn(self.model_name, batch, model, inference_args)]
+    responses = self.request_fn(self.model_name, batch, model, inference_args)
+    if len(batch) == 1 and not isinstance(responses, (list, tuple)):
+      # Request functions written for single-input batches may return the
+      # response itself rather than a list of responses.
+      responses = [responses]
+    if not isinstance(responses, (list, tuple)):
+      raise ValueError(
+          'request_fn must return a list with one response per input, but '
+          f'returned a single {type(responses).__name__} for a batch of '
+          f'{len(batch)} inputs.')
+    if len(responses) != len(batch):
+      raise ValueError(
+          'request_fn must return one response per input, but returned '
+          f'{len(responses)} responses for a batch of {len(batch)} inputs.')
     return utils._convert_to_result(batch, responses, self.model_name)

sdks/python/apache_beam/ml/inference/gemini_inference_test.py

@@ -16,18 +16,59 @@
 #
 # pytype: skip-file
 
+import threading
+import time
 import unittest
 from unittest import mock
 
 try:
   from google.genai import errors
+  from google.genai import types
 
   from apache_beam.ml.inference.gemini_inference import GeminiModelHandler
   from apache_beam.ml.inference.gemini_inference import _retry_on_appropriate_service_error
   from apache_beam.ml.inference.gemini_inference import generate_from_string
+  from apache_beam.ml.inference.gemini_inference import generate_image_from_strings_and_images
 except ImportError:
   raise unittest.SkipTest('Gemini dependencies are not installed')
 
+import apache_beam as beam
+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
+
+_TEST_MODEL = 'gemini-model-123'
+
+
+def _make_response(text):
+  return types.GenerateContentResponse(
+      candidates=[
+          types.Candidate(
+              content=types.Content(
+                  role='model', parts=[types.Part(text=text)]))
+      ])
+
+
+def _make_mock_client() -> mock.MagicMock:
+  client = mock.MagicMock()
+  client.models.generate_content = mock.MagicMock()
+  return client
+
+
+def _echo_contents(**kwargs):
+  """Returns a fake response derived from the request's contents.
+
+  Requests within a batch are issued concurrently, so call order is
+  nondeterministic; keying the response on the request contents lets tests
+  check that results are returned in batch order.
+  """
+  return _make_response(f"answer for: {kwargs['contents']}")
+
+
+def _fake_request_fn(model_name, batch, model, inference_args):
+  return [_make_response(f'answer for: {prompt}') for prompt in batch]
+
 
 class RetryOnClientErrorTest(unittest.TestCase):
   def test_retry_on_client_error_positive(self):
@@ -135,5 +176,150 @@ def test_create_client_with_flex_api(
         http_options=mock_http_options.return_value)
 
 
+class GenerateFromStringTest(unittest.TestCase):
+  def test_sends_each_prompt_as_its_own_request(self):
+    client = _make_mock_client()
+    client.models.generate_content.side_effect = _echo_contents
+    results = generate_from_string(_TEST_MODEL, ['hello', 'world'], client, {})
+    self.assertEqual([r.text for r in results],
+                     ['answer for: hello', 'answer for: world'])
+    self.assertEqual(client.models.generate_content.call_count, 2)
+    calls = client.models.generate_content.call_args_list
+    self.assertCountEqual([c.kwargs['contents'] for c in calls],
+                          ['hello', 'world'])
+    for c in calls:
+      self.assertEqual(c.kwargs['model'], _TEST_MODEL)
+
+  def test_requests_run_concurrently_and_preserve_order(self):
+    lock = threading.Lock()
+    state = {'in_flight': 0, 'max_in_flight': 0}
+    completion_order = []
+    prompts = [f'p{i}' for i in range(4)]
+
+    def slow_generate(**kwargs):
+      prompt = kwargs['contents']
+      with lock:
+        state['in_flight'] += 1
+        state['max_in_flight'] = max(state['max_in_flight'], state['in_flight'])
+      # Earlier prompts sleep longer so they finish last.
+      time.sleep(0.05 * (len(prompts) - prompts.index(prompt)))
+      with lock:
+        state['in_flight'] -= 1
+        completion_order.append(prompt)
+      return _make_response(f'answer for: {prompt}')
+
+    client = _make_mock_client()
+    client.models.generate_content.side_effect = slow_generate
+    results = generate_from_string(_TEST_MODEL, prompts, client, {})
+
+    self.assertGreater(state['max_in_flight'], 1)
+    self.assertNotEqual(completion_order, prompts)
+    self.assertEqual([r.text for r in results],
+                     [f'answer for: {p}' for p in prompts])
+
+  def test_concurrency_is_capped(self):
+    lock = threading.Lock()
+    state = {'in_flight': 0, 'max_in_flight': 0}
+
+    def slow_generate(**kwargs):
+      with lock:
+        state['in_flight'] += 1
+        state['max_in_flight'] = max(state['max_in_flight'], state['in_flight'])
+      time.sleep(0.02)
+      with lock:
+        state['in_flight'] -= 1
+      return _make_response('ok')
+
+    client = _make_mock_client()
+    client.models.generate_content.side_effect = slow_generate
+    with mock.patch('apache_beam.ml.inference.gemini_inference.'
+                    '_MAX_CONCURRENT_REQUESTS_PER_BATCH',
+                    2):
+      results = generate_from_string(
+          _TEST_MODEL, [f'p{i}' for i in range(6)], client, {})
+
+    self.assertEqual(len(results), 6)
+    self.assertLessEqual(state['max_in_flight'], 2)
+
+  def test_propagates_request_error(self):
+    def fail_on_second_prompt(**kwargs):
+      if kwargs['contents'] == 'p1':
+        raise errors.APIError(code=429, response_json={})
+      return _make_response('ok')
+
+    client = _make_mock_client()
+    client.models.generate_content.side_effect = fail_on_second_prompt
+    with self.assertRaises(errors.APIError):
+      generate_from_string(_TEST_MODEL, ['p0', 'p1', 'p2'], client, {})
+
+
+class GenerateImageFromStringsAndImagesTest(unittest.TestCase):
+  def test_sends_each_element_as_its_own_request(self):
+    client = _make_mock_client()
+    client.models.generate_content.side_effect = _echo_contents
+    elements = [['make it blue', 'context a'], ['make it red', 'context b']]
+    results = generate_image_from_strings_and_images(
+        _TEST_MODEL, elements, client, {})
+    self.assertEqual([r.text for r in results],
+                     [f'answer for: {e}' for e in elements])
+    calls = client.models.generate_content.call_args_list
+    self.assertCountEqual([c.kwargs['contents'] for c in calls], elements)
+
+
+class GeminiModelHandlerRequestTest(unittest.TestCase):
+  def _handler(self, request_fn):
+    return GeminiModelHandler(
+        model_name=_TEST_MODEL, request_fn=request_fn, api_key='123456789')
+
+  def test_returns_one_result_per_input(self):
+    client = _make_mock_client()
+    client.models.generate_content.side_effect = _echo_contents
+    prompts = ['p0', 'p1', 'p2']
+    results = list(self._handler(generate_from_string).request(prompts, client))
+    self.assertEqual([(r.example, r.inference.text) for r in results],
+                     [(p, f'answer for: {p}') for p in prompts])
+    self.assertEqual([r.model_id for r in results], [_TEST_MODEL] * 3)
+
+  def test_accepts_single_response_for_single_input(self):
+    def request_fn(model_name, batch, model, inference_args):
+      return _make_response('only answer')
+
+    results = list(self._handler(request_fn).request(['p0'], None))
+    self.assertEqual([r.inference.text for r in results], ['only answer'])
+
+  def test_rejects_single_response_for_multiple_inputs(self):
+    def request_fn(model_name, batch, model, inference_args):
+      return _make_response('one merged answer')
+
+    with self.assertRaisesRegex(ValueError, 'one response per input'):
+      self._handler(request_fn).request(['p0', 'p1', 'p2'], None)
+
+  def test_rejects_wrong_number_of_responses(self):
+    def request_fn(model_name, batch, model, inference_args):
+      return [_make_response('a'), _make_response('b')]
+
+    with self.assertRaisesRegex(ValueError, 'one response per input'):
+      self._handler(request_fn).request(['p0', 'p1', 'p2'], None)
+
+
+class GeminiModelHandlerPipelineTest(unittest.TestCase):
+  def test_every_input_gets_its_own_result(self):
+    handler = GeminiModelHandler(
+        model_name=_TEST_MODEL,
+        request_fn=_fake_request_fn,
+        api_key='123456789',
+        min_batch_size=5,
+        max_batch_size=5)
+    prompts = [f'prompt {i}' for i in range(20)]
+
+    with TestPipeline() as p:
+      results = (
+          p
+          | beam.Create(prompts)
+          | RunInference(handler)
+          | beam.Map(lambda r: (r.example, r.inference.text)))
+      assert_that(results, equal_to([(x, f'answer for: {x}') for x in prompts]))
+
+
 if __name__ == '__main__':
   unittest.main()