api_server.py 36.6 KB
Newer Older
1
2
# SPDX-License-Identifier: Apache-2.0

3
import asyncio
4
import atexit
5
import gc
6
7
import importlib
import inspect
8
import multiprocessing
9
import os
10
import re
11
import signal
12
import socket
13
import tempfile
14
import uuid
15
from argparse import Namespace
16
from contextlib import asynccontextmanager
17
from functools import partial
18
from http import HTTPStatus
19
from typing import Annotated, AsyncIterator, Dict, Optional, Set, Tuple, Union
20

21
import uvloop
22
from fastapi import APIRouter, Depends, FastAPI, Form, HTTPException, Request
Zhuohan Li's avatar
Zhuohan Li committed
23
24
from fastapi.exceptions import RequestValidationError
from fastapi.middleware.cors import CORSMiddleware
25
from fastapi.responses import JSONResponse, Response, StreamingResponse
26
from starlette.datastructures import State
27
from starlette.routing import Mount
28
from typing_extensions import assert_never
Zhuohan Li's avatar
Zhuohan Li committed
29

30
import vllm.envs as envs
31
from vllm.config import ModelConfig
Woosuk Kwon's avatar
Woosuk Kwon committed
32
from vllm.engine.arg_utils import AsyncEngineArgs
33
from vllm.engine.async_llm_engine import AsyncLLMEngine  # type: ignore
34
35
36
from vllm.engine.multiprocessing.client import MQLLMEngineClient
from vllm.engine.multiprocessing.engine import run_mp_engine
from vllm.engine.protocol import EngineClient
37
from vllm.entrypoints.chat_utils import load_chat_template
38
from vllm.entrypoints.launcher import serve_http
39
from vllm.entrypoints.logger import RequestLogger
40
41
from vllm.entrypoints.openai.cli_args import (make_arg_parser,
                                              validate_parsed_serve_args)
42
43
# yapf conflicts with isort for this block
# yapf: disable
44
from vllm.entrypoints.openai.protocol import (ChatCompletionRequest,
45
                                              ChatCompletionResponse,
46
                                              CompletionRequest,
47
                                              CompletionResponse,
48
49
                                              DetokenizeRequest,
                                              DetokenizeResponse,
50
51
                                              EmbeddingChatRequest,
                                              EmbeddingCompletionRequest,
52
                                              EmbeddingRequest,
53
54
55
                                              EmbeddingResponse,
                                              EmbeddingResponseData,
                                              ErrorResponse,
56
                                              LoadLoraAdapterRequest,
57
58
                                              PoolingChatRequest,
                                              PoolingCompletionRequest,
59
                                              PoolingRequest, PoolingResponse,
60
                                              RerankRequest, RerankResponse,
61
                                              ScoreRequest, ScoreResponse,
62
                                              TokenizeRequest,
63
                                              TokenizeResponse,
64
65
                                              TranscriptionRequest,
                                              TranscriptionResponse,
66
                                              UnloadLoraAdapterRequest)
67
from vllm.entrypoints.openai.reasoning_parsers import ReasoningParserManager
68
# yapf: enable
69
70
from vllm.entrypoints.openai.serving_chat import OpenAIServingChat
from vllm.entrypoints.openai.serving_completion import OpenAIServingCompletion
71
from vllm.entrypoints.openai.serving_embedding import OpenAIServingEmbedding
72
73
74
from vllm.entrypoints.openai.serving_engine import OpenAIServing
from vllm.entrypoints.openai.serving_models import (BaseModelPath,
                                                    OpenAIServingModels)
75
from vllm.entrypoints.openai.serving_pooling import OpenAIServingPooling
76
from vllm.entrypoints.openai.serving_rerank import JinaAIServingRerank
77
from vllm.entrypoints.openai.serving_score import OpenAIServingScores
78
79
from vllm.entrypoints.openai.serving_tokenization import (
    OpenAIServingTokenization)
80
81
from vllm.entrypoints.openai.serving_transcription import (
    OpenAIServingTranscription)
82
from vllm.entrypoints.openai.tool_parsers import ToolParserManager
83
from vllm.entrypoints.utils import with_cancellation
84
from vllm.logger import init_logger
yhu422's avatar
yhu422 committed
85
from vllm.usage.usage_lib import UsageContext
86
from vllm.utils import (FlexibleArgumentParser, get_open_zmq_ipc_path,
87
                        is_valid_ipv6_address, set_ulimit)
88
from vllm.version import __version__ as VLLM_VERSION
Zhuohan Li's avatar
Zhuohan Li committed
89

90
TIMEOUT_KEEP_ALIVE = 5  # seconds
Zhuohan Li's avatar
Zhuohan Li committed
91

92
prometheus_multiproc_dir: tempfile.TemporaryDirectory
93

94
# Cannot use __name__ (https://github.com/vllm-project/vllm/pull/4765)
95
logger = init_logger('vllm.entrypoints.openai.api_server')
96

97
_running_tasks: Set[asyncio.Task] = set()
98

99

100
@asynccontextmanager
101
async def lifespan(app: FastAPI):
102
103
    try:
        if app.state.log_stats:
104
            engine_client: EngineClient = app.state.engine_client
105
106
107

            async def _force_log():
                while True:
108
109
                    await asyncio.sleep(10.)
                    await engine_client.do_log_stats()
110
111
112
113
114
115

            task = asyncio.create_task(_force_log())
            _running_tasks.add(task)
            task.add_done_callback(_running_tasks.remove)
        else:
            task = None
116
117
118
119
120

        # Mark the startup heap as static so that it's ignored by GC.
        # Reduces pause times of oldest generation collections.
        gc.collect()
        gc.freeze()
121
122
123
124
125
126
127
128
        try:
            yield
        finally:
            if task is not None:
                task.cancel()
    finally:
        # Ensure app state including engine ref is gc'd
        del app.state
129
130


131
@asynccontextmanager
132
async def build_async_engine_client(
133
        args: Namespace) -> AsyncIterator[EngineClient]:
134

135
    # Context manager to handle engine_client lifecycle
136
137
138
    # Ensures everything is shutdown and cleaned up on error/exit
    engine_args = AsyncEngineArgs.from_cli_args(args)

139
140
141
142
143
144
145
146
147
    async with build_async_engine_client_from_engine_args(
            engine_args, args.disable_frontend_multiprocessing) as engine:
        yield engine


@asynccontextmanager
async def build_async_engine_client_from_engine_args(
    engine_args: AsyncEngineArgs,
    disable_frontend_multiprocessing: bool = False,
148
) -> AsyncIterator[EngineClient]:
149
    """
150
    Create EngineClient, either:
151
152
153
154
155
156
        - in-process using the AsyncLLMEngine Directly
        - multiprocess using AsyncLLMEngine RPC

    Returns the Client or None if the creation failed.
    """

157
    # AsyncLLMEngine.
158
    if (MQLLMEngineClient.is_unsupported_config(engine_args)
159
            or envs.VLLM_USE_V1 or disable_frontend_multiprocessing):
160

161
162
163
164
165
166
167
168
169
        engine_client: Optional[EngineClient] = None
        try:
            engine_client = AsyncLLMEngine.from_engine_args(
                engine_args=engine_args,
                usage_context=UsageContext.OPENAI_API_SERVER)
            yield engine_client
        finally:
            if engine_client and hasattr(engine_client, "shutdown"):
                engine_client.shutdown()
170

171
    # MQLLMEngine.
172
    else:
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
        if "PROMETHEUS_MULTIPROC_DIR" not in os.environ:
            # Make TemporaryDirectory for prometheus multiprocessing
            # Note: global TemporaryDirectory will be automatically
            #   cleaned up upon exit.
            global prometheus_multiproc_dir
            prometheus_multiproc_dir = tempfile.TemporaryDirectory()
            os.environ[
                "PROMETHEUS_MULTIPROC_DIR"] = prometheus_multiproc_dir.name
        else:
            logger.warning(
                "Found PROMETHEUS_MULTIPROC_DIR was set by user. "
                "This directory must be wiped between vLLM runs or "
                "you will find inaccurate metrics. Unset the variable "
                "and vLLM will properly handle cleanup.")

188
        # Select random path for IPC.
189
        ipc_path = get_open_zmq_ipc_path()
190
191
        logger.debug("Multiprocessing frontend to use %s for IPC Path.",
                     ipc_path)
192

193
        # Start RPCServer in separate process (holds the LLMEngine).
194
195
        # the current process might have CUDA context,
        # so we need to spawn a new process
196
197
        context = multiprocessing.get_context("spawn")

198
199
200
201
        # The Process can raise an exception during startup, which may
        # not actually result in an exitcode being reported. As a result
        # we use a shared variable to communicate the information.
        engine_alive = multiprocessing.Value('b', True, lock=False)
202
203
204
        engine_process = context.Process(target=run_mp_engine,
                                         args=(engine_args,
                                               UsageContext.OPENAI_API_SERVER,
205
                                               ipc_path, engine_alive))
206
        engine_process.start()
207
        engine_pid = engine_process.pid
208
        assert engine_pid is not None, "Engine process failed to start."
209
        logger.info("Started engine process with PID %d", engine_pid)
210

211
212
213
214
215
216
217
218
        def _cleanup_ipc_path():
            socket_path = ipc_path.replace("ipc://", "")
            if os.path.exists(socket_path):
                os.remove(socket_path)

        # Ensure we clean up the local IPC socket file on exit.
        atexit.register(_cleanup_ipc_path)

219
220
        # Build RPCClient, which conforms to EngineClient Protocol.
        engine_config = engine_args.create_engine_config()
221
222
223
224
        build_client = partial(MQLLMEngineClient, ipc_path, engine_config,
                               engine_pid)
        mq_engine_client = await asyncio.get_running_loop().run_in_executor(
            None, build_client)
225
        try:
226
227
            while True:
                try:
228
                    await mq_engine_client.setup()
229
                    break
230
                except TimeoutError:
231
232
                    if (not engine_process.is_alive()
                            or not engine_alive.value):
233
                        raise RuntimeError(
234
235
                            "Engine process failed to start. See stack "
                            "trace for the root cause.") from None
236

237
            yield mq_engine_client  # type: ignore[misc]
238
239
        finally:
            # Ensure rpc server process was terminated
240
            engine_process.terminate()
241
242

            # Close all open connections to the backend
243
            mq_engine_client.close()
244

245
246
247
248
249
            # Wait for engine process to join
            engine_process.join(4)
            if engine_process.exitcode is None:
                # Kill if taking longer than 5 seconds to stop
                engine_process.kill()
250

251
252
253
254
255
            # Lazy import for prometheus multiprocessing.
            # We need to set PROMETHEUS_MULTIPROC_DIR environment variable
            # before prometheus_client is imported.
            # See https://prometheus.github.io/client_python/multiprocess/
            from prometheus_client import multiprocess
256
            multiprocess.mark_process_dead(engine_process.pid)
257

258

259
260
async def validate_json_request(raw_request: Request):
    content_type = raw_request.headers.get("content-type", "").lower()
261
262
    media_type = content_type.split(";", maxsplit=1)[0]
    if media_type != "application/json":
263
264
265
266
267
268
        raise HTTPException(
            status_code=HTTPStatus.UNSUPPORTED_MEDIA_TYPE,
            detail="Unsupported Media Type: Only 'application/json' is allowed"
        )


Ethan Xu's avatar
Ethan Xu committed
269
router = APIRouter()
Zhuohan Li's avatar
Zhuohan Li committed
270

271

272
def mount_metrics(app: FastAPI):
273
274
275
276
277
278
279
280
281
    # Lazy import for prometheus multiprocessing.
    # We need to set PROMETHEUS_MULTIPROC_DIR environment variable
    # before prometheus_client is imported.
    # See https://prometheus.github.io/client_python/multiprocess/
    from prometheus_client import (CollectorRegistry, make_asgi_app,
                                   multiprocess)

    prometheus_multiproc_dir_path = os.getenv("PROMETHEUS_MULTIPROC_DIR", None)
    if prometheus_multiproc_dir_path is not None:
282
283
        logger.debug("vLLM to use %s as PROMETHEUS_MULTIPROC_DIR",
                     prometheus_multiproc_dir_path)
284
285
286
287
288
289
290
291
292
        registry = CollectorRegistry()
        multiprocess.MultiProcessCollector(registry)

        # Add prometheus asgi middleware to route /metrics requests
        metrics_route = Mount("/metrics", make_asgi_app(registry=registry))
    else:
        # Add prometheus asgi middleware to route /metrics requests
        metrics_route = Mount("/metrics", make_asgi_app())

293
    # Workaround for 307 Redirect for /metrics
294
    metrics_route.path_regex = re.compile("^/metrics(?P<path>.*)$")
295
    app.routes.append(metrics_route)
296
297


298
299
300
301
302
def base(request: Request) -> OpenAIServing:
    # Reuse the existing instance
    return tokenization(request)


303
304
305
306
def models(request: Request) -> OpenAIServingModels:
    return request.app.state.openai_serving_models


307
def chat(request: Request) -> Optional[OpenAIServingChat]:
308
309
310
    return request.app.state.openai_serving_chat


311
def completion(request: Request) -> Optional[OpenAIServingCompletion]:
312
313
314
    return request.app.state.openai_serving_completion


315
316
317
318
def pooling(request: Request) -> Optional[OpenAIServingPooling]:
    return request.app.state.openai_serving_pooling


319
320
def embedding(request: Request) -> Optional[OpenAIServingEmbedding]:
    return request.app.state.openai_serving_embedding
321
322


323
324
325
326
def score(request: Request) -> Optional[OpenAIServingScores]:
    return request.app.state.openai_serving_scores


327
328
329
330
def rerank(request: Request) -> Optional[JinaAIServingRerank]:
    return request.app.state.jinaai_serving_reranking


331
332
def tokenization(request: Request) -> OpenAIServingTokenization:
    return request.app.state.openai_serving_tokenization
333
334


335
336
337
338
def transcription(request: Request) -> OpenAIServingTranscription:
    return request.app.state.openai_serving_transcription


339
def engine_client(request: Request) -> EngineClient:
340
341
342
    return request.app.state.engine_client


Ethan Xu's avatar
Ethan Xu committed
343
@router.get("/health")
344
async def health(raw_request: Request) -> Response:
345
    """Health check."""
346
    await engine_client(raw_request).check_health()
347
348
349
    return Response(status_code=200)


350
351
352
353
354
355
@router.api_route("/ping", methods=["GET", "POST"])
async def ping(raw_request: Request) -> Response:
    """Ping check. Endpoint required for SageMaker"""
    return await health(raw_request)


356
@router.post("/tokenize", dependencies=[Depends(validate_json_request)])
357
@with_cancellation
358
async def tokenize(request: TokenizeRequest, raw_request: Request):
359
360
    handler = tokenization(raw_request)

361
    generator = await handler.create_tokenize(request, raw_request)
362
363
364
    if isinstance(generator, ErrorResponse):
        return JSONResponse(content=generator.model_dump(),
                            status_code=generator.code)
365
    elif isinstance(generator, TokenizeResponse):
366
367
        return JSONResponse(content=generator.model_dump())

368
369
    assert_never(generator)

370

371
@router.post("/detokenize", dependencies=[Depends(validate_json_request)])
372
@with_cancellation
373
async def detokenize(request: DetokenizeRequest, raw_request: Request):
374
375
    handler = tokenization(raw_request)

376
    generator = await handler.create_detokenize(request, raw_request)
377
378
379
    if isinstance(generator, ErrorResponse):
        return JSONResponse(content=generator.model_dump(),
                            status_code=generator.code)
380
    elif isinstance(generator, DetokenizeResponse):
381
382
        return JSONResponse(content=generator.model_dump())

383
384
    assert_never(generator)

385

Ethan Xu's avatar
Ethan Xu committed
386
@router.get("/v1/models")
387
async def show_available_models(raw_request: Request):
388
    handler = models(raw_request)
389

390
391
    models_ = await handler.show_available_models()
    return JSONResponse(content=models_.model_dump())
Zhuohan Li's avatar
Zhuohan Li committed
392
393


Ethan Xu's avatar
Ethan Xu committed
394
@router.get("/version")
395
async def show_version():
396
    ver = {"version": VLLM_VERSION}
397
398
399
    return JSONResponse(content=ver)


400
401
@router.post("/v1/chat/completions",
             dependencies=[Depends(validate_json_request)])
402
@with_cancellation
403
404
async def create_chat_completion(request: ChatCompletionRequest,
                                 raw_request: Request):
405
406
407
408
    handler = chat(raw_request)
    if handler is None:
        return base(raw_request).create_error_response(
            message="The model does not support Chat Completions API")
409

410
    generator = await handler.create_chat_completion(request, raw_request)
411

412
413
414
    if isinstance(generator, ErrorResponse):
        return JSONResponse(content=generator.model_dump(),
                            status_code=generator.code)
415

416
    elif isinstance(generator, ChatCompletionResponse):
417
        return JSONResponse(content=generator.model_dump())
418

419
420
    return StreamingResponse(content=generator, media_type="text/event-stream")

421

422
@router.post("/v1/completions", dependencies=[Depends(validate_json_request)])
423
@with_cancellation
424
async def create_completion(request: CompletionRequest, raw_request: Request):
425
426
427
428
429
430
    handler = completion(raw_request)
    if handler is None:
        return base(raw_request).create_error_response(
            message="The model does not support Completions API")

    generator = await handler.create_completion(request, raw_request)
431
432
433
    if isinstance(generator, ErrorResponse):
        return JSONResponse(content=generator.model_dump(),
                            status_code=generator.code)
434
    elif isinstance(generator, CompletionResponse):
435
        return JSONResponse(content=generator.model_dump())
Zhuohan Li's avatar
Zhuohan Li committed
436

437
438
    return StreamingResponse(content=generator, media_type="text/event-stream")

Zhuohan Li's avatar
Zhuohan Li committed
439

440
@router.post("/v1/embeddings", dependencies=[Depends(validate_json_request)])
441
@with_cancellation
442
async def create_embedding(request: EmbeddingRequest, raw_request: Request):
443
444
    handler = embedding(raw_request)
    if handler is None:
445
446
447
448
449
450
451
452
453
454
455
        fallback_handler = pooling(raw_request)
        if fallback_handler is None:
            return base(raw_request).create_error_response(
                message="The model does not support Embeddings API")

        logger.warning(
            "Embeddings API will become exclusive to embedding models "
            "in a future release. To return the hidden states directly, "
            "use the Pooling API (`/pooling`) instead.")

        res = await fallback_handler.create_pooling(request, raw_request)
456
457

        generator: Union[ErrorResponse, EmbeddingResponse]
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
        if isinstance(res, PoolingResponse):
            generator = EmbeddingResponse(
                id=res.id,
                object=res.object,
                created=res.created,
                model=res.model,
                data=[
                    EmbeddingResponseData(
                        index=d.index,
                        embedding=d.data,  # type: ignore
                    ) for d in res.data
                ],
                usage=res.usage,
            )
        else:
            generator = res
    else:
        generator = await handler.create_embedding(request, raw_request)
476

477
478
479
    if isinstance(generator, ErrorResponse):
        return JSONResponse(content=generator.model_dump(),
                            status_code=generator.code)
480
    elif isinstance(generator, EmbeddingResponse):
481
482
        return JSONResponse(content=generator.model_dump())

483
484
    assert_never(generator)

485

486
@router.post("/pooling", dependencies=[Depends(validate_json_request)])
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
@with_cancellation
async def create_pooling(request: PoolingRequest, raw_request: Request):
    handler = pooling(raw_request)
    if handler is None:
        return base(raw_request).create_error_response(
            message="The model does not support Pooling API")

    generator = await handler.create_pooling(request, raw_request)
    if isinstance(generator, ErrorResponse):
        return JSONResponse(content=generator.model_dump(),
                            status_code=generator.code)
    elif isinstance(generator, PoolingResponse):
        return JSONResponse(content=generator.model_dump())

    assert_never(generator)


504
@router.post("/score", dependencies=[Depends(validate_json_request)])
505
@with_cancellation
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
async def create_score(request: ScoreRequest, raw_request: Request):
    handler = score(raw_request)
    if handler is None:
        return base(raw_request).create_error_response(
            message="The model does not support Score API")

    generator = await handler.create_score(request, raw_request)
    if isinstance(generator, ErrorResponse):
        return JSONResponse(content=generator.model_dump(),
                            status_code=generator.code)
    elif isinstance(generator, ScoreResponse):
        return JSONResponse(content=generator.model_dump())

    assert_never(generator)


522
@router.post("/v1/score", dependencies=[Depends(validate_json_request)])
523
@with_cancellation
524
525
526
527
528
529
530
531
async def create_score_v1(request: ScoreRequest, raw_request: Request):
    logger.warning(
        "To indicate that Score API is not part of standard OpenAI API, we "
        "have moved it to `/score`. Please update your client accordingly.")

    return await create_score(request, raw_request)


532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
@router.post("/v1/audio/transcriptions")
@with_cancellation
async def create_transcriptions(request: Annotated[TranscriptionRequest,
                                                   Form()],
                                raw_request: Request):

    handler = transcription(raw_request)
    if handler is None:
        return base(raw_request).create_error_response(
            message="The model does not support Transcriptions API")

    audio_data = await request.file.read()
    generator = await handler.create_transcription(audio_data, request,
                                                   raw_request)

    if isinstance(generator, ErrorResponse):
        return JSONResponse(content=generator.model_dump(),
                            status_code=generator.code)

    elif isinstance(generator, TranscriptionResponse):
        return JSONResponse(content=generator.model_dump())

    return StreamingResponse(content=generator, media_type="text/event-stream")


557
@router.post("/rerank", dependencies=[Depends(validate_json_request)])
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
@with_cancellation
async def do_rerank(request: RerankRequest, raw_request: Request):
    handler = rerank(raw_request)
    if handler is None:
        return base(raw_request).create_error_response(
            message="The model does not support Rerank (Score) API")
    generator = await handler.do_rerank(request, raw_request)
    if isinstance(generator, ErrorResponse):
        return JSONResponse(content=generator.model_dump(),
                            status_code=generator.code)
    elif isinstance(generator, RerankResponse):
        return JSONResponse(content=generator.model_dump())

    assert_never(generator)


574
@router.post("/v1/rerank", dependencies=[Depends(validate_json_request)])
575
576
@with_cancellation
async def do_rerank_v1(request: RerankRequest, raw_request: Request):
577
    logger.warning_once(
578
579
580
581
582
583
584
        "To indicate that the rerank API is not part of the standard OpenAI"
        " API, we have located it at `/rerank`. Please update your client"
        "accordingly. (Note: Conforms to JinaAI rerank API)")

    return await do_rerank(request, raw_request)


585
@router.post("/v2/rerank", dependencies=[Depends(validate_json_request)])
586
587
588
589
590
@with_cancellation
async def do_rerank_v2(request: RerankRequest, raw_request: Request):
    return await do_rerank(request, raw_request)


591
TASK_HANDLERS: Dict[str, Dict[str, tuple]] = {
592
593
594
595
596
597
598
599
600
    "generate": {
        "messages": (ChatCompletionRequest, create_chat_completion),
        "default": (CompletionRequest, create_completion),
    },
    "embed": {
        "messages": (EmbeddingChatRequest, create_embedding),
        "default": (EmbeddingCompletionRequest, create_embedding),
    },
    "score": {
601
602
603
604
        "default": (RerankRequest, do_rerank)
    },
    "rerank": {
        "default": (RerankRequest, do_rerank)
605
606
607
608
609
610
611
612
613
614
615
    },
    "reward": {
        "messages": (PoolingChatRequest, create_pooling),
        "default": (PoolingCompletionRequest, create_pooling),
    },
    "classify": {
        "messages": (PoolingChatRequest, create_pooling),
        "default": (PoolingCompletionRequest, create_pooling),
    },
}

616
617
618
619
620
621
622
623
624
625
626
627
if envs.VLLM_SERVER_DEV_MODE:

    @router.post("/reset_prefix_cache")
    async def reset_prefix_cache(raw_request: Request):
        """
        Reset the prefix cache. Note that we currently do not check if the
        prefix cache is successfully reset in the API server.
        """
        logger.info("Resetting prefix cache...")
        await engine_client(raw_request).reset_prefix_cache()
        return Response(status_code=200)

628

629
@router.post("/invocations", dependencies=[Depends(validate_json_request)])
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
async def invocations(raw_request: Request):
    """
    For SageMaker, routes requests to other handlers based on model `task`.
    """
    body = await raw_request.json()
    task = raw_request.app.state.task

    if task not in TASK_HANDLERS:
        raise HTTPException(
            status_code=400,
            detail=f"Unsupported task: '{task}' for '/invocations'. "
            f"Expected one of {set(TASK_HANDLERS.keys())}")

    handler_config = TASK_HANDLERS[task]
    if "messages" in body:
        request_model, handler = handler_config["messages"]
    else:
        request_model, handler = handler_config["default"]

    # this is required since we lose the FastAPI automatic casting
    request = request_model.model_validate(body)
    return await handler(request, raw_request)


654
655
656
657
658
659
if envs.VLLM_TORCH_PROFILER_DIR:
    logger.warning(
        "Torch Profiler is enabled in the API server. This should ONLY be "
        "used for local development!")

    @router.post("/start_profile")
660
    async def start_profile(raw_request: Request):
661
        logger.info("Starting profiler...")
662
        await engine_client(raw_request).start_profile()
663
664
665
666
        logger.info("Profiler started.")
        return Response(status_code=200)

    @router.post("/stop_profile")
667
    async def stop_profile(raw_request: Request):
668
        logger.info("Stopping profiler...")
669
        await engine_client(raw_request).stop_profile()
670
671
672
673
        logger.info("Profiler stopped.")
        return Response(status_code=200)


674
675
676
677
678
if envs.VLLM_ALLOW_RUNTIME_LORA_UPDATING:
    logger.warning(
        "Lora dynamic loading & unloading is enabled in the API server. "
        "This should ONLY be used for local development!")

679
680
    @router.post("/v1/load_lora_adapter",
                 dependencies=[Depends(validate_json_request)])
681
682
    async def load_lora_adapter(request: LoadLoraAdapterRequest,
                                raw_request: Request):
683
684
685
686
687
        handler = models(raw_request)
        response = await handler.load_lora_adapter(request)
        if isinstance(response, ErrorResponse):
            return JSONResponse(content=response.model_dump(),
                                status_code=response.code)
688
689
690

        return Response(status_code=200, content=response)

691
692
    @router.post("/v1/unload_lora_adapter",
                 dependencies=[Depends(validate_json_request)])
693
694
    async def unload_lora_adapter(request: UnloadLoraAdapterRequest,
                                  raw_request: Request):
695
696
697
698
699
        handler = models(raw_request)
        response = await handler.unload_lora_adapter(request)
        if isinstance(response, ErrorResponse):
            return JSONResponse(content=response.model_dump(),
                                status_code=response.code)
700
701
702
703

        return Response(status_code=200, content=response)


704
def build_app(args: Namespace) -> FastAPI:
705
706
707
708
709
710
711
    if args.disable_fastapi_docs:
        app = FastAPI(openapi_url=None,
                      docs_url=None,
                      redoc_url=None,
                      lifespan=lifespan)
    else:
        app = FastAPI(lifespan=lifespan)
Ethan Xu's avatar
Ethan Xu committed
712
713
    app.include_router(router)
    app.root_path = args.root_path
Zhuohan Li's avatar
Zhuohan Li committed
714

715
716
    mount_metrics(app)

Zhuohan Li's avatar
Zhuohan Li committed
717
718
719
720
721
722
723
724
    app.add_middleware(
        CORSMiddleware,
        allow_origins=args.allowed_origins,
        allow_credentials=args.allow_credentials,
        allow_methods=args.allowed_methods,
        allow_headers=args.allowed_headers,
    )

Ethan Xu's avatar
Ethan Xu committed
725
726
    @app.exception_handler(RequestValidationError)
    async def validation_exception_handler(_, exc):
727
728
729
        err = ErrorResponse(message=str(exc),
                            type="BadRequestError",
                            code=HTTPStatus.BAD_REQUEST)
Ethan Xu's avatar
Ethan Xu committed
730
731
732
        return JSONResponse(err.model_dump(),
                            status_code=HTTPStatus.BAD_REQUEST)

733
    if token := envs.VLLM_API_KEY or args.api_key:
734
735
736

        @app.middleware("http")
        async def authentication(request: Request, call_next):
737
738
            if request.method == "OPTIONS":
                return await call_next(request)
739
740
741
742
            url_path = request.url.path
            if app.root_path and url_path.startswith(app.root_path):
                url_path = url_path[len(app.root_path):]
            if not url_path.startswith("/v1"):
743
744
745
746
747
748
                return await call_next(request)
            if request.headers.get("Authorization") != "Bearer " + token:
                return JSONResponse(content={"error": "Unauthorized"},
                                    status_code=401)
            return await call_next(request)

749
750
751
752
753
754
755
756
757
758
759
760
    if args.enable_request_id_headers:
        logger.warning(
            "CAUTION: Enabling X-Request-Id headers in the API Server. "
            "This can harm performance at high QPS.")

        @app.middleware("http")
        async def add_request_id(request: Request, call_next):
            request_id = request.headers.get(
                "X-Request-Id") or uuid.uuid4().hex
            response = await call_next(request)
            response.headers["X-Request-Id"] = request_id
            return response
761

762
763
764
765
    for middleware in args.middleware:
        module_path, object_name = middleware.rsplit(".", 1)
        imported = getattr(importlib.import_module(module_path), object_name)
        if inspect.isclass(imported):
766
            app.add_middleware(imported)  # type: ignore[arg-type]
767
768
769
        elif inspect.iscoroutinefunction(imported):
            app.middleware("http")(imported)
        else:
770
771
            raise ValueError(f"Invalid middleware {middleware}. "
                             f"Must be a function or a class.")
772

Ethan Xu's avatar
Ethan Xu committed
773
774
775
    return app


776
async def init_app_state(
777
    engine_client: EngineClient,
778
779
    model_config: ModelConfig,
    state: State,
780
    args: Namespace,
781
) -> None:
782
    if args.served_model_name is not None:
783
        served_model_names = args.served_model_name
784
    else:
785
        served_model_names = [args.model]
786

787
788
789
790
791
    if args.disable_log_requests:
        request_logger = None
    else:
        request_logger = RequestLogger(max_log_len=args.max_log_len)

792
793
794
795
796
    base_model_paths = [
        BaseModelPath(name=name, model_path=args.model)
        for name in served_model_names
    ]

797
    state.engine_client = engine_client
798
    state.log_stats = not args.disable_log_stats
Ethan Xu's avatar
Ethan Xu committed
799

800
    resolved_chat_template = load_chat_template(args.chat_template)
801
802
803
    if resolved_chat_template is not None:
        logger.info("Using supplied chat template:\n%s",
                    resolved_chat_template)
804

805
    state.openai_serving_models = OpenAIServingModels(
806
        engine_client=engine_client,
807
808
809
810
811
        model_config=model_config,
        base_model_paths=base_model_paths,
        lora_modules=args.lora_modules,
        prompt_adapters=args.prompt_adapters,
    )
812
    await state.openai_serving_models.init_static_loras()
813
    state.openai_serving_chat = OpenAIServingChat(
814
        engine_client,
815
        model_config,
816
        state.openai_serving_models,
817
818
        args.response_role,
        request_logger=request_logger,
819
820
        chat_template=resolved_chat_template,
        chat_template_content_format=args.chat_template_content_format,
821
        return_tokens_as_token_ids=args.return_tokens_as_token_ids,
822
        enable_auto_tools=args.enable_auto_tool_choice,
823
        tool_parser=args.tool_call_parser,
824
825
        enable_reasoning=args.enable_reasoning,
        reasoning_parser=args.reasoning_parser,
826
        enable_prompt_tokens_details=args.enable_prompt_tokens_details,
827
    ) if model_config.runner_type == "generate" else None
828
    state.openai_serving_completion = OpenAIServingCompletion(
829
        engine_client,
830
        model_config,
831
        state.openai_serving_models,
832
        request_logger=request_logger,
833
        return_tokens_as_token_ids=args.return_tokens_as_token_ids,
834
    ) if model_config.runner_type == "generate" else None
835
    state.openai_serving_pooling = OpenAIServingPooling(
836
        engine_client,
837
        model_config,
838
        state.openai_serving_models,
839
        request_logger=request_logger,
840
841
        chat_template=resolved_chat_template,
        chat_template_content_format=args.chat_template_content_format,
842
    ) if model_config.runner_type == "pooling" else None
843
844
845
    state.openai_serving_embedding = OpenAIServingEmbedding(
        engine_client,
        model_config,
846
        state.openai_serving_models,
847
848
849
850
        request_logger=request_logger,
        chat_template=resolved_chat_template,
        chat_template_content_format=args.chat_template_content_format,
    ) if model_config.task == "embed" else None
851
852
853
    state.openai_serving_scores = OpenAIServingScores(
        engine_client,
        model_config,
854
        state.openai_serving_models,
855
        request_logger=request_logger
856
    ) if model_config.task == "score" else None
857
858
859
860
861
862
    state.jinaai_serving_reranking = JinaAIServingRerank(
        engine_client,
        model_config,
        state.openai_serving_models,
        request_logger=request_logger
    ) if model_config.task == "score" else None
863
    state.openai_serving_tokenization = OpenAIServingTokenization(
864
        engine_client,
865
        model_config,
866
        state.openai_serving_models,
867
        request_logger=request_logger,
868
869
        chat_template=resolved_chat_template,
        chat_template_content_format=args.chat_template_content_format,
870
    )
871
872
873
874
875
876
    state.openai_serving_transcription = OpenAIServingTranscription(
        engine_client,
        model_config,
        state.openai_serving_models,
        request_logger=request_logger,
    ) if model_config.runner_type == "transcription" else None
877
    state.task = model_config.task
878
879


880
881
882
883
884
885
886
def create_server_socket(addr: Tuple[str, int]) -> socket.socket:
    family = socket.AF_INET
    if is_valid_ipv6_address(addr[0]):
        family = socket.AF_INET6

    sock = socket.socket(family=family, type=socket.SOCK_STREAM)
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
887
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1)
888
889
890
891
892
    sock.bind(addr)

    return sock


893
async def run_server(args, **uvicorn_kwargs) -> None:
894
895
896
    logger.info("vLLM API server version %s", VLLM_VERSION)
    logger.info("args: %s", args)

897
898
899
    if args.tool_parser_plugin and len(args.tool_parser_plugin) > 3:
        ToolParserManager.import_tool_parser(args.tool_parser_plugin)

900
    valid_tool_parses = ToolParserManager.tool_parsers.keys()
901
    if args.enable_auto_tool_choice \
902
        and args.tool_call_parser not in valid_tool_parses:
903
        raise KeyError(f"invalid tool call parser: {args.tool_call_parser} "
904
                       f"(chose from {{ {','.join(valid_tool_parses)} }})")
905

906
907
908
909
910
911
912
    valid_reasoning_parses = ReasoningParserManager.reasoning_parsers.keys()
    if args.enable_reasoning \
        and args.reasoning_parser not in valid_reasoning_parses:
        raise KeyError(
            f"invalid reasoning parser: {args.reasoning_parser} "
            f"(chose from {{ {','.join(valid_reasoning_parses)} }})")

913
914
915
    # workaround to make sure that we bind the port before the engine is set up.
    # This avoids race conditions with ray.
    # see https://github.com/vllm-project/vllm/issues/8204
916
917
    sock_addr = (args.host or "", args.port)
    sock = create_server_socket(sock_addr)
918

919
920
921
922
    # workaround to avoid footguns where uvicorn drops requests with too
    # many concurrent requests active
    set_ulimit()

923
924
925
926
927
928
    def signal_handler(*_) -> None:
        # Interrupt server on sigterm while initializing
        raise KeyboardInterrupt("terminated")

    signal.signal(signal.SIGTERM, signal_handler)

929
    async with build_async_engine_client(args) as engine_client:
930
931
        app = build_app(args)

932
        model_config = await engine_client.get_model_config()
933
        await init_app_state(engine_client, model_config, app.state, args)
934

935
936
937
938
939
940
941
942
        def _listen_addr(a: str) -> str:
            if is_valid_ipv6_address(a):
                return '[' + a + ']'
            return a or "0.0.0.0"

        logger.info("Starting vLLM API server on http://%s:%d",
                    _listen_addr(sock_addr[0]), sock_addr[1])

943
944
        shutdown_task = await serve_http(
            app,
945
            sock=sock,
946
947
948
949
950
951
952
953
            host=args.host,
            port=args.port,
            log_level=args.uvicorn_log_level,
            timeout_keep_alive=TIMEOUT_KEEP_ALIVE,
            ssl_keyfile=args.ssl_keyfile,
            ssl_certfile=args.ssl_certfile,
            ssl_ca_certs=args.ssl_ca_certs,
            ssl_cert_reqs=args.ssl_cert_reqs,
954
955
956
            **uvicorn_kwargs,
        )

957
958
    # NB: Await server shutdown only after the backend context is exited
    await shutdown_task
959

960
961
    sock.close()

Ethan Xu's avatar
Ethan Xu committed
962
963
964

if __name__ == "__main__":
    # NOTE(simon):
965
966
    # This section should be in sync with vllm/entrypoints/cli/main.py for CLI
    # entrypoints.
Ethan Xu's avatar
Ethan Xu committed
967
968
969
970
    parser = FlexibleArgumentParser(
        description="vLLM OpenAI-Compatible RESTful API server.")
    parser = make_arg_parser(parser)
    args = parser.parse_args()
971
    validate_parsed_serve_args(args)
972

973
    uvloop.run(run_server(args))