rustytorch/crates/training/rtx-federated/benches/federated_bench.rs

//! Benchmarks for RTX Federated Learning

use criterion::{BenchmarkId, Criterion, black_box, criterion_group, criterion_main};
use rtx_federated::*;
use std::time::Duration;
use tokio::runtime::Runtime;

fn benchmark_aggregation_algorithms(c: &mut Criterion) {
    let rt = Runtime::new().unwrap();

    let mut group = c.benchmark_group("aggregation_algorithms");
    group.measurement_time(Duration::from_secs(10));

    // Create test updates
    let updates = rt.block_on(async {
        let mut updates = Vec::new();
        for i in 0..100 {
            let mut update = aggregation::ModelUpdate::new(uuid::Uuid::new_v4());
            update.add_parameter("layer1", vec![1.0 + i as f64 * 0.01; 1000]);
            update.add_parameter("layer2", vec![0.5 + i as f64 * 0.005; 500]);
            update.sample_count = 100 + i;
            updates.push(update);
        }
        updates
    });

    // Benchmark FedAvg
    group.bench_function("fedavg", |b| {
        b.to_async(&rt).iter(|| async {
            let fedavg = aggregation::FedAvg::new(Some(0.9), false).await.unwrap();
            let result = fedavg.aggregate(black_box(&updates)).await.unwrap();
            black_box(result);
        });
    });

    // Benchmark FedProx
    group.bench_function("fedprox", |b| {
        b.to_async(&rt).iter(|| async {
            let fedprox = aggregation::FedProx::new(0.01, 5).await.unwrap();
            let result = fedprox.aggregate(black_box(&updates)).await.unwrap();
            black_box(result);
        });
    });

    group.finish();
}

fn benchmark_privacy_mechanisms(c: &mut Criterion) {
    let rt = Runtime::new().unwrap();

    let mut group = c.benchmark_group("privacy_mechanisms");

    let update = rt.block_on(async {
        let mut update = aggregation::ModelUpdate::new(uuid::Uuid::new_v4());
        update.add_parameter("layer1", vec![1.0; 10000]); // Large parameter set
        update.sample_count = 1000;
        update
    });

    // Benchmark Differential Privacy
    group.bench_function("differential_privacy", |b| {
        b.to_async(&rt).iter(|| async {
            let dp = privacy::DifferentialPrivacy::new(1.0, 1e-5).await.unwrap();
            let result = dp.apply_privacy(black_box(&update)).await.unwrap();
            black_box(result);
        });
    });

    group.finish();
}

fn benchmark_byzantine_protection(c: &mut Criterion) {
    let rt = Runtime::new().unwrap();

    let mut group = c.benchmark_group("byzantine_protection");

    // Create updates with some outliers
    let updates = rt.block_on(async {
        let mut updates = Vec::new();
        for i in 0..50 {
            let mut update = aggregation::ModelUpdate::new(uuid::Uuid::new_v4());
            if i < 45 {
                // Normal updates
                update.add_parameter("layer1", vec![1.0 + i as f64 * 0.01; 1000]);
            } else {
                // Outlier updates
                update.add_parameter("layer1", vec![100.0; 1000]);
            }
            update.sample_count = 100;
            updates.push(update);
        }
        updates
    });

    // Benchmark Krum
    group.bench_function("krum", |b| {
        b.to_async(&rt).iter(|| async {
            let krum = byzantine::Krum::new(2.0).await.unwrap();
            let result = krum.filter_updates(black_box(&updates)).await.unwrap();
            black_box(result);
        });
    });

    group.finish();
}

fn benchmark_federated_rounds(c: &mut Criterion) {
    let rt = Runtime::new().unwrap();

    let mut group = c.benchmark_group("federated_rounds");
    group.sample_size(10); // Fewer samples for longer benchmarks

    for num_clients in [10, 50, 100].iter() {
        group.bench_with_input(
            BenchmarkId::new("federated_round", num_clients),
            num_clients,
            |b, &num_clients| {
                b.to_async(&rt).iter(|| async {
                    let config = FederatedConfig::new();
                    let mut fed_system = FederatedSystem::new(config).await.unwrap();

                    // Register clients
                    for i in 0..num_clients {
                        let client = Client::new(format!("bench_client_{}", i)).await.unwrap();
                        fed_system.register_client(client).await.unwrap();
                    }

                    // Run one round
                    let metrics = fed_system.run_round().await.unwrap();
                    black_box(metrics);
                });
            },
        );
    }

    group.finish();
}

fn benchmark_client_management(c: &mut Criterion) {
    let rt = Runtime::new().unwrap();

    let mut group = c.benchmark_group("client_management");

    // Benchmark client registration
    group.bench_function("client_registration", |b| {
        b.to_async(&rt).iter(|| async {
            let manager = infrastructure::ClientManager::new().await.unwrap();

            for i in 0..100 {
                let client = Client::new(format!("bench_client_{}", i)).await.unwrap();
                manager.register_client(black_box(client)).await.unwrap();
            }
        });
    });

    group.finish();
}

fn benchmark_simulation(c: &mut Criterion) {
    let rt = Runtime::new().unwrap();

    let mut group = c.benchmark_group("simulation");
    group.sample_size(10);
    group.measurement_time(Duration::from_secs(30));

    group.bench_function("small_simulation", |b| {
        b.to_async(&rt).iter(|| async {
            let sim_config = simulation::SimulationConfig {
                num_clients: 20,
                num_rounds: 3,
                participation_rate: 0.5,
                data_distribution: simulation::DataDistribution::NonIID {
                    heterogeneity_level: 0.3,
                },
                ..Default::default()
            };

            let fed_config = FederatedConfig::new();
            let mut simulation = simulation::FederatedSimulation::new(sim_config, fed_config)
                .await
                .unwrap();

            let results = simulation.run_simulation().await.unwrap();
            black_box(results);
        });
    });

    group.finish();
}

criterion_group!(
    benches,
    benchmark_aggregation_algorithms,
    benchmark_privacy_mechanisms,
    benchmark_byzantine_protection,
    benchmark_federated_rounds,
    benchmark_client_management,
    benchmark_simulation
);

criterion_main!(benches);