[Core][Bugfix][Perf] Introduce MQLLMEngine to avoid asyncio OH (#8157)

Co-authored-by: Nick Hill <nickhill@us.ibm.com>
Co-authored-by: rshaw@neuralmagic.com <rshaw@neuralmagic.com>
Co-authored-by: Robert Shaw <114415538+robertgshaw2-neuralmagic@users.noreply.github.com>
Co-authored-by: Simon Mo <simon.mo@hey.com>

tests/entrypoints/openai/rpc/test_zmq_client.py

@@ -1,120 +0,0 @@
-import asyncio
-import tempfile
-import unittest
-import unittest.mock
-import uuid
-
-import pytest
-import pytest_asyncio
-
-from vllm.engine.async_llm_engine import AsyncLLMEngine
-from vllm.entrypoints.openai.rpc.client import (AsyncEngineRPCClient,
-                                                RPCClientClosedError)
-from vllm.entrypoints.openai.rpc.server import AsyncEngineRPCServer
-
-
-@pytest.fixture(scope="function")
-def tmp_socket():
-    with tempfile.TemporaryDirectory() as td:
-        yield f"ipc://{td}/{uuid.uuid4()}"
-
-
-@pytest_asyncio.fixture(scope="function")
-async def dummy_server(tmp_socket, monkeypatch):
-    dummy_engine = unittest.mock.AsyncMock()
-
-    def dummy_engine_builder(*args, **kwargs):
-        return dummy_engine
-
-    with monkeypatch.context() as m:
-        m.setattr(AsyncLLMEngine, "from_engine_args", dummy_engine_builder)
-        server = AsyncEngineRPCServer(None, None, rpc_path=tmp_socket)
-
-    loop = asyncio.get_running_loop()
-    server_task = loop.create_task(server.run_server_loop())
-
-    try:
-        yield server
-    finally:
-        server_task.cancel()
-        server.cleanup()
-
-
-@pytest_asyncio.fixture(scope="function")
-async def client(tmp_socket):
-    client = AsyncEngineRPCClient(rpc_path=tmp_socket)
-    # Sanity check: the server is connected
-    await client._wait_for_server_rpc()
-
-    try:
-        yield client
-    finally:
-        client.close()
-
-
-@pytest.mark.asyncio
-async def test_client_data_methods_use_timeouts(monkeypatch, dummy_server,
-                                                client: AsyncEngineRPCClient):
-    with monkeypatch.context() as m:
-        # Make the server _not_ reply with a model config
-        m.setattr(dummy_server, "get_config", lambda x: None)
-        m.setattr(client, "_data_timeout", 10)
-
-        # And ensure the task completes anyway
-        # (client.setup() invokes server.get_config())
-        client_task = asyncio.get_running_loop().create_task(client.setup())
-        with pytest.raises(TimeoutError, match="Server didn't reply within"):
-            await asyncio.wait_for(client_task, timeout=0.05)
-
-
-@pytest.mark.asyncio
-async def test_client_aborts_use_timeouts(monkeypatch, dummy_server,
-                                          client: AsyncEngineRPCClient):
-    with monkeypatch.context() as m:
-        # Hang all abort requests
-        m.setattr(dummy_server, "abort", lambda x: None)
-        m.setattr(client, "_data_timeout", 10)
-
-        # The client should suppress timeouts on `abort`s
-        # and return normally, assuming the server will eventually
-        # abort the request.
-        client_task = asyncio.get_running_loop().create_task(
-            client.abort("test request id"))
-        await asyncio.wait_for(client_task, timeout=0.05)
-
-
-@pytest.mark.asyncio
-async def test_client_data_methods_reraise_exceptions(
-        monkeypatch, dummy_server, client: AsyncEngineRPCClient):
-    with monkeypatch.context() as m:
-        # Make the server raise some random exception
-        exception = RuntimeError("Client test exception")
-
-        def raiser():
-            raise exception
-
-        m.setattr(dummy_server.engine, "get_model_config", raiser)
-        m.setattr(client, "_data_timeout", 10)
-
-        client_task = asyncio.get_running_loop().create_task(client.setup())
-        # And ensure the task completes, raising the exception
-        with pytest.raises(RuntimeError, match=str(exception)):
-            await asyncio.wait_for(client_task, timeout=0.05)
-
-
-@pytest.mark.asyncio
-async def test_client_errors_after_closing(monkeypatch, dummy_server,
-                                           client: AsyncEngineRPCClient):
-
-    client.close()
-
-    # Healthchecks and generate requests will fail with explicit errors
-    with pytest.raises(RPCClientClosedError):
-        await client.check_health()
-    with pytest.raises(RPCClientClosedError):
-        async for _ in client.generate(None, None, None):
-            pass
-
-    # But no-ops like aborting will pass
-    await client.abort("test-request-id")
-    await client.do_log_stats()

tests/entrypoints/openai/test_accuracy.py

@@ -18,38 +18,32 @@ TASK = "gsm8k"
 FILTER = "exact_match,strict-match"
 RTOL = 0.03
 EXPECTED_VALUE = 0.58
+DEFAULT_ARGS = ["--max-model-len", "4096", "--disable-log-requests"]
+MORE_ARGS_LIST = [["--enable-chunked-prefill"], ["--num-scheduler-steps", "8"]]
 
 
-@pytest.fixture(scope="module")
-def server():
-    args = [
-        "--max-model-len", "4096", "--enable-chunked-prefill",
-        "--disable-log-requests", "--enforce-eager"
-    ]
+@pytest.mark.parametrize("more_args", MORE_ARGS_LIST)
+def test_lm_eval_accuracy(more_args):
+    args = list(DEFAULT_ARGS)
+    args.extend(more_args)
+
+    print(f"Running with: {args}")
 
     with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
-        yield remote_server
+        url = f"{remote_server.url_for('v1')}/completions"
 
+        model_args = (
+            f"model={MODEL_NAME},"
+            f"base_url={url},"
+            f"num_concurrent={NUM_CONCURRENT},tokenized_requests=False")
 
-@pytest.fixture(scope="module")
-def server_data(server):
-    return {
-        "url": f"{server.url_for('v1')}/completions",
-    }
+        results = lm_eval.simple_evaluate(
+            model="local-completions",
+            model_args=model_args,
+            tasks=TASK,
+        )
 
-
-def test_lm_eval_accuracy(server_data):
-    model_args = (f"model={MODEL_NAME},"
-                  f"base_url={server_data['url']},"
-                  f"num_concurrent={NUM_CONCURRENT},tokenized_requests=False")
-
-    results = lm_eval.simple_evaluate(
-        model="local-completions",
-        model_args=model_args,
-        tasks=TASK,
-    )
-
-    measured_value = results["results"][TASK][FILTER]
-    assert (measured_value - RTOL < EXPECTED_VALUE
-            and measured_value + RTOL > EXPECTED_VALUE
-            ), f"Expected: {EXPECTED_VALUE} |  Measured: {measured_value}"
+        measured_value = results["results"][TASK][FILTER]
+        assert (measured_value - RTOL < EXPECTED_VALUE
+                and measured_value + RTOL > EXPECTED_VALUE
+                ), f"Expected: {EXPECTED_VALUE} |  Measured: {measured_value}"

tests/entrypoints/openai/test_chat_template.py

@@ -0,0 +1,101 @@
+import pytest
+
+from vllm.entrypoints.chat_utils import (apply_hf_chat_template,
+                                         load_chat_template)
+from vllm.entrypoints.openai.protocol import ChatCompletionRequest
+from vllm.transformers_utils.tokenizer import get_tokenizer
+
+from ...utils import VLLM_PATH
+
+chatml_jinja_path = VLLM_PATH / "examples/template_chatml.jinja"
+assert chatml_jinja_path.exists()
+
+# Define models, templates, and their corresponding expected outputs
+MODEL_TEMPLATE_GENERATON_OUTPUT = [
+    ("facebook/opt-125m", chatml_jinja_path, True, """<|im_start|>user
+Hello<|im_end|>
+<|im_start|>assistant
+Hi there!<|im_end|>
+<|im_start|>user
+What is the capital of<|im_end|>
+<|im_start|>assistant
+"""),
+    ("facebook/opt-125m", chatml_jinja_path, False, """<|im_start|>user
+Hello<|im_end|>
+<|im_start|>assistant
+Hi there!<|im_end|>
+<|im_start|>user
+What is the capital of""")
+]
+
+TEST_MESSAGES = [
+    {
+        'role': 'user',
+        'content': 'Hello'
+    },
+    {
+        'role': 'assistant',
+        'content': 'Hi there!'
+    },
+    {
+        'role': 'user',
+        'content': 'What is the capital of'
+    },
+]
+
+
+def test_load_chat_template():
+    # Testing chatml template
+    template_content = load_chat_template(chat_template=chatml_jinja_path)
+
+    # Test assertions
+    assert template_content is not None
+    # Hard coded value for template_chatml.jinja
+    assert template_content == """{% for message in messages %}{{'<|im_start|>' + message['role'] + '\\n' + message['content']}}{% if (loop.last and add_generation_prompt) or not loop.last %}{{ '<|im_end|>' + '\\n'}}{% endif %}{% endfor %}
+{% if add_generation_prompt and messages[-1]['role'] != 'assistant' %}{{ '<|im_start|>assistant\\n' }}{% endif %}"""  # noqa: E501
+
+
+def test_no_load_chat_template_filelike():
+    # Testing chatml template
+    template = "../../examples/does_not_exist"
+
+    with pytest.raises(ValueError, match="looks like a file path"):
+        load_chat_template(chat_template=template)
+
+
+def test_no_load_chat_template_literallike():
+    # Testing chatml template
+    template = "{{ messages }}"
+
+    template_content = load_chat_template(chat_template=template)
+
+    assert template_content == template
+
+
+@pytest.mark.parametrize(
+    "model,template,add_generation_prompt,expected_output",
+    MODEL_TEMPLATE_GENERATON_OUTPUT)
+def test_get_gen_prompt(model, template, add_generation_prompt,
+                        expected_output):
+    # Initialize the tokenizer
+    tokenizer = get_tokenizer(tokenizer_name=model)
+    template_content = load_chat_template(chat_template=template)
+
+    # Create a mock request object using keyword arguments
+    mock_request = ChatCompletionRequest(
+        model=model,
+        messages=TEST_MESSAGES,
+        add_generation_prompt=add_generation_prompt)
+
+    # Call the function and get the result
+    result = apply_hf_chat_template(
+        tokenizer,
+        conversation=mock_request.messages,
+        chat_template=mock_request.chat_template or template_content,
+        add_generation_prompt=mock_request.add_generation_prompt,
+    )
+
+    # Test assertion
+    assert result == expected_output, (
+        f"The generated prompt does not match the expected output for "
+        f"model {model} and template {template}")

tests/entrypoints/openai/test_mp_api_server.py

@@ -1,40 +0,0 @@
-import time
-
-import pytest
-
-from vllm.entrypoints.openai.api_server import build_async_engine_client
-from vllm.entrypoints.openai.cli_args import make_arg_parser
-from vllm.utils import FlexibleArgumentParser
-
-
-@pytest.mark.asyncio
-async def test_mp_crash_detection():
-
-    parser = FlexibleArgumentParser(description="vLLM's remote OpenAI server.")
-    parser = make_arg_parser(parser)
-    args = parser.parse_args([])
-    # use an invalid tensor_parallel_size to trigger the
-    # error in the server
-    args.tensor_parallel_size = 65536
-
-    start = time.perf_counter()
-    async with build_async_engine_client(args):
-        pass
-    end = time.perf_counter()
-
-    assert end - start < 60, ("Expected vLLM to gracefully shutdown in <60s "
-                              "if there is an error in the startup.")
-
-
-@pytest.mark.asyncio
-async def test_mp_cuda_init():
-    # it should not crash, when cuda is initialized
-    # in the API server process
-    import torch
-    torch.cuda.init()
-    parser = FlexibleArgumentParser(description="vLLM's remote OpenAI server.")
-    parser = make_arg_parser(parser)
-    args = parser.parse_args([])
-
-    async with build_async_engine_client(args):
-        pass

tests/entrypoints/openai/test_serving_chat.py

@@ -4,7 +4,7 @@ from dataclasses import dataclass
 from unittest.mock import MagicMock
 
 from vllm.config import MultiModalConfig
-from vllm.engine.async_llm_engine import AsyncLLMEngine
+from vllm.engine.multiprocessing.client import MQLLMEngineClient
 from vllm.entrypoints.openai.protocol import ChatCompletionRequest
 from vllm.entrypoints.openai.serving_chat import OpenAIServingChat
 from vllm.transformers_utils.tokenizer import get_tokenizer
@@ -52,8 +52,9 @@ def test_async_serving_chat_init():
 
 
 def test_serving_chat_should_set_correct_max_tokens():
-    mock_engine = MagicMock(spec=AsyncLLMEngine)
+    mock_engine = MagicMock(spec=MQLLMEngineClient)
     mock_engine.get_tokenizer.return_value = get_tokenizer(MODEL_NAME)
+    mock_engine.errored = False
 
     serving_chat = OpenAIServingChat(mock_engine,
                                      MockModelConfig(),

tests/entrypoints/openai/test_serving_engine.py

@@ -4,7 +4,7 @@ from unittest.mock import MagicMock
 import pytest
 
 from vllm.config import ModelConfig
-from vllm.engine.protocol import AsyncEngineClient
+from vllm.engine.protocol import EngineClient
 from vllm.entrypoints.openai.protocol import (ErrorResponse,
                                               LoadLoraAdapterRequest,
                                               UnloadLoraAdapterRequest)
@@ -18,7 +18,7 @@ LORA_UNLOADING_SUCCESS_MESSAGE = (
 
 
 async def _async_serving_engine_init():
-    mock_engine_client = MagicMock(spec=AsyncEngineClient)
+    mock_engine_client = MagicMock(spec=EngineClient)
     mock_model_config = MagicMock(spec=ModelConfig)
     # Set the max_model_len attribute to avoid missing attribute
     mock_model_config.max_model_len = 2048

tests/entrypoints/openai/test_shutdown.py

@@ -44,5 +44,5 @@ async def test_shutdown_on_engine_failure(tmp_path):
                                                 prompt="Hello, my name is")
 
             # Now the server should shut down
-            return_code = remote_server.proc.wait(timeout=3)
+            return_code = remote_server.proc.wait(timeout=8)
             assert return_code is not None