@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
rs 234 lines 7.73 KB
Raw
1 // SPDX-License-Identifier: Apache-2.0
2 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 use std::collections::HashMap;
5 use std::time::Duration;
6
7 use anyhow::{Context, Result, bail};
8 use async_openai::Client;
9 use async_openai::config::OpenAIConfig;
10 use async_openai::types::chat::{
11 ChatCompletionRequestUserMessageArgs, CreateChatCompletionRequest,
12 CreateChatCompletionRequestArgs,
13 };
14 use async_openai::types::models::ListModelResponse;
15 use clap::Parser;
16 use futures::StreamExt as _;
17 use tokio_util::sync::CancellationToken;
18 use tracing_subscriber::EnvFilter;
19 use vllm_chat::ToolStrictLevel;
20 use vllm_engine_core_client::TransportMode;
21 use vllm_server::{
22 ApiServerOptions, ChatTemplateContentFormatOption, Config, CoordinatorMode, CorsConfig,
23 HttpListenerMode, ParserSelection, RendererSelection, serve,
24 };
25
26 #[derive(Debug, Parser)]
27 #[command(
28 about = "Smoke-test the Rust OpenAI server with async-openai against an external Qwen vLLM engine."
29 )]
30 struct Args {
31 #[arg(long)]
32 handshake_address: String,
33 #[arg(long, default_value_t = 1)]
34 engine_count: usize,
35 #[arg(long, default_value = "Qwen/Qwen3-0.6B")]
36 model: String,
37 #[arg(long, default_value = "127.0.0.1")]
38 host: String,
39 #[arg(long, default_value_t = 30)]
40 ready_timeout_secs: u64,
41 #[arg(
42 long,
43 default_value = "What is the capital of France? Answer with one word."
44 )]
45 prompt: String,
46 }
47
48 #[tokio::main(flavor = "multi_thread")]
49 async fn main() -> Result<()> {
50 init_tracing();
51 let args = Args::parse();
52 let port = unique_local_port()?;
53 let config = Config {
54 transport_mode: TransportMode::HandshakeOwner {
55 handshake_address: args.handshake_address,
56 advertised_host: args.host,
57 engine_count: args.engine_count,
58 ready_timeout: Duration::from_secs(args.ready_timeout_secs),
59 local_input_address: None,
60 local_output_address: None,
61 },
62 coordinator_mode: CoordinatorMode::MaybeInProc,
63 model: args.model,
64 revision: None,
65 hf_overrides: Default::default(),
66 generation_config: Default::default(),
67 served_model_name: vec![],
68 listener_mode: HttpListenerMode::BindTcp {
69 host: "127.0.0.1".to_string(),
70 port,
71 },
72 grpc_listener_mode: None,
73 tool_call_parser: ParserSelection::Auto,
74 reasoning_parser: ParserSelection::Auto,
75 tool_strict_level: ToolStrictLevel::Auto,
76 renderer: RendererSelection::Auto,
77 language_model_only: false,
78 chat_template: None,
79 default_chat_template_kwargs: None,
80 limit_mm_per_prompt: HashMap::new(),
81 lora_modules: Vec::new(),
82 chat_template_content_format: ChatTemplateContentFormatOption::Auto,
83 max_logprobs: None,
84 api_server_options: ApiServerOptions::default(),
85 cors: CorsConfig::default(),
86 tls: None,
87 api_keys: Vec::new(),
88 disable_log_stats: false,
89 shutdown_timeout: Duration::ZERO,
90 manages_engine: false,
91 keep_alive_timeout: Duration::from_secs(5),
92 profiler: None,
93 };
94
95 let bind_address = format!("127.0.0.1:{port}");
96 let shutdown = CancellationToken::new();
97 let server_config = config.clone();
98 let server_shutdown = shutdown.clone();
99 let server_task = tokio::spawn(async move { serve(server_config, server_shutdown).await });
100
101 let client = Client::with_config(
102 OpenAIConfig::new()
103 .with_api_key("unused")
104 .with_api_base(format!("http://{bind_address}/v1")),
105 );
106
107 print_models(&client).await?;
108 let final_text = stream_completion(&client, &config.model, &args.prompt).await?;
109
110 println!();
111 println!("final_text={final_text:?}");
112
113 shutdown_server(server_task, shutdown).await
114 }
115
116 fn init_tracing() {
117 let filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info"));
118 let _ = tracing_subscriber::fmt().with_env_filter(filter).try_init();
119 }
120
121 fn unique_local_port() -> Result<u16> {
122 let listener = std::net::TcpListener::bind("127.0.0.1:0")
123 .context("failed to allocate local smoke-test port")?;
124 let port = listener.local_addr().context("failed to read local smoke-test port")?.port();
125 drop(listener);
126 Ok(port)
127 }
128
129 async fn print_models(client: &Client<OpenAIConfig>) -> Result<()> {
130 let models = wait_for_models(client).await?;
131 let model_ids = models.data.into_iter().map(|model| model.id).collect::<Vec<_>>();
132 println!("models={model_ids:?}");
133 Ok(())
134 }
135
136 async fn wait_for_models(client: &Client<OpenAIConfig>) -> Result<ListModelResponse> {
137 let mut last_error = None;
138 for _ in 0..240 {
139 match client.models().list().await {
140 Ok(models) => return Ok(models),
141 Err(error) => {
142 last_error = Some(error);
143 tokio::time::sleep(Duration::from_millis(500)).await;
144 }
145 }
146 }
147
148 match last_error {
149 Some(error) => Err(error).context("OpenAI server did not become ready in time"),
150 None => bail!("OpenAI server readiness loop finished without a result"),
151 }
152 }
153
154 async fn stream_completion(
155 client: &Client<OpenAIConfig>,
156 model: &str,
157 prompt: &str,
158 ) -> Result<String> {
159 // Keep this smoke test on async-openai's standard `create_stream` path so it
160 // exercises the ordinary typed chat-completions client without BYOT
161 // request/response types.
162 //
163 // The current async-openai chat-completions stream delta type does not expose
164 // our OpenAI-compatible `reasoning_content` extension field, so this
165 // example only validates the assistant role chunk, visible `content`
166 // deltas, and terminal finish chunk. Reasoning coverage lives in our own
167 // route tests and in the `vllm-chat` smoke example.
168 let request: CreateChatCompletionRequest = CreateChatCompletionRequestArgs::default()
169 .model(model)
170 .stream(true)
171 .temperature(0.0)
172 .max_completion_tokens(128u32)
173 .messages([ChatCompletionRequestUserMessageArgs::default()
174 .content(prompt)
175 .build()
176 .context("failed to build user chat message")?
177 .into()])
178 .build()
179 .context("failed to build chat completion request")?;
180
181 let mut stream = client
182 .chat()
183 .create_stream(request)
184 .await
185 .context("failed to create streaming chat completion")?;
186
187 let mut final_text = String::new();
188 let mut saw_role = false;
189 let mut saw_finish_reason = false;
190 let mut saw_text = false;
191
192 while let Some(chunk) = stream.next().await {
193 let chunk = chunk.context("streaming chat completion failed")?;
194 for choice in chunk.choices {
195 if choice.delta.role.is_some() {
196 saw_role = true;
197 }
198 if let Some(delta) = choice.delta.content {
199 if !saw_text {
200 print!("[answer] ");
201 }
202 print!("{delta}");
203 final_text.push_str(&delta);
204 saw_text = true;
205 }
206 if choice.finish_reason.is_some() {
207 saw_finish_reason = true;
208 }
209 }
210 }
211
212 if !saw_role {
213 bail!("stream ended without an assistant role chunk");
214 }
215 if !saw_finish_reason {
216 bail!("stream ended without a terminal finish reason");
217 }
218 if final_text.is_empty() {
219 bail!("stream ended without any content deltas");
220 }
221
222 Ok(final_text)
223 }
224
225 async fn shutdown_server(
226 server_task: tokio::task::JoinHandle<anyhow::Result<()>>,
227 shutdown: CancellationToken,
228 ) -> Result<()> {
229 shutdown.cancel();
230 server_task
231 .await
232 .context("server task join failed")?
233 .context("server task failed")
234 }