diff --git a/Cargo.lock b/Cargo.lock index 74b718f5b..ebd004e60 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11,6 +11,12 @@ dependencies = [ "memchr", ] +[[package]] +name = "allocator-api2" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923" + [[package]] name = "anstream" version = "1.0.0" @@ -426,6 +432,12 @@ version = "1.0.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" +[[package]] +name = "foldhash" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" + [[package]] name = "form_urlencoded" version = "1.2.2" @@ -597,6 +609,17 @@ dependencies = [ "tracing", ] +[[package]] +name = "hashbrown" +version = "0.15.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" +dependencies = [ + "allocator-api2", + "equivalent", + "foldhash", +] + [[package]] name = "hashbrown" version = "0.17.1" @@ -831,7 +854,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" dependencies = [ "equivalent", - "hashbrown", + "hashbrown 0.17.1", ] [[package]] @@ -912,6 +935,15 @@ version = "0.4.33" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" +[[package]] +name = "lru" +version = "0.12.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "234cf4f4a04dc1f57e24b96cc0cd600cf2af460d4161ac5ecdd0af8e1f3b2a38" +dependencies = [ + "hashbrown 0.15.5", +] + [[package]] name = "lru-slab" version = "0.1.2" @@ -1555,6 +1587,16 @@ version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" +[[package]] +name = "safetensors" +version = "0.4.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "44560c11236a6130a46ce36c836a62936dc81ebf8c36a37947423571be0e55b6" +dependencies = [ + "serde", + "serde_json", +] + [[package]] name = "scopeguard" version = "1.2.0" @@ -1728,10 +1770,12 @@ version = "0.1.0" dependencies = [ "async-trait", "futures", + "lru", "opentelemetry", "opentelemetry_sdk", "parking_lot", "rand 0.8.7", + "safetensors", "serde", "serde_json", "switchyard-llm-client", @@ -1752,9 +1796,13 @@ dependencies = [ "futures-util", "opentelemetry", "reqwest", + "safetensors", + "serde", "serde_json", "switchyard-protocol", "switchyard-translation", + "tempfile", + "thiserror 2.0.18", "tokio", "wiremock", ] @@ -1806,6 +1854,7 @@ dependencies = [ "opentelemetry_sdk", "prometheus", "rustls", + "safetensors", "serde", "serde_json", "switchyard-libsy", diff --git a/crates/libsy-llm-client/Cargo.toml b/crates/libsy-llm-client/Cargo.toml index 19b2470fe..800975452 100644 --- a/crates/libsy-llm-client/Cargo.toml +++ b/crates/libsy-llm-client/Cargo.toml @@ -18,9 +18,14 @@ reqwest.workspace = true async-trait.workspace = true futures-util.workspace = true opentelemetry = { version = "0.32", default-features = false, features = ["metrics"] } +safetensors = "0.4" +serde.workspace = true serde_json.workspace = true +thiserror.workspace = true +tokio.workspace = true [dev-dependencies] futures.workspace = true +tempfile = "3" tokio.workspace = true wiremock = "0.6" diff --git a/crates/libsy-llm-client/src/lib.rs b/crates/libsy-llm-client/src/lib.rs index 79fb84e3a..157f574c3 100644 --- a/crates/libsy-llm-client/src/lib.rs +++ b/crates/libsy-llm-client/src/lib.rs @@ -16,9 +16,14 @@ pub mod client; pub mod error; pub mod metrics; pub mod raw; +pub mod vllm_hidden_state_probe; pub use backend::{Backend, HttpBackendConfig}; pub use client::{ModelConfig, TranslatingLlmClient}; pub use error::{LlmClientError, Result}; pub use raw::RawResponse; pub use switchyard_translation::RawEventStream; +pub use vllm_hidden_state_probe::{ + HiddenStateFeatures, VllmHiddenStateProbe, VllmHiddenStateProbeConfig, + VllmHiddenStateProbeError, VllmHiddenStateProbeResult, +}; diff --git a/crates/libsy-llm-client/src/vllm_hidden_state_probe.rs b/crates/libsy-llm-client/src/vllm_hidden_state_probe.rs new file mode 100644 index 000000000..59c19ec88 --- /dev/null +++ b/crates/libsy-llm-client/src/vllm_hidden_state_probe.rs @@ -0,0 +1,1092 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! vLLM hidden-state extraction for learned prefill routing. +//! +//! The probe calls an OpenAI-compatible chat endpoint configured with vLLM's +//! `ExampleHiddenStatesConnector`, waits for the returned safetensors artifact, +//! reduces `[prompt_tokens, layers, hidden_size]` to one token-mean vector per +//! layer, and removes the consumed artifact. + +use std::fmt; +use std::fs::{File, OpenOptions}; +use std::io::Read; +use std::path::{Path, PathBuf}; +use std::time::{Duration, Instant, SystemTime}; + +use safetensors::{Dtype, SafeTensors}; +use serde::Deserialize; +use thiserror::Error; + +const ARTIFACT_WAIT_INTERVAL: Duration = Duration::from_millis(50); +const ARTIFACT_WAIT_TIMEOUT: Duration = Duration::from_secs(1); +const STALE_ARTIFACT_RETENTION: Duration = Duration::from_secs(5 * 60); + +/// Token-mean hidden-state features returned by a vLLM probe. +#[derive(Clone, Debug, PartialEq)] +pub struct HiddenStateFeatures { + /// Number of independently extracted hidden-state layers. + pub layer_count: usize, + /// Hidden width of each extracted layer. + pub hidden_size: usize, + /// Layer-major token-mean features. + pub values: Vec, +} + +/// Construction inputs for [`VllmHiddenStateProbe`]. +#[derive(Clone, Debug)] +pub struct VllmHiddenStateProbeConfig { + /// OpenAI-compatible vLLM base URL, normally ending in `/v1`. + pub base_url: String, + /// Probe model served by the vLLM endpoint. + pub model: String, + /// Dedicated shared directory used by `ExampleHiddenStatesConnector`. + pub hidden_states_dir: PathBuf, + /// Maximum duration of the HTTP request, including response-body decoding. + pub request_timeout: Duration, +} + +/// Failures returned while requesting or decoding vLLM hidden states. +#[derive(Debug, Error)] +pub enum VllmHiddenStateProbeError { + /// The probe cannot be constructed from its configuration. + #[error("invalid vLLM hidden-state probe configuration: {message}")] + Configuration { + /// Invalid field or invariant. + message: String, + }, + + /// The vLLM HTTP request exceeded its configured timeout. + #[error("vLLM hidden-state probe request timed out: {source}")] + Timeout { + /// Underlying HTTP timeout. + #[source] + source: reqwest::Error, + }, + + /// The vLLM HTTP request could not be completed. + #[error("vLLM hidden-state probe transport failed: {source}")] + Transport { + /// Underlying HTTP transport failure. + #[source] + source: reqwest::Error, + }, + + /// The vLLM endpoint returned a non-success status. + #[error("vLLM hidden-state probe returned HTTP {status}")] + Http { + /// Upstream HTTP status code. + status: u16, + }, + + /// The vLLM response did not contain the expected connector payload. + #[error("invalid vLLM hidden-state probe response: {message}")] + Response { + /// Response decoding or validation failure. + message: String, + }, + + /// The returned hidden-state artifact was unsafe or malformed. + #[error("vLLM hidden-state artifact error: {message}")] + Artifact { + /// Filesystem, safetensors, or feature validation failure. + message: String, + }, + + /// Tokio could not complete a blocking filesystem task. + #[error("vLLM hidden-state artifact task failed: {source}")] + BlockingTask { + /// Blocking task join failure. + #[source] + source: tokio::task::JoinError, + }, +} + +/// Result type for vLLM hidden-state probe operations. +pub type VllmHiddenStateProbeResult = Result; + +/// HTTP and filesystem client for vLLM prompt hidden-state extraction. +/// +/// The configured hidden-state directory must be dedicated to probe artifacts. +/// Each extraction performs a bounded stale-file sweep, and a returned artifact +/// is removed after reading even when tensor parsing fails. +pub struct VllmHiddenStateProbe { + completions_url: String, + model: String, + hidden_states_dir: PathBuf, + client: reqwest::Client, + artifact_wait_timeout: Duration, + stale_artifact_retention: Duration, +} + +impl VllmHiddenStateProbe { + /// Validates the probe configuration and constructs its bounded HTTP client. + pub fn new(config: VllmHiddenStateProbeConfig) -> VllmHiddenStateProbeResult { + let base_url = config.base_url.trim(); + if base_url.is_empty() { + return Err(configuration_error("base_url must not be empty")); + } + let model = config.model.trim(); + if model.is_empty() { + return Err(configuration_error("model must not be empty")); + } + if config.request_timeout.is_zero() { + return Err(configuration_error("request_timeout must be positive")); + } + + let hidden_states_dir = config.hidden_states_dir.canonicalize().map_err(|error| { + configuration_error(format!( + "hidden_states_dir {} is not accessible: {error}", + config.hidden_states_dir.display() + )) + })?; + if !hidden_states_dir.is_dir() { + return Err(configuration_error(format!( + "hidden_states_dir {} is not a directory", + hidden_states_dir.display() + ))); + } + let client = reqwest::Client::builder() + .timeout(config.request_timeout) + .build() + .map_err(map_reqwest_error)?; + + Ok(Self { + completions_url: completions_url(base_url), + model: model.to_string(), + hidden_states_dir, + client, + artifact_wait_timeout: ARTIFACT_WAIT_TIMEOUT, + stale_artifact_retention: STALE_ARTIFACT_RETENTION, + }) + } + + /// Requests and token-mean pools hidden states for one task instruction. + /// + /// The task is sent only as the single user message. Filesystem parsing and + /// cleanup run on Tokio's blocking pool. No task content is logged. + pub async fn extract(&self, task: &str) -> VllmHiddenStateProbeResult { + self.reap_stale_artifacts().await?; + + let response = self + .client + .post(&self.completions_url) + .json(&serde_json::json!({ + "model": self.model, + "messages": [{"role": "user", "content": task}], + "max_tokens": 1, + "kv_transfer_params": { + "include_output_tokens": false, + }, + })) + .send() + .await + .map_err(map_reqwest_error)?; + let status = response.status(); + if !status.is_success() { + return Err(VllmHiddenStateProbeError::Http { + status: status.as_u16(), + }); + } + let response: CompletionResponse = + response + .json() + .await + .map_err(|error| match map_reqwest_error(error) { + VllmHiddenStateProbeError::Timeout { source } => { + VllmHiddenStateProbeError::Timeout { source } + } + other => VllmHiddenStateProbeError::Response { + message: other.to_string(), + }, + })?; + let reported_path = response + .kv_transfer_params + .ok_or_else(|| response_error("missing kv_transfer_params"))? + .hidden_states_path; + let artifact_path = resolve_reported_path(&self.hidden_states_dir, &reported_path)?; + + let root = self.hidden_states_dir.clone(); + let artifact_wait_timeout = self.artifact_wait_timeout; + tokio::task::spawn_blocking(move || { + read_and_cleanup_hidden_states(&root, &artifact_path, artifact_wait_timeout) + }) + .await + .map_err(|source| VllmHiddenStateProbeError::BlockingTask { source })? + } + + async fn reap_stale_artifacts(&self) -> VllmHiddenStateProbeResult<()> { + let root = self.hidden_states_dir.clone(); + let retention = self.stale_artifact_retention; + tokio::task::spawn_blocking(move || cleanup_stale_artifacts(&root, retention)) + .await + .map_err(|source| VllmHiddenStateProbeError::BlockingTask { source })? + } +} + +impl fmt::Debug for VllmHiddenStateProbe { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("VllmHiddenStateProbe") + .field("completions_url", &self.completions_url) + .field("model", &self.model) + .field("hidden_states_dir", &self.hidden_states_dir) + .field("artifact_wait_timeout", &self.artifact_wait_timeout) + .field("stale_artifact_retention", &self.stale_artifact_retention) + .finish_non_exhaustive() + } +} + +#[derive(Deserialize)] +struct CompletionResponse { + kv_transfer_params: Option, +} + +#[derive(Deserialize)] +struct KvTransferParams { + hidden_states_path: String, +} + +fn completions_url(base_url: &str) -> String { + let base_url = base_url.trim_end_matches('/'); + if base_url.ends_with("/chat/completions") { + base_url.to_string() + } else { + format!("{base_url}/chat/completions") + } +} + +fn map_reqwest_error(error: reqwest::Error) -> VllmHiddenStateProbeError { + if error.is_timeout() { + VllmHiddenStateProbeError::Timeout { source: error } + } else { + VllmHiddenStateProbeError::Transport { source: error } + } +} + +fn configuration_error(message: impl Into) -> VllmHiddenStateProbeError { + VllmHiddenStateProbeError::Configuration { + message: message.into(), + } +} + +fn response_error(message: impl Into) -> VllmHiddenStateProbeError { + VllmHiddenStateProbeError::Response { + message: message.into(), + } +} + +fn artifact_error(message: impl Into) -> VllmHiddenStateProbeError { + VllmHiddenStateProbeError::Artifact { + message: message.into(), + } +} + +fn resolve_reported_path(root: &Path, reported: &str) -> VllmHiddenStateProbeResult { + if reported.trim().is_empty() { + return Err(response_error("hidden_states_path must not be empty")); + } + let reported = Path::new(reported); + if !has_safetensors_extension(reported) { + return Err(artifact_error(format!( + "hidden-state artifact must be a .safetensors file: {}", + reported.display() + ))); + } + let candidate = if reported.is_absolute() { + reported.to_path_buf() + } else { + root.join(reported) + }; + let parent = candidate + .parent() + .ok_or_else(|| artifact_error("hidden-state artifact has no parent directory"))?; + let canonical_parent = parent.canonicalize().map_err(|error| { + artifact_error(format!( + "hidden-state artifact parent {} is not accessible: {error}", + parent.display() + )) + })?; + if !canonical_parent.starts_with(root) { + return Err(artifact_error(format!( + "hidden-state artifact parent {} is outside configured directory {}", + canonical_parent.display(), + root.display() + ))); + } + let file_name = candidate + .file_name() + .ok_or_else(|| artifact_error("hidden-state artifact has no file name"))?; + Ok(canonical_parent.join(file_name)) +} + +fn validate_hidden_states_path(root: &Path, path: &Path) -> VllmHiddenStateProbeResult { + if !has_safetensors_extension(path) { + return Err(artifact_error(format!( + "hidden-state artifact must be a .safetensors file: {}", + path.display() + ))); + } + let actual = path.canonicalize().map_err(|error| { + artifact_error(format!( + "hidden-state artifact {} is not accessible: {error}", + path.display() + )) + })?; + if !actual.starts_with(root) { + return Err(artifact_error(format!( + "hidden-state artifact {} is outside configured directory {}", + actual.display(), + root.display() + ))); + } + if !has_safetensors_extension(&actual) { + return Err(artifact_error(format!( + "canonical hidden-state artifact must be a .safetensors file: {}", + actual.display() + ))); + } + let metadata = actual.metadata().map_err(|error| { + artifact_error(format!( + "hidden-state artifact metadata error for {}: {error}", + actual.display() + )) + })?; + if !metadata.is_file() { + return Err(artifact_error(format!( + "hidden-state artifact is not a regular file: {}", + actual.display() + ))); + } + Ok(actual) +} + +fn has_safetensors_extension(path: &Path) -> bool { + path.extension().and_then(|extension| extension.to_str()) == Some("safetensors") +} + +fn companion_lock_path(path: &Path) -> PathBuf { + let mut lock_path = path.as_os_str().to_os_string(); + lock_path.push(".lock"); + PathBuf::from(lock_path) +} + +fn open_synchronized_artifact( + path: &Path, + timeout: Duration, +) -> VllmHiddenStateProbeResult<(File, File)> { + let lock_path = companion_lock_path(path); + let deadline = Instant::now() + timeout; + loop { + match OpenOptions::new().read(true).open(&lock_path) { + Ok(lock_file) => { + let metadata = lock_file.metadata().map_err(|error| { + artifact_error(format!( + "hidden-state synchronization lock metadata error for {}: {error}", + lock_path.display() + )) + })?; + if !metadata.is_file() { + return Err(artifact_error(format!( + "hidden-state synchronization lock is not a regular file: {}", + lock_path.display() + ))); + } + match lock_file.try_lock_shared() { + Ok(()) => { + let artifact = File::open(path).map_err(|error| { + artifact_error(format!( + "hidden-state artifact open error for {} after writer completed: \ + {error}", + path.display() + )) + })?; + return Ok((artifact, lock_file)); + } + Err(std::fs::TryLockError::WouldBlock) => {} + Err(error) => { + return Err(artifact_error(format!( + "hidden-state synchronization lock error for {}: {error}", + lock_path.display() + ))); + } + } + } + Err(error) if error.kind() == std::io::ErrorKind::NotFound => {} + Err(error) => { + return Err(artifact_error(format!( + "hidden-state synchronization lock open error for {}: {error}", + lock_path.display() + ))); + } + } + let now = Instant::now(); + if now >= deadline { + return Err(artifact_error(format!( + "hidden-state artifact {} did not become readable within {} ms; ensure vLLM \ + use_synchronization_lock is enabled", + path.display(), + timeout.as_millis() + ))); + } + std::thread::sleep(ARTIFACT_WAIT_INTERVAL.min(deadline - now)); + } +} + +fn read_and_cleanup_hidden_states( + root: &Path, + path: &Path, + timeout: Duration, +) -> VllmHiddenStateProbeResult { + let (mut artifact, synchronization_lock) = open_synchronized_artifact(path, timeout)?; + let artifact_path = validate_hidden_states_path(root, path)?; + let mut bytes = Vec::new(); + let features = artifact + .read_to_end(&mut bytes) + .map_err(|error| { + artifact_error(format!( + "hidden-state artifact read error for {}: {error}", + artifact_path.display() + )) + }) + .and_then(|_| parse_hidden_state_features(&bytes)); + drop(artifact); + drop(synchronization_lock); + let cleanup = cleanup_artifact_files(&artifact_path); + match (features, cleanup) { + (Ok(features), Ok(())) => Ok(features), + (Ok(_), Err(cleanup_error)) => Err(cleanup_error), + (Err(feature_error), Ok(())) => Err(feature_error), + (Err(feature_error), Err(cleanup_error)) => Err(artifact_error(format!( + "{feature_error}; cleanup also failed: {cleanup_error}" + ))), + } +} + +fn cleanup_artifact_files(path: &Path) -> VllmHiddenStateProbeResult<()> { + let lock_path = companion_lock_path(path); + std::fs::remove_file(path).map_err(|error| { + artifact_error(format!( + "hidden-state artifact cleanup error for {}: {error}", + path.display() + )) + })?; + match std::fs::remove_file(&lock_path) { + Ok(()) => Ok(()), + Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()), + Err(error) => Err(artifact_error(format!( + "hidden-state synchronization lock cleanup error for {}: {error}", + lock_path.display() + ))), + } +} + +fn cleanup_stale_artifacts(root: &Path, retention: Duration) -> VllmHiddenStateProbeResult<()> { + let entries = std::fs::read_dir(root).map_err(|error| { + artifact_error(format!( + "failed to scan hidden-state directory {}: {error}", + root.display() + )) + })?; + let now = SystemTime::now(); + for entry in entries { + let entry = entry.map_err(|error| { + artifact_error(format!( + "failed to inspect hidden-state directory {}: {error}", + root.display() + )) + })?; + let path = entry.path(); + if !has_safetensors_extension(&path) { + continue; + } + let file_type = entry.file_type().map_err(|error| { + artifact_error(format!( + "failed to inspect hidden-state artifact {}: {error}", + path.display() + )) + })?; + if !file_type.is_file() { + continue; + } + let metadata = entry.metadata().map_err(|error| { + artifact_error(format!( + "failed to inspect hidden-state artifact {}: {error}", + path.display() + )) + })?; + let modified = metadata.modified().map_err(|error| { + artifact_error(format!( + "failed to read hidden-state artifact timestamp {}: {error}", + path.display() + )) + })?; + if now.duration_since(modified).unwrap_or_default() < retention { + continue; + } + let lock_path = companion_lock_path(&path); + let cleanup_lock = match OpenOptions::new().read(true).open(&lock_path) { + Ok(lock_file) => match lock_file.try_lock_shared() { + Ok(()) => lock_file, + Err(std::fs::TryLockError::WouldBlock) => continue, + Err(error) => { + return Err(artifact_error(format!( + "failed to lock stale hidden-state synchronization file {}: {error}", + lock_path.display() + ))); + } + }, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => { + let artifact = match OpenOptions::new().read(true).write(true).open(&path) { + Ok(file) => file, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => continue, + Err(error) => { + return Err(artifact_error(format!( + "failed to open stale hidden-state artifact {}: {error}", + path.display() + ))); + } + }; + match artifact.try_lock() { + Ok(()) => artifact, + Err(std::fs::TryLockError::WouldBlock) => continue, + Err(error) => { + return Err(artifact_error(format!( + "failed to lock stale hidden-state artifact {}: {error}", + path.display() + ))); + } + } + } + Err(error) => { + return Err(artifact_error(format!( + "failed to open stale hidden-state synchronization file {}: {error}", + lock_path.display() + ))); + } + }; + std::fs::remove_file(&path).map_err(|error| { + artifact_error(format!( + "failed to remove stale hidden-state artifact {}: {error}", + path.display() + )) + })?; + drop(cleanup_lock); + match std::fs::remove_file(&lock_path) { + Ok(()) => {} + Err(error) if error.kind() == std::io::ErrorKind::NotFound => {} + Err(error) => { + return Err(artifact_error(format!( + "failed to remove stale hidden-state synchronization file {}: {error}", + lock_path.display() + ))); + } + } + } + Ok(()) +} + +fn parse_hidden_state_features(bytes: &[u8]) -> VllmHiddenStateProbeResult { + let tensors = SafeTensors::deserialize(bytes) + .map_err(|error| artifact_error(format!("safetensors parse error: {error}")))?; + let hidden_states = tensors + .tensor("hidden_states") + .map_err(|error| artifact_error(format!("hidden_states tensor not found: {error}")))?; + let prompt_tokens = match hidden_states.shape() { + [prompt_tokens, _, _] => *prompt_tokens, + _ => { + return Err(artifact_error( + "expected hidden_states shape [prompt_tokens, layers, hidden_size]", + )); + } + }; + validate_token_ids(&tensors, prompt_tokens)?; + token_mean_per_layer( + hidden_states.data(), + hidden_states.dtype(), + hidden_states.shape(), + ) +} + +fn token_mean_per_layer( + data: &[u8], + dtype: Dtype, + shape: &[usize], +) -> VllmHiddenStateProbeResult { + if shape.len() != 3 { + return Err(artifact_error( + "expected hidden_states shape [prompt_tokens, layers, hidden_size]", + )); + } + let (prompt_tokens, layer_count, hidden_size) = (shape[0], shape[1], shape[2]); + if prompt_tokens == 0 { + return Err(artifact_error( + "hidden_states token dimension must be non-zero", + )); + } + if layer_count == 0 || hidden_size == 0 { + return Err(artifact_error( + "hidden_states layer and hidden dimensions must be non-zero", + )); + } + + let (element_size, decode): (usize, fn(&[u8]) -> f32) = match dtype { + Dtype::F32 => (size_of::(), decode_f32), + Dtype::BF16 => (size_of::(), decode_bf16), + other => { + return Err(artifact_error(format!( + "unsupported hidden_states dtype: {other:?}" + ))); + } + }; + let features_per_token = layer_count + .checked_mul(hidden_size) + .ok_or_else(|| artifact_error("hidden_states shape is too large"))?; + let bytes_per_token = features_per_token + .checked_mul(element_size) + .ok_or_else(|| artifact_error("hidden_states byte length is too large"))?; + let expected_bytes = prompt_tokens + .checked_mul(bytes_per_token) + .ok_or_else(|| artifact_error("hidden_states byte length is too large"))?; + if data.len() != expected_bytes { + return Err(artifact_error(format!( + "hidden_states byte length {} does not match shape byte length {expected_bytes}", + data.len() + ))); + } + + let mut pooled = vec![0.0f32; features_per_token]; + for token in data.chunks_exact(bytes_per_token) { + for (index, bytes) in token.chunks_exact(element_size).enumerate() { + accumulate(&mut pooled[index], decode(bytes))?; + } + } + let token_count = prompt_tokens as f32; + for value in &mut pooled { + *value /= token_count; + if !value.is_finite() { + return Err(artifact_error( + "hidden-state token mean produced a non-finite value", + )); + } + } + Ok(HiddenStateFeatures { + layer_count, + hidden_size, + values: pooled, + }) +} + +fn decode_f32(bytes: &[u8]) -> f32 { + f32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]) +} + +fn decode_bf16(bytes: &[u8]) -> f32 { + f32::from_bits(u32::from(u16::from_le_bytes([bytes[0], bytes[1]])) << 16) +} + +fn accumulate(sum: &mut f32, value: f32) -> VllmHiddenStateProbeResult<()> { + if !value.is_finite() { + return Err(artifact_error("hidden_states contains non-finite values")); + } + *sum += value; + if !sum.is_finite() { + return Err(artifact_error( + "hidden-state token accumulation produced a non-finite value", + )); + } + Ok(()) +} + +fn validate_token_ids( + tensors: &SafeTensors<'_>, + prompt_tokens: usize, +) -> VllmHiddenStateProbeResult<()> { + if !tensors + .names() + .iter() + .any(|name| name.as_str() == "token_ids") + { + return Ok(()); + } + let token_ids = tensors + .tensor("token_ids") + .map_err(|error| artifact_error(format!("token_ids tensor error: {error}")))?; + if token_ids.dtype() != Dtype::I64 { + return Err(artifact_error(format!( + "token_ids must use I64; got {:?}", + token_ids.dtype() + ))); + } + if token_ids.shape() != [prompt_tokens] { + return Err(artifact_error(format!( + "token_ids shape {:?} does not match hidden_states token count {prompt_tokens}", + token_ids.shape() + ))); + } + for bytes in token_ids.data().chunks_exact(size_of::()) { + let token_id = i64::from_le_bytes([ + bytes[0], bytes[1], bytes[2], bytes[3], bytes[4], bytes[5], bytes[6], bytes[7], + ]); + if token_id < 0 { + return Err(artifact_error("token_ids contains a negative token ID")); + } + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use std::collections::BTreeMap; + use std::error::Error; + + use safetensors::tensor::{serialize, TensorView}; + use serde_json::json; + use tempfile::TempDir; + use wiremock::matchers::{body_json, method, path}; + use wiremock::{Mock, MockServer, ResponseTemplate}; + + use super::*; + + type TestResult = Result>; + + fn f32_bytes(values: &[f32]) -> Vec { + values + .iter() + .flat_map(|value| value.to_le_bytes()) + .collect() + } + + fn bf16_bytes(values: &[f32]) -> Vec { + values + .iter() + .map(|value| (value.to_bits() >> 16) as u16) + .flat_map(|value| value.to_le_bytes()) + .collect() + } + + fn i64_bytes(values: &[i64]) -> Vec { + values + .iter() + .flat_map(|value| value.to_le_bytes()) + .collect() + } + + fn serialize_hidden_states( + dtype: Dtype, + shape: Vec, + hidden_data: &[u8], + token_ids: Option<(Dtype, Vec, Vec)>, + ) -> TestResult> { + let hidden_view = TensorView::new(dtype, shape, hidden_data)?; + let serialized = if let Some((dtype, shape, data)) = token_ids.as_ref() { + let token_view = TensorView::new(*dtype, shape.clone(), data)?; + serialize( + BTreeMap::from([("hidden_states", hidden_view), ("token_ids", token_view)]), + &None, + ) + } else { + serialize(BTreeMap::from([("hidden_states", hidden_view)]), &None) + }; + Ok(serialized?) + } + + fn valid_artifact_bytes() -> TestResult> { + serialize_hidden_states( + Dtype::F32, + vec![2, 2, 2], + &f32_bytes(&[ + 1.0, 3.0, 5.0, 7.0, // token 0, layers 0 and 1 + 2.0, 4.0, 6.0, 8.0, // token 1, layers 0 and 1 + ]), + None, + ) + } + + fn config(server: &MockServer, directory: &TempDir) -> VllmHiddenStateProbeConfig { + VllmHiddenStateProbeConfig { + base_url: format!("{}/v1", server.uri()), + model: "probe/model".into(), + hidden_states_dir: directory.path().to_path_buf(), + request_timeout: Duration::from_secs(1), + } + } + + #[test] + fn token_mean_pooling_supports_f32_and_bf16() -> TestResult { + let values = [1.0, 3.0, 5.0, 7.0, 2.0, 4.0, 6.0, 8.0]; + let f32_features = token_mean_per_layer(&f32_bytes(&values), Dtype::F32, &[2, 2, 2])?; + let bf16_features = token_mean_per_layer(&bf16_bytes(&values), Dtype::BF16, &[2, 2, 2])?; + + assert_eq!(f32_features.layer_count, 2); + assert_eq!(f32_features.hidden_size, 2); + assert_eq!(f32_features.values, vec![1.5, 3.5, 5.5, 7.5]); + assert_eq!(bf16_features, f32_features); + Ok(()) + } + + #[test] + fn token_mean_pooling_rejects_invalid_layout_dtype_and_values() -> TestResult { + let shape_error = token_mean_per_layer(&[], Dtype::F32, &[1, 2]) + .err() + .ok_or_else(|| artifact_error("invalid shape should fail"))?; + assert!(shape_error + .to_string() + .contains("expected hidden_states shape")); + + let dtype_error = token_mean_per_layer(&i64_bytes(&[1]), Dtype::I64, &[1, 1, 1]) + .err() + .ok_or_else(|| artifact_error("invalid dtype should fail"))?; + assert!(dtype_error.to_string().contains("unsupported")); + + let value_error = token_mean_per_layer(&f32_bytes(&[f32::NAN]), Dtype::F32, &[1, 1, 1]) + .err() + .ok_or_else(|| artifact_error("non-finite value should fail"))?; + assert!(value_error.to_string().contains("non-finite")); + Ok(()) + } + + #[test] + fn optional_token_ids_are_validated() -> TestResult { + let valid = serialize_hidden_states( + Dtype::F32, + vec![2, 1, 2], + &f32_bytes(&[1.0, 2.0, 3.0, 4.0]), + Some((Dtype::I64, vec![2], i64_bytes(&[101, 102]))), + )?; + assert_eq!(parse_hidden_state_features(&valid)?.values, vec![2.0, 3.0]); + + let negative = serialize_hidden_states( + Dtype::F32, + vec![2, 1, 2], + &f32_bytes(&[1.0, 2.0, 3.0, 4.0]), + Some((Dtype::I64, vec![2], i64_bytes(&[101, -1]))), + )?; + let error = parse_hidden_state_features(&negative) + .err() + .ok_or_else(|| artifact_error("negative token ID should fail"))?; + assert!(error.to_string().contains("negative token ID")); + Ok(()) + } + + #[tokio::test] + async fn extract_sends_exact_contract_and_removes_artifact() -> TestResult { + let server = MockServer::start().await; + let directory = TempDir::new()?; + let artifact_path = directory.path().join("hidden.safetensors"); + std::fs::write(&artifact_path, valid_artifact_bytes()?)?; + std::fs::write(companion_lock_path(&artifact_path), b"")?; + let expected_body = json!({ + "model": "probe/model", + "messages": [{"role": "user", "content": "Explain the failure."}], + "max_tokens": 1, + "kv_transfer_params": { + "include_output_tokens": false, + }, + }); + Mock::given(method("POST")) + .and(path("/v1/chat/completions")) + .and(body_json(expected_body)) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "kv_transfer_params": { + "hidden_states_path": artifact_path, + }, + }))) + .mount(&server) + .await; + let probe = VllmHiddenStateProbe::new(config(&server, &directory))?; + + let features = probe.extract("Explain the failure.").await?; + + assert_eq!(features.values, vec![1.5, 3.5, 5.5, 7.5]); + assert!(!artifact_path.exists()); + assert!(!companion_lock_path(&artifact_path).exists()); + Ok(()) + } + + #[tokio::test] + async fn malformed_artifact_is_removed() -> TestResult { + let server = MockServer::start().await; + let directory = TempDir::new()?; + let artifact_path = directory.path().join("malformed.safetensors"); + std::fs::write(&artifact_path, b"not safetensors")?; + std::fs::write(companion_lock_path(&artifact_path), b"")?; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "kv_transfer_params": { + "hidden_states_path": artifact_path, + }, + }))) + .mount(&server) + .await; + let probe = VllmHiddenStateProbe::new(config(&server, &directory))?; + + let error = probe + .extract("task") + .await + .err() + .ok_or_else(|| artifact_error("malformed artifact should fail"))?; + + assert!(error.to_string().contains("safetensors parse error")); + assert!(!artifact_path.exists()); + assert!(!companion_lock_path(&artifact_path).exists()); + Ok(()) + } + + #[tokio::test] + async fn path_outside_configured_directory_is_rejected_without_cleanup() -> TestResult { + let server = MockServer::start().await; + let directory = TempDir::new()?; + let outside = TempDir::new()?; + let artifact_path = outside.path().join("outside.safetensors"); + std::fs::write(&artifact_path, valid_artifact_bytes()?)?; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "kv_transfer_params": { + "hidden_states_path": artifact_path, + }, + }))) + .mount(&server) + .await; + let probe = VllmHiddenStateProbe::new(config(&server, &directory))?; + + let error = probe + .extract("task") + .await + .err() + .ok_or_else(|| artifact_error("outside path should fail"))?; + + assert!(error.to_string().contains("outside configured directory")); + assert!(artifact_path.exists()); + Ok(()) + } + + #[tokio::test] + async fn response_and_http_failures_are_typed() -> TestResult { + let missing_server = MockServer::start().await; + let directory = TempDir::new()?; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({}))) + .mount(&missing_server) + .await; + let probe = VllmHiddenStateProbe::new(config(&missing_server, &directory))?; + assert!(matches!( + probe.extract("task").await, + Err(VllmHiddenStateProbeError::Response { .. }) + )); + + let unavailable_server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(503)) + .mount(&unavailable_server) + .await; + let probe = VllmHiddenStateProbe::new(config(&unavailable_server, &directory))?; + assert!(matches!( + probe.extract("task").await, + Err(VllmHiddenStateProbeError::Http { status: 503 }) + )); + Ok(()) + } + + #[tokio::test] + async fn request_timeout_is_enforced() -> TestResult { + let server = MockServer::start().await; + let directory = TempDir::new()?; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(200).set_delay(Duration::from_millis(100))) + .mount(&server) + .await; + let mut config = config(&server, &directory); + config.request_timeout = Duration::from_millis(10); + let probe = VllmHiddenStateProbe::new(config)?; + + assert!(matches!( + probe.extract("task").await, + Err(VllmHiddenStateProbeError::Timeout { .. }) + )); + Ok(()) + } + + #[test] + fn stale_reaper_removes_only_unlocked_safetensors_files() -> TestResult { + let directory = TempDir::new()?; + let stale = directory.path().join("stale.safetensors"); + let locked = directory.path().join("locked.safetensors"); + let locked_path = companion_lock_path(&locked); + let unrelated = directory.path().join("keep.txt"); + std::fs::write(&stale, b"stale")?; + std::fs::write(&locked, b"locked")?; + std::fs::write(&locked_path, b"")?; + std::fs::write(&unrelated, b"keep")?; + let locked_file = OpenOptions::new() + .read(true) + .write(true) + .open(&locked_path)?; + locked_file.lock()?; + + cleanup_stale_artifacts(directory.path(), Duration::ZERO)?; + + assert!(!stale.exists()); + assert!(locked.exists()); + assert!(locked_path.exists()); + assert!(unrelated.exists()); + drop(locked_file); + cleanup_stale_artifacts(directory.path(), Duration::ZERO)?; + assert!(!locked.exists()); + assert!(!locked_path.exists()); + Ok(()) + } + + #[test] + fn artifact_reader_waits_for_connector_lock_and_cleans_both_files() -> TestResult { + let directory = TempDir::new()?; + let root = directory.path().canonicalize()?; + let artifact_path = directory.path().join("hidden.safetensors"); + let lock_path = companion_lock_path(&artifact_path); + std::fs::write(&artifact_path, valid_artifact_bytes()?)?; + std::fs::write(&lock_path, b"")?; + let writer_lock = OpenOptions::new().read(true).write(true).open(&lock_path)?; + writer_lock.lock()?; + let reader_path = artifact_path.clone(); + let reader = std::thread::spawn(move || { + read_and_cleanup_hidden_states(&root, &reader_path, Duration::from_secs(1)) + }); + + std::thread::sleep(Duration::from_millis(20)); + assert!(!reader.is_finished()); + drop(writer_lock); + let features = reader + .join() + .map_err(|_| artifact_error("artifact reader test thread panicked"))??; + + assert_eq!(features.values, vec![1.5, 3.5, 5.5, 7.5]); + assert!(!artifact_path.exists()); + assert!(!lock_path.exists()); + Ok(()) + } + + #[test] + fn configuration_rejects_invalid_values() -> TestResult { + let directory = TempDir::new()?; + let valid = VllmHiddenStateProbeConfig { + base_url: "https://example.test/v1".into(), + model: "probe/model".into(), + hidden_states_dir: directory.path().to_path_buf(), + request_timeout: Duration::from_secs(1), + }; + let mut invalid = valid.clone(); + invalid.model = " ".into(); + assert!(matches!( + VllmHiddenStateProbe::new(invalid), + Err(VllmHiddenStateProbeError::Configuration { .. }) + )); + + let mut invalid = valid; + invalid.request_timeout = Duration::ZERO; + assert!(matches!( + VllmHiddenStateProbe::new(invalid), + Err(VllmHiddenStateProbeError::Configuration { .. }) + )); + Ok(()) + } +} diff --git a/crates/libsy/Cargo.toml b/crates/libsy/Cargo.toml index c9e58e186..333d3bf4a 100644 --- a/crates/libsy/Cargo.toml +++ b/crates/libsy/Cargo.toml @@ -16,11 +16,13 @@ async-trait.workspace = true serde.workspace = true serde_json.workspace = true futures.workspace = true +lru = "0.12" # Metrics-only OTel API: instruments record through the host-installed global # meter provider. Pinned to the 0.32 line used across the workspace. opentelemetry = { version = "0.32", default-features = false, features = ["metrics"] } parking_lot.workspace = true rand.workspace = true +safetensors = "0.4" switchyard-protocol.workspace = true thiserror.workspace = true tokio.workspace = true diff --git a/crates/libsy/src/algorithms.rs b/crates/libsy/src/algorithms.rs index b36b542f8..78436cc54 100644 --- a/crates/libsy/src/algorithms.rs +++ b/crates/libsy/src/algorithms.rs @@ -10,12 +10,17 @@ mod fall_through; pub mod llm_class; pub mod noop; pub mod passthrough; +pub mod prefill_probe; pub mod rand; pub use fall_through::{FallThrough, FallThroughDecision}; pub use llm_class::{LlmTaskClassifier, TaskClassifierConfig}; pub use noop::{Noop, NoopDecision}; pub use passthrough::{Passthrough, PassthroughDecision}; +pub use prefill_probe::{ + PrefillFeatures, PrefillProbe, PrefillProbeClassifier, PrefillProbeClassifierConfig, + DEFAULT_PREFILL_PROBE_CACHE_CAPACITY, +}; pub use rand::{Random, RandomClassifier, RandomDecision}; pub use util::{AffinityRouter, SubagentOverride}; diff --git a/crates/libsy/src/algorithms/prefill_probe.rs b/crates/libsy/src/algorithms/prefill_probe.rs new file mode 100644 index 000000000..fab3cdb10 --- /dev/null +++ b/crates/libsy/src/algorithms/prefill_probe.rs @@ -0,0 +1,600 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Learned task routing from prompt hidden-state features. +//! +//! The classifier owns checkpoint inference, policy, and bounded task-level +//! decisions. A [`PrefillProbe`] supplies features without coupling `libsy` to +//! an HTTP client, provider SDK, or hidden-state transport. + +use std::collections::hash_map::RandomState; +use std::fmt; +use std::hash::BuildHasher; +use std::num::NonZeroUsize; +use std::path::PathBuf; +use std::sync::Arc; + +use async_trait::async_trait; +use lru::LruCache; +use parking_lot::Mutex; +use switchyard_protocol::Role; + +use self::artifact::InferenceArtifact; +use self::policy::{CostAwareRoutingPolicy, PrefillTier}; +use crate::{Classification, Classifier, Driver, LibsyError, Request, Result, Score}; + +mod artifact; +mod policy; + +const TERMINUS_TASK_DESCRIPTION_HEADER: &str = "Task Description:\n"; +const TERMINUS_TERMINAL_STATE_HEADER: &str = "\n\nCurrent terminal state:\n"; + +/// Default maximum number of successful task decisions retained by one classifier. +pub const DEFAULT_PREFILL_PROBE_CACHE_CAPACITY: usize = 4_096; + +/// Token-mean hidden-state features produced by a prefill probe. +#[derive(Clone, Debug, PartialEq)] +pub struct PrefillFeatures { + /// Number of independently extracted hidden-state layers. + pub layer_count: usize, + /// Hidden width of each extracted layer. + pub hidden_size: usize, + /// Layer-major token-mean features. + pub values: Vec, +} + +impl PrefillFeatures { + /// Creates one feature vector with its source layout. + pub fn new(layer_count: usize, hidden_size: usize, values: Vec) -> Self { + Self { + layer_count, + hidden_size, + values, + } + } +} + +/// Supplies prompt hidden-state features to [`PrefillProbeClassifier`]. +/// +/// Implementations own transport concerns such as endpoint timeouts and +/// temporary-artifact lifecycle. Errors are treated as an unavailable routing +/// optimization: the classifier selects strong without caching the failure. +#[async_trait] +pub trait PrefillProbe: Send + Sync { + /// Extracts token-mean features for one task instruction. + async fn extract(&self, task: &str) -> Result; +} + +/// Construction inputs for [`PrefillProbeClassifier`]. +#[derive(Clone, Debug)] +pub struct PrefillProbeClassifierConfig { + /// Probe model whose hidden-state layout matches the checkpoint metadata. + pub probe_model: String, + /// Directory containing `router.json` and `router.safetensors`. + pub checkpoint_dir: PathBuf, + /// Checkpoint output head corresponding to the strong completion target. + pub strong_checkpoint_head: String, + /// Checkpoint output head corresponding to the weak completion target. + pub weak_checkpoint_head: String, + /// Semantic name returned when the strong tier is selected. + pub strong_target: String, + /// Semantic name returned when the weak tier is selected. + pub weak_target: String, + /// Correctness weight in the cost-aware policy. + pub lambda: f64, + /// Non-negative weak-target cost in the same units as `strong_cost`. + pub weak_cost: f64, + /// Non-negative strong-target cost in the same units as `weak_cost`. + pub strong_cost: f64, + /// Maximum successful task decisions retained in memory. + pub cache_capacity: usize, +} + +struct LearnedRouting { + artifact: InferenceArtifact, + weak_head_index: usize, + strong_head_index: usize, + policy: CostAwareRoutingPolicy, +} + +impl LearnedRouting { + fn select(&self, features: PrefillFeatures) -> Result { + if features.layer_count != self.artifact.layer_count() { + return Err(inference_error(format!( + "feature layer count {} does not match checkpoint layer count {}", + features.layer_count, + self.artifact.layer_count(), + ))); + } + if features.hidden_size != self.artifact.hidden_size() { + return Err(inference_error(format!( + "feature hidden size {} does not match checkpoint hidden size {}", + features.hidden_size, + self.artifact.hidden_size(), + ))); + } + if features.values.len() != self.artifact.raw_feature_dim() { + return Err(inference_error(format!( + "feature length {} does not match checkpoint raw_feature_dim {}", + features.values.len(), + self.artifact.raw_feature_dim(), + ))); + } + + let projected = self.artifact.project(&features.values)?; + let logits = self.artifact.ensemble_logits(&projected)?; + let probabilities = self.artifact.ensemble_probabilities(&logits)?; + let weak_probability = probabilities.get(self.weak_head_index).ok_or_else(|| { + inference_error(format!( + "weak checkpoint head index {} is outside prediction length {}", + self.weak_head_index, + probabilities.len(), + )) + })?; + let strong_probability = probabilities.get(self.strong_head_index).ok_or_else(|| { + inference_error(format!( + "strong checkpoint head index {} is outside prediction length {}", + self.strong_head_index, + probabilities.len(), + )) + })?; + self.policy + .select(f64::from(*weak_probability), f64::from(*strong_probability)) + } +} + +/// Classifies a task as strong or weak from learned prompt features. +/// +/// Successful decisions are cached under process-randomized hashes, so raw +/// task text is not retained. The LRU bound prevents task cardinality from +/// growing memory without limit. Probe and inference failures select strong and +/// are deliberately not cached. +pub struct PrefillProbeClassifier { + probe: Arc, + routing: Arc, + strong_target: String, + weak_target: String, + decision_cache: Mutex>, + cache_hasher: RandomState, +} + +impl PrefillProbeClassifier { + /// Loads the learned checkpoint and constructs a transport-independent classifier. + pub fn new(config: PrefillProbeClassifierConfig, probe: Arc) -> Result { + let cache_capacity = validate_config(&config)?; + let artifact = InferenceArtifact::load(&config.checkpoint_dir, &config.probe_model)?; + Self::from_artifact(config, probe, artifact, cache_capacity) + } + + fn from_artifact( + config: PrefillProbeClassifierConfig, + probe: Arc, + artifact: InferenceArtifact, + cache_capacity: NonZeroUsize, + ) -> Result { + let strong_head_index = checkpoint_head_index( + &artifact, + "strong_checkpoint_head", + &config.strong_checkpoint_head, + )?; + let weak_head_index = checkpoint_head_index( + &artifact, + "weak_checkpoint_head", + &config.weak_checkpoint_head, + )?; + if strong_head_index == weak_head_index { + return Err(config_error( + "strong_checkpoint_head and weak_checkpoint_head must map to distinct outputs", + )); + } + + Ok(Self { + probe, + routing: Arc::new(LearnedRouting { + artifact, + weak_head_index, + strong_head_index, + policy: CostAwareRoutingPolicy::new( + config.lambda, + config.weak_cost, + config.strong_cost, + )?, + }), + strong_target: config.strong_target, + weak_target: config.weak_target, + decision_cache: Mutex::new(LruCache::new(cache_capacity)), + cache_hasher: RandomState::new(), + }) + } + + async fn select_for_task(&self, task: &str) -> String { + let cache_key = self.cache_hasher.hash_one(task); + if let Some(target) = self.decision_cache.lock().get(&cache_key).cloned() { + return target; + } + + let result = match self.probe.extract(task).await { + Ok(features) => { + let routing = Arc::clone(&self.routing); + tokio::task::spawn_blocking(move || routing.select(features)) + .await + .map_err(|error| inference_error(format!("inference task failed: {error}"))) + .and_then(|result| result) + } + Err(error) => Err(error), + }; + + match result { + Ok(tier) => { + let target = match tier { + PrefillTier::Weak => self.weak_target.clone(), + PrefillTier::Strong => self.strong_target.clone(), + }; + self.decision_cache.lock().put(cache_key, target.clone()); + target + } + Err(error) => { + tracing::warn!( + target: "libsy", + error = %error, + fallback_target = %self.strong_target, + "prefill probe unavailable; using uncached strong fallback" + ); + self.strong_target.clone() + } + } + } + + fn classification(&self, target: String) -> Classification { + Classification::Scores(vec![Score { + confidence: 1.0, + target, + }]) + } +} + +impl fmt::Debug for PrefillProbeClassifier { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("PrefillProbeClassifier") + .field("strong_target", &self.strong_target) + .field("weak_target", &self.weak_target) + .field("cache_capacity", &self.decision_cache.lock().cap()) + .finish_non_exhaustive() + } +} + +#[async_trait] +impl Classifier for PrefillProbeClassifier +where + S: Send + 'static, +{ + fn routing_tier(&self, selected_model: &str) -> Option<&'static str> { + if selected_model == self.weak_target { + Some("weak") + } else if selected_model == self.strong_target { + Some("strong") + } else { + None + } + } + + async fn score( + &self, + _state: &mut S, + request: &mut Request, + _driver: Option<&Driver>, + ) -> Result { + let Some(task) = probe_input(request) else { + tracing::warn!( + target: "libsy", + fallback_target = %self.strong_target, + "prefill probe request has no text user instruction; using strong fallback" + ); + return Ok(self.classification(self.strong_target.clone())); + }; + Ok(self.classification(self.select_for_task(&task).await)) + } +} + +fn validate_config(config: &PrefillProbeClassifierConfig) -> Result { + for (field, value) in [ + ("probe_model", config.probe_model.as_str()), + ( + "strong_checkpoint_head", + config.strong_checkpoint_head.as_str(), + ), + ("weak_checkpoint_head", config.weak_checkpoint_head.as_str()), + ("strong_target", config.strong_target.as_str()), + ("weak_target", config.weak_target.as_str()), + ] { + if value.trim().is_empty() { + return Err(config_error(format!("{field} must not be empty"))); + } + } + if config.strong_target == config.weak_target { + return Err(config_error( + "strong_target and weak_target must be distinct", + )); + } + NonZeroUsize::new(config.cache_capacity) + .ok_or_else(|| config_error("cache_capacity must be positive")) +} + +fn checkpoint_head_index( + artifact: &InferenceArtifact, + field: &str, + checkpoint_head: &str, +) -> Result { + artifact + .output_names() + .iter() + .position(|name| name == checkpoint_head) + .ok_or_else(|| { + config_error(format!( + "{field} {checkpoint_head:?} is not present in checkpoint output_names {:?}", + artifact.output_names(), + )) + }) +} + +/// Returns the first text-bearing user message, reduced to a benchmark task when recognized. +fn probe_input(request: &Request) -> Option { + let instruction = request + .llm_request + .messages + .iter() + .filter(|message| message.role == Role::User) + .find_map(|message| message.text_content("").filter(|text| !text.is_empty()))?; + Some( + terminus_task_instruction(&instruction) + .unwrap_or(&instruction) + .to_owned(), + ) +} + +fn terminus_task_instruction(instruction: &str) -> Option<&str> { + let (_, task_and_terminal) = instruction.split_once(TERMINUS_TASK_DESCRIPTION_HEADER)?; + let (task, _) = task_and_terminal.split_once(TERMINUS_TERMINAL_STATE_HEADER)?; + (!task.is_empty()).then_some(task) +} + +fn config_error(message: impl Into) -> LibsyError { + LibsyError::AlgorithmError { + message: format!("invalid prefill-probe config: {}", message.into()), + } +} + +fn inference_error(message: impl Into) -> LibsyError { + LibsyError::AlgorithmError { + message: format!("prefill-probe inference error: {}", message.into()), + } +} + +#[cfg(test)] +mod tests { + use std::collections::VecDeque; + + use parking_lot::Mutex; + use switchyard_protocol::{LlmRequest, Message}; + + use super::*; + + enum ProbeResult { + Features(PrefillFeatures), + Failure, + } + + struct RecordingProbe { + results: Mutex>, + inputs: Mutex>, + } + + impl RecordingProbe { + fn new(results: impl IntoIterator) -> Self { + Self { + results: Mutex::new(results.into_iter().collect()), + inputs: Mutex::new(Vec::new()), + } + } + + fn inputs(&self) -> Vec { + self.inputs.lock().clone() + } + } + + #[async_trait] + impl PrefillProbe for RecordingProbe { + async fn extract(&self, task: &str) -> Result { + self.inputs.lock().push(task.to_string()); + match self.results.lock().pop_front() { + Some(ProbeResult::Features(features)) => Ok(features), + Some(ProbeResult::Failure) => Err(inference_error("test probe failure")), + None => Err(inference_error("test probe has no result")), + } + } + } + + fn features() -> PrefillFeatures { + PrefillFeatures::new(2, 2, vec![0.0; 4]) + } + + fn config(cache_capacity: usize) -> PrefillProbeClassifierConfig { + PrefillProbeClassifierConfig { + probe_model: "probe/model".into(), + checkpoint_dir: "/unused/test/checkpoint".into(), + strong_checkpoint_head: "opus-4.7".into(), + weak_checkpoint_head: "nemotron-3-super".into(), + strong_target: "strong/model".into(), + weak_target: "weak/model".into(), + lambda: 1.0, + weak_cost: 0.0, + strong_cost: 1.0, + cache_capacity, + } + } + + fn classifier( + probe: Arc, + cache_capacity: usize, + ) -> Result { + let config = config(cache_capacity); + let capacity = validate_config(&config)?; + PrefillProbeClassifier::from_artifact( + config, + probe, + InferenceArtifact::with_test_probabilities([0.1, 0.8, 0.2, 0.1]), + capacity, + ) + } + + fn request(messages: Vec) -> Request { + Request { + llm_request: LlmRequest { + model: Some("auto".into()), + messages, + ..LlmRequest::default() + }, + raw_request: None, + metadata: None, + } + } + + async fn selected( + classifier: &PrefillProbeClassifier, + request: &mut Request, + ) -> Result { + classifier + .score(&mut (), request, None) + .await? + .argmax(false)? + .map(|score| score.target) + .ok_or_else(|| inference_error("classifier abstained")) + } + + #[tokio::test] + async fn successful_decision_is_cached_by_task_hash() -> Result<()> { + let probe = Arc::new(RecordingProbe::new([ProbeResult::Features(features())])); + let classifier = classifier(probe.clone(), 2)?; + let mut first = request(vec![Message::text(Role::User, "same task")]); + let mut second = request(vec![Message::text(Role::User, "same task")]); + + assert_eq!(selected(&classifier, &mut first).await?, "weak/model"); + assert_eq!(selected(&classifier, &mut second).await?, "weak/model"); + assert_eq!(probe.inputs(), ["same task"]); + assert_eq!(classifier.decision_cache.lock().len(), 1); + Ok(()) + } + + #[tokio::test] + async fn cache_evicts_at_capacity() -> Result<()> { + let probe = Arc::new(RecordingProbe::new([ + ProbeResult::Features(features()), + ProbeResult::Features(features()), + ProbeResult::Features(features()), + ])); + let classifier = classifier(probe.clone(), 1)?; + + for task in ["first task", "second task", "first task"] { + let mut request = request(vec![Message::text(Role::User, task)]); + assert_eq!(selected(&classifier, &mut request).await?, "weak/model"); + } + + assert_eq!(probe.inputs(), ["first task", "second task", "first task"]); + assert_eq!(classifier.decision_cache.lock().len(), 1); + Ok(()) + } + + #[tokio::test] + async fn probe_failure_falls_back_to_strong_without_caching() -> Result<()> { + let probe = Arc::new(RecordingProbe::new([ + ProbeResult::Failure, + ProbeResult::Features(features()), + ])); + let classifier = classifier(probe.clone(), 2)?; + let mut first = request(vec![Message::text(Role::User, "retry task")]); + let mut retry = first.clone(); + + assert_eq!(selected(&classifier, &mut first).await?, "strong/model"); + assert_eq!(selected(&classifier, &mut retry).await?, "weak/model"); + assert_eq!(probe.inputs(), ["retry task", "retry task"]); + Ok(()) + } + + #[tokio::test] + async fn malformed_features_fall_back_to_strong_without_caching() -> Result<()> { + let malformed = PrefillFeatures::new(1, 4, vec![0.0; 4]); + let probe = Arc::new(RecordingProbe::new([ + ProbeResult::Features(malformed), + ProbeResult::Features(features()), + ])); + let classifier = classifier(probe.clone(), 2)?; + let mut first = request(vec![Message::text(Role::User, "retry shape")]); + let mut retry = first.clone(); + + assert_eq!(selected(&classifier, &mut first).await?, "strong/model"); + assert_eq!(selected(&classifier, &mut retry).await?, "weak/model"); + assert_eq!(probe.inputs(), ["retry shape", "retry shape"]); + Ok(()) + } + + #[tokio::test] + async fn terminus_envelope_sends_only_task_text() -> Result<()> { + let probe = Arc::new(RecordingProbe::new([ProbeResult::Features(features())])); + let classifier = classifier(probe.clone(), 2)?; + let mut request = request(vec![Message::text( + Role::User, + concat!( + "\nTask Description:\n", + "repair the package", + "\n\nCurrent terminal state:\n", + "terminal output\n" + ), + )]); + + assert_eq!(selected(&classifier, &mut request).await?, "weak/model"); + assert_eq!(probe.inputs(), ["repair the package"]); + Ok(()) + } + + #[tokio::test] + async fn missing_user_text_uses_strong_without_probing() -> Result<()> { + let probe = Arc::new(RecordingProbe::new([])); + let classifier = classifier(probe.clone(), 2)?; + let mut request = request(vec![Message::text(Role::System, "system only")]); + + assert_eq!(selected(&classifier, &mut request).await?, "strong/model"); + assert!(probe.inputs().is_empty()); + Ok(()) + } + + #[test] + fn config_rejects_invalid_targets_heads_and_capacity() -> Result<()> { + let mut invalid = config(1); + invalid.weak_target = invalid.strong_target.clone(); + let error = validate_config(&invalid) + .err() + .ok_or_else(|| config_error("duplicate targets should fail"))?; + assert!(error.to_string().contains("must be distinct")); + + let mut invalid = config(0); + let error = validate_config(&invalid) + .err() + .ok_or_else(|| config_error("zero cache capacity should fail"))?; + assert!(error.to_string().contains("cache_capacity")); + + invalid.cache_capacity = 1; + invalid.weak_checkpoint_head = "missing".into(); + let capacity = validate_config(&invalid)?; + let error = PrefillProbeClassifier::from_artifact( + invalid, + Arc::new(RecordingProbe::new([])), + InferenceArtifact::with_test_probabilities([0.1, 0.8, 0.2, 0.1]), + capacity, + ) + .err() + .ok_or_else(|| config_error("missing checkpoint head should fail"))?; + assert!(error.to_string().contains("output_names")); + Ok(()) + } +} diff --git a/crates/libsy/src/algorithms/prefill_probe/artifact.rs b/crates/libsy/src/algorithms/prefill_probe/artifact.rs new file mode 100644 index 000000000..4de303b14 --- /dev/null +++ b/crates/libsy/src/algorithms/prefill_probe/artifact.rs @@ -0,0 +1,871 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Loading, validation, and CPU inference for learned prefill-probe artifacts. + +use std::collections::{BTreeMap, BTreeSet}; +use std::path::Path; + +use safetensors::{Dtype, SafeTensors}; +use serde::Deserialize; + +use crate::{LibsyError, Result}; + +const METADATA_FILE: &str = "router.json"; +const TENSOR_FILE: &str = "router.safetensors"; +const FORMAT_VERSION: u64 = 1; +const TRAINING_MODE: &str = "single_pca_block"; +const REPRESENTATION: &str = "token_mean_per_layer_concat"; +const PCA_DIM: usize = 200; +const TRUNK_HIDDEN: [usize; 2] = [256, 128]; +const ENSEMBLE_SIZE: usize = 5; +const OUTPUT_NAMES: [&str; 4] = ["qwen-122b", "nemotron-3-super", "opus-4.7", "gpt-5.5"]; +const PROBABILITY_LINK: &str = "independent_sigmoid"; +const ENSEMBLE_REDUCTION: &str = "probability_mean"; + +/// Immutable learned-router metadata and decoded tensors. +pub(super) struct InferenceArtifact { + metadata: ArtifactMetadata, + tensors: BTreeMap>, +} + +impl InferenceArtifact { + /// Loads and validates an exported artifact against the configured probe model. + pub(super) fn load(directory: impl AsRef, probe_model: &str) -> Result { + let directory = directory.as_ref(); + let metadata_path = directory.join(METADATA_FILE); + let metadata_json = std::fs::read_to_string(&metadata_path).map_err(|error| { + invalid_artifact(format!( + "failed to read {}: {error}", + metadata_path.display() + )) + })?; + let metadata: ArtifactMetadata = serde_json::from_str(&metadata_json).map_err(|error| { + invalid_artifact(format!( + "failed to parse {}: {error}", + metadata_path.display() + )) + })?; + metadata.validate(probe_model)?; + + let tensor_path = directory.join(&metadata.tensor_file); + let tensor_bytes = std::fs::read(&tensor_path).map_err(|error| { + invalid_artifact(format!("failed to read {}: {error}", tensor_path.display())) + })?; + let tensors = { + let tensors = SafeTensors::deserialize(&tensor_bytes).map_err(|error| { + invalid_artifact(format!( + "failed to parse {}: {error}", + tensor_path.display() + )) + })?; + validate_tensors(&tensors, &metadata)?; + decode_tensors(&tensors)? + }; + + Ok(Self { metadata, tensors }) + } + + /// Returns checkpoint output names in learned probability order. + pub(super) fn output_names(&self) -> &[String] { + &self.metadata.output_names + } + + /// Returns the expected number of hidden-state layers. + pub(super) fn layer_count(&self) -> usize { + self.metadata.extraction_layer_ids.len() + } + + /// Returns the expected hidden width for each layer. + pub(super) fn hidden_size(&self) -> usize { + self.metadata.hidden_size + } + + /// Returns the flattened token-mean feature dimension. + pub(super) fn raw_feature_dim(&self) -> usize { + self.metadata.raw_feature_dim + } + + /// Applies the fitted scaler and PCA projection. + pub(super) fn project(&self, raw_features: &[f32]) -> Result> { + let standardized = standardize( + raw_features, + self.tensor("transform.scaler_mean")?, + self.tensor("transform.scaler_scale")?, + )?; + project_pca( + &standardized, + self.tensor("transform.pca_mean")?, + self.tensor("transform.pca_components")?, + self.metadata.pca_dim, + ) + } + + /// Runs all shared-trunk ensemble members in checkpoint order. + pub(super) fn ensemble_logits(&self, pca_features: &[f32]) -> Result>> { + if pca_features.len() != self.metadata.pca_dim { + return Err(trunk_error(format!( + "input dimension {} does not match pca_dim {}", + pca_features.len(), + self.metadata.pca_dim, + ))); + } + + let mut ensemble_logits = Vec::with_capacity(self.metadata.ensemble_size); + for index in 0..self.metadata.ensemble_size { + let prefix = format!("ensemble.{index}"); + let hidden1 = dense_layer( + pca_features, + self.tensor(&format!("{prefix}.linear1.weight"))?, + self.tensor(&format!("{prefix}.linear1.bias"))?, + TRUNK_HIDDEN[0], + true, + )?; + let hidden2 = dense_layer( + &hidden1, + self.tensor(&format!("{prefix}.linear2.weight"))?, + self.tensor(&format!("{prefix}.linear2.bias"))?, + TRUNK_HIDDEN[1], + true, + )?; + let logits = dense_layer( + &hidden2, + self.tensor(&format!("{prefix}.output.weight"))?, + self.tensor(&format!("{prefix}.output.bias"))?, + self.metadata.output_names.len(), + false, + )?; + ensemble_logits.push(logits); + } + Ok(ensemble_logits) + } + + /// Applies independent sigmoid links and averages probabilities across members. + pub(super) fn ensemble_probabilities(&self, ensemble_logits: &[Vec]) -> Result> { + if ensemble_logits.len() != self.metadata.ensemble_size { + return Err(trunk_error(format!( + "logit member count {} does not match ensemble_size {}", + ensemble_logits.len(), + self.metadata.ensemble_size, + ))); + } + + let output_count = self.metadata.output_names.len(); + let mut probability_sums = vec![0.0f32; output_count]; + for (member, logits) in ensemble_logits.iter().enumerate() { + if logits.len() != output_count { + return Err(trunk_error(format!( + "member {member} logit count {} does not match output count {output_count}", + logits.len(), + ))); + } + for (sum, logit) in probability_sums.iter_mut().zip(logits) { + *sum += sigmoid_probability(*logit)?; + } + } + + let member_count = self.metadata.ensemble_size as f32; + probability_sums + .iter_mut() + .for_each(|probability| *probability /= member_count); + Ok(probability_sums) + } + + fn tensor(&self, name: &str) -> Result<&[f32]> { + self.tensors + .get(name) + .map(Vec::as_slice) + .ok_or_else(|| invalid_artifact(format!("missing decoded tensor {name}"))) + } + + #[cfg(test)] + pub(super) fn with_test_probabilities(probabilities: [f32; OUTPUT_NAMES.len()]) -> Self { + let metadata = ArtifactMetadata { + format_version: FORMAT_VERSION, + training_mode: TRAINING_MODE.into(), + encoder: "probe/model".into(), + representation: REPRESENTATION.into(), + extraction_layer_ids: vec![0, 1], + hidden_size: 2, + raw_feature_dim: 4, + feature_block_count: 1, + pca_dim: PCA_DIM, + pca_whiten: false, + output_names: OUTPUT_NAMES.iter().map(|name| (*name).into()).collect(), + trunk_hidden: TRUNK_HIDDEN.to_vec(), + ensemble_size: ENSEMBLE_SIZE, + probability_link: PROBABILITY_LINK.into(), + ensemble_reduction: ENSEMBLE_REDUCTION.into(), + tensor_file: TENSOR_FILE.into(), + }; + let mut tensors = BTreeMap::from([ + ("transform.scaler_mean".into(), vec![0.0; 4]), + ("transform.scaler_scale".into(), vec![1.0; 4]), + ("transform.pca_mean".into(), vec![0.0; 4]), + ("transform.pca_components".into(), vec![0.0; PCA_DIM * 4]), + ]); + let logits = probabilities.map(|probability| (probability / (1.0 - probability)).ln()); + for index in 0..ENSEMBLE_SIZE { + let prefix = format!("ensemble.{index}"); + tensors.insert( + format!("{prefix}.linear1.weight"), + vec![0.0; TRUNK_HIDDEN[0] * PCA_DIM], + ); + tensors.insert(format!("{prefix}.linear1.bias"), vec![0.0; TRUNK_HIDDEN[0]]); + tensors.insert( + format!("{prefix}.linear2.weight"), + vec![0.0; TRUNK_HIDDEN[1] * TRUNK_HIDDEN[0]], + ); + tensors.insert(format!("{prefix}.linear2.bias"), vec![0.0; TRUNK_HIDDEN[1]]); + tensors.insert( + format!("{prefix}.output.weight"), + vec![0.0; OUTPUT_NAMES.len() * TRUNK_HIDDEN[1]], + ); + tensors.insert(format!("{prefix}.output.bias"), logits.to_vec()); + } + Self { metadata, tensors } + } +} + +#[derive(Deserialize)] +#[cfg_attr(test, derive(serde::Serialize))] +struct ArtifactMetadata { + format_version: u64, + training_mode: String, + encoder: String, + representation: String, + extraction_layer_ids: Vec, + hidden_size: usize, + raw_feature_dim: usize, + feature_block_count: usize, + pca_dim: usize, + pca_whiten: bool, + output_names: Vec, + trunk_hidden: Vec, + ensemble_size: usize, + probability_link: String, + ensemble_reduction: String, + tensor_file: String, +} + +impl ArtifactMetadata { + fn validate(&self, probe_model: &str) -> Result<()> { + require( + self.format_version == FORMAT_VERSION, + format!( + "unsupported format_version {}; expected {FORMAT_VERSION}", + self.format_version + ), + )?; + require( + self.training_mode == TRAINING_MODE, + format!( + "training_mode must be {TRAINING_MODE}; got {}", + self.training_mode + ), + )?; + require( + self.encoder == probe_model, + format!( + "artifact encoder {} does not match probe model {probe_model}", + self.encoder + ), + )?; + require( + self.representation == REPRESENTATION, + format!( + "representation must be {REPRESENTATION}; got {}", + self.representation + ), + )?; + require( + !self.extraction_layer_ids.is_empty(), + "extraction_layer_ids must not be empty", + )?; + let expected_layer_ids = (0..self.extraction_layer_ids.len()).collect::>(); + require( + self.extraction_layer_ids == expected_layer_ids, + "extraction_layer_ids must be contiguous and ordered from zero", + )?; + require(self.hidden_size > 0, "hidden_size must be positive")?; + let expected_raw_dim = self + .extraction_layer_ids + .len() + .checked_mul(self.hidden_size) + .ok_or_else(|| invalid_artifact("raw feature dimension overflow"))?; + require( + self.raw_feature_dim == expected_raw_dim, + format!( + "raw_feature_dim {} does not equal layer count {} * hidden_size {}", + self.raw_feature_dim, + self.extraction_layer_ids.len(), + self.hidden_size + ), + )?; + require( + self.feature_block_count == 1, + format!( + "feature_block_count must be 1; got {}", + self.feature_block_count + ), + )?; + require( + self.pca_dim == PCA_DIM, + format!("pca_dim must be {PCA_DIM}; got {}", self.pca_dim), + )?; + require(!self.pca_whiten, "pca_whiten must be false")?; + require( + self.output_names + .iter() + .map(String::as_str) + .eq(OUTPUT_NAMES), + format!("output_names must be ordered as {OUTPUT_NAMES:?}"), + )?; + require( + self.trunk_hidden == TRUNK_HIDDEN, + format!("trunk_hidden must be {TRUNK_HIDDEN:?}"), + )?; + require( + self.ensemble_size == ENSEMBLE_SIZE, + format!( + "ensemble_size must be {ENSEMBLE_SIZE}; got {}", + self.ensemble_size + ), + )?; + require( + self.probability_link == PROBABILITY_LINK, + format!( + "probability_link must be {PROBABILITY_LINK}; got {}", + self.probability_link + ), + )?; + require( + self.ensemble_reduction == ENSEMBLE_REDUCTION, + format!( + "ensemble_reduction must be {ENSEMBLE_REDUCTION}; got {}", + self.ensemble_reduction + ), + )?; + require( + self.tensor_file == TENSOR_FILE, + format!( + "tensor_file must be {TENSOR_FILE}; got {}", + self.tensor_file + ), + )?; + Ok(()) + } +} + +struct TensorSpec { + name: String, + shape: Vec, +} + +fn expected_tensors(metadata: &ArtifactMetadata) -> Vec { + let mut expected = vec![ + TensorSpec { + name: "transform.scaler_mean".into(), + shape: vec![metadata.raw_feature_dim], + }, + TensorSpec { + name: "transform.scaler_scale".into(), + shape: vec![metadata.raw_feature_dim], + }, + TensorSpec { + name: "transform.pca_mean".into(), + shape: vec![metadata.raw_feature_dim], + }, + TensorSpec { + name: "transform.pca_components".into(), + shape: vec![metadata.pca_dim, metadata.raw_feature_dim], + }, + ]; + for index in 0..metadata.ensemble_size { + let prefix = format!("ensemble.{index}"); + expected.extend([ + TensorSpec { + name: format!("{prefix}.linear1.weight"), + shape: vec![TRUNK_HIDDEN[0], metadata.pca_dim], + }, + TensorSpec { + name: format!("{prefix}.linear1.bias"), + shape: vec![TRUNK_HIDDEN[0]], + }, + TensorSpec { + name: format!("{prefix}.linear2.weight"), + shape: vec![TRUNK_HIDDEN[1], TRUNK_HIDDEN[0]], + }, + TensorSpec { + name: format!("{prefix}.linear2.bias"), + shape: vec![TRUNK_HIDDEN[1]], + }, + TensorSpec { + name: format!("{prefix}.output.weight"), + shape: vec![metadata.output_names.len(), TRUNK_HIDDEN[1]], + }, + TensorSpec { + name: format!("{prefix}.output.bias"), + shape: vec![metadata.output_names.len()], + }, + ]); + } + expected +} + +fn validate_tensors(tensors: &SafeTensors<'_>, metadata: &ArtifactMetadata) -> Result<()> { + let expected = expected_tensors(metadata); + for spec in &expected { + let tensor = tensors + .tensor(&spec.name) + .map_err(|_| invalid_artifact(format!("missing tensor {}", spec.name)))?; + require( + tensor.dtype() == Dtype::F32, + format!("tensor {} must use F32", spec.name), + )?; + require( + tensor.shape() == spec.shape, + format!( + "tensor {} has shape {:?}; expected {:?}", + spec.name, + tensor.shape(), + spec.shape + ), + )?; + validate_finite_f32(&spec.name, tensor.data())?; + if spec.name == "transform.scaler_scale" { + validate_positive_f32(&spec.name, tensor.data())?; + } + } + + let expected_names = expected + .iter() + .map(|spec| spec.name.as_str()) + .collect::>(); + let actual_names = tensors + .names() + .into_iter() + .map(String::as_str) + .collect::>(); + require( + actual_names == expected_names, + "artifact contains unexpected tensors", + )?; + Ok(()) +} + +/// Decodes artifact tensors once so requests never reparse the checkpoint. +fn decode_tensors(tensors: &SafeTensors<'_>) -> Result>> { + let mut decoded = BTreeMap::new(); + for name in tensors.names() { + let tensor = tensors + .tensor(name) + .map_err(|error| invalid_artifact(format!("failed to read tensor {name}: {error}")))?; + let values = tensor + .data() + .chunks_exact(size_of::()) + .map(|bytes| f32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]])) + .collect(); + decoded.insert(name.clone(), values); + } + Ok(decoded) +} + +/// Applies the fitted standard scaler elementwise. +fn standardize(raw: &[f32], mean: &[f32], scale: &[f32]) -> Result> { + if raw.len() != mean.len() || raw.len() != scale.len() { + return Err(transform_error(format!( + "scaler dimensions do not match: raw={}, mean={}, scale={}", + raw.len(), + mean.len(), + scale.len(), + ))); + } + + raw.iter() + .zip(mean) + .zip(scale) + .map(|((value, mean), scale)| { + if !value.is_finite() { + return Err(transform_error("raw features contain non-finite values")); + } + let standardized = (*value - *mean) / *scale; + if standardized.is_finite() { + Ok(standardized) + } else { + Err(transform_error( + "standardization produced a non-finite value", + )) + } + }) + .collect() +} + +/// Projects standardized features using row-major sklearn PCA components. +fn project_pca( + standardized: &[f32], + mean: &[f32], + components: &[f32], + output_dim: usize, +) -> Result> { + if standardized.is_empty() || standardized.len() != mean.len() { + return Err(transform_error(format!( + "PCA input dimensions do not match: standardized={}, mean={}", + standardized.len(), + mean.len(), + ))); + } + let expected_components = output_dim + .checked_mul(standardized.len()) + .ok_or_else(|| transform_error("PCA component dimensions overflow"))?; + if output_dim == 0 || components.len() != expected_components { + return Err(transform_error(format!( + "PCA component dimensions do not match: values={}, expected={expected_components}", + components.len(), + ))); + } + + let centered = standardized + .iter() + .zip(mean) + .map(|(value, mean)| { + let centered = *value - *mean; + if centered.is_finite() { + Ok(centered) + } else { + Err(transform_error("PCA centering produced a non-finite value")) + } + }) + .collect::>>()?; + + components + .chunks_exact(standardized.len()) + .map(|component| { + let projected = component + .iter() + .zip(¢ered) + .map(|(weight, value)| f64::from(*weight) * f64::from(*value)) + .sum::() as f32; + if projected.is_finite() { + Ok(projected) + } else { + Err(transform_error( + "PCA projection produced a non-finite value", + )) + } + }) + .collect() +} + +/// Applies a row-major dense layer and optional ReLU activation. +fn dense_layer( + input: &[f32], + weights: &[f32], + bias: &[f32], + output_dim: usize, + relu: bool, +) -> Result> { + if input.is_empty() || output_dim == 0 || bias.len() != output_dim { + return Err(trunk_error(format!( + "dense dimensions do not match: input={}, output={output_dim}, bias={}", + input.len(), + bias.len(), + ))); + } + let expected_weights = output_dim + .checked_mul(input.len()) + .ok_or_else(|| trunk_error("dense weight dimensions overflow"))?; + if weights.len() != expected_weights { + return Err(trunk_error(format!( + "dense weight dimensions do not match: values={}, expected={expected_weights}", + weights.len(), + ))); + } + if input.iter().any(|value| !value.is_finite()) { + return Err(trunk_error("dense input contains non-finite values")); + } + + weights + .chunks_exact(input.len()) + .zip(bias) + .map(|(row, bias)| { + let value = row + .iter() + .zip(input) + .map(|(weight, input)| *weight * *input) + .sum::() + + *bias; + if !value.is_finite() { + return Err(trunk_error("dense layer produced a non-finite value")); + } + Ok(if relu { value.max(0.0) } else { value }) + }) + .collect() +} + +fn sigmoid_probability(logit: f32) -> Result { + if !logit.is_finite() { + return Err(trunk_error("sigmoid input contains a non-finite value")); + } + Ok(if logit >= 0.0 { + 1.0 / (1.0 + (-logit).exp()) + } else { + let exp = logit.exp(); + exp / (1.0 + exp) + }) +} + +fn validate_finite_f32(name: &str, data: &[u8]) -> Result<()> { + for bytes in data.chunks_exact(size_of::()) { + let value = f32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]); + require( + value.is_finite(), + format!("tensor {name} contains non-finite values"), + )?; + } + Ok(()) +} + +fn validate_positive_f32(name: &str, data: &[u8]) -> Result<()> { + for bytes in data.chunks_exact(size_of::()) { + let value = f32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]); + require( + value > 0.0, + format!("tensor {name} contains non-positive values"), + )?; + } + Ok(()) +} + +fn require(condition: bool, message: impl Into) -> Result<()> { + if condition { + Ok(()) + } else { + Err(invalid_artifact(message)) + } +} + +fn invalid_artifact(message: impl Into) -> LibsyError { + LibsyError::AlgorithmError { + message: format!("invalid prefill-probe artifact: {}", message.into()), + } +} + +fn transform_error(message: impl Into) -> LibsyError { + LibsyError::AlgorithmError { + message: format!("prefill-probe feature transform error: {}", message.into()), + } +} + +fn trunk_error(message: impl Into) -> LibsyError { + LibsyError::AlgorithmError { + message: format!("prefill-probe trunk inference error: {}", message.into()), + } +} + +#[cfg(test)] +mod tests { + use std::path::{Path, PathBuf}; + use std::sync::atomic::{AtomicU64, Ordering}; + + use safetensors::tensor::{serialize, TensorView}; + + use super::*; + + static NEXT_TEST_DIRECTORY: AtomicU64 = AtomicU64::new(0); + + struct TestArtifactDirectory(PathBuf); + + impl TestArtifactDirectory { + fn create() -> Result { + let sequence = NEXT_TEST_DIRECTORY.fetch_add(1, Ordering::Relaxed); + let path = std::env::temp_dir().join(format!( + "switchyard-libsy-prefill-artifact-{}-{sequence}", + std::process::id() + )); + std::fs::create_dir(&path).map_err(|error| { + invalid_artifact(format!( + "failed to create test directory {}: {error}", + path.display() + )) + })?; + Ok(Self(path)) + } + + fn write_metadata(&self, metadata: &ArtifactMetadata) -> Result<()> { + let bytes = serde_json::to_vec(metadata).map_err(|error| { + invalid_artifact(format!("failed to serialize test metadata: {error}")) + })?; + self.write(METADATA_FILE, &bytes) + } + + fn write(&self, name: &str, bytes: &[u8]) -> Result<()> { + let path = self.0.join(name); + std::fs::write(&path, bytes).map_err(|error| { + invalid_artifact(format!("failed to write {}: {error}", path.display())) + }) + } + + fn path(&self) -> &Path { + &self.0 + } + } + + impl Drop for TestArtifactDirectory { + fn drop(&mut self) { + let _ = std::fs::remove_dir_all(&self.0); + } + } + + fn test_metadata() -> ArtifactMetadata { + ArtifactMetadata { + format_version: FORMAT_VERSION, + training_mode: TRAINING_MODE.into(), + encoder: "probe/model".into(), + representation: REPRESENTATION.into(), + extraction_layer_ids: vec![0, 1], + hidden_size: 2, + raw_feature_dim: 4, + feature_block_count: 1, + pca_dim: PCA_DIM, + pca_whiten: false, + output_names: OUTPUT_NAMES.iter().map(|name| (*name).into()).collect(), + trunk_hidden: TRUNK_HIDDEN.to_vec(), + ensemble_size: ENSEMBLE_SIZE, + probability_link: PROBABILITY_LINK.into(), + ensemble_reduction: ENSEMBLE_REDUCTION.into(), + tensor_file: TENSOR_FILE.into(), + } + } + + fn repeated_f32_bytes(value: f32, count: usize) -> Vec { + let mut bytes = Vec::with_capacity(count * size_of::()); + for _ in 0..count { + bytes.extend_from_slice(&value.to_le_bytes()); + } + bytes + } + + fn serialize_valid_artifact(metadata: &ArtifactMetadata) -> Result> { + let storage = expected_tensors(metadata) + .into_iter() + .map(|spec| { + let value_count = spec.shape.iter().product(); + let fill = if spec.name == "transform.scaler_scale" { + 1.0 + } else { + 0.0 + }; + (spec.name, spec.shape, repeated_f32_bytes(fill, value_count)) + }) + .collect::>(); + let tensors = storage + .iter() + .map(|(name, shape, bytes)| { + TensorView::new(Dtype::F32, shape.clone(), bytes) + .map(|tensor| (name.as_str(), tensor)) + .map_err(|error| { + invalid_artifact(format!("failed to create test tensor {name}: {error}")) + }) + }) + .collect::>>()?; + serialize(tensors, &None).map_err(|error| { + invalid_artifact(format!("failed to serialize test artifact: {error}")) + }) + } + + #[test] + fn artifact_loads_and_executes_the_exported_shape() -> Result<()> { + let directory = TestArtifactDirectory::create()?; + let metadata = test_metadata(); + directory.write_metadata(&metadata)?; + directory.write(TENSOR_FILE, &serialize_valid_artifact(&metadata)?)?; + + let artifact = InferenceArtifact::load(directory.path(), "probe/model")?; + assert_eq!(artifact.layer_count(), 2); + assert_eq!(artifact.hidden_size(), 2); + assert_eq!(artifact.raw_feature_dim(), 4); + assert!(artifact + .output_names() + .iter() + .map(String::as_str) + .eq(OUTPUT_NAMES)); + + let projected = artifact.project(&[0.0; 4])?; + assert_eq!(projected, vec![0.0; PCA_DIM]); + let logits = artifact.ensemble_logits(&projected)?; + assert_eq!(logits, vec![vec![0.0; OUTPUT_NAMES.len()]; ENSEMBLE_SIZE]); + assert_eq!( + artifact.ensemble_probabilities(&logits)?, + vec![0.5; OUTPUT_NAMES.len()] + ); + Ok(()) + } + + #[test] + fn metadata_rejects_encoder_and_dimension_mismatches() -> Result<()> { + let encoder_error = test_metadata() + .validate("different/probe") + .err() + .ok_or_else(|| invalid_artifact("encoder mismatch should fail"))?; + assert!(encoder_error + .to_string() + .contains("does not match probe model")); + + let mut metadata = test_metadata(); + metadata.extraction_layer_ids = vec![0, 2]; + let layer_error = metadata + .validate("probe/model") + .err() + .ok_or_else(|| invalid_artifact("layer ordering mismatch should fail"))?; + assert!(layer_error.to_string().contains("contiguous and ordered")); + + let mut metadata = test_metadata(); + metadata.raw_feature_dim += 1; + let dimension_error = metadata + .validate("probe/model") + .err() + .ok_or_else(|| invalid_artifact("raw dimension mismatch should fail"))?; + assert!(dimension_error + .to_string() + .contains("does not equal layer count")); + Ok(()) + } + + #[test] + fn scaler_and_pca_match_exported_row_major_math() -> Result<()> { + let standardized = standardize(&[3.0, 6.0, 11.0], &[1.0, 2.0, 3.0], &[2.0, 2.0, 4.0])?; + assert_eq!(standardized, vec![1.0, 2.0, 2.0]); + + let projected = project_pca( + &standardized, + &[0.5, 1.0, 1.5], + &[ + 1.0, 10.0, 100.0, // PCA component 0 + -2.0, 0.5, 4.0, // PCA component 1 + ], + 2, + )?; + assert_eq!(projected, vec![60.5, 1.5]); + Ok(()) + } + + #[test] + fn inference_rejects_malformed_dimensions_and_non_finite_values() -> Result<()> { + let scaler_error = standardize(&[1.0, 2.0], &[0.0], &[1.0, 1.0]) + .err() + .ok_or_else(|| transform_error("scaler mismatch should fail"))?; + assert!(scaler_error.to_string().contains("scaler dimensions")); + + let dense_error = dense_layer(&[f32::MAX], &[2.0], &[0.0], 1, false) + .err() + .ok_or_else(|| trunk_error("non-finite dense output should fail"))?; + assert!(dense_error.to_string().contains("dense layer produced")); + + let sigmoid_error = sigmoid_probability(f32::NAN) + .err() + .ok_or_else(|| trunk_error("non-finite sigmoid input should fail"))?; + assert!(sigmoid_error.to_string().contains("sigmoid input")); + Ok(()) + } +} diff --git a/crates/libsy/src/algorithms/prefill_probe/policy.rs b/crates/libsy/src/algorithms/prefill_probe/policy.rs new file mode 100644 index 000000000..ec9875784 --- /dev/null +++ b/crates/libsy/src/algorithms/prefill_probe/policy.rs @@ -0,0 +1,151 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Cost-aware selection between the learned router's weak and strong heads. + +use crate::{LibsyError, Result}; + +/// Completion tier selected by the learned utility policy. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(super) enum PrefillTier { + Weak, + Strong, +} + +/// Validated utility policy for two completion targets. +#[derive(Clone, Copy, Debug)] +pub(super) struct CostAwareRoutingPolicy { + lambda: f64, + normalized_weak_cost: f64, + normalized_strong_cost: f64, +} + +impl CostAwareRoutingPolicy { + /// Validates the policy and min-max normalizes costs across weak and strong. + /// + /// With two targets, normalization makes routing depend on cost ordering, + /// not the magnitude of the difference between the configured costs. + pub(super) fn new(lambda: f64, weak_cost: f64, strong_cost: f64) -> Result { + if !lambda.is_finite() || !(0.0..=1.0).contains(&lambda) { + return Err(policy_error( + "routing policy lambda must be finite and in [0.0, 1.0]", + )); + } + validate_cost("weak_cost", weak_cost)?; + validate_cost("strong_cost", strong_cost)?; + + let (normalized_weak_cost, normalized_strong_cost) = if weak_cost == strong_cost { + (0.0, 0.0) + } else { + let minimum = weak_cost.min(strong_cost); + let range = weak_cost.max(strong_cost) - minimum; + ( + (weak_cost - minimum) / range, + (strong_cost - minimum) / range, + ) + }; + + Ok(Self { + lambda, + normalized_weak_cost, + normalized_strong_cost, + }) + } + + /// Selects the tier with the greater cost-adjusted correctness utility. + /// + /// Equal utilities deterministically select weak. + pub(super) fn select( + &self, + weak_probability: f64, + strong_probability: f64, + ) -> Result { + validate_probability("weak", weak_probability)?; + validate_probability("strong", strong_probability)?; + + let cost_weight = 1.0 - self.lambda; + let weak_utility = self.lambda * weak_probability - cost_weight * self.normalized_weak_cost; + let strong_utility = + self.lambda * strong_probability - cost_weight * self.normalized_strong_cost; + Ok(if weak_utility >= strong_utility { + PrefillTier::Weak + } else { + PrefillTier::Strong + }) + } +} + +fn validate_cost(field: &str, cost: f64) -> Result<()> { + if !cost.is_finite() || cost < 0.0 { + return Err(policy_error(format!( + "routing policy {field} must be finite and non-negative" + ))); + } + Ok(()) +} + +fn validate_probability(head: &str, probability: f64) -> Result<()> { + if !probability.is_finite() || !(0.0..=1.0).contains(&probability) { + return Err(policy_error(format!( + "{head} checkpoint probability must be finite and in [0.0, 1.0]" + ))); + } + Ok(()) +} + +fn policy_error(message: impl Into) -> LibsyError { + LibsyError::AlgorithmError { + message: format!("prefill-probe policy error: {}", message.into()), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn lambda_zero_uses_cost_and_weak_wins_equal_utility() -> Result<()> { + let weak_cheaper = CostAwareRoutingPolicy::new(0.0, 1.0, 10.0)?; + let strong_cheaper = CostAwareRoutingPolicy::new(0.0, 10.0, 1.0)?; + let equal_cost = CostAwareRoutingPolicy::new(0.0, 3.0, 3.0)?; + + assert_eq!(weak_cheaper.select(0.0, 1.0)?, PrefillTier::Weak); + assert_eq!(strong_cheaper.select(1.0, 0.0)?, PrefillTier::Strong); + assert_eq!(equal_cost.select(0.0, 1.0)?, PrefillTier::Weak); + Ok(()) + } + + #[test] + fn lambda_one_uses_correctness_probabilities() -> Result<()> { + let policy = CostAwareRoutingPolicy::new(1.0, 100.0, 1.0)?; + + assert_eq!(policy.select(0.8, 0.6)?, PrefillTier::Weak); + assert_eq!(policy.select(0.4, 0.6)?, PrefillTier::Strong); + assert_eq!(policy.select(0.6, 0.6)?, PrefillTier::Weak); + Ok(()) + } + + #[test] + fn invalid_values_are_rejected() -> Result<()> { + for (lambda, weak_cost, strong_cost, expected) in [ + (f64::NAN, 0.0, 1.0, "lambda"), + (-0.1, 0.0, 1.0, "lambda"), + (1.1, 0.0, 1.0, "lambda"), + (0.5, -1.0, 1.0, "weak_cost"), + (0.5, 1.0, f64::INFINITY, "strong_cost"), + ] { + let error = CostAwareRoutingPolicy::new(lambda, weak_cost, strong_cost) + .err() + .ok_or_else(|| policy_error("invalid policy value should fail"))?; + assert!(error.to_string().contains(expected)); + } + + let policy = CostAwareRoutingPolicy::new(0.5, 0.0, 1.0)?; + let error = policy + .select(f64::NAN, 0.5) + .err() + .ok_or_else(|| policy_error("invalid probability should fail"))?; + assert!(error.to_string().contains("weak")); + Ok(()) + } +} diff --git a/crates/libsy/src/lib.rs b/crates/libsy/src/lib.rs index 0c54c3015..26bbe2b1c 100644 --- a/crates/libsy/src/lib.rs +++ b/crates/libsy/src/lib.rs @@ -75,6 +75,9 @@ //! [`algorithms::Random`] provides uniform or weighted random routing. //! //! [`algorithms::LlmTaskClassifier`] uses one model to classify and route to its selected target. +//! +//! [`algorithms::PrefillProbeClassifier`] classifies from prompt hidden-state +//! features supplied by a transport-independent probe. mod core; pub use core::*; diff --git a/crates/switchyard-server/CONFIGURATION.md b/crates/switchyard-server/CONFIGURATION.md index bc9e955c4..b5942a06e 100644 --- a/crates/switchyard-server/CONFIGURATION.md +++ b/crates/switchyard-server/CONFIGURATION.md @@ -21,6 +21,42 @@ To support another wire format, add its `ClientFormat` variant and explicit cons ## Add an algorithm 1. Implement and export the algorithm from `libsy`. -2. Add its TOML fields as an `AlgorithmConfig` variant in `src/config.rs`. +2. Add its TOML fields as a `RouteConfig` variant in `src/config.rs`. 3. Construct it in the `build_algorithm` match, resolving target names with `resolve_targets`. 4. Add a parsing test and an end-to-end server test when the algorithm makes LLM calls. + +## Configure a learned prefill-probe route + +`prefill_probe` uses prompt hidden states from a separate vLLM endpoint to select one of two +completion targets: + +```toml +[routes.learned] +id = "switchyard/learned" +type = "prefill_probe" +strong_target = "strong" +weak_target = "weak" +probe_base_url = "http://127.0.0.1:8000/v1" +probe_model = "Qwen/Qwen3-8B" +hidden_states_dir = "/dev/shm/switchyard-prefill" +checkpoint_dir = "/opt/switchyard/router" +strong_checkpoint_head = "opus-4.7" +weak_checkpoint_head = "nemotron-3-super" +lambda = 0.75 +weak_cost = 0.25 +strong_cost = 1.0 +probe_timeout_secs = 30.0 +cache_capacity = 4096 +``` + +The strong and weak values reference entries under `targets`. `probe_model` must match both the +model served by the probe endpoint and the checkpoint's `encoder` metadata. The checkpoint +directory must contain `router.json` and `router.safetensors`. + +`probe_timeout_secs` defaults to 30 seconds and must be positive. `cache_capacity` defaults to +4096 successful task decisions and must be positive. Probe or checkpoint inference failures select +the strong target and are not cached. + +See the +[vLLM hidden-state probe guide](../../docs/operations/vllm_hidden_state_probe.md) +for the required vLLM connector configuration, artifact lifecycle, and cost-policy semantics. diff --git a/crates/switchyard-server/Cargo.toml b/crates/switchyard-server/Cargo.toml index 753333d32..cd7ef7d3d 100644 --- a/crates/switchyard-server/Cargo.toml +++ b/crates/switchyard-server/Cargo.toml @@ -36,6 +36,7 @@ tracing-subscriber = { version = "0.3", features = ["env-filter"] } [dev-dependencies] async-trait.workspace = true http-body-util = "0.1" +safetensors = "0.4" tempfile = "3" tokio.workspace = true tower = { version = "0.5", features = ["util"] } diff --git a/crates/switchyard-server/README.md b/crates/switchyard-server/README.md index f83760b0b..9fb17eb41 100644 --- a/crates/switchyard-server/README.md +++ b/crates/switchyard-server/README.md @@ -58,8 +58,8 @@ upstream, and a route's `id` is the model clients send to select that algorithm. Each target references an entry under `llm_clients`. All configured clients use `TranslatingLlmClient`; supported formats are `openai_chat`, `openai_responses`, and `anthropic_messages`. Supported algorithms are `noop`, `random`, `passthrough`, and -`llm_classifier`. An `api_key_env` value names an environment variable; the TOML -never contains the secret itself. If omitted, the client sends no authentication. +`llm_classifier`, and `prefill_probe`. An `api_key_env` value names an environment variable; the +TOML never contains the secret itself. If omitted, the client sends no authentication. Random-route `weights` are relative, follow target order, and do not need to sum to one. Omit them for equal weighting. The optional `seed` reproduces the selection sequence for the same call order. @@ -102,3 +102,6 @@ The `tier` label is `strong` or `weak` for a distinguishable built-in LLM-classi is omitted for untiered algorithms. Classifier calls are excluded from these families. See [CONFIGURATION.md](CONFIGURATION.md) to add an LLM client, target, or algorithm. +See the +[vLLM hidden-state probe guide](../../docs/operations/vllm_hidden_state_probe.md) +to deploy a learned `prefill_probe` route. diff --git a/crates/switchyard-server/src/config.rs b/crates/switchyard-server/src/config.rs index 5074b4ff5..71fab85b9 100644 --- a/crates/switchyard-server/src/config.rs +++ b/crates/switchyard-server/src/config.rs @@ -5,17 +5,27 @@ use std::collections::BTreeMap; use std::fs; -use std::path::Path; +use std::path::{Path, PathBuf}; use std::sync::Arc; - -use libsy::algorithms::{LlmTaskClassifier, Noop, Passthrough, Random, TaskClassifierConfig}; -use libsy::{Algorithm, LlmTarget, LlmTargetSet, RoutedLlmClient}; +use std::time::Duration; + +use async_trait::async_trait; +use libsy::algorithms::{ + FallThrough, LlmTaskClassifier, Noop, Passthrough, PrefillFeatures, PrefillProbe, + PrefillProbeClassifier, PrefillProbeClassifierConfig, Random, TaskClassifierConfig, + DEFAULT_PREFILL_PROBE_CACHE_CAPACITY, +}; +use libsy::{Algorithm, LibsyError, LlmTarget, LlmTargetSet, RoutedLlmClient, State}; use serde::Deserialize; -use switchyard_llm_client::{Backend, HttpBackendConfig, ModelConfig, TranslatingLlmClient}; +use switchyard_llm_client::{ + Backend, HttpBackendConfig, ModelConfig, TranslatingLlmClient, VllmHiddenStateProbe, + VllmHiddenStateProbeConfig, +}; use crate::{ServerError, ServerResult, ServerState}; const SUPPORTED_SCHEMA_VERSION: u32 = 1; +const DEFAULT_PREFILL_PROBE_TIMEOUT_SECS: f64 = 30.0; /// Loads a TOML deployment file and constructs the complete server state. pub fn load_server_state(path: impl AsRef) -> ServerResult { @@ -183,19 +193,64 @@ enum RouteConfig { #[serde(flatten)] classifier_config: TaskClassifierConfig, }, + PrefillProbe { + id: String, + strong_target: String, + weak_target: String, + probe_base_url: String, + probe_model: String, + hidden_states_dir: PathBuf, + checkpoint_dir: PathBuf, + strong_checkpoint_head: String, + weak_checkpoint_head: String, + lambda: f64, + weak_cost: f64, + strong_cost: f64, + #[serde(default = "default_prefill_probe_timeout_secs")] + probe_timeout_secs: f64, + #[serde(default = "default_prefill_probe_cache_capacity")] + cache_capacity: usize, + }, } impl RouteConfig { fn id(&self) -> &str { use RouteConfig::*; match self { - Noop { id } | Random { id, .. } | LlmClassifier { id, .. } | Passthrough { id, .. } => { - id - } + Noop { id } + | Random { id, .. } + | LlmClassifier { id, .. } + | PrefillProbe { id, .. } + | Passthrough { id, .. } => id, } } } +fn default_prefill_probe_timeout_secs() -> f64 { + DEFAULT_PREFILL_PROBE_TIMEOUT_SECS +} + +fn default_prefill_probe_cache_capacity() -> usize { + DEFAULT_PREFILL_PROBE_CACHE_CAPACITY +} + +struct ServerVllmPrefillProbe { + inner: VllmHiddenStateProbe, +} + +#[async_trait] +impl PrefillProbe for ServerVllmPrefillProbe { + async fn extract(&self, task: &str) -> libsy::Result { + self.inner + .extract(task) + .await + .map(|features| { + PrefillFeatures::new(features.layer_count, features.hidden_size, features.values) + }) + .map_err(|error| LibsyError::external("vLLM hidden-state probe", error)) + } +} + fn build_backend(client_name: &str, config: &LlmClientConfig) -> ServerResult { let base_url = config.base_url.trim(); if base_url.is_empty() { @@ -278,6 +333,68 @@ fn build_algorithm( })?; Ok(Arc::new(algorithm)) } + RouteConfig::PrefillProbe { + strong_target, + weak_target, + probe_base_url, + probe_model, + hidden_states_dir, + checkpoint_dir, + strong_checkpoint_head, + weak_checkpoint_head, + lambda, + weak_cost, + strong_cost, + probe_timeout_secs, + cache_capacity, + .. + } => { + let strong = resolve_target(route_name, strong_target, targets)?; + let weak = resolve_target(route_name, weak_target, targets)?; + let request_timeout = + Duration::try_from_secs_f64(*probe_timeout_secs).map_err(|error| { + ServerError::new(format!( + "prefill_probe route {route_name}: invalid probe_timeout_secs \ + {probe_timeout_secs}: {error}" + )) + })?; + if request_timeout.is_zero() { + return Err(ServerError::new(format!( + "prefill_probe route {route_name}: probe_timeout_secs must be positive" + ))); + } + let probe = VllmHiddenStateProbe::new(VllmHiddenStateProbeConfig { + base_url: probe_base_url.clone(), + model: probe_model.clone(), + hidden_states_dir: hidden_states_dir.clone(), + request_timeout, + }) + .map_err(|error| { + ServerError::new(format!("prefill_probe route {route_name}: {error}")) + })?; + let classifier = PrefillProbeClassifier::new( + PrefillProbeClassifierConfig { + probe_model: probe_model.clone(), + checkpoint_dir: checkpoint_dir.clone(), + strong_checkpoint_head: strong_checkpoint_head.clone(), + weak_checkpoint_head: weak_checkpoint_head.clone(), + strong_target: strong.semantic_name.clone(), + weak_target: weak.semantic_name.clone(), + lambda: *lambda, + weak_cost: *weak_cost, + strong_cost: *strong_cost, + cache_capacity: *cache_capacity, + }, + Arc::new(ServerVllmPrefillProbe { inner: probe }), + ) + .map_err(|error| { + ServerError::new(format!("prefill_probe route {route_name}: {error}")) + })?; + let target_set = LlmTargetSet::new(vec![strong, weak]); + let router = FallThrough::::new_with_state(target_set) + .with_classifier(Arc::new(classifier)); + Ok(Arc::new(router)) + } } } @@ -316,8 +433,17 @@ fn validate_value(label: &str, value: &str) -> ServerResult<()> { #[cfg(test)] mod tests { + use std::error::Error; + + use safetensors::tensor::{serialize, TensorView}; + use safetensors::Dtype; + use serde_json::json; + use tempfile::TempDir; + use super::*; + type TestResult = Result>; + const VALID_CONFIG: &str = r#" schema_version = 1 @@ -375,8 +501,102 @@ target = "weak" } } + fn repeated_f32_bytes(value: f32, count: usize) -> Vec { + let mut bytes = Vec::with_capacity(count * size_of::()); + for _ in 0..count { + bytes.extend_from_slice(&value.to_le_bytes()); + } + bytes + } + + fn write_test_checkpoint(directory: &Path) -> TestResult { + let metadata = json!({ + "format_version": 1, + "training_mode": "single_pca_block", + "encoder": "probe/model", + "representation": "token_mean_per_layer_concat", + "extraction_layer_ids": [0, 1], + "hidden_size": 2, + "raw_feature_dim": 4, + "feature_block_count": 1, + "pca_dim": 200, + "pca_whiten": false, + "output_names": ["qwen-122b", "nemotron-3-super", "opus-4.7", "gpt-5.5"], + "trunk_hidden": [256, 128], + "ensemble_size": 5, + "probability_link": "independent_sigmoid", + "ensemble_reduction": "probability_mean", + "tensor_file": "router.safetensors", + }); + std::fs::write( + directory.join("router.json"), + serde_json::to_vec(&metadata)?, + )?; + + let mut specs = vec![ + ("transform.scaler_mean".to_string(), vec![4], 0.0), + ("transform.scaler_scale".to_string(), vec![4], 1.0), + ("transform.pca_mean".to_string(), vec![4], 0.0), + ("transform.pca_components".to_string(), vec![200, 4], 0.0), + ]; + for index in 0..5 { + let prefix = format!("ensemble.{index}"); + specs.extend([ + (format!("{prefix}.linear1.weight"), vec![256, 200], 0.0), + (format!("{prefix}.linear1.bias"), vec![256], 0.0), + (format!("{prefix}.linear2.weight"), vec![128, 256], 0.0), + (format!("{prefix}.linear2.bias"), vec![128], 0.0), + (format!("{prefix}.output.weight"), vec![4, 128], 0.0), + (format!("{prefix}.output.bias"), vec![4], 0.0), + ]); + } + let storage = specs + .into_iter() + .map(|(name, shape, fill)| { + let value_count = shape.iter().product(); + (name, shape, repeated_f32_bytes(fill, value_count)) + }) + .collect::>(); + let tensors = storage + .iter() + .map(|(name, shape, bytes)| { + TensorView::new(Dtype::F32, shape.clone(), bytes) + .map(|tensor| (name.as_str(), tensor)) + }) + .collect::, _>>()?; + std::fs::write( + directory.join("router.safetensors"), + serialize(tensors, &None)?, + )?; + Ok(()) + } + + fn prefill_config(checkpoint_dir: &Path, hidden_states_dir: &Path) -> String { + format!( + r#"{VALID_CONFIG} + +[routes.prefill] +id = "switchyard/prefill" +type = "prefill_probe" +strong_target = "strong" +weak_target = "weak" +probe_base_url = "http://127.0.0.1:8000/v1" +probe_model = "probe/model" +hidden_states_dir = "{}" +checkpoint_dir = "{}" +strong_checkpoint_head = "opus-4.7" +weak_checkpoint_head = "nemotron-3-super" +lambda = 0.75 +weak_cost = 0.25 +strong_cost = 1.0 +"#, + hidden_states_dir.display(), + checkpoint_dir.display(), + ) + } + #[test] - fn builds_all_supported_algorithm_types() -> ServerResult<()> { + fn builds_stateless_and_llm_classifier_algorithm_types() -> ServerResult<()> { let state = server_state_from_toml(VALID_CONFIG)?; // The model id array is sorted alphabetically assert_eq!( @@ -391,6 +611,38 @@ target = "weak" Ok(()) } + #[test] + fn builds_prefill_probe_route_with_bounded_defaults() -> TestResult { + let checkpoint_dir = TempDir::new()?; + let hidden_states_dir = TempDir::new()?; + write_test_checkpoint(checkpoint_dir.path())?; + + let state = server_state_from_toml(&prefill_config( + checkpoint_dir.path(), + hidden_states_dir.path(), + ))?; + + assert!(state.models().any(|model| model == "switchyard/prefill")); + Ok(()) + } + + #[test] + fn prefill_probe_errors_include_route_context() -> TestResult { + let checkpoint_dir = TempDir::new()?; + let hidden_states_dir = TempDir::new()?; + write_test_checkpoint(checkpoint_dir.path())?; + let invalid = prefill_config(checkpoint_dir.path(), hidden_states_dir.path()).replace( + "strong_cost = 1.0", + "strong_cost = 1.0\nprobe_timeout_secs = 0", + ); + + let message = error_message(&invalid); + + assert!(message.contains("prefill_probe route prefill")); + assert!(message.contains("probe_timeout_secs must be positive")); + Ok(()) + } + #[test] fn rejects_unknown_fields_and_algorithm_types() { let unknown_field = diff --git a/docs/operations/vllm_hidden_state_probe.md b/docs/operations/vllm_hidden_state_probe.md new file mode 100644 index 000000000..747c75eb0 --- /dev/null +++ b/docs/operations/vllm_hidden_state_probe.md @@ -0,0 +1,161 @@ +# vLLM Hidden-State Probe + +The Rust server's `prefill_probe` route uses a small vLLM prefill request to extract prompt hidden +states, runs a learned checkpoint on CPU, and selects a weak or strong completion target. Use it +when a trained router can reduce completion cost enough to justify a separate probe-model prefill. + +This route is available through `switchyard-server` TOML configuration. It is not a Python +route-bundle algorithm. + +## Before you start + +You need: + +- A vLLM release with + [hidden-state extraction](https://docs.vllm.ai/en/v0.24.0/features/speculative_decoding/extract_hidden_states/) + and `ExampleHiddenStatesConnector`. +- A dedicated directory visible at the same path to vLLM and Switchyard. A RAM-backed filesystem + such as `/dev/shm` avoids persistent disk I/O. +- An exported checkpoint directory containing `router.json` and `router.safetensors`. +- Strong and weak completion targets configured in the Rust server. + +The checkpoint metadata's `encoder` must exactly match the configured probe model. Its extracted +layer count and hidden size must also match the tensors produced by vLLM. + +## Start the probe endpoint + +Create a directory used only for probe artifacts: + +```bash +mkdir -p /dev/shm/switchyard-prefill +``` + +Start vLLM with the hidden-state extraction method and disk connector. Replace the model and layer +IDs with the values used to train the checkpoint: + +```bash +vllm serve Qwen/Qwen3-8B \ + --speculative_config '{ + "method": "extract_hidden_states", + "num_speculative_tokens": 1, + "draft_model_config": { + "hf_config": { + "eagle_aux_hidden_state_layer_ids": [1, 2, 3, 4] + } + } + }' \ + --kv_transfer_config '{ + "kv_connector": "ExampleHiddenStatesConnector", + "kv_role": "kv_producer", + "kv_connector_extra_config": { + "shared_storage_path": "/dev/shm/switchyard-prefill", + "use_synchronization_lock": true + } + }' +``` + +Keep `use_synchronization_lock` enabled. Switchyard acquires the companion `.lock` file before +reading, so it cannot parse a partially written safetensors artifact. Chunked prefill is +incompatible with vLLM hidden-state extraction and must be disabled. + +Switchyard does not submit a custom output path. vLLM generates the filename under +`shared_storage_path` and returns it in `kv_transfer_params`, so +`allow_custom_save_path` can remain disabled. + +## Configure the Rust server + +Define the completion clients and targets as usual, then add the learned route: + +```toml +schema_version = 1 + +[llm_clients.completions] +format = "openai_chat" +base_url = "https://completion-provider.example/v1" +api_key_env = "COMPLETION_API_KEY" + +[targets.strong] +id = "provider/strong-model" +llm_client = "completions" + +[targets.weak] +id = "provider/weak-model" +llm_client = "completions" + +[routes.learned] +id = "switchyard/learned" +type = "prefill_probe" +strong_target = "strong" +weak_target = "weak" +probe_base_url = "http://127.0.0.1:8000/v1" +probe_model = "Qwen/Qwen3-8B" +hidden_states_dir = "/dev/shm/switchyard-prefill" +checkpoint_dir = "/opt/switchyard/router" +strong_checkpoint_head = "opus-4.7" +weak_checkpoint_head = "nemotron-3-super" +lambda = 0.75 +weak_cost = 0.25 +strong_cost = 1.0 +probe_timeout_secs = 30.0 +cache_capacity = 4096 +``` + +Validate construction before binding a port: + +```bash +cargo run -p switchyard-server -- --config routes.toml --dry-run +``` + +Then start the server: + +```bash +export COMPLETION_API_KEY="..." +cargo run -p switchyard-server -- --config routes.toml +``` + +Clients select the route by sending `switchyard/learned` as the model. + +## Tune the policy + +The checkpoint produces a correctness probability for each named output head. The policy selects +the tier with the greater utility: + +```text +utility = lambda * correctness_probability - (1 - lambda) * normalized_cost +``` + +`lambda` must be between `0.0` and `1.0`. At `1.0`, only predicted correctness matters; at `0.0`, +only configured cost matters. Equal utility selects weak. + +With exactly two targets, costs are min-max normalized. Only their ordering matters, not the +magnitude of the difference: + +- Set `weak_cost < strong_cost` when weak is cheaper. +- Set `weak_cost > strong_cost` when strong is cheaper. +- Equal costs remove the cost penalty from both tiers. + +Costs must be finite and non-negative and use the same units. The checkpoint head names must be +distinct entries in `router.json`'s `output_names`. + +`cache_capacity` bounds an in-memory LRU of successful decisions. Keys are process-randomized +hashes, so raw task text is not retained. Repeated identical tasks use the cached tier without +calling the probe endpoint. Probe and inference failures select strong and are not cached. + +## Artifact lifecycle and failures + +For each uncached task, Switchyard: + +1. Sends one user message to the probe endpoint with `max_tokens = 1` and prompt-only hidden-state + extraction. +2. Accepts only a returned `.safetensors` path inside `hidden_states_dir`. +3. Waits up to one second for vLLM's synchronization lock, then reads and validates + `hidden_states` and optional `token_ids` tensors on Tokio's blocking pool. +4. Token-mean pools `[tokens, layers, hidden]` into one vector per layer. +5. Removes the safetensors file and its companion `.lock`, including when tensor parsing fails. + +The HTTP request is bounded by `probe_timeout_secs`. A stale-file sweep runs before each probe and +removes unlocked `.safetensors` files older than five minutes from the dedicated directory. Do not +place unrelated safetensors files there. + +The probe task text is not written to Switchyard logs. Treat the temporary hidden-state artifacts +as sensitive and restrict access to the shared directory. diff --git a/mkdocs.yml b/mkdocs.yml index 0409a52e7..506aeaff5 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -35,6 +35,7 @@ nav: - Escalation-Router Routing: routing_algorithms/escalation_router_routing.md - Operations: - Context-Window Handling: operations/context_window.md + - vLLM Hidden-State Probe: operations/vllm_hidden_state_probe.md - Reference: - CLI Reference: cli_reference.md