| 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::output::EngineCoreFinishReason; |
| 12 | use vllm_engine_core_client::protocol::request::EngineCoreRequest; |
| 13 | use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams; |
| 14 | use vllm_engine_core_client::{ |
| 15 | EngineCoreClient, EngineCoreClientConfig, EngineCoreStreamOutput, TransportMode, |
| 16 | }; |
| 17 | |
| 18 | const BASE_PROMPT_TOKEN_IDS: &[u32] = &[20841, 448, 6896, 25, 23811]; |
| 19 | |
| 20 | #[derive(Debug, Parser)] |
| 21 | #[command(about = "Smoke-test engine-core sample logprobs against an external vLLM engine.")] |
| 22 | struct Args { |
| 23 | #[arg(long)] |
| 24 | handshake_address: String, |
| 25 | #[arg(long, default_value_t = 1)] |
| 26 | engine_count: usize, |
| 27 | #[arg(long, default_value = "Qwen/Qwen3-0.6B")] |
| 28 | model: String, |
| 29 | #[arg(long, default_value = "127.0.0.1")] |
| 30 | host: String, |
| 31 | #[arg(long, default_value_t = 0)] |
| 32 | client_index: u32, |
| 33 | #[arg(long, default_value_t = 30)] |
| 34 | ready_timeout_secs: u64, |
| 35 | #[arg(long, default_value_t = 120)] |
| 36 | output_timeout_secs: u64, |
| 37 | #[arg(long, default_value_t = 1)] |
| 38 | max_tokens: u32, |
| 39 | #[arg(long, default_value_t = 2)] |
| 40 | logprobs: i32, |
| 41 | #[arg(long, default_value_t = 1)] |
| 42 | prompt_logprobs: i32, |
| 43 | #[arg(long, default_value_t = 96)] |
| 44 | prompt_repeats: usize, |
| 45 | } |
| 46 | |
| 47 | fn init_tracing() { |
| 48 | let filter = EnvFilter::try_from_default_env() |
| 49 | .unwrap_or_else(|_| EnvFilter::new("vllm_engine_core_client=debug")); |
| 50 | let _ = tracing_subscriber::fmt().with_env_filter(filter).try_init(); |
| 51 | } |
| 52 | |
| 53 | fn unique_request_id() -> String { |
| 54 | let nanos = std::time::SystemTime::now() |
| 55 | .duration_since(std::time::UNIX_EPOCH) |
| 56 | .expect("system clock should be after unix epoch") |
| 57 | .as_nanos(); |
| 58 | format!("rust-engine-core-logprobs-{nanos}") |
| 59 | } |
| 60 | |
| 61 | fn build_prompt_token_ids(prompt_repeats: usize) -> Vec<u32> { |
| 62 | let repeats = prompt_repeats.max(1); |
| 63 | BASE_PROMPT_TOKEN_IDS.repeat(repeats) |
| 64 | } |
| 65 | |
| 66 | fn build_request( |
| 67 | request_id: String, |
| 68 | prompt_token_ids: Vec<u32>, |
| 69 | max_tokens: u32, |
| 70 | logprobs: i32, |
| 71 | prompt_logprobs: i32, |
| 72 | client_index: u32, |
| 73 | ) -> EngineCoreRequest { |
| 74 | EngineCoreRequest { |
| 75 | request_id, |
| 76 | prompt_token_ids: Some(prompt_token_ids), |
| 77 | sampling_params: Some(EngineCoreSamplingParams { |
| 78 | max_tokens, |
| 79 | logprobs: Some(logprobs), |
| 80 | prompt_logprobs: Some(prompt_logprobs), |
| 81 | ..EngineCoreSamplingParams::for_test() |
| 82 | }), |
| 83 | arrival_time: 0.0, |
| 84 | client_index, |
| 85 | ..EngineCoreRequest::default() |
| 86 | } |
| 87 | } |
| 88 | |
| 89 | async fn wait_for_final_output( |
| 90 | mut stream: vllm_engine_core_client::EngineCoreOutputStream, |
| 91 | ) -> Result<EngineCoreStreamOutput> { |
| 92 | while let Some(output) = stream.next().await { |
| 93 | let output = output.context("failed to receive engine-core output")?; |
| 94 | if output.finished() { |
| 95 | return Ok(output); |
| 96 | } |
| 97 | } |
| 98 | bail!("request stream ended without a final output") |
| 99 | } |
| 100 | |
| 101 | #[tokio::main(flavor = "multi_thread")] |
| 102 | async fn main() -> Result<()> { |
| 103 | init_tracing(); |
| 104 | let args = Args::parse(); |
| 105 | let ready_timeout = Duration::from_secs(args.ready_timeout_secs); |
| 106 | let output_timeout = Duration::from_secs(args.output_timeout_secs); |
| 107 | let request_id = unique_request_id(); |
| 108 | let prompt_token_ids = build_prompt_token_ids(args.prompt_repeats); |
| 109 | let client = EngineCoreClient::connect(EngineCoreClientConfig { |
| 110 | transport_mode: TransportMode::HandshakeOwner { |
| 111 | handshake_address: args.handshake_address.clone(), |
| 112 | advertised_host: args.host.clone(), |
| 113 | engine_count: args.engine_count, |
| 114 | ready_timeout, |
| 115 | local_input_address: None, |
| 116 | local_output_address: None, |
| 117 | }, |
| 118 | coordinator_mode: None, |
| 119 | model_name: args.model.clone(), |
| 120 | client_index: args.client_index, |
| 121 | engine_stats_enabled: true, |
| 122 | }) |
| 123 | .await |
| 124 | .context("failed to connect to external vLLM engine")?; |
| 125 | |
| 126 | println!("model={}", args.model); |
| 127 | println!("handshake_address={}", args.handshake_address); |
| 128 | println!("engine_count={}", args.engine_count); |
| 129 | println!("input_address={}", client.input_address()); |
| 130 | println!("output_address={}", client.output_address()); |
| 131 | println!("engine_identities={:x?}", client.engine_identities()); |
| 132 | |
| 133 | let request = build_request( |
| 134 | request_id.clone(), |
| 135 | prompt_token_ids.clone(), |
| 136 | args.max_tokens, |
| 137 | args.logprobs, |
| 138 | args.prompt_logprobs, |
| 139 | args.client_index, |
| 140 | ); |
| 141 | println!("request_id={request_id}"); |
| 142 | println!("prompt_len={}", prompt_token_ids.len()); |
| 143 | println!("base_prompt_len={}", BASE_PROMPT_TOKEN_IDS.len()); |
| 144 | println!("prompt_repeats={}", args.prompt_repeats); |
| 145 | println!("requested_logprobs={}", args.logprobs); |
| 146 | println!("requested_prompt_logprobs={}", args.prompt_logprobs); |
| 147 | |
| 148 | let stream = client.call(request).await.context("failed to submit engine-core request")?; |
| 149 | let output = timeout(output_timeout, wait_for_final_output(stream)) |
| 150 | .await |
| 151 | .context("timed out waiting for final output")??; |
| 152 | |
| 153 | let finish_reason = output.finish_reason; |
| 154 | let token_ids = output.new_token_ids.clone(); |
| 155 | let logprobs = output |
| 156 | .new_logprobs |
| 157 | .as_ref() |
| 158 | .and_then(|value| value.as_direct()) |
| 159 | .context("engine output did not include decoded sample logprobs")?; |
| 160 | let prompt_logprobs = output |
| 161 | .new_prompt_logprobs_tensors |
| 162 | .as_ref() |
| 163 | .and_then(|value| value.as_direct()) |
| 164 | .context("engine output did not include decoded prompt logprobs")?; |
| 165 | |
| 166 | println!("token_ids={token_ids:?}"); |
| 167 | println!("finish_reason={finish_reason:?}"); |
| 168 | println!("new_logprobs={logprobs:#?}"); |
| 169 | println!("new_prompt_logprobs_tensors={prompt_logprobs:#?}"); |
| 170 | |
| 171 | client.shutdown().await.context("failed to shut down engine-core client")?; |
| 172 | |
| 173 | if finish_reason != Some(EngineCoreFinishReason::Length) { |
| 174 | bail!("unexpected finish_reason: expected Length, got {finish_reason:?}"); |
| 175 | } |
| 176 | if token_ids.is_empty() { |
| 177 | bail!("engine returned no generated token ids"); |
| 178 | } |
| 179 | if logprobs.is_empty() { |
| 180 | bail!("decoded logprobs payload is unexpectedly empty"); |
| 181 | } |
| 182 | if prompt_logprobs.is_empty() { |
| 183 | bail!("decoded prompt logprobs payload is unexpectedly empty"); |
| 184 | } |
| 185 | if prompt_logprobs.len() + 1 < prompt_token_ids.len() { |
| 186 | bail!( |
| 187 | "prompt logprobs rows look too short: prompt_len={}, rows={}", |
| 188 | prompt_token_ids.len(), |
| 189 | prompt_logprobs.len() |
| 190 | ); |
| 191 | } |
| 192 | |
| 193 | Ok(()) |
| 194 | } |