vllm-omni/tests/edge/local/test_e2e_serve_text.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""M0 serving end-to-end: the OpenAI-compatible server on the local text plan.

``test_e2e_local_text.py`` drives ``LocalTextEngine`` in-process. This module
drives the *server* path instead -- ``python -m vllm_omni.entrypoints.cli.main
serve <model> --omni`` as a subprocess, through ``tests.helpers.runtime.OmniServer``
(the same launcher the online-serving suite uses) -- and talks to it over HTTP
with plain ``requests``: health, model listing, non-streaming chat, streaming
chat, and a client that disconnects mid-stream after which the server must
still answer. Everything runs on one local device; no network access.

Runs on the same checkpoint rule as the engine test (``VLLM_OMNI_LOCAL_TEST_MODEL``
overrides). Marked ``local_model``: it needs the real model and a real device.
"""

from __future__ import annotations

import json
import os
import time
from collections.abc import Iterator

import pytest
import requests

from tests.edge.local.test_e2e_local_text import MAX_MODEL_LEN, _find_model
from vllm_omni.edge.local.prompts import acceptance_prompts

pytestmark = [pytest.mark.local_model]

MAX_TOKENS = 64
REQUEST_TIMEOUT = 600  # s; the first request also pays any lazy warmup


def _chat_prompts() -> list[tuple[str, str]]:
    """The three chat-shaped acceptance prompts as (name, user text).

    The acceptance file stores fully rendered prompts; for the chat endpoint we
    send the user turn only and let the server apply the chat template.
    """
    rendered = dict(acceptance_prompts())
    out = []
    for name, text in rendered.items():
        if not name.startswith("chat_"):
            continue
        # The rendered form is "...<|User|>{text}<...end...><...start...><|Bot|>..."
        # Strip to the user turn; if the markers are not there, send the whole thing.
        user = text.split("<|User|>", 1)[-1]
        user = user.split("<｜end▁of▁sentence｜>", 1)[0]
        out.append((name, user.strip() or text))
    return out


@pytest.fixture(scope="module")
def model_dir() -> str:
    return _find_model()


@pytest.fixture(scope="module")
def server(model_dir: str) -> Iterator[tuple[str, str]]:
    """Start the Omni server on the local text plan's settings; yield (base_url, model)."""
    from tests.helpers.runtime import OmniServer

    env = {
        "VLLM_NO_USAGE_STATS": "1",
        "DO_NOT_TRACK": "1",
        "HF_HUB_OFFLINE": "1",
        # The same sampler setting LocalTextEngine applies to itself
        # (vllm_omni/edge/local/engine.py). Without it the serve path picks the
        # flashinfer top-k/top-p sampler, whose JIT needs nvcc at runtime; a
        # venv without a CUDA toolkit (the WSL reference venv) dies in
        # profile_run with "Could not find nvcc". The local plan never wants a
        # runtime compiler on the device it admits.
        "VLLM_USE_FLASHINFER_SAMPLER": "0",
    }
    serve_args = [
        "--max-model-len",
        str(MAX_MODEL_LEN),
        "--max-num-seqs",
        "2",
        "--enforce-eager",
        "--gpu-memory-utilization",
        "0.55",
        "--stage-init-timeout",
        "600",
        "--init-timeout",
        "900",
        "--disable-log-stats",
    ]
    srv = OmniServer(model_dir, serve_args, env_dict=env, use_omni=True)
    with srv:
        yield f"http://{srv.host}:{srv.port}", model_dir


def _wait_healthy(base_url: str, timeout: float = 300) -> None:
    deadline = time.time() + timeout
    last = None
    while time.time() < deadline:
        try:
            r = requests.get(f"{base_url}/health", timeout=5)
            if r.status_code == 200:
                return
            last = r.status_code
        except requests.RequestException as e:  # server socket open but app not ready
            last = repr(e)
        time.sleep(2)
    pytest.fail(f"/health never returned 200 within {timeout}s (last: {last})")


def _chat(base_url: str, model: str, user: str, *, stream: bool, max_tokens: int = MAX_TOKENS):
    body = {
        "model": model,
        "messages": [{"role": "user", "content": user}],
        "max_tokens": max_tokens,
        "temperature": 0,
        "stream": stream,
    }
    return requests.post(
        f"{base_url}/v1/chat/completions", json=body, stream=stream, timeout=REQUEST_TIMEOUT
    )


def _sse_chunks(resp: requests.Response) -> Iterator[dict]:
    for raw in resp.iter_lines():
        if not raw:
            continue
        line = raw.decode("utf-8") if isinstance(raw, bytes) else raw
        if not line.startswith("data:"):
            continue
        payload = line[len("data:") :].strip()
        if payload == "[DONE]":
            return
        yield json.loads(payload)


def test_health_and_model_listing(server):
    base_url, model = server
    _wait_healthy(base_url)
    r = requests.get(f"{base_url}/v1/models", timeout=30)
    assert r.status_code == 200, r.text
    ids = [m["id"] for m in r.json()["data"]]
    assert model in ids, ids


@pytest.mark.parametrize("name,user", _chat_prompts(), ids=[n for n, _ in _chat_prompts()])
def test_chat_completion_reaches_max_tokens_or_stops(server, name, user):
    base_url, model = server
    _wait_healthy(base_url)
    r = _chat(base_url, model, user, stream=False)
    assert r.status_code == 200, r.text
    out = r.json()
    choice = out["choices"][0]
    assert choice["finish_reason"] in ("length", "stop"), choice
    assert out["usage"]["completion_tokens"] > 0, out["usage"]
    assert out["usage"]["completion_tokens"] <= MAX_TOKENS, out["usage"]
    assert choice["message"]["content"] or choice["message"].get("reasoning_content"), choice


def test_streaming_delivers_incremental_chunks_and_a_finish(server):
    base_url, model = server
    _wait_healthy(base_url)
    r = _chat(base_url, model, "Count from one to twenty in words.", stream=True)
    assert r.status_code == 200, r.text
    n_chunks, finish, text = 0, None, []
    for chunk in _sse_chunks(r):
        n_chunks += 1
        choice = chunk["choices"][0]
        delta = choice.get("delta", {})
        text.append(delta.get("content") or delta.get("reasoning_content") or "")
        if choice.get("finish_reason"):
            finish = choice["finish_reason"]
    assert n_chunks >= 4, n_chunks  # incremental, not one blob
    assert finish in ("length", "stop"), finish
    assert "".join(text).strip()


def test_a_client_disconnect_mid_stream_does_not_take_the_server_down(server):
    base_url, model = server
    _wait_healthy(base_url)
    r = _chat(base_url, model, "Write a long story about a lighthouse.", stream=True, max_tokens=256)
    assert r.status_code == 200, r.text
    seen = 0
    for _ in _sse_chunks(r):
        seen += 1
        if seen >= 5:
            break
    r.close()  # drop the connection with the request still generating
    assert seen >= 5
    # The server must still be healthy and still answer.
    _wait_healthy(base_url, timeout=60)
    r2 = _chat(base_url, model, "Say hello.", stream=False, max_tokens=16)
    assert r2.status_code == 200, r2.text
    assert r2.json()["usage"]["completion_tokens"] > 0