A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 19:45 UTC
[Frontend] track server_load (#13950)
daniel-salib committed
Mar 14, 2025 at 09:53 UTC
73deea2fdba03690b9067e71819d5d385d2afa54
4 files changed
+131
-4
tests/entrypoints/openai/test_basic.py
+48
@@ -171,3 +171,51 @@ async def test_request_wrong_content_type(server: RemoteOpenAIServer):
171
extra_headers={
172
"Content-Type": "application/x-www-form-urlencoded"
173
})
174
+
175
+
176
+@pytest.mark.parametrize(
177
+ "server_args",
178
+ [
179
+ pytest.param(["--enable-server-load-tracking"],
180
+ id="enable-server-load-tracking")
181
+ ],
182
+ indirect=True,
183
+)
184
+@pytest.mark.asyncio
185
+async def test_server_load(server: RemoteOpenAIServer):
186
+ # Check initial server load
187
+ response = requests.get(server.url_for("load"))
188
+ assert response.status_code == HTTPStatus.OK
189
+ assert response.json().get("server_load") == 0
190
+
191
+ def make_long_completion_request():
192
+ return requests.post(
193
+ server.url_for("v1/completions"),
194
+ headers={"Content-Type": "application/json"},
195
+ json={
196
+ "prompt": "Give me a long story",
197
+ "max_tokens": 1000,
198
+ "temperature": 0,
199
+ },
200
+ )
201
+
202
+ # Start the completion request in a background thread.
203
+ completion_future = asyncio.create_task(
204
+ asyncio.to_thread(make_long_completion_request))
205
+
206
+ # Give a short delay to ensure the request has started.
207
+ await asyncio.sleep(0.1)
208
+
209
+ # Check server load while the completion request is running.
210
+ response = requests.get(server.url_for("load"))
211
+ assert response.status_code == HTTPStatus.OK
212
+ assert response.json().get("server_load") == 1
213
+
214
+ # Wait for the completion request to finish.
215
+ await completion_future
216
+ await asyncio.sleep(0.1)
217
+
218
+ # Check server load after the completion request has finished.
219
+ response = requests.get(server.url_for("load"))
220
+ assert response.status_code == HTTPStatus.OK
221
+ assert response.json().get("server_load") == 0
vllm/entrypoints/openai/api_server.py
+30
-2
@@ -80,7 +80,7 @@ from vllm.entrypoints.openai.serving_tokenization import (
80
from vllm.entrypoints.openai.serving_transcription import (
81
OpenAIServingTranscription)
82
from vllm.entrypoints.openai.tool_parsers import ToolParserManager
83
-from vllm.entrypoints.utils import with_cancellation
83
+from vllm.entrypoints.utils import load_aware_call, with_cancellation
84
from vllm.logger import init_logger
85
from vllm.usage.usage_lib import UsageContext
86
from vllm.utils import (FlexibleArgumentParser, get_open_zmq_ipc_path,
@@ -347,6 +347,24 @@ async def health(raw_request: Request) -> Response:
347
return Response(status_code=200)
348
349
350
+@router.get("/load")
351
+async def get_server_load_metrics(request: Request):
352
+ # This endpoint returns the current server load metrics.
353
+ # It tracks requests utilizing the GPU from the following routes:
354
+ # - /v1/chat/completions
355
+ # - /v1/completions
356
+ # - /v1/audio/transcriptions
357
+ # - /v1/embeddings
358
+ # - /pooling
359
+ # - /score
360
+ # - /v1/score
361
+ # - /rerank
362
+ # - /v1/rerank
363
+ # - /v2/rerank
364
+ return JSONResponse(
365
+ content={'server_load': request.app.state.server_load_metrics})
366
+
367
+
368
@router.api_route("/ping", methods=["GET", "POST"])
369
async def ping(raw_request: Request) -> Response:
370
"""Ping check. Endpoint required for SageMaker"""
@@ -400,6 +418,7 @@ async def show_version():
418
@router.post("/v1/chat/completions",
419
dependencies=[Depends(validate_json_request)])
420
@with_cancellation
421
+@load_aware_call
422
async def create_chat_completion(request: ChatCompletionRequest,
423
raw_request: Request):
424
handler = chat(raw_request)
@@ -421,6 +440,7 @@ async def create_chat_completion(request: ChatCompletionRequest,
440
441
@router.post("/v1/completions", dependencies=[Depends(validate_json_request)])
442
@with_cancellation
443
+@load_aware_call
444
async def create_completion(request: CompletionRequest, raw_request: Request):
445
handler = completion(raw_request)
446
if handler is None:
@@ -439,6 +459,7 @@ async def create_completion(request: CompletionRequest, raw_request: Request):
459
460
@router.post("/v1/embeddings", dependencies=[Depends(validate_json_request)])
461
@with_cancellation
462
+@load_aware_call
463
async def create_embedding(request: EmbeddingRequest, raw_request: Request):
464
handler = embedding(raw_request)
465
if handler is None:
@@ -485,6 +506,7 @@ async def create_embedding(request: EmbeddingRequest, raw_request: Request):
506
507
@router.post("/pooling", dependencies=[Depends(validate_json_request)])
508
@with_cancellation
509
+@load_aware_call
510
async def create_pooling(request: PoolingRequest, raw_request: Request):
511
handler = pooling(raw_request)
512
if handler is None:
@@ -503,6 +525,7 @@ async def create_pooling(request: PoolingRequest, raw_request: Request):
525
526
@router.post("/score", dependencies=[Depends(validate_json_request)])
527
@with_cancellation
528
+@load_aware_call
529
async def create_score(request: ScoreRequest, raw_request: Request):
530
handler = score(raw_request)
531
if handler is None:
@@ -521,6 +544,7 @@ async def create_score(request: ScoreRequest, raw_request: Request):
544
545
@router.post("/v1/score", dependencies=[Depends(validate_json_request)])
546
@with_cancellation
547
+@load_aware_call
548
async def create_score_v1(request: ScoreRequest, raw_request: Request):
549
logger.warning(
550
"To indicate that Score API is not part of standard OpenAI API, we "
@@ -531,10 +555,10 @@ async def create_score_v1(request: ScoreRequest, raw_request: Request):
555
556
@router.post("/v1/audio/transcriptions")
557
@with_cancellation
558
+@load_aware_call
559
async def create_transcriptions(request: Annotated[TranscriptionRequest,
560
Form()],
561
raw_request: Request):
537
-
562
handler = transcription(raw_request)
563
if handler is None:
564
return base(raw_request).create_error_response(
@@ -556,6 +580,7 @@ async def create_transcriptions(request: Annotated[TranscriptionRequest,
580
581
@router.post("/rerank", dependencies=[Depends(validate_json_request)])
582
@with_cancellation
583
+@load_aware_call
584
async def do_rerank(request: RerankRequest, raw_request: Request):
585
handler = rerank(raw_request)
586
if handler is None:
@@ -894,6 +919,9 @@ async def init_app_state(
919
) if model_config.runner_type == "transcription" else None
920
state.task = model_config.task
921
922
+ state.enable_server_load_tracking = args.enable_server_load_tracking
923
+ state.server_load_metrics = 0
924
+
925
926
def create_server_socket(addr: tuple[str, int]) -> socket.socket:
927
family = socket.AF_INET
vllm/entrypoints/openai/cli_args.py
+7
@@ -257,6 +257,13 @@ def make_arg_parser(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
257
action='store_true',
258
default=False,
259
help="If set to True, enable prompt_tokens_details in usage.")
260
+ parser.add_argument(
261
+ "--enable-server-load-tracking",
262
+ action='store_true',
263
+ default=False,
264
+ help=
265
+ "If set to True, enable tracking server_load_metrics in the app state."
266
+ )
267
268
return parser
269
vllm/entrypoints/utils.py
+46
-2
@@ -4,6 +4,8 @@ import asyncio
4
import functools
5
6
from fastapi import Request
7
+from fastapi.responses import JSONResponse, StreamingResponse
8
+from starlette.background import BackgroundTask, BackgroundTasks
9
10
11
async def listen_for_disconnect(request: Request) -> None:
@@ -17,9 +19,9 @@ async def listen_for_disconnect(request: Request) -> None:
19
def with_cancellation(handler_func):
20
"""Decorator that allows a route handler to be cancelled by client
21
disconnections.
20
-
22
+
23
This does _not_ use request.is_disconnected, which does not work with
22
- middleware. Instead this follows the pattern from
24
+ middleware. Instead this follows the pattern from
25
starlette.StreamingResponse, which simultaneously awaits on two tasks- one
26
to wait for an http disconnect message, and the other to do the work that we
27
want done. When the first task finishes, the other is cancelled.
@@ -57,3 +59,45 @@ def with_cancellation(handler_func):
59
return None
60
61
return wrapper
62
+
63
+
64
+def decrement_server_load(request: Request):
65
+ request.app.state.server_load_metrics -= 1
66
+
67
+
68
+def load_aware_call(func):
69
+
70
+ @functools.wraps(func)
71
+ async def wrapper(*args, raw_request: Request, **kwargs):
72
+ if not raw_request.app.state.enable_server_load_tracking:
73
+ return await func(*args, raw_request=raw_request, **kwargs)
74
+
75
+ raw_request.app.state.server_load_metrics += 1
76
+ try:
77
+ response = await func(*args, raw_request=raw_request, **kwargs)
78
+ except Exception:
79
+ raw_request.app.state.server_load_metrics -= 1
80
+ raise
81
+
82
+ if isinstance(response, (JSONResponse, StreamingResponse)):
83
+ if response.background is None:
84
+ response.background = BackgroundTask(decrement_server_load,
85
+ raw_request)
86
+ elif isinstance(response.background, BackgroundTasks):
87
+ response.background.add_task(decrement_server_load,
88
+ raw_request)
89
+ elif isinstance(response.background, BackgroundTask):
90
+ # Convert the single BackgroundTask to BackgroundTasks
91
+ # and chain the decrement_server_load task to it
92
+ tasks = BackgroundTasks()
93
+ tasks.add_task(response.background.func,
94
+ *response.background.args,
95
+ **response.background.kwargs)
96
+ tasks.add_task(decrement_server_load, raw_request)
97
+ response.background = tasks
98
+ else:
99
+ raw_request.app.state.server_load_metrics -= 1
100
+
101
+ return response
102
+
103
+ return wrapper