@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 163 lines 6.36 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import warnings
4 from argparse import Namespace
5 from typing import cast
6
7 from starlette.datastructures import State
8
9 from vllm.engine.protocol import EngineClient
10 from vllm.entrypoints.chat_utils import load_chat_template
11 from vllm.entrypoints.launchers.cli_args import resolve_default_chat_template_kwargs
12 from vllm.entrypoints.mcp.tool_server import init_tool_server
13 from vllm.entrypoints.openai.models.protocol import BaseModelPath
14 from vllm.entrypoints.openai.models.serving import OpenAIServingModels
15 from vllm.entrypoints.serve.tokenize.serving import ServingTokenization
16 from vllm.entrypoints.serve.utils.api_utils import process_lora_modules
17 from vllm.entrypoints.serve.utils.request_logger import RequestLogger
18 from vllm.plugins.endpoint_plugins.interface import init_endpoint_plugins_state
19 from vllm.renderers.online_derenderer import OnlineDerenderer
20 from vllm.renderers.online_renderer import OnlineRenderer
21 from vllm.tasks import FALLBACK_SUPPORTED_TASKS, POOLING_TASKS, SupportedTask
22
23
24 async def init_app_state(
25 engine_client: EngineClient,
26 state: State,
27 args: Namespace,
28 supported_tasks: tuple["SupportedTask", ...] | None = None,
29 ) -> None:
30 vllm_config = engine_client.vllm_config
31
32 if args.tool_call_parser is not None:
33 from vllm.parser.metrics import init_parser_metrics
34
35 init_parser_metrics(
36 model_name=cast(str, vllm_config.model_config.served_model_name)
37 )
38
39 if supported_tasks is None:
40 warnings.warn(
41 "The 'supported_tasks' parameter was not provided to "
42 "init_app_state and will be required in a future version. "
43 "Please pass 'supported_tasks' explicitly.",
44 DeprecationWarning,
45 stacklevel=2,
46 )
47 supported_tasks = FALLBACK_SUPPORTED_TASKS
48
49 if args.served_model_name is not None:
50 served_model_names = args.served_model_name
51 else:
52 served_model_names = [args.model]
53
54 if args.enable_log_requests:
55 request_logger = RequestLogger(max_log_len=args.max_log_len)
56 else:
57 request_logger = None
58
59 base_model_paths = [
60 BaseModelPath(name=name, model_path=args.model) for name in served_model_names
61 ]
62
63 state.engine_client = engine_client
64 state.log_stats = not args.disable_log_stats
65 state.vllm_config = vllm_config
66 state.args = args
67 resolved_chat_template = load_chat_template(args.chat_template)
68 default_chat_template_kwargs = resolve_default_chat_template_kwargs(args)
69 state.tool_server = (
70 await init_tool_server(args) if "generate" in supported_tasks else None
71 )
72
73 # Merge default_mm_loras into the static lora_modules
74 default_mm_loras = (
75 vllm_config.lora_config.default_mm_loras
76 if vllm_config.lora_config is not None
77 else {}
78 )
79 lora_modules = process_lora_modules(args.lora_modules, default_mm_loras)
80
81 state.openai_serving_models = OpenAIServingModels(
82 engine_client=engine_client,
83 base_model_paths=base_model_paths,
84 lora_modules=lora_modules,
85 )
86 await state.openai_serving_models.init_static_loras()
87
88 state.online_renderer = OnlineRenderer(
89 model_config=engine_client.model_config,
90 renderer=engine_client.renderer,
91 request_logger=request_logger,
92 chat_template=resolved_chat_template,
93 chat_template_content_format=args.chat_template_content_format,
94 trust_request_chat_template=args.trust_request_chat_template,
95 trust_request_mm_kwargs=args.trust_request_mm_kwargs,
96 enable_auto_tools=args.enable_auto_tool_choice,
97 exclude_tools_when_tool_choice_none=args.exclude_tools_when_tool_choice_none,
98 tool_parser=args.tool_call_parser,
99 tool_strict_level=args.tool_strict_level,
100 reasoning_parser=args.structured_outputs_config.reasoning_parser,
101 default_chat_template_kwargs=default_chat_template_kwargs,
102 log_error_stack=args.log_error_stack,
103 )
104 state.online_renderer.warmup()
105
106 state.online_derenderer = OnlineDerenderer(
107 model_config=engine_client.model_config,
108 renderer=engine_client.renderer,
109 request_logger=request_logger,
110 chat_template=resolved_chat_template,
111 chat_template_content_format=args.chat_template_content_format,
112 trust_request_chat_template=args.trust_request_chat_template,
113 enable_auto_tools=args.enable_auto_tool_choice,
114 exclude_tools_when_tool_choice_none=args.exclude_tools_when_tool_choice_none,
115 tool_parser=args.tool_call_parser,
116 tool_strict_level=args.tool_strict_level,
117 reasoning_parser=args.structured_outputs_config.reasoning_parser,
118 default_chat_template_kwargs=default_chat_template_kwargs,
119 log_error_stack=args.log_error_stack,
120 )
121
122 state.serving_tokenization = ServingTokenization(
123 state.openai_serving_models,
124 state.online_renderer,
125 request_logger=request_logger,
126 chat_template=resolved_chat_template,
127 chat_template_content_format=args.chat_template_content_format,
128 default_chat_template_kwargs=default_chat_template_kwargs,
129 trust_request_chat_template=args.trust_request_chat_template,
130 )
131
132 if "generate" in supported_tasks:
133 from vllm.entrypoints.generate.api_router import init_generate_state
134
135 await init_generate_state(
136 engine_client,
137 state,
138 args,
139 request_logger,
140 supported_tasks,
141 default_chat_template_kwargs,
142 )
143
144 from vllm.entrypoints.scale_out.factories import init_scale_out_state
145
146 init_scale_out_state(state, args, engine_client, request_logger)
147
148 if "transcription" in supported_tasks or "realtime" in supported_tasks:
149 from vllm.entrypoints.speech_to_text.factories import init_speech_to_text_state
150
151 init_speech_to_text_state(
152 engine_client, state, args, request_logger, supported_tasks
153 )
154
155 if any(task in POOLING_TASKS for task in supported_tasks):
156 from vllm.entrypoints.pooling.factories import init_pooling_state
157
158 init_pooling_state(engine_client, state, args, request_logger, supported_tasks)
159
160 await init_endpoint_plugins_state(engine_client, state, args)
161
162 state.enable_server_load_tracking = args.enable_server_load_tracking
163 state.server_load_metrics = 0