use super::{InferenceEngine, Precision};
use crate::config::model::{DeviceType, ModelConfig};
use crate::device::memory_limit::{MemoryLimitController, MemoryLimitStatus};
use crate::error::VecboostError;
use crate::monitor::MemoryMonitor;
use crate::utils::hash::verify_sha256;
use crate::utils::hf_hub::build_hf_repo;
use async_trait::async_trait;
use ndarray::{Array1, Array2};
use ort::session::{Session, builder::GraphOptimizationLevel};
use ort::value::Tensor;
use std::sync::{Arc, Mutex};
#[cfg(target_os = "macos")]
use tokenizers::Tokenizer;
#[cfg(target_os = "macos")]
use tokenizers::{PaddingParams, PaddingStrategy};
#[cfg(not(target_os = "macos"))]
type Tokenizer = crate::text::Tokenizer;
pub struct OnnxEngine {
session: Arc<Mutex<Session>>,
tokenizer: Tokenizer,
hidden_size: usize,
max_input_length: usize,
precision: Precision,
memory_monitor: Option<Arc<MemoryMonitor>>,
memory_limit_controller: Option<Arc<MemoryLimitController>>,
fallback_triggered: bool,
fallback_lock: Arc<Mutex<()>>, device_type: DeviceType,
supports_cuda: bool,
}
impl OnnxEngine {
pub fn new(config: &ModelConfig, precision: Precision) -> Result<Self, VecboostError> {
Self::with_device(config, precision, config.device.clone())
}
pub fn with_device(
config: &ModelConfig,
precision: Precision,
device_type: DeviceType,
) -> Result<Self, VecboostError> {
let model_path = &config.model_path;
let is_local_path = model_path.exists() && model_path.is_dir();
let (onnx_filename, tokenizer_filename) = if is_local_path {
log::info!("Using local model path: {:?}", model_path);
let onnx_filename = if model_path.join("model_quantized.onnx").exists() {
model_path.join("model_quantized.onnx")
} else if model_path.join("model.onnx").exists() {
model_path.join("model.onnx")
} else {
return Err(VecboostError::ModelLoadError(format!(
"No ONNX model found in {:?}",
model_path
)));
};
let tokenizer_filename = if model_path.join("tokenizer.json").exists() {
model_path.join("tokenizer.json")
} else {
let parent = model_path.parent().and_then(|p| p.parent());
let cache_path = parent.ok_or_else(|| {
VecboostError::ModelLoadError(format!(
"Cannot determine HuggingFace cache path from {:?}",
model_path
))
})?;
let cache_tokenizer = cache_path.join("tokenizer.json");
if cache_tokenizer.exists() {
log::info!(
"Using tokenizer from HuggingFace cache: {:?}",
cache_tokenizer
);
cache_tokenizer
} else {
return Err(VecboostError::ModelLoadError(format!(
"Tokenizer not found at {:?} or {:?}",
model_path.join("tokenizer.json"),
cache_tokenizer
)));
}
};
(onnx_filename, tokenizer_filename)
} else {
log::info!("Using HuggingFace Hub for model: {:?}", model_path);
let repo_id = model_path.to_string_lossy().into_owned();
let repo = build_hf_repo(&repo_id)?;
log::info!("Downloading/Loading ONNX model files...");
let onnx_filename = repo
.download_file()
.filename("model.onnx")
.send()
.or_else(|_| repo.download_file().filename("model_quantized.onnx").send())
.map_err(|e| VecboostError::ModelLoadError(e.to_string()))?;
let tokenizer_filename = repo
.download_file()
.filename("tokenizer.json")
.send()
.map_err(|e| VecboostError::ModelLoadError(e.to_string()))?;
(onnx_filename, tokenizer_filename)
};
let num_threads = std::thread::available_parallelism()
.map(|p| p.get())
.unwrap_or(4);
let supports_cuda = device_type == DeviceType::Cuda;
let supports_amd = matches!(device_type, DeviceType::Amd | DeviceType::OpenCL);
log::info!("Initializing ONNX Runtime session...");
let session = Session::builder()
.map_err(|e| VecboostError::ModelLoadError(e.to_string()))?
.with_optimization_level(GraphOptimizationLevel::Level3)
.map_err(|e| VecboostError::ModelLoadError(e.to_string()))?
.with_intra_threads(num_threads)
.map_err(|e| VecboostError::ModelLoadError(e.to_string()))?;
let mut session = if supports_cuda {
log::info!("Attempting to configure CUDA execution provider for ONNX Runtime");
#[cfg(feature = "cuda")]
{
let cuda_provider =
ort::execution_providers::CUDAExecutionProvider::default().build();
session
.with_execution_providers([cuda_provider])
.map_err(|e| VecboostError::ModelLoadError(e.to_string()))?
}
#[cfg(not(feature = "cuda"))]
{
log::warn!(
"CUDA execution provider not available. ONNX Runtime CUDA support requires cuda feature flag. Using CPU execution provider."
);
session
}
} else if supports_amd {
log::info!(
"AMD GPU detected but ROCm execution provider is not configured in this build. Using CPU execution provider."
);
session
} else {
session
};
if let Some(ref expected_hash) = config.model_sha256 {
log::info!("Verifying model file SHA256 hash...");
let is_valid = verify_sha256(&onnx_filename, expected_hash).map_err(|e| {
VecboostError::ModelLoadError(format!("Failed to verify SHA256: {}", e))
})?;
if !is_valid {
return Err(VecboostError::ModelLoadError(format!(
"Model file SHA256 verification failed. Expected: {}, File: {:?}",
expected_hash, onnx_filename
)));
}
log::info!("Model file SHA256 verification passed");
}
let session = session
.commit_from_file(onnx_filename)
.map_err(|e| VecboostError::ModelLoadError(e.to_string()))?;
log::info!("Loading tokenizer...");
#[allow(unused_mut)]
let mut tokenizer = Tokenizer::from_file(&tokenizer_filename.to_string_lossy())
.map_err(|e| VecboostError::ModelLoadError(e.to_string()))?;
#[cfg(target_os = "macos")]
{
if let Some(pp) = tokenizer.get_padding_mut() {
pp.strategy = PaddingStrategy::BatchLongest;
} else {
let pp = PaddingParams {
strategy: PaddingStrategy::BatchLongest,
..Default::default()
};
tokenizer.with_padding(Some(pp));
}
}
let vocab_size = tokenizer.get_vocab_size();
let hidden_size = config.expected_dimension.unwrap_or(1024);
log::info!(
"Using hidden_size from configuration: {:?}",
config.expected_dimension
);
log::info!("Final hidden_size value: {}", hidden_size);
let max_input_length = std::cmp::min(vocab_size, 512);
let actual_precision = match precision {
Precision::Fp16 => {
if supports_cuda || supports_amd {
log::info!("Using FP16 precision with GPU acceleration");
Precision::Fp16
} else {
log::warn!("FP16 not supported without GPU acceleration, falling back to FP32");
Precision::Fp32
}
}
_ => {
log::info!("Using {} precision", precision);
precision
}
};
log::info!(
"ONNX Engine initialized: hidden_size={}, max_input_length={}, vocab_size={}, precision={:?}",
hidden_size,
max_input_length,
vocab_size,
actual_precision
);
let memory_monitor = if supports_cuda || supports_amd {
Some(Arc::new(MemoryMonitor::new()))
} else {
None
};
Ok(Self {
session: Arc::new(Mutex::new(session)),
tokenizer,
hidden_size,
max_input_length,
precision: actual_precision,
memory_monitor,
memory_limit_controller: None,
fallback_triggered: false,
fallback_lock: Arc::new(Mutex::new(())),
device_type,
supports_cuda,
})
}
pub fn set_memory_limit_controller(&mut self, controller: Arc<MemoryLimitController>) {
self.memory_limit_controller = Some(controller);
}
pub fn device_type(&self) -> DeviceType {
self.device_type.clone()
}
pub fn is_fallback_triggered(&self) -> bool {
self.fallback_triggered
}
pub fn check_memory_pressure(&self, threshold_percent: u64) -> bool {
if let Some(ref monitor) = self.memory_monitor {
if let Ok(handle) = tokio::runtime::Handle::try_current() {
handle.block_on(async {
let stats = monitor.get_memory_stats().await;
let usage_percent = (stats.current_bytes * 100)
.checked_div(stats.total_bytes)
.unwrap_or(0);
usage_percent >= threshold_percent
})
} else {
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.unwrap_or_else(|_| {
log::error!("Failed to create Tokio runtime for memory check");
std::process::exit(1);
});
rt.block_on(async {
let stats = monitor.get_memory_stats().await;
let usage_percent = (stats.current_bytes * 100)
.checked_div(stats.total_bytes)
.unwrap_or(0);
usage_percent >= threshold_percent
})
}
} else {
false
}
}
pub async fn check_memory_limit_and_fallback(
&mut self,
config: &ModelConfig,
) -> Result<bool, VecboostError> {
if self.fallback_triggered {
return Ok(false);
}
if let Some(ref controller) = self.memory_limit_controller {
let status = controller.check_limit().await;
if status == MemoryLimitStatus::Exceeded {
log::warn!("Memory limit exceeded for ONNX engine, attempting fallback to CPU");
self.try_fallback_to_cpu(config).await?;
return Ok(true);
} else if status == MemoryLimitStatus::Critical {
log::warn!(
"Memory limit critical for ONNX engine, checking memory pressure for fallback"
);
if self.check_memory_pressure(90) {
self.try_fallback_to_cpu(config).await?;
return Ok(true);
}
}
}
Ok(false)
}
pub async fn update_memory_limit(&self, used_bytes: u64) {
if let Some(ref controller) = self.memory_limit_controller {
controller.update_usage(used_bytes).await;
}
}
pub async fn get_memory_status(&self) -> Option<MemoryLimitStatus> {
if let Some(ref controller) = self.memory_limit_controller {
Some(controller.check_limit().await)
} else {
None
}
}
pub async fn update_gpu_memory(&self) {
if let Some(ref monitor) = self.memory_monitor {
#[cfg(feature = "onnx")]
monitor.update_gpu_memory_from_ort().await;
}
}
fn forward_pass(&self, text: &str) -> Result<Vec<f32>, VecboostError> {
let tokens = self
.tokenizer
.encode(text, true)
.map_err(|e| VecboostError::TokenizationError(e.to_string()))?;
let input_ids: Vec<i64> = tokens
.get_ids()
.iter()
.take(self.max_input_length)
.map(|&id| id as i64)
.collect();
let attention_mask: Vec<i64> = tokens
.get_attention_mask()
.iter()
.take(self.max_input_length)
.map(|&v| v as i64)
.collect();
let input_ids_array = Array1::from(input_ids);
let attention_mask_array = Array1::from(attention_mask.clone());
let input_ids_tensor = Tensor::from_array(input_ids_array.into_dyn())
.map_err(|e| VecboostError::InferenceError(e.to_string()))?;
let attention_mask_tensor = Tensor::from_array(attention_mask_array.into_dyn())
.map_err(|e: ort::Error| VecboostError::InferenceError(e.to_string()))?;
let last_hidden_state = {
let mut session_guard = self
.session
.lock()
.map_err(|e| VecboostError::InferenceError(e.to_string()))?;
let outputs = session_guard
.run(ort::inputs![
"input_ids" => input_ids_tensor,
"attention_mask" => attention_mask_tensor
])
.map_err(|e| VecboostError::InferenceError(e.to_string()))?;
let output_array = outputs["last_hidden_state"]
.try_extract_array::<f32>()
.map_err(|e| VecboostError::InferenceError(e.to_string()))?
.to_owned();
log::debug!("ONNX model output shape: {:?}", output_array.shape());
output_array
};
let seq_len = attention_mask.iter().filter(|&&v| v == 1).count();
if seq_len == 0 {
return Err(VecboostError::InferenceError(
"Empty sequence after mask".to_string(),
));
}
let mut weighted_sum = vec![0.0f32; self.hidden_size];
let mut mask_sum = 0.0f32;
for (seq_idx, &mask_val) in attention_mask.iter().enumerate().take(seq_len) {
if mask_val == 1 {
for h in 0..self.hidden_size {
let token_embedding = last_hidden_state[[seq_idx, h]];
weighted_sum[h] += token_embedding * mask_val as f32;
mask_sum += mask_val as f32;
}
}
}
if mask_sum > 0.0 {
weighted_sum.iter_mut().for_each(|v| *v /= mask_sum);
}
Ok(weighted_sum)
}
pub async fn try_fallback_to_cpu(&mut self, config: &ModelConfig) -> Result<(), VecboostError> {
let _lock = self.fallback_lock.lock().map_err(|e| {
VecboostError::InferenceError(format!("Failed to acquire fallback lock: {}", e))
})?;
if self.fallback_triggered {
return Ok(());
}
log::info!("Attempting fallback from GPU to CPU for ONNX engine");
self.memory_monitor = None;
self.device_type = DeviceType::Cpu;
self.supports_cuda = false;
self.fallback_triggered = true;
let repo_id = config.model_path.to_string_lossy().into_owned();
let repo = build_hf_repo(&repo_id)?;
let onnx_filename = repo
.download_file()
.filename("model.onnx")
.send()
.or_else(|_| repo.download_file().filename("model_quantized.onnx").send())
.map_err(|e| VecboostError::ModelLoadError(e.to_string()))?;
let num_threads = std::thread::available_parallelism()
.map(|p| p.get())
.unwrap_or(4);
let new_session = Session::builder()
.map_err(|e| VecboostError::ModelLoadError(e.to_string()))?
.with_optimization_level(GraphOptimizationLevel::Level3)
.map_err(|e| VecboostError::ModelLoadError(e.to_string()))?
.with_intra_threads(num_threads)
.map_err(|e| VecboostError::ModelLoadError(e.to_string()))?
.commit_from_file(onnx_filename)
.map_err(|e| VecboostError::ModelLoadError(e.to_string()))?;
let mut session_guard = self
.session
.lock()
.map_err(|e| VecboostError::ModelLoadError(e.to_string()))?;
*session_guard = new_session;
drop(session_guard);
self.precision = Precision::Fp32;
log::info!("Successfully fell back to CPU for ONNX engine");
Ok(())
}
fn forward_pass_batch(&self, texts: &[String]) -> Result<Vec<Vec<f32>>, VecboostError> {
let batch_size = texts.len();
if batch_size == 0 {
return Ok(vec![]);
}
let mut all_input_ids: Vec<Vec<i64>> = Vec::with_capacity(batch_size);
let mut all_attention_masks: Vec<Vec<i64>> = Vec::with_capacity(batch_size);
let mut max_seq_len = 0;
for text in texts {
let text_ref: &str = text.as_str();
let tokens = self
.tokenizer
.encode(text_ref, true)
.map_err(|e| VecboostError::TokenizationError(e.to_string()))?;
let input_ids: Vec<i64> = tokens
.get_ids()
.iter()
.take(self.max_input_length)
.map(|&id| id as i64)
.collect();
let attention_mask: Vec<i64> = tokens
.get_attention_mask()
.iter()
.take(self.max_input_length)
.map(|&v| v as i64)
.collect();
if input_ids.len() > max_seq_len {
max_seq_len = input_ids.len();
}
all_input_ids.push(input_ids);
all_attention_masks.push(attention_mask);
}
let padded_batch_size = all_input_ids.len();
let mut batch_input_ids = vec![0i64; padded_batch_size * max_seq_len];
let mut batch_attention_mask = vec![0i64; padded_batch_size * max_seq_len];
for (batch_idx, (input_ids, attention_mask)) in all_input_ids
.iter()
.zip(all_attention_masks.iter())
.enumerate()
{
for (seq_idx, (&id, &mask)) in input_ids.iter().zip(attention_mask.iter()).enumerate() {
let pos = batch_idx * max_seq_len + seq_idx;
batch_input_ids[pos] = id;
batch_attention_mask[pos] = mask;
}
}
let input_ids_array =
Array2::from_shape_vec((padded_batch_size, max_seq_len), batch_input_ids)
.map_err(|e| VecboostError::InferenceError(e.to_string()))?;
let attention_mask_array =
Array2::from_shape_vec((padded_batch_size, max_seq_len), batch_attention_mask)
.map_err(|e| VecboostError::InferenceError(e.to_string()))?;
let input_ids_tensor = Tensor::from_array(input_ids_array.into_dyn())
.map_err(|e| VecboostError::InferenceError(e.to_string()))?;
let attention_mask_tensor = Tensor::from_array(attention_mask_array.into_dyn())
.map_err(|e| VecboostError::InferenceError(e.to_string()))?;
let last_hidden_state = {
let mut session_guard = self
.session
.lock()
.map_err(|e| VecboostError::InferenceError(e.to_string()))?;
let outputs = session_guard
.run(ort::inputs![
"input_ids" => input_ids_tensor,
"attention_mask" => attention_mask_tensor
])
.map_err(|e| VecboostError::InferenceError(e.to_string()))?;
outputs["last_hidden_state"]
.try_extract_array::<f32>()
.map_err(|e| VecboostError::InferenceError(e.to_string()))?
.to_owned()
};
let mut results = Vec::with_capacity(padded_batch_size);
for batch_idx in 0..padded_batch_size {
let attention_mask = &all_attention_masks[batch_idx];
let actual_seq_len = attention_mask.len();
log::debug!(
"batch_idx: {}, max_seq_len: {}, actual_seq_len: {}",
batch_idx,
max_seq_len,
actual_seq_len
);
let mut weighted_sum = vec![0.0f32; self.hidden_size];
let mut mask_sum = 0.0f32;
let effective_max_seq = std::cmp::min(max_seq_len, actual_seq_len);
for seq_idx in 0..effective_max_seq {
let mask_val = attention_mask[seq_idx];
if mask_val == 1 {
for h in 0..self.hidden_size {
let token_embedding = last_hidden_state[[batch_idx, seq_idx, h]];
weighted_sum[h] += token_embedding * mask_val as f32;
mask_sum += mask_val as f32;
}
}
}
if mask_sum > 0.0 {
weighted_sum
.iter_mut()
.take(self.hidden_size)
.for_each(|v| *v /= mask_sum);
}
results.push(weighted_sum);
}
Ok(results)
}
}
#[async_trait]
impl InferenceEngine for OnnxEngine {
fn embed(&self, text: &str) -> Result<Vec<f32>, VecboostError> {
self.forward_pass(text)
}
fn embed_batch(&self, texts: &[String]) -> Result<Vec<Vec<f32>>, VecboostError> {
self.forward_pass_batch(texts)
}
fn precision(&self) -> &Precision {
&self.precision
}
fn supports_mixed_precision(&self) -> bool {
self.memory_monitor.is_some()
}
fn is_fallback_triggered(&self) -> bool {
self.fallback_triggered
}
async fn try_fallback_to_cpu(&mut self, config: &ModelConfig) -> Result<(), VecboostError> {
OnnxEngine::try_fallback_to_cpu(self, config).await
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::config::model::{EngineType, ModelConfig};
use std::path::PathBuf;
fn test_config() -> ModelConfig {
ModelConfig {
name: "test-onnx".to_string(),
engine_type: EngineType::Onnx,
model_path: PathBuf::from("/nonexistent/onnx-model"),
tokenizer_path: None,
device: DeviceType::Cpu,
max_batch_size: 32,
pooling_mode: None,
expected_dimension: Some(1024),
memory_limit_bytes: None,
oom_fallback_enabled: true,
model_sha256: None,
}
}
#[test]
fn test_onnx_engine_new_returns_error_for_empty_dir() {
let temp_dir = tempfile::tempdir().expect("Failed to create temp dir");
let mut config = test_config();
config.model_path = temp_dir.path().to_path_buf();
let result = OnnxEngine::new(&config, Precision::Fp32);
assert!(result.is_err());
if let Err(e) = result {
assert!(
matches!(e, VecboostError::ModelLoadError(_)),
"Expected ModelLoadError, got {:?}",
e
);
}
}
#[test]
fn test_onnx_engine_with_device_empty_dir_cpu() {
let temp_dir = tempfile::tempdir().expect("Failed to create temp dir");
let mut config = test_config();
config.model_path = temp_dir.path().to_path_buf();
let result = OnnxEngine::with_device(&config, Precision::Fp32, DeviceType::Cpu);
assert!(result.is_err());
if let Err(VecboostError::ModelLoadError(msg)) = result {
assert!(
msg.contains("No ONNX model found"),
"Expected 'No ONNX model found', got: {}",
msg
);
}
}
#[test]
fn test_onnx_engine_with_device_fp16_empty_dir() {
let temp_dir = tempfile::tempdir().expect("Failed to create temp dir");
let mut config = test_config();
config.model_path = temp_dir.path().to_path_buf();
let result = OnnxEngine::with_device(&config, Precision::Fp16, DeviceType::Cpu);
assert!(result.is_err());
}
#[test]
fn test_onnx_engine_with_device_int8_empty_dir() {
let temp_dir = tempfile::tempdir().expect("Failed to create temp dir");
let mut config = test_config();
config.model_path = temp_dir.path().to_path_buf();
let result = OnnxEngine::with_device(&config, Precision::Int8, DeviceType::Cpu);
assert!(result.is_err());
}
#[test]
fn test_onnx_engine_with_device_cuda_empty_dir() {
let temp_dir = tempfile::tempdir().expect("Failed to create temp dir");
let mut config = test_config();
config.model_path = temp_dir.path().to_path_buf();
let result = OnnxEngine::with_device(&config, Precision::Fp32, DeviceType::Cuda);
assert!(result.is_err());
if let Err(e) = result {
assert!(
matches!(e, VecboostError::ModelLoadError(_)),
"Expected ModelLoadError, got {:?}",
e
);
}
}
#[test]
fn test_onnx_engine_with_device_amd_empty_dir() {
let temp_dir = tempfile::tempdir().expect("Failed to create temp dir");
let mut config = test_config();
config.model_path = temp_dir.path().to_path_buf();
let result = OnnxEngine::with_device(&config, Precision::Fp32, DeviceType::Amd);
assert!(result.is_err());
}
#[test]
fn test_onnx_engine_with_device_opencl_empty_dir() {
let temp_dir = tempfile::tempdir().expect("Failed to create temp dir");
let mut config = test_config();
config.model_path = temp_dir.path().to_path_buf();
let result = OnnxEngine::with_device(&config, Precision::Fp32, DeviceType::OpenCL);
assert!(result.is_err());
}
#[test]
fn test_onnx_engine_with_model_but_no_tokenizer() {
let temp_dir = tempfile::tempdir().expect("Failed to create temp dir");
std::fs::write(temp_dir.path().join("model.onnx"), b"fake onnx")
.expect("Failed to write fake model.onnx");
let mut config = test_config();
config.model_path = temp_dir.path().to_path_buf();
let result = OnnxEngine::with_device(&config, Precision::Fp32, DeviceType::Cpu);
assert!(result.is_err());
if let Err(VecboostError::ModelLoadError(msg)) = result {
assert!(
msg.contains("Tokenizer not found"),
"Expected 'Tokenizer not found', got: {}",
msg
);
}
}
#[test]
fn test_onnx_engine_prefers_quantized_model_but_no_tokenizer() {
let temp_dir = tempfile::tempdir().expect("Failed to create temp dir");
std::fs::write(
temp_dir.path().join("model_quantized.onnx"),
b"fake quantized",
)
.expect("Failed to write fake model_quantized.onnx");
std::fs::write(temp_dir.path().join("model.onnx"), b"fake onnx")
.expect("Failed to write fake model.onnx");
let mut config = test_config();
config.model_path = temp_dir.path().to_path_buf();
let result = OnnxEngine::with_device(&config, Precision::Fp32, DeviceType::Cpu);
assert!(result.is_err());
if let Err(VecboostError::ModelLoadError(msg)) = result {
assert!(
msg.contains("Tokenizer not found"),
"Expected 'Tokenizer not found', got: {}",
msg
);
}
}
#[test]
fn test_precision_display() {
assert_eq!(Precision::Fp32.to_string(), "fp32");
assert_eq!(Precision::Fp16.to_string(), "fp16");
assert_eq!(Precision::Int8.to_string(), "int8");
}
#[test]
fn test_device_type_serialization() {
assert_eq!(
serde_json::to_string(&DeviceType::Cpu).expect("serialize Cpu"),
"\"cpu\""
);
assert_eq!(
serde_json::to_string(&DeviceType::Amd).expect("serialize Amd"),
"\"amd\""
);
assert_eq!(
serde_json::to_string(&DeviceType::OpenCL).expect("serialize OpenCL"),
"\"opencl\""
);
}
#[ignore = "ort crate environment cleanup crashes test runner when Session::builder() is called"]
#[test]
fn test_onnx_engine_model_and_tokenizer_fails_at_session() {
let temp_dir = tempfile::tempdir().expect("Failed to create temp dir");
std::fs::write(temp_dir.path().join("model.onnx"), b"fake onnx")
.expect("Failed to write fake model.onnx");
std::fs::write(temp_dir.path().join("tokenizer.json"), b"{}")
.expect("Failed to write fake tokenizer.json");
let mut config = test_config();
config.model_path = temp_dir.path().to_path_buf();
let result = OnnxEngine::with_device(&config, Precision::Fp32, DeviceType::Cpu);
assert!(result.is_err());
if let Err(e) = result {
assert!(
matches!(e, VecboostError::ModelLoadError(_)),
"Expected ModelLoadError from session commit, got {:?}",
e
);
}
}
#[ignore = "ort crate environment cleanup crashes test runner when Session::builder() is called"]
#[test]
fn test_onnx_engine_quantized_preferred_with_tokenizer() {
let temp_dir = tempfile::tempdir().expect("Failed to create temp dir");
std::fs::write(
temp_dir.path().join("model_quantized.onnx"),
b"fake quantized",
)
.expect("Failed to write fake model_quantized.onnx");
std::fs::write(temp_dir.path().join("model.onnx"), b"fake onnx")
.expect("Failed to write fake model.onnx");
std::fs::write(temp_dir.path().join("tokenizer.json"), b"{}")
.expect("Failed to write fake tokenizer.json");
let mut config = test_config();
config.model_path = temp_dir.path().to_path_buf();
let result = OnnxEngine::with_device(&config, Precision::Fp32, DeviceType::Cpu);
assert!(result.is_err());
if let Err(e) = result {
assert!(
matches!(e, VecboostError::ModelLoadError(_)),
"Expected ModelLoadError from session commit, got {:?}",
e
);
}
}
#[ignore = "ort crate environment cleanup crashes test runner when Session::builder() is called"]
#[test]
fn test_onnx_engine_tokenizer_from_cache_path() {
let temp_dir = tempfile::tempdir().expect("Failed to create temp dir");
let cache_dir = temp_dir.path();
let intermediate = cache_dir.join("intermediate");
let model_dir = intermediate.join("model_dir");
std::fs::create_dir_all(&model_dir).expect("Failed to create model_dir");
std::fs::write(model_dir.join("model.onnx"), b"fake onnx")
.expect("Failed to write fake model.onnx");
std::fs::write(cache_dir.join("tokenizer.json"), b"{}")
.expect("Failed to write cache tokenizer.json");
let mut config = test_config();
config.model_path = model_dir;
let result = OnnxEngine::with_device(&config, Precision::Fp32, DeviceType::Cpu);
assert!(result.is_err());
if let Err(e) = result {
assert!(
matches!(e, VecboostError::ModelLoadError(_)),
"Expected ModelLoadError after cache tokenizer load, got {:?}",
e
);
}
}
#[ignore = "ort crate environment cleanup crashes test runner when Session::builder() is called"]
#[test]
fn test_onnx_engine_sha256_verification_mismatch() {
let temp_dir = tempfile::tempdir().expect("Failed to create temp dir");
std::fs::write(temp_dir.path().join("model.onnx"), b"fake onnx")
.expect("Failed to write fake model.onnx");
std::fs::write(temp_dir.path().join("tokenizer.json"), b"{}")
.expect("Failed to write fake tokenizer.json");
let mut config = test_config();
config.model_path = temp_dir.path().to_path_buf();
config.model_sha256 =
Some("0000000000000000000000000000000000000000000000000000000000000000".to_string());
let result = OnnxEngine::with_device(&config, Precision::Fp32, DeviceType::Cpu);
assert!(result.is_err());
if let Err(VecboostError::ModelLoadError(msg)) = result {
assert!(
msg.contains("SHA256 verification failed"),
"Expected SHA256 verification failure, got: {}",
msg
);
}
}
#[ignore = "ort crate environment cleanup crashes test runner when Session::builder() is called"]
#[test]
fn test_onnx_engine_cuda_device_with_model_and_tokenizer() {
let temp_dir = tempfile::tempdir().expect("Failed to create temp dir");
std::fs::write(temp_dir.path().join("model.onnx"), b"fake onnx")
.expect("Failed to write fake model.onnx");
std::fs::write(temp_dir.path().join("tokenizer.json"), b"{}")
.expect("Failed to write fake tokenizer.json");
let mut config = test_config();
config.model_path = temp_dir.path().to_path_buf();
let result = OnnxEngine::with_device(&config, Precision::Fp32, DeviceType::Cuda);
assert!(result.is_err());
if let Err(e) = result {
assert!(
matches!(e, VecboostError::ModelLoadError(_)),
"Expected ModelLoadError from CUDA path, got {:?}",
e
);
}
}
#[ignore = "ort crate environment cleanup crashes test runner when Session::builder() is called"]
#[test]
fn test_onnx_engine_amd_device_with_model_and_tokenizer() {
let temp_dir = tempfile::tempdir().expect("Failed to create temp dir");
std::fs::write(temp_dir.path().join("model.onnx"), b"fake onnx")
.expect("Failed to write fake model.onnx");
std::fs::write(temp_dir.path().join("tokenizer.json"), b"{}")
.expect("Failed to write fake tokenizer.json");
let mut config = test_config();
config.model_path = temp_dir.path().to_path_buf();
let result = OnnxEngine::with_device(&config, Precision::Fp32, DeviceType::Amd);
assert!(result.is_err());
if let Err(e) = result {
assert!(
matches!(e, VecboostError::ModelLoadError(_)),
"Expected ModelLoadError from AMD path, got {:?}",
e
);
}
}
#[ignore = "ort crate environment cleanup crashes test runner when Session::builder() is called"]
#[test]
fn test_onnx_engine_opencl_device_with_model_and_tokenizer() {
let temp_dir = tempfile::tempdir().expect("Failed to create temp dir");
std::fs::write(temp_dir.path().join("model.onnx"), b"fake onnx")
.expect("Failed to write fake model.onnx");
std::fs::write(temp_dir.path().join("tokenizer.json"), b"{}")
.expect("Failed to write fake tokenizer.json");
let mut config = test_config();
config.model_path = temp_dir.path().to_path_buf();
let result = OnnxEngine::with_device(&config, Precision::Fp32, DeviceType::OpenCL);
assert!(result.is_err());
if let Err(e) = result {
assert!(
matches!(e, VecboostError::ModelLoadError(_)),
"Expected ModelLoadError from OpenCL path, got {:?}",
e
);
}
}
#[ignore = "ort crate environment cleanup crashes test runner when Session::builder() is called"]
#[test]
fn test_onnx_engine_fp16_cuda_device_with_model() {
let temp_dir = tempfile::tempdir().expect("Failed to create temp dir");
std::fs::write(temp_dir.path().join("model.onnx"), b"fake onnx")
.expect("Failed to write fake model.onnx");
std::fs::write(temp_dir.path().join("tokenizer.json"), b"{}")
.expect("Failed to write fake tokenizer.json");
let mut config = test_config();
config.model_path = temp_dir.path().to_path_buf();
let result = OnnxEngine::with_device(&config, Precision::Fp16, DeviceType::Cuda);
assert!(result.is_err());
if let Err(e) = result {
assert!(
matches!(e, VecboostError::ModelLoadError(_)),
"Expected ModelLoadError from FP16+CUDA path, got {:?}",
e
);
}
}
#[test]
fn test_onnx_engine_no_cache_path_determinable() {
let temp_dir = tempfile::tempdir().expect("Failed to create temp dir");
std::fs::write(temp_dir.path().join("model.onnx"), b"fake onnx")
.expect("Failed to write fake model.onnx");
let mut config = test_config();
config.model_path = temp_dir.path().to_path_buf();
let result = OnnxEngine::with_device(&config, Precision::Fp32, DeviceType::Cpu);
assert!(result.is_err());
if let Err(VecboostError::ModelLoadError(msg)) = result {
assert!(
msg.contains("Tokenizer not found") || msg.contains("Cannot determine"),
"Expected tokenizer/cache path error, got: {}",
msg
);
}
}
#[test]
fn test_onnx_engine_cache_path_determinable_no_tokenizer() {
let temp_dir = tempfile::tempdir().expect("Failed to create temp dir");
let intermediate = temp_dir.path().join("intermediate");
let model_dir = intermediate.join("model_dir");
std::fs::create_dir_all(&model_dir).expect("Failed to create model_dir");
std::fs::write(model_dir.join("model.onnx"), b"fake onnx")
.expect("Failed to write fake model.onnx");
let mut config = test_config();
config.model_path = model_dir;
let result = OnnxEngine::with_device(&config, Precision::Fp32, DeviceType::Cpu);
assert!(result.is_err());
if let Err(VecboostError::ModelLoadError(msg)) = result {
assert!(
msg.contains("Tokenizer not found"),
"Expected 'Tokenizer not found', got: {}",
msg
);
assert!(
!msg.contains("Cannot determine"),
"Should not be 'Cannot determine' since cache path exists, got: {}",
msg
);
}
}
#[test]
fn test_onnx_engine_only_quantized_no_model_onnx() {
let temp_dir = tempfile::tempdir().expect("Failed to create temp dir");
std::fs::write(
temp_dir.path().join("model_quantized.onnx"),
b"fake quantized",
)
.expect("Failed to write fake model_quantized.onnx");
let mut config = test_config();
config.model_path = temp_dir.path().to_path_buf();
let result = OnnxEngine::with_device(&config, Precision::Fp32, DeviceType::Cpu);
assert!(result.is_err());
if let Err(VecboostError::ModelLoadError(msg)) = result {
assert!(
msg.contains("Tokenizer not found") || msg.contains("Cannot determine"),
"Expected tokenizer/cache error, got: {}",
msg
);
}
}
#[test]
fn test_onnx_engine_fp16_amd_empty_dir() {
let temp_dir = tempfile::tempdir().expect("Failed to create temp dir");
let mut config = test_config();
config.model_path = temp_dir.path().to_path_buf();
let result = OnnxEngine::with_device(&config, Precision::Fp16, DeviceType::Amd);
assert!(result.is_err());
}
#[test]
fn test_onnx_engine_fp16_opencl_empty_dir() {
let temp_dir = tempfile::tempdir().expect("Failed to create temp dir");
let mut config = test_config();
config.model_path = temp_dir.path().to_path_buf();
let result = OnnxEngine::with_device(&config, Precision::Fp16, DeviceType::OpenCL);
assert!(result.is_err());
}
#[test]
fn test_onnx_engine_int8_cuda_empty_dir() {
let temp_dir = tempfile::tempdir().expect("Failed to create temp dir");
let mut config = test_config();
config.model_path = temp_dir.path().to_path_buf();
let result = OnnxEngine::with_device(&config, Precision::Int8, DeviceType::Cuda);
assert!(result.is_err());
}
#[test]
fn test_onnx_engine_int8_amd_empty_dir() {
let temp_dir = tempfile::tempdir().expect("Failed to create temp dir");
let mut config = test_config();
config.model_path = temp_dir.path().to_path_buf();
let result = OnnxEngine::with_device(&config, Precision::Int8, DeviceType::Amd);
assert!(result.is_err());
}
#[test]
fn test_onnx_engine_model_path_is_file_not_dir() {
let temp_dir = tempfile::tempdir().expect("Failed to create temp dir");
let file_path = temp_dir.path().join("not_a_dir");
std::fs::write(&file_path, b"file").expect("Failed to write file");
let mut config = test_config();
config.model_path = file_path;
let result = OnnxEngine::with_device(&config, Precision::Fp32, DeviceType::Cpu);
assert!(result.is_err());
}
#[test]
fn test_onnx_engine_metal_empty_dir() {
let temp_dir = tempfile::tempdir().expect("Failed to create temp dir");
let mut config = test_config();
config.model_path = temp_dir.path().to_path_buf();
let result = OnnxEngine::with_device(&config, Precision::Fp32, DeviceType::Metal);
assert!(result.is_err());
}
}