diff --git a/.github/workflows/docker-precheck.yml b/.github/workflows/docker-precheck.yml index aebb41d8e..d88065414 100644 --- a/.github/workflows/docker-precheck.yml +++ b/.github/workflows/docker-precheck.yml @@ -96,6 +96,7 @@ jobs: crates/model-artifact/ \ crates/model-hf/ \ crates/skippy-protocol/ \ + crates/skippy-tokenizer/ \ crates/skippy-topology/ \ crates/skippy-ffi/ \ crates/skippy-runtime/ \ diff --git a/Cargo.lock b/Cargo.lock index 6bb81f87c..5f46d77cc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7442,6 +7442,7 @@ dependencies = [ "protoc-bin-vendored", "serde", "serde_json", + "skippy-tokenizer", ] [[package]] @@ -7497,6 +7498,7 @@ dependencies = [ "skippy-metrics", "skippy-protocol", "skippy-runtime", + "skippy-tokenizer", "socket2", "tempfile", "tokio", @@ -7505,6 +7507,14 @@ dependencies = [ "tower", ] +[[package]] +name = "skippy-tokenizer" +version = "0.72.1" +dependencies = [ + "serde", + "serde_json", +] + [[package]] name = "skippy-topology" version = "0.72.1" diff --git a/Cargo.toml b/Cargo.toml index 440cf6889..6145380c0 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -39,6 +39,7 @@ members = [ "crates/model-artifact", "crates/model-hf", "crates/model-resolver", + "crates/skippy-tokenizer", "crates/skippy-protocol", "crates/skippy-coordinator", "crates/skippy-topology", diff --git a/crates/mesh-llm-host-runtime/src/runtime/local_split/tests.rs b/crates/mesh-llm-host-runtime/src/runtime/local_split/tests.rs index b1376e693..e7bcc581d 100644 --- a/crates/mesh-llm-host-runtime/src/runtime/local_split/tests.rs +++ b/crates/mesh-llm-host-runtime/src/runtime/local_split/tests.rs @@ -1018,8 +1018,26 @@ fn split_recovery_candidate_participants_excludes_unavailable_stage_nodes() { ); } -#[tokio::test] -async fn load_split_runtime_generation_stops_candidate_stages_after_partial_load_failure() { +#[test] +fn load_split_runtime_generation_stops_candidate_stages_after_partial_load_failure() { + std::thread::Builder::new() + .name("local-split-test".to_owned()) + .stack_size(8 * 1024 * 1024) + .spawn(|| { + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("build local split test runtime") + .block_on( + load_split_runtime_generation_stops_candidate_stages_after_partial_load_failure_inner(), + ); + }) + .expect("spawn local split test thread") + .join() + .expect("local split test thread panicked"); +} + +async fn load_split_runtime_generation_stops_candidate_stages_after_partial_load_failure_inner() { let node = mesh::Node::new_for_tests(NodeRole::Host { http_port: 9337 }) .await .unwrap(); diff --git a/crates/skippy-protocol/Cargo.toml b/crates/skippy-protocol/Cargo.toml index e937b2d7a..13da59d39 100644 --- a/crates/skippy-protocol/Cargo.toml +++ b/crates/skippy-protocol/Cargo.toml @@ -10,6 +10,7 @@ homepage = "https://github.com/Mesh-LLM/mesh-llm" [dependencies] prost = "0.14" serde.workspace = true +skippy-tokenizer = { path = "../skippy-tokenizer", version = "0.72.1" } [dev-dependencies] serde_json.workspace = true diff --git a/crates/skippy-protocol/src/lib.rs b/crates/skippy-protocol/src/lib.rs index d7fe61660..b06cd5e8e 100644 --- a/crates/skippy-protocol/src/lib.rs +++ b/crates/skippy-protocol/src/lib.rs @@ -1,6 +1,13 @@ use serde::{Deserialize, Serialize}; -pub mod tokenizer; +/// Compatibility namespace for tokenizer contracts. +/// +/// New consumers should depend on `skippy-tokenizer` directly. Keeping this +/// re-export avoids breaking older protocol users while the contract moves +/// out of the wire-protocol crate. +pub mod tokenizer { + pub use skippy_tokenizer::*; +} pub mod binary; pub mod proto { diff --git a/crates/skippy-protocol/src/tokenizer.rs b/crates/skippy-protocol/src/tokenizer.rs deleted file mode 100644 index ff0494f3b..000000000 --- a/crates/skippy-protocol/src/tokenizer.rs +++ /dev/null @@ -1,67 +0,0 @@ -use serde::{Deserialize, Serialize}; - -pub const MAX_TOKENIZE_INPUT_BYTES: usize = 1_048_576; - -#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] -pub struct TokenizerIdentity { - pub model_id: String, - pub source_model_sha256: String, - pub tokenizer_id: String, -} - -#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] -pub struct TokenizeRequest { - pub expected_identity: TokenizerIdentity, - pub text: String, - #[serde(default)] - pub add_special: bool, - #[serde(default)] - pub include_token_pieces: bool, -} - -#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] -pub struct TokenizeResponse { - pub identity: TokenizerIdentity, - pub token_ids: Vec, - #[serde(skip_serializing_if = "Option::is_none")] - pub token_pieces: Option>>, -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn tokenizer_wire_contract_is_exact() { - let identity = TokenizerIdentity { - model_id: "model".to_owned(), - source_model_sha256: "a".repeat(64), - tokenizer_id: format!("gguf-source-sha256:{}", "a".repeat(64)), - }; - let request: TokenizeRequest = serde_json::from_value(serde_json::json!({ - "expected_identity": identity, - "text": "hello", - })) - .unwrap(); - assert!(!request.add_special); - assert!(!request.include_token_pieces); - - let response = TokenizeResponse { - identity: request.expected_identity, - token_ids: vec![1, 2], - token_pieces: None, - }; - assert_eq!( - serde_json::to_value(response).unwrap(), - serde_json::json!({ - "identity": { - "model_id": "model", - "source_model_sha256": "a".repeat(64), - "tokenizer_id": format!("gguf-source-sha256:{}", "a".repeat(64)), - }, - "token_ids": [1, 2], - }) - ); - assert_eq!(MAX_TOKENIZE_INPUT_BYTES, 1_048_576); - } -} diff --git a/crates/skippy-runtime/src/native.rs b/crates/skippy-runtime/src/native.rs index dbbb3e2c8..4879d3b7d 100644 --- a/crates/skippy-runtime/src/native.rs +++ b/crates/skippy-runtime/src/native.rs @@ -325,6 +325,21 @@ impl StageModel { } pub fn tokenize(&self, text: &str, add_special: bool) -> Result> { + self.tokenize_bounded(text, add_special, usize::MAX)? + .ok_or_else(|| anyhow!("tokenizer output exceeds the requested limit")) + } + + /// Tokenize without allocating a token buffer larger than `max_tokens`. + /// + /// The native ABI reports the required count during its sizing call. When + /// that count exceeds the bound, this returns `Ok(None)` before allocating + /// the output vector. + pub fn tokenize_bounded( + &self, + text: &str, + add_special: bool, + max_tokens: usize, + ) -> Result>> { let text = CString::new(text).context("text contains an interior NUL byte")?; let mut count = 0usize; let mut error = ptr::null_mut(); @@ -345,6 +360,10 @@ impl StageModel { free_error(error); } + if count > max_tokens { + return Ok(None); + } + let mut tokens = vec![0_i32; count]; let mut error = ptr::null_mut(); let status = unsafe { @@ -358,9 +377,13 @@ impl StageModel { &mut error, ) }; + if status == Status::BufferTooSmall { + free_error(error); + return Ok(None); + } ensure_ok(status, error)?; tokens.truncate(count); - Ok(tokens) + Ok(Some(tokens)) } pub fn detokenize(&self, tokens: &[i32]) -> Result { diff --git a/crates/skippy-server/Cargo.toml b/crates/skippy-server/Cargo.toml index 5c0b40371..b4f413f2a 100644 --- a/crates/skippy-server/Cargo.toml +++ b/crates/skippy-server/Cargo.toml @@ -25,6 +25,7 @@ blake3.workspace = true clap.workspace = true futures-util = "0.3" skippy-runtime = { path = "../skippy-runtime", version = "0.72.1" } +skippy-tokenizer = { path = "../skippy-tokenizer", version = "0.72.1" } model-artifact = { path = "../model-artifact", version = "0.72.1" } mesh-native-serving-plugin-api = { path = "../mesh-native-serving-plugin-api", version = "0.72.1" } skippy-protocol = { path = "../skippy-protocol", version = "0.72.1" } diff --git a/crates/skippy-server/README.md b/crates/skippy-server/README.md index d3b4eb0dc..0d61b1a69 100644 --- a/crates/skippy-server/README.md +++ b/crates/skippy-server/README.md @@ -77,6 +77,20 @@ topology, builds the stage configs, loads/starts handles, watches readiness and status, withdraws routes before shutdown, and then calls handle shutdown during unload or replan. +### In-process tokenizer capability + +`SkippyRuntimeHandle::tokenizer_capability()` returns a model-bound +`skippy_tokenizer::Tokenizer` backed by the already-loaded stage-zero runtime. +Consumers can call `tokenize_batch` for bounded, ordered results without an HTTP +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 `/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 +deadline handling. + ## Notes - `serve-binary` is the tuned binary stage-to-stage path. diff --git a/crates/skippy-server/src/embedded.rs b/crates/skippy-server/src/embedded.rs index 2b959f97b..0a981ac57 100644 --- a/crates/skippy-server/src/embedded.rs +++ b/crates/skippy-server/src/embedded.rs @@ -1,6 +1,7 @@ use std::{ net::SocketAddr, - sync::{Arc, Mutex, TryLockError}, + sync::atomic::{AtomicBool, Ordering}, + sync::{Arc, Mutex, OnceLock, TryLockError}, }; use anyhow::{Context, Result}; @@ -84,6 +85,8 @@ pub struct SkippyRuntimeHandle { runtime: Arc>, telemetry: Telemetry, status: Arc>, + tokenizer_active: Arc, + tokenizer_capability: OnceLock>, /// Last session stats read out of [`Self::runtime`], and when. /// /// A native call (long prefill, decode batch) holds the runtime lock while @@ -140,6 +143,8 @@ impl SkippyRuntimeHandle { stopped_at_unix_nanos: None, last_error: None, })), + tokenizer_active: Arc::new(AtomicBool::new(true)), + tokenizer_capability: OnceLock::new(), last_session_stats: Arc::new(Mutex::new(initial_session_stats)), } } @@ -241,7 +246,15 @@ impl SkippyRuntimeHandle { /// Returns the stateless tokenizer capability backed by this already-loaded /// stage-zero runtime. This never opens a second model. pub fn tokenizer_capability(&self) -> Result { - TokenizerCapability::from_stage_zero(&self.config, self.runtime.clone()) + self.tokenizer_capability + .get_or_init(|| { + TokenizerCapability::from_stage_zero_with_lifecycle( + &self.config, + self.runtime.clone(), + self.tokenizer_active.clone(), + ) + }) + .clone() } pub fn status(&self) -> EmbeddedRuntimeStatus { @@ -270,6 +283,9 @@ impl SkippyRuntimeHandle { } pub fn shutdown(&self) { + self.tokenizer_active.store(false, Ordering::Release); + let runtime = self.runtime.lock().expect("runtime lock poisoned"); + drop(runtime); let mut status = self.status.lock().expect("runtime status lock poisoned"); if status.state == EmbeddedState::Stopped { return; @@ -689,6 +705,28 @@ mod tests { assert!(status.runtime_loaded); } + #[test] + fn shutdown_waits_for_runtime_lock_after_invalidating_tokenizer() { + let handle = Arc::new(test_handle(1)); + let held = handle.runtime.lock().expect("runtime lock"); + let (shutdown_tx, shutdown_rx) = mpsc::channel(); + let shutdown_handle = Arc::clone(&handle); + thread::spawn(move || { + shutdown_handle.shutdown(); + shutdown_tx.send(()).expect("send shutdown result"); + }); + + assert!( + shutdown_rx.recv_timeout(Duration::from_millis(50)).is_err(), + "shutdown should synchronize with an in-flight runtime operation" + ); + drop(held); + shutdown_rx + .recv_timeout(Duration::from_secs(1)) + .expect("shutdown should complete after the runtime lock is released"); + assert_eq!(handle.status().state, EmbeddedState::Stopped); + } + fn empty_cache(value: u32) -> Mutex> { Mutex::new(Captured { value, diff --git a/crates/skippy-server/src/tokenizer.rs b/crates/skippy-server/src/tokenizer.rs index 917b0bfff..e0eef02a1 100644 --- a/crates/skippy-server/src/tokenizer.rs +++ b/crates/skippy-server/src/tokenizer.rs @@ -1,7 +1,6 @@ use std::{ - error::Error, - fmt, path::Path, + sync::atomic::{AtomicBool, Ordering}, sync::{Arc, Mutex}, }; @@ -15,46 +14,19 @@ use axum::{ 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, - tokenizer::{MAX_TOKENIZE_INPUT_BYTES, TokenizeRequest, TokenizeResponse, TokenizerIdentity}, +use skippy_protocol::StageConfig; +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; -pub const MAX_TOKENIZE_TOKENS: usize = 262_144; - -#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] -#[serde(rename_all = "snake_case")] -pub enum TokenizerCapabilityError { - StageZeroRequired, - IdentityUnavailable, - IdentityMismatch, - InputTooLarge, - TooManyTokens, - BackendFailure, -} - -impl TokenizerCapabilityError { - pub const fn code(self) -> &'static str { - match self { - Self::StageZeroRequired => "stage_zero_required", - Self::IdentityUnavailable => "identity_unavailable", - Self::IdentityMismatch => "identity_mismatch", - Self::InputTooLarge => "input_too_large", - Self::TooManyTokens => "too_many_tokens", - Self::BackendFailure => "backend_failure", - } - } -} - -impl fmt::Display for TokenizerCapabilityError { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - formatter.write_str(self.code()) - } -} - -impl Error for TokenizerCapabilityError {} +pub type TokenizerCapabilityError = TokenizerError; pub(crate) fn tokenizer_identity_from_stage( stage_index: u32, @@ -75,21 +47,56 @@ pub(crate) fn tokenizer_identity_from_stage( model_id: model_id.to_owned(), tokenizer_id: format!("gguf-source-sha256:{source_model_sha256}"), source_model_sha256, + tokenizer_version: Some(TOKENIZER_VERSION.to_owned()), + stage_index, + serving_profile: Some("stage-zero".to_owned()), }) } +fn identity_matches(expected: &TokenizerIdentity, actual: &TokenizerIdentity) -> bool { + expected.model_id == actual.model_id + && expected.source_model_sha256 == actual.source_model_sha256 + && expected.tokenizer_id == actual.tokenizer_id + && expected.stage_index == actual.stage_index + && expected + .tokenizer_version + .as_ref() + .is_none_or(|value| actual.tokenizer_version.as_ref() == Some(value)) + && expected + .serving_profile + .as_ref() + .is_none_or(|value| actual.serving_profile.as_ref() == Some(value)) +} + fn is_sha256(value: &str) -> bool { value.len() == 64 && value.bytes().all(|byte| byte.is_ascii_hexdigit()) } trait TokenizerSource: Send + Sync { - fn tokenize(&self, text: &str, add_special: bool) - -> Result, TokenizerCapabilityError>; + fn tokenize( + &self, + text: &str, + add_special: bool, + max_tokens: usize, + ) -> Result, TokenizerCapabilityError>; fn token_pieces(&self, token_ids: &[i32]) -> Result>, TokenizerCapabilityError>; } struct LoadedStageZeroTokenizer { runtime: Arc>, + active: Arc, + #[cfg(test)] + initial_check_signal: Option>, +} + +impl LoadedStageZeroTokenizer { + fn ensure_active(&self) -> Result<(), TokenizerCapabilityError> { + if self.active.load(Ordering::Acquire) { + Ok(()) + } else { + Err(TokenizerCapabilityError::RuntimeUnavailable) + } + } } impl TokenizerSource for LoadedStageZeroTokenizer { @@ -97,27 +104,46 @@ impl TokenizerSource for LoadedStageZeroTokenizer { &self, text: &str, add_special: bool, + max_tokens: usize, ) -> Result, TokenizerCapabilityError> { - self.runtime - .lock() - .map_err(|_| TokenizerCapabilityError::BackendFailure)? + self.ensure_active()?; + #[cfg(test)] + if let Some(signal) = &self.initial_check_signal { + signal.wait(); + } + let runtime = + self.runtime + .lock() + .map_err(|_| TokenizerCapabilityError::BackendFailure { + message: "runtime lock poisoned".to_owned(), + })?; + self.ensure_active()?; + let tokens = runtime .model - .tokenize(text, add_special) - .map_err(|_| TokenizerCapabilityError::BackendFailure) + .tokenize_bounded(text, add_special, max_tokens) + .map_err(|error| TokenizerCapabilityError::BackendFailure { + message: error.to_string(), + })?; + tokens.ok_or(TokenizerCapabilityError::TooManyTokens { limit: max_tokens }) } fn token_pieces(&self, token_ids: &[i32]) -> Result>, TokenizerCapabilityError> { - let runtime = self - .runtime - .lock() - .map_err(|_| TokenizerCapabilityError::BackendFailure)?; + self.ensure_active()?; + let runtime = + self.runtime + .lock() + .map_err(|_| TokenizerCapabilityError::BackendFailure { + message: "runtime lock poisoned".to_owned(), + })?; + self.ensure_active()?; token_ids .iter() .map(|token_id| { - runtime - .model - .detokenize_bytes(&[*token_id]) - .map_err(|_| TokenizerCapabilityError::BackendFailure) + runtime.model.detokenize_bytes(&[*token_id]).map_err(|_| { + TokenizerCapabilityError::BackendFailure { + message: "detokenization failed".to_owned(), + } + }) }) .collect() } @@ -131,16 +157,29 @@ pub struct TokenizerCapability { } impl TokenizerCapability { - pub(crate) fn from_stage_zero( + pub fn from_stage_zero( + config: &StageConfig, + runtime: Arc>, + ) -> Result { + Self::from_stage_zero_with_lifecycle(config, runtime, Arc::new(AtomicBool::new(true))) + } + + pub(crate) fn from_stage_zero_with_lifecycle( config: &StageConfig, runtime: Arc>, + active: Arc, ) -> Result { let identity = tokenizer_identity_from_stage( config.stage_index, &config.model_id, config.source_model_sha256.as_deref(), )?; - let source: Arc = Arc::new(LoadedStageZeroTokenizer { runtime }); + let source: Arc = Arc::new(LoadedStageZeroTokenizer { + runtime, + active, + #[cfg(test)] + initial_check_signal: None, + }); let inventory = inventory_from_stage(config, &identity, source.as_ref()).map(Arc::new); Ok(Self { identity, @@ -164,15 +203,87 @@ impl TokenizerCapability { &self, request: TokenizeRequest, ) -> Result { - if request.expected_identity != self.identity { - return Err(TokenizerCapabilityError::IdentityMismatch); + self.tokenize_batch(&[request])? + .into_iter() + .next() + .expect("single-request tokenizer batch must return one item") + .result + } +} + +impl Tokenizer for TokenizerCapability { + fn identity(&self) -> &TokenizerIdentity { + &self.identity + } + + fn limits(&self) -> TokenizerLimits { + TokenizerLimits::default() + } + + fn tokenize_batch( + &self, + requests: &[TokenizeRequest], + ) -> Result, TokenizerCapabilityError> { + let limits = self.limits(); + if requests.len() > limits.max_batch_size { + return Err(TokenizerCapabilityError::BatchTooLarge { + limit: limits.max_batch_size, + }); } - if request.text.len() > MAX_TOKENIZE_INPUT_BYTES { - return Err(TokenizerCapabilityError::InputTooLarge); + let batch_input_bytes = requests + .iter() + .map(|request| request.text.len()) + .try_fold(0usize, usize::checked_add) + .ok_or(TokenizerCapabilityError::BatchInputTooLarge { + limit: limits.max_batch_input_bytes, + })?; + if batch_input_bytes > limits.max_batch_input_bytes { + return Err(TokenizerCapabilityError::BatchInputTooLarge { + limit: limits.max_batch_input_bytes, + }); } - let token_ids = self.source.tokenize(&request.text, request.add_special)?; - if token_ids.len() > MAX_TOKENIZE_TOKENS { - return Err(TokenizerCapabilityError::TooManyTokens); + + let items = requests + .iter() + .enumerate() + .map(|(request_index, request)| { + let result = self.tokenize_one(request, limits); + TokenizeBatchItem { + request_index, + result, + } + }) + .collect(); + Ok(items) + } +} + +impl TokenizerCapability { + fn tokenize_one( + &self, + request: &TokenizeRequest, + limits: TokenizerLimits, + ) -> Result { + if !identity_matches(&request.expected_identity, &self.identity) { + return Err(TokenizerCapabilityError::IdentityMismatch { + expected: Box::new(request.expected_identity.clone()), + actual: Box::new(self.identity.clone()), + }); + } + if request.text.len() > limits.max_input_bytes { + return Err(TokenizerCapabilityError::InputTooLarge { + limit: limits.max_input_bytes, + }); + } + let token_ids = self.source.tokenize( + &request.text, + request.special_tokens == SpecialTokenPolicy::Add, + limits.max_output_tokens, + )?; + if token_ids.len() > limits.max_output_tokens { + return Err(TokenizerCapabilityError::TooManyTokens { + limit: limits.max_output_tokens, + }); } let token_pieces = request .include_token_pieces @@ -237,17 +348,51 @@ struct TokenizerErrorBody { error: &'static str, } +#[derive(Deserialize)] +struct HttpTokenizeRequest { + expected_identity: TokenizerIdentity, + text: String, + #[serde(default)] + special_tokens: SpecialTokenPolicy, + #[serde(default)] + add_special: Option, + #[serde(default)] + include_token_pieces: bool, +} + +impl HttpTokenizeRequest { + fn into_request(self) -> TokenizeRequest { + let special_tokens = self.add_special.map_or(self.special_tokens, |add_special| { + if add_special { + SpecialTokenPolicy::Add + } else { + SpecialTokenPolicy::Omit + } + }); + TokenizeRequest { + expected_identity: self.expected_identity, + text: self.text, + special_tokens, + include_token_pieces: self.include_token_pieces, + } + } +} + struct TokenizerHttpError(TokenizerCapabilityError); impl IntoResponse for TokenizerHttpError { fn into_response(self) -> Response { let status = match self.0 { - TokenizerCapabilityError::InputTooLarge => StatusCode::PAYLOAD_TOO_LARGE, - TokenizerCapabilityError::TooManyTokens => StatusCode::UNPROCESSABLE_ENTITY, - TokenizerCapabilityError::BackendFailure => StatusCode::INTERNAL_SERVER_ERROR, + TokenizerCapabilityError::InputTooLarge { .. } + | TokenizerCapabilityError::BatchInputTooLarge { .. } + | TokenizerCapabilityError::BatchTooLarge { .. } => StatusCode::PAYLOAD_TOO_LARGE, + TokenizerCapabilityError::TooManyTokens { .. } => StatusCode::UNPROCESSABLE_ENTITY, + TokenizerCapabilityError::BackendFailure { .. } => StatusCode::INTERNAL_SERVER_ERROR, TokenizerCapabilityError::StageZeroRequired + | TokenizerCapabilityError::UnsupportedStage { .. } | TokenizerCapabilityError::IdentityUnavailable - | TokenizerCapabilityError::IdentityMismatch => StatusCode::CONFLICT, + | TokenizerCapabilityError::IdentityMismatch { .. } => StatusCode::CONFLICT, + TokenizerCapabilityError::RuntimeUnavailable => StatusCode::SERVICE_UNAVAILABLE, }; ( status, @@ -272,17 +417,27 @@ pub(crate) fn tokenizer_http_router(capability: TokenizerCapability) -> Router { async fn tokenize_entrypoint( State(capability): State, - Json(request): Json, + Json(request): Json, ) -> Result, TokenizerHttpError> { - tokio::task::spawn_blocking(move || capability.tokenize(request)) + tokio::task::spawn_blocking(move || capability.tokenize(request.into_request())) .await - .map_err(|_| TokenizerHttpError(TokenizerCapabilityError::BackendFailure))? + .map_err(|_| { + TokenizerHttpError(TokenizerCapabilityError::BackendFailure { + message: "tokenizer task failed".to_owned(), + }) + })? .map(Json) .map_err(TokenizerHttpError) } #[cfg(test)] mod tests { + use std::{ + sync::{Barrier, atomic::AtomicBool}, + thread, + time::Duration, + }; + use axum::{ body::{Body, to_bytes}, http::{Request, StatusCode, header::CONTENT_TYPE}, @@ -366,11 +521,15 @@ mod tests { &self, _text: &str, add_special: bool, + max_tokens: usize, ) -> Result, TokenizerCapabilityError> { let mut tokens = self.tokens.clone(); if add_special { tokens.insert(0, 1); } + if tokens.len() > max_tokens { + return Err(TokenizerCapabilityError::TooManyTokens { limit: max_tokens }); + } Ok(tokens) } @@ -405,11 +564,46 @@ mod tests { TokenizeRequest { expected_identity: identity(), text, - add_special: false, + special_tokens: SpecialTokenPolicy::Omit, include_token_pieces: false, } } + #[test] + fn queued_tokenization_rechecks_lifecycle_after_runtime_lock() { + let runtime = Arc::new(Mutex::new(RuntimeState::new_modelless_for_test(1))); + let active = Arc::new(AtomicBool::new(true)); + let source = Arc::new(LoadedStageZeroTokenizer { + runtime: Arc::clone(&runtime), + active: Arc::clone(&active), + initial_check_signal: Some(Arc::new(Barrier::new(2))), + }); + let held = runtime.lock().expect("runtime lock"); + let (result_tx, result_rx) = std::sync::mpsc::channel(); + let request_source = Arc::clone(&source); + thread::spawn(move || { + result_tx + .send(request_source.tokenize("hello", false, 8)) + .expect("send tokenizer result"); + }); + + source + .initial_check_signal + .as_ref() + .expect("test synchronization signal") + .wait(); + active.store(false, Ordering::Release); + drop(held); + + assert_eq!( + result_rx + .recv_timeout(Duration::from_secs(1)) + .expect("receive tokenizer result") + .unwrap_err(), + TokenizerCapabilityError::RuntimeUnavailable + ); + } + async fn post_tokenize(capability: TokenizerCapability, request: &TokenizeRequest) -> Response { tokenizer_http_router(capability) .oneshot( @@ -434,8 +628,57 @@ mod tests { #[test] fn tokenization_returns_source_tokens() { let (capability, _) = capability(vec![4, 5]); - let response = capability.tokenize(request("hello".to_string())).unwrap(); - assert_eq!(response.token_ids, vec![4, 5]); + let mut request = request("hello".to_string()); + request.special_tokens = SpecialTokenPolicy::Add; + let response = capability.tokenize(request).unwrap(); + assert_eq!(response.token_ids, vec![1, 4, 5]); + } + + #[test] + fn batch_results_keep_request_indexes_and_attribute_identity_errors() { + let (capability, _) = capability(vec![4, 5]); + let mut mismatched = request("second".to_owned()); + mismatched.expected_identity.model_id = "other-model".to_owned(); + let expected_identity = mismatched.expected_identity.clone(); + + let results = capability + .tokenize_batch(&[request("first".to_owned()), mismatched.clone()]) + .unwrap(); + assert_eq!(results.len(), 2); + assert_eq!(results[0].request_index, 0); + assert!(results[0].result.is_ok()); + assert_eq!(results[1].request_index, 1); + assert_eq!( + results[1].result.as_ref().unwrap_err().clone(), + TokenizerCapabilityError::IdentityMismatch { + expected: Box::new(expected_identity), + actual: Box::new(identity()), + } + ); + } + + #[test] + fn batch_size_and_input_limits_are_rejected_before_tokenization() { + let (capability, _) = capability(vec![4]); + let too_many = (0..=MAX_TOKENIZE_BATCH_SIZE) + .map(|index| request(index.to_string())) + .collect::>(); + assert_eq!( + capability.tokenize_batch(&too_many).unwrap_err(), + TokenizerCapabilityError::BatchTooLarge { + limit: MAX_TOKENIZE_BATCH_SIZE + } + ); + + let too_many_bytes = (0..=MAX_TOKENIZE_BATCH_INPUT_BYTES / MAX_TOKENIZE_INPUT_BYTES) + .map(|_| request("x".repeat(MAX_TOKENIZE_INPUT_BYTES))) + .collect::>(); + assert_eq!( + capability.tokenize_batch(&too_many_bytes).unwrap_err(), + TokenizerCapabilityError::BatchInputTooLarge { + limit: MAX_TOKENIZE_BATCH_INPUT_BYTES + } + ); } #[test] @@ -444,13 +687,23 @@ mod tests { let error = input_bounded .tokenize(request("x".repeat(MAX_TOKENIZE_INPUT_BYTES + 1))) .unwrap_err(); - assert_eq!(error, TokenizerCapabilityError::InputTooLarge); + assert_eq!( + error, + TokenizerCapabilityError::InputTooLarge { + limit: MAX_TOKENIZE_INPUT_BYTES + } + ); let (output_bounded, _) = capability(vec![7; MAX_TOKENIZE_TOKENS + 1]); let error = output_bounded .tokenize(request("x".to_string())) .unwrap_err(); - assert_eq!(error, TokenizerCapabilityError::TooManyTokens); + assert_eq!( + error, + TokenizerCapabilityError::TooManyTokens { + limit: MAX_TOKENIZE_TOKENS + } + ); } #[test] @@ -466,12 +719,26 @@ mod tests { let (capability, _) = capability(Vec::new()); let mut request = request("x".to_string()); request.expected_identity.model_id = "another-model".to_string(); + let expected_identity = request.expected_identity.clone(); assert_eq!( capability.tokenize(request).unwrap_err(), - TokenizerCapabilityError::IdentityMismatch + TokenizerCapabilityError::IdentityMismatch { + expected: Box::new(expected_identity), + actual: Box::new(identity()), + } ); } + #[test] + fn legacy_identity_without_optional_provenance_fields_remains_usable() { + let (capability, _) = capability(vec![4, 5]); + let mut request = request("legacy".to_string()); + request.expected_identity.tokenizer_version = None; + request.expected_identity.serving_profile = None; + + assert_eq!(capability.tokenize(request).unwrap().token_ids, vec![4, 5]); + } + #[test] fn source_gguf_path_falls_back_when_source_model_path_is_unavailable() { let temp_dir = tempfile::tempdir().expect("create temporary GGUF directory"); diff --git a/crates/skippy-tokenizer/Cargo.toml b/crates/skippy-tokenizer/Cargo.toml new file mode 100644 index 000000000..e1570a810 --- /dev/null +++ b/crates/skippy-tokenizer/Cargo.toml @@ -0,0 +1,14 @@ +[package] +name = "skippy-tokenizer" +edition.workspace = true +license.workspace = true +version.workspace = true +description = "Model-bound tokenizer capability contract for Skippy consumers" +repository = "https://github.com/Mesh-LLM/mesh-llm" +homepage = "https://github.com/Mesh-LLM/mesh-llm" + +[dependencies] +serde.workspace = true + +[dev-dependencies] +serde_json.workspace = true diff --git a/crates/skippy-tokenizer/README.md b/crates/skippy-tokenizer/README.md new file mode 100644 index 000000000..9564eb53c --- /dev/null +++ b/crates/skippy-tokenizer/README.md @@ -0,0 +1,6 @@ +# 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. diff --git a/crates/skippy-tokenizer/src/lib.rs b/crates/skippy-tokenizer/src/lib.rs new file mode 100644 index 000000000..248ce63c0 --- /dev/null +++ b/crates/skippy-tokenizer/src/lib.rs @@ -0,0 +1,294 @@ +//! A model-bound tokenizer capability for in-process Skippy consumers. +//! +//! This crate intentionally contains only the stable contract. Runtime +//! implementations bind it to an already-loaded tokenizer; they must not open +//! another model on behalf of a caller. + +use std::{error::Error, fmt}; + +use serde::{Deserialize, Deserializer, Serialize, Serializer, ser::SerializeStruct}; + +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 TOKENIZER_VERSION: &str = "gguf-native-v1"; + +#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] +pub struct TokenizerIdentity { + pub model_id: String, + pub source_model_sha256: String, + pub tokenizer_id: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub tokenizer_version: Option, + #[serde(default)] + pub stage_index: u32, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub serving_profile: Option, +} + +#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, PartialEq, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum SpecialTokenPolicy { + #[default] + Omit, + Add, +} + +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct TokenizeRequest { + pub expected_identity: TokenizerIdentity, + pub text: String, + pub special_tokens: SpecialTokenPolicy, + pub include_token_pieces: bool, +} + +#[derive(Deserialize)] +struct TokenizeRequestWire { + expected_identity: TokenizerIdentity, + text: String, + #[serde(default)] + special_tokens: Option, + #[serde(default)] + add_special: Option, + #[serde(default)] + include_token_pieces: bool, +} + +impl<'de> Deserialize<'de> for TokenizeRequest { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + let wire = TokenizeRequestWire::deserialize(deserializer)?; + let special_tokens = wire + .add_special + .map(|add_special| { + if add_special { + SpecialTokenPolicy::Add + } else { + SpecialTokenPolicy::Omit + } + }) + .or(wire.special_tokens) + .unwrap_or_default(); + Ok(Self { + expected_identity: wire.expected_identity, + text: wire.text, + special_tokens, + include_token_pieces: wire.include_token_pieces, + }) + } +} + +impl Serialize for TokenizeRequest { + fn serialize(&self, serializer: S) -> Result + where + S: Serializer, + { + let mut state = serializer.serialize_struct("TokenizeRequest", 5)?; + state.serialize_field("expected_identity", &self.expected_identity)?; + state.serialize_field("text", &self.text)?; + state.serialize_field("special_tokens", &self.special_tokens)?; + // Keep the legacy REST spelling in serialized requests so a new + // consumer can still request special tokens from an older server. + state.serialize_field( + "add_special", + &(self.special_tokens == SpecialTokenPolicy::Add), + )?; + state.serialize_field("include_token_pieces", &self.include_token_pieces)?; + state.end() + } +} + +impl TokenizeRequest { + pub fn with_special_tokens( + expected_identity: TokenizerIdentity, + text: String, + special_tokens: SpecialTokenPolicy, + ) -> Self { + Self { + expected_identity, + text, + special_tokens, + include_token_pieces: false, + } + } +} + +#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] +pub struct TokenizeResponse { + pub identity: TokenizerIdentity, + pub token_ids: Vec, + #[serde(skip_serializing_if = "Option::is_none")] + pub token_pieces: Option>>, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct TokenizerLimits { + pub max_batch_size: usize, + pub max_batch_input_bytes: usize, + pub max_input_bytes: usize, + pub max_output_tokens: usize, +} + +impl Default for TokenizerLimits { + fn default() -> Self { + Self { + max_batch_size: MAX_TOKENIZE_BATCH_SIZE, + max_batch_input_bytes: MAX_TOKENIZE_BATCH_INPUT_BYTES, + max_input_bytes: MAX_TOKENIZE_INPUT_BYTES, + max_output_tokens: MAX_TOKENIZE_TOKENS, + } + } +} + +#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] +#[serde(rename_all = "snake_case", tag = "code")] +pub enum TokenizerError { + StageZeroRequired, + UnsupportedStage { + stage_index: u32, + }, + IdentityUnavailable, + RuntimeUnavailable, + IdentityMismatch { + expected: Box, + actual: Box, + }, + BatchTooLarge { + limit: usize, + }, + BatchInputTooLarge { + limit: usize, + }, + InputTooLarge { + limit: usize, + }, + TooManyTokens { + limit: usize, + }, + BackendFailure { + message: String, + }, +} + +impl TokenizerError { + pub fn code(&self) -> &'static str { + match self { + Self::StageZeroRequired => "stage_zero_required", + Self::UnsupportedStage { .. } => "unsupported_stage", + Self::IdentityUnavailable => "identity_unavailable", + Self::RuntimeUnavailable => "runtime_unavailable", + Self::IdentityMismatch { .. } => "identity_mismatch", + Self::BatchTooLarge { .. } => "batch_too_large", + Self::BatchInputTooLarge { .. } => "batch_input_too_large", + Self::InputTooLarge { .. } => "input_too_large", + Self::TooManyTokens { .. } => "too_many_tokens", + Self::BackendFailure { .. } => "backend_failure", + } + } +} + +impl fmt::Display for TokenizerError { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::StageZeroRequired => formatter.write_str("stage-zero tokenizer required"), + Self::UnsupportedStage { stage_index } => { + write!(formatter, "tokenizer is unavailable on stage {stage_index}") + } + Self::IdentityUnavailable => formatter.write_str("tokenizer identity unavailable"), + Self::RuntimeUnavailable => formatter.write_str("tokenizer runtime unavailable"), + Self::IdentityMismatch { .. } => formatter.write_str("tokenizer identity mismatch"), + Self::BatchTooLarge { limit } => { + write!(formatter, "tokenizer batch exceeds {limit} requests") + } + Self::BatchInputTooLarge { limit } => { + write!(formatter, "tokenizer batch input exceeds {limit} bytes") + } + Self::InputTooLarge { limit } => { + write!(formatter, "tokenizer input exceeds {limit} bytes") + } + Self::TooManyTokens { limit } => { + write!(formatter, "tokenizer output exceeds {limit} tokens") + } + Self::BackendFailure { message } => { + write!(formatter, "tokenizer backend failure: {message}") + } + } + } +} + +impl Error for TokenizerError {} + +#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] +pub struct TokenizeBatchItem { + pub request_index: usize, + pub result: Result, +} + +pub trait Tokenizer: Send + Sync { + fn identity(&self) -> &TokenizerIdentity; + + fn limits(&self) -> TokenizerLimits { + TokenizerLimits::default() + } + + fn tokenize_batch( + &self, + requests: &[TokenizeRequest], + ) -> Result, TokenizerError>; +} + +#[cfg(test)] +mod tests { + use super::*; + + fn identity() -> TokenizerIdentity { + TokenizerIdentity { + model_id: "model".to_owned(), + source_model_sha256: "a".repeat(64), + tokenizer_id: "gguf-source-sha256:".to_owned() + &"a".repeat(64), + tokenizer_version: Some("gguf-v1".to_owned()), + stage_index: 0, + serving_profile: Some("stage-zero".to_owned()), + } + } + + #[test] + fn request_defaults_are_safe_and_identity_is_serialized() { + let request: TokenizeRequest = serde_json::from_value(serde_json::json!({ + "expected_identity": identity(), + "text": "hello" + })) + .unwrap(); + assert_eq!(request.special_tokens, SpecialTokenPolicy::Omit); + assert!(!request.include_token_pieces); + assert_eq!( + serde_json::from_value::(serde_json::to_value(identity()).unwrap()) + .unwrap(), + identity() + ); + let serialized = serde_json::to_value(&request).unwrap(); + assert_eq!(serialized["add_special"], false); + let legacy: TokenizeRequest = serde_json::from_value(serde_json::json!({ + "expected_identity": identity(), + "text": "hello", + "add_special": true + })) + .unwrap(); + assert_eq!(legacy.special_tokens, SpecialTokenPolicy::Add); + } + + #[test] + fn error_codes_are_machine_readable() { + assert_eq!( + TokenizerError::BatchTooLarge { limit: 1 }.code(), + "batch_too_large" + ); + assert_eq!( + TokenizerError::UnsupportedStage { stage_index: 1 }.code(), + "unsupported_stage" + ); + } +} diff --git a/scripts/affected-crates.sh b/scripts/affected-crates.sh index ceffec188..fb95ea919 100755 --- a/scripts/affected-crates.sh +++ b/scripts/affected-crates.sh @@ -49,6 +49,7 @@ WORKSPACE_MEMBERS=( "model-hf" "model-resolver" "skippy-protocol" + "skippy-tokenizer" "skippy-coordinator" "skippy-topology" "skippy-cache" diff --git a/scripts/plan-clippy-batches.sh b/scripts/plan-clippy-batches.sh index d6c7bead2..9db5d00d9 100644 --- a/scripts/plan-clippy-batches.sh +++ b/scripts/plan-clippy-batches.sh @@ -49,6 +49,7 @@ WORKSPACE_MEMBERS=( "model-hf" "model-resolver" "skippy-protocol" + "skippy-tokenizer" "skippy-coordinator" "skippy-topology" "skippy-cache" diff --git a/scripts/publish-crates.sh b/scripts/publish-crates.sh index 32164a515..d7def3344 100755 --- a/scripts/publish-crates.sh +++ b/scripts/publish-crates.sh @@ -391,7 +391,8 @@ unpublished_registry_deps() { skippy-cache \ skippy-metrics \ skippy-protocol \ - skippy-runtime + skippy-runtime \ + skippy-tokenizer ;; mesh-native-serving-plugin-host) printf '%s\n' \ @@ -473,6 +474,7 @@ should_skip_initial_dry_run() { publish_crates=( mesh-llm-identity + skippy-tokenizer mesh-llm-protocol mesh-llm-routing mesh-llm-types