| 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 tokio::time::timeout; |
| 10 | use tracing_subscriber::EnvFilter; |
| 11 | use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams; |
| 12 | use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, TransportMode}; |
| 13 | use vllm_llm::{FinishReason, GenerateOutputStream, GenerateRequest, Llm}; |
| 14 | |
| 15 | const PROMPT_TOKEN_IDS: &[u32] = &[20841, 448, 6896, 25, 23811]; |
| 16 | |
| 17 | #[derive(Debug, Parser)] |
| 18 | #[command(about = "Smoke-test the Rust LLM facade against an external vLLM engine.")] |
| 19 | struct Args { |
| 20 | #[arg(long)] |
| 21 | handshake_address: String, |
| 22 | #[arg(long, default_value_t = 1)] |
| 23 | engine_count: usize, |
| 24 | #[arg(long, default_value = "Qwen/Qwen3-0.6B")] |
| 25 | model: String, |
| 26 | #[arg(long, default_value = "127.0.0.1")] |
| 27 | host: String, |
| 28 | #[arg(long, default_value_t = 0)] |
| 29 | client_index: u32, |
| 30 | #[arg(long, default_value_t = 30)] |
| 31 | ready_timeout_secs: u64, |
| 32 | #[arg(long, default_value_t = 120)] |
| 33 | output_timeout_secs: u64, |
| 34 | #[arg(long, default_value_t = 5)] |
| 35 | max_tokens: u32, |
| 36 | } |
| 37 | |
| 38 | fn unique_request_id() -> String { |
| 39 | format!("rust-llm-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 | fn build_request(request_id: String, max_tokens: u32) -> GenerateRequest { |
| 49 | GenerateRequest { |
| 50 | request_id, |
| 51 | prompt_token_ids: PROMPT_TOKEN_IDS.to_vec(), |
| 52 | sampling_params: EngineCoreSamplingParams { |
| 53 | max_tokens, |
| 54 | ..EngineCoreSamplingParams::for_test() |
| 55 | }, |
| 56 | mm_features: None, |
| 57 | arrival_time: None, |
| 58 | cache_salt: None, |
| 59 | trace_headers: None, |
| 60 | priority: 0, |
| 61 | data_parallel_rank: None, |
| 62 | session_id: None, |
| 63 | kv_hints: None, |
| 64 | reasoning_parser_kwargs: None, |
| 65 | reasoning_ended: None, |
| 66 | lora_request: None, |
| 67 | } |
| 68 | } |
| 69 | |
| 70 | #[derive(Debug)] |
| 71 | struct CompletedRequest { |
| 72 | token_ids: Vec<u32>, |
| 73 | finish_reason: FinishReason, |
| 74 | } |
| 75 | |
| 76 | async fn wait_for_request_completion(mut stream: GenerateOutputStream) -> Result<CompletedRequest> { |
| 77 | let output = match stream.next().await { |
| 78 | Some(output) => output.context("failed to receive request output")?, |
| 79 | None => bail!("request stream ended without a final output"), |
| 80 | }; |
| 81 | |
| 82 | let none = stream.next().await; |
| 83 | assert!( |
| 84 | none.is_none(), |
| 85 | "expected final-only stream to end after the final output" |
| 86 | ); |
| 87 | |
| 88 | let finish_reason = output.finish_reason.expect("final-only output must have a finish reason"); |
| 89 | let token_ids = output.token_ids; |
| 90 | |
| 91 | Ok(CompletedRequest { |
| 92 | token_ids, |
| 93 | finish_reason, |
| 94 | }) |
| 95 | } |
| 96 | |
| 97 | async fn wait_for_timeout( |
| 98 | stream: GenerateOutputStream, |
| 99 | output_timeout: Duration, |
| 100 | ) -> Result<CompletedRequest> { |
| 101 | timeout(output_timeout, wait_for_request_completion(stream)) |
| 102 | .await |
| 103 | .context("timed out waiting for request output")? |
| 104 | } |
| 105 | |
| 106 | #[tokio::main(flavor = "multi_thread")] |
| 107 | async fn main() -> Result<()> { |
| 108 | init_tracing(); |
| 109 | let args = Args::parse(); |
| 110 | let ready_timeout = Duration::from_secs(args.ready_timeout_secs); |
| 111 | let output_timeout = Duration::from_secs(args.output_timeout_secs); |
| 112 | let request_id = unique_request_id(); |
| 113 | let client = EngineCoreClient::connect(EngineCoreClientConfig { |
| 114 | transport_mode: TransportMode::HandshakeOwner { |
| 115 | handshake_address: args.handshake_address.clone(), |
| 116 | advertised_host: args.host.clone(), |
| 117 | engine_count: args.engine_count, |
| 118 | ready_timeout, |
| 119 | local_input_address: None, |
| 120 | local_output_address: None, |
| 121 | }, |
| 122 | coordinator_mode: None, |
| 123 | model_name: args.model.clone(), |
| 124 | client_index: args.client_index, |
| 125 | engine_stats_enabled: true, |
| 126 | }) |
| 127 | .await |
| 128 | .context("failed to connect to external vLLM engine")?; |
| 129 | |
| 130 | println!("model={}", args.model); |
| 131 | println!("handshake_address={}", args.handshake_address); |
| 132 | println!("engine_count={}", args.engine_count); |
| 133 | println!("input_address={}", client.input_address()); |
| 134 | println!("output_address={}", client.output_address()); |
| 135 | println!("engine_identities={:x?}", client.engine_identities()); |
| 136 | |
| 137 | let llm = Llm::new(client); |
| 138 | let request = build_request(request_id.clone(), args.max_tokens); |
| 139 | println!("request_id={request_id}"); |
| 140 | println!("prompt_token_ids={PROMPT_TOKEN_IDS:?}"); |
| 141 | |
| 142 | let stream = llm.generate(request).await.context("failed to submit generate request")?; |
| 143 | let output = wait_for_timeout(stream, output_timeout).await?; |
| 144 | |
| 145 | llm.shutdown().await.context("failed to shut down llm client")?; |
| 146 | |
| 147 | println!("token_ids={:?}", output.token_ids); |
| 148 | println!("finish_reason={:?}", output.finish_reason); |
| 149 | |
| 150 | Ok(()) |
| 151 | } |