api_server.py 22.5 KB
Newer Older
1
import asyncio
2
3
import importlib
import inspect
4
import multiprocessing
5
import os
6
import re
7
import signal
8
import socket
9
import tempfile
10
from argparse import Namespace
11
from contextlib import asynccontextmanager
12
from functools import partial
13
from http import HTTPStatus
14
from typing import AsyncIterator, Optional, Set
15

16
import uvloop
17
from fastapi import APIRouter, FastAPI, Request
Zhuohan Li's avatar
Zhuohan Li committed
18
19
from fastapi.exceptions import RequestValidationError
from fastapi.middleware.cors import CORSMiddleware
20
from fastapi.responses import JSONResponse, Response, StreamingResponse
21
from starlette.datastructures import State
22
from starlette.routing import Mount
23
from typing_extensions import assert_never
Zhuohan Li's avatar
Zhuohan Li committed
24

25
import vllm.envs as envs
26
from vllm.config import ModelConfig
Woosuk Kwon's avatar
Woosuk Kwon committed
27
28
from vllm.engine.arg_utils import AsyncEngineArgs
from vllm.engine.async_llm_engine import AsyncLLMEngine
29
30
31
from vllm.engine.multiprocessing.client import MQLLMEngineClient
from vllm.engine.multiprocessing.engine import run_mp_engine
from vllm.engine.protocol import EngineClient
32
from vllm.entrypoints.launcher import serve_http
33
from vllm.entrypoints.logger import RequestLogger
34
35
from vllm.entrypoints.openai.cli_args import (make_arg_parser,
                                              validate_parsed_serve_args)
36
37
# yapf conflicts with isort for this block
# yapf: disable
38
from vllm.entrypoints.openai.protocol import (ChatCompletionRequest,
39
                                              ChatCompletionResponse,
40
                                              CompletionRequest,
41
                                              CompletionResponse,
42
43
                                              DetokenizeRequest,
                                              DetokenizeResponse,
44
45
                                              EmbeddingRequest,
                                              EmbeddingResponse, ErrorResponse,
46
                                              LoadLoraAdapterRequest,
47
                                              TokenizeRequest,
48
49
50
                                              TokenizeResponse,
                                              UnloadLoraAdapterRequest)
# yapf: enable
51
52
from vllm.entrypoints.openai.serving_chat import OpenAIServingChat
from vllm.entrypoints.openai.serving_completion import OpenAIServingCompletion
53
from vllm.entrypoints.openai.serving_embedding import OpenAIServingEmbedding
54
from vllm.entrypoints.openai.serving_engine import BaseModelPath, OpenAIServing
55
56
from vllm.entrypoints.openai.serving_tokenization import (
    OpenAIServingTokenization)
57
from vllm.entrypoints.openai.tool_parsers import ToolParserManager
58
from vllm.logger import init_logger
yhu422's avatar
yhu422 committed
59
from vllm.usage.usage_lib import UsageContext
60
from vllm.utils import FlexibleArgumentParser, get_open_zmq_ipc_path
61
from vllm.version import __version__ as VLLM_VERSION
Zhuohan Li's avatar
Zhuohan Li committed
62

63
TIMEOUT_KEEP_ALIVE = 5  # seconds
Zhuohan Li's avatar
Zhuohan Li committed
64

65
prometheus_multiproc_dir: tempfile.TemporaryDirectory
66

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

70
_running_tasks: Set[asyncio.Task] = set()
71

72

73
@asynccontextmanager
74
async def lifespan(app: FastAPI):
75
76
    try:
        if app.state.log_stats:
77
            engine_client: EngineClient = app.state.engine_client
78
79
80

            async def _force_log():
                while True:
81
82
                    await asyncio.sleep(10.)
                    await engine_client.do_log_stats()
83
84
85
86
87
88
89
90
91
92
93
94
95
96

            task = asyncio.create_task(_force_log())
            _running_tasks.add(task)
            task.add_done_callback(_running_tasks.remove)
        else:
            task = None
        try:
            yield
        finally:
            if task is not None:
                task.cancel()
    finally:
        # Ensure app state including engine ref is gc'd
        del app.state
97
98


99
@asynccontextmanager
100
async def build_async_engine_client(
101
        args: Namespace) -> AsyncIterator[EngineClient]:
102

103
    # Context manager to handle engine_client lifecycle
104
105
106
    # Ensures everything is shutdown and cleaned up on error/exit
    engine_args = AsyncEngineArgs.from_cli_args(args)

107
108
109
110
111
112
113
114
115
    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,
116
) -> AsyncIterator[EngineClient]:
117
    """
118
    Create EngineClient, either:
119
120
121
122
123
124
        - in-process using the AsyncLLMEngine Directly
        - multiprocess using AsyncLLMEngine RPC

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

125
126
127
    # Fall back
    # TODO: fill out feature matrix.
    if (MQLLMEngineClient.is_unsupported_config(engine_args)
128
            or disable_frontend_multiprocessing):
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
        engine_config = engine_args.create_engine_config()
        uses_ray = getattr(AsyncLLMEngine._get_executor_cls(engine_config),
                           "uses_ray", False)

        build_engine = partial(AsyncLLMEngine.from_engine_args,
                               engine_args=engine_args,
                               engine_config=engine_config,
                               usage_context=UsageContext.OPENAI_API_SERVER)
        if uses_ray:
            # Must run in main thread with ray for its signal handlers to work
            engine_client = build_engine()
        else:
            engine_client = await asyncio.get_running_loop().run_in_executor(
                None, build_engine)

        yield engine_client
145
146
147
148
        return

    # Otherwise, use the multiprocessing AsyncLLMEngine.
    else:
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
        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.")

164
        # Select random path for IPC.
165
166
167
        ipc_path = get_open_zmq_ipc_path()
        logger.info("Multiprocessing frontend to use %s for IPC Path.",
                    ipc_path)
168

169
        # Start RPCServer in separate process (holds the LLMEngine).
170
171
        # the current process might have CUDA context,
        # so we need to spawn a new process
172
173
174
175
176
177
178
        context = multiprocessing.get_context("spawn")

        engine_process = context.Process(target=run_mp_engine,
                                         args=(engine_args,
                                               UsageContext.OPENAI_API_SERVER,
                                               ipc_path))
        engine_process.start()
179
180
181
        engine_pid = engine_process.pid
        assert engine_pid is not None, "Engine process failed to start"
        logger.info("Started engine process with PID %d", engine_pid)
182
183
184
185
186

        # Build RPCClient, which conforms to EngineClient Protocol.
        # NOTE: Actually, this is not true yet. We still need to support
        # embedding models via RPC (see TODO above)
        engine_config = engine_args.create_engine_config()
187
188
        mp_engine_client = MQLLMEngineClient(ipc_path, engine_config,
                                             engine_pid)
189
190

        try:
191
192
            while True:
                try:
193
                    await mp_engine_client.setup()
194
                    break
195
                except TimeoutError:
196
                    if not engine_process.is_alive():
197
198
                        raise RuntimeError(
                            "Engine process failed to start") from None
199

200
            yield mp_engine_client  # type: ignore[misc]
201
202
        finally:
            # Ensure rpc server process was terminated
203
            engine_process.terminate()
204
205

            # Close all open connections to the backend
206
            mp_engine_client.close()
207

208
209
210
211
212
            # 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()
213

214
215
216
217
218
            # 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
219
            multiprocess.mark_process_dead(engine_process.pid)
220

221

Ethan Xu's avatar
Ethan Xu committed
222
router = APIRouter()
Zhuohan Li's avatar
Zhuohan Li committed
223

224

225
def mount_metrics(app: FastAPI):
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
    # 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:
        logger.info("vLLM to use %s as PROMETHEUS_MULTIPROC_DIR",
                    prometheus_multiproc_dir_path)
        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())

246
    # Workaround for 307 Redirect for /metrics
247
    metrics_route.path_regex = re.compile("^/metrics(?P<path>.*)$")
248
    app.routes.append(metrics_route)
249
250


251
252
253
254
255
256
def base(request: Request) -> OpenAIServing:
    # Reuse the existing instance
    return tokenization(request)


def chat(request: Request) -> Optional[OpenAIServingChat]:
257
258
259
    return request.app.state.openai_serving_chat


260
def completion(request: Request) -> Optional[OpenAIServingCompletion]:
261
262
263
    return request.app.state.openai_serving_completion


264
265
def embedding(request: Request) -> Optional[OpenAIServingEmbedding]:
    return request.app.state.openai_serving_embedding
266
267


268
269
def tokenization(request: Request) -> OpenAIServingTokenization:
    return request.app.state.openai_serving_tokenization
270
271


272
def engine_client(request: Request) -> EngineClient:
273
274
275
    return request.app.state.engine_client


Ethan Xu's avatar
Ethan Xu committed
276
@router.get("/health")
277
async def health(raw_request: Request) -> Response:
278
    """Health check."""
279
    await engine_client(raw_request).check_health()
280
281
282
    return Response(status_code=200)


Ethan Xu's avatar
Ethan Xu committed
283
@router.post("/tokenize")
284
async def tokenize(request: TokenizeRequest, raw_request: Request):
285
286
287
    handler = tokenization(raw_request)

    generator = await handler.create_tokenize(request)
288
289
290
    if isinstance(generator, ErrorResponse):
        return JSONResponse(content=generator.model_dump(),
                            status_code=generator.code)
291
    elif isinstance(generator, TokenizeResponse):
292
293
        return JSONResponse(content=generator.model_dump())

294
295
    assert_never(generator)

296

Ethan Xu's avatar
Ethan Xu committed
297
@router.post("/detokenize")
298
async def detokenize(request: DetokenizeRequest, raw_request: Request):
299
300
301
    handler = tokenization(raw_request)

    generator = await handler.create_detokenize(request)
302
303
304
    if isinstance(generator, ErrorResponse):
        return JSONResponse(content=generator.model_dump(),
                            status_code=generator.code)
305
    elif isinstance(generator, DetokenizeResponse):
306
307
        return JSONResponse(content=generator.model_dump())

308
309
    assert_never(generator)

310

Ethan Xu's avatar
Ethan Xu committed
311
@router.get("/v1/models")
312
async def show_available_models(raw_request: Request):
313
314
315
    handler = base(raw_request)

    models = await handler.show_available_models()
316
    return JSONResponse(content=models.model_dump())
Zhuohan Li's avatar
Zhuohan Li committed
317
318


Ethan Xu's avatar
Ethan Xu committed
319
@router.get("/version")
320
async def show_version():
321
    ver = {"version": VLLM_VERSION}
322
323
324
    return JSONResponse(content=ver)


Ethan Xu's avatar
Ethan Xu committed
325
@router.post("/v1/chat/completions")
326
327
async def create_chat_completion(request: ChatCompletionRequest,
                                 raw_request: Request):
328
329
330
331
    handler = chat(raw_request)
    if handler is None:
        return base(raw_request).create_error_response(
            message="The model does not support Chat Completions API")
332

333
    generator = await handler.create_chat_completion(request, raw_request)
334

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

339
    elif isinstance(generator, ChatCompletionResponse):
340
        return JSONResponse(content=generator.model_dump())
341

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

344

Ethan Xu's avatar
Ethan Xu committed
345
@router.post("/v1/completions")
346
async def create_completion(request: CompletionRequest, raw_request: Request):
347
348
349
350
351
352
    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)
353
354
355
    if isinstance(generator, ErrorResponse):
        return JSONResponse(content=generator.model_dump(),
                            status_code=generator.code)
356
    elif isinstance(generator, CompletionResponse):
357
        return JSONResponse(content=generator.model_dump())
Zhuohan Li's avatar
Zhuohan Li committed
358

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

Zhuohan Li's avatar
Zhuohan Li committed
361

Ethan Xu's avatar
Ethan Xu committed
362
@router.post("/v1/embeddings")
363
async def create_embedding(request: EmbeddingRequest, raw_request: Request):
364
365
366
367
368
369
    handler = embedding(raw_request)
    if handler is None:
        return base(raw_request).create_error_response(
            message="The model does not support Embeddings API")

    generator = await handler.create_embedding(request, raw_request)
370
371
372
    if isinstance(generator, ErrorResponse):
        return JSONResponse(content=generator.model_dump(),
                            status_code=generator.code)
373
    elif isinstance(generator, EmbeddingResponse):
374
375
        return JSONResponse(content=generator.model_dump())

376
377
    assert_never(generator)

378

379
380
381
382
383
384
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")
385
    async def start_profile(raw_request: Request):
386
        logger.info("Starting profiler...")
387
        await engine_client(raw_request).start_profile()
388
389
390
391
        logger.info("Profiler started.")
        return Response(status_code=200)

    @router.post("/stop_profile")
392
    async def stop_profile(raw_request: Request):
393
        logger.info("Stopping profiler...")
394
        await engine_client(raw_request).stop_profile()
395
396
397
398
        logger.info("Profiler stopped.")
        return Response(status_code=200)


399
400
401
402
403
404
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!")

    @router.post("/v1/load_lora_adapter")
405
406
    async def load_lora_adapter(request: LoadLoraAdapterRequest,
                                raw_request: Request):
407
408
409
410
411
412
413
        for route in [chat, completion, embedding]:
            handler = route(raw_request)
            if handler is not None:
                response = await handler.load_lora_adapter(request)
                if isinstance(response, ErrorResponse):
                    return JSONResponse(content=response.model_dump(),
                                        status_code=response.code)
414
415
416
417

        return Response(status_code=200, content=response)

    @router.post("/v1/unload_lora_adapter")
418
419
    async def unload_lora_adapter(request: UnloadLoraAdapterRequest,
                                  raw_request: Request):
420
421
422
423
424
425
426
        for route in [chat, completion, embedding]:
            handler = route(raw_request)
            if handler is not None:
                response = await handler.unload_lora_adapter(request)
                if isinstance(response, ErrorResponse):
                    return JSONResponse(content=response.model_dump(),
                                        status_code=response.code)
427
428
429
430

        return Response(status_code=200, content=response)


431
def build_app(args: Namespace) -> FastAPI:
432
433
434
435
436
437
438
    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
439
440
    app.include_router(router)
    app.root_path = args.root_path
Zhuohan Li's avatar
Zhuohan Li committed
441

442
443
    mount_metrics(app)

Zhuohan Li's avatar
Zhuohan Li committed
444
445
446
447
448
449
450
451
    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
452
453
    @app.exception_handler(RequestValidationError)
    async def validation_exception_handler(_, exc):
454
455
        chat = app.state.openai_serving_chat
        err = chat.create_error_response(message=str(exc))
Ethan Xu's avatar
Ethan Xu committed
456
457
458
        return JSONResponse(err.model_dump(),
                            status_code=HTTPStatus.BAD_REQUEST)

459
    if token := envs.VLLM_API_KEY or args.api_key:
460
461
462

        @app.middleware("http")
        async def authentication(request: Request, call_next):
463
            root_path = "" if args.root_path is None else args.root_path
464
465
            if request.method == "OPTIONS":
                return await call_next(request)
466
            if not request.url.path.startswith(f"{root_path}/v1"):
467
468
469
470
471
472
473
474
475
476
477
478
479
480
                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)

    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):
            app.add_middleware(imported)
        elif inspect.iscoroutinefunction(imported):
            app.middleware("http")(imported)
        else:
481
482
            raise ValueError(f"Invalid middleware {middleware}. "
                             f"Must be a function or a class.")
483

Ethan Xu's avatar
Ethan Xu committed
484
485
486
    return app


487
def init_app_state(
488
    engine_client: EngineClient,
489
490
    model_config: ModelConfig,
    state: State,
491
    args: Namespace,
492
) -> None:
493
    if args.served_model_name is not None:
494
        served_model_names = args.served_model_name
495
    else:
496
        served_model_names = [args.model]
497

498
499
500
501
502
    if args.disable_log_requests:
        request_logger = None
    else:
        request_logger = RequestLogger(max_log_len=args.max_log_len)

503
504
505
506
507
    base_model_paths = [
        BaseModelPath(name=name, model_path=args.model)
        for name in served_model_names
    ]

508
    state.engine_client = engine_client
509
    state.log_stats = not args.disable_log_stats
Ethan Xu's avatar
Ethan Xu committed
510

511
    state.openai_serving_chat = OpenAIServingChat(
512
        engine_client,
513
        model_config,
514
        base_model_paths,
515
516
517
518
519
        args.response_role,
        lora_modules=args.lora_modules,
        prompt_adapters=args.prompt_adapters,
        request_logger=request_logger,
        chat_template=args.chat_template,
520
        return_tokens_as_token_ids=args.return_tokens_as_token_ids,
521
        enable_auto_tools=args.enable_auto_tool_choice,
522
523
        tool_parser=args.tool_call_parser,
    ) if model_config.task == "generate" else None
524
    state.openai_serving_completion = OpenAIServingCompletion(
525
        engine_client,
526
        model_config,
527
        base_model_paths,
528
529
530
        lora_modules=args.lora_modules,
        prompt_adapters=args.prompt_adapters,
        request_logger=request_logger,
531
        return_tokens_as_token_ids=args.return_tokens_as_token_ids,
532
    ) if model_config.task == "generate" else None
533
    state.openai_serving_embedding = OpenAIServingEmbedding(
534
        engine_client,
535
        model_config,
536
        base_model_paths,
537
        request_logger=request_logger,
538
539
        chat_template=args.chat_template,
    ) if model_config.task == "embedding" else None
540
    state.openai_serving_tokenization = OpenAIServingTokenization(
541
        engine_client,
542
        model_config,
543
        base_model_paths,
544
545
546
547
        lora_modules=args.lora_modules,
        request_logger=request_logger,
        chat_template=args.chat_template,
    )
548
549


550
async def run_server(args, **uvicorn_kwargs) -> None:
551
552
553
    logger.info("vLLM API server version %s", VLLM_VERSION)
    logger.info("args: %s", args)

554
555
556
557
558
559
560
561
562
    if args.tool_parser_plugin and len(args.tool_parser_plugin) > 3:
        ToolParserManager.import_tool_parser(args.tool_parser_plugin)

    valide_tool_parses = ToolParserManager.tool_parsers.keys()
    if args.enable_auto_tool_choice \
        and args.tool_call_parser not in valide_tool_parses:
        raise KeyError(f"invalid tool call parser: {args.tool_call_parser} "
                       f"(chose from {{ {','.join(valide_tool_parses)} }})")

563
564
565
566
    # 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
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
567
    sock.bind(("", args.port))
568

569
570
571
572
573
574
    def signal_handler(*_) -> None:
        # Interrupt server on sigterm while initializing
        raise KeyboardInterrupt("terminated")

    signal.signal(signal.SIGTERM, signal_handler)

575
    async with build_async_engine_client(args) as engine_client:
576
577
        app = build_app(args)

578
579
        model_config = await engine_client.get_model_config()
        init_app_state(engine_client, model_config, app.state, args)
580
581
582
583
584
585
586
587
588
589
590

        shutdown_task = await serve_http(
            app,
            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,
591
            fd=sock.fileno(),
592
593
594
            **uvicorn_kwargs,
        )

595
596
    # NB: Await server shutdown only after the backend context is exited
    await shutdown_task
597

Ethan Xu's avatar
Ethan Xu committed
598
599
600
601
602
603
604
605

if __name__ == "__main__":
    # NOTE(simon):
    # This section should be in sync with vllm/scripts.py for CLI entrypoints.
    parser = FlexibleArgumentParser(
        description="vLLM OpenAI-Compatible RESTful API server.")
    parser = make_arg_parser(parser)
    args = parser.parse_args()
606
    validate_parsed_serve_args(args)
607

608
    uvloop.run(run_server(args))