239 lines
8.2 KiB
Rust
239 lines
8.2 KiB
Rust
use rtx_auto::{
|
|
agents::{
|
|
DataEngineeringAgent, KernelSynthesizerAgent, ParallelPlannerAgent, QuantGuardianAgent,
|
|
},
|
|
error::AutoError,
|
|
proposal::{Proposal, ProposalType, ProposalValidator},
|
|
rollback::RollbackManager,
|
|
};
|
|
use rtx_runtime::Runtime;
|
|
use rtx_tensor::{DType, Device, Tensor};
|
|
use std::collections::HashMap;
|
|
use std::sync::Arc;
|
|
|
|
#[tokio::test]
|
|
async fn test_agent_collaboration() {
|
|
let runtime = Arc::new(Runtime::new().expect("Failed to create runtime"));
|
|
|
|
// Create all agents
|
|
let data_agent = DataEngineeringAgent::new(runtime.clone()).unwrap();
|
|
let parallel_agent = ParallelPlannerAgent::new(runtime.clone()).unwrap();
|
|
let quant_agent = QuantGuardianAgent::new(runtime.clone()).unwrap();
|
|
let kernel_agent = KernelSynthesizerAgent::new(runtime.clone()).unwrap();
|
|
|
|
// Create a tensor for optimization
|
|
let device = Device::cuda(0).unwrap_or(Device::Cpu);
|
|
let tensor = Tensor::randn(&[1024, 1024], &device).unwrap();
|
|
|
|
// Get proposals from each agent
|
|
let data_proposals = data_agent.generate_layout_proposals(&tensor).await.unwrap();
|
|
let quant_proposals = quant_agent
|
|
.generate_quantization_proposals(&tensor, 0.95)
|
|
.await
|
|
.unwrap();
|
|
|
|
assert!(!data_proposals.is_empty());
|
|
assert!(!quant_proposals.is_empty());
|
|
|
|
// Verify that proposals can be combined
|
|
let mut all_proposals = data_proposals;
|
|
all_proposals.extend(quant_proposals);
|
|
|
|
assert!(
|
|
all_proposals.len() >= 2,
|
|
"Should have proposals from multiple agents"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_end_to_end_optimization_pipeline() {
|
|
let runtime = Arc::new(Runtime::new().expect("Failed to create runtime"));
|
|
|
|
// Set up the pipeline components
|
|
let data_agent = DataEngineeringAgent::new(runtime.clone()).unwrap();
|
|
let validator = ProposalValidator::new(runtime.clone()).unwrap();
|
|
let mut rollback_manager = RollbackManager::new(runtime.clone()).unwrap();
|
|
|
|
// Create initial state
|
|
let mut initial_state = HashMap::new();
|
|
initial_state.insert("performance_baseline".to_string(), vec![1.0]);
|
|
initial_state.insert("accuracy".to_string(), vec![0.95]);
|
|
|
|
// Create checkpoint
|
|
let checkpoint = rollback_manager
|
|
.create_checkpoint(
|
|
rtx_auto::rollback::CheckpointType::BeforeOptimization,
|
|
initial_state.clone(),
|
|
"Baseline before optimization".to_string(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
rollback_manager.save_checkpoint(&checkpoint).await.unwrap();
|
|
|
|
// Generate proposals
|
|
let device = Device::cuda(0).unwrap_or(Device::Cpu);
|
|
let tensor = Tensor::randn(&[512, 512], &device).unwrap();
|
|
let proposals = data_agent.generate_layout_proposals(&tensor).await.unwrap();
|
|
|
|
// Validate and rank proposals
|
|
let ranked_proposals = validator.rank_proposals(&proposals).await.unwrap();
|
|
assert!(!ranked_proposals.is_empty(), "Should have ranked proposals");
|
|
|
|
// Select best proposal
|
|
let (best_proposal, _score) = &ranked_proposals[0];
|
|
assert_eq!(best_proposal.proposal_type(), ProposalType::DataLayout);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_multi_agent_consensus() {
|
|
let runtime = Arc::new(Runtime::new().expect("Failed to create runtime"));
|
|
|
|
let data_agent = DataEngineeringAgent::new(runtime.clone()).unwrap();
|
|
let kernel_agent = KernelSynthesizerAgent::new(runtime.clone()).unwrap();
|
|
let validator = ProposalValidator::new(runtime.clone()).unwrap();
|
|
|
|
// Get proposals from multiple agents for the same workload
|
|
let device = Device::cuda(0).unwrap_or(Device::Cpu);
|
|
let tensor = Tensor::randn(&[256, 256], &device).unwrap();
|
|
let data_proposals = data_agent.generate_layout_proposals(&tensor).await.unwrap();
|
|
|
|
// Also get quantization proposals to have multiple proposal types
|
|
let quant_agent = QuantGuardianAgent::new(runtime.clone()).unwrap();
|
|
let quant_proposals = quant_agent
|
|
.generate_quantization_proposals(&tensor, 0.95)
|
|
.await
|
|
.unwrap();
|
|
|
|
// Combine and rank all proposals
|
|
let mut all_proposals = data_proposals;
|
|
all_proposals.extend(quant_proposals);
|
|
|
|
let ranked = validator.rank_proposals(&all_proposals).await.unwrap();
|
|
assert!(!ranked.is_empty(), "Should rank combined proposals");
|
|
|
|
// Verify diversity of proposal types
|
|
let proposal_types: std::collections::HashSet<_> = ranked
|
|
.iter()
|
|
.map(|(proposal, _)| proposal.proposal_type())
|
|
.collect();
|
|
assert!(
|
|
proposal_types.len() >= 2,
|
|
"Should have proposals from different agents"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_adaptive_optimization_cycle() {
|
|
let runtime = Arc::new(Runtime::new().expect("Failed to create runtime"));
|
|
|
|
let quant_agent = QuantGuardianAgent::new(runtime.clone()).unwrap();
|
|
let mut rollback_manager = RollbackManager::new(runtime.clone()).unwrap();
|
|
|
|
// Set up adaptive thresholds
|
|
rollback_manager
|
|
.set_rollback_threshold("accuracy", 0.90)
|
|
.await
|
|
.unwrap();
|
|
|
|
// Initial state with good performance
|
|
let mut good_state = HashMap::new();
|
|
good_state.insert("accuracy".to_string(), vec![0.95]);
|
|
good_state.insert("latency".to_string(), vec![100.0]);
|
|
|
|
let checkpoint = rollback_manager
|
|
.create_checkpoint(
|
|
rtx_auto::rollback::CheckpointType::Automatic,
|
|
good_state,
|
|
"Good performance state".to_string(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
rollback_manager.save_checkpoint(&checkpoint).await.unwrap();
|
|
|
|
// Generate aggressive quantization proposal
|
|
let device = Device::cuda(0).unwrap_or(Device::Cpu);
|
|
let tensor = Tensor::randn(&[128, 128], &device).unwrap();
|
|
let proposals = quant_agent
|
|
.generate_quantization_proposals(&tensor, 0.85)
|
|
.await
|
|
.unwrap();
|
|
|
|
assert!(
|
|
!proposals.is_empty(),
|
|
"Should generate quantization proposals"
|
|
);
|
|
|
|
// Simulate applying proposal with mixed results
|
|
let mut result_state = HashMap::new();
|
|
result_state.insert("accuracy".to_string(), vec![0.88]); // Below threshold
|
|
result_state.insert("latency".to_string(), vec![50.0]); // Better latency
|
|
|
|
let should_rollback = rollback_manager
|
|
.should_trigger_automatic_rollback(&result_state)
|
|
.await
|
|
.unwrap();
|
|
assert!(
|
|
should_rollback,
|
|
"Should trigger rollback due to accuracy degradation"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_proposal_conflict_resolution() {
|
|
let runtime = Arc::new(Runtime::new().expect("Failed to create runtime"));
|
|
let validator = ProposalValidator::new(runtime.clone()).unwrap();
|
|
|
|
// Create conflicting proposals
|
|
let proposal1 = Proposal::new(
|
|
ProposalType::DataLayout,
|
|
"Optimize for sequential access".to_string(),
|
|
1.5,
|
|
);
|
|
|
|
let proposal2 = Proposal::new(
|
|
ProposalType::DataLayout,
|
|
"Optimize for random access".to_string(),
|
|
1.3,
|
|
);
|
|
|
|
let proposals = vec![proposal1, proposal2];
|
|
|
|
// Detect conflicts
|
|
let conflicts = validator.detect_conflicts(&proposals).await;
|
|
assert!(conflicts.is_ok());
|
|
|
|
let conflicts = conflicts.unwrap();
|
|
assert!(!conflicts.is_empty(), "Should detect conflicting proposals");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_performance_regression_detection() {
|
|
let runtime = Arc::new(Runtime::new().expect("Failed to create runtime"));
|
|
let quant_agent = QuantGuardianAgent::new(runtime.clone()).unwrap();
|
|
|
|
// Create reference tensor with known values
|
|
let device = Device::cuda(0).unwrap_or(Device::Cpu);
|
|
let reference = Tensor::ones(&[64, 64], &device).unwrap();
|
|
|
|
// Create degraded version
|
|
let degraded_data = vec![0.5f32; 64 * 64];
|
|
let degraded = Tensor::from_vec(degraded_data, &[64, 64], &device).unwrap();
|
|
|
|
// Monitor accuracy should detect significant degradation
|
|
let metrics = quant_agent
|
|
.monitor_accuracy(&reference, °raded)
|
|
.await
|
|
.unwrap();
|
|
assert!(metrics.mse > 0.0, "Should detect mean squared error");
|
|
|
|
// Use threshold 0.8 - degradation detected when mse > (1.0 - 0.8) = 0.2
|
|
// MSE between [1,1,1,...] and [0.5,0.5,0.5,...] is 0.25, which is > 0.2
|
|
let is_degraded = quant_agent
|
|
.detect_accuracy_degradation(&reference, °raded, 0.8)
|
|
.await
|
|
.unwrap();
|
|
assert!(is_degraded, "Should detect accuracy degradation");
|
|
}
|