main
py 74 lines 2.32 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from argparse import Namespace
4
5 from fastapi import FastAPI
6
7 from vllm import envs
8 from vllm.config import ModelConfig
9 from vllm.tasks import POOLING_TASKS, SupportedTask
10
11
12 def register_api_routers(
13 args: Namespace,
14 app: FastAPI,
15 supported_tasks: tuple["SupportedTask", ...],
16 model_config: ModelConfig | None = None,
17 ):
18 from vllm.entrypoints.serve import register_vllm_serve_api_routers
19
20 register_vllm_serve_api_routers(app)
21
22 from vllm.entrypoints.openai.models.api_router import (
23 attach_router as register_models_api_router,
24 )
25
26 register_models_api_router(app)
27
28 from vllm.entrypoints.serve.sagemaker.api_router import (
29 attach_router as register_sagemaker_api_router,
30 )
31
32 register_sagemaker_api_router(app, supported_tasks, model_config)
33
34 if envs.VLLM_SERVER_DEV_MODE:
35 from vllm.entrypoints.serve import register_vllm_dev_api_routers
36
37 register_vllm_dev_api_routers(app)
38
39 if "generate" in supported_tasks:
40 from vllm.entrypoints.generate.api_router import (
41 register_generate_api_routers,
42 )
43
44 register_generate_api_routers(app)
45
46 from vllm.entrypoints.serve.elastic_ep.api_router import (
47 attach_router as elastic_ep_attach_router,
48 )
49
50 elastic_ep_attach_router(app)
51
52 if "generate" in supported_tasks or "render" in supported_tasks:
53 from vllm.entrypoints.scale_out.factories import register_scale_out_api_routers
54
55 register_scale_out_api_routers(app, supported_tasks)
56
57 if "transcription" in supported_tasks or "realtime" in supported_tasks:
58 from vllm.entrypoints.speech_to_text.factories import (
59 register_speech_to_text_api_routers,
60 )
61
62 register_speech_to_text_api_routers(app, supported_tasks)
63
64 if any(task in POOLING_TASKS for task in supported_tasks):
65 from vllm.entrypoints.pooling.factories import register_pooling_api_routers
66
67 register_pooling_api_routers(app, supported_tasks, model_config)
68
69 if getattr(args, "enable_fault_tolerance", False):
70 from vllm.entrypoints.serve.fault_tolerance.api_router import (
71 register_fault_tolerance_api_router,
72 )
73
74 register_fault_tolerance_api_router(app)