[Frontend] Implement Batched Chat Completions (#5317)

Signed-off-by: Alex Brooks <albrooks@redhat.com>

tests/e2e/online_serving/test_flux_2_dev_expansion.py

@@ -136,3 +136,36 @@ def test_flux_2_dev(
     }
 
     openai_client.send_diffusion_request(request_config)
+
+
+@pytest.mark.parametrize(
+    "omni_server",
+    _get_flux_2_dev_feature_cases(MODEL),
+    indirect=True,
+)
+def test_flux_2_dev_batched_chat_completions(
+    omni_server: OmniServer,
+    openai_client: OpenAIClientHandler,
+):
+    """Validate batched chat completions for diffusion models."""
+    messages = [
+        [{"role": "user", "content": PROMPT}],
+        [{"role": "user", "content": "A sunset over the ocean."}],
+    ]
+    responses = openai_client.send_batched_chat_completions_http_request(
+        {
+            "json": {
+                "model": omni_server.model,
+                "messages": messages,
+                "num_inference_steps": 2,
+                "height": 512,
+                "width": 512,
+                "seed": 42,
+            },
+        },
+    )
+    assert responses and len(responses) == 1
+    resp = responses[0]
+    assert resp.success
+    choices = resp.json_body["choices"]
+    assert len(choices) == len(messages)

tests/e2e/online_serving/test_qwen3_omni.py

@@ -250,3 +250,101 @@ def test_completions_rejected_for_thinker_talker(omni_server, openai_client) ->
         err_code=400,
     )
     assert not responses[0].success
+
+
+@pytest.mark.advanced_model
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_batched_completions_text(omni_server, openai_client) -> None:
+    """Ensure that we can make a batch chat completions request (text only)."""
+    responses = openai_client.send_batched_chat_completions_http_request(
+        {
+            "json": {
+                "model": omni_server.model,
+                "messages": [
+                    [{"role": "user", "content": "What color is the sky? Answer in one word."}],
+                    [{"role": "user", "content": "What is 2+2? Answer in one word."}],
+                ],
+                "modalities": ["text"],
+                "max_tokens": 20,
+            },
+        },
+    )
+    assert responses
+    resp = responses[0]
+    assert resp.success
+    body = resp.json_body
+    choices = body["choices"]
+    assert len(choices) == 2
+    for choice in choices:
+        assert choice["message"]["content"]
+
+
+@pytest.mark.advanced_model
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_batched_completions_audio_out(omni_server, openai_client) -> None:
+    """Ensure that we can make a batch chat completions request (audio + text)."""
+    messages = [
+        [{"role": "user", "content": "Say hello."}],
+        [{"role": "user", "content": "Say goodbye."}],
+    ]
+    num_messages = len(messages)
+    responses = openai_client.send_batched_chat_completions_http_request(
+        {
+            "json": {
+                "model": omni_server.model,
+                "messages": messages,
+                "modalities": ["text", "audio"],
+                "max_tokens": 200,
+            },
+        },
+    )
+    assert responses and len(responses) == 1
+    resp = responses[0]
+    assert resp.success
+    choices = resp.json_body["choices"]
+
+    assert len(choices) == num_messages
+
+    # Every message in the batch should contain both the audio and text content
+    for choice in choices:
+        assert choice["message"]["audio"]["data"]
+        assert choice["message"]["content"]
+
+
+@pytest.mark.advanced_model
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+@pytest.mark.parametrize(
+    "sad_opts",
+    [
+        {"stream": "not-a-bool"},  # Fails validation
+        {"n": 2},  # vLLM Batched endpoints doesn't support n > 1
+    ],
+)
+def test_batched_completions_with_bad_values(omni_server, openai_client, sad_opts) -> None:
+    """Ensure that a bad values are correctly handled as 400s."""
+    messages = [
+        [{"role": "user", "content": "Say hello."}],
+        [{"role": "user", "content": "Say goodbye."}],
+    ]
+    responses = openai_client.send_batched_chat_completions_http_request(
+        {
+            "json": {
+                "model": omni_server.model,
+                "messages": messages,
+                **sad_opts,
+            },
+        },
+    )
+    assert responses and len(responses) == 1
+    resp = responses[0]
+    assert not responses[0].success
+    assert resp.status_code == 400

tests/entrypoints/openai/test_batch_chat_completions.py

@@ -0,0 +1,150 @@
+import asyncio
+from unittest.mock import AsyncMock
+
+import pytest
+from openai.types.chat import ChatCompletionUserMessageParam
+from openai.types.chat.chat_completion_audio import ChatCompletionAudio
+from starlette.requests import Request
+from vllm.entrypoints.chat_utils import ChatCompletionMessageParam
+from vllm.entrypoints.openai.chat_completion.protocol import (
+    BatchChatCompletionRequest,
+    ChatCompletionResponse,
+    ChatCompletionResponseChoice,
+    ChatMessage,
+)
+from vllm.entrypoints.openai.engine.protocol import ErrorResponse, UsageInfo
+
+from vllm_omni.entrypoints.openai.batch_serving import OmniOpenAIServingChatBatch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+MSG_BATCH: list[list[ChatCompletionMessageParam]] = [
+    [ChatCompletionUserMessageParam(role="user", content="What color is the sky?")],
+    [ChatCompletionUserMessageParam(role="user", content="What is 2+2?")],
+]
+collapse = OmniOpenAIServingChatBatch._maybe_collapse_choices
+
+
+# Helpers for testing creating a __new__ serving instance and ensuring IDs are correct
+def _make_raw_request(headers: list[tuple[bytes, bytes]]):
+    scope = {
+        "type": "http",
+        "method": "POST",
+        "path": "/v1/chat/completions/batch",
+        "headers": headers,
+    }
+    return Request(scope)
+
+
+def _make_completion(text):
+    return ChatCompletionResponse(
+        model="test-model",
+        choices=[
+            ChatCompletionResponseChoice(
+                index=0,
+                message=ChatMessage(role="assistant", content=text),
+            )
+        ],
+        usage=UsageInfo(prompt_tokens=0, completion_tokens=0, total_tokens=0),
+    )
+
+
+def _make_handler(responses):
+    handler = OmniOpenAIServingChatBatch.__new__(OmniOpenAIServingChatBatch)
+    handler.create_chat_completion = AsyncMock(side_effect=responses)
+    return handler
+
+
+# Helpers for creating choices of different modality types
+def _text_choice(text="hello"):
+    return ChatCompletionResponseChoice(
+        index=0,
+        message=ChatMessage(role="assistant", content=text),
+    )
+
+
+def _audio_choice():
+    audio = ChatCompletionAudio(id="a1", data="base64audio", expires_at=0, transcript="")
+    return ChatCompletionResponseChoice(
+        index=0,
+        message=ChatMessage(role="assistant", content=None, audio=audio),
+    )
+
+
+def _image_choice():
+    content = [{"type": "image_url", "image_url": {"url": "data:image/png;base64,abc"}}]
+    msg = ChatMessage.model_construct(role="assistant")
+    object.__setattr__(msg, "content", content)
+    return ChatCompletionResponseChoice.model_construct(
+        index=0,
+        message=msg,
+    )
+
+
+def _get_subrequest_ids(handler):
+    return [call.args[0].request_id for call in handler.create_chat_completion.call_args_list]
+
+
+@pytest.mark.parametrize("has_header", [True, False])
+def test_subrequest_ids_are_unique(has_header: bool):
+    """Ensure that a submitted batch request creates unique subrequest IDs"""
+    header = None if not has_header else b"user-123"
+    headers = [] if not header else [(b"x-request-id", header)]
+
+    handler = _make_handler([_make_completion("hello"), _make_completion("world")])
+    request = BatchChatCompletionRequest(
+        model="test-model",
+        messages=MSG_BATCH,
+        seed=42,
+    )
+    result = asyncio.run(handler.create_batch_chat_completion(request, _make_raw_request(headers)))
+    assert not isinstance(result, ErrorResponse)
+    assert result.id.startswith("chatcmpl-batch")
+
+    # Ensure request IDs are unique and that if we had a header, it's in all reqs
+    req_ids = _get_subrequest_ids(handler)
+    assert len(set(req_ids)) == len(req_ids) == 2
+
+    if header is not None:
+        assert all(header.decode("utf-8") in req_id for req_id in req_ids)
+
+
+### Tests for ensuring that the (currently batch specific) choice collapsing is correct
+def test_single_text_passthrough():
+    result = collapse([_text_choice()])
+    assert result.message.content == "hello"
+
+
+def test_single_image_passthrough():
+    result = collapse([_image_choice()])
+    assert isinstance(result.message.content, list)
+    assert result.message.content[0]["type"] == "image_url"
+
+
+def test_text_plus_audio():
+    result = collapse([_text_choice(), _audio_choice()])
+    assert result.message.content == "hello"
+    assert result.message.audio.data == "base64audio"
+
+
+def test_audio_plus_text_order_independent():
+    result = collapse([_audio_choice(), _text_choice()])
+    assert result.message.content == "hello"
+    assert result.message.audio.data == "base64audio"
+
+
+def test_image_plus_audio():
+    result = collapse([_image_choice(), _audio_choice()])
+    assert isinstance(result.message.content, list)
+    assert result.message.content[0]["type"] == "image_url"
+    assert result.message.audio.data == "base64audio"
+
+
+def test_two_content_choices_raises():
+    with pytest.raises(ValueError, match="Multiple content choices cannot be set"):
+        collapse([_text_choice(), _text_choice()])
+
+
+def test_three_choices_raises():
+    with pytest.raises(ValueError, match="got 3"):
+        collapse([_text_choice(), _audio_choice(), _text_choice()])

tests/helpers/runtime.py

@@ -1252,6 +1252,28 @@ class OpenAIClientHandler:
         )
         return [resp]
 
+    def send_batched_chat_completions_http_request(
+        self,
+        request_config: dict[str, Any],
+        *,
+        err_code: int | tuple[int, ...] | list[int] | None = None,
+        err_message: str | tuple[str, ...] | list[str] | None = None,
+    ) -> list[HttpResponse]:
+        """Post to batched chat completions."""
+        cfg = _merge_http_expectation_kwargs(
+            request_config,
+            err_code=err_code,
+            err_message=err_message,
+        )
+        r = self._post_json_endpoint("/v1/chat/completions/batch", cfg, default_timeout=120.0)
+        resp = self._http_response_from_requests(r)
+        assert_http_error(
+            resp,
+            err_code=cfg.get("err_code"),
+            err_message=cfg.get("err_message"),
+        )
+        return [resp]
+
     def send_completions_http_request(
         self,
         request_config: dict[str, Any],