374 lines
12 KiB
Rust
374 lines
12 KiB
Rust
//! RTX Federated Learning Demo
|
|
//!
|
|
//! This example demonstrates the key capabilities of RTX Federated Learning:
|
|
//! - Advanced aggregation algorithms (FedAvg, FedProx)
|
|
//! - Privacy-preserving mechanisms (Differential Privacy)
|
|
//! - Byzantine fault tolerance (Krum)
|
|
//! - Production infrastructure components
|
|
//! - Simulation environment
|
|
|
|
use rtx_federated::*;
|
|
use std::sync::Arc;
|
|
use tracing::{info, warn};
|
|
|
|
#[tokio::main]
|
|
async fn main() -> std::result::Result<(), Box<dyn std::error::Error>> {
|
|
// Initialize logging
|
|
tracing_subscriber::fmt::init();
|
|
|
|
info!("🚀 RTX Federated Learning Demo Starting");
|
|
|
|
// Demo 1: Basic Federated Learning
|
|
demo_basic_federated_learning().await?;
|
|
|
|
// Demo 2: Privacy-Preserving Federated Learning
|
|
demo_privacy_preserving_learning().await?;
|
|
|
|
// Demo 3: Byzantine-Robust Federated Learning
|
|
demo_byzantine_robust_learning().await?;
|
|
|
|
// Demo 4: Advanced Aggregation Algorithms
|
|
demo_advanced_aggregation().await?;
|
|
|
|
// Demo 5: Simulation Environment
|
|
demo_simulation_environment().await?;
|
|
|
|
info!("✅ RTX Federated Learning Demo Complete");
|
|
Ok(())
|
|
}
|
|
|
|
async fn demo_basic_federated_learning() -> std::result::Result<(), Box<dyn std::error::Error>> {
|
|
info!("📚 Demo 1: Basic Federated Learning");
|
|
|
|
// Create federated learning configuration
|
|
let config = FederatedConfig::new().with_aggregation(aggregation::AggregationConfig::FedAvg {
|
|
momentum: Some(0.9),
|
|
adaptive_learning_rate: true,
|
|
weight_decay: None,
|
|
});
|
|
|
|
// Initialize federated system
|
|
let mut fed_system = FederatedSystem::new(config).await?;
|
|
|
|
// Register clients with diverse characteristics
|
|
for i in 0..10 {
|
|
let mut client = Client::new(format!("basic_client_{}", i)).await?;
|
|
|
|
// Vary client capabilities
|
|
client.capabilities.compute_score = 0.3 + (i as f64 * 0.07);
|
|
client.capabilities.bandwidth_mbps = 10.0 + (i as f64 * 2.0);
|
|
client.data_profile.sample_count = 800 + (i * 50);
|
|
|
|
fed_system.register_client(client).await?;
|
|
}
|
|
|
|
// Run federated learning rounds
|
|
for round in 1..=5 {
|
|
info!("🔄 Running federated learning round {}", round);
|
|
let metrics = fed_system.run_round().await?;
|
|
|
|
info!("📊 Round {} Results:", round);
|
|
info!(" - Participating clients: {}", metrics.active_clients);
|
|
info!(" - Average accuracy: {:.4}", metrics.average_accuracy);
|
|
info!(" - Training time: {}ms", metrics.training_time_ms);
|
|
}
|
|
|
|
let final_metrics = fed_system.get_metrics().await;
|
|
info!("🎯 Final Results:");
|
|
info!(
|
|
" - Total rounds completed: {}",
|
|
final_metrics.rounds_completed
|
|
);
|
|
info!(
|
|
" - Final average accuracy: {:.4}",
|
|
final_metrics.average_accuracy
|
|
);
|
|
|
|
fed_system.shutdown().await?;
|
|
info!("✅ Basic federated learning demo completed\n");
|
|
Ok(())
|
|
}
|
|
|
|
async fn demo_privacy_preserving_learning() -> std::result::Result<(), Box<dyn std::error::Error>> {
|
|
info!("🔒 Demo 2: Privacy-Preserving Federated Learning");
|
|
|
|
// Configure with differential privacy
|
|
let config = FederatedConfig::new().with_aggregation(aggregation::AggregationConfig::FedAvg {
|
|
momentum: Some(0.9),
|
|
adaptive_learning_rate: true,
|
|
weight_decay: None,
|
|
});
|
|
|
|
let mut fed_system = FederatedSystem::new(config).await?;
|
|
|
|
// Register privacy-conscious clients
|
|
for i in 0..6 {
|
|
let mut client = Client::new(format!("private_client_{}", i)).await?;
|
|
client.data_profile.privacy_level = if i % 2 == 0 {
|
|
PrivacyLevel::Confidential
|
|
} else {
|
|
PrivacyLevel::Restricted
|
|
};
|
|
|
|
fed_system.register_client(client).await?;
|
|
}
|
|
|
|
// Note: Differential privacy mechanism is configured in the system
|
|
info!("🔒 Privacy mechanisms are active in the federated system");
|
|
|
|
// Run privacy-preserving federated learning
|
|
for round in 1..=3 {
|
|
let metrics = fed_system.run_round().await?;
|
|
info!(
|
|
"🔒 Private Round {}: {} clients, accuracy: {:.4}",
|
|
round, metrics.active_clients, metrics.average_accuracy
|
|
);
|
|
}
|
|
|
|
fed_system.shutdown().await?;
|
|
info!("✅ Privacy-preserving learning demo completed\n");
|
|
Ok(())
|
|
}
|
|
|
|
async fn demo_byzantine_robust_learning() -> std::result::Result<(), Box<dyn std::error::Error>> {
|
|
info!("🛡️ Demo 3: Byzantine-Robust Federated Learning");
|
|
|
|
// Configure with Byzantine tolerance
|
|
let mut config = FederatedConfig::new();
|
|
config.byzantine_tolerance = true;
|
|
|
|
let mut fed_system = FederatedSystem::new(config).await?;
|
|
|
|
// Register honest clients
|
|
for i in 0..8 {
|
|
let mut client = Client::new(format!("honest_client_{}", i)).await?;
|
|
client.status = ClientStatus::Available;
|
|
fed_system.register_client(client).await?;
|
|
}
|
|
|
|
// Register potentially malicious clients
|
|
for i in 0..2 {
|
|
let mut malicious_client = Client::new(format!("suspicious_client_{}", i)).await?;
|
|
malicious_client.status = ClientStatus::Available;
|
|
fed_system.register_client(malicious_client).await?;
|
|
}
|
|
|
|
// Demonstrate Krum algorithm directly
|
|
let krum = byzantine::Krum::new(0.2).await?;
|
|
|
|
// Create test updates with one outlier
|
|
let mut updates = Vec::new();
|
|
for i in 0..5 {
|
|
let mut update = aggregation::ModelUpdate::new(uuid::Uuid::new_v4());
|
|
if i < 4 {
|
|
update.add_parameter("layer1", vec![1.0 + i as f64 * 0.1, 2.0 + i as f64 * 0.1]);
|
|
} else {
|
|
// Malicious update
|
|
update.add_parameter("layer1", vec![100.0, 100.0]);
|
|
}
|
|
update.sample_count = 100;
|
|
updates.push(update);
|
|
}
|
|
|
|
let filtered_updates = krum.filter_updates(&updates).await?;
|
|
info!(
|
|
"🛡️ Krum filtered {} updates to {} honest updates",
|
|
updates.len(),
|
|
filtered_updates.len()
|
|
);
|
|
|
|
// Run Byzantine-robust federated learning
|
|
for round in 1..=3 {
|
|
let metrics = fed_system.run_round().await?;
|
|
info!(
|
|
"🛡️ Robust Round {}: {} clients, {} Byzantine attacks detected",
|
|
round, metrics.active_clients, metrics.byzantine_attacks_detected
|
|
);
|
|
}
|
|
|
|
fed_system.shutdown().await?;
|
|
info!("✅ Byzantine-robust learning demo completed\n");
|
|
Ok(())
|
|
}
|
|
|
|
async fn demo_advanced_aggregation() -> std::result::Result<(), Box<dyn std::error::Error>> {
|
|
info!("🧠 Demo 4: Advanced Aggregation Algorithms");
|
|
|
|
// Test different aggregation algorithms
|
|
let algorithms = vec![
|
|
(
|
|
"FedAvg",
|
|
aggregation::AggregationConfig::FedAvg {
|
|
momentum: Some(0.9),
|
|
adaptive_learning_rate: true,
|
|
weight_decay: None,
|
|
},
|
|
),
|
|
(
|
|
"FedProx",
|
|
aggregation::AggregationConfig::FedProx {
|
|
proximal_mu: 0.01,
|
|
local_epochs: 5,
|
|
adaptive_proximal: false,
|
|
},
|
|
),
|
|
(
|
|
"SCAFFOLD",
|
|
aggregation::AggregationConfig::Scaffold {
|
|
learning_rate: 0.01,
|
|
local_steps: 100,
|
|
variance_reduction: true,
|
|
},
|
|
),
|
|
(
|
|
"FedNova",
|
|
aggregation::AggregationConfig::FedNova {
|
|
tau_effective: 10.0,
|
|
momentum_factor: 0.9,
|
|
normalize_weights: true,
|
|
},
|
|
),
|
|
];
|
|
|
|
for (name, algorithm_config) in algorithms {
|
|
info!("🔬 Testing {} algorithm", name);
|
|
|
|
let config = FederatedConfig::new().with_aggregation(algorithm_config);
|
|
|
|
let mut fed_system = FederatedSystem::new(config).await?;
|
|
|
|
// Register clients
|
|
for i in 0..5 {
|
|
let client = Client::new(format!("{}_client_{}", name.to_lowercase(), i)).await?;
|
|
fed_system.register_client(client).await?;
|
|
}
|
|
|
|
// Run one round
|
|
let metrics = fed_system.run_round().await?;
|
|
info!(
|
|
" - {} result: {} clients, accuracy: {:.4}, time: {}ms",
|
|
name, metrics.active_clients, metrics.average_accuracy, metrics.training_time_ms
|
|
);
|
|
|
|
fed_system.shutdown().await?;
|
|
}
|
|
|
|
info!("✅ Advanced aggregation algorithms demo completed\n");
|
|
Ok(())
|
|
}
|
|
|
|
async fn demo_simulation_environment() -> std::result::Result<(), Box<dyn std::error::Error>> {
|
|
info!("🎮 Demo 5: Simulation Environment");
|
|
|
|
// Configure simulation
|
|
let sim_config = simulation::SimulationConfig {
|
|
num_clients: 50,
|
|
num_rounds: 10,
|
|
participation_rate: 0.3,
|
|
data_distribution: simulation::DataDistribution::NonIID {
|
|
heterogeneity_level: 0.4,
|
|
},
|
|
environment: simulation::SimulationEnvironment {
|
|
byzantine_clients: true,
|
|
byzantine_fraction: 0.1,
|
|
network_conditions: simulation::NetworkConditions {
|
|
bandwidth_mean: 15.0,
|
|
bandwidth_std: 8.0,
|
|
latency_mean: 45.0,
|
|
latency_std: 15.0,
|
|
packet_loss_rate: 0.02,
|
|
},
|
|
system_heterogeneity: simulation::SystemHeterogeneity {
|
|
compute_heterogeneity: 0.4,
|
|
memory_heterogeneity: 0.3,
|
|
mobile_fraction: 0.6,
|
|
},
|
|
},
|
|
};
|
|
|
|
// Configure federated learning
|
|
let mut fed_config = FederatedConfig::new();
|
|
fed_config.client_selection_ratio = 0.3;
|
|
fed_config.byzantine_tolerance = true;
|
|
|
|
// Run simulation
|
|
info!(
|
|
"🎮 Starting simulation with {} clients for {} rounds",
|
|
sim_config.num_clients, sim_config.num_rounds
|
|
);
|
|
|
|
let mut simulation = simulation::FederatedSimulation::new(sim_config, fed_config).await?;
|
|
let results = simulation.run_simulation().await?;
|
|
|
|
// Display results
|
|
info!("🎯 Simulation Results:");
|
|
info!(" - Total rounds: {}", results.num_rounds);
|
|
info!(" - Final accuracy: {:.4}", results.final_accuracy);
|
|
info!(
|
|
" - Total communication cost: {:.2}",
|
|
results.total_communication_cost
|
|
);
|
|
info!(" - Convergence round: {:?}", results.convergence_round);
|
|
|
|
// Show round-by-round progress
|
|
info!("📈 Round-by-round progress:");
|
|
for (i, round_result) in results.round_results.iter().enumerate() {
|
|
if i % 2 == 0 || i == results.round_results.len() - 1 {
|
|
info!(
|
|
" Round {}: {} participants, accuracy: {:.4}, cost: {:.2}",
|
|
round_result.round_number,
|
|
round_result.participating_clients,
|
|
round_result.average_accuracy,
|
|
round_result.communication_cost
|
|
);
|
|
}
|
|
}
|
|
|
|
info!("✅ Simulation environment demo completed\n");
|
|
Ok(())
|
|
}
|
|
|
|
// Utility function to demonstrate infrastructure components
|
|
async fn _demo_infrastructure_components() -> std::result::Result<(), Box<dyn std::error::Error>> {
|
|
info!("🏗️ Infrastructure Components Demo");
|
|
|
|
// Client Manager
|
|
let manager = infrastructure::ClientManager::new().await?;
|
|
for i in 0..5 {
|
|
let client = Client::new(format!("infra_client_{}", i)).await?;
|
|
|
|
// 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: client.id,
|
|
endpoint: client
|
|
.endpoint
|
|
.as_ref()
|
|
.and_then(|e| e.parse::<SocketAddr>().ok())
|
|
.unwrap_or_else(|| "127.0.0.1:8080".parse().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?;
|
|
}
|
|
|
|
let total_clients = manager.connected_clients();
|
|
info!(
|
|
"📊 Client Manager Stats: {} total clients connected",
|
|
total_clients
|
|
);
|
|
|
|
info!("✅ Infrastructure components demo completed");
|
|
Ok(())
|
|
}
|