tannnmayy opened a new issue, #40913: URL: https://github.com/apache/beam/issues/40913
### What happened? **Environment:** Apache Beam master at b88840e0d3 (same code in 2.77.0), Python SDK, google-genai 2.33.0. `GeminiModelHandler`'s request functions send the whole batch in a single `generate_content` call ([gemini_inference.py L77-L78](https://github.com/apache/beam/blob/b88840e0d3624dc2bec79a5b6f0292b640761ef2/sdks/python/apache_beam/ml/inference/gemini_inference.py#L77-L78)), and `request()` wraps the one response in a list ([L263-L266](https://github.com/apache/beam/blob/b88840e0d3624dc2bec79a5b6f0292b640761ef2/sdks/python/apache_beam/ml/inference/gemini_inference.py#L263-L266)): ```python responses = [self.request_fn(self.model_name, batch, model, inference_args)] return utils._convert_to_result(batch, responses, self.model_name) ``` google-genai turns a list of strings passed as `contents` into a single user message with one part per string, so Gemini returns one response for the whole batch. `_convert_to_result` zips the batch with that one-item list, so only the first element of each batch gets a `PredictionResult`. The other elements are dropped without any error, and the one result answers the merged message rather than its own prompt. **With the real Gemini API** (`gemini-2.5-flash`, a batch of 3 prompts asking for 2+2, the capital of France and the color of a clear sky): | Call | Output | |---|---| | One `generate_content` call with the 3 prompts | one response: `4 Paris Blue` | | `RunInference(GeminiModelHandler(...))` with a batch of 3 | 1 result: the first prompt gets `4 Paris Blue`; the other two get nothing | **Minimal reproduction** (no API key needed): ```python from apache_beam.ml.inference.gemini_inference import GeminiModelHandler def fake_request_fn(model_name, batch, model, inference_args): # Stands in for generate_from_string: one generate_content call per batch. return f"one response for {len(batch)} prompts" handler = GeminiModelHandler( model_name="gemini-2.5-flash", request_fn=fake_request_fn, api_key="fake-key") results = list(handler.request(["prompt A", "prompt B", "prompt C"], model=None)) print(len(results)) # 1 print(results[0].example, "->", results[0].inference) # prompt A -> one response for 3 prompts ``` **In a pipeline** (default batching, the same fake `request_fn`, local runner), 200 prompts produced 11 results, while RunInference's `num_inferences` metric reported 200 and `with_exception_handling()` sent nothing to the failure output. BatchElements starts with batches of one element, so small runs such as the `gemini_text_classification.py` example (3 prompts) don't show the problem. `generate_image_from_strings_and_images` follows the same pattern. The batching path isn't covered by the unit tests, and the examples and integration tests use 1–3 prompts. The list wrapping was added in #36177; before that (#35036), `zip` iterated over the response object's fields. **Proposed fix:** send one `generate_content` request per input, concurrently on a bounded thread pool (up to 32 in flight per batch) as `OpenAIModelHandler` does (#40239), and have `request()` check that the request function returned one response per input. I have this working with tests and will open a PR. Questions: 1. Should a mismatched response count raise an error (my current approach) or only log a warning? 2. Is copying the small fan-out helper from `openai_inference.py` fine, or should it move to a shared module? 3. Should `GeminiModelHandler` get a default `max_batch_size`? ### Issue Priority Priority: 1 (data loss / total loss of function) ### 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]
