103 lines
3.6 KiB
Rust
103 lines
3.6 KiB
Rust
use rtx_auto::{
|
|
agents::DataEngineeringAgent,
|
|
error::AutoError,
|
|
proposal::{Proposal, ProposalStatus, ProposalType},
|
|
};
|
|
use rtx_runtime::Runtime;
|
|
use rtx_tensor::{DType, Device, Tensor};
|
|
use std::sync::Arc;
|
|
|
|
#[tokio::test]
|
|
async fn test_data_engineering_agent_creation() {
|
|
let runtime = Arc::new(Runtime::new().expect("Failed to create runtime"));
|
|
let agent = DataEngineeringAgent::new(runtime.clone());
|
|
assert!(agent.is_ok(), "Failed to create DataEngineeringAgent");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_analyze_data_access_patterns() {
|
|
let runtime = Arc::new(Runtime::new().expect("Failed to create runtime"));
|
|
let agent = DataEngineeringAgent::new(runtime.clone()).unwrap();
|
|
|
|
let device = Device::cuda(0).unwrap_or(Device::Cpu);
|
|
let tensors = vec![
|
|
Tensor::zeros_typed(&[1000, 512], DType::F32, &device).unwrap(),
|
|
Tensor::zeros_typed(&[512, 256], DType::F32, &device).unwrap(),
|
|
];
|
|
|
|
let access_patterns = agent.analyze_access_patterns(&tensors).await;
|
|
assert!(access_patterns.is_ok());
|
|
|
|
let patterns = access_patterns.unwrap();
|
|
assert!(!patterns.is_empty(), "Should detect access patterns");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_generate_data_layout_proposals() {
|
|
let runtime = Arc::new(Runtime::new().expect("Failed to create runtime"));
|
|
let agent = DataEngineeringAgent::new(runtime.clone()).unwrap();
|
|
|
|
let device = Device::cuda(0).unwrap_or(Device::Cpu);
|
|
let tensor = Tensor::zeros_typed(&[1024, 768], DType::F32, &device).unwrap();
|
|
let proposals = agent.generate_layout_proposals(&tensor).await;
|
|
|
|
assert!(proposals.is_ok());
|
|
let proposals = proposals.unwrap();
|
|
assert!(
|
|
!proposals.is_empty(),
|
|
"Should generate at least one layout proposal"
|
|
);
|
|
|
|
for proposal in &proposals {
|
|
assert_eq!(proposal.proposal_type(), ProposalType::DataLayout);
|
|
assert_eq!(proposal.status(), ProposalStatus::Pending);
|
|
assert!(proposal.expected_performance_gain() > 0.0);
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_memory_coalescing_proposals() {
|
|
let runtime = Arc::new(Runtime::new().expect("Failed to create runtime"));
|
|
let agent = DataEngineeringAgent::new(runtime.clone()).unwrap();
|
|
|
|
let device = Device::cuda(0).unwrap_or(Device::Cpu);
|
|
let tensors = vec![
|
|
Tensor::zeros_typed(&[100, 100], DType::F32, &device).unwrap(),
|
|
Tensor::zeros_typed(&[100, 100], DType::F32, &device).unwrap(),
|
|
Tensor::zeros_typed(&[100, 100], DType::F32, &device).unwrap(),
|
|
];
|
|
|
|
let proposals = agent.generate_coalescing_proposals(&tensors).await;
|
|
assert!(proposals.is_ok());
|
|
|
|
let proposals = proposals.unwrap();
|
|
assert!(
|
|
!proposals.is_empty(),
|
|
"Should generate coalescing proposals"
|
|
);
|
|
|
|
for proposal in &proposals {
|
|
assert_eq!(proposal.proposal_type(), ProposalType::MemoryCoalescing);
|
|
assert!(proposal.description().to_lowercase().contains("coalesce"));
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_cache_optimization_proposals() {
|
|
let runtime = Arc::new(Runtime::new().expect("Failed to create runtime"));
|
|
let agent = DataEngineeringAgent::new(runtime.clone()).unwrap();
|
|
|
|
let device = Device::cuda(0).unwrap_or(Device::Cpu);
|
|
let tensor = Tensor::zeros_typed(&[2048, 2048], DType::F32, &device).unwrap();
|
|
let proposals = agent.generate_cache_optimization_proposals(&tensor).await;
|
|
|
|
assert!(proposals.is_ok());
|
|
let proposals = proposals.unwrap();
|
|
|
|
for proposal in &proposals {
|
|
assert_eq!(proposal.proposal_type(), ProposalType::CacheOptimization);
|
|
assert!(proposal.expected_performance_gain() > 0.0);
|
|
assert!(proposal.description().len() > 10);
|
|
}
|
|
}
|