Merge bf92779eaacea46e113038f71409ce083df05ea3 into 553741557a57df44a46fb3ab72b6aafeaf945337

deepinfer/entrypoints/llm.py

@@ -119,7 +119,6 @@ class LLM:
         if isinstance(prompts, str):
             # Convert a single prompt to a list.
             prompts = [prompts]
-            multi_modal_data = [multi_modal_data]
         if prompts is not None and prompt_token_ids is not None and len(prompts) != len(prompt_token_ids):
             raise ValueError("The lengths of prompts and prompt_token_ids must be the same.")
         if sampling_params is None:
@@ -128,6 +127,23 @@ class LLM:
 
         # Add requests to the engine.
         num_requests = len(prompts) if prompts is not None else len(prompt_token_ids)
+
+        # Fail loudly instead of silently misaligning results[i] with prompts[i]:
+        # per-prompt lists (sampling_params / multi_modal_data) must have exactly
+        # one entry per prompt, and a bare multi_modal_data dict can only pair
+        # with a single prompt. (vLLM upstream enforces the same invariants.)
+        if isinstance(sampling_params, list) and len(sampling_params) != num_requests:
+            raise ValueError(f"len(sampling_params) ({len(sampling_params)}) must match the number of prompts ({num_requests}).")
+        if isinstance(multi_modal_data, dict):
+            if num_requests != 1:
+                raise ValueError(
+                    f"A single multi_modal_data dict can only be paired with exactly one prompt; "
+                    f"got {num_requests} prompts. Pass a list of dicts (one per prompt) instead."
+                )
+            multi_modal_data = [multi_modal_data]
+        elif isinstance(multi_modal_data, list) and len(multi_modal_data) != num_requests:
+            raise ValueError(f"len(multi_modal_data) ({len(multi_modal_data)}) must match the number of prompts ({num_requests}).")
+
         for i in range(num_requests):
             prompt = prompts[i] if prompts is not None else None
             token_ids = None if prompt_token_ids is None else prompt_token_ids[i]

tests/entrypoints/test_llm_generate_list_lengths.py

@@ -0,0 +1,89 @@
+# SPDX-License-Identifier: Apache-2.0
+"""generate() must fail loudly when per-prompt list lengths are mismatched.
+
+Otherwise results[i] silently no longer corresponds to prompts[i] (e.g. an
+extra sampling_params / multi_modal_data entry is silently dropped, or a bare
+multi_modal_data dict is paired with several prompts and triggers an obscure
+TypeError/KeyError deep inside the request loop).
+"""
+
+import itertools
+from types import SimpleNamespace
+
+import pytest
+
+from deepinfer.entrypoints.llm import LLM
+from deepinfer.sampling_params import SamplingParams
+
+SENTINEL_SEQUENCE = object()
+
+
+def _llm() -> LLM:
+    """A bare LLM whose engine accepts every request, without loading a model."""
+    llm = object.__new__(LLM)
+    llm.request_counter = itertools.count()
+    llm.llm_engine = SimpleNamespace(
+        tokenizer=SimpleNamespace(encode=lambda prompt: [1, 2, 3]),
+        multi_modal_processor=None,
+        add_request=lambda **kwargs: SENTINEL_SEQUENCE,
+    )
+    llm.decode_sequence = False
+    return llm
+
+
+def _run_engine(_use_tqdm: bool = False) -> list[object]:
+    return []
+
+
+def test_generate_accepts_single_prompt_with_single_sampling_params() -> None:
+    llm = _llm()
+    llm._run_engine = _run_engine
+
+    results = llm.generate("hello", sampling_params=SamplingParams(max_tokens=1))
+
+    assert results == []
+
+
+def test_generate_rejects_sampling_params_length_mismatch() -> None:
+    llm = _llm()
+
+    with pytest.raises(ValueError, match="len\\(sampling_params\\) .* must match"):
+        llm.generate(
+            ["hello", "world"],
+            sampling_params=[SamplingParams(max_tokens=1), SamplingParams(max_tokens=2), SamplingParams(max_tokens=3)],
+        )
+
+
+def test_generate_rejects_multi_modal_data_length_mismatch() -> None:
+    llm = _llm()
+
+    with pytest.raises(ValueError, match="len\\(multi_modal_data\\) .* must match"):
+        llm.generate(
+            ["hello", "world"],
+            sampling_params=SamplingParams(max_tokens=1),
+            multi_modal_data=[dict(image_bytes=[b"a"]), dict(image_bytes=[b"b"]), dict(image_bytes=[b"c"])],
+        )
+
+
+def test_generate_rejects_bare_dict_with_multiple_prompts() -> None:
+    llm = _llm()
+
+    with pytest.raises(ValueError, match="single multi_modal_data dict"):
+        llm.generate(
+            ["hello", "world"],
+            sampling_params=SamplingParams(max_tokens=1),
+            multi_modal_data=dict(image_bytes=[b"a"]),
+        )
+
+
+def test_generate_accepts_bare_dict_with_single_prompt() -> None:
+    llm = _llm()
+    llm._run_engine = _run_engine
+
+    results = llm.generate(
+        "hello",
+        sampling_params=SamplingParams(max_tokens=1),
+        multi_modal_data=dict(image_bytes=[b"a"]),
+    )
+
+    assert results == []