561 lines
18 KiB
Rust
561 lines
18 KiB
Rust
//! STRICT TDD - Federation Manager Tests
|
|
//! These tests define the complete federated learning system behavior before implementation
|
|
|
|
use chrono::Utc;
|
|
use std::collections::HashMap;
|
|
use uuid::Uuid;
|
|
|
|
use rtx_platform::federation::{
|
|
AggregationStrategy, FederatedTrainingJob, FederationManager, ModelParameters,
|
|
ModelSharingRequest, ModelUpdate, ParticipantInfo, PrivacyBudget, PrivacyConfig,
|
|
};
|
|
use rtx_platform::{PlatformConfig, PlatformResult};
|
|
|
|
fn create_test_privacy_config() -> PrivacyConfig {
|
|
PrivacyConfig {
|
|
epsilon: 1.0, // Differential privacy parameter
|
|
delta: 1e-5, // Differential privacy parameter
|
|
noise_multiplier: 1.1,
|
|
max_grad_norm: 4.0,
|
|
secure_aggregation: true,
|
|
homomorphic_encryption: true,
|
|
minimum_participants: 3,
|
|
consent_required: true,
|
|
}
|
|
}
|
|
|
|
fn create_test_platform_config() -> PlatformConfig {
|
|
PlatformConfig {
|
|
database_url: "postgresql://test:test@localhost/federation_test".to_string(),
|
|
redis_urls: vec!["redis://localhost:6379".to_string()],
|
|
kafka_brokers: vec!["localhost:9092".to_string()],
|
|
metrics_endpoint: "localhost:9090".to_string(),
|
|
regions: HashMap::new(),
|
|
sla_targets: rtx_platform::slo::SlaTargets::default(),
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_federation_manager_creation() -> PlatformResult<()> {
|
|
let config = create_test_platform_config();
|
|
let privacy_config = create_test_privacy_config();
|
|
|
|
let federation_manager = FederationManager::new(&config, privacy_config).await?;
|
|
|
|
assert!(federation_manager.is_healthy().await?);
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_participant_registration() -> PlatformResult<()> {
|
|
let config = create_test_platform_config();
|
|
let privacy_config = create_test_privacy_config();
|
|
|
|
let mut federation_manager = FederationManager::new(&config, privacy_config).await?;
|
|
federation_manager.start().await?;
|
|
|
|
let participant_info = ParticipantInfo {
|
|
id: Uuid::new_v4(),
|
|
tenant_id: Uuid::new_v4(),
|
|
name: "test-participant".to_string(),
|
|
capabilities: vec!["torch".to_string(), "tensorflow".to_string()],
|
|
data_size: 10000,
|
|
compute_power: 100.0,
|
|
bandwidth: 1000.0,
|
|
privacy_level: "high".to_string(),
|
|
};
|
|
|
|
federation_manager
|
|
.register_participant(participant_info.clone())
|
|
.await?;
|
|
|
|
let participants = federation_manager.get_participants().await?;
|
|
assert_eq!(participants.len(), 1);
|
|
assert_eq!(participants[0].id, participant_info.id);
|
|
|
|
federation_manager.shutdown().await?;
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_federated_model_creation() -> PlatformResult<()> {
|
|
let config = create_test_platform_config();
|
|
let privacy_config = create_test_privacy_config();
|
|
|
|
let mut federation_manager = FederationManager::new(&config, privacy_config).await?;
|
|
federation_manager.start().await?;
|
|
|
|
let model_params = ModelParameters {
|
|
model_type: "neural_network".to_string(),
|
|
architecture: "resnet50".to_string(),
|
|
parameters: HashMap::from([
|
|
("layers".to_string(), "50".to_string()),
|
|
("input_size".to_string(), "224".to_string()),
|
|
]),
|
|
weights: vec![0.1, 0.2, 0.3, 0.4, 0.5], // Simplified weights
|
|
};
|
|
|
|
let federated_model = federation_manager
|
|
.create_federated_model("image_classification", model_params)
|
|
.await?;
|
|
|
|
assert_eq!(federated_model.name, "image_classification");
|
|
assert_eq!(federated_model.parameters.model_type, "neural_network");
|
|
|
|
federation_manager.shutdown().await?;
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_secure_aggregation_with_encryption() -> PlatformResult<()> {
|
|
let config = create_test_platform_config();
|
|
let privacy_config = create_test_privacy_config();
|
|
|
|
let federation_manager = FederationManager::new(&config, privacy_config).await?;
|
|
let secure_aggregator = federation_manager.secure_aggregator();
|
|
|
|
// Create model updates from multiple participants
|
|
let updates = vec![
|
|
ModelUpdate {
|
|
participant_id: Uuid::new_v4(),
|
|
model_id: Uuid::new_v4(),
|
|
round: 1,
|
|
parameters: vec![1.0, 2.0, 3.0],
|
|
gradient_norm: 2.5,
|
|
timestamp: Utc::now(),
|
|
encrypted: true,
|
|
},
|
|
ModelUpdate {
|
|
participant_id: Uuid::new_v4(),
|
|
model_id: Uuid::new_v4(),
|
|
round: 1,
|
|
parameters: vec![2.0, 3.0, 4.0],
|
|
gradient_norm: 3.2,
|
|
timestamp: Utc::now(),
|
|
encrypted: true,
|
|
},
|
|
ModelUpdate {
|
|
participant_id: Uuid::new_v4(),
|
|
model_id: Uuid::new_v4(),
|
|
round: 1,
|
|
parameters: vec![1.5, 2.5, 3.5],
|
|
gradient_norm: 2.8,
|
|
timestamp: Utc::now(),
|
|
encrypted: true,
|
|
},
|
|
];
|
|
|
|
let aggregated_update = secure_aggregator
|
|
.aggregate_updates(&updates, AggregationStrategy::FederatedAverage)
|
|
.await?;
|
|
|
|
// Verify secure aggregation (average of encrypted updates)
|
|
assert_eq!(aggregated_update.parameters.len(), 3);
|
|
assert!((aggregated_update.parameters[0] - 1.5).abs() < 0.1); // (1.0+2.0+1.5)/3 = 1.5
|
|
assert!((aggregated_update.parameters[1] - 2.5).abs() < 0.1); // (2.0+3.0+2.5)/3 = 2.5
|
|
assert!((aggregated_update.parameters[2] - 3.5).abs() < 0.1); // (3.0+4.0+3.5)/3 = 3.5
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_differential_privacy_noise_injection() -> PlatformResult<()> {
|
|
let config = create_test_platform_config();
|
|
let privacy_config = create_test_privacy_config();
|
|
|
|
let federation_manager = FederationManager::new(&config, privacy_config).await?;
|
|
let dp_engine = federation_manager.differential_privacy_engine();
|
|
|
|
let original_update = ModelUpdate {
|
|
participant_id: Uuid::new_v4(),
|
|
model_id: Uuid::new_v4(),
|
|
round: 1,
|
|
parameters: vec![1.0, 2.0, 3.0, 4.0, 5.0],
|
|
gradient_norm: 3.0,
|
|
timestamp: Utc::now(),
|
|
encrypted: false,
|
|
};
|
|
|
|
let privacy_budget = PrivacyBudget {
|
|
epsilon_consumed: 0.5,
|
|
epsilon_remaining: 0.5,
|
|
delta_consumed: 5e-6,
|
|
delta_remaining: 5e-6,
|
|
};
|
|
|
|
let private_update = dp_engine
|
|
.add_noise(original_update.clone(), privacy_budget)
|
|
.await?;
|
|
|
|
// Verify noise was added (parameters should be different but similar)
|
|
assert_eq!(
|
|
private_update.parameters.len(),
|
|
original_update.parameters.len()
|
|
);
|
|
|
|
let mut differences = 0;
|
|
for (original, private) in original_update
|
|
.parameters
|
|
.iter()
|
|
.zip(private_update.parameters.iter())
|
|
{
|
|
if (original - private).abs() > 0.001 {
|
|
differences += 1;
|
|
}
|
|
}
|
|
|
|
// Most parameters should have noise added
|
|
assert!(differences >= 3);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_homomorphic_encryption_operations() -> PlatformResult<()> {
|
|
let config = create_test_platform_config();
|
|
let privacy_config = create_test_privacy_config();
|
|
|
|
let federation_manager = FederationManager::new(&config, privacy_config).await?;
|
|
let he_engine = federation_manager.homomorphic_encryption();
|
|
|
|
// Generate encryption keys
|
|
let keypair = he_engine.generate_keypair().await?;
|
|
|
|
// Encrypt some model parameters
|
|
let plaintext_params = vec![1.5, 2.7, 3.9];
|
|
let encrypted_params = he_engine
|
|
.encrypt(&plaintext_params, &keypair.public_key)
|
|
.await?;
|
|
|
|
// Perform homomorphic addition (simulating secure aggregation)
|
|
let another_encrypted = he_engine
|
|
.encrypt(&vec![0.5, 0.3, 0.1], &keypair.public_key)
|
|
.await?;
|
|
let sum_encrypted = he_engine
|
|
.add_encrypted(&encrypted_params, &another_encrypted)
|
|
.await?;
|
|
|
|
// Decrypt the result
|
|
let decrypted_sum = he_engine
|
|
.decrypt(&sum_encrypted, &keypair.private_key)
|
|
.await?;
|
|
|
|
// Verify homomorphic addition worked correctly
|
|
assert_eq!(decrypted_sum.len(), 3);
|
|
assert!((decrypted_sum[0] - 2.0).abs() < 0.1); // 1.5 + 0.5 = 2.0
|
|
assert!((decrypted_sum[1] - 3.0).abs() < 0.1); // 2.7 + 0.3 = 3.0
|
|
assert!((decrypted_sum[2] - 4.0).abs() < 0.1); // 3.9 + 0.1 = 4.0
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_consent_management_system() -> PlatformResult<()> {
|
|
let config = create_test_platform_config();
|
|
let privacy_config = create_test_privacy_config();
|
|
|
|
let mut federation_manager = FederationManager::new(&config, privacy_config).await?;
|
|
federation_manager.start().await?;
|
|
|
|
let consent_manager = federation_manager.consent_manager();
|
|
let tenant_id = Uuid::new_v4();
|
|
let model_id = Uuid::new_v4();
|
|
|
|
// Request consent for model sharing
|
|
let sharing_request = ModelSharingRequest {
|
|
id: Uuid::new_v4(),
|
|
requester_tenant_id: tenant_id,
|
|
target_model_id: model_id,
|
|
purpose: "collaborative_training".to_string(),
|
|
data_usage: "aggregation_only".to_string(),
|
|
duration_days: 30,
|
|
privacy_guarantees: vec![
|
|
"differential_privacy".to_string(),
|
|
"secure_aggregation".to_string(),
|
|
],
|
|
};
|
|
|
|
consent_manager
|
|
.request_consent(sharing_request.clone())
|
|
.await?;
|
|
|
|
// Verify consent request was stored
|
|
let pending_requests = consent_manager.get_pending_requests(tenant_id).await?;
|
|
assert_eq!(pending_requests.len(), 1);
|
|
assert_eq!(pending_requests[0].id, sharing_request.id);
|
|
|
|
// Grant consent
|
|
consent_manager
|
|
.grant_consent(sharing_request.id, tenant_id)
|
|
.await?;
|
|
|
|
// Verify consent was granted
|
|
let has_consent = consent_manager.has_consent(tenant_id, model_id).await?;
|
|
assert!(has_consent);
|
|
|
|
federation_manager.shutdown().await?;
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_federated_training_job_lifecycle() -> PlatformResult<()> {
|
|
let config = create_test_platform_config();
|
|
let privacy_config = create_test_privacy_config();
|
|
|
|
let mut federation_manager = FederationManager::new(&config, privacy_config).await?;
|
|
federation_manager.start().await?;
|
|
|
|
// Register participants
|
|
let participants = vec![
|
|
ParticipantInfo {
|
|
id: Uuid::new_v4(),
|
|
tenant_id: Uuid::new_v4(),
|
|
name: "participant-1".to_string(),
|
|
capabilities: vec!["torch".to_string()],
|
|
data_size: 5000,
|
|
compute_power: 50.0,
|
|
bandwidth: 500.0,
|
|
privacy_level: "high".to_string(),
|
|
},
|
|
ParticipantInfo {
|
|
id: Uuid::new_v4(),
|
|
tenant_id: Uuid::new_v4(),
|
|
name: "participant-2".to_string(),
|
|
capabilities: vec!["torch".to_string()],
|
|
data_size: 7000,
|
|
compute_power: 70.0,
|
|
bandwidth: 700.0,
|
|
privacy_level: "high".to_string(),
|
|
},
|
|
ParticipantInfo {
|
|
id: Uuid::new_v4(),
|
|
tenant_id: Uuid::new_v4(),
|
|
name: "participant-3".to_string(),
|
|
capabilities: vec!["torch".to_string()],
|
|
data_size: 3000,
|
|
compute_power: 30.0,
|
|
bandwidth: 300.0,
|
|
privacy_level: "high".to_string(),
|
|
},
|
|
];
|
|
|
|
for participant in &participants {
|
|
federation_manager
|
|
.register_participant(participant.clone())
|
|
.await?;
|
|
}
|
|
|
|
// Create federated training job
|
|
let job = FederatedTrainingJob {
|
|
id: Uuid::new_v4(),
|
|
name: "image_classification_job".to_string(),
|
|
model_id: Uuid::new_v4(),
|
|
participants: participants.iter().map(|p| p.id).collect(),
|
|
target_rounds: 5,
|
|
current_round: 0,
|
|
convergence_threshold: 0.001,
|
|
privacy_budget: PrivacyBudget {
|
|
epsilon_consumed: 0.0,
|
|
epsilon_remaining: 1.0,
|
|
delta_consumed: 0.0,
|
|
delta_remaining: 1e-5,
|
|
},
|
|
status: rtx_platform::federation::TrainingStatus::Initializing,
|
|
created_at: Utc::now(),
|
|
updated_at: Utc::now(),
|
|
};
|
|
|
|
let started_job = federation_manager.start_training_job(job).await?;
|
|
|
|
assert_eq!(
|
|
started_job.status,
|
|
rtx_platform::federation::TrainingStatus::Running
|
|
);
|
|
assert_eq!(started_job.participants.len(), 3);
|
|
|
|
// Simulate training round completion
|
|
federation_manager
|
|
.advance_training_round(started_job.id)
|
|
.await?;
|
|
|
|
let updated_job = federation_manager.get_training_job(started_job.id).await?;
|
|
assert_eq!(updated_job.current_round, 1);
|
|
|
|
federation_manager.shutdown().await?;
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_privacy_budget_tracking() -> PlatformResult<()> {
|
|
let config = create_test_platform_config();
|
|
let mut privacy_config = create_test_privacy_config();
|
|
privacy_config.epsilon = 2.0; // Higher budget for testing
|
|
|
|
let federation_manager = FederationManager::new(&config, privacy_config).await?;
|
|
let dp_engine = federation_manager.differential_privacy_engine();
|
|
|
|
let tenant_id = Uuid::new_v4();
|
|
|
|
// Initial privacy budget should be full
|
|
let initial_budget = dp_engine.get_privacy_budget(tenant_id).await?;
|
|
assert_eq!(initial_budget.epsilon_remaining, 2.0);
|
|
assert_eq!(initial_budget.epsilon_consumed, 0.0);
|
|
|
|
// Consume some privacy budget
|
|
let update = ModelUpdate {
|
|
participant_id: Uuid::new_v4(),
|
|
model_id: Uuid::new_v4(),
|
|
round: 1,
|
|
parameters: vec![1.0, 2.0, 3.0],
|
|
gradient_norm: 2.0,
|
|
timestamp: Utc::now(),
|
|
encrypted: false,
|
|
};
|
|
|
|
let consumed_budget = PrivacyBudget {
|
|
epsilon_consumed: 0.5,
|
|
epsilon_remaining: 1.5,
|
|
delta_consumed: 2e-6,
|
|
delta_remaining: 8e-6,
|
|
};
|
|
|
|
dp_engine.add_noise(update, consumed_budget.clone()).await?;
|
|
|
|
// Update privacy budget tracking
|
|
dp_engine
|
|
.update_privacy_budget(tenant_id, consumed_budget.clone())
|
|
.await?;
|
|
|
|
// Check updated budget
|
|
let updated_budget = dp_engine.get_privacy_budget(tenant_id).await?;
|
|
assert_eq!(updated_budget.epsilon_consumed, 0.5);
|
|
assert_eq!(updated_budget.epsilon_remaining, 1.5);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_cross_tenant_model_sharing() -> PlatformResult<()> {
|
|
let config = create_test_platform_config();
|
|
let privacy_config = create_test_privacy_config();
|
|
|
|
let mut federation_manager = FederationManager::new(&config, privacy_config).await?;
|
|
federation_manager.start().await?;
|
|
|
|
let tenant_a = Uuid::new_v4();
|
|
let tenant_b = Uuid::new_v4();
|
|
|
|
// Create a model for tenant A
|
|
let model_params = ModelParameters {
|
|
model_type: "neural_network".to_string(),
|
|
architecture: "transformer".to_string(),
|
|
parameters: HashMap::from([
|
|
("layers".to_string(), "12".to_string()),
|
|
("hidden_size".to_string(), "768".to_string()),
|
|
]),
|
|
weights: vec![0.1, 0.2, 0.3],
|
|
};
|
|
|
|
let model = federation_manager
|
|
.create_federated_model("nlp_model", model_params)
|
|
.await?;
|
|
|
|
// Request sharing consent between tenants
|
|
let sharing_request = ModelSharingRequest {
|
|
id: Uuid::new_v4(),
|
|
requester_tenant_id: tenant_b,
|
|
target_model_id: model.id,
|
|
purpose: "transfer_learning".to_string(),
|
|
data_usage: "parameter_sharing".to_string(),
|
|
duration_days: 14,
|
|
privacy_guarantees: vec![
|
|
"differential_privacy".to_string(),
|
|
"encrypted_transfer".to_string(),
|
|
],
|
|
};
|
|
|
|
federation_manager
|
|
.request_model_sharing(sharing_request.clone())
|
|
.await?;
|
|
|
|
// Grant consent (simulating tenant A approval)
|
|
federation_manager
|
|
.approve_model_sharing(sharing_request.id, tenant_a)
|
|
.await?;
|
|
|
|
// Verify tenant B can now access the model
|
|
let shared_model = federation_manager
|
|
.get_shared_model(model.id, tenant_b)
|
|
.await?;
|
|
|
|
assert_eq!(shared_model.id, model.id);
|
|
assert_eq!(shared_model.name, "nlp_model");
|
|
|
|
federation_manager.shutdown().await?;
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_byzantine_fault_tolerance() -> PlatformResult<()> {
|
|
let config = create_test_platform_config();
|
|
let privacy_config = create_test_privacy_config();
|
|
|
|
let federation_manager = FederationManager::new(&config, privacy_config).await?;
|
|
let secure_aggregator = federation_manager.secure_aggregator();
|
|
|
|
// Create mix of honest and malicious updates
|
|
let updates = vec![
|
|
// Honest updates
|
|
ModelUpdate {
|
|
participant_id: Uuid::new_v4(),
|
|
model_id: Uuid::new_v4(),
|
|
round: 1,
|
|
parameters: vec![1.0, 2.0, 3.0],
|
|
gradient_norm: 2.5,
|
|
timestamp: Utc::now(),
|
|
encrypted: true,
|
|
},
|
|
ModelUpdate {
|
|
participant_id: Uuid::new_v4(),
|
|
model_id: Uuid::new_v4(),
|
|
round: 1,
|
|
parameters: vec![1.1, 2.1, 3.1],
|
|
gradient_norm: 2.6,
|
|
timestamp: Utc::now(),
|
|
encrypted: true,
|
|
},
|
|
ModelUpdate {
|
|
participant_id: Uuid::new_v4(),
|
|
model_id: Uuid::new_v4(),
|
|
round: 1,
|
|
parameters: vec![0.9, 1.9, 2.9],
|
|
gradient_norm: 2.4,
|
|
timestamp: Utc::now(),
|
|
encrypted: true,
|
|
},
|
|
// Malicious update (extreme values)
|
|
ModelUpdate {
|
|
participant_id: Uuid::new_v4(),
|
|
model_id: Uuid::new_v4(),
|
|
round: 1,
|
|
parameters: vec![100.0, 200.0, 300.0], // Outlier values
|
|
gradient_norm: 50.0,
|
|
timestamp: Utc::now(),
|
|
encrypted: true,
|
|
},
|
|
];
|
|
|
|
let aggregated_update = secure_aggregator
|
|
.aggregate_updates(&updates, AggregationStrategy::ByzantineRobust)
|
|
.await?;
|
|
|
|
// Verify Byzantine-robust aggregation filtered out outliers
|
|
assert_eq!(aggregated_update.parameters.len(), 3);
|
|
// Should be close to average of honest updates (~1.0, ~2.0, ~3.0)
|
|
assert!(aggregated_update.parameters[0] < 10.0); // Much less than outlier
|
|
assert!(aggregated_update.parameters[1] < 20.0);
|
|
assert!(aggregated_update.parameters[2] < 30.0);
|
|
|
|
Ok(())
|
|
}
|