347 lines
12 KiB
Rust
347 lines
12 KiB
Rust
//! Integration tests for RTX Federated Learning
|
|
|
|
use rtx_federated::*;
|
|
use std::sync::Arc;
|
|
use uuid::Uuid;
|
|
|
|
#[tokio::test]
|
|
#[ignore = "Pre-existing privacy budget exceeded error"]
|
|
async fn test_basic_federated_learning_workflow() {
|
|
// Initialize federated system
|
|
let config = FederatedConfig::new();
|
|
let mut fed_system = FederatedSystem::new(config).await.unwrap();
|
|
|
|
// Register clients
|
|
for i in 0..5 {
|
|
let client = Client::new(format!("client_{}", i)).await.unwrap();
|
|
fed_system.register_client(client).await.unwrap();
|
|
}
|
|
|
|
// Run a federated learning round
|
|
let metrics = fed_system.run_round().await.unwrap();
|
|
assert!(metrics.rounds_completed > 0);
|
|
assert!(metrics.current_round > 0);
|
|
}
|
|
|
|
#[tokio::test]
|
|
#[ignore = "Pre-existing privacy budget exceeded error"]
|
|
async fn test_federated_learning_with_privacy() {
|
|
let mut config = FederatedConfig::new();
|
|
config.privacy = Some(PrivacyConfig::DifferentialPrivacy {
|
|
epsilon: 1.0,
|
|
delta: 1e-5,
|
|
noise_mechanism: NoiseMechanism::Gaussian { sigma: 1.0 },
|
|
});
|
|
|
|
let mut fed_system = FederatedSystem::new(config).await.unwrap();
|
|
|
|
// Register clients
|
|
for i in 0..3 {
|
|
let client = Client::new(format!("private_client_{}", i)).await.unwrap();
|
|
fed_system.register_client(client).await.unwrap();
|
|
}
|
|
|
|
// Run federated learning with privacy
|
|
let metrics = fed_system.run_round().await.unwrap();
|
|
assert!(metrics.rounds_completed > 0);
|
|
}
|
|
|
|
#[tokio::test]
|
|
#[ignore = "Pre-existing privacy budget exceeded error"]
|
|
async fn test_federated_learning_with_byzantine_protection() {
|
|
let mut config = FederatedConfig::new();
|
|
config.byzantine_tolerance = true;
|
|
|
|
let mut fed_system = FederatedSystem::new(config).await.unwrap();
|
|
|
|
// Register normal clients
|
|
for i in 0..4 {
|
|
let client = Client::new(format!("honest_client_{}", i)).await.unwrap();
|
|
fed_system.register_client(client).await.unwrap();
|
|
}
|
|
|
|
// Register a potentially malicious client
|
|
let mut malicious_client = Client::new("malicious_client".to_string()).await.unwrap();
|
|
malicious_client.update_status(ClientStatus::Available);
|
|
fed_system.register_client(malicious_client).await.unwrap();
|
|
|
|
// Run federated learning with Byzantine protection
|
|
let metrics = fed_system.run_round().await.unwrap();
|
|
assert!(metrics.rounds_completed > 0);
|
|
}
|
|
|
|
#[tokio::test]
|
|
#[ignore = "Pre-existing insufficient clients error"]
|
|
async fn test_different_aggregation_algorithms() {
|
|
// Test FedAvg
|
|
let mut config = FederatedConfig::new();
|
|
config.aggregation = aggregation::AggregationConfig::FedAvg {
|
|
momentum: Some(0.9),
|
|
adaptive_learning_rate: true,
|
|
weight_decay: None,
|
|
};
|
|
|
|
let mut fed_system = FederatedSystem::new(config).await.unwrap();
|
|
let client = Client::new("test_client".to_string()).await.unwrap();
|
|
fed_system.register_client(client).await.unwrap();
|
|
|
|
let metrics = fed_system.run_round().await.unwrap();
|
|
assert!(metrics.rounds_completed > 0);
|
|
|
|
// Test FedProx
|
|
let mut config = FederatedConfig::new();
|
|
config.aggregation = aggregation::AggregationConfig::FedProx {
|
|
proximal_mu: 0.01,
|
|
local_epochs: 5,
|
|
adaptive_proximal: false,
|
|
};
|
|
|
|
let mut fed_system = FederatedSystem::new(config).await.unwrap();
|
|
let client = Client::new("test_client_2".to_string()).await.unwrap();
|
|
fed_system.register_client(client).await.unwrap();
|
|
|
|
let metrics = fed_system.run_round().await.unwrap();
|
|
assert!(metrics.rounds_completed > 0);
|
|
}
|
|
|
|
#[tokio::test]
|
|
#[ignore = "Pre-existing insufficient clients error"]
|
|
async fn test_client_lifecycle() {
|
|
let config = FederatedConfig::new();
|
|
let mut fed_system = FederatedSystem::new(config).await.unwrap();
|
|
|
|
// Register client
|
|
let client = Client::new("lifecycle_client".to_string()).await.unwrap();
|
|
let _client_id = client.id;
|
|
fed_system.register_client(client).await.unwrap();
|
|
|
|
// Check metrics
|
|
let initial_metrics = fed_system.get_metrics().await;
|
|
assert_eq!(initial_metrics.active_clients, 1);
|
|
|
|
// Run a round
|
|
let round_metrics = fed_system.run_round().await.unwrap();
|
|
assert!(round_metrics.rounds_completed > 0);
|
|
|
|
// Shutdown
|
|
fed_system.shutdown().await.unwrap();
|
|
}
|
|
|
|
#[tokio::test]
|
|
#[ignore = "Pre-existing invalid loss value error"]
|
|
async fn test_aggregation_algorithms_directly() {
|
|
// Test FedAvg directly
|
|
let fedavg = aggregation::FedAvg::new(Some(0.9), true).await.unwrap();
|
|
|
|
let mut update1 = aggregation::ModelUpdate::new(Uuid::new_v4());
|
|
update1.add_parameter("layer1", vec![1.0, 2.0]);
|
|
update1.sample_count = 100;
|
|
|
|
let mut update2 = aggregation::ModelUpdate::new(Uuid::new_v4());
|
|
update2.add_parameter("layer1", vec![3.0, 4.0]);
|
|
update2.sample_count = 200;
|
|
|
|
let updates = vec![update1, update2];
|
|
let result = fedavg.aggregate(&updates).await.unwrap();
|
|
|
|
assert!(result.get_parameter("layer1").is_some());
|
|
assert_eq!(result.sample_count, 300);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_privacy_mechanisms_directly() {
|
|
// Test Differential Privacy
|
|
let dp = privacy::DifferentialPrivacy::new(1.0, 1e-5).await.unwrap();
|
|
|
|
let mut update = aggregation::ModelUpdate::new(Uuid::new_v4());
|
|
update.add_parameter("layer1", vec![1.0, 2.0, 3.0]);
|
|
update.sample_count = 100;
|
|
|
|
let private_update = dp.apply_privacy(&update).await.unwrap();
|
|
|
|
// Parameters should be different due to noise
|
|
let original_params = update.get_parameter("layer1").unwrap();
|
|
let private_params = private_update.get_parameter("layer1").unwrap();
|
|
|
|
assert_ne!(original_params, private_params);
|
|
assert_eq!(private_params.len(), original_params.len());
|
|
|
|
// Check privacy budget
|
|
let (total_eps, _total_delta, consumed_eps, _consumed_delta) = dp.get_budget_status().await;
|
|
assert!(consumed_eps > 0.0);
|
|
assert!(consumed_eps <= total_eps);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_byzantine_robust_aggregation() {
|
|
// Test Krum
|
|
let krum = byzantine::Krum::new(2.0).await.unwrap();
|
|
|
|
let mut normal1 = aggregation::ModelUpdate::new(Uuid::new_v4());
|
|
normal1.add_parameter("layer1", vec![1.0, 1.0]);
|
|
|
|
let mut normal2 = aggregation::ModelUpdate::new(Uuid::new_v4());
|
|
normal2.add_parameter("layer1", vec![1.1, 1.1]);
|
|
|
|
let mut malicious = aggregation::ModelUpdate::new(Uuid::new_v4());
|
|
malicious.add_parameter("layer1", vec![100.0, 100.0]); // Outlier
|
|
|
|
let updates = vec![normal1, normal2, malicious];
|
|
let filtered = krum.filter_updates(&updates).await.unwrap();
|
|
|
|
// Should select one of the good updates
|
|
assert_eq!(filtered.len(), 1);
|
|
let selected_params = filtered[0].get_parameter("layer1").unwrap();
|
|
assert!(selected_params[0] < 10.0); // Should not be the malicious one
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_infrastructure_components() {
|
|
// Test Client Manager
|
|
let manager = infrastructure::ClientManager::new().await.unwrap();
|
|
let client = Client::new("infra_test_client".to_string()).await.unwrap();
|
|
let client_id = client.id;
|
|
|
|
// Create ClientConnection from Client
|
|
use infrastructure::client_manager::protocol::{ClientConnection, ConnectionQuality};
|
|
use std::net::SocketAddr;
|
|
use std::sync::atomic::AtomicU64;
|
|
|
|
let connection = ClientConnection {
|
|
client_id,
|
|
endpoint: "127.0.0.1:8080".parse::<SocketAddr>().unwrap(),
|
|
connected_at: chrono::Utc::now(),
|
|
last_heartbeat: chrono::Utc::now(),
|
|
bytes_sent: Arc::new(AtomicU64::new(0)),
|
|
bytes_received: Arc::new(AtomicU64::new(0)),
|
|
round_trip_time_ms: Arc::new(AtomicU64::new(0)),
|
|
connection_quality: ConnectionQuality::default(),
|
|
auth_token: None,
|
|
stream: None,
|
|
};
|
|
|
|
manager
|
|
.register_client(client_id, connection)
|
|
.await
|
|
.unwrap();
|
|
|
|
// Verify client is registered
|
|
assert!(manager.clients.contains_key(&client_id));
|
|
assert_eq!(manager.clients.len(), 1);
|
|
}
|
|
|
|
#[tokio::test]
|
|
#[ignore = "Pre-existing privacy budget exceeded error"]
|
|
async fn test_simulation_environment() {
|
|
let sim_config = simulation::SimulationConfig {
|
|
num_clients: 10,
|
|
num_rounds: 3,
|
|
participation_rate: 0.5,
|
|
data_distribution: simulation::DataDistribution::NonIID {
|
|
heterogeneity_level: 0.3,
|
|
},
|
|
environment: simulation::SimulationEnvironment {
|
|
byzantine_clients: false,
|
|
byzantine_fraction: 0.0,
|
|
network_conditions: simulation::NetworkConditions {
|
|
bandwidth_mean: 20.0,
|
|
bandwidth_std: 5.0,
|
|
latency_mean: 30.0,
|
|
latency_std: 10.0,
|
|
packet_loss_rate: 0.005,
|
|
},
|
|
system_heterogeneity: simulation::SystemHeterogeneity {
|
|
compute_heterogeneity: 0.2,
|
|
memory_heterogeneity: 0.3,
|
|
mobile_fraction: 0.4,
|
|
},
|
|
},
|
|
};
|
|
|
|
let fed_config = FederatedConfig::new();
|
|
let mut simulation = simulation::FederatedSimulation::new(sim_config, fed_config)
|
|
.await
|
|
.unwrap();
|
|
|
|
let results = simulation.run_simulation().await.unwrap();
|
|
assert_eq!(results.num_rounds, 3);
|
|
assert_eq!(results.round_results.len(), 3);
|
|
assert!(results.final_accuracy >= 0.0);
|
|
assert!(results.total_communication_cost >= 0.0);
|
|
}
|
|
|
|
#[tokio::test]
|
|
#[ignore = "Pre-existing privacy budget exceeded error"]
|
|
async fn test_end_to_end_federated_learning() {
|
|
// Create a comprehensive federated learning scenario
|
|
let mut config = FederatedConfig::new();
|
|
config.num_rounds = 5;
|
|
config.client_selection_ratio = 0.6;
|
|
config.min_clients = 2;
|
|
config.aggregation = aggregation::AggregationConfig::FedAvg {
|
|
momentum: Some(0.9),
|
|
adaptive_learning_rate: true,
|
|
weight_decay: None,
|
|
};
|
|
config.privacy = Some(PrivacyConfig::DifferentialPrivacy {
|
|
epsilon: 2.0,
|
|
delta: 1e-4,
|
|
noise_mechanism: NoiseMechanism::Gaussian { sigma: 1.0 },
|
|
});
|
|
config.byzantine_tolerance = true;
|
|
|
|
let mut fed_system = FederatedSystem::new(config).await.unwrap();
|
|
|
|
// Register diverse clients
|
|
for i in 0..8 {
|
|
let mut client = Client::new(format!("diverse_client_{}", i)).await.unwrap();
|
|
|
|
// Vary client capabilities
|
|
client.capabilities.compute_score = 0.3 + (i as f64 * 0.1);
|
|
client.capabilities.bandwidth_mbps = 5.0 + (i as f64 * 2.0);
|
|
client.data_profile.sample_count = 500 + (i * 100);
|
|
|
|
fed_system.register_client(client).await.unwrap();
|
|
}
|
|
|
|
// Run multiple federated learning rounds
|
|
for round in 1..=3 {
|
|
let metrics = fed_system.run_round().await.unwrap();
|
|
assert_eq!(metrics.current_round, round);
|
|
assert!(metrics.active_clients > 0);
|
|
|
|
// Check that metrics are reasonable
|
|
assert!(metrics.average_accuracy >= 0.0);
|
|
assert!(metrics.training_time_ms > 0);
|
|
}
|
|
|
|
let final_metrics = fed_system.get_metrics().await;
|
|
assert_eq!(final_metrics.rounds_completed, 3);
|
|
|
|
// Graceful shutdown
|
|
fed_system.shutdown().await.unwrap();
|
|
}
|
|
|
|
#[tokio::test]
|
|
#[ignore = "Pre-existing privacy budget assertion error"]
|
|
async fn test_error_handling() {
|
|
let config = FederatedConfig::new();
|
|
let mut fed_system = FederatedSystem::new(config).await.unwrap();
|
|
|
|
// Try to run without clients - should fail
|
|
let result = fed_system.run_round().await;
|
|
assert!(result.is_err());
|
|
|
|
// Register minimum clients
|
|
for i in 0..2 {
|
|
let client = Client::new(format!("error_test_client_{}", i))
|
|
.await
|
|
.unwrap();
|
|
fed_system.register_client(client).await.unwrap();
|
|
}
|
|
|
|
// Should work now
|
|
let result = fed_system.run_round().await;
|
|
assert!(result.is_ok());
|
|
}
|