| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import subprocess |
| 4 | import tempfile |
| 5 | import time |
| 6 | from pathlib import Path |
| 7 | |
| 8 | import pytest |
| 9 | import requests |
| 10 | import urllib3 |
| 11 | |
| 12 | from ..utils import RemoteOpenAIServer |
| 13 | |
| 14 | MODEL_NAME = "meta-llama/Llama-3.2-1B-Instruct" |
| 15 | |
| 16 | |
| 17 | def generate_self_signed_cert(cert_dir: Path) -> tuple[Path, Path]: |
| 18 | """Generate a self-signed certificate for testing.""" |
| 19 | cert_file = cert_dir / "cert.pem" |
| 20 | key_file = cert_dir / "key.pem" |
| 21 | |
| 22 | # Generate self-signed certificate using openssl |
| 23 | subprocess.run( |
| 24 | [ |
| 25 | "openssl", |
| 26 | "req", |
| 27 | "-x509", |
| 28 | "-newkey", |
| 29 | "rsa:2048", |
| 30 | "-keyout", |
| 31 | str(key_file), |
| 32 | "-out", |
| 33 | str(cert_file), |
| 34 | "-days", |
| 35 | "1", |
| 36 | "-nodes", |
| 37 | "-subj", |
| 38 | "/CN=localhost", |
| 39 | ], |
| 40 | check=True, |
| 41 | capture_output=True, |
| 42 | ) |
| 43 | return cert_file, key_file |
| 44 | |
| 45 | |
| 46 | class RemoteOpenAIServerSSL(RemoteOpenAIServer): |
| 47 | """RemoteOpenAIServer subclass that supports SSL with self-signed certs.""" |
| 48 | |
| 49 | @property |
| 50 | def url_root(self) -> str: |
| 51 | return f"https://{self.host}:{self.port}" |
| 52 | |
| 53 | def _wait_for_server(self, *, url: str, timeout: float): |
| 54 | """Override to use HTTPS with SSL verification disabled.""" |
| 55 | # Suppress InsecureRequestWarning for self-signed certs |
| 56 | urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) |
| 57 | |
| 58 | start = time.time() |
| 59 | while True: |
| 60 | try: |
| 61 | if requests.get(url, verify=False).status_code == 200: |
| 62 | break |
| 63 | except Exception: |
| 64 | result = self._poll() |
| 65 | if result is not None and result != 0: |
| 66 | raise RuntimeError("Server exited unexpectedly.") from None |
| 67 | |
| 68 | time.sleep(0.5) |
| 69 | if time.time() - start > timeout: |
| 70 | raise RuntimeError("Server failed to start in time.") from None |
| 71 | |
| 72 | |
| 73 | @pytest.fixture(scope="function") |
| 74 | def server(): |
| 75 | args = ["--max-model-len", "1024", "--enforce-eager", "--load-format", "dummy"] |
| 76 | |
| 77 | with RemoteOpenAIServer(MODEL_NAME, args) as remote_server: |
| 78 | yield remote_server |
| 79 | |
| 80 | |
| 81 | @pytest.fixture(scope="function") |
| 82 | def ssl_server(): |
| 83 | """Start a vLLM server with SSL enabled using a self-signed certificate.""" |
| 84 | with tempfile.TemporaryDirectory() as cert_dir: |
| 85 | cert_file, key_file = generate_self_signed_cert(Path(cert_dir)) |
| 86 | args = [ |
| 87 | "--max-model-len", |
| 88 | "1024", |
| 89 | "--enforce-eager", |
| 90 | "--load-format", |
| 91 | "dummy", |
| 92 | "--ssl-certfile", |
| 93 | str(cert_file), |
| 94 | "--ssl-keyfile", |
| 95 | str(key_file), |
| 96 | ] |
| 97 | |
| 98 | with RemoteOpenAIServerSSL(MODEL_NAME, args) as remote_server: |
| 99 | yield remote_server |
| 100 | |
| 101 | |
| 102 | @pytest.mark.benchmark |
| 103 | def test_bench_serve(server): |
| 104 | # Test default model detection and input/output len |
| 105 | command = [ |
| 106 | "vllm", |
| 107 | "bench", |
| 108 | "serve", |
| 109 | "--host", |
| 110 | server.host, |
| 111 | "--port", |
| 112 | str(server.port), |
| 113 | "--input-len", |
| 114 | "32", |
| 115 | "--output-len", |
| 116 | "4", |
| 117 | "--num-prompts", |
| 118 | "5", |
| 119 | ] |
| 120 | result = subprocess.run(command, capture_output=True, text=True) |
| 121 | print(result.stdout) |
| 122 | print(result.stderr) |
| 123 | |
| 124 | assert result.returncode == 0, f"Benchmark failed: {result.stderr}" |
| 125 | |
| 126 | |
| 127 | @pytest.mark.benchmark |
| 128 | def test_bench_serve_insecure(ssl_server): |
| 129 | """Test --insecure flag with an HTTPS server using a self-signed certificate.""" |
| 130 | base_url = f"https://{ssl_server.host}:{ssl_server.port}" |
| 131 | command = [ |
| 132 | "vllm", |
| 133 | "bench", |
| 134 | "serve", |
| 135 | "--base-url", |
| 136 | base_url, |
| 137 | "--input-len", |
| 138 | "32", |
| 139 | "--output-len", |
| 140 | "4", |
| 141 | "--num-prompts", |
| 142 | "5", |
| 143 | "--insecure", |
| 144 | ] |
| 145 | result = subprocess.run(command, capture_output=True, text=True) |
| 146 | print(result.stdout) |
| 147 | print(result.stderr) |
| 148 | |
| 149 | assert result.returncode == 0, f"Benchmark failed: {result.stderr}" |
| 150 | |
| 151 | |
| 152 | @pytest.mark.benchmark |
| 153 | def test_bench_serve_chat(server): |
| 154 | command = [ |
| 155 | "vllm", |
| 156 | "bench", |
| 157 | "serve", |
| 158 | "--model", |
| 159 | MODEL_NAME, |
| 160 | "--host", |
| 161 | server.host, |
| 162 | "--port", |
| 163 | str(server.port), |
| 164 | "--dataset-name", |
| 165 | "random", |
| 166 | "--random-input-len", |
| 167 | "32", |
| 168 | "--random-output-len", |
| 169 | "4", |
| 170 | "--num-prompts", |
| 171 | "5", |
| 172 | "--endpoint", |
| 173 | "/v1/chat/completions", |
| 174 | "--backend", |
| 175 | "openai-chat", |
| 176 | ] |
| 177 | result = subprocess.run(command, capture_output=True, text=True) |
| 178 | print(result.stdout) |
| 179 | print(result.stderr) |
| 180 | |
| 181 | assert result.returncode == 0, f"Benchmark failed: {result.stderr}" |