| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import asyncio |
| 5 | from types import SimpleNamespace |
| 6 | from unittest.mock import AsyncMock |
| 7 | |
| 8 | import pytest |
| 9 | import requests |
| 10 | from fastapi import FastAPI |
| 11 | from fastapi.testclient import TestClient |
| 12 | from prometheus_client import CollectorRegistry |
| 13 | from prometheus_client.parser import text_string_to_metric_families |
| 14 | |
| 15 | from tests.utils import RemoteOpenAIServer |
| 16 | from vllm.entrypoints.serve.dev.sleep import metrics as sleep_metrics |
| 17 | from vllm.entrypoints.serve.dev.sleep.api_router import attach_router |
| 18 | from vllm.entrypoints.serve.dev.sleep.metrics import SleepModeOperationMetrics |
| 19 | from vllm.entrypoints.serve.exception_handling.register import init_exception_handler |
| 20 | from vllm.entrypoints.serve.instrumentator.metrics import ( |
| 21 | attach_router as attach_metrics_router, |
| 22 | ) |
| 23 | from vllm.v1.metrics import prometheus as prometheus_metrics |
| 24 | |
| 25 | MODEL_NAME = "meta-llama/Llama-3.2-1B" |
| 26 | |
| 27 | |
| 28 | @pytest.fixture |
| 29 | def sleep_route_app(monkeypatch): |
| 30 | app = FastAPI() |
| 31 | app.state.args = SimpleNamespace(log_error_stack=False) |
| 32 | app.state.engine_client = AsyncMock() |
| 33 | metrics = SleepModeOperationMetrics(CollectorRegistry()) |
| 34 | monkeypatch.setattr( |
| 35 | "vllm.entrypoints.serve.dev.sleep.api_router.sleep_mode_operation_metrics", |
| 36 | lambda: metrics, |
| 37 | ) |
| 38 | |
| 39 | attach_router(app) |
| 40 | init_exception_handler(app) |
| 41 | return app, metrics |
| 42 | |
| 43 | |
| 44 | @pytest.mark.cpu_test |
| 45 | @pytest.mark.parametrize("level", [0, 1, 2]) |
| 46 | def test_sleep_route_response_and_engine_arguments(sleep_route_app, level): |
| 47 | app, _ = sleep_route_app |
| 48 | with TestClient(app) as client: |
| 49 | response = client.post("/sleep", params={"level": level}) |
| 50 | assert response.status_code == 200 |
| 51 | assert response.json() == {"status": "sleeping", "level": level} |
| 52 | app.state.engine_client.sleep.assert_awaited_once_with(level, "abort") |
| 53 | |
| 54 | |
| 55 | @pytest.mark.cpu_test |
| 56 | @pytest.mark.parametrize( |
| 57 | ("query", "expected_param"), |
| 58 | [ |
| 59 | ("level=invalid", "query.level"), |
| 60 | ("level=-1", "query.level"), |
| 61 | ("level=3", "query.level"), |
| 62 | ("mode=invalid", "query.mode"), |
| 63 | ], |
| 64 | ) |
| 65 | def test_sleep_route_rejects_invalid_query_before_dispatch( |
| 66 | sleep_route_app, query, expected_param |
| 67 | ): |
| 68 | app, metrics = sleep_route_app |
| 69 | with TestClient(app) as client: |
| 70 | response = client.post(f"/sleep?{query}") |
| 71 | assert response.status_code == 400 |
| 72 | assert response.json()["error"]["param"] == expected_param |
| 73 | app.state.engine_client.sleep.assert_not_awaited() |
| 74 | assert list(metrics.duration.collect()[0].samples) == [] |
| 75 | assert list(metrics.in_flight.collect()[0].samples) == [] |
| 76 | |
| 77 | |
| 78 | @pytest.mark.cpu_test |
| 79 | @pytest.mark.parametrize( |
| 80 | ("fully_awake", "tags", "expected"), |
| 81 | [ |
| 82 | (True, "", {"status": "awake", "tags": None}), |
| 83 | (False, "?tags=weights", {"status": "sleeping", "tags": ["weights"]}), |
| 84 | ], |
| 85 | ) |
| 86 | def test_wake_route_maps_engine_result_without_state_query( |
| 87 | sleep_route_app, fully_awake, tags, expected |
| 88 | ): |
| 89 | app, _ = sleep_route_app |
| 90 | app.state.engine_client.wake_up.return_value = fully_awake |
| 91 | with TestClient(app) as client: |
| 92 | response = client.post(f"/wake_up{tags}") |
| 93 | assert response.status_code == 200 |
| 94 | assert response.json() == expected |
| 95 | app.state.engine_client.wake_up.assert_awaited_once_with(expected["tags"]) |
| 96 | app.state.engine_client.is_sleeping.assert_not_awaited() |
| 97 | |
| 98 | |
| 99 | @pytest.mark.cpu_test |
| 100 | @pytest.mark.parametrize("fails", [False, True]) |
| 101 | def test_release_kv_cache_memory_route(sleep_route_app, fails): |
| 102 | app, metrics = sleep_route_app |
| 103 | release = app.state.engine_client.release_kv_cache_memory |
| 104 | if fails: |
| 105 | release.side_effect = RuntimeError("requires a completed pause first") |
| 106 | with TestClient(app, raise_server_exceptions=False) as client: |
| 107 | response = client.post("/release_kv_cache_memory") |
| 108 | |
| 109 | assert response.status_code == (500 if fails else 200) |
| 110 | if not fails: |
| 111 | assert response.json() == {"status": "kv_cache_released"} |
| 112 | release.assert_awaited_once_with() |
| 113 | assert metrics.in_flight.labels("release_kv_cache_memory")._value.get() == 0 |
| 114 | assert _duration_count(metrics, "release_kv_cache_memory") == 1 |
| 115 | |
| 116 | |
| 117 | @pytest.mark.cpu_test |
| 118 | @pytest.mark.parametrize("operation", ["sleep", "release_kv_cache_memory", "wake"]) |
| 119 | @pytest.mark.parametrize("exception", [RuntimeError, asyncio.CancelledError]) |
| 120 | def test_sleep_mode_recorder_tracks_duration_and_cleans_up(operation, exception): |
| 121 | metrics = SleepModeOperationMetrics(CollectorRegistry()) |
| 122 | with metrics.record(operation): |
| 123 | assert metrics.in_flight.labels(operation)._value.get() == 1 |
| 124 | with pytest.raises(exception), metrics.record(operation): |
| 125 | raise exception() |
| 126 | assert metrics.in_flight.labels(operation)._value.get() == 0 |
| 127 | assert metrics.duration.labels(operation)._sum.get() >= 0 |
| 128 | assert _duration_count(metrics, operation) == 2 |
| 129 | |
| 130 | |
| 131 | def _duration_count(metrics, operation): |
| 132 | return next( |
| 133 | sample.value |
| 134 | for sample in metrics.duration.collect()[0].samples |
| 135 | if sample.name.endswith("_count") and sample.labels == {"operation": operation} |
| 136 | ) |
| 137 | |
| 138 | |
| 139 | @pytest.mark.cpu_test |
| 140 | @pytest.mark.parametrize( |
| 141 | ("path", "method", "operation", "status", "fails"), |
| 142 | [ |
| 143 | ("/sleep", "POST", "sleep", 200, False), |
| 144 | ("/release_kv_cache_memory", "POST", "release_kv_cache_memory", 200, False), |
| 145 | ("/wake_up", "POST", "wake", 200, False), |
| 146 | ("/is_sleeping", "GET", None, 200, False), |
| 147 | ("/sleep?level=3", "POST", None, 400, False), |
| 148 | ("/release_kv_cache_memory", "POST", "release_kv_cache_memory", 500, True), |
| 149 | ], |
| 150 | ) |
| 151 | def test_sleep_routes_visible_on_production_metrics_endpoint( |
| 152 | monkeypatch, path, method, operation, status, fails |
| 153 | ): |
| 154 | registry = CollectorRegistry() |
| 155 | monkeypatch.delenv("PROMETHEUS_MULTIPROC_DIR", raising=False) |
| 156 | monkeypatch.setattr(sleep_metrics, "REGISTRY", registry) |
| 157 | monkeypatch.setattr(sleep_metrics, "_metrics", None) |
| 158 | monkeypatch.setattr(prometheus_metrics, "REGISTRY", registry) |
| 159 | |
| 160 | app = FastAPI() |
| 161 | app.state.args = SimpleNamespace(log_error_stack=False) |
| 162 | app.state.engine_client = AsyncMock() |
| 163 | app.state.engine_client.wake_up.return_value = True |
| 164 | app.state.engine_client.is_sleeping.return_value = False |
| 165 | if fails: |
| 166 | app.state.engine_client.release_kv_cache_memory.side_effect = RuntimeError( |
| 167 | "engine failed" |
| 168 | ) |
| 169 | attach_router(app) |
| 170 | init_exception_handler(app) |
| 171 | attach_metrics_router(app) |
| 172 | |
| 173 | with TestClient(app, raise_server_exceptions=False) as client: |
| 174 | assert client.request(method, path).status_code == status |
| 175 | response = client.get("/metrics") |
| 176 | |
| 177 | assert response.status_code == 200 |
| 178 | samples = [ |
| 179 | sample |
| 180 | for family in text_string_to_metric_families(response.text) |
| 181 | for sample in family.samples |
| 182 | ] |
| 183 | assert any( |
| 184 | sample.name == "http_requests_total" |
| 185 | and sample.labels |
| 186 | == { |
| 187 | "handler": path.split("?")[0], |
| 188 | "method": method, |
| 189 | "status": f"{status // 100}xx", |
| 190 | } |
| 191 | and sample.value == 1 |
| 192 | for sample in samples |
| 193 | ) |
| 194 | assert not any("sleep_mode_operations_total" in sample.name for sample in samples) |
| 195 | assert any( |
| 196 | sample.name == "vllm:sleep_mode_operations_in_flight" |
| 197 | and sample.labels == {"operation": operation} |
| 198 | and sample.value == 0 |
| 199 | for sample in samples |
| 200 | ) == (operation is not None) |
| 201 | assert any( |
| 202 | sample.name == "vllm:sleep_mode_operation_duration_seconds_count" |
| 203 | and sample.labels == {"operation": operation} |
| 204 | and sample.value == 1 |
| 205 | for sample in samples |
| 206 | ) == (operation is not None) |
| 207 | |
| 208 | |
| 209 | def test_sleep_mode(): |
| 210 | # dtype, max-len etc set so that this can run in CI |
| 211 | args = [ |
| 212 | "--dtype", |
| 213 | "bfloat16", |
| 214 | "--max-model-len", |
| 215 | "8192", |
| 216 | "--max-num-seqs", |
| 217 | "128", |
| 218 | "--enable-sleep-mode", |
| 219 | ] |
| 220 | |
| 221 | with RemoteOpenAIServer( |
| 222 | MODEL_NAME, |
| 223 | args, |
| 224 | env_dict={"VLLM_SERVER_DEV_MODE": "1", "CUDA_VISIBLE_DEVICES": "0"}, |
| 225 | ) as remote_server: |
| 226 | response = requests.post(remote_server.url_for("sleep"), params={"level": "1"}) |
| 227 | assert response.status_code == 200 |
| 228 | response = requests.get(remote_server.url_for("is_sleeping")) |
| 229 | assert response.status_code == 200 |
| 230 | assert response.json().get("is_sleeping") is True |
| 231 | |
| 232 | # check sleep metrics |
| 233 | response = requests.get(remote_server.url_for("metrics")) |
| 234 | assert response.status_code == 200 |
| 235 | awake, weights_offloaded, discard_all = _get_sleep_metrics_from_api(response) |
| 236 | assert awake == 0 |
| 237 | assert weights_offloaded == 1 |
| 238 | assert discard_all == 0 |
| 239 | |
| 240 | response = requests.post(remote_server.url_for("wake_up")) |
| 241 | assert response.status_code == 200 |
| 242 | response = requests.get(remote_server.url_for("is_sleeping")) |
| 243 | assert response.status_code == 200 |
| 244 | assert response.json().get("is_sleeping") is False |
| 245 | |
| 246 | # check sleep metrics |
| 247 | response = requests.get(remote_server.url_for("metrics")) |
| 248 | assert response.status_code == 200 |
| 249 | awake, weights_offloaded, discard_all = _get_sleep_metrics_from_api(response) |
| 250 | assert awake == 1 |
| 251 | assert weights_offloaded == 0 |
| 252 | assert discard_all == 0 |
| 253 | |
| 254 | # test wake up with tags |
| 255 | response = requests.post(remote_server.url_for("sleep"), params={"level": "1"}) |
| 256 | assert response.status_code == 200 |
| 257 | |
| 258 | response = requests.post( |
| 259 | remote_server.url_for("wake_up"), params={"tags": ["weights"]} |
| 260 | ) |
| 261 | assert response.status_code == 200 |
| 262 | |
| 263 | # Partial wake keeps the engine sleeping. |
| 264 | response = requests.get(remote_server.url_for("is_sleeping")) |
| 265 | assert response.status_code == 200 |
| 266 | assert response.json().get("is_sleeping") is True |
| 267 | |
| 268 | response = requests.post( |
| 269 | remote_server.url_for("wake_up"), params={"tags": ["kv_cache"]} |
| 270 | ) |
| 271 | assert response.status_code == 200 |
| 272 | |
| 273 | response = requests.get(remote_server.url_for("is_sleeping")) |
| 274 | assert response.status_code == 200 |
| 275 | assert response.json().get("is_sleeping") is False |
| 276 | |
| 277 | # check sleep metrics |
| 278 | response = requests.get(remote_server.url_for("metrics")) |
| 279 | assert response.status_code == 200 |
| 280 | awake, weights_offloaded, discard_all = _get_sleep_metrics_from_api(response) |
| 281 | assert awake == 1 |
| 282 | assert weights_offloaded == 0 |
| 283 | assert discard_all == 0 |
| 284 | |
| 285 | |
| 286 | def _get_sleep_metrics_from_api(response: requests.Response): |
| 287 | """Return (awake, weights_offloaded, discard_all).""" |
| 288 | awake, weights_offloaded, discard_all = None, None, None |
| 289 | |
| 290 | for family in text_string_to_metric_families(response.text): |
| 291 | if family.name == "vllm:engine_sleep_state": |
| 292 | for sample in family.samples: |
| 293 | if sample.name == "vllm:engine_sleep_state": |
| 294 | for label_name, label_value in sample.labels.items(): |
| 295 | if label_value == "awake": |
| 296 | awake = sample.value |
| 297 | elif label_value == "weights_offloaded": |
| 298 | weights_offloaded = sample.value |
| 299 | elif label_value == "discard_all": |
| 300 | discard_all = sample.value |
| 301 | |
| 302 | assert awake is not None |
| 303 | assert weights_offloaded is not None |
| 304 | assert discard_all is not None |
| 305 | |
| 306 | return awake, weights_offloaded, discard_all |