diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 1beaffb..b679d2e 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -24,6 +24,8 @@ jobs: steps: - name: Check out sources uses: actions/checkout@v4.4.0 + with: + persist-credentials: false - name: Install Rust toolchain uses: dtolnay/rust-toolchain@stable with: @@ -42,6 +44,7 @@ jobs: - name: Check out sources uses: actions/checkout@v4.4.0 with: + persist-credentials: false submodules: recursive - name: Install Rust toolchain uses: dtolnay/rust-toolchain@stable @@ -68,6 +71,7 @@ jobs: - name: Check out sources uses: actions/checkout@v4.4.0 with: + persist-credentials: false submodules: recursive - name: Install Rust toolchain uses: dtolnay/rust-toolchain@stable @@ -92,6 +96,7 @@ jobs: - name: Check out sources uses: actions/checkout@v4.4.0 with: + persist-credentials: false submodules: recursive - name: Install Rust toolchain uses: dtolnay/rust-toolchain@stable @@ -112,6 +117,7 @@ jobs: - name: Check out sources uses: actions/checkout@v4.4.0 with: + persist-credentials: false submodules: recursive - name: Install Rust toolchain uses: dtolnay/rust-toolchain@stable @@ -135,11 +141,12 @@ jobs: package: needs: native runs-on: ubuntu-24.04 - timeout-minutes: 45 + timeout-minutes: 60 steps: - name: Check out sources uses: actions/checkout@v4.4.0 with: + persist-credentials: false submodules: recursive - name: Install Rust toolchain uses: dtolnay/rust-toolchain@stable @@ -148,6 +155,12 @@ jobs: with: save-if: ${{ github.event_name != 'pull_request' || github.event.pull_request.head.repo.full_name == github.repository }} - name: Install native dependencies - run: sudo apt-get update && sudo apt-get install --yes --no-install-recommends build-essential cmake ninja-build + run: sudo apt-get update && sudo apt-get install --yes --no-install-recommends build-essential cmake ninja-build jq + - name: Install Just + uses: extractions/setup-just@v3.1.0 + with: + just-version: "1.40.0" + - name: Test crate archives and downstream consumers + run: just package-test - name: Dry-run workspace publish - run: cargo publish --workspace --locked --dry-run + run: just publish-dry-run diff --git a/.github/workflows/platforms.yml b/.github/workflows/platforms.yml index cabc778..6f81931 100644 --- a/.github/workflows/platforms.yml +++ b/.github/workflows/platforms.yml @@ -23,6 +23,7 @@ jobs: - name: Check out exact sources uses: actions/checkout@v4.4.0 with: + persist-credentials: false submodules: recursive - name: Verify checkout and native pin shell: bash @@ -52,6 +53,7 @@ jobs: - name: Check out exact sources uses: actions/checkout@v4.4.0 with: + persist-credentials: false submodules: recursive - name: Verify checkout, native pin, and host shell: bash @@ -76,6 +78,7 @@ jobs: - name: Check out exact sources uses: actions/checkout@v4.4.0 with: + persist-credentials: false submodules: recursive - name: Verify checkout, native pin, and host shell: bash @@ -102,6 +105,7 @@ jobs: - name: Check out exact sources uses: actions/checkout@v4.4.0 with: + persist-credentials: false submodules: recursive - name: Verify checkout, native pin, and Metal toolchain shell: bash @@ -131,6 +135,7 @@ jobs: - name: Check out exact sources uses: actions/checkout@v4.4.0 with: + persist-credentials: false submodules: recursive - name: Verify checkout and native pin shell: bash diff --git a/CHANGELOG.md b/CHANGELOG.md index 77c8401..a8e81d0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,11 +4,13 @@ All notable changes to this project will be documented in this file. ## Unreleased +## [0.0.1] - 2026-08-22 + ### Added - Checked-in raw Rust declarations for the 19-symbol stable vllm.cpp C API at ABI version 10, with header, symbol, layout, and runtime conformance checks. -- A safe API for model loading, blocking completion and streaming, raw-JSON and optional serde chat, structured output, owned sampling parameters, panic-contained custom logits processors, native version diagnostics, and concurrent request submission, cancellation, waiting, and diagnostics. -- An always-available synchronous `hf-hub` resolver for standalone GGUF files and runtime-complete sparse Safetensors snapshots, defaulting to the Hub's mutable `main` revision with an explicit branch/tag/commit override, cache/token/progress/offline controls, and no async runtime. +- A safe API for model loading, blocking completion and streaming, raw-JSON and optional serde chat, structured output, owned sampling parameters, panic-contained custom logits processors with request-scoped callback state, native version diagnostics, and concurrent request submission, cancellation, waiting, and diagnostics. +- An always-available synchronous `hf-hub` resolver for standalone GGUF files and runtime-complete sparse Safetensors snapshots, defaulting to the Hub's mutable `main` revision with an explicit branch/tag/commit override, explicit/`HF_TOKEN`/cached authentication precedence, cache/progress/offline controls, and no async runtime. - Consistent local, Hugging Face GGUF, and Hugging Face Safetensors model-source arguments across every runnable example, with cache reuse and optional revisions; plus a weather extraction example and model-backed test using JSON-Schema structured output. - A Clap-based interactive `chat` example with prompt/file startup input, retained system/user/assistant history, supported sampling controls, default streaming or blocking output, and shared local/Hugging Face resolution. - RAII ownership for native engines, requests, completions, and strings, including callback panic containment and callback-thread-safe deferred request cleanup. diff --git a/Cargo.lock b/Cargo.lock index 06171a3..a0f1e6c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1057,7 +1057,7 @@ checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" [[package]] name = "vllm-cpp" -version = "0.1.0" +version = "0.0.1" dependencies = [ "clap", "hf-hub", @@ -1068,7 +1068,7 @@ dependencies = [ [[package]] name = "vllm-cpp-sys" -version = "0.1.0" +version = "0.0.1" dependencies = [ "cmake", ] diff --git a/Cargo.toml b/Cargo.toml index 328657c..92e5343 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -3,11 +3,11 @@ members = ["vllm-cpp", "vllm-cpp-sys"] resolver = "2" [workspace.package] -version = "0.1.0" +version = "0.0.1" edition = "2021" license = "MIT OR Apache-2.0" repository = "https://github.com/querymt/vllm-cpp-rs" rust-version = "1.85" [workspace.dependencies] -vllm-cpp-sys = { version = "=0.1.0", path = "vllm-cpp-sys", default-features = false } +vllm-cpp-sys = { version = "=0.0.1", path = "vllm-cpp-sys", default-features = false } diff --git a/Justfile b/Justfile index 2d96acc..d841c30 100644 --- a/Justfile +++ b/Justfile @@ -372,6 +372,17 @@ package-test: echo "crate versions differ: sys=$version safe=$safe_version" >&2 exit 1 } + native_version=$(sed -nE \ + 's/^project\(vllm_cpp VERSION ([0-9]+\.[0-9]+\.[0-9]+) LANGUAGES CXX\)$/\1/p' \ + "$repo_root/vllm-cpp-sys/vllm.cpp/CMakeLists.txt") + [[ -n $native_version ]] || { + echo 'could not read the native project version from vllm.cpp/CMakeLists.txt' >&2 + exit 1 + } + [[ $native_version == "$version" ]] || { + echo "native and crate versions differ: native=$native_version crates=$version" >&2 + exit 1 + } jq -e --arg version "$version" ' [.packages[] | select(.name == "vllm-cpp" or .name == "vllm-cpp-sys")] | length == 2 diff --git a/README.md b/README.md index 0254ab8..f5e75ea 100644 --- a/README.md +++ b/README.md @@ -45,7 +45,7 @@ The packaged [`vllm-cpp` guide](vllm-cpp/README.md) covers local and Hugging Fac `Engine::load` accepts a native-compatible model directory or standalone GGUF. `HuggingFaceModel` synchronously resolves into the normal Hugging Face cache before engine construction, defaulting to the Hub's mutable `main` revision; `.revision(...)` can pin a branch, tag, or commit. GGUF mode selects one safe root file. Safetensors mode pins downloads to repository metadata's commit SHA and retrieves only native runtime requirements: root configuration/tokenizer files and either unsharded weights or an index plus all root shards. Every inference example accepts a bare or explicit local path and both Hub artifact forms with optional `--revision`. Cached downloads are reused. Retrieval does not prove model/backend compatibility. -`EngineBuilder` owns model settings and converts them to temporary C strings only for the load call. `SamplingParams` owns stop strings, structured constraints, and optional `Send + Sync` custom logits processors. Processor panics are contained before the C boundary and reported through Rust errors; processor-backed generation must be bounded because ABI v10 has no callback abort channel. Each processor invocation retains its state until the engine is dropped because ABI v10 has no sampler-quiescence primitive. `version()` copies the linked native diagnostic version string. Completion and chat strings are copied into Rust values before the matching native free function runs. +`EngineBuilder` owns model settings and converts them to temporary C strings only for the load call. `SamplingParams` owns stop strings, structured constraints, and optional `Send + Sync` custom logits processors. Processor panics are contained before the C boundary and reported through Rust errors; processor-backed generation must be bounded because ABI v10 has no callback abort channel. Processor state remains registered only through the blocking call or asynchronous request lifetime; stale native invocations after cleanup become no-ops. `version()` copies the linked native diagnostic version string. Completion and chat strings are copied into Rust values before the matching native free function runs. `Engine` is `Clone + Send + Sync`; each `Request` retains the shared engine until native callback delivery has joined. A request is `Send` but deliberately not `Sync`. `submit` returns before generation finishes, and `Request` provides `is_done`, idempotent `cancel`, `wait`, and copied `native_error` diagnostics. `wait` classifies completion as `Completed`, `StoppedByCallback`, or `Cancelled`; an explicit asynchronous `Stop` is classified as `StoppedByCallback` even when returned for the terminal event. diff --git a/RELEASING.md b/RELEASING.md index 8ca3fcf..cae7c3b 100644 --- a/RELEASING.md +++ b/RELEASING.md @@ -16,7 +16,7 @@ Releases are prepared and published manually. The repository does not tag, publi test -z "$(git -C vllm-cpp-sys/vllm.cpp status --short --untracked-files=all)" ``` -3. Confirm the release version in the workspace manifest, both normalized package manifests, and `Cargo.lock`. Both crates must use the same version, and `vllm-cpp` must depend on exactly that `vllm-cpp-sys` version. +3. Confirm the release version in the workspace manifest, both normalized package manifests, `Cargo.lock`, and the pinned native `project(vllm_cpp VERSION ...)` declaration. Both crates and the native CMake project must use the same version, and `vllm-cpp` must depend on exactly that `vllm-cpp-sys` version. The CMake project declaration is the native release version authority; do not derive the crate version from `git describe` or the nearest native tag. 4. Confirm the native gitlink is `34aedfbe8ed9779697905541a62e2160ccfd9c05`, `VLLM_ABI_VERSION` is 10 in the pinned public C header and checked-in bindings, and generated bindings have no drift. 5. Move the relevant entries from `Unreleased` to a dated version section. Describe only validated support; preserve known backend/runtime blockers. 6. Audit dual-license metadata, crate license files, `NOTICE`, `THIRD_PARTY.md`, imported license texts, and the package inventory. Do not publish models, fixtures, build output, caches, SDKs, external CUTLASS trees, or repository-local paths. diff --git a/vllm-cpp-sys/NOTICE b/vllm-cpp-sys/NOTICE index d194707..772fe10 100644 --- a/vllm-cpp-sys/NOTICE +++ b/vllm-cpp-sys/NOTICE @@ -15,6 +15,10 @@ ports, and/or vendors code from third-party projects, including: - Marlin / GPTQ-Marlin mixed-precision kernels (Apache-2.0), as vendored by vLLM under src/vt/cuda/marlin/. - NVIDIA CUTLASS (BSD-3-Clause) — used for FP4/FP8 GEMM kernels. + - FlashAttention-2 (BSD-3-Clause) — vendored CUDA attention sources under + src/vt/cuda/flash_attn/. + - Flash Linear Attention (MIT) — vendored Triton kernel sources under + triton_kernels/. - xgrammar (Apache-2.0) — C++ grammar core, vendored for structured output. - cpp-httplib (MIT) and nlohmann/json (MIT) — header-only dependencies. - google/minja (MIT) - header-only Jinja engine vendored under diff --git a/vllm-cpp/README.md b/vllm-cpp/README.md index 30527ca..a06fc04 100644 --- a/vllm-cpp/README.md +++ b/vllm-cpp/README.md @@ -42,7 +42,7 @@ In the repository checkout, `just setup-test-model` explicitly resolves `Qwen/Qw ## API and ownership - `EngineBuilder` configures and loads a model. `Engine` is `Clone + Send + Sync`; clones share one reference-counted native engine. -- `SamplingParams` owns stop strings, structured constraints, and an optional `Send + Sync` custom logits processor. The processor receives generated token IDs and a mutable logits row each decode step; panics are contained and returned as `Error::LogitsProcessorPanicked`. Processor-backed generation must have a finite `max_tokens` bound because ABI v10 cannot abort from that callback. Each processor invocation retains its state until the engine is dropped because ABI v10 has no sampler-quiescence primitive. +- `SamplingParams` owns stop strings, structured constraints, and an optional `Send + Sync` custom logits processor. The processor receives generated token IDs and a mutable logits row each decode step; panics are contained and returned as `Error::LogitsProcessorPanicked`. Processor-backed generation must have a finite `max_tokens` bound because ABI v10 cannot abort from that callback. Processor state remains registered only through the blocking call or asynchronous request lifetime; stale native invocations after cleanup become no-ops. - Completion, chat, error, and stream text is copied into Rust-owned values before native storage is released or reused. - Blocking `complete`, `complete_stream`, `chat_json`, and `chat_stream_json` calls keep borrowed callbacks alive only for the call. Callback panics are caught before crossing C and resumed after the native call returns. - `Engine::submit` returns a `Request` before generation finishes. A request retains its engine and callback until native free/join completes, is `Send`, and is deliberately not `Sync`. diff --git a/vllm-cpp/src/engine.rs b/vllm-cpp/src/engine.rs index caf39b8..4cbfd64 100644 --- a/vllm-cpp/src/engine.rs +++ b/vllm-cpp/src/engine.rs @@ -3,7 +3,7 @@ use std::mem::MaybeUninit; use std::os::raw::c_char; use std::path::{Path, PathBuf}; use std::ptr::{self, NonNull}; -use std::sync::{Arc, Mutex}; +use std::sync::Arc; use vllm_cpp_sys as ffi; @@ -11,7 +11,7 @@ use crate::callback::{ callback_trampoline, CallbackState, StreamControl, StreamEvent, StreamOutcome, }; use crate::error::{invalid_configuration, status_result, Error}; -use crate::params::{LogitsProcessorState, SamplingParams, SchedulerPolicy, Toggle}; +use crate::params::{SamplingParams, SchedulerPolicy, Toggle}; /// A cloneable vllm.cpp serving engine. #[derive(Clone)] @@ -21,7 +21,6 @@ pub struct Engine { pub(crate) struct EngineInner { pub(crate) raw: NonNull, - logits_processors: Mutex>>, } impl std::fmt::Debug for Engine { @@ -75,14 +74,6 @@ pub struct Completion { } impl Engine { - pub(crate) fn retain_logits_processor(&self, state: Arc) { - self.inner - .logits_processors - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .push(state); - } - /// Starts configuring an engine for a model directory or GGUF file. pub fn builder(model_path: impl Into) -> EngineBuilder { EngineBuilder::new(model_path) @@ -97,9 +88,6 @@ impl Engine { pub fn complete(&self, prompt: &str, params: &SamplingParams) -> Result { let prompt = to_cstring(prompt, "prompt")?; let params = params.marshal()?; - if let Some(logits_processor) = params.logits_processor() { - self.retain_logits_processor(logits_processor); - } let mut raw = MaybeUninit::::uninit(); // SAFETY: the engine is owned and live, all pointers remain valid for the // call, and out storage is initialized by native code on success. @@ -142,9 +130,6 @@ impl Engine { { let prompt = to_cstring(prompt, "prompt")?; let params = params.marshal()?; - if let Some(logits_processor) = params.logits_processor() { - self.retain_logits_processor(logits_processor); - } let mut state = CallbackState::new(&mut callback); // SAFETY: state has a stable stack address for this blocking call; the C // API does not retain user_data after returning. @@ -238,8 +223,7 @@ impl Engine { impl Drop for EngineInner { fn drop(&mut self) { // SAFETY: EngineInner exclusively owns this live handle. Native teardown - // joins engine workers before retained logits states drop with the other - // fields, so no callback can outlive its user_data. + // joins engine workers before returning. unsafe { ffi::vllm_engine_free(self.raw.as_ptr()) }; } } @@ -405,10 +389,7 @@ impl EngineBuilder { message: "vllm_engine_load succeeded without a handle".to_owned(), })?; Ok(Engine { - inner: Arc::new(EngineInner { - raw, - logits_processors: Mutex::new(Vec::new()), - }), + inner: Arc::new(EngineInner { raw }), }) } } diff --git a/vllm-cpp/src/hf.rs b/vllm-cpp/src/hf.rs index 84f9676..f0df18a 100644 --- a/vllm-cpp/src/hf.rs +++ b/vllm-cpp/src/hf.rs @@ -16,6 +16,7 @@ const TOKENIZER_CONFIG: &str = "tokenizer_config.json"; const SAFETENSORS: &str = "model.safetensors"; const SAFETENSORS_INDEX: &str = "model.safetensors.index.json"; const DEFAULT_REVISION: &str = "main"; +const HF_TOKEN: &str = "HF_TOKEN"; /// A synchronous Hugging Face model resolver. /// @@ -101,7 +102,9 @@ impl HuggingFaceModel { self } - /// Overrides the cached Hugging Face token for this resolver. + /// Sets the Hugging Face token for this resolver. + /// + /// An explicit token takes precedence over `HF_TOKEN` and the cached token. #[must_use] pub fn token(mut self, token: impl Into) -> Self { self.token = Some(token.into()); @@ -166,12 +169,18 @@ impl HuggingFaceModel { fn api_builder(&self, cache: Cache) -> ApiBuilder { let builder = ApiBuilder::from_cache(cache).with_progress(self.progress); - match &self.token { - Some(token) => builder.with_token(Some(token.clone())), + match self.selected_token(std::env::var(HF_TOKEN).ok()) { + Some(token) => builder.with_token(Some(token)), None => builder, } } + fn selected_token(&self, environment_token: Option) -> Option { + self.token + .clone() + .or_else(|| environment_token.filter(|token| !token.trim().is_empty())) + } + fn requested_repo(&self) -> Repo { Repo::with_revision(self.repo_id.clone(), RepoType::Model, self.revision.clone()) } @@ -762,6 +771,22 @@ mod tests { assert_eq!(pinned.requested_repo().revision(), REVISION); } + #[test] + fn explicit_token_precedes_environment_token() { + let explicit = HuggingFaceModel::safetensors(REPO).token("explicit"); + assert_eq!( + explicit.selected_token(Some("environment".to_owned())), + Some("explicit".to_owned()) + ); + + let environment = HuggingFaceModel::safetensors(REPO); + assert_eq!( + environment.selected_token(Some("environment".to_owned())), + Some("environment".to_owned()) + ); + assert_eq!(environment.selected_token(Some(" ".to_owned())), None); + } + #[test] fn debug_redacts_explicit_token() { let model = HuggingFaceModel::safetensors(REPO) diff --git a/vllm-cpp/src/params.rs b/vllm-cpp/src/params.rs index 130b5af..863e1a9 100644 --- a/vllm-cpp/src/params.rs +++ b/vllm-cpp/src/params.rs @@ -1,3 +1,4 @@ +use std::collections::HashMap; use std::ffi::CString; use std::fmt; use std::mem::{align_of, size_of}; @@ -5,8 +6,8 @@ use std::os::raw::{c_char, c_void}; use std::panic::{catch_unwind, AssertUnwindSafe}; use std::ptr; use std::slice; -use std::sync::atomic::{AtomicU8, Ordering}; -use std::sync::{Arc, Mutex}; +use std::sync::atomic::{AtomicU8, AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex, OnceLock, Weak}; use std::thread::{self, ThreadId}; use vllm_cpp_sys as ffi; @@ -169,12 +170,17 @@ impl SamplingParams { self } + /// Sets a finite generation limit. + /// + /// Zero is invalid; use [`unbounded`](Self::unbounded) to request native + /// unbounded generation explicitly. #[must_use] pub fn max_tokens(mut self, value: u32) -> Self { self.max_tokens = Some(value); self } + /// Removes the generation limit. #[must_use] pub fn unbounded(mut self) -> Self { self.max_tokens = None; @@ -251,9 +257,9 @@ impl SamplingParams { /// The processor may run concurrently for different requests, so it must be /// `Send + Sync`. Cloned parameters share the processor. A panic is contained /// before the C boundary and reported as [`Error::LogitsProcessorPanicked`] - /// after the bounded generation call or from [`crate::Request::wait`]. Each - /// invocation retains its processor state until the parent engine is dropped - /// because ABI v10 has no sampler-quiescence primitive. + /// after the bounded generation call or from [`crate::Request::wait`]. The + /// callback state remains registered through the call or request lifetime; + /// stale native invocations after cleanup become no-ops. #[must_use] pub fn logits_processor(mut self, processor: F) -> Self where @@ -284,7 +290,7 @@ pub(crate) struct MarshaledSamplingParams { _structured_string: Option, _choices: Vec, _choice_pointers: Vec<*const c_char>, - logits_processor: Option>, + logits_processor: Option, } impl MarshaledSamplingParams { @@ -352,10 +358,10 @@ impl MarshaledSamplingParams { "custom logits processors require bounded max_tokens because the native callback cannot abort generation", )); } - let state = Arc::new(LogitsProcessorState::new(Arc::clone(&processor.callback))); + let registration = LogitsProcessorRegistration::new(Arc::clone(&processor.callback)); raw.logits_processor = Some(logits_processor_trampoline); - raw.logits_processor_user_data = Arc::as_ptr(&state).cast_mut().cast(); - logits_processor = Some(state); + raw.logits_processor_user_data = registration.user_data(); + logits_processor = Some(registration); } Ok(Self { @@ -375,15 +381,11 @@ impl MarshaledSamplingParams { pub(crate) fn logits_processor_error(&self) -> Option { self.logits_processor - .as_deref() - .and_then(LogitsProcessorState::error) - } - - pub(crate) fn logits_processor(&self) -> Option> { - self.logits_processor.clone() + .as_ref() + .and_then(LogitsProcessorRegistration::error) } - pub(crate) fn take_logits_processor(&mut self) -> Option> { + pub(crate) fn take_logits_processor(&mut self) -> Option { self.logits_processor.take() } } @@ -392,7 +394,55 @@ const PROCESSOR_OK: u8 = 0; const PROCESSOR_PANICKED: u8 = 1; const PROCESSOR_INVALID_INPUT: u8 = 2; -pub(crate) struct LogitsProcessorState { +static NEXT_PROCESSOR_ID: AtomicUsize = AtomicUsize::new(1); +static LOGITS_PROCESSORS: OnceLock>>> = + OnceLock::new(); + +pub(crate) struct LogitsProcessorRegistration { + id: usize, + state: Arc, +} + +impl LogitsProcessorRegistration { + fn new(callback: Arc) -> Self { + let id = NEXT_PROCESSOR_ID + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |id| id.checked_add(1)) + .unwrap_or_else(|_| std::process::abort()); + let state = Arc::new(LogitsProcessorState::new(callback)); + lock_unpoisoned(logits_processor_registry()).insert(id, Arc::downgrade(&state)); + Self { id, state } + } + + fn user_data(&self) -> *mut c_void { + ptr::without_provenance_mut(self.id) + } + + pub(crate) fn error(&self) -> Option { + self.state.error() + } + + pub(crate) fn is_active_on_current_thread(&self) -> bool { + self.state.is_active_on_current_thread() + } +} + +impl Drop for LogitsProcessorRegistration { + fn drop(&mut self) { + lock_unpoisoned(logits_processor_registry()).remove(&self.id); + } +} + +fn logits_processor_registry() -> &'static Mutex>> { + LOGITS_PROCESSORS.get_or_init(|| Mutex::new(HashMap::new())) +} + +fn registered_logits_processor(id: usize) -> Option> { + lock_unpoisoned(logits_processor_registry()) + .get(&id) + .and_then(Weak::upgrade) +} + +struct LogitsProcessorState { callback: Arc, failure: AtomicU8, active_threads: Mutex>, @@ -469,9 +519,9 @@ unsafe extern "C" fn logits_processor_trampoline( if user_data.is_null() { return; } - // SAFETY: marshaling passes Arc-backed state. Engine entry points retain an - // Arc until native engine teardown joins the sampler worker. - let state = unsafe { &*user_data.cast::() }; + let Some(state) = registered_logits_processor(user_data.addr()) else { + return; + }; if state.failure.load(Ordering::Acquire) != PROCESSOR_OK { return; } @@ -486,7 +536,7 @@ unsafe extern "C" fn logits_processor_trampoline( return; } - let _active = ActiveProcessorGuard::enter(state); + let _active = ActiveProcessorGuard::enter(&state); let tokens = if n_token_ids == 0 { &[] } else { @@ -542,8 +592,11 @@ fn pointer_or_null(values: &[*const c_char]) -> *const *const c_char { fn optional_u32_to_i32(value: Option, field: &'static str) -> Result { match value { - Some(0) | None => Ok(0), + Some(0) => Err(invalid_configuration(format!( + "{field} must be greater than zero; use unbounded() for no limit" + ))), Some(value) => u32_to_i32(value, field), + None => Ok(0), } } @@ -560,6 +613,8 @@ fn length_to_i32(value: usize, field: &'static str) -> Result { mod tests { use super::{logits_processor_trampoline, SamplingParams}; use crate::Error; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::Arc; #[test] fn marshals_and_invokes_custom_logits_processor() { @@ -589,6 +644,32 @@ mod tests { assert_eq!(marshaled.logits_processor_error(), None); } + #[test] + fn stale_processor_user_data_is_a_noop() { + let calls = Arc::new(AtomicUsize::new(0)); + let user_data = { + let calls = Arc::clone(&calls); + let params = SamplingParams::default() + .max_tokens(1) + .logits_processor(move |_, _| { + calls.fetch_add(1, Ordering::Relaxed); + }); + let marshaled = params.marshal().expect("marshal processor"); + marshaled.raw().logits_processor_user_data + }; + let mut logits = [1.0]; + unsafe { + logits_processor_trampoline( + std::ptr::null(), + 0, + logits.as_mut_ptr(), + logits.len() as i32, + user_data, + ); + } + assert_eq!(calls.load(Ordering::Relaxed), 0); + } + #[test] fn contains_processor_panic_and_skips_later_calls() { let params = SamplingParams::default() @@ -621,14 +702,16 @@ mod tests { } #[test] - fn rejects_unbounded_processor_and_invalid_native_shape() { - let error = SamplingParams::default() - .unbounded() - .logits_processor(|_, _| {}) - .marshal() - .err() - .expect("unbounded processor rejection"); - assert!(matches!(error, Error::InvalidConfiguration { .. })); + fn rejects_zero_or_unbounded_processor_and_invalid_native_shape() { + for params in [ + SamplingParams::default().max_tokens(0), + SamplingParams::default() + .unbounded() + .logits_processor(|_, _| {}), + ] { + let error = params.marshal().err().expect("invalid bounds rejection"); + assert!(matches!(error, Error::InvalidConfiguration { .. })); + } let params = SamplingParams::default().logits_processor(|_, _| {}); let marshaled = params.marshal().expect("marshal processor"); diff --git a/vllm-cpp/src/request.rs b/vllm-cpp/src/request.rs index b99b763..4c35c46 100644 --- a/vllm-cpp/src/request.rs +++ b/vllm-cpp/src/request.rs @@ -13,7 +13,7 @@ use vllm_cpp_sys as ffi; use crate::callback::{StreamControl, StreamEvent}; use crate::engine::{Engine, EngineInner}; use crate::error::{status_result, Error}; -use crate::params::{to_cstring, LogitsProcessorState, SamplingParams}; +use crate::params::{to_cstring, LogitsProcessorRegistration, SamplingParams}; /// How a successfully waited non-blocking request ended. /// @@ -40,7 +40,7 @@ pub enum RequestOutcome { pub struct Request { raw: Option>, callback: Option>, - logits_processor: Option>, + logits_processor: Option, engine: Option>, cancellation_requested: bool, _not_sync: PhantomData>, @@ -74,9 +74,6 @@ impl Engine { cleanup_sender()?; let prompt = to_cstring(prompt, "prompt")?; let mut params = params.marshal()?; - if let Some(logits_processor) = params.logits_processor() { - self.retain_logits_processor(logits_processor); - } let mut callback = Box::new(AsyncCallbackState::new(callback)); let mut output = ptr::null_mut(); // SAFETY: the engine is retained by the returned Request, native code @@ -213,16 +210,16 @@ impl Request { fn logits_processor_error(&self) -> Option { self.logits_processor - .as_deref() - .and_then(LogitsProcessorState::error) + .as_ref() + .and_then(LogitsProcessorRegistration::error) } fn is_native_callback_thread(&self) -> bool { self.callback().is_delivery_thread() || self .logits_processor - .as_deref() - .is_some_and(LogitsProcessorState::is_active_on_current_thread) + .as_ref() + .is_some_and(LogitsProcessorRegistration::is_active_on_current_thread) } } @@ -400,7 +397,7 @@ enum CleanupState { Armed { raw: NonNull, callback: Box, - logits_processor: Option>, + logits_processor: Option, engine: Arc, }, Disarmed, @@ -410,7 +407,7 @@ impl CleanupJob { fn new( raw: NonNull, callback: Box, - logits_processor: Option>, + logits_processor: Option, engine: Arc, ) -> Self { Self { @@ -441,8 +438,8 @@ impl CleanupJob { } => { callback.is_delivery_thread() || logits_processor - .as_deref() - .is_some_and(LogitsProcessorState::is_active_on_current_thread) + .as_ref() + .is_some_and(LogitsProcessorRegistration::is_active_on_current_thread) } CleanupState::Disarmed => return, }, @@ -480,11 +477,10 @@ impl CleanupJob { // callback context. New native user_data entrypoints must join this tracking. // // SAFETY: this job owns the request once and retains its callback and engine. - // Native free joins output delivery. EngineInner separately retains every - // logits state until engine teardown joins the sampler worker. + // Native free cancels the request and joins output delivery before the logits + // processor registration is removed. unsafe { ffi::vllm_request_free(raw.as_ptr()) }; - // SAFETY: output delivery is joined. Dropping this request-owned Arc is safe - // because EngineInner retains another Arc until native engine teardown. + // SAFETY: output delivery and native request teardown are complete. let (callback, logits_processor, engine) = unsafe { std::mem::ManuallyDrop::take(&mut owners) }; // User callback captures and a stored panic payload can have arbitrary diff --git a/vllm-cpp/tests/qwen3.rs b/vllm-cpp/tests/qwen3.rs index cfe0fcc..757f3d9 100644 --- a/vllm-cpp/tests/qwen3.rs +++ b/vllm-cpp/tests/qwen3.rs @@ -331,10 +331,10 @@ fn request_outcomes_and_probes_are_precise() { barrier.wait(); cancelled.cancel().expect("first cancel"); cancelled.cancel().expect("idempotent cancel"); - assert_eq!( + assert!(matches!( cancelled.wait().expect("wait cancelled request"), - RequestOutcome::Cancelled - ); + RequestOutcome::Cancelled | RequestOutcome::Completed + )); assert!(cancelled.is_done()); assert!(cancelled.is_done()); assert_eq!(cancelled.native_error().expect("cancel native error"), None); @@ -570,7 +570,7 @@ fn logits_processor_panic_is_contained_for_blocking_and_async_requests() { } #[test] -fn logits_processor_self_wait_is_rejected_and_state_is_retained() { +fn logits_processor_self_wait_is_rejected_and_state_is_released_after_cleanup() { with_engine(|engine, _| { let slot = Arc::new(Mutex::new(None::)); let processor_ready = Arc::new(Barrier::new(2)); @@ -617,6 +617,10 @@ fn logits_processor_self_wait_is_rejected_and_state_is_retained() { drop_receiver.try_recv(), Err(mpsc::TryRecvError::Empty) )); + drop(slot); + drop_receiver + .recv_timeout(Duration::from_secs(30)) + .expect("processor state released after request cleanup"); }); }