| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import asyncio |
| 4 | import os |
| 5 | import subprocess |
| 6 | import sys |
| 7 | from contextlib import nullcontext |
| 8 | from pathlib import Path |
| 9 | from types import SimpleNamespace |
| 10 | |
| 11 | import pytest |
| 12 | from fastapi import FastAPI |
| 13 | from fastapi.testclient import TestClient |
| 14 | from prometheus_client import CollectorRegistry |
| 15 | |
| 16 | from vllm.entrypoints.serve.dev.rlhf import api_router |
| 17 | from vllm.entrypoints.serve.dev.rlhf import metrics as rlhf_metrics |
| 18 | from vllm.entrypoints.serve.dev.rlhf.metrics import WeightOperationMetrics |
| 19 | from vllm.entrypoints.serve.exception_handling.register import init_exception_handler |
| 20 | from vllm.v1.metrics.prometheus import get_prometheus_registry |
| 21 | |
| 22 | pytestmark = [pytest.mark.cpu_test, pytest.mark.skip_global_cleanup] |
| 23 | |
| 24 | DURATION = "vllm:rl_weight_update_operation_duration_seconds_count" |
| 25 | IN_FLIGHT = "vllm:rl_weight_update_operations_in_flight" |
| 26 | |
| 27 | |
| 28 | def sample(registry, name, operation): |
| 29 | return registry.get_sample_value(name, {"operation": operation}) |
| 30 | |
| 31 | |
| 32 | class Request: |
| 33 | def __init__(self, engine): |
| 34 | self.app = SimpleNamespace(state=SimpleNamespace(engine_client=engine)) |
| 35 | |
| 36 | |
| 37 | class Engine: |
| 38 | """Holds every engine call open until the test releases it.""" |
| 39 | |
| 40 | def __init__(self): |
| 41 | self.started = asyncio.Event() |
| 42 | self.release = asyncio.Event() |
| 43 | |
| 44 | def __getattr__(self, name): |
| 45 | async def call(*args, **kwargs): |
| 46 | self.started.set() |
| 47 | await self.release.wait() |
| 48 | |
| 49 | return call |
| 50 | |
| 51 | |
| 52 | @pytest.fixture |
| 53 | def registry(monkeypatch): |
| 54 | registry = CollectorRegistry() |
| 55 | monkeypatch.setattr(rlhf_metrics, "_metrics", WeightOperationMetrics(registry)) |
| 56 | return registry |
| 57 | |
| 58 | |
| 59 | @pytest.mark.parametrize("fail", [False, True]) |
| 60 | def test_record_releases_in_flight_and_observes_duration(registry, fail): |
| 61 | metrics = rlhf_metrics.weight_operation_metrics() |
| 62 | with ( |
| 63 | pytest.raises(RuntimeError) if fail else nullcontext(), |
| 64 | metrics.record("update"), |
| 65 | ): |
| 66 | assert sample(registry, IN_FLIGHT, "update") == 1 |
| 67 | if fail: |
| 68 | raise RuntimeError |
| 69 | assert sample(registry, IN_FLIGHT, "update") == 0 |
| 70 | assert sample(registry, DURATION, "update") == 1 |
| 71 | |
| 72 | |
| 73 | @pytest.mark.asyncio |
| 74 | @pytest.mark.parametrize( |
| 75 | "operation,route,body", |
| 76 | [ |
| 77 | ("init", api_router.init_weight_transfer_engine, {"init_info": {}}), |
| 78 | ("start", api_router.start_weight_update, None), |
| 79 | ("start_draft", api_router.start_draft_weight_update, None), |
| 80 | ("update", api_router.update_weights, {"update_info": {}}), |
| 81 | ("finish", api_router.finish_weight_update, None), |
| 82 | ], |
| 83 | ) |
| 84 | async def test_route_records_its_operation(registry, operation, route, body): |
| 85 | engine = Engine() |
| 86 | task = asyncio.create_task(route(Request(engine), **(body or {}))) |
| 87 | await engine.started.wait() |
| 88 | assert sample(registry, IN_FLIGHT, operation) == 1 |
| 89 | engine.release.set() |
| 90 | await task |
| 91 | assert sample(registry, IN_FLIGHT, operation) == 0 |
| 92 | assert sample(registry, DURATION, operation) == 1 |
| 93 | |
| 94 | |
| 95 | @pytest.mark.asyncio |
| 96 | async def test_cancelled_route_releases_in_flight(registry): |
| 97 | engine = Engine() |
| 98 | task = asyncio.create_task( |
| 99 | api_router.update_weights(Request(engine), update_info={}) |
| 100 | ) |
| 101 | await engine.started.wait() |
| 102 | task.cancel() |
| 103 | with pytest.raises(asyncio.CancelledError): |
| 104 | await task |
| 105 | assert sample(registry, IN_FLIGHT, "update") == 0 |
| 106 | assert sample(registry, DURATION, "update") == 1 |
| 107 | |
| 108 | |
| 109 | HOLD_OPERATION = """ |
| 110 | import sys |
| 111 | from vllm.entrypoints.serve.dev.rlhf.metrics import weight_operation_metrics |
| 112 | with weight_operation_metrics().record("update"): |
| 113 | print("entered", flush=True) |
| 114 | sys.stdin.readline() |
| 115 | """ |
| 116 | |
| 117 | |
| 118 | def test_in_flight_is_summed_across_api_server_processes(tmp_path, monkeypatch): |
| 119 | """Two API-server processes each hold an operation; /metrics sums them.""" |
| 120 | monkeypatch.setenv("PROMETHEUS_MULTIPROC_DIR", str(tmp_path)) |
| 121 | servers = [ |
| 122 | subprocess.Popen( |
| 123 | [sys.executable, "-c", HOLD_OPERATION], |
| 124 | stdin=subprocess.PIPE, |
| 125 | stdout=subprocess.PIPE, |
| 126 | text=True, |
| 127 | env=os.environ.copy(), |
| 128 | ) |
| 129 | for _ in range(2) |
| 130 | ] |
| 131 | try: |
| 132 | for server in servers: |
| 133 | assert server.stdout is not None |
| 134 | assert server.stdout.readline().strip() == "entered" |
| 135 | assert sample(get_prometheus_registry(), IN_FLIGHT, "update") == 2 |
| 136 | finally: |
| 137 | for server in servers: |
| 138 | server.communicate("\n", timeout=60) |
| 139 | |
| 140 | scraped = get_prometheus_registry() |
| 141 | assert sample(scraped, IN_FLIGHT, "update") == 0 |
| 142 | assert sample(scraped, DURATION, "update") == 2 |
| 143 | |
| 144 | |
| 145 | def test_rejected_request_is_not_recorded(registry): |
| 146 | app = FastAPI() |
| 147 | app.include_router(api_router.router) |
| 148 | init_exception_handler(app) |
| 149 | app.state.engine_client = Engine() |
| 150 | app.state.args = SimpleNamespace(log_error_stack=False) |
| 151 | response = TestClient(app).post("/update_weights", json={}) |
| 152 | assert response.status_code == 400 |
| 153 | assert sample(registry, DURATION, "update") is None |
| 154 | |
| 155 | |
| 156 | def test_default_recorder_is_created_lazily_on_the_default_registry(): |
| 157 | """The default recorder registers on the process registry only when first used. |
| 158 | |
| 159 | Runs in a fresh interpreter: the registry is process-global, so a route test |
| 160 | elsewhere in the session may already have registered these collectors, and |
| 161 | "not registered yet" would then be unobservable. |
| 162 | """ |
| 163 | probe = """ |
| 164 | from prometheus_client import REGISTRY |
| 165 | |
| 166 | from vllm.entrypoints.serve.dev.rlhf import metrics as rlhf_metrics |
| 167 | |
| 168 | NAMES = [ |
| 169 | "vllm:rl_weight_update_operation_duration_seconds", |
| 170 | "vllm:rl_weight_update_operations_in_flight", |
| 171 | ] |
| 172 | assert rlhf_metrics._metrics is None |
| 173 | assert not set(NAMES) & set(REGISTRY._names_to_collectors) |
| 174 | |
| 175 | metrics = rlhf_metrics.weight_operation_metrics() |
| 176 | assert REGISTRY._names_to_collectors[NAMES[0]] is metrics.duration |
| 177 | assert REGISTRY._names_to_collectors[NAMES[1]] is metrics.in_flight |
| 178 | assert rlhf_metrics.weight_operation_metrics() is metrics |
| 179 | print("lazy registration ok") |
| 180 | """ |
| 181 | result = subprocess.run( |
| 182 | [sys.executable, "-c", probe], |
| 183 | cwd=Path(__file__).parents[3], |
| 184 | capture_output=True, |
| 185 | text=True, |
| 186 | ) |
| 187 | assert result.returncode == 0, result.stderr |
| 188 | assert "lazy registration ok" in result.stdout |