From 6513263e05508ef4ce1a1441261aad99b9d91b64 Mon Sep 17 00:00:00 2001 From: asurla1998 Date: Tue, 28 Jul 2026 22:58:27 +0000 Subject: [PATCH] feat(libsy): add prefill-probe classifier Signed-off-by: asurla1998 --- Cargo.lock | 46 +- crates/libsy/Cargo.toml | 2 + crates/libsy/src/algorithms.rs | 5 + crates/libsy/src/algorithms/prefill_probe.rs | 600 ++++++++++++ .../src/algorithms/prefill_probe/artifact.rs | 871 ++++++++++++++++++ .../src/algorithms/prefill_probe/policy.rs | 151 +++ crates/libsy/src/lib.rs | 3 + 7 files changed, 1677 insertions(+), 1 deletion(-) create mode 100644 crates/libsy/src/algorithms/prefill_probe.rs create mode 100644 crates/libsy/src/algorithms/prefill_probe/artifact.rs create mode 100644 crates/libsy/src/algorithms/prefill_probe/policy.rs diff --git a/Cargo.lock b/Cargo.lock index 74b718f5b..8efc9b130 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", 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::*;