diff --git a/deepinfer/entrypoints/llm.py b/deepinfer/entrypoints/llm.py index aa8d5e8..859abb0 100644 --- a/deepinfer/entrypoints/llm.py +++ b/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] diff --git a/tests/entrypoints/test_llm_generate_list_lengths.py b/tests/entrypoints/test_llm_generate_list_lengths.py new file mode 100644 index 0000000..43108bd --- /dev/null +++ b/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 == []