Adding max queue time parameter

docs/source/models/engine_args.rst

@@ -112,6 +112,10 @@ Below, you can find an explanation of every engine argument for vLLM:
 
     Maximum number of paddings in a batch.
 
+.. option:: --max-queue-length <size>
+
+    Maximum number of requests that can be present across all queues.
+
 .. option:: --disable-log-stats
 
     Disable logging statistics.

tests/async_engine/test_max_queue_length.py

@@ -0,0 +1,171 @@
+import asyncio
+import os
+import subprocess
+import time
+from vllm.logger import init_logger
+
+import sys
+from fastapi.responses import JSONResponse
+import pytest
+import requests
+# using Ray for overall ease of process management, parallel requests,
+# and debugging.
+import ray
+import openai  # use the official client for correctness check
+# downloading lora to test lora requests
+from huggingface_hub import snapshot_download
+
+logger = init_logger(__name__)
+
+MAX_SERVER_START_WAIT_S = 600  # wait for server to start for 60 seconds
+# any model with a chat template should work here
+MODEL_NAME = "HuggingFaceH4/zephyr-7b-beta"
+# technically this needs Mistral-7B-v0.1 as base, but we're not testing
+# generation quality here
+LORA_NAME = "typeof/zephyr-7b-beta-lora"
+
+pytestmark = pytest.mark.asyncio
+
+
+@ray.remote(num_gpus=1)
+class ServerRunner:
+
+    def __init__(self, args):
+        env = os.environ.copy()
+        env["PYTHONUNBUFFERED"] = "1"
+        self.proc = subprocess.Popen(
+            ["python3", "-m", "vllm.entrypoints.openai.api_server"] + args,
+            env=env,
+            stdout=sys.stdout,
+            stderr=sys.stderr,
+        )
+        self._wait_for_server()
+
+    def ready(self):
+        return True
+
+    def _wait_for_server(self):
+        # run health check
+        start = time.time()
+        while True:
+            try:
+                if requests.get(
+                        "http://localhost:8000/health").status_code == 200:
+                    break
+            except Exception as err:
+                if self.proc.poll() is not None:
+                    raise RuntimeError("Server exited unexpectedly.") from err
+
+                time.sleep(0.5)
+                if time.time() - start > MAX_SERVER_START_WAIT_S:
+                    raise RuntimeError(
+                        "Server failed to start in time.") from err
+
+    def __del__(self):
+        if hasattr(self, "proc"):
+            self.proc.terminate()
+
+
+@pytest.fixture(scope="session")
+def zephyr_lora_files():
+    return snapshot_download(repo_id=LORA_NAME)
+
+
+@pytest.fixture(scope="session")
+def server():
+    ray.init()
+    server_runner = ServerRunner.remote([
+        "--model",
+        MODEL_NAME,
+        # use half precision for speed and memory savings in CI environment
+        "--dtype",
+        "half",
+        "--max-model-len",
+        "1024",
+        "--enforce-eager",
+        "--max-num-seqs",
+        "1",
+        "--max-queue-length",
+        "3",
+        "--max-num-batched-tokens",
+        "2048",
+        "--gpu-memory-utilization",
+        "1"
+    ])
+    ray.get(server_runner.ready.remote())
+    yield server_runner
+    ray.shutdown()
+
+
+@pytest.fixture(scope="session")
+def client():
+    client = openai.AsyncOpenAI(base_url="http://localhost:8000/v1",
+                                api_key="token-abc123",
+                                max_retries=0)
+    yield client
+
+
+@pytest.mark.parametrize("model_name", [MODEL_NAME])
+async def test_max_queue_length(server, client: openai.AsyncOpenAI,
+                                model_name: str):
+    sample_chats = [[{
+        "role": "system",
+        "content": "You are a helpful assistant."
+    }, {
+        "role": "user",
+        "content": "Who won the world series in 2020?"
+    }],
+                    [{
+                        "role": "system",
+                        "content": "You are a helpful assistant."
+                    }, {
+                        "role": "user",
+                        "content": "Where was the 2020 world series played?"
+                    }],
+                    [{
+                        "role": "system",
+                        "content": "You are a helpful assistant."
+                    }, {
+                        "role": "user",
+                        "content": "How long did the 2020 world series last?"
+                    }],
+                    [{
+                        "role": "system",
+                        "content": "You are a helpful assistant."
+                    }, {
+                        "role":
+                        "user",
+                        "content":
+                        "What were some television viewership statistics?"
+                    }],
+                    [{
+                        "role": "system",
+                        "content": "You are a helpful assistant."
+                    }, {
+                        "role": "user",
+                        "content": "Why was the 2020 world series so popular?"
+                    }]]
+
+    async def make_api_call(sample_chat):
+        chat_completion = await client.chat.completions.create(
+            messages=sample_chat,
+            model=model_name,
+            temperature=0.8,
+            presence_penalty=0.2,
+            max_tokens=400,
+        )
+        return chat_completion
+
+    async def main():
+        coroutines = [
+            make_api_call(sample_chat) for sample_chat in sample_chats
+        ]
+
+        responses = await asyncio.gather(*coroutines, return_exceptions=True)
+
+        for response in responses:
+            logger.info(response)
+            if isinstance(response, JSONResponse):
+                assert response.status_code == 503
+
+    await main()

tests/engine/test_max_queue_length.py

@@ -0,0 +1,108 @@
+import pytest
+import argparse
+from typing import List, Tuple
+from vllm.logger import init_logger
+
+from vllm import EngineArgs, LLMEngine, SamplingParams, RequestOutput
+
+# initialize constants
+logger = init_logger(__name__)
+
+
+class QueueOverflowError(Exception):
+    pass
+
+
+@pytest.fixture
+def test_prompts() -> List[Tuple[str, SamplingParams]]:
+    """Create a list of test prompts with their sampling parameters."""
+    return [
+        ("A robot may not injure a human being",
+         SamplingParams(temperature=0.8,
+                        top_k=5,
+                        presence_penalty=0.2,
+                        ignore_eos=True,
+                        max_tokens=1000)),
+        ("To be or not to be,",
+         SamplingParams(temperature=0.8,
+                        top_k=5,
+                        presence_penalty=0.2,
+                        ignore_eos=True,
+                        max_tokens=1000)),
+        ("What is the meaning of life?",
+         SamplingParams(temperature=0.8,
+                        top_k=5,
+                        presence_penalty=0.2,
+                        ignore_eos=True,
+                        max_tokens=1000)),
+        ("It is only with the heart that one can see rightly",
+         SamplingParams(temperature=0.8,
+                        top_k=5,
+                        presence_penalty=0.2,
+                        ignore_eos=True,
+                        max_tokens=1000)),
+    ]
+
+
+def process_requests(engine: LLMEngine,
+                     test_prompts: List[Tuple[str, SamplingParams]]):
+    """Continuously process a list of prompts and handle the outputs."""
+    request_id = 0
+    while test_prompts or engine.has_unfinished_requests():
+        if test_prompts:
+            prompt, sampling_params = test_prompts.pop(0)
+            try:
+                engine.add_request(str(request_id), prompt, sampling_params)
+            except ValueError as e:
+                # Log error, cleanup, end test
+                logger.info(f"{e}")
+                for i in range(request_id):
+                    engine.abort_request(str(i))
+                raise QueueOverflowError(
+                    f"Queue exceeded max length: {e}") from e
+            request_id += 1
+
+        request_outputs: List[RequestOutput] = engine.step()
+
+        for request_output in request_outputs:
+            if request_output.finished:
+                print(request_output)
+
+
+@pytest.mark.parametrize(
+    "max_wait_q_len, expect_error",
+    [
+        (1, True),  # error expected 
+        (2, True),
+        (3, False),  # No error expected 
+        (4, False),
+    ])
+def test_max_queue_length(max_wait_q_len, expect_error, test_prompts):
+
+    # Setup engine with appropriate max_queue_length value
+    parser = argparse.ArgumentParser(
+        description='Demo on using the LLMEngine class directly')
+    parser = EngineArgs.add_cli_args(parser)
+    args_to_test = [
+        "--dtype",
+        "float",
+        '--max-num-seqs',
+        str(1), '--max-queue-length',
+        str(max_wait_q_len),
+        "--max-num-batched-tokens",
+        "2048",
+        "--gpu-memory-utilization",
+        "1",
+        "--max-model-len",
+        "1024",
+    ]
+    args = parser.parse_args(args_to_test)
+    engine_args = EngineArgs.from_cli_args(args)
+    engine = LLMEngine.from_engine_args(engine_args)
+
+    # Test engine against request
+    try:
+        process_requests(engine, test_prompts)
+        assert not expect_error, "QueueOverflowError did not occur as expected."
+    except QueueOverflowError as e:
+        assert expect_error, f" QueueOverflowError occurred as expected: {e}"

tests/engine/tmql.py

@@ -0,0 +1,103 @@
+import pytest
+import argparse
+from typing import List, Tuple
+from vllm.logger import init_logger
+
+from vllm import EngineArgs, LLMEngine, SamplingParams, RequestOutput
+
+# init variables
+max_wait_q_len = 2
+
+logger = init_logger(__name__)
+
+
+class QueueOverflowError(Exception):
+    pass
+
+
+def create_test_prompts() -> List[Tuple[str, SamplingParams]]:
+    """Create a list of test prompts with their sampling parameters."""
+    return [
+        ("A robot may not injure a human being",
+         SamplingParams(temperature=0.8,
+                        top_k=5,
+                        presence_penalty=0.2,
+                        ignore_eos=True,
+                        max_tokens=1000)),
+        ("To be or not to be,",
+         SamplingParams(temperature=0.8,
+                        top_k=5,
+                        presence_penalty=0.2,
+                        ignore_eos=True,
+                        max_tokens=1000)),
+        ("What is the meaning of life?",
+         SamplingParams(temperature=0.8,
+                        top_k=5,
+                        presence_penalty=0.2,
+                        ignore_eos=True,
+                        max_tokens=1000)),
+        ("It is only with the heart that one can see rightly",
+         SamplingParams(temperature=0.8,
+                        top_k=5,
+                        presence_penalty=0.2,
+                        ignore_eos=True,
+                        max_tokens=1000)),
+    ]
+
+
+def process_requests(engine: LLMEngine,
+                     test_prompts: List[Tuple[str, SamplingParams]]):
+    """Continuously process a list of prompts and handle the outputs."""
+    request_id = 0
+    # make sure to set something like max_num_seq to ONE
+    while test_prompts or engine.has_unfinished_requests():
+        if test_prompts:
+            prompt, sampling_params = test_prompts.pop(0)
+            try:
+                engine.add_request(str(request_id), prompt, sampling_params)
+            except ValueError as e:
+                # Log error, cleanup, end test
+                logger.info(f"{e}")
+                for i in range(request_id):
+                    engine.abort_request(str(i))
+                raise QueueOverflowError(
+                    f"Queue exceeded max length: {e}") from e
+            request_id += 1
+
+        request_outputs: List[RequestOutput] = engine.step()
+
+        for request_output in request_outputs:
+            if request_output.finished:
+                print(request_output)
+
+
+def initialize_engine(args: argparse.Namespace) -> LLMEngine:
+    """Initialize the LLMEngine from the command line arguments."""
+    engine_args = EngineArgs.from_cli_args(args)
+    return LLMEngine.from_engine_args(engine_args)
+
+
+def main(args: argparse.Namespace):
+    """Main function that sets up and runs the prompt processing."""
+    engine = initialize_engine(args)
+    test_prompts = create_test_prompts()
+    with pytest.raises(QueueOverflowError,
+                       match="Queue exceeded max length: .*"):
+        process_requests(engine, test_prompts)
+
+
+# def test_max_queue_length():
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(
+        description='Demo on using the LLMEngine class directly')
+    parser = EngineArgs.add_cli_args(parser)
+    args_to_test = [
+        '--max-num-seqs',
+        str(1), '--max-queue-length',
+        str(max_wait_q_len)
+    ]
+    args = parser.parse_args(args_to_test)
+    main(args)
+
+
+