From b9b5c54053e708328fcedf0c9e12d94e63f38a5f Mon Sep 17 00:00:00 2001 From: Ivan Zatevakhin Date: Thu, 27 Aug 2026 01:58:51 +0100 Subject: [PATCH] feat: add token, transcription, and embedding APIs --- vllm-cpp/src/abi.rs | 28 +- vllm-cpp/src/engine.rs | 1139 ++++++++++++++++++++++++++++++++++-- vllm-cpp/src/error.rs | 8 + vllm-cpp/src/lib.rs | 28 +- vllm-cpp/tests/qwen3.rs | 187 +++++- vllm-cpp/tests/safe_api.rs | 36 +- 6 files changed, 1375 insertions(+), 51 deletions(-) diff --git a/vllm-cpp/src/abi.rs b/vllm-cpp/src/abi.rs index d8cf2ec..9bd826f 100644 --- a/vllm-cpp/src/abi.rs +++ b/vllm-cpp/src/abi.rs @@ -26,6 +26,12 @@ impl Compatibility { unsafe { ffi::vllm_sampling_params_default() } } + pub(crate) fn transcription_params_default(&self) -> ffi::vllm_transcription_params { + // SAFETY: the engine retained this token after exact ABI equality was + // established, so this versioned struct may be returned by value. + unsafe { ffi::vllm_transcription_params_default() } + } + fn from_actual(actual: i32) -> Result { let expected = ffi::VLLM_ABI_VERSION as i32; if actual != expected { @@ -54,6 +60,14 @@ impl Compatibility { ) -> ffi::vllm_sampling_params { default() } + + #[cfg(test)] + fn transcription_params_default_with( + &self, + default: impl FnOnce() -> ffi::vllm_transcription_params, + ) -> ffi::vllm_transcription_params { + default() + } } #[cfg(test)] @@ -82,7 +96,7 @@ mod tests { } #[test] - fn compatibility_precedes_both_by_value_defaults() { + fn compatibility_precedes_all_by_value_defaults() { let calls = RefCell::new(Vec::new()); let compatibility = Compatibility::check_with(|| { calls.borrow_mut().push("abi"); @@ -100,10 +114,20 @@ mod tests { // SAFETY: every field in this generated C struct permits zero. unsafe { std::mem::zeroed() } }); + compatibility.transcription_params_default_with(|| { + calls.borrow_mut().push("transcription_default"); + // SAFETY: every field in this generated C struct permits zero. + unsafe { std::mem::zeroed() } + }); assert_eq!( *calls.borrow(), - ["abi", "model_default", "sampling_default"] + [ + "abi", + "model_default", + "sampling_default", + "transcription_default" + ] ); } } diff --git a/vllm-cpp/src/engine.rs b/vllm-cpp/src/engine.rs index cafde32..ca3e3df 100644 --- a/vllm-cpp/src/engine.rs +++ b/vllm-cpp/src/engine.rs @@ -1,6 +1,6 @@ use std::ffi::{CStr, CString}; use std::marker::PhantomData; -use std::mem::MaybeUninit; +use std::mem::{align_of, size_of, MaybeUninit}; use std::os::raw::c_char; use std::path::{Path, PathBuf}; use std::ptr::{self, NonNull}; @@ -41,24 +41,25 @@ struct LoadedEngine { pub(crate) type EngineInner = OwnedEngine; -/// A thread-local RAII owner for a native transcription-task engine. +/// A thread-local RAII owner for blocking native transcription. /// /// ABI 17 has no task query, so loading cannot prove that a checkpoint supports -/// transcription. Native task selection and future wrong-task diagnostics remain -/// authoritative. This owner intentionally exposes no operation yet and is -/// neither `Send` nor `Sync`. +/// transcription. Native task selection and wrong-task [`Error::InvalidArgument`] +/// diagnostics remain authoritative. This owner is neither `Send` nor `Sync` and +/// operations require exclusive access. pub struct TranscriptionEngine { - _inner: OwnedEngine, + inner: OwnedEngine, } -/// A thread-local RAII owner for a native embedding-task engine. +/// A thread-local RAII owner for blocking native embeddings. /// /// ABI 17 has no task query, so loading cannot prove that a checkpoint supports -/// embeddings. Native task selection and future wrong-task diagnostics remain -/// authoritative. This owner intentionally exposes no operation yet and is -/// neither `Send` nor `Sync`. +/// embeddings. Native task selection and wrong-task [`Error::InvalidArgument`] +/// diagnostics remain authoritative. Native embedding batches are serialized; +/// this Rust owner is neither `Send` nor `Sync` and operations require exclusive +/// access. pub struct EmbeddingEngine { - _inner: OwnedEngine, + inner: OwnedEngine, } impl std::fmt::Debug for Engine { @@ -287,6 +288,86 @@ pub struct Completion { pub completion_tokens: u32, } +/// A Rust-owned pre-tokenized completion result. +/// +/// [`Self::token_ids`] contains only IDs that fit the caller's reporting buffer. +/// [`Self::truncated`] reports whether native generation produced additional IDs. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct TokenCompletion { + pub token_ids: Vec, + pub completion: Option, + pub truncated: bool, +} + +/// Borrowed audio for one blocking transcription call. +/// +/// The input is borrowed only until [`TranscriptionEngine::transcribe`] returns. +/// Rust does not inspect, decode, or resample either input form. WAV decoding and +/// the requirement for 16 kHz mono audio remain native behavior. +#[derive(Debug, Clone, Copy)] +pub enum TranscriptionInput<'a> { + WavFile(&'a Path), + Pcm { + samples: &'a [f32], + sample_rate: u32, + }, +} + +/// A Rust-owned transcription result. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Transcription { + pub text: Option, + pub token_ids: Vec, +} + +/// A Rust-owned row-major embedding batch. +/// +/// Rows preserve input order. The flattened values and all row views remain +/// independent of native result storage. +#[derive(Debug, Clone, PartialEq)] +pub struct EmbeddingResult { + values: Vec, + dimension: usize, + prompt_tokens: u32, +} + +impl EmbeddingResult { + /// Returns all row-major values. + #[must_use] + pub fn values(&self) -> &[f32] { + &self.values + } + + /// Returns the number of embedding rows. + #[must_use] + pub fn n_embeddings(&self) -> usize { + self.values.len() / self.dimension + } + + /// Returns the number of values in each row. + #[must_use] + pub fn dimension(&self) -> usize { + self.dimension + } + + /// Returns the total number of native input tokens for the batch. + #[must_use] + pub fn prompt_tokens(&self) -> u32 { + self.prompt_tokens + } + + /// Returns one row, or `None` when `index` is out of bounds. + #[must_use] + pub fn row(&self, index: usize) -> Option<&[f32]> { + self.rows().nth(index) + } + + /// Iterates over rows in input order. + pub fn rows(&self) -> std::slice::ChunksExact<'_, f32> { + self.values.chunks_exact(self.dimension) + } +} + impl Engine { /// Starts configuring an engine for a model directory or GGUF file. pub fn builder(model_path: impl Into) -> EngineBuilder { @@ -322,11 +403,53 @@ impl Engine { } // SAFETY: VLLM_OK initializes every completion field. let raw = unsafe { raw.assume_init() }; - let guard = CompletionGuard(raw); + let guard = NativeResultGuard::new(raw, ffi::vllm_completion_free); if let Some(error) = params.logits_processor_error() { return Err(error); } - completion_from_raw(&guard.0) + completion_from_raw(guard.raw()) + } + + /// Runs one blocking completion from a pre-tokenized prompt. + /// + /// `prompt_tokens` is borrowed only for this call. `max_output_tokens` limits + /// how many generated IDs are reported; it does not limit generation. Native + /// completion metadata is always requested to determine truncation accurately, + /// even when `include_completion` omits it from the Rust result. All returned + /// data is Rust-owned. + pub fn complete_tokens( + &self, + prompt_tokens: &[i32], + params: &SamplingParams, + max_output_tokens: usize, + include_completion: bool, + ) -> Result { + validate_token_input(prompt_tokens.len(), max_output_tokens)?; + let params = params.marshal(&self.inner.compatibility)?; + complete_tokens_with( + prompt_tokens, + params.raw(), + max_output_tokens, + include_completion, + |prompt, n_prompt, params, output, capacity, written, completion| { + // SAFETY: all borrowed and output storage remains live for this + // blocking call, and the native result is initialized on success. + unsafe { + ffi::vllm_complete_tokens( + self.inner.raw.as_ptr(), + prompt, + n_prompt, + params, + output, + capacity, + written, + completion, + ) + } + }, + ffi::vllm_completion_free, + || params.logits_processor_error(), + ) } /// Runs one blocking streaming text completion. @@ -597,26 +720,67 @@ impl TranscriptionEngine { /// Loads a native engine owner with a transcription-only Rust method surface. /// /// ABI 17 cannot inspect the resolved task at load time. This constructor does - /// not probe or infer checkpoint architecture; native task selection remains - /// authoritative for future operations and diagnostics. + /// not probe or infer checkpoint architecture; native task selection and + /// wrong-task diagnostics remain authoritative. pub fn load(model_path: impl Into) -> Result { Ok(Self { - _inner: load_engine::(ModelConfig::new(model_path))?, + inner: load_engine::(ModelConfig::new(model_path))?, }) } + + /// Runs one blocking transcription and returns Rust-owned text and token IDs. + /// + /// Input storage is borrowed only for this call. Rust does not decode WAV + /// files, inspect PCM values, or resample audio; those checks and diagnostics + /// are native behavior. + pub fn transcribe(&mut self, input: TranscriptionInput<'_>) -> Result { + let input = MarshaledTranscriptionInput::new(input, &self.inner.compatibility)?; + transcribe_with( + &input, + |params, output| { + // SAFETY: the engine and marshaled input are live for this + // blocking call, and output is initialized on success. + unsafe { ffi::vllm_transcribe(self.inner.raw.as_ptr(), params, output) } + }, + ffi::vllm_transcription_free, + ) + } } impl EmbeddingEngine { /// Loads a native engine owner with an embedding-only Rust method surface. /// /// ABI 17 cannot inspect the resolved task at load time. This constructor does - /// not probe or infer checkpoint architecture; native task selection remains - /// authoritative for future operations and diagnostics. + /// not probe or infer checkpoint architecture; native task selection and + /// wrong-task diagnostics remain authoritative. pub fn load(model_path: impl Into) -> Result { Ok(Self { - _inner: load_engine::(ModelConfig::new(model_path))?, + inner: load_engine::(ModelConfig::new(model_path))?, }) } + + /// Runs one blocking native embedding batch. + /// + /// Input strings are borrowed only until this call returns. Native execution + /// is serialized per engine; exclusive `&mut self` access also keeps this + /// thread-local owner from overlapping Rust calls. The Rust-owned row-major + /// result preserves input order. + pub fn embed(&mut self, texts: I) -> Result + where + I: IntoIterator, + S: AsRef, + { + let texts = MarshaledEmbeddingInput::new(texts)?; + embed_with( + &texts, + |pointers, count, output| { + // SAFETY: all C strings and the pointer array remain live for this + // blocking call, and output is initialized on success. + unsafe { ffi::vllm_embed(self.inner.raw.as_ptr(), pointers, count, output) } + }, + ffi::vllm_embedding_result_free, + ) + } } fn load_engine(config: ModelConfig) -> Result, Error> { @@ -656,11 +820,383 @@ fn load_engine_with( }) } +fn validate_token_input(prompt_len: usize, capacity: usize) -> Result<(i32, i32), Error> { + if prompt_len == 0 { + return Err(invalid_configuration("prompt_tokens must not be empty")); + } + let prompt_len = i32::try_from(prompt_len) + .map_err(|_| invalid_configuration("prompt_tokens length exceeds native i32 range"))?; + let capacity = i32::try_from(capacity) + .map_err(|_| invalid_configuration("max_output_tokens exceeds native i32 range"))?; + Ok((prompt_len, capacity)) +} + +fn complete_tokens_with( + prompt_tokens: &[i32], + params: &ffi::vllm_sampling_params, + max_output_tokens: usize, + include_completion: bool, + call: impl FnOnce( + *const i32, + i32, + *const ffi::vllm_sampling_params, + *mut i32, + i32, + *mut i32, + *mut ffi::vllm_completion, + ) -> ffi::vllm_status, + free: unsafe extern "C" fn(*mut ffi::vllm_completion), + processor_error: impl FnOnce() -> Option, +) -> Result { + let (prompt_len, capacity) = validate_token_input(prompt_tokens.len(), max_output_tokens)?; + let mut token_ids = Vec::new(); + token_ids + .try_reserve_exact(max_output_tokens) + .map_err(|_| invalid_configuration("max_output_tokens cannot be allocated"))?; + token_ids.resize(max_output_tokens, 0); + let output = if token_ids.is_empty() { + ptr::null_mut() + } else { + token_ids.as_mut_ptr() + }; + let mut written = 0; + let mut raw = MaybeUninit::::uninit(); + let status = call( + prompt_tokens.as_ptr(), + prompt_len, + params, + output, + capacity, + &mut written, + raw.as_mut_ptr(), + ); + if status != ffi::vllm_status_VLLM_OK { + let native_error = status_result(status) + .expect_err("a non-OK native completion status must produce an error"); + if let Some(error) = processor_error() { + return Err(error); + } + return Err(native_error); + } + + // SAFETY: the successful native call initialized every completion field. + let raw = unsafe { raw.assume_init() }; + let guard = NativeResultGuard::new(raw, free); + if let Some(error) = processor_error() { + return Err(error); + } + token_completion_from_raw( + token_ids, + written, + prompt_tokens.len(), + include_completion, + guard.raw(), + ) +} + +fn token_completion_from_raw( + mut token_ids: Vec, + written: i32, + prompt_len: usize, + include_completion: bool, + raw: &ffi::vllm_completion, +) -> Result { + let written = usize::try_from(written) + .map_err(|_| invalid_native_output("written token count", "count is negative"))?; + if written > token_ids.len() { + return Err(invalid_native_output( + "written token count", + "count exceeds output capacity", + )); + } + let native_prompt = usize::try_from(raw.prompt_tokens) + .map_err(|_| invalid_native_output("prompt token count", "count is negative"))?; + if native_prompt != prompt_len { + return Err(invalid_native_output( + "prompt token count", + "count does not match the input prompt", + )); + } + let total = usize::try_from(raw.completion_tokens) + .map_err(|_| invalid_native_output("completion token count", "count is negative"))?; + if written != total.min(token_ids.len()) { + return Err(invalid_native_output( + "written token count", + "count does not match completion metadata and capacity", + )); + } + + let completion = include_completion + .then(|| completion_from_raw(raw)) + .transpose()?; + token_ids.truncate(written); + Ok(TokenCompletion { + token_ids, + completion, + truncated: written < total, + }) +} + +struct MarshaledTranscriptionInput<'a> { + raw: ffi::vllm_transcription_params, + _path: Option, + _samples: PhantomData<&'a [f32]>, +} + +impl<'a> MarshaledTranscriptionInput<'a> { + fn new(input: TranscriptionInput<'a>, compatibility: &Compatibility) -> Result { + Self::new_with(input, || compatibility.transcription_params_default()) + } + + fn new_with( + input: TranscriptionInput<'a>, + defaults: impl FnOnce() -> ffi::vllm_transcription_params, + ) -> Result { + match input { + TranscriptionInput::WavFile(path) => { + let path = path_to_cstring(path, "WAV path")?; + let mut raw = defaults(); + raw.audio_path = path.as_ptr(); + Ok(Self { + raw, + _path: Some(path), + _samples: PhantomData, + }) + } + TranscriptionInput::Pcm { + samples, + sample_rate, + } => { + let (n_samples, sample_rate) = validate_pcm_input(samples.len(), sample_rate)?; + let mut raw = defaults(); + raw.pcm = samples.as_ptr(); + raw.n_samples = n_samples; + raw.sample_rate = sample_rate; + Ok(Self { + raw, + _path: None, + _samples: PhantomData, + }) + } + } + } + + fn raw(&self) -> &ffi::vllm_transcription_params { + &self.raw + } +} + +fn validate_pcm_input(sample_count: usize, sample_rate: u32) -> Result<(i64, i32), Error> { + if sample_count == 0 { + return Err(invalid_configuration("PCM samples must not be empty")); + } + let sample_count = i64::try_from(sample_count) + .map_err(|_| invalid_configuration("PCM sample count exceeds native i64 range"))?; + if sample_rate == 0 { + return Err(invalid_configuration( + "PCM sample rate must be greater than zero", + )); + } + let sample_rate = i32::try_from(sample_rate) + .map_err(|_| invalid_configuration("PCM sample rate exceeds native i32 range"))?; + Ok((sample_count, sample_rate)) +} + +fn transcribe_with( + input: &MarshaledTranscriptionInput<'_>, + call: impl FnOnce( + *const ffi::vllm_transcription_params, + *mut ffi::vllm_transcription, + ) -> ffi::vllm_status, + free: unsafe extern "C" fn(*mut ffi::vllm_transcription), +) -> Result { + let mut raw = MaybeUninit::::uninit(); + let status = call(input.raw(), raw.as_mut_ptr()); + status_result(status)?; + // SAFETY: the successful native call initialized every result field. + let raw = unsafe { raw.assume_init() }; + let guard = NativeResultGuard::new(raw, free); + transcription_from_raw(guard.raw()) +} + +fn transcription_from_raw(raw: &ffi::vllm_transcription) -> Result { + let text = match raw.has_text { + 0 if raw.text.is_null() => None, + 0 => { + return Err(invalid_native_output( + "transcription text", + "pointer is non-null when has_text is zero", + )); + } + 1 if raw.text.is_null() => { + return Err(invalid_native_output( + "transcription text", + "pointer is null when has_text is one", + )); + } + 1 => Some(c_string_to_owned(raw.text, "transcription text")?), + _ => { + return Err(invalid_native_output( + "transcription has_text", + "value is not zero or one", + )); + } + }; + let token_count = usize::try_from(raw.n_token_ids) + .map_err(|_| invalid_native_output("transcription token count", "count is negative"))?; + validate_pointer_count(raw.token_ids, token_count, "transcription token IDs")?; + let token_ids = checked_copy_slice(raw.token_ids, token_count, "transcription token IDs")?; + Ok(Transcription { text, token_ids }) +} + +struct MarshaledEmbeddingInput { + strings: Vec, + pointers: Vec<*const c_char>, + count: i32, +} + +impl MarshaledEmbeddingInput { + fn new(texts: I) -> Result + where + I: IntoIterator, + S: AsRef, + { + let strings = texts + .into_iter() + .map(|text| to_cstring(text.as_ref(), "embedding text")) + .collect::, _>>()?; + let count = validate_embedding_count(strings.len())?; + let pointers = strings.iter().map(|text| text.as_ptr()).collect(); + Ok(Self { + strings, + pointers, + count, + }) + } + + fn pointers(&self) -> *const *const c_char { + debug_assert!(!self.strings.is_empty()); + self.pointers.as_ptr() + } +} + +fn validate_embedding_count(count: usize) -> Result { + if count == 0 { + return Err(invalid_configuration("embedding batch must not be empty")); + } + i32::try_from(count) + .map_err(|_| invalid_configuration("embedding batch exceeds native i32 range")) +} + +fn embed_with( + input: &MarshaledEmbeddingInput, + call: impl FnOnce(*const *const c_char, i32, *mut ffi::vllm_embedding_result) -> ffi::vllm_status, + free: unsafe extern "C" fn(*mut ffi::vllm_embedding_result), +) -> Result { + let mut raw = MaybeUninit::::uninit(); + let status = call(input.pointers(), input.count, raw.as_mut_ptr()); + status_result(status)?; + // SAFETY: the successful native call initialized every result field. + let raw = unsafe { raw.assume_init() }; + let guard = NativeResultGuard::new(raw, free); + embedding_from_raw(guard.raw(), input.strings.len()) +} + +fn embedding_from_raw( + raw: &ffi::vllm_embedding_result, + expected_rows: usize, +) -> Result { + let rows = usize::try_from(raw.n_embeddings) + .map_err(|_| invalid_native_output("embedding row count", "count is negative"))?; + if rows == 0 { + return Err(invalid_native_output( + "embedding row count", + "count is zero", + )); + } + if rows != expected_rows { + return Err(invalid_native_output( + "embedding row count", + "count does not match the input batch", + )); + } + let dimension = usize::try_from(raw.dim) + .map_err(|_| invalid_native_output("embedding dimension", "dimension is negative"))?; + if dimension == 0 { + return Err(invalid_native_output( + "embedding dimension", + "dimension is zero", + )); + } + let prompt_tokens = native_count_to_u32(raw.prompt_tokens, "embedding prompt token count")?; + let value_count = checked_product(rows, dimension, "embedding values")?; + validate_pointer_count(raw.values, value_count, "embedding values")?; + let values = checked_copy_slice(raw.values, value_count, "embedding values")?; + Ok(EmbeddingResult { + values, + dimension, + prompt_tokens, + }) +} + +fn checked_product(left: usize, right: usize, field: &'static str) -> Result { + left.checked_mul(right) + .ok_or_else(|| invalid_native_output(field, "element count overflows usize")) +} + +fn validate_pointer_count( + pointer: *const T, + length: usize, + field: &'static str, +) -> Result<(), Error> { + if length == 0 { + if pointer.is_null() { + return Ok(()); + } + return Err(invalid_native_output( + field, + "pointer is non-null for an empty result", + )); + } + if pointer.is_null() { + return Err(invalid_native_output( + field, + "pointer is null for a non-empty result", + )); + } + if (pointer as usize) % align_of::() != 0 { + return Err(invalid_native_output(field, "pointer is not aligned")); + } + if length > isize::MAX as usize / size_of::() { + return Err(invalid_native_output( + field, + "element count exceeds addressable slice size", + )); + } + Ok(()) +} + +fn checked_copy_slice( + pointer: *const T, + length: usize, + field: &'static str, +) -> Result, Error> { + validate_pointer_count(pointer, length, field)?; + if length == 0 { + return Ok(Vec::new()); + } + // SAFETY: validation established a non-null, aligned, addressable range and + // native keeps it live until the result guard is dropped. + Ok(unsafe { std::slice::from_raw_parts(pointer, length) }.to_vec()) +} + +fn invalid_native_output(field: &'static str, message: &'static str) -> Error { + Error::InvalidNativeOutput { field, message } +} + fn completion_from_raw(raw: &ffi::vllm_completion) -> Result { if raw.text.is_null() { - return Err(Error::Runtime { - message: "vllm_complete succeeded without text".to_owned(), - }); + return Err(invalid_native_output("completion text", "pointer is null")); } let text = c_string_to_owned(raw.text, "completion text")?; let finish_reason = if raw.finish_reason.is_null() { @@ -674,8 +1210,8 @@ fn completion_from_raw(raw: &ffi::vllm_completion) -> Result Ok(Completion { text, finish_reason, - prompt_tokens: count_to_u32(raw.prompt_tokens, "prompt token count")?, - completion_tokens: count_to_u32(raw.completion_tokens, "completion token count")?, + prompt_tokens: native_count_to_u32(raw.prompt_tokens, "prompt token count")?, + completion_tokens: native_count_to_u32(raw.completion_tokens, "completion token count")?, }) } @@ -691,8 +1227,8 @@ fn parse_finish_reason(value: String) -> FinishReason { } } -fn count_to_u32(value: i32, field: &'static str) -> Result { - u32::try_from(value).map_err(|_| invalid_configuration(format!("native {field} was negative"))) +fn native_count_to_u32(value: i32, field: &'static str) -> Result { + u32::try_from(value).map_err(|_| invalid_native_output(field, "count is negative")) } fn optional_u32_to_i32(value: Option, field: &'static str) -> Result, Error> { @@ -737,13 +1273,26 @@ fn c_string_to_owned(pointer: *const c_char, field: &'static str) -> Result { + raw: T, + free: unsafe extern "C" fn(*mut T), +} + +impl NativeResultGuard { + fn new(raw: T, free: unsafe extern "C" fn(*mut T)) -> Self { + Self { raw, free } + } + + fn raw(&self) -> &T { + &self.raw + } +} -impl Drop for CompletionGuard { +impl Drop for NativeResultGuard { fn drop(&mut self) { - // SAFETY: the native function initialized this completion and the guard - // releases its owned members exactly once. - unsafe { ffi::vllm_completion_free(&mut self.0) }; + // SAFETY: guards are armed only after successful native initialization + // and uniquely release their result storage exactly once. + unsafe { (self.free)(&mut self.raw) }; } } @@ -760,13 +1309,19 @@ impl Drop for NativeStringGuard { mod tests { use std::cell::RefCell; use std::ffi::CStr; + use std::mem::{align_of, size_of}; + use std::path::Path; use std::ptr::{self, NonNull}; + use std::sync::atomic::{AtomicUsize, Ordering}; use vllm_cpp_sys as ffi; use super::{ - load_engine_with, Device, EmbeddingTask, MarshaledModelParams, ModelConfig, - SchedulerPolicy, TextTask, Toggle, TranscriptionTask, + checked_product, complete_tokens_with, embed_with, embedding_from_raw, load_engine_with, + token_completion_from_raw, transcribe_with, transcription_from_raw, + validate_embedding_count, validate_pcm_input, validate_token_input, Device, EmbeddingTask, + MarshaledEmbeddingInput, MarshaledModelParams, MarshaledTranscriptionInput, ModelConfig, + SchedulerPolicy, TextTask, Toggle, TranscriptionInput, TranscriptionTask, }; use crate::abi::Compatibility; use crate::Error; @@ -968,4 +1523,522 @@ mod tests { assert!(matches!(result, Err(Error::ModelLoad { .. }))); } + + static COMPLETION_FREES: AtomicUsize = AtomicUsize::new(0); + static TRANSCRIPTION_FREES: AtomicUsize = AtomicUsize::new(0); + static EMBEDDING_FREES: AtomicUsize = AtomicUsize::new(0); + + unsafe extern "C" fn count_completion_free(_: *mut ffi::vllm_completion) { + COMPLETION_FREES.fetch_add(1, Ordering::SeqCst); + } + + unsafe extern "C" fn count_transcription_free(_: *mut ffi::vllm_transcription) { + TRANSCRIPTION_FREES.fetch_add(1, Ordering::SeqCst); + } + + unsafe extern "C" fn count_embedding_free(_: *mut ffi::vllm_embedding_result) { + EMBEDDING_FREES.fetch_add(1, Ordering::SeqCst); + } + + fn zeroed_sampling_params() -> ffi::vllm_sampling_params { + // SAFETY: zero is a valid bit pattern for every generated C field. + unsafe { std::mem::zeroed() } + } + + fn raw_completion( + text: *mut std::os::raw::c_char, + prompt_tokens: i32, + completion_tokens: i32, + ) -> ffi::vllm_completion { + ffi::vllm_completion { + text, + finish_reason: ptr::null(), + prompt_tokens, + completion_tokens, + } + } + + #[test] + fn token_input_validation_rejects_empty_and_native_range_overflow() { + assert!(matches!( + validate_token_input(0, 0), + Err(Error::InvalidConfiguration { .. }) + )); + assert!(matches!( + validate_token_input(1, i32::MAX as usize + 1), + Err(Error::InvalidConfiguration { .. }) + )); + #[cfg(target_pointer_width = "64")] + assert!(matches!( + validate_token_input(i32::MAX as usize + 1, 0), + Err(Error::InvalidConfiguration { .. }) + )); + } + + #[test] + fn token_zero_capacity_uses_null_and_hidden_completion_metadata() { + COMPLETION_FREES.store(0, Ordering::SeqCst); + let text = b"generated\0"; + let params = zeroed_sampling_params(); + let result = complete_tokens_with( + &[9707], + ¶ms, + 0, + false, + |prompt, n_prompt, _, output, capacity, written, completion| { + assert!(!prompt.is_null()); + assert_eq!(n_prompt, 1); + assert!(output.is_null()); + assert_eq!(capacity, 0); + // SAFETY: all pointers target writable caller storage. + unsafe { + *written = 0; + *completion = raw_completion(text.as_ptr().cast_mut().cast(), 1, 3); + } + ffi::vllm_status_VLLM_OK + }, + count_completion_free, + || None, + ) + .expect("zero-capacity completion"); + + assert!(result.token_ids.is_empty()); + assert!(result.completion.is_none()); + assert!(result.truncated); + assert_eq!(COMPLETION_FREES.load(Ordering::SeqCst), 1); + } + + #[test] + fn token_completion_copies_truncated_ids_and_optional_completion() { + COMPLETION_FREES.store(0, Ordering::SeqCst); + let text = b"ok\0"; + let finish = b"length\0"; + let params = zeroed_sampling_params(); + let result = complete_tokens_with( + &[1, 2], + ¶ms, + 2, + true, + |_, _, _, output, _, written, completion| { + // SAFETY: all pointers target writable caller storage. + unsafe { + *output = 10; + *output.add(1) = 11; + *written = 2; + *completion = ffi::vllm_completion { + text: text.as_ptr().cast_mut().cast(), + finish_reason: finish.as_ptr().cast(), + prompt_tokens: 2, + completion_tokens: 4, + }; + } + ffi::vllm_status_VLLM_OK + }, + count_completion_free, + || None, + ) + .expect("truncated completion"); + + assert_eq!(result.token_ids, [10, 11]); + assert!(result.truncated); + let completion = result.completion.expect("included completion"); + assert_eq!(completion.text, "ok"); + assert_eq!(completion.prompt_tokens, 2); + assert_eq!(completion.completion_tokens, 4); + assert_eq!(COMPLETION_FREES.load(Ordering::SeqCst), 1); + } + + #[test] + fn token_metadata_relationships_are_validated() { + let text = b"ok\0"; + let base = raw_completion(text.as_ptr().cast_mut().cast(), 1, 2); + for (written, capacity, prompt, total) in [ + (-1, 2, 1, 2), + (3, 2, 1, 3), + (2, 2, 9, 2), + (0, 2, 1, -1), + (1, 2, 1, 2), + ] { + let raw = raw_completion(base.text, prompt, total); + let result = token_completion_from_raw(vec![0; capacity], written, 1, false, &raw); + assert!(matches!(result, Err(Error::InvalidNativeOutput { .. }))); + } + } + + #[test] + fn token_native_failure_discards_partial_output_without_arming_guard() { + COMPLETION_FREES.store(0, Ordering::SeqCst); + let params = zeroed_sampling_params(); + let result = complete_tokens_with( + &[1], + ¶ms, + 1, + false, + |_, _, _, output, _, written, _| { + // SAFETY: output and written point to caller-owned storage. + unsafe { + *output = 99; + *written = 1; + } + ffi::vllm_status_VLLM_ERR_INVALID_ARGUMENT + }, + count_completion_free, + || None, + ); + assert!(matches!(result, Err(Error::InvalidArgument { .. }))); + assert_eq!(COMPLETION_FREES.load(Ordering::SeqCst), 0); + } + + #[test] + fn token_conversion_error_still_frees_once() { + COMPLETION_FREES.store(0, Ordering::SeqCst); + let invalid_utf8 = [0xff_u8, 0]; + let params = zeroed_sampling_params(); + let result = complete_tokens_with( + &[1], + ¶ms, + 1, + true, + |_, _, _, output, _, written, completion| { + // SAFETY: all pointers target writable caller storage. + unsafe { + *output = 2; + *written = 1; + *completion = raw_completion(invalid_utf8.as_ptr().cast_mut().cast(), 1, 1); + } + ffi::vllm_status_VLLM_OK + }, + count_completion_free, + || None, + ); + assert_eq!( + result, + Err(Error::InvalidUtf8 { + field: "completion text" + }) + ); + assert_eq!(COMPLETION_FREES.load(Ordering::SeqCst), 1); + } + + fn transcription_defaults() -> ffi::vllm_transcription_params { + ffi::vllm_transcription_params { + audio_path: ptr::null(), + pcm: ptr::null(), + n_samples: 37, + sample_rate: 38, + } + } + + #[test] + fn transcription_marshaling_selects_and_retains_one_pointer_family() { + let path = MarshaledTranscriptionInput::new_with( + TranscriptionInput::WavFile(Path::new("audio.wav")), + transcription_defaults, + ) + .expect("path input"); + assert!(!path.raw.audio_path.is_null()); + assert!(path.raw.pcm.is_null()); + assert_eq!(c_string(path.raw.audio_path), "audio.wav"); + assert_eq!(path.raw.n_samples, 37); + assert_eq!(path.raw.sample_rate, 38); + + let samples = [0.25, -0.5]; + let pcm = MarshaledTranscriptionInput::new_with( + TranscriptionInput::Pcm { + samples: &samples, + sample_rate: 44_100, + }, + transcription_defaults, + ) + .expect("PCM input"); + assert!(pcm.raw.audio_path.is_null()); + assert_eq!(pcm.raw.pcm, samples.as_ptr()); + assert_eq!(pcm.raw.n_samples, 2); + assert_eq!(pcm.raw.sample_rate, 44_100); + } + + #[test] + fn transcription_input_validation_rejects_invalid_pcm_and_path() { + assert!(validate_pcm_input(0, 16_000).is_err()); + assert!(validate_pcm_input(1, 0).is_err()); + assert!(validate_pcm_input(1, i32::MAX as u32 + 1).is_err()); + #[cfg(target_pointer_width = "64")] + assert!(validate_pcm_input(i64::MAX as usize + 1, 16_000).is_err()); + #[cfg(unix)] + assert!(matches!( + MarshaledTranscriptionInput::new_with( + TranscriptionInput::WavFile(Path::new("bad\0path")), + transcription_defaults, + ), + Err(Error::InteriorNul { field: "WAV path" }) + )); + } + + #[test] + fn transcription_metadata_and_optional_outputs_are_validated() { + let text = b"text\0"; + let id = 7; + let cases = [ + ffi::vllm_transcription { + text: ptr::null_mut(), + token_ids: ptr::null_mut(), + n_token_ids: 0, + has_text: 2, + }, + ffi::vllm_transcription { + text: text.as_ptr().cast_mut().cast(), + token_ids: ptr::null_mut(), + n_token_ids: 0, + has_text: 0, + }, + ffi::vllm_transcription { + text: ptr::null_mut(), + token_ids: ptr::null_mut(), + n_token_ids: 0, + has_text: 1, + }, + ffi::vllm_transcription { + text: ptr::null_mut(), + token_ids: ptr::from_ref(&id).cast_mut(), + n_token_ids: 0, + has_text: 0, + }, + ffi::vllm_transcription { + text: ptr::null_mut(), + token_ids: ptr::null_mut(), + n_token_ids: 1, + has_text: 0, + }, + ffi::vllm_transcription { + text: ptr::null_mut(), + token_ids: ptr::null_mut(), + n_token_ids: -1, + has_text: 0, + }, + ]; + for raw in cases { + assert!(matches!( + transcription_from_raw(&raw), + Err(Error::InvalidNativeOutput { .. }) + )); + } + + let empty = ffi::vllm_transcription { + text: ptr::null_mut(), + token_ids: ptr::null_mut(), + n_token_ids: 0, + has_text: 0, + }; + assert_eq!( + transcription_from_raw(&empty).expect("IDs-only empty result"), + super::Transcription { + text: None, + token_ids: vec![] + } + ); + } + + #[test] + fn transcription_copies_outputs_and_frees_once_on_conversion_error() { + TRANSCRIPTION_FREES.store(0, Ordering::SeqCst); + let samples = [0.0]; + let input = MarshaledTranscriptionInput::new_with( + TranscriptionInput::Pcm { + samples: &samples, + sample_rate: 16_000, + }, + transcription_defaults, + ) + .expect("PCM input"); + let text = b"copied\0"; + let mut ids = vec![3, 4, 3]; + let copied = transcribe_with( + &input, + |_, output| { + // SAFETY: output targets writable caller storage and test data stays live. + unsafe { + *output = ffi::vllm_transcription { + text: text.as_ptr().cast_mut().cast(), + token_ids: ids.as_mut_ptr(), + n_token_ids: 3, + has_text: 1, + }; + } + ffi::vllm_status_VLLM_OK + }, + count_transcription_free, + ) + .expect("copied transcription"); + ids.fill(9); + assert_eq!(copied.text.as_deref(), Some("copied")); + assert_eq!(copied.token_ids, [3, 4, 3]); + assert_eq!(TRANSCRIPTION_FREES.load(Ordering::SeqCst), 1); + + let invalid_utf8 = [0xff_u8, 0]; + let result = transcribe_with( + &input, + |_, output| { + // SAFETY: output targets writable caller storage. + unsafe { + *output = ffi::vllm_transcription { + text: invalid_utf8.as_ptr().cast_mut().cast(), + token_ids: ptr::null_mut(), + n_token_ids: 0, + has_text: 1, + }; + } + ffi::vllm_status_VLLM_OK + }, + count_transcription_free, + ); + assert!(matches!(result, Err(Error::InvalidUtf8 { .. }))); + assert_eq!(TRANSCRIPTION_FREES.load(Ordering::SeqCst), 2); + } + + #[test] + fn embedding_marshaling_retains_strings_and_allows_empty_text() { + let input = MarshaledEmbeddingInput::new(["first", "", "third"]).expect("embedding input"); + assert_eq!(input.count, 3); + for (index, expected) in ["first", "", "third"].iter().enumerate() { + // SAFETY: pointers refer to CStrings retained by input. + let actual = unsafe { CStr::from_ptr(*input.pointers.as_ptr().add(index)) }; + assert_eq!(actual.to_str().expect("UTF-8"), *expected); + } + assert!(matches!( + MarshaledEmbeddingInput::new(std::iter::empty::<&str>()), + Err(Error::InvalidConfiguration { .. }) + )); + assert!(matches!( + MarshaledEmbeddingInput::new(["bad\0text"]), + Err(Error::InteriorNul { + field: "embedding text" + }) + )); + assert!(validate_embedding_count(i32::MAX as usize + 1).is_err()); + } + + #[test] + fn embedding_metadata_rejects_shape_pointer_and_size_errors() { + let aligned = NonNull::::dangling().as_ptr(); + let cases = [ + ffi::vllm_embedding_result { + values: ptr::null_mut(), + n_embeddings: 0, + dim: 1, + prompt_tokens: 0, + }, + ffi::vllm_embedding_result { + values: ptr::null_mut(), + n_embeddings: 2, + dim: 1, + prompt_tokens: 0, + }, + ffi::vllm_embedding_result { + values: ptr::null_mut(), + n_embeddings: 1, + dim: 0, + prompt_tokens: 0, + }, + ffi::vllm_embedding_result { + values: aligned, + n_embeddings: 1, + dim: 1, + prompt_tokens: -1, + }, + ffi::vllm_embedding_result { + values: ptr::null_mut(), + n_embeddings: 1, + dim: 1, + prompt_tokens: 0, + }, + ffi::vllm_embedding_result { + values: (align_of::() - 1) as *mut f32, + n_embeddings: 1, + dim: 1, + prompt_tokens: 0, + }, + ffi::vllm_embedding_result { + values: aligned, + n_embeddings: i32::MAX, + dim: i32::MAX, + prompt_tokens: 0, + }, + ]; + for (index, raw) in cases.iter().enumerate() { + let expected_rows = if index == 1 { + 1 + } else { + raw.n_embeddings.max(1) as usize + }; + assert!(matches!( + embedding_from_raw(raw, expected_rows), + Err(Error::InvalidNativeOutput { .. }) + )); + } + assert!(checked_product(usize::MAX, 2, "test product").is_err()); + assert!(i32::MAX as usize * i32::MAX as usize > isize::MAX as usize / size_of::()); + } + + #[test] + fn embedding_rows_preserve_order_own_values_and_free_once() { + EMBEDDING_FREES.store(0, Ordering::SeqCst); + let input = MarshaledEmbeddingInput::new(["a", "b"]).expect("embedding input"); + let mut native_values = vec![1.0, 2.0, 3.0, 4.0]; + let result = embed_with( + &input, + |pointers, count, output| { + assert_eq!(count, 2); + assert!(!pointers.is_null()); + // SAFETY: output targets writable caller storage and values stays live. + unsafe { + *output = ffi::vllm_embedding_result { + values: native_values.as_mut_ptr(), + n_embeddings: 2, + dim: 2, + prompt_tokens: 5, + }; + } + ffi::vllm_status_VLLM_OK + }, + count_embedding_free, + ) + .expect("embedding result"); + native_values.fill(9.0); + + assert_eq!(result.values(), [1.0, 2.0, 3.0, 4.0]); + assert_eq!(result.n_embeddings(), 2); + assert_eq!(result.dimension(), 2); + assert_eq!(result.prompt_tokens(), 5); + assert_eq!(result.row(0), Some(&[1.0, 2.0][..])); + assert_eq!(result.row(2), None); + assert_eq!( + result.rows().collect::>(), + [&[1.0, 2.0][..], &[3.0, 4.0][..]] + ); + assert_eq!(EMBEDDING_FREES.load(Ordering::SeqCst), 1); + } + + #[test] + fn embedding_conversion_error_frees_once() { + EMBEDDING_FREES.store(0, Ordering::SeqCst); + let input = MarshaledEmbeddingInput::new(["a"]).expect("embedding input"); + let result = embed_with( + &input, + |_, _, output| { + // SAFETY: output targets writable caller storage. + unsafe { + *output = ffi::vllm_embedding_result { + values: ptr::null_mut(), + n_embeddings: 1, + dim: 1, + prompt_tokens: 1, + }; + } + ffi::vllm_status_VLLM_OK + }, + count_embedding_free, + ); + assert!(matches!(result, Err(Error::InvalidNativeOutput { .. }))); + assert_eq!(EMBEDDING_FREES.load(Ordering::SeqCst), 1); + } } diff --git a/vllm-cpp/src/error.rs b/vllm-cpp/src/error.rs index 70501bd..c3a4abf 100644 --- a/vllm-cpp/src/error.rs +++ b/vllm-cpp/src/error.rs @@ -60,6 +60,11 @@ pub enum Error { PathEncoding, /// Native code returned bytes that are not valid UTF-8. InvalidUtf8 { field: &'static str }, + /// Native code returned malformed count, pointer, shape, or overflow metadata. + InvalidNativeOutput { + field: &'static str, + message: &'static str, + }, /// An asynchronous output callback panicked. CallbackPanicked, /// A custom logits processor panicked. @@ -94,6 +99,9 @@ impl fmt::Display for Error { Self::InteriorNul { field } => write!(f, "{field} contains an interior NUL byte"), Self::PathEncoding => write!(f, "path cannot be represented by the native API"), Self::InvalidUtf8 { field } => write!(f, "native {field} is not valid UTF-8"), + Self::InvalidNativeOutput { field, message } => { + write!(f, "invalid native output for {field}: {message}") + } Self::CallbackPanicked => write!(f, "asynchronous request callback panicked"), Self::LogitsProcessorPanicked => write!(f, "custom logits processor panicked"), Self::RequestCallbackThread { operation } => { diff --git a/vllm-cpp/src/lib.rs b/vllm-cpp/src/lib.rs index 952d81b..594c921 100644 --- a/vllm-cpp/src/lib.rs +++ b/vllm-cpp/src/lib.rs @@ -4,19 +4,21 @@ //! //! Resolve a Hub model with [`HuggingFaceModel`] (default `main`, or an explicit //! revision), then create a text [`Engine`] with [`Engine::load`] or configure -//! native model settings through [`EngineBuilder`]. Path-only -//! [`TranscriptionEngine`] and [`EmbeddingEngine`] owners reserve task-specific -//! method surfaces for future operations. [`SamplingParams`] owns sampling, -//! stop-string, [`StructuredOutput`] settings, and an optional host-side logits -//! processor for completion calls. The text engine provides blocking completion, -//! streaming, raw-JSON chat, and [`Engine::submit`] for a concurrent [`Request`]. -//! Enable `serde` for `serde_json::Value` chat helpers. +//! native model settings through [`EngineBuilder`]. [`TranscriptionEngine`] and +//! [`EmbeddingEngine`] provide blocking task-specific operations. [`SamplingParams`] +//! owns sampling, stop-string, [`StructuredOutput`] settings, and an optional +//! host-side logits processor for completion calls. The text engine provides +//! blocking text and pre-tokenized completion, streaming, raw-JSON chat, and +//! [`Engine::submit`] for a concurrent [`Request`]. Enable `serde` for +//! `serde_json::Value` chat helpers. //! //! # Ownership and callbacks //! //! [`Engine`] is a cloneable RAII owner; clones share one reference-counted -//! native engine. Rust copies completion, stream, chat, and error text before -//! native storage is freed or reused. Blocking callbacks may borrow caller data. +//! native engine. Rust copies completion, transcription, embedding, stream, +//! chat, and error data before native storage is freed or reused. Pre-tokenized +//! prompts, audio, paths, and embedding strings are borrowed only for their +//! blocking calls. Blocking callbacks may borrow caller data. //! Their panics are caught before the C boundary and resumed after the native //! call returns. Custom logits processors are `Send + Sync`, may run concurrently //! on native worker threads, and report contained panic through @@ -25,8 +27,9 @@ //! A [`Request`] retains its engine and asynchronous callback until native //! free/join completes. Requests are `Send` but intentionally not `Sync`. The text //! [`Engine`] is `Send + Sync`; transcription and embedding owners are -//! conservatively neither and must remain on their creating thread. Asynchronous -//! callbacks run on a native delivery thread, must be `Send + 'static`, and surface panic through +//! conservatively neither, must remain on their creating thread, and require +//! exclusive access for operations. Native embedding execution is serialized. +//! Asynchronous callbacks run on a native delivery thread, must be `Send + 'static`, and surface panic through //! [`Error::CallbackPanicked`]. ABI version 17 forbids waiting for or freeing a //! request from its callback thread; callback-thread drop delegates ownership to //! a cleanup reaper instead. ABI 17 exposes no task-introspection API, so loading @@ -62,7 +65,8 @@ mod request; pub use callback::{StreamControl, StreamEvent, StreamOutcome}; pub use engine::{ - Completion, EmbeddingEngine, Engine, EngineBuilder, FinishReason, TranscriptionEngine, + Completion, EmbeddingEngine, EmbeddingResult, Engine, EngineBuilder, FinishReason, + TokenCompletion, Transcription, TranscriptionEngine, TranscriptionInput, }; pub use error::{Error, HuggingFaceError}; pub use hf::HuggingFaceModel; diff --git a/vllm-cpp/tests/qwen3.rs b/vllm-cpp/tests/qwen3.rs index 757f3d9..a6d0258 100644 --- a/vllm-cpp/tests/qwen3.rs +++ b/vllm-cpp/tests/qwen3.rs @@ -4,8 +4,8 @@ use std::thread; use std::time::{Duration, Instant}; use vllm_cpp::{ - Engine, Error, FinishReason, Request, RequestOutcome, SamplingParams, StreamControl, - StructuredOutput, + EmbeddingEngine, Engine, Error, FinishReason, Request, RequestOutcome, SamplingParams, + StreamControl, StructuredOutput, TranscriptionEngine, TranscriptionInput, }; fn model_path() -> Option { @@ -55,6 +55,189 @@ fn wait_until_done(request: &Request) { } } +fn native_fixture(relative: &str) -> Option { + let path = Path::new(env!("CARGO_MANIFEST_DIR")) + .join("../vllm-cpp-sys/vllm.cpp/tests/vllm/models/fixtures") + .join(relative); + if path.exists() { + Some(path) + } else { + eprintln!( + "skipping native fixture test; fixture is absent: {}", + path.display() + ); + None + } +} + +fn read_pcm16_mono_wav(path: &Path) -> Vec { + let bytes = std::fs::read(path).expect("read fixture WAV"); + assert_eq!(&bytes[0..4], b"RIFF"); + assert_eq!(&bytes[8..12], b"WAVE"); + let mut offset = 12; + while offset + 8 <= bytes.len() { + let name = &bytes[offset..offset + 4]; + let size = u32::from_le_bytes(bytes[offset + 4..offset + 8].try_into().unwrap()) as usize; + let start = offset + 8; + if name == b"data" { + return bytes[start..start + size] + .chunks_exact(2) + .map(|sample| i16::from_le_bytes([sample[0], sample[1]]) as f32 / 32768.0) + .collect(); + } + offset = start + size + (size % 2); + } + panic!("fixture WAV has no data chunk"); +} + +#[test] +fn pretokenized_completion_matches_qwen_hello_and_reports_truncation() { + with_engine(|engine, _| { + let params = SamplingParams::greedy().max_tokens(4); + let full = engine + .complete_tokens(&[9707], ¶ms, 8, true) + .expect("full pre-tokenized completion"); + assert_eq!(full.token_ids.len(), 4); + assert!(!full.truncated); + let details = full.completion.as_ref().expect("completion details"); + assert_eq!(details.prompt_tokens, 1); + assert_eq!(details.completion_tokens, 4); + assert_eq!( + details, + &engine + .complete("Hello", ¶ms) + .expect("string-prompt parity completion") + ); + + let small = engine + .complete_tokens(&[9707], ¶ms, 2, false) + .expect("truncated pre-tokenized completion"); + assert_eq!(small.token_ids, full.token_ids[..2]); + assert!(small.completion.is_none()); + assert!(small.truncated); + + let zero = engine + .complete_tokens(&[9707], ¶ms, 0, false) + .expect("zero-capacity pre-tokenized completion"); + assert!(zero.token_ids.is_empty()); + assert!(zero.completion.is_none()); + assert!(zero.truncated); + }); +} + +#[test] +fn pretokenized_logits_processor_controls_tokens_and_panic_leaves_engine_reusable() { + with_engine(|engine, _| { + let params = SamplingParams::greedy() + .max_tokens(3) + .logits_processor(|_, logits| { + logits.fill(f32::NEG_INFINITY); + logits[10] = f32::INFINITY; + }); + let forced = engine + .complete_tokens(&[9707], ¶ms, 3, false) + .expect("forced token completion"); + assert_eq!(forced.token_ids, [10, 10, 10]); + + let panicking = SamplingParams::greedy() + .max_tokens(1) + .logits_processor(|_, _| panic!("intentional token processor panic")); + assert_eq!( + engine + .complete_tokens(&[9707], &panicking, 1, true) + .expect_err("processor panic"), + Error::LogitsProcessorPanicked + ); + engine + .complete_tokens(&[9707], &SamplingParams::greedy().max_tokens(1), 1, false) + .expect("engine remains reusable"); + }); +} + +#[test] +fn committed_transcription_fixture_supports_path_pcm_and_wrong_task() { + let Some(root) = native_fixture("parakeet_e2e") else { + return; + }; + let model = root.join("ctc"); + let wav = root.join("audio.wav"); + let mut engine = TranscriptionEngine::load(&model).expect("load CTC fixture"); + let from_path = engine + .transcribe(TranscriptionInput::WavFile(&wav)) + .expect("transcribe fixture path"); + assert_eq!(from_path.token_ids, [3, 4, 3]); + assert_eq!(from_path.text.as_deref(), Some("atheat")); + + let samples = read_pcm16_mono_wav(&wav); + let from_pcm = engine + .transcribe(TranscriptionInput::Pcm { + samples: &samples, + sample_rate: 16_000, + }) + .expect("transcribe fixture PCM"); + assert_eq!(from_pcm, from_path); + + let text = Engine::load(&model).expect("task-neutral load of CTC fixture"); + assert!(matches!( + text.complete_tokens(&[0], &SamplingParams::greedy().max_tokens(1), 1, false), + Err(Error::InvalidArgument { .. }) + )); + + if let Some(embedding_model) = native_fixture("llama_embed_e2e") { + let mut wrong = + TranscriptionEngine::load(&embedding_model).expect("task-neutral embedding load"); + assert!(matches!( + wrong.transcribe(TranscriptionInput::WavFile(&wav)), + Err(Error::InvalidArgument { .. }) + )); + } +} + +#[test] +fn committed_embedding_fixture_preserves_shape_order_ownership_and_wrong_task() { + let Some(model) = native_fixture("llama_embed_e2e") else { + return; + }; + let mut engine = EmbeddingEngine::load(&model).expect("load embedding fixture"); + let result = engine + .embed(["the quick brown fox", "the lazy dog"]) + .expect("embed fixture inputs"); + assert_eq!(result.n_embeddings(), 2); + assert_eq!(result.dimension(), 64); + assert!(result.prompt_tokens() > 0); + assert_ne!(result.row(0), result.row(1)); + for row in result.rows() { + let l2 = row + .iter() + .map(|value| f64::from(*value).powi(2)) + .sum::() + .sqrt(); + assert!((l2 - 1.0).abs() < 1e-5); + } + drop(engine); + assert_eq!(result.values().len(), 128); + + let text = Engine::load(&model).expect("task-neutral embedding load"); + assert!(matches!( + text.complete_tokens(&[0], &SamplingParams::greedy().max_tokens(1), 1, false), + Err(Error::InvalidArgument { .. }) + )); + + if let Some(transcription_root) = native_fixture("parakeet_e2e") { + let mut wrong = EmbeddingEngine::load(transcription_root.join("ctc")) + .expect("task-neutral transcription load"); + assert!(matches!( + wrong.embed(["hello"]), + Err(Error::InvalidArgument { .. }) + )); + } + + let mut engine = EmbeddingEngine::load(&model).expect("reload embedding fixture"); + engine + .embed(["the fox"]) + .expect("embedding owner remains usable"); +} + #[test] fn greedy_completion_and_streaming_match() { with_engine(|engine, _| { diff --git a/vllm-cpp/tests/safe_api.rs b/vllm-cpp/tests/safe_api.rs index 94ced05..9a0b0ed 100644 --- a/vllm-cpp/tests/safe_api.rs +++ b/vllm-cpp/tests/safe_api.rs @@ -1,7 +1,8 @@ use static_assertions::{assert_impl_all, assert_not_impl_any}; use vllm_cpp::{ - Device, EmbeddingEngine, Engine, EngineBuilder, Error, HuggingFaceError, HuggingFaceModel, - Request, SchedulerPolicy, Toggle, TranscriptionEngine, + Device, EmbeddingEngine, EmbeddingResult, Engine, EngineBuilder, Error, HuggingFaceError, + HuggingFaceModel, Request, SchedulerPolicy, Toggle, TokenCompletion, Transcription, + TranscriptionEngine, TranscriptionInput, }; assert_impl_all!(Device: Clone, Copy, std::fmt::Debug, Default, Eq, PartialEq, Send, Sync); @@ -11,6 +12,10 @@ assert_impl_all!(HuggingFaceError: Clone, std::fmt::Debug, Eq, PartialEq); assert_impl_all!(HuggingFaceModel: Clone, std::fmt::Debug); assert_impl_all!(Request: Send); assert_impl_all!(vllm_cpp::SamplingParams: Clone, Send, Sync); +assert_impl_all!(TokenCompletion: Clone, std::fmt::Debug, Eq, PartialEq, Send, Sync); +assert_impl_all!(Transcription: Clone, std::fmt::Debug, Eq, PartialEq, Send, Sync); +assert_impl_all!(EmbeddingResult: Clone, std::fmt::Debug, PartialEq, Send, Sync); +assert_impl_all!(TranscriptionInput<'static>: Clone, Copy, std::fmt::Debug, Send, Sync); assert_not_impl_any!(Request: Sync); assert_not_impl_any!(TranscriptionEngine: Send, Sync); assert_not_impl_any!(EmbeddingEngine: Send, Sync); @@ -19,6 +24,33 @@ fn missing_model() -> &'static str { "/nonexistent/vllm-cpp-rs-safe-api-model" } +#[test] +fn constructs_both_borrowed_transcription_inputs() { + let path = std::path::Path::new("audio.wav"); + let samples = [0.0_f32, 0.25]; + let inputs = [ + TranscriptionInput::WavFile(path), + TranscriptionInput::Pcm { + samples: &samples, + sample_rate: 16_000, + }, + ]; + assert!(matches!(inputs[0], TranscriptionInput::WavFile(_))); + assert!(matches!(inputs[1], TranscriptionInput::Pcm { .. })); +} + +#[test] +fn invalid_native_output_display_is_stable() { + let error = Error::InvalidNativeOutput { + field: "embedding dimension", + message: "dimension is zero", + }; + assert_eq!( + error.to_string(), + "invalid native output for embedding dimension: dimension is zero" + ); +} + #[test] fn hugging_face_constructors_accept_default_and_explicit_revisions() { let gguf = HuggingFaceModel::gguf("owner/model", "model.gguf");