| 1 | // SPDX-License-Identifier: Apache-2.0 |
| 2 | // SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | use std::collections::BTreeSet; |
| 5 | use std::sync::{Arc, Once}; |
| 6 | use std::time::Duration; |
| 7 | |
| 8 | use futures::StreamExt as _; |
| 9 | use tokio::sync::oneshot; |
| 10 | use tokio::time::timeout; |
| 11 | use tracing_subscriber::EnvFilter; |
| 12 | use uuid::Uuid; |
| 13 | use vllm_engine_core_client::mock_engine::default_ready_response; |
| 14 | use vllm_engine_core_client::protocol::handshake::EngineCoreReadyResponse; |
| 15 | use vllm_engine_core_client::protocol::logprobs::{ |
| 16 | Logprobs, MaybeWireLogprobs, PositionLogprobs, TokenLogprob, |
| 17 | }; |
| 18 | use vllm_engine_core_client::protocol::lora::LoraRequest; |
| 19 | use vllm_engine_core_client::protocol::output::{ |
| 20 | EngineCoreEvent, EngineCoreEventType, EngineCoreFinishReason, EngineCoreOutput, |
| 21 | EngineCoreOutputs, RequestBatchOutputs, |
| 22 | }; |
| 23 | use vllm_engine_core_client::protocol::request::EngineCoreRequest; |
| 24 | use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams; |
| 25 | use vllm_engine_core_client::protocol::sampling_mask::{MaybeWireSamplingMask, SamplingMask}; |
| 26 | use vllm_engine_core_client::protocol::stats::PrefillStats; |
| 27 | use vllm_engine_core_client::test_utils::{ |
| 28 | IpcNamespace, spawn_mock_engine_task, spawn_mock_engine_task_with_ready, |
| 29 | }; |
| 30 | use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, EngineId}; |
| 31 | use vllm_llm::{ |
| 32 | Error, FinishReason, GenerateOutputStreamExt as _, GeneratePromptInfo, GenerateRequest, Llm, |
| 33 | }; |
| 34 | use vllm_metrics::METRICS; |
| 35 | use zeromq::prelude::{SocketRecv, SocketSend}; |
| 36 | use zeromq::{DealerSocket, PushSocket, ZmqMessage}; |
| 37 | |
| 38 | static TRACING: Once = Once::new(); |
| 39 | |
| 40 | fn request_output( |
| 41 | request_id: &str, |
| 42 | new_token_ids: Vec<u32>, |
| 43 | finish_reason: Option<EngineCoreFinishReason>, |
| 44 | ) -> EngineCoreOutput { |
| 45 | request_output_with_events(request_id, new_token_ids, finish_reason, None) |
| 46 | } |
| 47 | |
| 48 | fn request_output_with_events( |
| 49 | request_id: &str, |
| 50 | new_token_ids: Vec<u32>, |
| 51 | finish_reason: Option<EngineCoreFinishReason>, |
| 52 | events: Option<Vec<EngineCoreEvent>>, |
| 53 | ) -> EngineCoreOutput { |
| 54 | EngineCoreOutput { |
| 55 | request_id: request_id.to_string(), |
| 56 | new_token_ids, |
| 57 | finish_reason, |
| 58 | events, |
| 59 | ..Default::default() |
| 60 | } |
| 61 | } |
| 62 | |
| 63 | fn request_output_with_logprobs( |
| 64 | request_id: &str, |
| 65 | new_token_ids: Vec<u32>, |
| 66 | finish_reason: Option<EngineCoreFinishReason>, |
| 67 | new_logprobs: Option<Logprobs>, |
| 68 | prompt_logprobs: Option<Logprobs>, |
| 69 | ) -> EngineCoreOutput { |
| 70 | EngineCoreOutput { |
| 71 | request_id: request_id.to_string(), |
| 72 | new_token_ids, |
| 73 | new_logprobs: new_logprobs.map(MaybeWireLogprobs::Direct), |
| 74 | new_prompt_logprobs_tensors: prompt_logprobs.map(MaybeWireLogprobs::Direct), |
| 75 | finish_reason, |
| 76 | ..Default::default() |
| 77 | } |
| 78 | } |
| 79 | |
| 80 | fn request_output_with_logprobs_and_kv( |
| 81 | request_id: &str, |
| 82 | new_token_ids: Vec<u32>, |
| 83 | finish_reason: Option<EngineCoreFinishReason>, |
| 84 | new_logprobs: Option<Logprobs>, |
| 85 | prompt_logprobs: Option<Logprobs>, |
| 86 | kv_transfer_params: Option<serde_json::Value>, |
| 87 | ec_transfer_params: Option<serde_json::Value>, |
| 88 | ) -> EngineCoreOutput { |
| 89 | EngineCoreOutput { |
| 90 | request_id: request_id.to_string(), |
| 91 | new_token_ids, |
| 92 | new_logprobs: new_logprobs.map(MaybeWireLogprobs::Direct), |
| 93 | new_prompt_logprobs_tensors: prompt_logprobs.map(MaybeWireLogprobs::Direct), |
| 94 | finish_reason, |
| 95 | kv_transfer_params, |
| 96 | ec_transfer_params, |
| 97 | ..Default::default() |
| 98 | } |
| 99 | } |
| 100 | |
| 101 | fn logprobs_for_position( |
| 102 | sampled_token_id: u32, |
| 103 | sampled_logprob: f32, |
| 104 | sampled_rank: u32, |
| 105 | top_token_id: u32, |
| 106 | top_logprob: f32, |
| 107 | ) -> Logprobs { |
| 108 | Logprobs { |
| 109 | positions: vec![PositionLogprobs { |
| 110 | entries: vec![ |
| 111 | TokenLogprob { |
| 112 | token_id: sampled_token_id, |
| 113 | logprob: sampled_logprob, |
| 114 | rank: sampled_rank, |
| 115 | }, |
| 116 | TokenLogprob { |
| 117 | token_id: top_token_id, |
| 118 | logprob: top_logprob, |
| 119 | rank: 1, |
| 120 | }, |
| 121 | ], |
| 122 | }], |
| 123 | } |
| 124 | } |
| 125 | |
| 126 | fn prompt_logprobs() -> Logprobs { |
| 127 | Logprobs { |
| 128 | positions: vec![ |
| 129 | PositionLogprobs { |
| 130 | entries: vec![ |
| 131 | TokenLogprob { |
| 132 | token_id: 11, |
| 133 | logprob: -0.1, |
| 134 | rank: 2, |
| 135 | }, |
| 136 | TokenLogprob { |
| 137 | token_id: 7, |
| 138 | logprob: -0.05, |
| 139 | rank: 1, |
| 140 | }, |
| 141 | ], |
| 142 | }, |
| 143 | PositionLogprobs { |
| 144 | entries: vec![ |
| 145 | TokenLogprob { |
| 146 | token_id: 22, |
| 147 | logprob: -0.2, |
| 148 | rank: 3, |
| 149 | }, |
| 150 | TokenLogprob { |
| 151 | token_id: 8, |
| 152 | logprob: -0.1, |
| 153 | rank: 1, |
| 154 | }, |
| 155 | ], |
| 156 | }, |
| 157 | ], |
| 158 | } |
| 159 | } |
| 160 | |
| 161 | fn sample_generate_request(request_id: &str, max_tokens: u32) -> GenerateRequest { |
| 162 | GenerateRequest { |
| 163 | request_id: request_id.to_string(), |
| 164 | prompt_token_ids: vec![11, 22], |
| 165 | sampling_params: EngineCoreSamplingParams { |
| 166 | max_tokens, |
| 167 | ..EngineCoreSamplingParams::for_test() |
| 168 | }, |
| 169 | mm_features: None, |
| 170 | arrival_time: Some(42.5), |
| 171 | cache_salt: None, |
| 172 | trace_headers: None, |
| 173 | priority: 0, |
| 174 | data_parallel_rank: None, |
| 175 | session_id: None, |
| 176 | kv_hints: None, |
| 177 | reasoning_parser_kwargs: None, |
| 178 | reasoning_ended: None, |
| 179 | lora_request: None, |
| 180 | } |
| 181 | } |
| 182 | |
| 183 | async fn send_outputs(push: &mut PushSocket, outputs: EngineCoreOutputs) { |
| 184 | push.send(ZmqMessage::from(rmp_serde::to_vec_named(&outputs).unwrap())) |
| 185 | .await |
| 186 | .unwrap(); |
| 187 | } |
| 188 | |
| 189 | async fn recv_engine_message(dealer: &mut DealerSocket) -> Vec<bytes::Bytes> { |
| 190 | dealer.recv().await.unwrap().into_vec() |
| 191 | } |
| 192 | |
| 193 | async fn connect_async_llm_with_ipc( |
| 194 | handshake_address: String, |
| 195 | client_index: u32, |
| 196 | model_name: &str, |
| 197 | ipc: &IpcNamespace, |
| 198 | ) -> Llm { |
| 199 | let client = EngineCoreClient::connect( |
| 200 | EngineCoreClientConfig::new_single(handshake_address) |
| 201 | .with_model_name(model_name) |
| 202 | .with_client_index(client_index) |
| 203 | .with_local_input_output_addresses( |
| 204 | Some(ipc.input_endpoint()), |
| 205 | Some(ipc.output_endpoint()), |
| 206 | ), |
| 207 | ) |
| 208 | .await |
| 209 | .unwrap(); |
| 210 | Llm::new(client) |
| 211 | } |
| 212 | |
| 213 | fn request_metrics_model_name(prefix: &str) -> String { |
| 214 | format!("{prefix}-{}", Uuid::new_v4().simple()) |
| 215 | } |
| 216 | |
| 217 | fn init_tracing() { |
| 218 | TRACING.call_once(|| { |
| 219 | let filter = EnvFilter::try_from_default_env() |
| 220 | .unwrap_or_else(|_| EnvFilter::new("vllm_engine_core_client=debug")); |
| 221 | let _ = tracing_subscriber::fmt().with_test_writer().with_env_filter(filter).try_init(); |
| 222 | }); |
| 223 | } |
| 224 | |
| 225 | #[tokio::test(flavor = "multi_thread", worker_threads = 2)] |
| 226 | async fn generate_streams_outputs() { |
| 227 | init_tracing(); |
| 228 | let ipc = IpcNamespace::new().unwrap(); |
| 229 | let handshake_address = ipc.handshake_endpoint(); |
| 230 | let engine_id = b"engine-delta".to_vec(); |
| 231 | |
| 232 | let (shutdown_tx, engine_task) = spawn_mock_engine_task( |
| 233 | handshake_address.clone(), |
| 234 | engine_id.clone(), |
| 235 | |dealer, push| { |
| 236 | Box::pin(async move { |
| 237 | let add = recv_engine_message(dealer).await; |
| 238 | assert_eq!(add[0].as_ref(), &[0x00]); |
| 239 | let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap(); |
| 240 | assert_eq!(request.external_req_id.as_deref(), Some("req-delta")); |
| 241 | assert!(request.request_id.starts_with("req-delta-")); |
| 242 | assert_ne!(request.request_id, "req-delta"); |
| 243 | assert_eq!(request.client_index, 7); |
| 244 | assert_eq!(request.prompt_token_ids, Some(vec![11, 22])); |
| 245 | |
| 246 | send_outputs( |
| 247 | push, |
| 248 | RequestBatchOutputs { |
| 249 | outputs: vec![ |
| 250 | request_output_with_logprobs( |
| 251 | &request.request_id, |
| 252 | vec![1, 2], |
| 253 | None, |
| 254 | Some(logprobs_for_position(1, -0.3, 4, 9, -0.1)), |
| 255 | Some(prompt_logprobs()), |
| 256 | ), |
| 257 | request_output_with_logprobs( |
| 258 | &request.request_id, |
| 259 | vec![3], |
| 260 | Some(EngineCoreFinishReason::Length), |
| 261 | Some(logprobs_for_position(3, -0.4, 5, 10, -0.2)), |
| 262 | None, |
| 263 | ), |
| 264 | ], |
| 265 | finished_requests: Some(BTreeSet::from([request.request_id.clone()])), |
| 266 | ..Default::default() |
| 267 | } |
| 268 | .into(), |
| 269 | ) |
| 270 | .await; |
| 271 | }) |
| 272 | }, |
| 273 | ); |
| 274 | |
| 275 | let llm = connect_async_llm_with_ipc(handshake_address, 7, "test-model", &ipc).await; |
| 276 | let mut stream = llm.generate(sample_generate_request("req-delta", 3)).await.unwrap(); |
| 277 | let internal_id = stream.request_id().to_string(); |
| 278 | |
| 279 | let first = stream.next().await.unwrap().unwrap(); |
| 280 | assert_eq!(first.request_id, internal_id); |
| 281 | assert_eq!( |
| 282 | first.prompt_info, |
| 283 | Some(GeneratePromptInfo { |
| 284 | prompt_token_ids: vec![11, 22].into(), |
| 285 | prompt_logprobs: Some(prompt_logprobs()), |
| 286 | prompt_token_id_logprobs: None, |
| 287 | }) |
| 288 | ); |
| 289 | assert_eq!(first.token_ids, vec![1, 2]); |
| 290 | assert_eq!( |
| 291 | first.logprobs, |
| 292 | Some(logprobs_for_position(1, -0.3, 4, 9, -0.1)) |
| 293 | ); |
| 294 | assert_eq!(first.finish_reason, None); |
| 295 | |
| 296 | let second = stream.next().await.unwrap().unwrap(); |
| 297 | assert_eq!(second.prompt_info, None); |
| 298 | assert_eq!(second.token_ids, vec![3]); |
| 299 | assert_eq!( |
| 300 | second.logprobs, |
| 301 | Some(logprobs_for_position(3, -0.4, 5, 10, -0.2)) |
| 302 | ); |
| 303 | assert_eq!(second.finish_reason, Some(FinishReason::Length)); |
| 304 | assert!(stream.next().await.is_none()); |
| 305 | |
| 306 | let _ = shutdown_tx.send(()); |
| 307 | engine_task.await.unwrap(); |
| 308 | llm.shutdown().await.unwrap(); |
| 309 | } |
| 310 | |
| 311 | #[tokio::test(flavor = "multi_thread", worker_threads = 2)] |
| 312 | async fn collect_output_aggregates_raw_tokens_logprobs_and_terminal_metadata() { |
| 313 | init_tracing(); |
| 314 | let ipc = IpcNamespace::new().unwrap(); |
| 315 | let handshake_address = ipc.handshake_endpoint(); |
| 316 | let engine_id = b"engine-collect-output".to_vec(); |
| 317 | |
| 318 | let (shutdown_tx, engine_task) = spawn_mock_engine_task( |
| 319 | handshake_address.clone(), |
| 320 | engine_id.clone(), |
| 321 | |dealer, push| { |
| 322 | Box::pin(async move { |
| 323 | let add = recv_engine_message(dealer).await; |
| 324 | let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap(); |
| 325 | assert_eq!(request.external_req_id.as_deref(), Some("req-collect")); |
| 326 | assert!(request.request_id.starts_with("req-collect-")); |
| 327 | |
| 328 | send_outputs( |
| 329 | push, |
| 330 | RequestBatchOutputs { |
| 331 | outputs: vec![ |
| 332 | EngineCoreOutput { |
| 333 | prefill_stats: Some(PrefillStats { |
| 334 | num_prompt_tokens: 2, |
| 335 | num_cached_tokens: 1, |
| 336 | num_local_cached_tokens: 1, |
| 337 | ..Default::default() |
| 338 | }), |
| 339 | new_sampling_mask: Some(MaybeWireSamplingMask::Direct( |
| 340 | SamplingMask { |
| 341 | rows: vec![vec![1, 33, 99]], |
| 342 | }, |
| 343 | )), |
| 344 | ..request_output_with_logprobs( |
| 345 | &request.request_id, |
| 346 | vec![33], |
| 347 | None, |
| 348 | Some(logprobs_for_position(33, -0.1, 1, 99, -0.2)), |
| 349 | Some(prompt_logprobs()), |
| 350 | ) |
| 351 | }, |
| 352 | EngineCoreOutput { |
| 353 | new_sampling_mask: Some(MaybeWireSamplingMask::Direct( |
| 354 | SamplingMask { |
| 355 | rows: vec![vec![2, 44, 88]], |
| 356 | }, |
| 357 | )), |
| 358 | ..request_output_with_logprobs_and_kv( |
| 359 | &request.request_id, |
| 360 | vec![44], |
| 361 | Some(EngineCoreFinishReason::Stop), |
| 362 | Some(logprobs_for_position(44, -0.3, 1, 88, -0.4)), |
| 363 | None, |
| 364 | Some(serde_json::json!({"connector": "x"})), |
| 365 | None, |
| 366 | ) |
| 367 | }, |
| 368 | ], |
| 369 | ..Default::default() |
| 370 | } |
| 371 | .into(), |
| 372 | ) |
| 373 | .await; |
| 374 | }) |
| 375 | }, |
| 376 | ); |
| 377 | |
| 378 | let llm = connect_async_llm_with_ipc(handshake_address, 7, "test-model", &ipc).await; |
| 379 | let stream = llm.generate(sample_generate_request("req-collect", 4)).await.unwrap(); |
| 380 | let internal_id = stream.request_id().to_string(); |
| 381 | let collected = stream.collect_output().await.unwrap(); |
| 382 | |
| 383 | let _ = shutdown_tx.send(()); |
| 384 | engine_task.await.unwrap(); |
| 385 | |
| 386 | assert_eq!(collected.request_id, internal_id); |
| 387 | assert_eq!(collected.prompt_token_ids, vec![11, 22]); |
| 388 | assert_eq!(collected.token_ids, vec![33, 44]); |
| 389 | assert_eq!(collected.finish_reason, FinishReason::stop_eos()); |
| 390 | assert_eq!(collected.usage.cached_token_count, 1); |
| 391 | assert_eq!(collected.prompt_logprobs, Some(prompt_logprobs())); |
| 392 | assert_eq!( |
| 393 | collected.logprobs.as_ref().map(|lp| lp.positions.len()), |
| 394 | Some(2) |
| 395 | ); |
| 396 | assert_eq!( |
| 397 | collected.kv_transfer_params, |
| 398 | Some(serde_json::json!({"connector": "x"})) |
| 399 | ); |
| 400 | assert_eq!( |
| 401 | collected.sampling_mask, |
| 402 | Some(SamplingMask { |
| 403 | rows: vec![vec![1, 33, 99], vec![2, 44, 88]], |
| 404 | }) |
| 405 | ); |
| 406 | } |
| 407 | |
| 408 | #[tokio::test] |
| 409 | async fn collect_output_rejects_partial_sampling_mask() { |
| 410 | let output = vllm_llm::GenerateOutput { |
| 411 | request_id: "req-partial-mask".to_string(), |
| 412 | prompt_info: Some(GeneratePromptInfo { |
| 413 | prompt_token_ids: Arc::from([11_u32, 22]), |
| 414 | prompt_logprobs: None, |
| 415 | prompt_token_id_logprobs: None, |
| 416 | }), |
| 417 | token_ids: vec![33, 44], |
| 418 | logprobs: None, |
| 419 | finish_reason: Some(FinishReason::Length), |
| 420 | cached_token_count: 0, |
| 421 | kv_transfer_params: None, |
| 422 | ec_transfer_params: None, |
| 423 | sampling_mask: Some(SamplingMask { |
| 424 | rows: vec![vec![1, 33, 99]], |
| 425 | }), |
| 426 | spec_decode_metrics: None, |
| 427 | }; |
| 428 | |
| 429 | let error = futures::stream::iter([Ok(output)]).collect_output().await.unwrap_err(); |
| 430 | assert!(matches!( |
| 431 | error, |
| 432 | Error::SamplingMaskTokenCountMismatch { |
| 433 | token_count: 2, |
| 434 | row_count: 1, |
| 435 | .. |
| 436 | } |
| 437 | )); |
| 438 | } |
| 439 | |
| 440 | #[tokio::test(flavor = "multi_thread", worker_threads = 2)] |
| 441 | async fn generate_propagates_unexpected_close_errors() { |
| 442 | let ipc = IpcNamespace::new().unwrap(); |
| 443 | let handshake_address = ipc.handshake_endpoint(); |
| 444 | let engine_id = b"engine-close".to_vec(); |
| 445 | |
| 446 | let (shutdown_tx, engine_task) = spawn_mock_engine_task( |
| 447 | handshake_address.clone(), |
| 448 | engine_id.clone(), |
| 449 | |dealer, push| { |
| 450 | Box::pin(async move { |
| 451 | let add = recv_engine_message(dealer).await; |
| 452 | assert_eq!(add[0].as_ref(), &[0x00]); |
| 453 | let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap(); |
| 454 | |
| 455 | send_outputs( |
| 456 | push, |
| 457 | RequestBatchOutputs { |
| 458 | outputs: Vec::new(), |
| 459 | finished_requests: Some(BTreeSet::from([request.request_id])), |
| 460 | ..Default::default() |
| 461 | } |
| 462 | .into(), |
| 463 | ) |
| 464 | .await; |
| 465 | }) |
| 466 | }, |
| 467 | ); |
| 468 | |
| 469 | let llm = connect_async_llm_with_ipc(handshake_address, 0, "test-model", &ipc).await; |
| 470 | let mut stream = llm.generate(sample_generate_request("req-close", 1)).await.unwrap(); |
| 471 | let internal_id = stream.request_id().to_string(); |
| 472 | |
| 473 | let error = stream.next().await.unwrap().unwrap_err(); |
| 474 | assert!(matches!( |
| 475 | error, |
| 476 | Error::EngineCoreClient(vllm_engine_core_client::Error::RequestStreamClosed { |
| 477 | request_id |
| 478 | }) if request_id == internal_id |
| 479 | )); |
| 480 | assert!(stream.next().await.is_none()); |
| 481 | |
| 482 | let _ = shutdown_tx.send(()); |
| 483 | engine_task.await.unwrap(); |
| 484 | llm.shutdown().await.unwrap(); |
| 485 | } |
| 486 | |
| 487 | #[tokio::test(flavor = "multi_thread", worker_threads = 2)] |
| 488 | async fn dropping_a_live_generate_stream_triggers_abort() { |
| 489 | let ipc = IpcNamespace::new().unwrap(); |
| 490 | let handshake_address = ipc.handshake_endpoint(); |
| 491 | let engine_id = b"engine-drop".to_vec(); |
| 492 | |
| 493 | let (shutdown_tx, engine_task) = spawn_mock_engine_task( |
| 494 | handshake_address.clone(), |
| 495 | engine_id.clone(), |
| 496 | |dealer, push| { |
| 497 | Box::pin(async move { |
| 498 | let add = recv_engine_message(dealer).await; |
| 499 | assert_eq!(add[0].as_ref(), &[0x00]); |
| 500 | let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap(); |
| 501 | assert_eq!(request.external_req_id.as_deref(), Some("req-drop")); |
| 502 | assert!(request.request_id.starts_with("req-drop-")); |
| 503 | |
| 504 | send_outputs( |
| 505 | push, |
| 506 | RequestBatchOutputs { |
| 507 | outputs: vec![request_output(&request.request_id, vec![99], None)], |
| 508 | ..Default::default() |
| 509 | } |
| 510 | .into(), |
| 511 | ) |
| 512 | .await; |
| 513 | |
| 514 | let abort = |
| 515 | timeout(Duration::from_secs(1), recv_engine_message(dealer)).await.unwrap(); |
| 516 | assert_eq!(abort[0].as_ref(), &[0x01]); |
| 517 | let aborted_ids: Vec<String> = rmp_serde::from_slice(&abort[1]).unwrap(); |
| 518 | assert_eq!(aborted_ids, vec![request.request_id]); |
| 519 | }) |
| 520 | }, |
| 521 | ); |
| 522 | |
| 523 | let llm = connect_async_llm_with_ipc(handshake_address, 0, "test-model", &ipc).await; |
| 524 | let mut stream = llm.generate(sample_generate_request("req-drop", 4)).await.unwrap(); |
| 525 | |
| 526 | let output = stream.next().await.unwrap().unwrap(); |
| 527 | assert_eq!(output.token_ids, vec![99]); |
| 528 | drop(stream); |
| 529 | |
| 530 | let _ = shutdown_tx.send(()); |
| 531 | engine_task.await.unwrap(); |
| 532 | llm.shutdown().await.unwrap(); |
| 533 | } |
| 534 | |
| 535 | #[tokio::test(flavor = "multi_thread", worker_threads = 2)] |
| 536 | async fn duplicate_external_request_ids_are_randomized_before_reaching_engine_core_client() { |
| 537 | let ipc = IpcNamespace::new().unwrap(); |
| 538 | let handshake_address = ipc.handshake_endpoint(); |
| 539 | let engine_id = b"engine-dup".to_vec(); |
| 540 | |
| 541 | let (shutdown_tx, engine_task) = spawn_mock_engine_task( |
| 542 | handshake_address.clone(), |
| 543 | engine_id.clone(), |
| 544 | |dealer, push| { |
| 545 | Box::pin(async move { |
| 546 | let add_1 = recv_engine_message(dealer).await; |
| 547 | assert_eq!(add_1[0].as_ref(), &[0x00]); |
| 548 | let request_1: EngineCoreRequest = rmp_serde::from_slice(&add_1[1]).unwrap(); |
| 549 | assert_eq!(request_1.external_req_id.as_deref(), Some("req-dup")); |
| 550 | assert!(request_1.request_id.starts_with("req-dup-")); |
| 551 | |
| 552 | let add_2 = recv_engine_message(dealer).await; |
| 553 | assert_eq!(add_2[0].as_ref(), &[0x00]); |
| 554 | let request_2: EngineCoreRequest = rmp_serde::from_slice(&add_2[1]).unwrap(); |
| 555 | assert_eq!(request_2.external_req_id.as_deref(), Some("req-dup")); |
| 556 | assert!(request_2.request_id.starts_with("req-dup-")); |
| 557 | assert_ne!(request_1.request_id, request_2.request_id); |
| 558 | |
| 559 | send_outputs( |
| 560 | push, |
| 561 | RequestBatchOutputs { |
| 562 | outputs: vec![request_output( |
| 563 | &request_1.request_id, |
| 564 | vec![], |
| 565 | Some(EngineCoreFinishReason::Length), |
| 566 | )], |
| 567 | finished_requests: Some(BTreeSet::from([request_1.request_id.clone()])), |
| 568 | ..Default::default() |
| 569 | } |
| 570 | .into(), |
| 571 | ) |
| 572 | .await; |
| 573 | |
| 574 | send_outputs( |
| 575 | push, |
| 576 | RequestBatchOutputs { |
| 577 | outputs: vec![request_output( |
| 578 | &request_2.request_id, |
| 579 | vec![], |
| 580 | Some(EngineCoreFinishReason::Length), |
| 581 | )], |
| 582 | finished_requests: Some(BTreeSet::from([request_2.request_id])), |
| 583 | ..Default::default() |
| 584 | } |
| 585 | .into(), |
| 586 | ) |
| 587 | .await; |
| 588 | }) |
| 589 | }, |
| 590 | ); |
| 591 | |
| 592 | let llm = connect_async_llm_with_ipc(handshake_address, 0, "test-model", &ipc).await; |
| 593 | let stream_1 = llm.generate(sample_generate_request("req-dup", 1)).await.unwrap(); |
| 594 | let stream_2 = llm.generate(sample_generate_request("req-dup", 1)).await.unwrap(); |
| 595 | let internal_id_1 = stream_1.request_id().to_string(); |
| 596 | let internal_id_2 = stream_2.request_id().to_string(); |
| 597 | let collected_1 = stream_1.collect_output().await.unwrap(); |
| 598 | let collected_2 = stream_2.collect_output().await.unwrap(); |
| 599 | assert_eq!(collected_1.request_id, internal_id_1); |
| 600 | assert_eq!(collected_2.request_id, internal_id_2); |
| 601 | let _ = shutdown_tx.send(()); |
| 602 | engine_task.await.unwrap(); |
| 603 | llm.shutdown().await.unwrap(); |
| 604 | } |
| 605 | |
| 606 | #[tokio::test(flavor = "multi_thread", worker_threads = 2)] |
| 607 | async fn abort_resolves_external_request_id_to_internal_before_reaching_engine() { |
| 608 | let ipc = IpcNamespace::new().unwrap(); |
| 609 | let handshake_address = ipc.handshake_endpoint(); |
| 610 | let engine_id = b"engine-abort".to_vec(); |
| 611 | |
| 612 | let (shutdown_tx, engine_task) = spawn_mock_engine_task( |
| 613 | handshake_address.clone(), |
| 614 | engine_id.clone(), |
| 615 | |dealer, push| { |
| 616 | Box::pin(async move { |
| 617 | let add = recv_engine_message(dealer).await; |
| 618 | assert_eq!(add[0].as_ref(), &[0x00]); |
| 619 | let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap(); |
| 620 | assert_eq!(request.external_req_id.as_deref(), Some("req-abort")); |
| 621 | assert!(request.request_id.starts_with("req-abort-")); |
| 622 | assert_ne!(request.request_id, "req-abort"); |
| 623 | |
| 624 | send_outputs( |
| 625 | push, |
| 626 | RequestBatchOutputs { |
| 627 | outputs: vec![request_output(&request.request_id, vec![7], None)], |
| 628 | ..Default::default() |
| 629 | } |
| 630 | .into(), |
| 631 | ) |
| 632 | .await; |
| 633 | |
| 634 | // The abort frame must carry the internal engine id, not the |
| 635 | // external "req-abort" id the caller aborted by. |
| 636 | let abort = |
| 637 | timeout(Duration::from_secs(1), recv_engine_message(dealer)).await.unwrap(); |
| 638 | assert_eq!(abort[0].as_ref(), &[0x01]); |
| 639 | let aborted_ids: Vec<String> = rmp_serde::from_slice(&abort[1]).unwrap(); |
| 640 | assert_eq!(aborted_ids, vec![request.request_id]); |
| 641 | }) |
| 642 | }, |
| 643 | ); |
| 644 | |
| 645 | let llm = connect_async_llm_with_ipc(handshake_address, 0, "test-model", &ipc).await; |
| 646 | let mut stream = llm.generate(sample_generate_request("req-abort", 4)).await.unwrap(); |
| 647 | let internal_id = stream.request_id().to_string(); |
| 648 | assert_ne!(internal_id, "req-abort"); |
| 649 | |
| 650 | assert_eq!(stream.next().await.unwrap().unwrap().token_ids, vec![7]); |
| 651 | |
| 652 | // Abort by the external id; engine-core only knows the internal id. |
| 653 | llm.abort(&["req-abort".to_string()]).await.unwrap(); |
| 654 | |
| 655 | // The consumer stream is finalized locally with a clean abort terminal |
| 656 | // rather than hanging or surfacing as RequestStreamClosed. The engine sends |
| 657 | // no final output for a client abort, so this output is synthesized. |
| 658 | let terminal = stream.next().await.unwrap().unwrap(); |
| 659 | assert_eq!(terminal.finish_reason, Some(FinishReason::Abort)); |
| 660 | assert!(terminal.token_ids.is_empty()); |
| 661 | assert!(stream.next().await.is_none()); |
| 662 | |
| 663 | let _ = shutdown_tx.send(()); |
| 664 | engine_task.await.unwrap(); |
| 665 | drop(stream); |
| 666 | llm.shutdown().await.unwrap(); |
| 667 | } |
| 668 | |
| 669 | #[tokio::test(flavor = "multi_thread", worker_threads = 2)] |
| 670 | async fn abort_by_external_id_aborts_all_internal_requests() { |
| 671 | let ipc = IpcNamespace::new().unwrap(); |
| 672 | let handshake_address = ipc.handshake_endpoint(); |
| 673 | let engine_id = b"engine-abort-many".to_vec(); |
| 674 | |
| 675 | let (shutdown_tx, engine_task) = spawn_mock_engine_task( |
| 676 | handshake_address.clone(), |
| 677 | engine_id.clone(), |
| 678 | |dealer, push| { |
| 679 | Box::pin(async move { |
| 680 | let add_1 = recv_engine_message(dealer).await; |
| 681 | assert_eq!(add_1[0].as_ref(), &[0x00]); |
| 682 | let request_1: EngineCoreRequest = rmp_serde::from_slice(&add_1[1]).unwrap(); |
| 683 | |
| 684 | let add_2 = recv_engine_message(dealer).await; |
| 685 | assert_eq!(add_2[0].as_ref(), &[0x00]); |
| 686 | let request_2: EngineCoreRequest = rmp_serde::from_slice(&add_2[1]).unwrap(); |
| 687 | |
| 688 | assert_eq!(request_1.external_req_id.as_deref(), Some("req-dup-abort")); |
| 689 | assert_eq!(request_2.external_req_id.as_deref(), Some("req-dup-abort")); |
| 690 | assert_ne!(request_1.request_id, request_2.request_id); |
| 691 | |
| 692 | send_outputs( |
| 693 | push, |
| 694 | RequestBatchOutputs { |
| 695 | outputs: vec![ |
| 696 | request_output(&request_1.request_id, vec![7], None), |
| 697 | request_output(&request_2.request_id, vec![8], None), |
| 698 | ], |
| 699 | ..Default::default() |
| 700 | } |
| 701 | .into(), |
| 702 | ) |
| 703 | .await; |
| 704 | |
| 705 | // A single abort by the shared external id must abort both |
| 706 | // internal engine ids it expanded into. |
| 707 | let abort = |
| 708 | timeout(Duration::from_secs(1), recv_engine_message(dealer)).await.unwrap(); |
| 709 | assert_eq!(abort[0].as_ref(), &[0x01]); |
| 710 | let mut aborted_ids: Vec<String> = rmp_serde::from_slice(&abort[1]).unwrap(); |
| 711 | aborted_ids.sort(); |
| 712 | let mut expected = vec![request_1.request_id, request_2.request_id]; |
| 713 | expected.sort(); |
| 714 | assert_eq!(aborted_ids, expected); |
| 715 | }) |
| 716 | }, |
| 717 | ); |
| 718 | |
| 719 | let llm = connect_async_llm_with_ipc(handshake_address, 0, "test-model", &ipc).await; |
| 720 | let mut stream_1 = llm.generate(sample_generate_request("req-dup-abort", 4)).await.unwrap(); |
| 721 | let mut stream_2 = llm.generate(sample_generate_request("req-dup-abort", 4)).await.unwrap(); |
| 722 | assert_ne!(stream_1.request_id(), stream_2.request_id()); |
| 723 | |
| 724 | assert_eq!(stream_1.next().await.unwrap().unwrap().token_ids, vec![7]); |
| 725 | assert_eq!(stream_2.next().await.unwrap().unwrap().token_ids, vec![8]); |
| 726 | |
| 727 | llm.abort(&["req-dup-abort".to_string()]).await.unwrap(); |
| 728 | |
| 729 | // Both internal requests the external id expanded into are finalized with a |
| 730 | // clean abort terminal. |
| 731 | for stream in [&mut stream_1, &mut stream_2] { |
| 732 | let terminal = stream.next().await.unwrap().unwrap(); |
| 733 | assert_eq!(terminal.finish_reason, Some(FinishReason::Abort)); |
| 734 | assert!(stream.next().await.is_none()); |
| 735 | } |
| 736 | |
| 737 | let _ = shutdown_tx.send(()); |
| 738 | engine_task.await.unwrap(); |
| 739 | drop(stream_1); |
| 740 | drop(stream_2); |
| 741 | llm.shutdown().await.unwrap(); |
| 742 | } |
| 743 | |
| 744 | #[tokio::test(flavor = "multi_thread", worker_threads = 2)] |
| 745 | async fn generate_records_request_metrics_in_prometheus_output() { |
| 746 | let ipc = IpcNamespace::new().unwrap(); |
| 747 | let handshake_address = ipc.handshake_endpoint(); |
| 748 | let engine_id = EngineId::from_engine_index(4); |
| 749 | let model_name = request_metrics_model_name("metrics-model"); |
| 750 | |
| 751 | let (shutdown_tx, engine_task) = spawn_mock_engine_task( |
| 752 | handshake_address.clone(), |
| 753 | engine_id.clone(), |
| 754 | |dealer, push| { |
| 755 | Box::pin(async move { |
| 756 | let add = recv_engine_message(dealer).await; |
| 757 | assert_eq!(add[0].as_ref(), &[0x00]); |
| 758 | let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap(); |
| 759 | |
| 760 | send_outputs( |
| 761 | push, |
| 762 | RequestBatchOutputs { |
| 763 | engine_index: 4, |
| 764 | timestamp: 10.0, |
| 765 | outputs: vec![EngineCoreOutput { |
| 766 | prefill_stats: Some(PrefillStats { |
| 767 | num_prompt_tokens: 2, |
| 768 | num_computed_tokens: 2, |
| 769 | ..Default::default() |
| 770 | }), |
| 771 | ..request_output_with_events( |
| 772 | &request.request_id, |
| 773 | vec![1], |
| 774 | None, |
| 775 | Some(vec![ |
| 776 | EngineCoreEvent { |
| 777 | r#type: EngineCoreEventType::Queued, |
| 778 | timestamp: 8.0, |
| 779 | }, |
| 780 | EngineCoreEvent { |
| 781 | r#type: EngineCoreEventType::Scheduled, |
| 782 | timestamp: 9.0, |
| 783 | }, |
| 784 | ]), |
| 785 | ) |
| 786 | }], |
| 787 | ..Default::default() |
| 788 | } |
| 789 | .into(), |
| 790 | ) |
| 791 | .await; |
| 792 | |
| 793 | send_outputs( |
| 794 | push, |
| 795 | RequestBatchOutputs { |
| 796 | engine_index: 4, |
| 797 | timestamp: 11.5, |
| 798 | outputs: vec![request_output_with_events( |
| 799 | &request.request_id, |
| 800 | vec![2, 3], |
| 801 | Some(EngineCoreFinishReason::Length), |
| 802 | Some(vec![EngineCoreEvent { |
| 803 | r#type: EngineCoreEventType::Preempted, |
| 804 | timestamp: 10.5, |
| 805 | }]), |
| 806 | )], |
| 807 | finished_requests: Some(BTreeSet::from([request.request_id])), |
| 808 | ..Default::default() |
| 809 | } |
| 810 | .into(), |
| 811 | ) |
| 812 | .await; |
| 813 | }) |
| 814 | }, |
| 815 | ); |
| 816 | |
| 817 | let llm = connect_async_llm_with_ipc(handshake_address, 0, &model_name, &ipc).await; |
| 818 | let mut request = sample_generate_request("req-metrics", 8); |
| 819 | request.arrival_time = None; |
| 820 | let mut stream = llm.generate(request).await.unwrap(); |
| 821 | |
| 822 | assert_eq!(stream.next().await.unwrap().unwrap().token_ids, vec![1]); |
| 823 | let final_output = stream.next().await.unwrap().unwrap(); |
| 824 | assert_eq!(final_output.token_ids, vec![2, 3]); |
| 825 | assert_eq!(final_output.finish_reason, Some(FinishReason::Length)); |
| 826 | assert!(stream.next().await.is_none()); |
| 827 | |
| 828 | let rendered = METRICS.render().unwrap(); |
| 829 | assert!(rendered.contains(&format!( |
| 830 | "vllm:request_success_total{{model_name=\"{model_name}\",engine=\"4\",finished_reason=\"length\"}} 1" |
| 831 | ))); |
| 832 | assert!(rendered.contains(&format!( |
| 833 | "vllm:prompt_tokens_total{{model_name=\"{model_name}\",engine=\"4\"}} 2" |
| 834 | ))); |
| 835 | assert!(rendered.contains(&format!( |
| 836 | "vllm:prompt_tokens_by_source_total{{model_name=\"{model_name}\",engine=\"4\",source=\"local_compute\"}} 2" |
| 837 | ))); |
| 838 | assert!(rendered.contains(&format!( |
| 839 | "vllm:prompt_tokens_by_source_total{{model_name=\"{model_name}\",engine=\"4\",source=\"local_cache_hit\"}} 0" |
| 840 | ))); |
| 841 | assert!(rendered.contains(&format!( |
| 842 | "vllm:prompt_tokens_by_source_total{{model_name=\"{model_name}\",engine=\"4\",source=\"external_kv_transfer\"}} 0" |
| 843 | ))); |
| 844 | assert!(rendered.contains(&format!( |
| 845 | "vllm:prompt_tokens_cached_total{{model_name=\"{model_name}\",engine=\"4\"}} 0" |
| 846 | ))); |
| 847 | assert!(rendered.contains(&format!( |
| 848 | "vllm:generation_tokens_total{{model_name=\"{model_name}\",engine=\"4\"}} 3" |
| 849 | ))); |
| 850 | assert!(rendered.contains(&format!( |
| 851 | "vllm:num_preemptions_total{{model_name=\"{model_name}\",engine=\"4\"}} 1" |
| 852 | ))); |
| 853 | assert!(rendered.contains(&format!( |
| 854 | "vllm:request_num_preemptions_sum{{model_name=\"{model_name}\",engine=\"4\"}} 1.0" |
| 855 | ))); |
| 856 | assert!(rendered.contains(&format!( |
| 857 | "vllm:request_num_preemptions_count{{model_name=\"{model_name}\",engine=\"4\"}} 1" |
| 858 | ))); |
| 859 | assert!(rendered.contains(&format!( |
| 860 | "vllm:time_to_first_token_seconds_count{{model_name=\"{model_name}\",engine=\"4\"}} 1" |
| 861 | ))); |
| 862 | assert!(rendered.contains(&format!( |
| 863 | "vllm:inter_token_latency_seconds_count{{model_name=\"{model_name}\",engine=\"4\"}} 1" |
| 864 | ))); |
| 865 | assert!(rendered.contains(&format!( |
| 866 | "vllm:e2e_request_latency_seconds_count{{model_name=\"{model_name}\",engine=\"4\"}} 1" |
| 867 | ))); |
| 868 | assert!(rendered.contains(&format!( |
| 869 | "vllm:request_prompt_tokens_count{{model_name=\"{model_name}\",engine=\"4\"}} 1" |
| 870 | ))); |
| 871 | assert!(rendered.contains(&format!( |
| 872 | "vllm:request_generation_tokens_count{{model_name=\"{model_name}\",engine=\"4\"}} 1" |
| 873 | ))); |
| 874 | assert!(rendered.contains(&format!( |
| 875 | "vllm:request_prefill_kv_computed_tokens_count{{model_name=\"{model_name}\",engine=\"4\"}} 1" |
| 876 | ))); |
| 877 | assert!(rendered.contains(&format!( |
| 878 | "vllm:request_queue_time_seconds_sum{{model_name=\"{model_name}\",engine=\"4\"}} 1.0" |
| 879 | ))); |
| 880 | |
| 881 | let _ = shutdown_tx.send(()); |
| 882 | engine_task.await.unwrap(); |
| 883 | llm.shutdown().await.unwrap(); |
| 884 | } |
| 885 | |
| 886 | #[tokio::test(flavor = "multi_thread", worker_threads = 2)] |
| 887 | async fn generate_without_engine_stats_exports_only_frontend_observed_metrics() { |
| 888 | let ipc = IpcNamespace::new().unwrap(); |
| 889 | let handshake_address = ipc.handshake_endpoint(); |
| 890 | let engine_id = EngineId::from_engine_index(6); |
| 891 | let model_name = request_metrics_model_name("metrics-no-log-stats-model"); |
| 892 | let lora_name = request_metrics_model_name("metrics-no-log-stats-lora"); |
| 893 | let (finish_tx, finish_rx) = oneshot::channel::<()>(); |
| 894 | |
| 895 | let (shutdown_tx, engine_task) = spawn_mock_engine_task_with_ready( |
| 896 | handshake_address.clone(), |
| 897 | engine_id, |
| 898 | EngineCoreReadyResponse { |
| 899 | supports_lora: true, |
| 900 | max_loras: 1, |
| 901 | ..default_ready_response() |
| 902 | }, |
| 903 | move |dealer, push| { |
| 904 | Box::pin(async move { |
| 905 | let add = recv_engine_message(dealer).await; |
| 906 | assert_eq!(add[0].as_ref(), &[0x00]); |
| 907 | let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap(); |
| 908 | |
| 909 | // Engines started with `--disable-log-stats` attach neither |
| 910 | // lifecycle events nor scheduler stats to their outputs. |
| 911 | for token in [1, 2] { |
| 912 | send_outputs( |
| 913 | push, |
| 914 | RequestBatchOutputs { |
| 915 | engine_index: 6, |
| 916 | timestamp: 10.0 + f64::from(token), |
| 917 | outputs: vec![request_output(&request.request_id, vec![token], None)], |
| 918 | ..Default::default() |
| 919 | } |
| 920 | .into(), |
| 921 | ) |
| 922 | .await; |
| 923 | } |
| 924 | |
| 925 | finish_rx.await.unwrap(); |
| 926 | send_outputs( |
| 927 | push, |
| 928 | RequestBatchOutputs { |
| 929 | engine_index: 6, |
| 930 | timestamp: 13.0, |
| 931 | outputs: vec![request_output( |
| 932 | &request.request_id, |
| 933 | vec![3], |
| 934 | Some(EngineCoreFinishReason::Length), |
| 935 | )], |
| 936 | finished_requests: Some(BTreeSet::from([request.request_id])), |
| 937 | ..Default::default() |
| 938 | } |
| 939 | .into(), |
| 940 | ) |
| 941 | .await; |
| 942 | }) |
| 943 | }, |
| 944 | ); |
| 945 | |
| 946 | let client = EngineCoreClient::connect( |
| 947 | EngineCoreClientConfig::new_single(handshake_address) |
| 948 | .with_model_name(&model_name) |
| 949 | .with_engine_stats_enabled(false) |
| 950 | .with_local_input_output_addresses( |
| 951 | Some(ipc.input_endpoint()), |
| 952 | Some(ipc.output_endpoint()), |
| 953 | ), |
| 954 | ) |
| 955 | .await |
| 956 | .unwrap(); |
| 957 | let llm = Llm::new(client); |
| 958 | let mut request = sample_generate_request("req-no-log-stats", 8); |
| 959 | request.arrival_time = None; |
| 960 | request.lora_request = Some(LoraRequest { |
| 961 | lora_name: lora_name.clone(), |
| 962 | lora_int_id: 1, |
| 963 | lora_path: "/unused".to_string(), |
| 964 | base_model_name: None, |
| 965 | tensorizer_config_dict: None, |
| 966 | load_inplace: false, |
| 967 | is_3d_lora_weight: false, |
| 968 | }); |
| 969 | let mut stream = llm.generate(request).await.unwrap(); |
| 970 | |
| 971 | assert_eq!(stream.next().await.unwrap().unwrap().token_ids, vec![1]); |
| 972 | // The dispatcher handles the second batch only after finishing the first |
| 973 | // one, including its LoRA info update. |
| 974 | assert_eq!(stream.next().await.unwrap().unwrap().token_ids, vec![2]); |
| 975 | assert!(!METRICS.render().unwrap().contains(&lora_name)); |
| 976 | |
| 977 | finish_tx.send(()).unwrap(); |
| 978 | let final_output = stream.next().await.unwrap().unwrap(); |
| 979 | assert_eq!(final_output.finish_reason, Some(FinishReason::Length)); |
| 980 | assert!(stream.next().await.is_none()); |
| 981 | |
| 982 | let rendered = METRICS.render().unwrap(); |
| 983 | for metric in [ |
| 984 | "vllm:time_to_first_token_seconds_count", |
| 985 | "vllm:e2e_request_latency_seconds_count", |
| 986 | "vllm:request_decode_time_seconds_count", |
| 987 | ] { |
| 988 | assert!( |
| 989 | rendered.contains(&format!( |
| 990 | "{metric}{{model_name=\"{model_name}\",engine=\"6\"}} 1" |
| 991 | )), |
| 992 | "{metric} should be exported" |
| 993 | ); |
| 994 | } |
| 995 | for metric in [ |
| 996 | "vllm:num_requests_running", |
| 997 | "vllm:kv_cache_usage_perc", |
| 998 | "vllm:num_preemptions_total", |
| 999 | "vllm:request_num_preemptions_count", |
| 1000 | "vllm:request_queue_time_seconds_count", |
| 1001 | "vllm:request_prefill_time_seconds_count", |
| 1002 | "vllm:request_inference_time_seconds_count", |
| 1003 | ] { |
| 1004 | assert!( |
| 1005 | !rendered.contains(&format!("{metric}{{model_name=\"{model_name}\"")), |
| 1006 | "{metric} should not be exported" |
| 1007 | ); |
| 1008 | } |
| 1009 | |
| 1010 | let _ = shutdown_tx.send(()); |
| 1011 | engine_task.await.unwrap(); |
| 1012 | llm.shutdown().await.unwrap(); |
| 1013 | } |
| 1014 | |
| 1015 | #[tokio::test(flavor = "multi_thread", worker_threads = 2)] |
| 1016 | async fn dropping_stream_records_abort_terminal_request_metrics() { |
| 1017 | let ipc = IpcNamespace::new().unwrap(); |
| 1018 | let handshake_address = ipc.handshake_endpoint(); |
| 1019 | let engine_id = EngineId::from_engine_index(5); |
| 1020 | let model_name = request_metrics_model_name("metrics-drop-model"); |
| 1021 | |
| 1022 | let (shutdown_tx, engine_task) = spawn_mock_engine_task( |
| 1023 | handshake_address.clone(), |
| 1024 | engine_id.clone(), |
| 1025 | |dealer, push| { |
| 1026 | Box::pin(async move { |
| 1027 | let add = recv_engine_message(dealer).await; |
| 1028 | assert_eq!(add[0].as_ref(), &[0x00]); |
| 1029 | let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap(); |
| 1030 | assert_eq!(request.external_req_id.as_deref(), Some("req-metrics-drop")); |
| 1031 | assert!(request.request_id.starts_with("req-metrics-drop-")); |
| 1032 | |
| 1033 | send_outputs( |
| 1034 | push, |
| 1035 | RequestBatchOutputs { |
| 1036 | engine_index: 5, |
| 1037 | timestamp: 10.0, |
| 1038 | outputs: vec![request_output_with_events( |
| 1039 | &request.request_id, |
| 1040 | vec![99], |
| 1041 | None, |
| 1042 | Some(vec![ |
| 1043 | EngineCoreEvent { |
| 1044 | r#type: EngineCoreEventType::Queued, |
| 1045 | timestamp: 8.0, |
| 1046 | }, |
| 1047 | EngineCoreEvent { |
| 1048 | r#type: EngineCoreEventType::Scheduled, |
| 1049 | timestamp: 9.0, |
| 1050 | }, |
| 1051 | ]), |
| 1052 | )], |
| 1053 | ..Default::default() |
| 1054 | } |
| 1055 | .into(), |
| 1056 | ) |
| 1057 | .await; |
| 1058 | |
| 1059 | let abort = |
| 1060 | timeout(Duration::from_secs(1), recv_engine_message(dealer)).await.unwrap(); |
| 1061 | assert_eq!(abort[0].as_ref(), &[0x01]); |
| 1062 | let aborted_ids: Vec<String> = rmp_serde::from_slice(&abort[1]).unwrap(); |
| 1063 | assert_eq!(aborted_ids, vec![request.request_id]); |
| 1064 | }) |
| 1065 | }, |
| 1066 | ); |
| 1067 | |
| 1068 | let llm = connect_async_llm_with_ipc(handshake_address, 0, &model_name, &ipc).await; |
| 1069 | let mut request = sample_generate_request("req-metrics-drop", 8); |
| 1070 | request.arrival_time = None; |
| 1071 | let mut stream = llm.generate(request).await.unwrap(); |
| 1072 | assert_eq!(stream.next().await.unwrap().unwrap().token_ids, vec![99]); |
| 1073 | drop(stream); |
| 1074 | |
| 1075 | let _ = shutdown_tx.send(()); |
| 1076 | engine_task.await.unwrap(); |
| 1077 | let rendered = METRICS.render().unwrap(); |
| 1078 | assert!(rendered.contains(&format!( |
| 1079 | "vllm:request_success_total{{model_name=\"{model_name}\",engine=\"5\",finished_reason=\"abort\"}} 1" |
| 1080 | ))); |
| 1081 | assert!(rendered.contains(&format!( |
| 1082 | "vllm:e2e_request_latency_seconds_count{{model_name=\"{model_name}\",engine=\"5\"}} 1" |
| 1083 | ))); |
| 1084 | assert!(rendered.contains(&format!( |
| 1085 | "vllm:request_num_preemptions_sum{{model_name=\"{model_name}\",engine=\"5\"}} 0.0" |
| 1086 | ))); |
| 1087 | assert!(rendered.contains(&format!( |
| 1088 | "vllm:request_num_preemptions_count{{model_name=\"{model_name}\",engine=\"5\"}} 1" |
| 1089 | ))); |
| 1090 | |
| 1091 | llm.shutdown().await.unwrap(); |
| 1092 | } |