684 lines
22 KiB
Rust
684 lines
22 KiB
Rust
//! Comprehensive tests for the core InferenceEngine
|
|
//!
|
|
//! This test suite validates the complete inference engine functionality
|
|
//! including model loading, graph optimization, batch processing, and
|
|
//! integration with all subsystems.
|
|
|
|
use futures::future;
|
|
use rtx_inference::*;
|
|
use rtx_tensor::{DType, Device, Shape, Tensor};
|
|
use std::collections::HashMap;
|
|
use std::sync::Arc;
|
|
use std::time::{Duration, Instant};
|
|
use tokio::sync::Mutex;
|
|
use tokio_stream::StreamExt;
|
|
use uuid::Uuid;
|
|
|
|
/// Mock model weights for testing
|
|
#[derive(Debug, Clone)]
|
|
struct MockModelWeights {
|
|
layers: HashMap<String, Tensor>,
|
|
config: ModelConfig,
|
|
}
|
|
|
|
fn create_test_config() -> InferenceEngineConfig {
|
|
InferenceEngineConfig {
|
|
model_path: "test_model.safetensors".to_string(),
|
|
device: Device::cpu(),
|
|
max_batch_size: 32,
|
|
max_sequence_length: 2048,
|
|
memory_pool_size: 1024 * 1024 * 1024, // 1GB
|
|
optimization_level: 2,
|
|
enable_caching: true,
|
|
enable_quantization: false,
|
|
request_manager_config: RequestManagerConfig::default(),
|
|
scheduler_config: BatchSchedulerConfig::default(),
|
|
cache_config: KvCacheConfig::default(),
|
|
}
|
|
}
|
|
|
|
impl MockModelWeights {
|
|
fn new() -> InferenceResult<Self> {
|
|
let config = ModelConfig {
|
|
vocab_size: 32000,
|
|
hidden_size: 4096,
|
|
num_layers: 32,
|
|
num_heads: 32,
|
|
max_position_embeddings: 2048,
|
|
layer_norm_epsilon: 1e-6,
|
|
};
|
|
let mut layers = HashMap::new();
|
|
|
|
// Create mock embedding layer
|
|
let embedding_weight =
|
|
Tensor::randn(&[config.vocab_size, config.hidden_size], &Device::cpu())?;
|
|
layers.insert("embedding.weight".to_string(), embedding_weight);
|
|
|
|
// Create mock transformer layers
|
|
for layer_idx in 0..config.num_layers {
|
|
// Attention weights
|
|
let attn_q = Tensor::randn(&[config.hidden_size, config.hidden_size], &Device::cpu())?;
|
|
let attn_k = Tensor::randn(&[config.hidden_size, config.hidden_size], &Device::cpu())?;
|
|
let attn_v = Tensor::randn(&[config.hidden_size, config.hidden_size], &Device::cpu())?;
|
|
let attn_o = Tensor::randn(&[config.hidden_size, config.hidden_size], &Device::cpu())?;
|
|
|
|
layers.insert(
|
|
format!("layers.{}.attention.q_proj.weight", layer_idx),
|
|
attn_q,
|
|
);
|
|
layers.insert(
|
|
format!("layers.{}.attention.k_proj.weight", layer_idx),
|
|
attn_k,
|
|
);
|
|
layers.insert(
|
|
format!("layers.{}.attention.v_proj.weight", layer_idx),
|
|
attn_v,
|
|
);
|
|
layers.insert(
|
|
format!("layers.{}.attention.o_proj.weight", layer_idx),
|
|
attn_o,
|
|
);
|
|
|
|
// MLP weights
|
|
let mlp_gate = Tensor::randn(
|
|
&[config.hidden_size, config.hidden_size * 4],
|
|
&Device::cpu(),
|
|
)?;
|
|
let mlp_up = Tensor::randn(
|
|
&[config.hidden_size, config.hidden_size * 4],
|
|
&Device::cpu(),
|
|
)?;
|
|
let mlp_down = Tensor::randn(
|
|
&[config.hidden_size * 4, config.hidden_size],
|
|
&Device::cpu(),
|
|
)?;
|
|
|
|
layers.insert(
|
|
format!("layers.{}.mlp.gate_proj.weight", layer_idx),
|
|
mlp_gate,
|
|
);
|
|
layers.insert(format!("layers.{}.mlp.up_proj.weight", layer_idx), mlp_up);
|
|
layers.insert(
|
|
format!("layers.{}.mlp.down_proj.weight", layer_idx),
|
|
mlp_down,
|
|
);
|
|
|
|
// Layer norm weights
|
|
let input_layernorm = Tensor::ones(&[config.hidden_size], &Device::cpu())?;
|
|
let post_attention_layernorm = Tensor::ones(&[config.hidden_size], &Device::cpu())?;
|
|
|
|
layers.insert(
|
|
format!("layers.{}.input_layernorm.weight", layer_idx),
|
|
input_layernorm,
|
|
);
|
|
layers.insert(
|
|
format!("layers.{}.post_attention_layernorm.weight", layer_idx),
|
|
post_attention_layernorm,
|
|
);
|
|
}
|
|
|
|
// Output layer
|
|
let lm_head = Tensor::randn(&[config.hidden_size, config.vocab_size], &Device::cpu())?;
|
|
layers.insert("lm_head.weight".to_string(), lm_head);
|
|
|
|
Ok(Self { layers, config })
|
|
}
|
|
|
|
fn get_layer(&self, name: &str) -> Option<&Tensor> {
|
|
self.layers.get(name)
|
|
}
|
|
|
|
fn layer_count(&self) -> usize {
|
|
self.layers.len()
|
|
}
|
|
}
|
|
|
|
// Core InferenceEngine tests - these should fail initially
|
|
|
|
#[tokio::test]
|
|
async fn test_inference_engine_creation() {
|
|
let config = create_test_config();
|
|
|
|
let engine = InferenceEngine::new(config.clone()).await;
|
|
assert!(
|
|
engine.is_ok(),
|
|
"Failed to create inference engine: {:?}",
|
|
engine.err()
|
|
);
|
|
|
|
let engine = engine.unwrap();
|
|
assert_eq!(engine.device(), &config.device);
|
|
assert_eq!(engine.max_batch_size(), config.max_batch_size);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_model_loading_and_registration() {
|
|
let config = create_test_config();
|
|
let mut engine = InferenceEngine::new(config).await.unwrap();
|
|
|
|
// Create mock model weights
|
|
let mock_weights = MockModelWeights::new().unwrap();
|
|
|
|
// Load model should fail initially (not implemented)
|
|
let result = engine
|
|
.load_model("test_model", &mock_weights.layers, &mock_weights.config)
|
|
.await;
|
|
assert!(result.is_ok(), "Model loading should succeed");
|
|
|
|
// Check model is registered
|
|
let models = engine.list_models().await;
|
|
assert!(models.contains(&"test_model".to_string()));
|
|
|
|
// Get model info
|
|
let info = engine.get_model_info("test_model").await.unwrap();
|
|
assert_eq!(info.name, "test_model");
|
|
assert_eq!(info.layer_count, mock_weights.layer_count());
|
|
assert!(info.memory_usage > 0);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_single_inference_request() {
|
|
let config = create_test_config();
|
|
let mut engine = InferenceEngine::new(config).await.unwrap();
|
|
|
|
// Load model
|
|
let mock_weights = MockModelWeights::new().unwrap();
|
|
engine
|
|
.load_model("test_model", &mock_weights.layers, &mock_weights.config)
|
|
.await
|
|
.unwrap();
|
|
|
|
// Create inference request
|
|
let input_tokens = vec![1, 2, 3, 4, 5];
|
|
let request = InferenceRequest::new("test_model".to_string(), input_tokens.clone(), 10);
|
|
let request_id = request.id.clone();
|
|
|
|
// Run inference
|
|
let start_time = Instant::now();
|
|
let result = engine.infer(request).await;
|
|
let inference_time = start_time.elapsed();
|
|
|
|
assert!(result.is_ok(), "Inference should succeed");
|
|
let result = result.unwrap();
|
|
|
|
// Validate output
|
|
assert_eq!(result.request_id, request_id);
|
|
assert!(!result.output_tokens.is_empty());
|
|
assert!(result.output_tokens.len() <= 10);
|
|
assert!(matches!(
|
|
result.finish_reason,
|
|
FinishReason::Length | FinishReason::EndOfSequence
|
|
));
|
|
|
|
// Validate performance metrics
|
|
let metrics = result.metrics.unwrap();
|
|
assert!(metrics.total_time > Duration::ZERO);
|
|
assert!(metrics.generation_time > Duration::ZERO);
|
|
assert_eq!(metrics.input_token_count, input_tokens.len());
|
|
assert_eq!(metrics.output_token_count, result.output_tokens.len());
|
|
assert!(metrics.tokens_per_second > 0.0);
|
|
|
|
println!(
|
|
"Inference completed in {:?}, {:.2} tokens/sec",
|
|
inference_time, metrics.tokens_per_second
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_batch_inference() {
|
|
let mut config = create_test_config();
|
|
config.max_batch_size = 4;
|
|
let mut engine = InferenceEngine::new(config).await.unwrap();
|
|
|
|
// Load model
|
|
let mock_weights = MockModelWeights::new().unwrap();
|
|
engine
|
|
.load_model("test_model", &mock_weights.layers, &mock_weights.config)
|
|
.await
|
|
.unwrap();
|
|
|
|
// Create multiple requests
|
|
let requests = vec![
|
|
InferenceRequest::new("test_model".to_string(), vec![1, 2, 3], 5),
|
|
InferenceRequest::new("test_model".to_string(), vec![4, 5, 6, 7], 5),
|
|
InferenceRequest::new("test_model".to_string(), vec![8, 9], 5),
|
|
InferenceRequest::new("test_model".to_string(), vec![10, 11, 12, 13, 14], 5),
|
|
];
|
|
|
|
// Submit all requests concurrently
|
|
let mut handles = vec![];
|
|
for request in requests {
|
|
let engine_clone = engine.clone();
|
|
let handle = tokio::spawn(async move { engine_clone.infer(request).await });
|
|
handles.push(handle);
|
|
}
|
|
|
|
// Wait for all results
|
|
let results: Vec<_> = future::join_all(handles).await;
|
|
|
|
// Validate all succeeded
|
|
for result in results {
|
|
let inference_result = result.unwrap().unwrap();
|
|
assert!(!inference_result.output_tokens.is_empty());
|
|
assert!(inference_result.metrics.is_some());
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_streaming_inference() {
|
|
let config = create_test_config();
|
|
let mut engine = InferenceEngine::new(config).await.unwrap();
|
|
|
|
// Load model
|
|
let mock_weights = MockModelWeights::new().unwrap();
|
|
engine
|
|
.load_model("test_model", &mock_weights.layers, &mock_weights.config)
|
|
.await
|
|
.unwrap();
|
|
|
|
// Create streaming request
|
|
let mut request = InferenceRequest::new("test_model".to_string(), vec![1, 2, 3], 20);
|
|
request.stream = true;
|
|
|
|
// Start streaming inference
|
|
let mut stream = engine.infer_stream(request).await.unwrap();
|
|
let mut token_count = 0;
|
|
let mut total_time = Duration::ZERO;
|
|
|
|
// Collect streaming tokens
|
|
while let Some(token_result) = stream.next().await {
|
|
assert!(token_result.is_ok());
|
|
let token = token_result.unwrap();
|
|
|
|
assert!(token.id > 0);
|
|
assert!(!token.text.is_empty());
|
|
assert!(token.probability > 0.0);
|
|
|
|
token_count += 1;
|
|
total_time += token.generation_time;
|
|
|
|
if token_count >= 20 {
|
|
break;
|
|
}
|
|
}
|
|
|
|
assert!(token_count > 0);
|
|
assert!(total_time > Duration::ZERO);
|
|
println!(
|
|
"Generated {} tokens via streaming in {:?}",
|
|
token_count, total_time
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_memory_management() {
|
|
let mut config = create_test_config();
|
|
config.memory_pool_size = 100 * 1024 * 1024; // 100MB
|
|
let mut engine = InferenceEngine::new(config).await.unwrap();
|
|
|
|
// Load model
|
|
let mock_weights = MockModelWeights::new().unwrap();
|
|
engine
|
|
.load_model("test_model", &mock_weights.layers, &mock_weights.config)
|
|
.await
|
|
.unwrap();
|
|
|
|
// Get initial memory stats
|
|
let initial_stats = engine.get_memory_stats().await;
|
|
assert!(initial_stats.total_allocated > 0);
|
|
assert!(initial_stats.total_allocated <= 100 * 1024 * 1024);
|
|
|
|
// Run inference to allocate memory
|
|
let request = InferenceRequest::new("test_model".to_string(), vec![1, 2, 3], 10);
|
|
let _result = engine.infer(request).await.unwrap();
|
|
|
|
// Check memory usage increased
|
|
let after_inference_stats = engine.get_memory_stats().await;
|
|
assert!(after_inference_stats.total_allocated >= initial_stats.total_allocated);
|
|
|
|
// Force garbage collection
|
|
engine.gc().await.unwrap();
|
|
|
|
// Memory should be reclaimed
|
|
let after_gc_stats = engine.get_memory_stats().await;
|
|
assert!(after_gc_stats.total_allocated <= after_inference_stats.total_allocated);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_graph_optimization() {
|
|
let mut config = create_test_config();
|
|
config.optimization_level = 3; // Maximum optimization
|
|
let mut engine = InferenceEngine::new(config).await.unwrap();
|
|
|
|
// Load model
|
|
let mock_weights = MockModelWeights::new().unwrap();
|
|
engine
|
|
.load_model("test_model", &mock_weights.layers, &mock_weights.config)
|
|
.await
|
|
.unwrap();
|
|
|
|
// Get optimization stats
|
|
let stats = engine.get_optimization_stats("test_model").await.unwrap();
|
|
assert!(stats.optimizations_applied > 0);
|
|
assert!(stats.nodes_eliminated > 0);
|
|
assert!(stats.memory_saved > 0);
|
|
assert!(stats.estimated_speedup > 1.0);
|
|
|
|
// Compare optimized vs unoptimized performance
|
|
let request = InferenceRequest::new("test_model".to_string(), vec![1, 2, 3, 4, 5], 10);
|
|
|
|
// Measure with optimization
|
|
let start = Instant::now();
|
|
let _optimized_result = engine.infer(request.clone()).await.unwrap();
|
|
let optimized_time = start.elapsed();
|
|
|
|
// Disable optimization temporarily
|
|
engine
|
|
.set_optimization_level("test_model", 0)
|
|
.await
|
|
.unwrap();
|
|
|
|
// Measure without optimization
|
|
let start = Instant::now();
|
|
let _unoptimized_result = engine.infer(request.clone()).await.unwrap();
|
|
let unoptimized_time = start.elapsed();
|
|
|
|
// Optimized should be faster
|
|
assert!(
|
|
optimized_time < unoptimized_time,
|
|
"Optimized inference ({:?}) should be faster than unoptimized ({:?})",
|
|
optimized_time,
|
|
unoptimized_time
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_quantization_integration() {
|
|
let mut config = create_test_config();
|
|
config.enable_quantization = true;
|
|
let mut engine = InferenceEngine::new(config).await.unwrap();
|
|
|
|
// Load model
|
|
let mock_weights = MockModelWeights::new().unwrap();
|
|
engine
|
|
.load_model("test_model", &mock_weights.layers, &mock_weights.config)
|
|
.await
|
|
.unwrap();
|
|
|
|
// Apply quantization
|
|
let quant_config = QuantizationConfig::new()
|
|
.with_scheme(QuantizationScheme::INT8)
|
|
.with_calibration_samples(100);
|
|
|
|
engine
|
|
.quantize_model("test_model", quant_config)
|
|
.await
|
|
.unwrap();
|
|
|
|
// Check quantization was applied
|
|
let model_info = engine.get_model_info("test_model").await.unwrap();
|
|
assert!(model_info.is_quantized);
|
|
assert!(model_info.memory_usage < model_info.original_memory_usage);
|
|
|
|
// Inference should still work
|
|
let request = InferenceRequest::new("test_model".to_string(), vec![1, 2, 3], 5);
|
|
let result = engine.infer(request).await.unwrap();
|
|
assert!(!result.output_tokens.is_empty());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_performance_monitoring() {
|
|
let config = create_test_config();
|
|
let mut engine = InferenceEngine::new(config).await.unwrap();
|
|
|
|
// Load model
|
|
let mock_weights = MockModelWeights::new().unwrap();
|
|
engine
|
|
.load_model("test_model", &mock_weights.layers, &mock_weights.config)
|
|
.await
|
|
.unwrap();
|
|
|
|
// Run multiple inferences to generate metrics
|
|
for i in 0..10 {
|
|
let request = InferenceRequest::new("test_model".to_string(), vec![1, 2, i as i32], 5);
|
|
let _result = engine.infer(request).await.unwrap();
|
|
}
|
|
|
|
// Get performance metrics
|
|
let metrics = engine.get_performance_metrics().await;
|
|
assert!(metrics.total_requests >= 10);
|
|
assert!(metrics.total_tokens_generated > 0);
|
|
assert!(metrics.average_latency > Duration::ZERO);
|
|
assert!(metrics.average_throughput > 0.0);
|
|
assert!(metrics.success_rate >= 0.9); // Should be high
|
|
|
|
// Get model-specific metrics
|
|
let model_metrics = engine.get_model_metrics("test_model").await.unwrap();
|
|
assert!(model_metrics.inference_count >= 10);
|
|
assert!(model_metrics.total_inference_time > Duration::ZERO);
|
|
assert!(model_metrics.average_tokens_per_request > 0.0);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_error_handling() {
|
|
let config = create_test_config();
|
|
let mut engine = InferenceEngine::new(config).await.unwrap();
|
|
|
|
// Test inference without model
|
|
let request = InferenceRequest::new("nonexistent_model".to_string(), vec![1, 2, 3], 5);
|
|
let result = engine.infer(request).await;
|
|
assert!(result.is_err());
|
|
assert!(matches!(
|
|
result.unwrap_err(),
|
|
InferenceError::ModelNotFound { .. }
|
|
));
|
|
|
|
// Load model
|
|
let mock_weights = MockModelWeights::new().unwrap();
|
|
engine
|
|
.load_model("test_model", &mock_weights.layers, &mock_weights.config)
|
|
.await
|
|
.unwrap();
|
|
|
|
// Test with invalid input
|
|
let invalid_request = InferenceRequest::new("test_model".to_string(), vec![], 5); // Empty input
|
|
let result = engine.infer(invalid_request).await;
|
|
assert!(result.is_err());
|
|
assert!(matches!(
|
|
result.unwrap_err(),
|
|
InferenceError::InvalidRequest { .. }
|
|
));
|
|
|
|
// Test with too large input
|
|
let large_input = vec![1i32; 10000]; // Exceed max sequence length
|
|
let large_request = InferenceRequest::new("test_model".to_string(), large_input, 5);
|
|
let result = engine.infer(large_request).await;
|
|
assert!(result.is_err());
|
|
assert!(matches!(
|
|
result.unwrap_err(),
|
|
InferenceError::InvalidRequest { .. }
|
|
));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_concurrent_inference() {
|
|
let mut config = create_test_config();
|
|
config.max_batch_size = 8;
|
|
let mut engine_temp = InferenceEngine::new(config).await.unwrap();
|
|
|
|
// Load model
|
|
let mock_weights = MockModelWeights::new().unwrap();
|
|
engine_temp
|
|
.load_model("test_model", &mock_weights.layers, &mock_weights.config)
|
|
.await
|
|
.unwrap();
|
|
|
|
let engine = Arc::new(engine_temp);
|
|
|
|
// Run concurrent inferences
|
|
let mut handles = vec![];
|
|
for i in 0..20 {
|
|
let engine_clone = engine.clone();
|
|
let handle = tokio::spawn(async move {
|
|
let request = InferenceRequest::new("test_model".to_string(), vec![i, i + 1, i + 2], 3);
|
|
engine_clone.infer(request).await
|
|
});
|
|
handles.push(handle);
|
|
}
|
|
|
|
// Wait for all to complete
|
|
let results = future::join_all(handles).await;
|
|
|
|
// Check all succeeded
|
|
let mut success_count = 0;
|
|
for result in results {
|
|
if result.unwrap().is_ok() {
|
|
success_count += 1;
|
|
}
|
|
}
|
|
|
|
assert_eq!(
|
|
success_count, 20,
|
|
"All concurrent inferences should succeed"
|
|
);
|
|
|
|
// Check performance metrics show concurrent processing
|
|
let metrics = engine.get_performance_metrics().await;
|
|
assert!(metrics.concurrent_requests_handled > 0);
|
|
assert!(metrics.peak_batch_size > 1);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_model_version_management() {
|
|
let config = create_test_config();
|
|
let mut engine = InferenceEngine::new(config).await.unwrap();
|
|
|
|
// Load model version 1
|
|
let mock_weights_v1 = MockModelWeights::new().unwrap();
|
|
engine
|
|
.load_model_version(
|
|
"test_model",
|
|
"v1.0",
|
|
&mock_weights_v1.layers,
|
|
&mock_weights_v1.config,
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
// Load model version 2
|
|
let mock_weights_v2 = MockModelWeights::new().unwrap();
|
|
engine
|
|
.load_model_version(
|
|
"test_model",
|
|
"v2.0",
|
|
&mock_weights_v2.layers,
|
|
&mock_weights_v2.config,
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
// List versions
|
|
let versions = engine.list_model_versions("test_model").await.unwrap();
|
|
assert!(versions.contains(&"v1.0".to_string()));
|
|
assert!(versions.contains(&"v2.0".to_string()));
|
|
|
|
// Switch to specific version
|
|
engine
|
|
.switch_model_version("test_model", "v1.0")
|
|
.await
|
|
.unwrap();
|
|
let current_version = engine
|
|
.get_current_model_version("test_model")
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(current_version, "v1.0");
|
|
|
|
// Test inference works with specific version
|
|
let request = InferenceRequest::new("test_model".to_string(), vec![1, 2, 3], 5);
|
|
let result = engine.infer(request).await.unwrap();
|
|
assert!(!result.output_tokens.is_empty());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_health_checks() {
|
|
let config = create_test_config();
|
|
let mut engine = InferenceEngine::new(config).await.unwrap();
|
|
|
|
// Initial health should be ok
|
|
let health = engine.health_check().await;
|
|
assert!(health.is_healthy);
|
|
assert_eq!(health.status, "Ready");
|
|
|
|
// Load model
|
|
let mock_weights = MockModelWeights::new().unwrap();
|
|
engine
|
|
.load_model("test_model", &mock_weights.layers, &mock_weights.config)
|
|
.await
|
|
.unwrap();
|
|
|
|
// Health should still be good
|
|
let health = engine.health_check().await;
|
|
assert!(health.is_healthy);
|
|
assert!(health.models_loaded > 0);
|
|
|
|
// Detailed health check
|
|
let detailed_health = engine.detailed_health_check().await;
|
|
assert!(detailed_health.is_healthy);
|
|
assert!(detailed_health.memory_usage.total_allocated > 0);
|
|
assert!(detailed_health.memory_usage.utilization < 1.0);
|
|
assert!(detailed_health.model_health.contains_key("test_model"));
|
|
}
|
|
|
|
// Integration tests combining multiple features
|
|
|
|
#[tokio::test]
|
|
async fn test_full_pipeline_integration() {
|
|
let mut config = create_test_config();
|
|
config.max_batch_size = 4;
|
|
config.enable_caching = true;
|
|
config.enable_quantization = true;
|
|
config.optimization_level = 2;
|
|
let mut engine = InferenceEngine::new(config).await.unwrap();
|
|
|
|
// Load and optimize model
|
|
let mock_weights = MockModelWeights::new().unwrap();
|
|
engine
|
|
.load_model("test_model", &mock_weights.layers, &mock_weights.config)
|
|
.await
|
|
.unwrap();
|
|
|
|
// Apply quantization
|
|
let quant_config = QuantizationConfig::new().with_scheme(QuantizationScheme::INT8);
|
|
engine
|
|
.quantize_model("test_model", quant_config)
|
|
.await
|
|
.unwrap();
|
|
|
|
// Run inference workload
|
|
let mut results = vec![];
|
|
for i in 0..10 {
|
|
let request = InferenceRequest::new("test_model".to_string(), vec![i, i + 1, i + 2], 8);
|
|
let result = engine.infer(request).await.unwrap();
|
|
results.push(result);
|
|
}
|
|
|
|
// Validate all succeeded
|
|
assert_eq!(results.len(), 10);
|
|
for result in &results {
|
|
assert!(!result.output_tokens.is_empty());
|
|
assert!(result.metrics.is_some());
|
|
}
|
|
|
|
// Check performance metrics
|
|
let metrics = engine.get_performance_metrics().await;
|
|
assert!(metrics.total_requests >= 10);
|
|
assert!(metrics.success_rate >= 0.9);
|
|
|
|
// Validate memory efficiency
|
|
let memory_stats = engine.get_memory_stats().await;
|
|
assert!(memory_stats.fragmentation < 0.2); // Low fragmentation
|
|
assert!(memory_stats.utilization > 0.1); // Some memory used
|
|
}
|
|
|
|
// Note: TensorError -> InferenceError conversion is handled internally by the crate
|