@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
rs 182 lines 6.45 KB
Raw
1 // SPDX-License-Identifier: Apache-2.0
2 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 use std::time::Duration;
5
6 use anyhow::{Context, Result, bail};
7 use clap::Parser;
8 use futures::StreamExt as _;
9 use tracing_subscriber::EnvFilter;
10 use vllm_chat::{
11 AssistantBlockKind, AssistantMessageExt as _, ChatEvent, ChatLlm, ChatMessage, ChatRequest,
12 ChatRole, SamplingParams, load_model_backends,
13 };
14 use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, TransportMode};
15 use vllm_llm::Llm;
16 use vllm_text::TextLlm;
17
18 #[derive(Debug, Parser)]
19 #[command(about = "Smoke-test the Rust chat facade against an external Qwen vLLM engine.")]
20 struct Args {
21 #[arg(long)]
22 handshake_address: String,
23 #[arg(long, default_value_t = 1)]
24 engine_count: usize,
25 #[arg(long, default_value = "Qwen/Qwen3-0.6B")]
26 model: String,
27 #[arg(long, default_value = "127.0.0.1")]
28 host: String,
29 #[arg(long, default_value_t = 30)]
30 ready_timeout_secs: u64,
31 #[arg(long)]
32 prompt: String,
33 }
34
35 const CLIENT_INDEX: u32 = 0;
36 const OUTPUT_TIMEOUT_SECS: u64 = 120;
37
38 fn unique_request_id() -> String {
39 format!("rust-chat-smoke-{}", uuid::Uuid::new_v4())
40 }
41
42 fn init_tracing() {
43 let filter = EnvFilter::try_from_default_env()
44 .unwrap_or_else(|_| EnvFilter::new("vllm_engine_core_client=debug"));
45 let _ = tracing_subscriber::fmt().with_env_filter(filter).try_init();
46 }
47
48 #[tokio::main(flavor = "multi_thread")]
49 async fn main() -> Result<()> {
50 init_tracing();
51 let args = Args::parse();
52 let loaded = load_model_backends(&args.model, Default::default())
53 .await
54 .with_context(|| format!("failed to load backends for {}", args.model))?;
55 let text_backend = loaded.text_backend;
56 let chat_backend = loaded.chat_backend;
57
58 let ready_timeout = Duration::from_secs(args.ready_timeout_secs);
59 let output_timeout = Duration::from_secs(OUTPUT_TIMEOUT_SECS);
60 let request_id = unique_request_id();
61 let client = EngineCoreClient::connect(EngineCoreClientConfig {
62 transport_mode: TransportMode::HandshakeOwner {
63 handshake_address: args.handshake_address.clone(),
64 advertised_host: args.host.clone(),
65 engine_count: args.engine_count,
66 ready_timeout,
67 local_input_address: None,
68 local_output_address: None,
69 },
70 coordinator_mode: None,
71 model_name: args.model.clone(),
72 client_index: CLIENT_INDEX,
73 engine_stats_enabled: true,
74 })
75 .await
76 .context("failed to connect to external vLLM engine")?;
77
78 println!("model={}", args.model);
79 println!("tokenizer_source=tokenizers + hf-hub");
80 println!("chat_template_source=tokenizer_config.json or adjacent chat template file");
81 println!("handshake_address={}", args.handshake_address);
82 println!("engine_count={}", args.engine_count);
83 println!("input_address={}", client.input_address());
84 println!("output_address={}", client.output_address());
85 println!("engine_identities={:x?}", client.engine_identities());
86
87 let llm = Llm::new(client);
88 let chat = ChatLlm::new(TextLlm::new(llm, text_backend), chat_backend);
89
90 let request = ChatRequest {
91 messages: vec![ChatMessage::text(ChatRole::User, args.prompt.clone())],
92 sampling_params: SamplingParams {
93 temperature: Some(0.0),
94 ..Default::default()
95 },
96 request_id: request_id.clone(),
97 ..ChatRequest::for_test()
98 };
99
100 println!("request_id={request_id}");
101 println!("prompt={}", args.prompt);
102
103 let mut stream = chat.chat(request).await.context("failed to submit chat request")?;
104 let output = tokio::time::timeout(output_timeout, async {
105 let mut final_reasoning = String::new();
106 let mut final_text = String::new();
107 let mut final_output_token_count = 0usize;
108 let mut finish_reason = None;
109 let mut saw_start = false;
110 let mut saw_stream_output = false;
111
112 while let Some(event) = stream.next().await.transpose()? {
113 match event {
114 ChatEvent::Start { .. } => {
115 saw_start = true;
116 }
117 ChatEvent::BlockStart { kind, .. } => {
118 if saw_stream_output {
119 println!();
120 }
121 match kind {
122 AssistantBlockKind::Reasoning => print!("[reasoning] "),
123 AssistantBlockKind::Text => print!("[answer] "),
124 AssistantBlockKind::ToolCall => {}
125 }
126 saw_stream_output = true;
127 }
128 ChatEvent::ToolCallStart { name, .. } => {
129 if saw_stream_output {
130 println!();
131 }
132 print!("[tool:{name}] ");
133 saw_stream_output = true;
134 }
135 ChatEvent::LogprobsDelta { .. } => {}
136 ChatEvent::Done {
137 message,
138 usage,
139 finish_reason: reason,
140 ..
141 } => {
142 final_reasoning = message.reasoning().unwrap_or_default();
143 final_text = message.text();
144 final_output_token_count = usage.output_token_count;
145 finish_reason = Some(reason);
146 break;
147 }
148 ChatEvent::BlockDelta { kind, delta, .. } => match kind {
149 AssistantBlockKind::Reasoning | AssistantBlockKind::Text => {
150 print!("{delta}");
151 }
152 AssistantBlockKind::ToolCall => {}
153 },
154 ChatEvent::ToolCallArgumentsDelta { delta, .. } => print!("{delta}"),
155 ChatEvent::BlockEnd { .. } | ChatEvent::ToolCallEnd { .. } => {}
156 }
157 }
158
159 println!();
160
161 if !saw_start {
162 bail!("chat stream ended without a start event");
163 }
164 Ok::<_, anyhow::Error>((
165 final_reasoning,
166 final_text,
167 final_output_token_count,
168 finish_reason,
169 ))
170 })
171 .await
172 .context("timed out waiting for chat output")??;
173
174 chat.shutdown().await.context("failed to shut down chat client")?;
175
176 println!("final_reasoning={:?}", output.0);
177 println!("final_text={:?}", output.1);
178 println!("final_output_token_count={:?}", output.2);
179 println!("finish_reason={:?}", output.3);
180
181 Ok(())
182 }