diff --git a/Cargo.lock b/Cargo.lock index 850c76fdc..1b5c3f5c8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4305,6 +4305,7 @@ dependencies = [ "libloading", "mesh-native-serving-plugin-api", "skippy-server", + "skippy-tokenizer", ] [[package]] diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/native_mtp_tests.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/native_mtp_tests.rs index 172acf54d..880624e08 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/native_mtp_tests.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/native_mtp_tests.rs @@ -1,12 +1,91 @@ +use super::super::{SkippyModelHandle, SkippyModelLoadOptions}; use super::test_support::*; use super::*; use crate::inference::skippy::SkippyTelemetryOptions; +use anyhow::Result; +use openai_frontend::OpenAiBackend; use skippy_protocol::LoadMode; use skippy_runtime::package::{ PackageExtensionPolicyInfo, PackageGenerationInfo, PackageSpeculativeDecodingInfo, PackageSpeculativeProposerInfo, PackageSpeculativeStrategyInfo, PackageWindowPolicyInfo, }; -use std::collections::BTreeMap; +use std::{ + collections::BTreeMap, + path::Path, + sync::{ + Arc, Mutex, + atomic::{AtomicUsize, Ordering}, + }, + time::Duration, +}; + +#[derive(Default)] +struct RecordingNativeProposalIngress { + proposals: AtomicUsize, + reports: Mutex>, +} + +impl skippy_server::LinearProposalIngress for RecordingNativeProposalIngress { + fn propose( + &self, + _query: skippy_server::LinearProposalQuery, + ) -> anyhow::Result { + self.proposals.fetch_add(1, Ordering::Relaxed); + let decision_id = skippy_server::OpaqueProposalDecisionId::new(vec![1])?; + Ok(skippy_server::LinearProposalSourceResponse::new(Some( + skippy_server::LinearProposal::new(decision_id, vec![0]), + ))) + } + + fn report(&self, receipt: &skippy_server::LinearProposalReceipt) -> anyhow::Result<()> { + self.reports.lock().unwrap().push(receipt.clone()); + Ok(()) + } +} + +struct NoopGenerationReceiptSink; + +impl skippy_server::frontend::GenerationReceiptSink for NoopGenerationReceiptSink { + fn begin(&self, _start: &skippy_server::frontend::GenerationStart) -> Result<()> { + Ok(()) + } + + fn committed(&self, _commit: &skippy_server::frontend::GenerationCommit) -> Result<()> { + Ok(()) + } + + fn abort(&self, _abort: &skippy_server::frontend::GenerationAbort) -> Result<()> { + Ok(()) + } + + fn record(&self, _receipt: &skippy_server::frontend::GenerationReceipt) -> Result<()> { + Ok(()) + } +} + +struct RecordingNativeHooksFactory { + ingress: Arc, +} + +impl skippy_server::serving_hooks::ModelServingHooksFactory for RecordingNativeHooksFactory { + fn create( + &self, + _tokenizer: skippy_server::TokenizerCapability, + ) -> Result { + let source: Arc = self.ingress.clone(); + let ingress = skippy_server::frontend::LinearProposalIngressConfig::new( + source, + Duration::from_millis(25), + 1, + )?; + Ok(skippy_server::serving_hooks::ModelServingHooks::new( + skippy_server::frontend::GenerationReceiptConfig::new(Arc::new( + NoopGenerationReceiptSink, + )), + ingress, + )) + } +} fn native_mtp_generation() -> PackageGenerationInfo { let mut proposers = BTreeMap::new(); @@ -761,7 +840,7 @@ strategy = "ngram-suffix" } #[test] -fn standalone_ngram_rejected_for_single_stage_serving() { +fn standalone_ngram_uses_native_plugin_verify_path_for_single_stage_serving() { let mesh_config = parse_config( r#" [defaults.speculative] @@ -789,17 +868,102 @@ ngram_max_proposal_tokens = 48 .to_embedded_openai_args(4096, true) .expect("staged serving should build OpenAI args"); - // Single-stage/direct serving has no N-gram verify path, so it must reject - // rather than silently run target-only while reporting a proposer. - let error = resolved + // Single-stage/direct serving uses the native plugin ingress, whose + // proposal is verified and repaired by the local target runtime. + let openai = resolved .to_embedded_openai_args(0, false) - .expect_err("single-stage standalone N-gram must be rejected"); + .expect("single-stage standalone N-gram should use the native plugin verify path"); + assert_eq!(openai.speculative_window, 48); + assert_eq!( + openai.speculative.ngram.as_ref().map(|ngram| ngram.kind), + Some(skippy_server::NgramProposerKind::Suffix) + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn direct_single_stage_serving_delivers_target_authoritative_native_receipts() -> Result<()> { + let Some(model_path) = std::env::var_os("SKIPPY_NATIVE_PLUGIN_MODEL") else { + eprintln!("skipping: SKIPPY_NATIVE_PLUGIN_MODEL is not set"); + return Ok(()); + }; + let model_path = Path::new(&model_path); + let model_bytes = std::fs::metadata(model_path)?.len(); + let model_id = "Qwen/Qwen3-0.6B:Q4_K_M"; + let mesh_config = parse_config( + r#" +[defaults.speculative] +strategy = "ngram-suffix" +ngram_min = 2 +ngram_max = 8 +ngram_max_proposal_tokens = 1 +"#, + ); + let resolved = resolve_skippy_config(SkippyConfigResolveRequest { + mesh_config: &mesh_config, + model_id, + model_path, + model_bytes, + allocatable_memory_bytes: None, + request_defaults: None, + package_generation: None, + })?; + let embedded_openai = resolved.to_embedded_openai_args(0, false)?; + let ingress = Arc::new(RecordingNativeProposalIngress::default()); + let factory: skippy_server::serving_hooks::SharedModelServingHooksFactory = + Arc::new(RecordingNativeHooksFactory { + ingress: Arc::clone(&ingress), + }); + let mut options = SkippyModelLoadOptions::for_direct_gguf(model_id, model_path) + .with_ctx_size(256) + .with_embedded_openai(embedded_openai) + .with_serving_hooks_factory(Some(factory)); + options.n_gpu_layers = 0; + + let handle = SkippyModelHandle::load_with_hooks( + options, + None, + crate::runtime::survey::SurveyTelemetry::disabled(), + )?; + let request = serde_json::from_value(serde_json::json!({ + "model": model_id, + "messages": [{"role": "user", "content": "Say hello."}], + "max_tokens": 2, + "temperature": 0.0 + }))?; + let response = handle.chat_completion(request).await; + handle.shutdown(); + response?; + + for _ in 0..100 { + if !ingress.reports.lock().unwrap().is_empty() { + break; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + assert!(ingress.proposals.load(Ordering::Relaxed) > 0); + let reports = ingress.reports.lock().unwrap(); assert!( - error - .to_string() - .contains("requires multi-stage split serving"), - "{error}" + !reports.is_empty(), + "native proposal report was not delivered" ); + for receipt in reports.iter() { + assert_eq!(receipt.proposal_token_count, 1); + assert!(receipt.verification_rows > 0); + assert!(receipt.accepted_proposal_tokens <= receipt.proposal_token_count); + assert!(!receipt.committed_tokens.is_empty()); + assert_eq!( + receipt.canonical_prediction_count, + receipt.committed_tokens.len() + ); + assert_eq!( + receipt.verification_rows, + receipt.verification_row_predictions.len() + ); + assert!(receipt.canonical_prediction_count <= receipt.verification_rows); + assert!(receipt.canonical_position >= receipt.base_position); + assert!(receipt.canonical_position <= receipt.position_after_verification); + } + Ok(()) } #[test] diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/translation.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/translation.rs index 59c273b84..f93311190 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/translation.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/translation.rs @@ -280,27 +280,14 @@ impl ResolvedSkippyConfig { if self.skippy.prefill_controls_explicit { bail!("skippy prefill chunk controls require staged serving"); } - // A stage with no downstream delegates to generate_local_tokens, which - // has no N-gram verification path. Reject a standalone N-gram plan here - // rather than silently running target-only while reporting a proposer. - if self.speculative.decode.ngram.is_some() - && !self.speculative.decode.native_mtp.enabled - { - bail!( - "standalone N-gram speculation ({}) requires multi-stage split serving; \ - single-stage and direct GGUF requests have no N-gram verification path", - self.speculative.decode.effective_strategy - ); - } } Ok(()) } - fn speculative_mode_for_embedded(&self, staged: bool) -> &'static str { + fn speculative_mode_for_embedded(&self, _staged: bool) -> &'static str { if self.speculative.mode == "draft" && self.speculative.draft_model_path.is_some() { "draft" - } else if staged - && self.speculative.decode.ngram.is_some() + } else if self.speculative.decode.ngram.is_some() && !self.speculative.decode.native_mtp.enabled { "ngram" diff --git a/crates/mesh-native-serving-plugin-api/README.md b/crates/mesh-native-serving-plugin-api/README.md index 44cdda055..54eb871d8 100644 --- a/crates/mesh-native-serving-plugin-api/README.md +++ b/crates/mesh-native-serving-plugin-api/README.md @@ -14,9 +14,27 @@ layout of the prefix understood by an older plugin. Removing or reordering fields, changing a field's type, or changing a callback signature requires a new versioned ABI table and continued support for the previous table. -The V1 activation contract lends the complete native tokenizer inventory through -`ActivationContext`. A plugin copies or transforms the inventory before -returning from `activate`. +The V2 activation contract lends a model-bound tokenizer capability through +`ActivationContext`. The capability includes the opaque native tokenizer +inventory, the model identity and binding digest, explicit input/output +limits, and a bounded encode callback for tagged ordinary-byte and opaque +control pieces. Mesh does not know or own Rosetta vocabulary: a consumer +builds any translation from the inventory and calls the capability outside the +proposal deadline. Unsupported or non-lossless input is rejected; it is never +decoded with replacement semantics. A plugin copies or transforms the +inventory before returning from `activate` and must not call the encode +callback from proposal callbacks. + +The V2 table is a breaking replacement for V1. Plugins must resolve +`mesh_native_serving_plugin_v2` and validate the V2 ABI/version and structure +sizes before use. + +For `TokenizerCapability::encode`, a non-null `output_length` is required. The +host initializes it to zero before validating the remaining arguments, so every +non-OK return other than `OUTPUT_TOO_SMALL` reports zero tokens; on +`OUTPUT_TOO_SMALL` it reports the required capacity, and on `OK` it reports the +number of written tokens. A null `output_length` cannot be written and therefore +returns `INVALID_ARGUMENT`. Each proposal query carries the capacity Skippy can verify at that exact decode position. The native host adapter also applies a bounded implementation cap so diff --git a/crates/mesh-native-serving-plugin-api/src/lib.rs b/crates/mesh-native-serving-plugin-api/src/lib.rs index 39eeb3231..318d1615f 100644 --- a/crates/mesh-native-serving-plugin-api/src/lib.rs +++ b/crates/mesh-native-serving-plugin-api/src/lib.rs @@ -14,10 +14,13 @@ use std::ffi::{c_char, c_void}; -pub const NATIVE_SERVING_PLUGIN_ABI_V1: u32 = 1; -pub const NATIVE_SERVING_PLUGIN_ENTRY_V1: &[u8] = b"mesh_native_serving_plugin_v1\0"; +pub const NATIVE_SERVING_PLUGIN_ABI_V2: u32 = 2; +pub const NATIVE_SERVING_PLUGIN_ENTRY_V2: &[u8] = b"mesh_native_serving_plugin_v2\0"; pub const MAX_DECISION_ID_BYTES: usize = 64; pub const TOKENIZER_INVENTORY_SCHEMA: u32 = 1; +pub const TOKENIZER_CAPABILITY_ABI: u32 = 1; +pub const MAX_TOKENIZER_INVENTORY_ENTRIES: usize = 1_000_000; +pub const MAX_TOKENIZER_INPUT_PIECES: usize = 4_096; /// Host-owned typed inventory. This Rust value never crosses the ABI directly. #[derive(Clone, Debug, Eq, PartialEq)] @@ -37,8 +40,14 @@ pub struct TokenizerInventoryToken { #[derive(Clone, Debug, Eq, PartialEq)] pub enum TokenizerInventoryPiece { - Bytes { bytes: Vec }, - Control { identity: String }, + Bytes { + bytes: Vec, + }, + /// Opaque bytes for a native special-token descriptor. Mesh never parses + /// or names this value; the plugin owns its interpretation. + Control { + descriptor: Vec, + }, } pub type PluginInstance = *mut c_void; @@ -92,6 +101,70 @@ pub struct TokenizerInventoryView { pub entry_count: usize, } +#[repr(C)] +#[derive(Clone, Copy, Debug, Default)] +pub struct TokenizerLimits { + pub max_input_bytes: usize, + pub max_output_tokens: usize, +} + +#[repr(transparent)] +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct TokenizerInputPieceKind(pub u32); + +impl TokenizerInputPieceKind { + pub const BYTES: Self = Self(0); + pub const CONTROL: Self = Self(1); +} + +impl Default for TokenizerInputPieceKind { + fn default() -> Self { + Self::BYTES + } +} + +#[repr(C)] +#[derive(Clone, Copy, Debug, Default)] +pub struct TokenizerInputPiece { + pub kind: TokenizerInputPieceKind, + pub bytes: ByteSlice, +} + +#[repr(transparent)] +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct TokenizerEncodeStatus(pub u32); + +impl TokenizerEncodeStatus { + pub const OK: Self = Self(0); + pub const INVALID_ARGUMENT: Self = Self(1); + pub const UNSUPPORTED_INPUT: Self = Self(2); + pub const OUTPUT_TOO_SMALL: Self = Self(3); + pub const LIMIT_EXCEEDED: Self = Self(4); + pub const UNAVAILABLE: Self = Self(5); + pub const INTERNAL_ERROR: Self = Self(6); +} + +/// A model-bound, host-owned tokenizer capability. The table and its +/// inventory view are lent only for `activate`; the `context` passed to +/// `encode` remains valid until plugin shutdown. Plugins must copy any +/// activation data they need and must perform preparation/encoding outside the +/// latency-sensitive proposal callbacks. Mesh never invokes `encode` from the +/// proposal path. +#[repr(C)] +#[derive(Clone, Copy, Debug)] +pub struct TokenizerCapability { + pub struct_size: usize, + pub abi_version: u32, + pub model_id: ByteSlice, + pub source_model_sha256: ByteSlice, + pub tokenizer_id: ByteSlice, + pub limits: TokenizerLimits, + pub binding_digest: [u8; 32], + pub inventory: *const TokenizerInventoryView, + pub context: *mut c_void, + pub encode: EncodeTokenizer, +} + #[repr(C)] #[derive(Clone, Copy, Debug, Default)] pub struct TokenSlice { @@ -183,6 +256,15 @@ impl ProposalDiscardReason { pub type MonotonicNowNs = unsafe extern "C" fn(context: *mut c_void) -> u64; +pub type EncodeTokenizer = unsafe extern "C" fn( + context: *mut c_void, + input_pieces: *const TokenizerInputPiece, + input_piece_count: usize, + output_tokens: *mut i32, + output_capacity: usize, + output_length: *mut usize, +) -> TokenizerEncodeStatus; + #[repr(C)] #[derive(Clone, Copy)] pub struct ActivationContext { @@ -190,7 +272,7 @@ pub struct ActivationContext { pub model_id: ByteSlice, pub source_model_sha256: ByteSlice, pub tokenizer_id: ByteSlice, - pub tokenizer_inventory: *const TokenizerInventoryView, + pub tokenizer_capability: *const TokenizerCapability, pub config_path: ByteSlice, pub state_directory: ByteSlice, pub proposal_deadline_ns: u64, @@ -336,7 +418,7 @@ pub type LastError = unsafe extern "C" fn(instance: PluginInstance, output: *mut c_char, capacity: usize) -> usize; #[repr(C)] -pub struct NativeServingPluginV1 { +pub struct NativeServingPluginV2 { pub abi_version: u32, pub struct_size: usize, pub plugin_name: ByteSlice, @@ -358,9 +440,9 @@ pub struct NativeServingPluginV1 { // `plugin_name` to remain immutable and valid for the loaded library's entire // lifetime. The host copies the name during load and only calls function // pointers afterward. -unsafe impl Sync for NativeServingPluginV1 {} +unsafe impl Sync for NativeServingPluginV2 {} -pub type NativeServingPluginEntryV1 = unsafe extern "C" fn() -> *const NativeServingPluginV1; +pub type NativeServingPluginEntryV2 = unsafe extern "C" fn() -> *const NativeServingPluginV2; #[cfg(test)] mod tests { @@ -379,8 +461,22 @@ mod tests { } #[test] - fn initial_contract_is_v1() { + fn initial_contract_is_v2() { assert_eq!(MAX_DECISION_ID_BYTES, 64); - assert_eq!(NATIVE_SERVING_PLUGIN_ABI_V1, 1); + assert_eq!(NATIVE_SERVING_PLUGIN_ABI_V2, 2); + assert_eq!(TOKENIZER_CAPABILITY_ABI, 1); + } + + #[test] + fn structured_input_piece_kinds_are_stable_and_opaque() { + assert_eq!(TokenizerInputPieceKind::BYTES.0, 0); + assert_eq!(TokenizerInputPieceKind::CONTROL.0, 1); + let descriptor = [0xff, 0x00]; + let piece = TokenizerInputPiece { + kind: TokenizerInputPieceKind::CONTROL, + bytes: ByteSlice::from_bytes(&descriptor), + }; + assert_eq!(piece.bytes.length, 2); + assert_eq!(piece.bytes.pointer, descriptor.as_ptr()); } } diff --git a/crates/mesh-native-serving-plugin-host/Cargo.toml b/crates/mesh-native-serving-plugin-host/Cargo.toml index 7b74c9c62..8955a39e5 100644 --- a/crates/mesh-native-serving-plugin-host/Cargo.toml +++ b/crates/mesh-native-serving-plugin-host/Cargo.toml @@ -13,6 +13,7 @@ anyhow.workspace = true libloading = "0.8" mesh-native-serving-plugin-api = { path = "../mesh-native-serving-plugin-api", version = "0.72.1" } skippy-server = { path = "../skippy-server", version = "0.72.1" } +skippy-tokenizer = { path = "../skippy-tokenizer", version = "0.72.1" } [lints] workspace = true diff --git a/crates/mesh-native-serving-plugin-host/README.md b/crates/mesh-native-serving-plugin-host/README.md index 2f8f3d82b..1ef18b4b1 100644 --- a/crates/mesh-native-serving-plugin-host/README.md +++ b/crates/mesh-native-serving-plugin-host/README.md @@ -6,6 +6,10 @@ Skippy's decode thread waits only until its own absolute deadline and never joins or drains plugin work. Lifecycle and authoritative verification outcomes remain ordered through the same bounded queue. -At activation, the host passes a borrowed view of the model's already-bound -native tokenizer inventory over the plugin ABI. The view and its byte slices -are valid only for that call; plugins must not retain them. +At activation, the host passes a model-bound tokenizer capability over the +plugin ABI. Its inventory view and input-piece slices are borrowed, while the +capability callback context remains owned by the active plugin host until +shutdown. The host bounds and copies structured input before calling the +loaded tokenizer. Opaque controls are never decoded as text: unsupported +controls return `UNSUPPORTED_INPUT`. Tokenizer calls are preparation-time +operations and are not part of proposal dispatch. diff --git a/crates/mesh-native-serving-plugin-host/src/lib.rs b/crates/mesh-native-serving-plugin-host/src/lib.rs index 4dabeb67a..ecdc15c68 100644 --- a/crates/mesh-native-serving-plugin-host/src/lib.rs +++ b/crates/mesh-native-serving-plugin-host/src/lib.rs @@ -3,6 +3,7 @@ mod plugin_dispatch; #[cfg(test)] mod test_support; +mod tokenizer_capability; use std::{ collections::HashMap, @@ -17,6 +18,7 @@ use std::{ use anyhow::{Context, Result, anyhow, bail}; use libloading::Library; use mesh_native_serving_plugin_api as abi; +use plugin_dispatch::{PluginCommand, PluginDriver}; use skippy_server::frontend::{ GenerationAbort, GenerationCommit, GenerationLifecycleIngress, GenerationLifecycleObservation, GenerationReceipt, GenerationReceiptConfig, GenerationStart, LinearProposal, @@ -26,13 +28,17 @@ use skippy_server::frontend::{ }; use skippy_server::serving_hooks::{ModelServingHooks, ModelServingHooksFactory}; use skippy_server::tokenizer::TokenizerCapability; - -use plugin_dispatch::{PluginCommand, PluginDriver}; +use tokenizer_capability::HostTokenizerCapability; const ERROR_BUFFER_BYTES: usize = 2_048; const MAX_NATIVE_PLUGIN_PROPOSAL_TOKENS: usize = 4_096; const PROPOSAL_POLL_INTERVAL: Duration = Duration::from_micros(50); +const _: () = assert!( + abi::MAX_TOKENIZER_INPUT_PIECES == skippy_tokenizer::MAX_TOKENIZE_PIECES, + "ABI and tokenizer piece bounds must match", +); + /// Mesh-owned factory for one independently built native serving plugin. #[derive(Clone)] pub struct NativeServingPluginFactory { @@ -86,13 +92,8 @@ impl NativeServingPluginFactory { impl ModelServingHooksFactory for NativeServingPluginFactory { fn create(&self, tokenizer: TokenizerCapability) -> Result { - let identity = tokenizer.identity(); - let inventory = ActivationInventory::from_inventory( - tokenizer - .inventory() - .context("bound model does not expose a tokenizer inventory")?, - )?; - let inventory_view = inventory.view(); + let tokenizer_capability = HostTokenizerCapability::new(tokenizer)?; + let identity = tokenizer_capability.tokenizer.identity(); let context = abi::ActivationContext { struct_size: size_of::(), model_id: abi::ByteSlice::from_bytes(identity.model_id.as_bytes()), @@ -100,7 +101,7 @@ impl ModelServingHooksFactory for NativeServingPluginFactory { identity.source_model_sha256.as_bytes(), ), tokenizer_id: abi::ByteSlice::from_bytes(identity.tokenizer_id.as_bytes()), - tokenizer_inventory: &raw const inventory_view, + tokenizer_capability: &tokenizer_capability.abi, config_path: path_slice(&self.config_path), state_directory: path_slice(&self.state_directory), proposal_deadline_ns: u64::try_from(self.proposal_deadline.as_nanos()) @@ -125,6 +126,7 @@ impl ModelServingHooksFactory for NativeServingPluginFactory { let active = ActivePlugin { definition: Arc::clone(&self.definition), instance: Some(instance), + _tokenizer_capability: Some(tokenizer_capability), proposal_token_buffer: Mutex::new(vec![0; MAX_NATIVE_PLUGIN_PROPOSAL_TOKENS]), committed_generated_tokens: Mutex::new(HashMap::new()), }; @@ -144,52 +146,9 @@ impl ModelServingHooksFactory for NativeServingPluginFactory { } } -struct ActivationInventory { - entries: Vec, -} - -impl ActivationInventory { - fn from_inventory(inventory: &abi::TokenizerInventory) -> Result { - if inventory.schema_version != abi::TOKENIZER_INVENTORY_SCHEMA - || inventory.tokens.is_empty() - { - bail!("bound model exposes an unsupported or empty tokenizer inventory"); - } - let entries = inventory - .tokens - .iter() - .map(|entry| { - let (piece_kind, bytes) = match &entry.piece { - abi::TokenizerInventoryPiece::Bytes { bytes } => { - (abi::TokenizerPieceKind::BYTES, bytes.as_slice()) - } - abi::TokenizerInventoryPiece::Control { identity } => { - (abi::TokenizerPieceKind::CONTROL, identity.as_bytes()) - } - }; - abi::TokenizerInventoryEntry { - id: entry.id, - piece_kind, - bytes: abi::ByteSlice::from_bytes(bytes), - } - }) - .collect(); - Ok(Self { entries }) - } - - fn view(&self) -> abi::TokenizerInventoryView { - abi::TokenizerInventoryView { - struct_size: size_of::(), - schema_version: abi::TOKENIZER_INVENTORY_SCHEMA, - entries: self.entries.as_ptr(), - entry_count: self.entries.len(), - } - } -} - struct LoadedDefinition { _library: Option, - api: NonNull, + api: NonNull, name: String, } @@ -204,7 +163,7 @@ impl LoadedDefinition { let library = unsafe { Library::new(path) } .with_context(|| format!("load native serving plugin {}", path.display()))?; let entry = unsafe { - library.get::(abi::NATIVE_SERVING_PLUGIN_ENTRY_V1) + library.get::(abi::NATIVE_SERVING_PLUGIN_ENTRY_V2) } .with_context(|| { format!( @@ -222,7 +181,7 @@ impl LoadedDefinition { }) } - fn api(&self) -> &abi::NativeServingPluginV1 { + fn api(&self) -> &abi::NativeServingPluginV2 { unsafe { self.api.as_ref() } } @@ -246,19 +205,19 @@ impl LoadedDefinition { } } -fn validate_table(table: &abi::NativeServingPluginV1) -> Result { - if table.abi_version != abi::NATIVE_SERVING_PLUGIN_ABI_V1 { +fn validate_table(table: &abi::NativeServingPluginV2) -> Result { + if table.abi_version != abi::NATIVE_SERVING_PLUGIN_ABI_V2 { bail!( "native serving plugin ABI {} is incompatible with host ABI {}", table.abi_version, - abi::NATIVE_SERVING_PLUGIN_ABI_V1 + abi::NATIVE_SERVING_PLUGIN_ABI_V2 ); } - if table.struct_size != size_of::() { + if table.struct_size != size_of::() { bail!( "native serving plugin table size {} does not match host size {}", table.struct_size, - size_of::() + size_of::() ); } let name = unsafe { read_utf8(table.plugin_name, "plugin name") }?; @@ -271,6 +230,7 @@ fn validate_table(table: &abi::NativeServingPluginV1) -> Result { struct ActivePlugin { definition: Arc, instance: Option>, + _tokenizer_capability: Option>, proposal_token_buffer: Mutex>, committed_generated_tokens: Mutex>, } @@ -840,32 +800,6 @@ mod tests { use super::*; use crate::test_support::fake_table; - #[test] - fn tokenizer_inventory_view_borrows_host_owned_bytes_for_activation() { - let inventory = abi::TokenizerInventory { - schema_version: abi::TOKENIZER_INVENTORY_SCHEMA, - model_id: "glm".to_string(), - source_model_sha256: "a".repeat(64), - tokenizer_id: "gguf-source-sha256:test".to_string(), - tokens: vec![abi::TokenizerInventoryToken { - id: 0, - piece: abi::TokenizerInventoryPiece::Bytes { - bytes: b"hello".to_vec(), - }, - }], - }; - let activation = ActivationInventory::from_inventory(&inventory).unwrap(); - let view = activation.view(); - assert_eq!(view.schema_version, abi::TOKENIZER_INVENTORY_SCHEMA); - assert_eq!(view.entry_count, 1); - let entry = unsafe { &*view.entries }; - assert_eq!(entry.piece_kind, abi::TokenizerPieceKind::BYTES); - assert_eq!( - unsafe { std::slice::from_raw_parts(entry.bytes.pointer, entry.bytes.length) }, - b"hello" - ); - } - #[test] fn output_validation_is_fail_closed() { let decision = [1_u8; abi::MAX_DECISION_ID_BYTES]; @@ -908,7 +842,7 @@ mod tests { .to_string() .contains("incompatible") ); - table.abi_version = abi::NATIVE_SERVING_PLUGIN_ABI_V1; + table.abi_version = abi::NATIVE_SERVING_PLUGIN_ABI_V2; table.struct_size -= 1; assert!( validate_table(&table) diff --git a/crates/mesh-native-serving-plugin-host/src/plugin_dispatch.rs b/crates/mesh-native-serving-plugin-host/src/plugin_dispatch.rs index 8913ca655..fecc49f08 100644 --- a/crates/mesh-native-serving-plugin-host/src/plugin_dispatch.rs +++ b/crates/mesh-native-serving-plugin-host/src/plugin_dispatch.rs @@ -436,11 +436,11 @@ fn plugin_worker( command, }) = queue.next() { - let (result, lifecycle) = match command { - PluginCommand::Begin(event) => (active.begin(&event), true), - PluginCommand::Committed(event) => (active.committed(&event), true), - PluginCommand::Abort(event) => (active.abort(&event), true), - PluginCommand::Finish(event) => (active.finish(&event), true), + let result = match command { + PluginCommand::Begin(event) => active.begin(&event), + PluginCommand::Committed(event) => active.committed(&event), + PluginCommand::Abort(event) => active.abort(&event), + PluginCommand::Finish(event) => active.finish(&event), PluginCommand::Proposal(query, reply) => { run_proposal(&active, &passive_queue, enqueued_at, query, &reply); continue; @@ -449,7 +449,8 @@ fn plugin_worker( unreachable!("passive plugin callbacks must use the passive worker queue") } }; - if lifecycle && result.is_err() { + if let Err(error) = &result { + eprintln!("native serving plugin lifecycle callback failed: {error:#}"); lifecycle_delivery_failures.fetch_add(1, Ordering::Relaxed); } } diff --git a/crates/mesh-native-serving-plugin-host/src/test_support.rs b/crates/mesh-native-serving-plugin-host/src/test_support.rs index a420dd69e..5aee87556 100644 --- a/crates/mesh-native-serving-plugin-host/src/test_support.rs +++ b/crates/mesh-native-serving-plugin-host/src/test_support.rs @@ -168,10 +168,10 @@ unsafe extern "C" fn fake_last_error( 0 } -pub(crate) fn fake_table() -> abi::NativeServingPluginV1 { - abi::NativeServingPluginV1 { - abi_version: abi::NATIVE_SERVING_PLUGIN_ABI_V1, - struct_size: size_of::(), +pub(crate) fn fake_table() -> abi::NativeServingPluginV2 { + abi::NativeServingPluginV2 { + abi_version: abi::NATIVE_SERVING_PLUGIN_ABI_V2, + struct_size: size_of::(), plugin_name: abi::ByteSlice::from_bytes(FAKE_NAME), activate: fake_activate, shutdown: fake_shutdown, @@ -257,6 +257,7 @@ pub(crate) fn fake_active_with_timing( ActivePlugin { definition, instance: NonNull::new(Box::into_raw(state).cast::()), + _tokenizer_capability: None, proposal_token_buffer: Mutex::new(vec![0; MAX_NATIVE_PLUGIN_PROPOSAL_TOKENS]), committed_generated_tokens: Mutex::new(HashMap::new()), }, diff --git a/crates/mesh-native-serving-plugin-host/src/tokenizer_capability.rs b/crates/mesh-native-serving-plugin-host/src/tokenizer_capability.rs new file mode 100644 index 000000000..8dd8fa60a --- /dev/null +++ b/crates/mesh-native-serving-plugin-host/src/tokenizer_capability.rs @@ -0,0 +1,290 @@ +use std::{ffi::c_void, mem::size_of, sync::Arc}; + +use anyhow::{Context, Result, anyhow, bail}; +use mesh_native_serving_plugin_api as abi; +use skippy_server::tokenizer::TokenizerCapability; +use skippy_tokenizer::{EncodeRequest, InputPiece}; + +pub(super) struct ActivationInventory { + entries: Vec, +} + +impl ActivationInventory { + fn from_inventory(inventory: &abi::TokenizerInventory) -> Result { + if inventory.schema_version != abi::TOKENIZER_INVENTORY_SCHEMA + || inventory.tokens.is_empty() + || inventory.tokens.len() > abi::MAX_TOKENIZER_INVENTORY_ENTRIES + { + bail!("bound model exposes an unsupported or empty tokenizer inventory"); + } + let mut previous_id = None; + let entries = inventory + .tokens + .iter() + .map(|entry| { + if entry.id > abi::MAX_TOKENIZER_INVENTORY_ENTRIES as u32 { + return Err(anyhow!( + "native tokenizer ID exceeds the bounded inventory limit" + )); + } + if previous_id.is_some_and(|previous| entry.id <= previous) { + return Err(anyhow!( + "native tokenizer inventory IDs must be strictly increasing" + )); + } + previous_id = Some(entry.id); + let (piece_kind, bytes) = match &entry.piece { + abi::TokenizerInventoryPiece::Bytes { bytes } => { + (abi::TokenizerPieceKind::BYTES, bytes.as_slice()) + } + abi::TokenizerInventoryPiece::Control { descriptor } => { + if descriptor.is_empty() { + return Err(anyhow!( + "native tokenizer control descriptor must not be empty" + )); + } + (abi::TokenizerPieceKind::CONTROL, descriptor.as_slice()) + } + }; + Ok(abi::TokenizerInventoryEntry { + id: entry.id, + piece_kind, + bytes: abi::ByteSlice::from_bytes(bytes), + }) + }) + .collect::>>()?; + Ok(Self { entries }) + } + + fn view(&self) -> abi::TokenizerInventoryView { + abi::TokenizerInventoryView { + struct_size: size_of::(), + schema_version: abi::TOKENIZER_INVENTORY_SCHEMA, + entries: self.entries.as_ptr(), + entry_count: self.entries.len(), + } + } +} + +pub(super) struct HostTokenizerCapability { + pub(super) tokenizer: TokenizerCapability, + model_id: Vec, + source_model_sha256: Vec, + tokenizer_id: Vec, + inventory: ActivationInventory, + inventory_view: abi::TokenizerInventoryView, + pub(super) abi: abi::TokenizerCapability, +} + +// SAFETY: the raw pointers are immutable views into fields of this boxed +// value, which is kept alive by ActivePlugin for the whole plugin lifetime. +unsafe impl Send for HostTokenizerCapability {} +// SAFETY: the callback delegates to the model-bound, synchronized tokenizer. +unsafe impl Sync for HostTokenizerCapability {} + +impl HostTokenizerCapability { + pub(super) fn new(tokenizer: TokenizerCapability) -> Result> { + let inventory = ActivationInventory::from_inventory( + tokenizer + .inventory() + .context("bound model does not expose a tokenizer inventory")?, + )?; + let binding_digest = tokenizer + .binding_digest() + .context("bound model does not expose a tokenizer binding digest")?; + let limits = tokenizer.limits(); + let identity = tokenizer.identity().clone(); + let mut capability = Arc::new(Self { + tokenizer, + model_id: identity.model_id.into_bytes(), + source_model_sha256: identity.source_model_sha256.into_bytes(), + tokenizer_id: identity.tokenizer_id.into_bytes(), + inventory, + inventory_view: abi::TokenizerInventoryView { + struct_size: 0, + schema_version: 0, + entries: std::ptr::null(), + entry_count: 0, + }, + abi: abi::TokenizerCapability { + struct_size: size_of::(), + abi_version: abi::TOKENIZER_CAPABILITY_ABI, + model_id: abi::ByteSlice::default(), + source_model_sha256: abi::ByteSlice::default(), + tokenizer_id: abi::ByteSlice::default(), + limits: abi::TokenizerLimits { + max_input_bytes: limits.max_input_bytes, + max_output_tokens: limits.max_output_tokens, + }, + binding_digest, + inventory: std::ptr::null(), + context: std::ptr::null_mut(), + encode: encode_tokenizer, + }, + }); + let inner = Arc::get_mut(&mut capability).expect("new tokenizer capability is unique"); + inner.abi.model_id = abi::ByteSlice::from_bytes(&inner.model_id); + inner.abi.source_model_sha256 = abi::ByteSlice::from_bytes(&inner.source_model_sha256); + inner.abi.tokenizer_id = abi::ByteSlice::from_bytes(&inner.tokenizer_id); + inner.inventory_view = inner.inventory.view(); + inner.abi.inventory = &raw const inner.inventory_view; + inner.abi.context = (&raw const inner.tokenizer).cast_mut().cast(); + Ok(capability) + } +} + +unsafe extern "C" fn encode_tokenizer( + context: *mut c_void, + input_pieces: *const abi::TokenizerInputPiece, + input_piece_count: usize, + output_tokens: *mut i32, + output_capacity: usize, + output_length: *mut usize, +) -> abi::TokenizerEncodeStatus { + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + if output_length.is_null() { + return abi::TokenizerEncodeStatus::INVALID_ARGUMENT; + } + unsafe { *output_length = 0 }; + if context.is_null() { + return abi::TokenizerEncodeStatus::INVALID_ARGUMENT; + } + if input_pieces.is_null() && input_piece_count != 0 { + return abi::TokenizerEncodeStatus::INVALID_ARGUMENT; + } + if output_tokens.is_null() && output_capacity != 0 { + return abi::TokenizerEncodeStatus::INVALID_ARGUMENT; + } + if input_piece_count > abi::MAX_TOKENIZER_INPUT_PIECES { + return abi::TokenizerEncodeStatus::LIMIT_EXCEEDED; + } + let pieces = if input_piece_count == 0 { + &[] + } else { + unsafe { std::slice::from_raw_parts(input_pieces, input_piece_count) } + }; + let mut total_input_bytes = 0usize; + for piece in pieces { + if piece.bytes.pointer.is_null() && piece.bytes.length != 0 { + return abi::TokenizerEncodeStatus::INVALID_ARGUMENT; + } + total_input_bytes = match total_input_bytes.checked_add(piece.bytes.length) { + Some(total) => total, + None => return abi::TokenizerEncodeStatus::LIMIT_EXCEEDED, + }; + } + let tokenizer = unsafe { &*context.cast::() }; + let limits = tokenizer.limits(); + if total_input_bytes > limits.max_input_bytes { + return abi::TokenizerEncodeStatus::LIMIT_EXCEEDED; + } + let mut owned_pieces = Vec::with_capacity(input_piece_count); + for piece in pieces { + let bytes = if piece.bytes.length == 0 { + &[] + } else { + unsafe { std::slice::from_raw_parts(piece.bytes.pointer, piece.bytes.length) } + }; + owned_pieces.push(match piece.kind { + kind if kind == abi::TokenizerInputPieceKind::BYTES => { + InputPiece::Bytes(bytes.to_vec()) + } + kind if kind == abi::TokenizerInputPieceKind::CONTROL => InputPiece::Control { + descriptor: bytes.to_vec(), + }, + _ => return abi::TokenizerEncodeStatus::UNSUPPORTED_INPUT, + }); + } + let request = EncodeRequest::new(tokenizer.identity().clone(), owned_pieces); + let encoded = match tokenizer.encode(request) { + Ok(encoded) => encoded, + Err(error) => return encode_status(error), + }; + write_encode_output( + &encoded.token_ids, + output_tokens, + output_capacity, + output_length, + ) + })); + result.unwrap_or(abi::TokenizerEncodeStatus::INTERNAL_ERROR) +} + +fn write_encode_output( + token_ids: &[i32], + output_tokens: *mut i32, + output_capacity: usize, + output_length: *mut usize, +) -> abi::TokenizerEncodeStatus { + if output_length.is_null() || (output_tokens.is_null() && output_capacity != 0) { + return abi::TokenizerEncodeStatus::INVALID_ARGUMENT; + } + unsafe { *output_length = token_ids.len() }; + if token_ids.len() > output_capacity { + return abi::TokenizerEncodeStatus::OUTPUT_TOO_SMALL; + } + if !token_ids.is_empty() { + unsafe { + std::ptr::copy_nonoverlapping(token_ids.as_ptr(), output_tokens, token_ids.len()); + } + } + abi::TokenizerEncodeStatus::OK +} + +fn encode_status(error: skippy_tokenizer::TokenizerError) -> abi::TokenizerEncodeStatus { + match error { + skippy_tokenizer::TokenizerError::UnsupportedInput { .. } => { + abi::TokenizerEncodeStatus::UNSUPPORTED_INPUT + } + skippy_tokenizer::TokenizerError::RuntimeUnavailable => { + abi::TokenizerEncodeStatus::UNAVAILABLE + } + skippy_tokenizer::TokenizerError::InputTooLarge { .. } + | skippy_tokenizer::TokenizerError::TooManyPieces { .. } + | skippy_tokenizer::TokenizerError::TooManyTokens { .. } => { + abi::TokenizerEncodeStatus::LIMIT_EXCEEDED + } + skippy_tokenizer::TokenizerError::IdentityMismatch { .. } + | skippy_tokenizer::TokenizerError::IdentityUnavailable + | skippy_tokenizer::TokenizerError::StageZeroRequired + | skippy_tokenizer::TokenizerError::UnsupportedStage { .. } => { + abi::TokenizerEncodeStatus::INVALID_ARGUMENT + } + skippy_tokenizer::TokenizerError::BackendFailure { .. } + | skippy_tokenizer::TokenizerError::BatchTooLarge { .. } + | skippy_tokenizer::TokenizerError::BatchInputTooLarge { .. } => { + abi::TokenizerEncodeStatus::INTERNAL_ERROR + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn tokenizer_inventory_view_borrows_host_owned_bytes_for_activation() { + let inventory = abi::TokenizerInventory { + schema_version: abi::TOKENIZER_INVENTORY_SCHEMA, + model_id: "glm".to_string(), + source_model_sha256: "a".repeat(64), + tokenizer_id: "gguf-source-sha256:test".to_string(), + tokens: vec![abi::TokenizerInventoryToken { + id: 0, + piece: abi::TokenizerInventoryPiece::Bytes { + bytes: b"hello".to_vec(), + }, + }], + }; + let activation = ActivationInventory::from_inventory(&inventory).unwrap(); + let view = activation.view(); + assert_eq!(view.schema_version, abi::TOKENIZER_INVENTORY_SCHEMA); + assert_eq!(view.entry_count, 1); + let entry = unsafe { &*view.entries }; + assert_eq!(entry.piece_kind, abi::TokenizerPieceKind::BYTES); + assert_eq!( + unsafe { std::slice::from_raw_parts(entry.bytes.pointer, entry.bytes.length) }, + b"hello" + ); + } +} diff --git a/crates/skippy-server/README.md b/crates/skippy-server/README.md index 0d61b1a69..a02e6a381 100644 --- a/crates/skippy-server/README.md +++ b/crates/skippy-server/README.md @@ -86,6 +86,15 @@ round trip or a second model load. Every request supplies the expected `TokenizerIdentity`; mismatches are returned as per-item errors. The identity includes the model, source digest, tokenizer id, stage, and serving profile. +The capability also exposes a bounded structured encode operation for ordinary +byte runs and opaque native control descriptors. Mesh does not interpret +Rosetta vocabulary or control identities. The loaded backend accepts only +lossless inputs it can preserve; unsupported controls, invalid UTF-8, interior +NULs, identity mismatches, and limit violations return explicit errors rather +than being decoded with replacement semantics. Native-serving plugins receive +the same capability and inventory during activation and must prepare outside +the proposal deadline. + The `/v1/tokenize` route is retained only as an explicit compatibility and out-of-band adapter. It accepts the legacy `add_special` field as well as the facade's `special_tokens` policy and is not part of generation or proposal diff --git a/crates/skippy-server/src/tokenizer.rs b/crates/skippy-server/src/tokenizer.rs index e0eef02a1..a6bace0a7 100644 --- a/crates/skippy-server/src/tokenizer.rs +++ b/crates/skippy-server/src/tokenizer.rs @@ -1,5 +1,4 @@ use std::{ - path::Path, sync::atomic::{AtomicBool, Ordering}, sync::{Arc, Mutex}, }; @@ -12,20 +11,24 @@ use axum::{ routing::post, }; use mesh_native_serving_plugin_api as native_plugin_api; -use model_artifact::gguf::scan_gguf_tokenizer_inventory; use serde::{Deserialize, Serialize}; use skippy_protocol::StageConfig; +use skippy_tokenizer::{ + EncodeRequest, EncodeResponse, InputPiece, SpecialTokenPolicy, TokenizeBatchItem, + TokenizeRequest, TokenizeResponse, Tokenizer, TokenizerError, TokenizerIdentity, + TokenizerLimits, +}; pub use skippy_tokenizer::{ MAX_TOKENIZE_BATCH_INPUT_BYTES, MAX_TOKENIZE_BATCH_SIZE, MAX_TOKENIZE_INPUT_BYTES, MAX_TOKENIZE_TOKENS, TOKENIZER_VERSION, }; -use skippy_tokenizer::{ - SpecialTokenPolicy, TokenizeBatchItem, TokenizeRequest, TokenizeResponse, Tokenizer, - TokenizerError, TokenizerIdentity, TokenizerLimits, -}; use crate::runtime_state::RuntimeState; +mod binding; + +use binding::{inventory_from_stage, tokenizer_binding_digest}; + pub type TokenizerCapabilityError = TokenizerError; pub(crate) fn tokenizer_identity_from_stage( @@ -79,6 +82,11 @@ trait TokenizerSource: Send + Sync { add_special: bool, max_tokens: usize, ) -> Result, TokenizerCapabilityError>; + fn encode( + &self, + pieces: &[InputPiece], + max_tokens: usize, + ) -> Result, TokenizerCapabilityError>; fn token_pieces(&self, token_ids: &[i32]) -> Result>, TokenizerCapabilityError>; } @@ -147,6 +155,42 @@ impl TokenizerSource for LoadedStageZeroTokenizer { }) .collect() } + + fn encode( + &self, + pieces: &[InputPiece], + max_tokens: usize, + ) -> Result, TokenizerCapabilityError> { + let mut bytes = Vec::new(); + for piece in pieces { + match piece { + InputPiece::Bytes(piece_bytes) => bytes.extend_from_slice(piece_bytes), + InputPiece::Control { .. } => { + return Err(TokenizerCapabilityError::UnsupportedInput { + reason: "native control descriptor is unsupported by this backend" + .to_owned(), + }); + } + } + } + if bytes.is_empty() { + return Ok(Vec::new()); + } + if max_tokens == 0 { + return Err(TokenizerCapabilityError::TooManyTokens { limit: 0 }); + } + let text = std::str::from_utf8(&bytes).map_err(|_| { + TokenizerCapabilityError::UnsupportedInput { + reason: "input is not valid UTF-8".to_owned(), + } + })?; + if text.as_bytes().contains(&0) { + return Err(TokenizerCapabilityError::UnsupportedInput { + reason: "input contains an interior NUL byte".to_owned(), + }); + } + self.tokenize(text, false, max_tokens) + } } #[derive(Clone)] @@ -154,6 +198,7 @@ pub struct TokenizerCapability { identity: TokenizerIdentity, source: Arc, inventory: Option>, + binding_digest: Option<[u8; 32]>, } impl TokenizerCapability { @@ -181,10 +226,14 @@ impl TokenizerCapability { initial_check_signal: None, }); let inventory = inventory_from_stage(config, &identity, source.as_ref()).map(Arc::new); + let binding_digest = inventory + .as_deref() + .map(|inventory| tokenizer_binding_digest(&identity, inventory)); Ok(Self { identity, source, inventory, + binding_digest, }) } @@ -199,6 +248,23 @@ impl TokenizerCapability { self.inventory.as_deref() } + /// A stable digest of the bound identity, inventory, and encode behavior. + /// It is absent when the model did not expose a complete native inventory. + pub fn binding_digest(&self) -> Option<[u8; 32]> { + self.binding_digest + } + + pub fn limits(&self) -> TokenizerLimits { + ::limits(self) + } + + pub fn encode( + &self, + request: EncodeRequest, + ) -> Result { + ::encode(self, request) + } + pub fn tokenize( &self, request: TokenizeRequest, @@ -256,6 +322,44 @@ impl Tokenizer for TokenizerCapability { .collect(); Ok(items) } + + fn encode(&self, request: EncodeRequest) -> Result { + if !identity_matches(&request.expected_identity, &self.identity) { + return Err(TokenizerCapabilityError::IdentityMismatch { + expected: Box::new(request.expected_identity), + actual: Box::new(self.identity.clone()), + }); + } + let limits = self.limits(); + if request.pieces.len() > skippy_tokenizer::MAX_TOKENIZE_PIECES { + return Err(TokenizerCapabilityError::TooManyPieces { + limit: skippy_tokenizer::MAX_TOKENIZE_PIECES, + }); + } + let input_bytes = + request + .total_input_bytes() + .ok_or(TokenizerCapabilityError::InputTooLarge { + limit: limits.max_input_bytes, + })?; + if input_bytes > limits.max_input_bytes { + return Err(TokenizerCapabilityError::InputTooLarge { + limit: limits.max_input_bytes, + }); + } + let token_ids = self + .source + .encode(&request.pieces, limits.max_output_tokens)?; + if token_ids.len() > limits.max_output_tokens { + return Err(TokenizerCapabilityError::TooManyTokens { + limit: limits.max_output_tokens, + }); + } + Ok(EncodeResponse { + identity: self.identity.clone(), + token_ids, + }) + } } impl TokenizerCapability { @@ -297,52 +401,6 @@ impl TokenizerCapability { } } -fn source_gguf_path(config: &StageConfig) -> Option<&Path> { - [ - config.source_model_path.as_deref(), - config.model_path.as_deref(), - ] - .into_iter() - .flatten() - .map(Path::new) - .find(|path| path.is_file()) -} - -fn inventory_from_stage( - config: &StageConfig, - identity: &TokenizerIdentity, - source: &dyn TokenizerSource, -) -> Option { - let source_path = source_gguf_path(config)?; - let vocabulary = scan_gguf_tokenizer_inventory(source_path)?; - let token_ids = (0..vocabulary.tokens.len()) - .map(|id| i32::try_from(id).ok()) - .collect::>>()?; - let token_pieces = source.token_pieces(&token_ids).ok()?; - if token_pieces.len() != vocabulary.tokens.len() { - return None; - } - let mut tokens = Vec::with_capacity(vocabulary.tokens.len()); - for (id, (token, bytes)) in vocabulary.tokens.into_iter().zip(token_pieces).enumerate() { - let id = u32::try_from(id).ok()?; - let piece = if token.is_control { - native_plugin_api::TokenizerInventoryPiece::Control { - identity: String::from_utf8(token.raw).ok()?, - } - } else { - native_plugin_api::TokenizerInventoryPiece::Bytes { bytes } - }; - tokens.push(native_plugin_api::TokenizerInventoryToken { id, piece }); - } - Some(native_plugin_api::TokenizerInventory { - schema_version: native_plugin_api::TOKENIZER_INVENTORY_SCHEMA, - model_id: identity.model_id.clone(), - source_model_sha256: identity.source_model_sha256.clone(), - tokenizer_id: identity.tokenizer_id.clone(), - tokens, - }) -} - #[derive(Debug, Serialize)] struct TokenizerErrorBody { error: &'static str, @@ -384,9 +442,11 @@ impl IntoResponse for TokenizerHttpError { fn into_response(self) -> Response { let status = match self.0 { TokenizerCapabilityError::InputTooLarge { .. } + | TokenizerCapabilityError::TooManyPieces { .. } | TokenizerCapabilityError::BatchInputTooLarge { .. } | TokenizerCapabilityError::BatchTooLarge { .. } => StatusCode::PAYLOAD_TOO_LARGE, TokenizerCapabilityError::TooManyTokens { .. } => StatusCode::UNPROCESSABLE_ENTITY, + TokenizerCapabilityError::UnsupportedInput { .. } => StatusCode::UNPROCESSABLE_ENTITY, TokenizerCapabilityError::BackendFailure { .. } => StatusCode::INTERNAL_SERVER_ERROR, TokenizerCapabilityError::StageZeroRequired | TokenizerCapabilityError::UnsupportedStage { .. } @@ -542,6 +602,32 @@ mod tests { .map(|token_id| token_id.to_string().into_bytes()) .collect()) } + + fn encode( + &self, + pieces: &[InputPiece], + max_tokens: usize, + ) -> Result, TokenizerCapabilityError> { + if pieces + .iter() + .any(|piece| matches!(piece, InputPiece::Control { .. })) + { + return Err(TokenizerCapabilityError::UnsupportedInput { + reason: "recording source does not support controls".to_owned(), + }); + } + for piece in pieces { + let InputPiece::Bytes(bytes) = piece else { + unreachable!(); + }; + std::str::from_utf8(bytes).map_err(|_| { + TokenizerCapabilityError::UnsupportedInput { + reason: "recording source only accepts UTF-8".to_owned(), + } + })?; + } + self.tokenize("", false, max_tokens) + } } fn identity() -> TokenizerIdentity { @@ -555,6 +641,7 @@ mod tests { identity: identity(), source: source.clone(), inventory: None, + binding_digest: None, }, source, ) @@ -634,6 +721,59 @@ mod tests { assert_eq!(response.token_ids, vec![1, 4, 5]); } + #[test] + fn encode_preserves_identity_and_rejects_non_utf8_without_fallback() { + let (capability, _) = capability(vec![4, 5]); + let response = capability + .encode(EncodeRequest::bytes(identity(), b"hello".to_vec())) + .unwrap(); + assert_eq!(response.identity, identity()); + assert_eq!(response.token_ids, vec![4, 5]); + + assert!(matches!( + capability.encode(EncodeRequest::bytes(identity(), vec![0xff])), + Err(TokenizerCapabilityError::UnsupportedInput { .. }) + )); + } + + #[test] + fn encode_enforces_the_input_bound_before_calling_the_source() { + let (capability, _) = capability(vec![4, 5]); + let request = EncodeRequest::new( + identity(), + vec![InputPiece::Bytes(vec![ + b'a'; + skippy_tokenizer::MAX_TOKENIZE_INPUT_BYTES + + 1 + ])], + ); + assert_eq!( + capability.encode(request).unwrap_err(), + TokenizerCapabilityError::InputTooLarge { + limit: skippy_tokenizer::MAX_TOKENIZE_INPUT_BYTES + } + ); + } + + #[test] + fn encode_rejects_opaque_controls_when_backend_cannot_preserve_them() { + let (capability, _) = capability(vec![4, 5]); + let request = EncodeRequest::new( + identity(), + vec![ + InputPiece::Bytes(b"before".to_vec()), + InputPiece::Control { + descriptor: vec![0xff, 0x00], + }, + InputPiece::Bytes(b"after".to_vec()), + ], + ); + assert!(matches!( + capability.encode(request), + Err(TokenizerCapabilityError::UnsupportedInput { .. }) + )); + } + #[test] fn batch_results_keep_request_indexes_and_attribute_identity_errors() { let (capability, _) = capability(vec![4, 5]); @@ -749,7 +889,10 @@ mod tests { Some(model_path.display().to_string()), ); - assert_eq!(source_gguf_path(&config), Some(model_path.as_path())); + assert_eq!( + binding::source_gguf_path(&config), + Some(model_path.as_path()) + ); } #[test] diff --git a/crates/skippy-server/src/tokenizer/binding.rs b/crates/skippy-server/src/tokenizer/binding.rs new file mode 100644 index 000000000..0501f1060 --- /dev/null +++ b/crates/skippy-server/src/tokenizer/binding.rs @@ -0,0 +1,89 @@ +use std::path::Path; + +use mesh_native_serving_plugin_api as native_plugin_api; +use model_artifact::gguf::scan_gguf_tokenizer_inventory; +use skippy_protocol::StageConfig; +use skippy_tokenizer::TokenizerIdentity; + +use super::{TOKENIZER_VERSION, TokenizerSource}; + +pub(super) fn source_gguf_path(config: &StageConfig) -> Option<&Path> { + [ + config.source_model_path.as_deref(), + config.model_path.as_deref(), + ] + .into_iter() + .flatten() + .map(Path::new) + .find(|path| path.is_file()) +} + +pub(super) fn inventory_from_stage( + config: &StageConfig, + identity: &TokenizerIdentity, + source: &dyn TokenizerSource, +) -> Option { + let source_path = source_gguf_path(config)?; + let vocabulary = scan_gguf_tokenizer_inventory(source_path)?; + let token_ids = (0..vocabulary.tokens.len()) + .map(|id| i32::try_from(id).ok()) + .collect::>>()?; + let token_pieces = source.token_pieces(&token_ids).ok()?; + if token_pieces.len() != vocabulary.tokens.len() { + return None; + } + let mut tokens = Vec::with_capacity(vocabulary.tokens.len()); + for (id, (token, bytes)) in vocabulary.tokens.into_iter().zip(token_pieces).enumerate() { + let id = u32::try_from(id).ok()?; + let piece = if token.is_control { + native_plugin_api::TokenizerInventoryPiece::Control { + descriptor: token.raw, + } + } else { + native_plugin_api::TokenizerInventoryPiece::Bytes { bytes } + }; + tokens.push(native_plugin_api::TokenizerInventoryToken { id, piece }); + } + Some(native_plugin_api::TokenizerInventory { + schema_version: native_plugin_api::TOKENIZER_INVENTORY_SCHEMA, + model_id: identity.model_id.clone(), + source_model_sha256: identity.source_model_sha256.clone(), + tokenizer_id: identity.tokenizer_id.clone(), + tokens, + }) +} + +pub(super) fn tokenizer_binding_digest( + identity: &TokenizerIdentity, + inventory: &native_plugin_api::TokenizerInventory, +) -> [u8; 32] { + let mut hasher = blake3::Hasher::new(); + hasher.update(b"mesh-native-tokenizer-binding-v1\0"); + update_string(&mut hasher, &identity.model_id); + update_string(&mut hasher, &identity.source_model_sha256); + update_string(&mut hasher, &identity.tokenizer_id); + update_string(&mut hasher, TOKENIZER_VERSION); + hasher.update(&inventory.schema_version.to_le_bytes()); + hasher.update(&(inventory.tokens.len() as u64).to_le_bytes()); + for token in &inventory.tokens { + hasher.update(&token.id.to_le_bytes()); + match &token.piece { + native_plugin_api::TokenizerInventoryPiece::Bytes { bytes } => { + hasher.update(&[0]); + hasher.update(&(bytes.len() as u64).to_le_bytes()); + hasher.update(bytes); + } + native_plugin_api::TokenizerInventoryPiece::Control { descriptor } => { + hasher.update(&[1]); + hasher.update(&(descriptor.len() as u64).to_le_bytes()); + hasher.update(descriptor); + } + } + } + *hasher.finalize().as_bytes() +} + +fn update_string(hasher: &mut blake3::Hasher, value: &str) { + hasher.update(&(value.len() as u64).to_le_bytes()); + hasher.update(value.as_bytes()); +} diff --git a/crates/skippy-tokenizer/README.md b/crates/skippy-tokenizer/README.md index 9564eb53c..7fdb26586 100644 --- a/crates/skippy-tokenizer/README.md +++ b/crates/skippy-tokenizer/README.md @@ -1,6 +1,8 @@ # skippy-tokenizer `skippy-tokenizer` defines the model-bound tokenizer capability contract used -by in-process Skippy consumers. It contains request, response, identity, limit, -and typed-error types plus the bounded batch-tokenization trait; runtime crates -provide the implementation backed by an already-loaded stage-zero model. +by in-process Skippy consumers. It contains identity, limits, typed errors, a +text REST facade, and a bounded structured encoder. Encoder pieces are either +ordinary bytes or opaque native control descriptors. Mesh does not interpret +control descriptors; a backend that cannot preserve one returns +`unsupported_input` instead of using replacement decoding. diff --git a/crates/skippy-tokenizer/src/lib.rs b/crates/skippy-tokenizer/src/lib.rs index 248ce63c0..190629dc0 100644 --- a/crates/skippy-tokenizer/src/lib.rs +++ b/crates/skippy-tokenizer/src/lib.rs @@ -12,8 +12,61 @@ pub const MAX_TOKENIZE_INPUT_BYTES: usize = 1_048_576; pub const MAX_TOKENIZE_BATCH_SIZE: usize = 64; pub const MAX_TOKENIZE_BATCH_INPUT_BYTES: usize = 8 * MAX_TOKENIZE_INPUT_BYTES; pub const MAX_TOKENIZE_TOKENS: usize = 262_144; +pub const MAX_TOKENIZE_PIECES: usize = 4_096; pub const TOKENIZER_VERSION: &str = "gguf-native-v1"; +/// One piece of a structured, lossless encode request. Control descriptors are +/// opaque to Mesh; a runtime may accept them only when it can preserve their +/// native meaning without decoding them as ordinary text. +#[derive(Clone, Debug, Eq, PartialEq)] +pub enum InputPiece { + Bytes(Vec), + Control { descriptor: Vec }, +} + +impl InputPiece { + fn len(&self) -> usize { + match self { + Self::Bytes(bytes) | Self::Control { descriptor: bytes } => bytes.len(), + } + } +} + +/// A bounded request to encode structured input. The byte run is not decoded +/// with replacement semantics: runtimes that cannot preserve it must return +/// [`TokenizerError::UnsupportedInput`]. +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct EncodeRequest { + pub expected_identity: TokenizerIdentity, + pub pieces: Vec, +} + +impl EncodeRequest { + pub fn new(expected_identity: TokenizerIdentity, pieces: Vec) -> Self { + Self { + expected_identity, + pieces, + } + } + + pub fn bytes(expected_identity: TokenizerIdentity, bytes: Vec) -> Self { + Self::new(expected_identity, vec![InputPiece::Bytes(bytes)]) + } + + pub fn total_input_bytes(&self) -> Option { + self.pieces + .iter() + .map(InputPiece::len) + .try_fold(0, usize::checked_add) + } +} + +#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] +pub struct EncodeResponse { + pub identity: TokenizerIdentity, + pub token_ids: Vec, +} + #[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] pub struct TokenizerIdentity { pub model_id: String, @@ -165,9 +218,15 @@ pub enum TokenizerError { InputTooLarge { limit: usize, }, + TooManyPieces { + limit: usize, + }, TooManyTokens { limit: usize, }, + UnsupportedInput { + reason: String, + }, BackendFailure { message: String, }, @@ -184,7 +243,9 @@ impl TokenizerError { Self::BatchTooLarge { .. } => "batch_too_large", Self::BatchInputTooLarge { .. } => "batch_input_too_large", Self::InputTooLarge { .. } => "input_too_large", + Self::TooManyPieces { .. } => "too_many_pieces", Self::TooManyTokens { .. } => "too_many_tokens", + Self::UnsupportedInput { .. } => "unsupported_input", Self::BackendFailure { .. } => "backend_failure", } } @@ -209,9 +270,15 @@ impl fmt::Display for TokenizerError { Self::InputTooLarge { limit } => { write!(formatter, "tokenizer input exceeds {limit} bytes") } + Self::TooManyPieces { limit } => { + write!(formatter, "tokenizer input exceeds {limit} pieces") + } Self::TooManyTokens { limit } => { write!(formatter, "tokenizer output exceeds {limit} tokens") } + Self::UnsupportedInput { reason } => { + write!(formatter, "tokenizer does not support this input: {reason}") + } Self::BackendFailure { message } => { write!(formatter, "tokenizer backend failure: {message}") } @@ -238,6 +305,13 @@ pub trait Tokenizer: Send + Sync { &self, requests: &[TokenizeRequest], ) -> Result, TokenizerError>; + + /// Encode one structured byte/control sequence with automatic special-token + /// insertion disabled. Implementations must reject input they cannot + /// preserve exactly instead of applying lossy decoding. Mesh does not + /// interpret Rosetta or control identities; control descriptors remain + /// opaque to this contract. + fn encode(&self, request: EncodeRequest) -> Result; } #[cfg(test)] @@ -290,5 +364,12 @@ mod tests { TokenizerError::UnsupportedStage { stage_index: 1 }.code(), "unsupported_stage" ); + assert_eq!( + TokenizerError::UnsupportedInput { + reason: "invalid UTF-8".to_owned() + } + .code(), + "unsupported_input" + ); } }