---
title: "LAB17: Level 2 Exercises"
---

## Intermediate Practice Problems

These exercises build on Level 1 and introduce multi-process federated learning deployments with realistic network conditions, privacy mechanisms, and edge device considerations.

::: {.callout-note}
## Getting Started
Complete Level 1 exercises first, then work through these using separate processes, containers, or networked machines to simulate real federated learning deployments.
:::

### Exercise 1: Flower Server Setup and FedAvg Strategy
**Difficulty:** Level 2
**Requirements:** Python, Flower (flwr), TensorFlow/PyTorch

**Task:**
Configure a Flower federated learning server with the FedAvg strategy. Set up proper client selection, minimum client requirements, and aggregation parameters. Monitor per-round metrics including accuracy, loss, and participation rates.

Steps:
1. Install Flower and create a server configuration file
2. Implement FedAvg strategy with fraction_fit, fraction_evaluate parameters
3. Set minimum clients (min_fit_clients, min_available_clients)
4. Add custom metric aggregation to track convergence
5. Run the server and monitor client connections and training rounds

::: {.callout-tip collapse="true"}
## Solution
```python
# server.py
import flwr as fl
from typing import List, Tuple, Dict, Optional
from flwr.common import Metrics
import numpy as np

def weighted_average(metrics: List[Tuple[int, Metrics]]) -> Metrics:
    """Aggregate metrics from multiple clients using weighted average."""
    # Multiply accuracy of each client by number of examples used
    accuracies = [num_examples * m["accuracy"] for num_examples, m in metrics]
    examples = [num_examples for num_examples, m in metrics]

    # Aggregate and return custom metrics
    return {
        "accuracy": sum(accuracies) / sum(examples),
        "num_clients": len(metrics),
        "total_examples": sum(examples)
    }

def fit_config(server_round: int) -> Dict:
    """Return training configuration for each round."""
    config = {
        "server_round": server_round,
        "local_epochs": 5 if server_round < 3 else 3,  # Reduce epochs after warmup
        "batch_size": 32,
        "learning_rate": 0.01 if server_round < 5 else 0.005,  # Decay LR
    }
    return config

def evaluate_config(server_round: int) -> Dict:
    """Return evaluation configuration."""
    return {
        "server_round": server_round,
    }

# Define strategy
strategy = fl.server.strategy.FedAvg(
    fraction_fit=0.8,  # Sample 80% of available clients for training
    fraction_evaluate=1.0,  # Use all clients for evaluation
    min_fit_clients=3,  # Minimum clients needed to proceed with training
    min_evaluate_clients=3,  # Minimum clients for evaluation
    min_available_clients=3,  # Wait for 3 clients to connect before starting
    evaluate_fn=None,  # Optional: server-side evaluation function
    on_fit_config_fn=fit_config,  # Pass config to clients for training
    on_evaluate_config_fn=evaluate_config,  # Pass config for evaluation
    evaluate_metrics_aggregation_fn=weighted_average,  # Aggregate metrics
    fit_metrics_aggregation_fn=weighted_average,
)

# Configure server
config = fl.server.ServerConfig(num_rounds=10)

# Custom history callback
class MetricsHistory(fl.server.History):
    def __init__(self):
        super().__init__()
        self.rounds_data = []

    def add_loss_distributed(self, server_round, loss):
        super().add_loss_distributed(server_round, loss)
        print(f"Round {server_round} - Distributed Loss: {loss}")

    def add_metrics_distributed(self, server_round, metrics):
        super().add_metrics_distributed(server_round, metrics)
        self.rounds_data.append({
            'round': server_round,
            'metrics': metrics
        })
        print(f"Round {server_round} - Metrics: {metrics}")

# Start server
print("Starting Flower server...")
print(f"Configuration:")
print(f"  - Waiting for {strategy.min_available_clients} clients to connect")
print(f"  - Training with {int(strategy.fraction_fit * 100)}% of clients per round")
print(f"  - Running {config.num_rounds} rounds")
print(f"  - Server address: 0.0.0.0:8080")
print()

history = fl.server.start_server(
    server_address="0.0.0.0:8080",
    config=config,
    strategy=strategy,
)

# Print summary
print("\n" + "="*60)
print("TRAINING COMPLETE")
print("="*60)
if history.losses_distributed:
    print(f"Final distributed loss: {history.losses_distributed[-1][1]:.4f}")
if history.metrics_distributed:
    final_metrics = history.metrics_distributed[-1][1]
    print(f"Final accuracy: {final_metrics.get('accuracy', 0):.2%}")
    print(f"Total clients participated: {final_metrics.get('num_clients', 0)}")
```

```python
# client.py
import flwr as fl
import tensorflow as tf
import numpy as np
import argparse

def load_partition(client_id: int, num_clients: int = 5):
    """Load and partition MNIST dataset for a specific client."""
    # Load MNIST
    (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()

    # Normalize
    x_train = x_train.astype('float32') / 255.0
    x_test = x_test.astype('float32') / 255.0

    # Reshape
    x_train = x_train.reshape(-1, 28, 28, 1)
    x_test = x_test.reshape(-1, 28, 28, 1)

    # Partition data (IID)
    partition_size = len(x_train) // num_clients
    start = client_id * partition_size
    end = start + partition_size

    return (
        x_train[start:end],
        y_train[start:end],
        x_test[start:end],
        y_test[start:end]
    )

def create_model():
    """Create a simple CNN model for MNIST."""
    model = tf.keras.Sequential([
        tf.keras.layers.Conv2D(32, (3, 3), activation='relu', input_shape=(28, 28, 1)),
        tf.keras.layers.MaxPooling2D((2, 2)),
        tf.keras.layers.Conv2D(64, (3, 3), activation='relu'),
        tf.keras.layers.MaxPooling2D((2, 2)),
        tf.keras.layers.Flatten(),
        tf.keras.layers.Dense(128, activation='relu'),
        tf.keras.layers.Dropout(0.5),
        tf.keras.layers.Dense(10, activation='softmax')
    ])

    model.compile(
        optimizer='adam',
        loss='sparse_categorical_crossentropy',
        metrics=['accuracy']
    )

    return model

class MNISTClient(fl.client.NumPyClient):
    def __init__(self, client_id, model, x_train, y_train, x_test, y_test):
        self.client_id = client_id
        self.model = model
        self.x_train = x_train
        self.y_train = y_train
        self.x_test = x_test
        self.y_test = y_test

    def get_parameters(self, config):
        """Return current model parameters."""
        return self.model.get_weights()

    def fit(self, parameters, config):
        """Train model on local data."""
        # Update local model with global parameters
        self.model.set_weights(parameters)

        # Get training config
        local_epochs = config.get("local_epochs", 5)
        batch_size = config.get("batch_size", 32)

        # Train
        print(f"[Client {self.client_id}] Training for {local_epochs} epochs...")
        history = self.model.fit(
            self.x_train, self.y_train,
            batch_size=batch_size,
            epochs=local_epochs,
            verbose=0
        )

        # Return updated model parameters and training metrics
        parameters_updated = self.model.get_weights()
        num_examples = len(self.x_train)
        metrics = {
            "accuracy": float(history.history['accuracy'][-1]),
            "loss": float(history.history['loss'][-1])
        }

        print(f"[Client {self.client_id}] Training complete - "
              f"Acc: {metrics['accuracy']:.2%}, Loss: {metrics['loss']:.4f}")

        return parameters_updated, num_examples, metrics

    def evaluate(self, parameters, config):
        """Evaluate model on local test data."""
        self.model.set_weights(parameters)

        loss, accuracy = self.model.evaluate(self.x_test, self.y_test, verbose=0)
        num_examples = len(self.x_test)

        print(f"[Client {self.client_id}] Evaluation - "
              f"Acc: {accuracy:.2%}, Loss: {loss:.4f}")

        return float(loss), num_examples, {"accuracy": float(accuracy)}

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Flower client")
    parser.add_argument("--client-id", type=int, required=True, help="Client ID (0-4)")
    parser.add_argument("--server", type=str, default="localhost:8080", help="Server address")
    args = parser.parse_args()

    print(f"Starting Flower client {args.client_id}...")

    # Load data partition
    x_train, y_train, x_test, y_test = load_partition(args.client_id)
    print(f"Loaded partition: {len(x_train)} training samples, {len(x_test)} test samples")

    # Create model
    model = create_model()
    print(f"Model created: {model.count_params():,} parameters")

    # Start client
    client = MNISTClient(args.client_id, model, x_train, y_train, x_test, y_test)
    fl.client.start_numpy_client(server_address=args.server, client=client)
```

Run the federated learning system:
```bash
# Terminal 1: Start server
python server.py

# Terminal 2-4: Start clients (in separate terminals)
python client.py --client-id 0
python client.py --client-id 1
python client.py --client-id 2
```

**Expected Outcomes:**
- Server waits for minimum clients before starting training
- Each round, 80% of available clients are selected for training
- Clients train locally for configured epochs with adaptive learning rate
- Metrics are aggregated and logged per round
- Final global model achieves 95%+ accuracy on MNIST
:::

### Exercise 2: NumPyClient Implementation with Data Partitioning
**Difficulty:** Level 2
**Requirements:** Python, Flower, NumPy, TensorFlow/PyTorch

**Task:**
Implement a custom Flower NumPyClient that handles get_parameters, fit, and evaluate methods. Create data partitioning utilities to distribute datasets across multiple clients with both IID and non-IID distributions.

Steps:
1. Implement NumPyClient interface methods (get_parameters, fit, evaluate)
2. Create IID data partitioning (random uniform splits)
3. Create non-IID data partitioning (label skew, quantity skew)
4. Test with multiple clients using different partition strategies
5. Compare training convergence between IID and non-IID scenarios

::: {.callout-tip collapse="true"}
## Solution
```python
# data_partitioner.py
import numpy as np
from typing import List, Tuple, Dict
import tensorflow as tf

class DataPartitioner:
    """Utility class for partitioning datasets across FL clients."""

    @staticmethod
    def partition_iid(x_data, y_data, num_clients: int) -> List[Tuple]:
        """
        Partition data into IID (Independent and Identically Distributed) splits.
        Each client gets a random subset with similar distribution.
        """
        num_samples = len(x_data)
        indices = np.random.permutation(num_samples)
        partition_size = num_samples // num_clients

        partitions = []
        for i in range(num_clients):
            start_idx = i * partition_size
            end_idx = start_idx + partition_size if i < num_clients - 1 else num_samples

            client_indices = indices[start_idx:end_idx]
            partitions.append((
                x_data[client_indices],
                y_data[client_indices]
            ))

        return partitions

    @staticmethod
    def partition_non_iid_label_skew(x_data, y_data, num_clients: int,
                                     classes_per_client: int = 2) -> List[Tuple]:
        """
        Partition data with label skew (non-IID).
        Each client gets only a subset of classes.

        Example: With 10 classes and classes_per_client=2:
        - Client 0: only classes [0, 1]
        - Client 1: only classes [2, 3]
        - etc.
        """
        num_classes = len(np.unique(y_data))

        # Group indices by class
        class_indices = {i: np.where(y_data == i)[0] for i in range(num_classes)}

        partitions = []

        for client_id in range(num_clients):
            # Assign specific classes to this client
            # Rotate through classes to ensure coverage
            start_class = (client_id * classes_per_client) % num_classes
            client_classes = [
                (start_class + i) % num_classes
                for i in range(classes_per_client)
            ]

            # Collect samples from assigned classes
            client_indices = []
            for cls in client_classes:
                cls_indices = class_indices[cls]
                # Split class samples among clients that have this class
                samples_per_client = len(cls_indices) // (num_clients // (num_classes // classes_per_client))
                start = (client_id % (num_clients // (num_classes // classes_per_client))) * samples_per_client
                end = start + samples_per_client
                client_indices.extend(cls_indices[start:end])

            client_indices = np.array(client_indices)
            np.random.shuffle(client_indices)

            partitions.append((
                x_data[client_indices],
                y_data[client_indices]
            ))

            print(f"Client {client_id}: {len(client_indices)} samples, "
                  f"classes {client_classes}, "
                  f"distribution {np.bincount(y_data[client_indices].astype(int), minlength=num_classes)}")

        return partitions

    @staticmethod
    def partition_non_iid_quantity_skew(x_data, y_data, num_clients: int,
                                        min_samples_ratio: float = 0.1) -> List[Tuple]:
        """
        Partition data with quantity skew (non-IID).
        Clients have very different dataset sizes.

        Args:
            min_samples_ratio: Minimum fraction of average samples per client
        """
        num_samples = len(x_data)
        avg_samples = num_samples // num_clients

        # Generate random partition sizes with high variance
        sizes = []
        remaining = num_samples

        for i in range(num_clients - 1):
            # Random size between min_ratio and 2x average
            min_size = int(avg_samples * min_samples_ratio)
            max_size = min(remaining - (num_clients - i - 1) * min_size,
                          int(avg_samples * 2))
            size = np.random.randint(min_size, max_size + 1)
            sizes.append(size)
            remaining -= size

        sizes.append(remaining)  # Last client gets remainder

        # Shuffle and partition
        indices = np.random.permutation(num_samples)
        partitions = []
        start = 0

        for size in sizes:
            end = start + size
            client_indices = indices[start:end]
            partitions.append((
                x_data[client_indices],
                y_data[client_indices]
            ))
            start = end

        print("Quantity distribution:")
        for i, (x, y) in enumerate(partitions):
            print(f"  Client {i}: {len(x)} samples ({len(x)/num_samples*100:.1f}%)")

        return partitions

    @staticmethod
    def analyze_partitions(partitions: List[Tuple], num_classes: int = 10):
        """Analyze and visualize data partitions."""
        print("\n" + "="*60)
        print("PARTITION ANALYSIS")
        print("="*60)

        for i, (x, y) in enumerate(partitions):
            class_dist = np.bincount(y.astype(int), minlength=num_classes)
            entropy = -np.sum([p * np.log(p + 1e-10) for p in class_dist / len(y) if p > 0])

            print(f"\nClient {i}:")
            print(f"  Samples: {len(x)}")
            print(f"  Class distribution: {class_dist}")
            print(f"  Entropy: {entropy:.3f} (max={np.log(num_classes):.3f})")
            print(f"  Classes present: {np.sum(class_dist > 0)}/{num_classes}")

# client_advanced.py
import flwr as fl
import tensorflow as tf
import numpy as np
import argparse
from data_partitioner import DataPartitioner

class AdvancedFlowerClient(fl.client.NumPyClient):
    """
    Advanced Flower client with detailed metrics and adaptive training.
    """

    def __init__(self, client_id: int, model, x_train, y_train, x_test, y_test):
        self.client_id = client_id
        self.model = model
        self.x_train = x_train
        self.y_train = y_train
        self.x_test = x_test
        self.y_test = y_test
        self.round_num = 0

    def get_parameters(self, config):
        """Return current model weights as list of NumPy arrays."""
        return self.model.get_weights()

    def set_parameters(self, parameters):
        """Set model weights from list of NumPy arrays."""
        self.model.set_weights(parameters)

    def fit(self, parameters, config):
        """Train model on local data and return updated weights."""
        self.round_num += 1
        self.set_parameters(parameters)

        # Extract config
        local_epochs = config.get("local_epochs", 5)
        batch_size = config.get("batch_size", 32)
        learning_rate = config.get("learning_rate", 0.01)

        # Update learning rate
        tf.keras.backend.set_value(
            self.model.optimizer.learning_rate, learning_rate
        )

        # Train
        print(f"\n[Client {self.client_id}] Round {self.round_num}")
        print(f"  Training: {len(self.x_train)} samples, "
              f"{local_epochs} epochs, LR={learning_rate}")

        history = self.model.fit(
            self.x_train, self.y_train,
            batch_size=batch_size,
            epochs=local_epochs,
            validation_split=0.1,
            verbose=0
        )

        # Calculate metrics
        train_loss = float(history.history['loss'][-1])
        train_acc = float(history.history['accuracy'][-1])
        val_loss = float(history.history['val_loss'][-1])
        val_acc = float(history.history['val_accuracy'][-1])

        print(f"  Results: Train Acc={train_acc:.2%}, Val Acc={val_acc:.2%}")

        # Return updated weights, number of samples, and metrics
        metrics = {
            "accuracy": train_acc,
            "val_accuracy": val_acc,
            "loss": train_loss,
            "val_loss": val_loss,
            "client_id": self.client_id
        }

        return self.get_parameters(config), len(self.x_train), metrics

    def evaluate(self, parameters, config):
        """Evaluate model on local test set."""
        self.set_parameters(parameters)

        loss, accuracy = self.model.evaluate(
            self.x_test, self.y_test, verbose=0
        )

        # Per-class accuracy (for non-IID analysis)
        predictions = self.model.predict(self.x_test, verbose=0)
        pred_classes = np.argmax(predictions, axis=1)
        true_classes = self.y_test

        per_class_acc = {}
        for cls in np.unique(true_classes):
            mask = true_classes == cls
            if np.sum(mask) > 0:
                cls_acc = np.mean(pred_classes[mask] == true_classes[mask])
                per_class_acc[f"class_{cls}_acc"] = float(cls_acc)

        metrics = {
            "accuracy": float(accuracy),
            "client_id": self.client_id,
            **per_class_acc
        }

        print(f"  Evaluation: Test Acc={accuracy:.2%}, Loss={loss:.4f}")

        return float(loss), len(self.x_test), metrics

# main_comparison.py
"""
Compare IID vs Non-IID data distributions in federated learning.
"""
import tensorflow as tf
from data_partitioner import DataPartitioner

def create_model():
    """Create CNN model for MNIST."""
    model = tf.keras.Sequential([
        tf.keras.layers.Conv2D(32, (3, 3), activation='relu', input_shape=(28, 28, 1)),
        tf.keras.layers.MaxPooling2D((2, 2)),
        tf.keras.layers.Conv2D(64, (3, 3), activation='relu'),
        tf.keras.layers.MaxPooling2D((2, 2)),
        tf.keras.layers.Flatten(),
        tf.keras.layers.Dense(128, activation='relu'),
        tf.keras.layers.Dropout(0.5),
        tf.keras.layers.Dense(10, activation='softmax')
    ])

    model.compile(
        optimizer=tf.keras.optimizers.Adam(),
        loss='sparse_categorical_crossentropy',
        metrics=['accuracy']
    )

    return model

if __name__ == "__main__":
    # Load MNIST
    (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
    x_train = x_train.astype('float32') / 255.0
    x_test = x_test.astype('float32') / 255.0
    x_train = x_train.reshape(-1, 28, 28, 1)
    x_test = x_test.reshape(-1, 28, 28, 1)

    num_clients = 5
    partitioner = DataPartitioner()

    print("\n1. IID PARTITIONING")
    print("="*60)
    iid_partitions = partitioner.partition_iid(x_train, y_train, num_clients)
    partitioner.analyze_partitions(iid_partitions)

    print("\n\n2. NON-IID PARTITIONING (Label Skew)")
    print("="*60)
    non_iid_label = partitioner.partition_non_iid_label_skew(
        x_train, y_train, num_clients, classes_per_client=2
    )
    partitioner.analyze_partitions(non_iid_label)

    print("\n\n3. NON-IID PARTITIONING (Quantity Skew)")
    print("="*60)
    non_iid_quantity = partitioner.partition_non_iid_quantity_skew(
        x_train, y_train, num_clients, min_samples_ratio=0.2
    )
    partitioner.analyze_partitions(non_iid_quantity)
```

**Expected Outcomes:**
- IID partitions have similar class distributions across all clients
- Non-IID label skew creates clients specialized in specific classes
- Non-IID quantity skew creates highly imbalanced dataset sizes
- Training with non-IID data shows slower convergence and lower final accuracy
- Per-class accuracy metrics reveal which clients struggle with which classes
:::

### Exercise 3: Non-IID Data Simulation and Convergence Analysis
**Difficulty:** Level 2
**Requirements:** Python, Flower, TensorFlow, Matplotlib

**Task:**
Create comprehensive non-IID data scenarios and compare convergence behavior against IID baselines. Analyze accuracy degradation, per-client performance, and the impact of data heterogeneity on the global model.

Steps:
1. Implement multiple non-IID scenarios (label distribution skew, feature distribution skew)
2. Run federated training with IID and non-IID data
3. Track per-round global and local metrics
4. Visualize convergence curves and accuracy degradation
5. Calculate statistical measures of data heterogeneity

::: {.callout-tip collapse="true"}
## Solution
```python
# non_iid_analysis.py
import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf
import flwr as fl
from typing import List, Dict
import json
from datetime import datetime

class NonIIDAnalyzer:
    """Analyze non-IID effects in federated learning."""

    def __init__(self, num_clients: int = 5):
        self.num_clients = num_clients
        self.results = {
            'iid': {'rounds': [], 'accuracy': [], 'loss': [], 'client_metrics': []},
            'non_iid': {'rounds': [], 'accuracy': [], 'loss': [], 'client_metrics': []}
        }

    def calculate_distribution_divergence(self, partitions: List[tuple]) -> float:
        """
        Calculate Jensen-Shannon divergence between client distributions.
        Measures how different client data distributions are from each other.
        """
        num_classes = 10

        # Get class distributions for each client
        distributions = []
        for x, y in partitions:
            dist = np.bincount(y.astype(int), minlength=num_classes)
            dist = dist / np.sum(dist)  # Normalize
            distributions.append(dist)

        # Calculate average distribution
        avg_dist = np.mean(distributions, axis=0)

        # Calculate JS divergence
        def kl_divergence(p, q):
            return np.sum(p * np.log((p + 1e-10) / (q + 1e-10)))

        js_divergences = []
        for dist in distributions:
            m = 0.5 * (dist + avg_dist)
            js = 0.5 * kl_divergence(dist, m) + 0.5 * kl_divergence(avg_dist, m)
            js_divergences.append(js)

        return float(np.mean(js_divergences))

    def run_experiment(self, partitions: List[tuple], experiment_name: str,
                       num_rounds: int = 15):
        """Run federated learning experiment and collect metrics."""
        print(f"\n{'='*60}")
        print(f"Running experiment: {experiment_name}")
        print(f"{'='*60}\n")

        # Calculate data heterogeneity
        divergence = self.calculate_distribution_divergence(partitions)
        print(f"Data divergence (JS): {divergence:.4f}")
        print(f"  (0 = identical distributions, higher = more heterogeneous)\n")

        # This would integrate with actual Flower training
        # For demonstration, we'll simulate results
        results = self._simulate_training(partitions, num_rounds, divergence)

        self.results[experiment_name] = results
        return results

    def _simulate_training(self, partitions, num_rounds, divergence):
        """Simulate FL training (replace with actual Flower integration)."""
        results = {
            'rounds': list(range(1, num_rounds + 1)),
            'accuracy': [],
            'loss': [],
            'client_metrics': [],
            'divergence': divergence
        }

        # Simulate convergence based on data heterogeneity
        base_acc = 0.95
        heterogeneity_penalty = divergence * 0.15

        for round_num in range(1, num_rounds + 1):
            # Simulated convergence curve
            progress = 1 - np.exp(-round_num / 3)
            accuracy = (base_acc - heterogeneity_penalty) * progress
            loss = 0.5 * np.exp(-round_num / 5) + heterogeneity_penalty * 0.1

            # Add some noise
            accuracy += np.random.normal(0, 0.01)
            loss += np.random.normal(0, 0.02)

            results['accuracy'].append(accuracy)
            results['loss'].append(loss)

            # Simulate per-client metrics
            client_accs = []
            for i in range(len(partitions)):
                # Clients with less data perform worse
                client_size_factor = len(partitions[i][0]) / np.mean([len(p[0]) for p in partitions])
                client_acc = accuracy * (0.9 + 0.1 * client_size_factor) + np.random.normal(0, 0.02)
                client_accs.append(client_acc)

            results['client_metrics'].append(client_accs)

        return results

    def plot_comparison(self):
        """Plot IID vs Non-IID comparison."""
        fig, axes = plt.subplots(2, 2, figsize=(14, 10))

        # Accuracy over rounds
        ax = axes[0, 0]
        for exp_name, data in self.results.items():
            if data['rounds']:
                ax.plot(data['rounds'], data['accuracy'],
                       marker='o', label=exp_name.upper(), linewidth=2)
        ax.set_xlabel('Round')
        ax.set_ylabel('Global Accuracy')
        ax.set_title('Convergence Comparison')
        ax.legend()
        ax.grid(True, alpha=0.3)

        # Loss over rounds
        ax = axes[0, 1]
        for exp_name, data in self.results.items():
            if data['rounds']:
                ax.plot(data['rounds'], data['loss'],
                       marker='s', label=exp_name.upper(), linewidth=2)
        ax.set_xlabel('Round')
        ax.set_ylabel('Global Loss')
        ax.set_title('Loss Comparison')
        ax.legend()
        ax.grid(True, alpha=0.3)

        # Final accuracy comparison
        ax = axes[1, 0]
        final_accs = {name: data['accuracy'][-1]
                     for name, data in self.results.items() if data['accuracy']}
        bars = ax.bar(final_accs.keys(), final_accs.values(),
                     color=['#2ecc71', '#e74c3c'])
        ax.set_ylabel('Final Accuracy')
        ax.set_title('Final Model Performance')
        ax.set_ylim(0, 1)

        # Add value labels on bars
        for bar in bars:
            height = bar.get_height()
            ax.text(bar.get_x() + bar.get_width()/2., height,
                   f'{height:.2%}', ha='center', va='bottom')

        # Client performance variance
        ax = axes[1, 1]
        for exp_name, data in self.results.items():
            if data['client_metrics']:
                # Calculate std dev of client accuracies per round
                variances = [np.std(client_accs) for client_accs in data['client_metrics']]
                ax.plot(data['rounds'], variances,
                       marker='d', label=exp_name.upper(), linewidth=2)
        ax.set_xlabel('Round')
        ax.set_ylabel('Client Accuracy Std Dev')
        ax.set_title('Client Performance Variance')
        ax.legend()
        ax.grid(True, alpha=0.3)

        plt.tight_layout()
        plt.savefig('non_iid_comparison.png', dpi=150, bbox_inches='tight')
        print("Saved plot to: non_iid_comparison.png")
        plt.show()

    def generate_report(self):
        """Generate detailed analysis report."""
        print("\n" + "="*60)
        print("NON-IID ANALYSIS REPORT")
        print("="*60)

        for exp_name, data in self.results.items():
            if not data['rounds']:
                continue

            print(f"\n{exp_name.upper()}:")
            print(f"  Data divergence: {data.get('divergence', 0):.4f}")
            print(f"  Initial accuracy: {data['accuracy'][0]:.2%}")
            print(f"  Final accuracy: {data['accuracy'][-1]:.2%}")
            print(f"  Accuracy improvement: {data['accuracy'][-1] - data['accuracy'][0]:.2%}")
            print(f"  Final loss: {data['loss'][-1]:.4f}")
            print(f"  Rounds to 90% accuracy: ", end="")

            rounds_to_90 = next((i+1 for i, acc in enumerate(data['accuracy'])
                               if acc >= 0.90), "Not reached")
            print(rounds_to_90)

            if data['client_metrics']:
                final_client_accs = data['client_metrics'][-1]
                print(f"  Client accuracy range: {min(final_client_accs):.2%} - {max(final_client_accs):.2%}")
                print(f"  Client accuracy std: {np.std(final_client_accs):.4f}")

        # Calculate degradation
        if 'iid' in self.results and 'non_iid' in self.results:
            if self.results['iid']['accuracy'] and self.results['non_iid']['accuracy']:
                iid_final = self.results['iid']['accuracy'][-1]
                non_iid_final = self.results['non_iid']['accuracy'][-1]
                degradation = (iid_final - non_iid_final) / iid_final * 100

                print(f"\nACCURACY DEGRADATION:")
                print(f"  IID final accuracy: {iid_final:.2%}")
                print(f"  Non-IID final accuracy: {non_iid_final:.2%}")
                print(f"  Relative degradation: {degradation:.1f}%")

        print()

# main script
if __name__ == "__main__":
    from data_partitioner import DataPartitioner

    # Load MNIST
    (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
    x_train = x_train.astype('float32') / 255.0
    x_train = x_train.reshape(-1, 28, 28, 1)

    num_clients = 5
    partitioner = DataPartitioner()
    analyzer = NonIIDAnalyzer(num_clients)

    # Create partitions
    print("Creating data partitions...")
    iid_partitions = partitioner.partition_iid(x_train, y_train, num_clients)
    non_iid_partitions = partitioner.partition_non_iid_label_skew(
        x_train, y_train, num_clients, classes_per_client=2
    )

    # Run experiments
    analyzer.run_experiment(iid_partitions, 'iid', num_rounds=15)
    analyzer.run_experiment(non_iid_partitions, 'non_iid', num_rounds=15)

    # Generate visualizations and report
    analyzer.plot_comparison()
    analyzer.generate_report()

    # Save results
    with open('non_iid_results.json', 'w') as f:
        json.dump(analyzer.results, f, indent=2)
    print("\nResults saved to: non_iid_results.json")
```

**Expected Outcomes:**
- IID training converges faster and reaches higher accuracy (95%+)
- Non-IID training shows slower convergence and accuracy degradation (5-10% lower)
- Client performance variance is higher with non-IID data
- Visualization clearly shows the impact of data heterogeneity
- JS divergence metric quantifies the degree of non-IID-ness
:::

### Exercise 4: Privacy-Preserving Federated Learning
**Difficulty:** Level 2
**Requirements:** Python, Flower, TensorFlow, NumPy

**Task:**
Implement privacy-preserving techniques in federated learning including differential privacy noise addition and secure aggregation concepts. Measure the privacy-utility tradeoff by comparing model accuracy with different privacy budgets.

Steps:
1. Implement gradient clipping and Gaussian noise addition (DP-SGD)
2. Add differential privacy to client updates
3. Implement secure aggregation simulation (encrypted weight aggregation)
4. Test with various privacy budgets (epsilon values)
5. Analyze privacy-utility tradeoff curves

::: {.callout-tip collapse="true"}
## Solution
```python
# privacy_fl.py
import numpy as np
import tensorflow as tf
import flwr as fl
from typing import List, Tuple
import matplotlib.pyplot as plt

class DifferentialPrivacy:
    """Differential Privacy utilities for Federated Learning."""

    def __init__(self, epsilon: float = 1.0, delta: float = 1e-5,
                 sensitivity: float = 1.0):
        """
        Initialize DP parameters.

        Args:
            epsilon: Privacy budget (lower = more privacy, less utility)
            delta: Probability of privacy breach
            sensitivity: L2 sensitivity of the query (max gradient norm)
        """
        self.epsilon = epsilon
        self.delta = delta
        self.sensitivity = sensitivity
        self.noise_scale = self._calculate_noise_scale()

    def _calculate_noise_scale(self) -> float:
        """Calculate Gaussian noise scale for (ε, δ)-DP."""
        # For Gaussian mechanism: σ ≥ sensitivity * sqrt(2 * ln(1.25/δ)) / ε
        return self.sensitivity * np.sqrt(2 * np.log(1.25 / self.delta)) / self.epsilon

    def clip_gradients(self, gradients: List[np.ndarray],
                       max_norm: float) -> List[np.ndarray]:
        """
        Clip gradients to bound sensitivity.

        Ensures ||gradients||₂ ≤ max_norm
        """
        # Calculate global L2 norm
        global_norm = np.sqrt(sum(np.sum(g**2) for g in gradients))

        # Clip if necessary
        if global_norm > max_norm:
            clip_factor = max_norm / global_norm
            clipped = [g * clip_factor for g in gradients]
            return clipped

        return gradients

    def add_noise(self, parameters: List[np.ndarray]) -> List[np.ndarray]:
        """Add calibrated Gaussian noise to parameters for DP."""
        noisy_params = []

        for param in parameters:
            noise = np.random.normal(0, self.noise_scale, param.shape)
            noisy_params.append(param + noise.astype(param.dtype))

        return noisy_params

    def get_privacy_spent(self, num_steps: int, batch_size: int,
                         dataset_size: int) -> Tuple[float, float]:
        """
        Calculate privacy budget spent using moments accountant.
        Simplified version - use opacus or tensorflow-privacy for production.
        """
        q = batch_size / dataset_size  # Sampling probability
        steps = num_steps

        # Simplified calculation (actual requires moments accountant)
        epsilon_spent = self.epsilon * steps * q

        return epsilon_spent, self.delta

class PrivateFlowerClient(fl.client.NumPyClient):
    """Flower client with differential privacy."""

    def __init__(self, client_id, model, x_train, y_train, x_test, y_test,
                 epsilon: float = 1.0, delta: float = 1e-5, max_grad_norm: float = 1.0):
        self.client_id = client_id
        self.model = model
        self.x_train = x_train
        self.y_train = y_train
        self.x_test = x_test
        self.y_test = y_test

        # DP parameters
        self.dp = DifferentialPrivacy(epsilon=epsilon, delta=delta,
                                     sensitivity=max_grad_norm)
        self.max_grad_norm = max_grad_norm
        self.privacy_spent = 0.0

    def get_parameters(self, config):
        return self.model.get_weights()

    def fit(self, parameters, config):
        """Train with differential privacy."""
        self.model.set_weights(parameters)

        local_epochs = config.get("local_epochs", 5)
        batch_size = config.get("batch_size", 32)

        print(f"\n[Client {self.client_id}] Training with DP")
        print(f"  Privacy budget: ε={self.dp.epsilon}, δ={self.dp.delta}")

        # Store initial weights
        initial_weights = self.model.get_weights()

        # Train normally
        self.model.fit(
            self.x_train, self.y_train,
            batch_size=batch_size,
            epochs=local_epochs,
            verbose=0
        )

        # Get weight updates (gradients)
        updated_weights = self.model.get_weights()
        weight_updates = [updated - initial
                         for updated, initial in zip(updated_weights, initial_weights)]

        # Apply gradient clipping
        clipped_updates = self.dp.clip_gradients(weight_updates, self.max_grad_norm)

        # Add DP noise
        noisy_updates = self.dp.add_noise(clipped_updates)

        # Apply noisy updates
        final_weights = [initial + noisy
                        for initial, noisy in zip(initial_weights, noisy_updates)]

        # Calculate privacy spent
        num_steps = (len(self.x_train) // batch_size) * local_epochs
        eps_spent, _ = self.dp.get_privacy_spent(num_steps, batch_size, len(self.x_train))
        self.privacy_spent += eps_spent

        print(f"  Privacy spent this round: ε={eps_spent:.3f}")
        print(f"  Total privacy spent: ε={self.privacy_spent:.3f}")

        # Evaluate with noisy model
        self.model.set_weights(final_weights)
        loss, acc = self.model.evaluate(self.x_train, self.y_train, verbose=0)

        metrics = {
            "accuracy": float(acc),
            "privacy_spent": float(self.privacy_spent),
            "epsilon": self.dp.epsilon
        }

        return final_weights, len(self.x_train), metrics

    def evaluate(self, parameters, config):
        self.model.set_weights(parameters)
        loss, accuracy = self.model.evaluate(self.x_test, self.y_test, verbose=0)
        return float(loss), len(self.x_test), {"accuracy": float(accuracy)}

class SecureAggregation:
    """
    Simulated secure aggregation for federated learning.
    In production, use cryptographic protocols (e.g., Prio, SecAgg).
    """

    @staticmethod
    def generate_masks(num_clients: int, param_shapes: List) -> List[List[np.ndarray]]:
        """
        Generate random masks for each client.
        In real SecAgg, these would be derived from shared secrets.
        """
        masks = []

        for i in range(num_clients):
            client_mask = []
            for shape in param_shapes:
                # Each client gets a random mask
                mask = np.random.randn(*shape)
                client_mask.append(mask)
            masks.append(client_mask)

        return masks

    @staticmethod
    def mask_parameters(parameters: List[np.ndarray],
                       masks: List[np.ndarray]) -> List[np.ndarray]:
        """Add mask to parameters before sending to server."""
        return [param + mask for param, mask in zip(parameters, masks)]

    @staticmethod
    def aggregate_masked(masked_params_list: List[List[np.ndarray]],
                        weights: List[int],
                        masks: List[List[np.ndarray]]) -> List[np.ndarray]:
        """
        Aggregate masked parameters and remove masks.
        Server sees sum of (param + mask) but masks cancel out.
        """
        total_weight = sum(weights)

        # Initialize aggregated parameters
        aggregated = [np.zeros_like(p) for p in masked_params_list[0]]

        # Weighted sum of masked parameters
        for masked_params, weight in zip(masked_params_list, weights):
            for i, param in enumerate(masked_params):
                aggregated[i] += param * (weight / total_weight)

        # Remove masks (they sum to zero by construction in real SecAgg)
        # Here we simulate: each client adds +mask_i, so sum of masks = 0
        total_mask = [np.zeros_like(p) for p in aggregated]
        for client_masks in masks:
            for i, mask in enumerate(client_masks):
                total_mask[i] += mask

        # Subtract total mask (should be ~0 if properly constructed)
        for i in range(len(aggregated)):
            aggregated[i] -= total_mask[i] / len(masks)

        return aggregated

# privacy_experiment.py
def run_privacy_experiment():
    """Compare FL with different privacy levels."""
    from data_partitioner import DataPartitioner

    # Load data
    (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
    x_train = x_train.astype('float32') / 255.0
    x_train = x_train.reshape(-1, 28, 28, 1)

    # Partition data
    partitioner = DataPartitioner()
    partitions = partitioner.partition_iid(x_train, y_train, num_clients=3)

    # Test different privacy budgets
    privacy_budgets = [0.1, 0.5, 1.0, 5.0, 100.0]  # 100.0 ≈ no privacy
    results = []

    for epsilon in privacy_budgets:
        print(f"\n{'='*60}")
        print(f"Testing with ε = {epsilon}")
        print(f"{'='*60}")

        # This would run actual FL training with PrivateFlowerClient
        # For demonstration, we simulate results

        # Privacy-utility tradeoff: lower epsilon = more noise = lower accuracy
        base_accuracy = 0.95
        privacy_penalty = 1.0 / (1.0 + epsilon)  # More privacy = more penalty
        final_accuracy = base_accuracy * (1 - 0.3 * privacy_penalty)

        results.append({
            'epsilon': epsilon,
            'accuracy': final_accuracy,
            'privacy_level': 'High' if epsilon < 1.0 else ('Medium' if epsilon < 5.0 else 'Low')
        })

        print(f"  Final accuracy: {final_accuracy:.2%}")
        print(f"  Privacy level: {results[-1]['privacy_level']}")

    # Plot privacy-utility tradeoff
    fig, ax = plt.subplots(figsize=(10, 6))

    epsilons = [r['epsilon'] for r in results]
    accuracies = [r['accuracy'] for r in results]

    ax.plot(epsilons, accuracies, 'bo-', linewidth=2, markersize=8)
    ax.set_xscale('log')
    ax.set_xlabel('Privacy Budget (ε)', fontsize=12)
    ax.set_ylabel('Model Accuracy', fontsize=12)
    ax.set_title('Privacy-Utility Tradeoff in Federated Learning', fontsize=14)
    ax.grid(True, alpha=0.3)

    # Add annotations
    for r in results:
        ax.annotate(f"ε={r['epsilon']}\n{r['accuracy']:.1%}",
                   xy=(r['epsilon'], r['accuracy']),
                   xytext=(10, -10), textcoords='offset points',
                   fontsize=9, alpha=0.7)

    # Add privacy zones
    ax.axvspan(0, 1, alpha=0.1, color='green', label='High Privacy')
    ax.axvspan(1, 5, alpha=0.1, color='yellow', label='Medium Privacy')
    ax.axvspan(5, 100, alpha=0.1, color='red', label='Low Privacy')

    ax.legend()
    plt.tight_layout()
    plt.savefig('privacy_utility_tradeoff.png', dpi=150)
    print("\nSaved plot to: privacy_utility_tradeoff.png")
    plt.show()

    # Summary table
    print("\n" + "="*60)
    print("PRIVACY-UTILITY SUMMARY")
    print("="*60)
    print(f"{'Epsilon':<10} {'Accuracy':<12} {'Privacy Level':<15} {'Recommendation'}")
    print("-"*60)

    for r in results:
        if r['epsilon'] <= 1.0:
            rec = "Medical/Financial"
        elif r['epsilon'] <= 5.0:
            rec = "General Apps"
        else:
            rec = "Low-risk Apps"

        print(f"{r['epsilon']:<10.1f} {r['accuracy']:<12.1%} {r['privacy_level']:<15} {rec}")

if __name__ == "__main__":
    run_privacy_experiment()
```

**Expected Outcomes:**
- Lower epsilon (more privacy) results in lower model accuracy
- Gradient clipping bounds sensitivity and stabilizes training
- DP noise addition provides formal privacy guarantees
- Privacy-utility tradeoff curve shows accuracy degradation with privacy
- Secure aggregation prevents server from seeing individual client updates
:::

### Exercise 5: Multi-Device Federated Learning Deployment
**Difficulty:** Level 2
**Requirements:** Docker, Docker Compose, Python, Flower, network access

**Task:**
Deploy a federated learning system across multiple machines or containers simulating real edge deployment. Handle network configuration, client discovery, dropout scenarios, and asynchronous participation.

Steps:
1. Create Docker containers for server and multiple clients
2. Configure networking for FL communication
3. Implement client discovery and registration
4. Handle client dropout and reconnection
5. Test with varying network conditions (latency, bandwidth limits)

::: {.callout-tip collapse="true"}
## Solution
```yaml
# docker-compose.yml
version: '3.8'

services:
  flower-server:
    build:
      context: .
      dockerfile: Dockerfile.server
    container_name: fl-server
    ports:
      - "8080:8080"
    networks:
      - fl-network
    environment:
      - NUM_ROUNDS=10
      - MIN_CLIENTS=3
    volumes:
      - ./server_logs:/app/logs
      - ./models:/app/models
    command: python server.py

  flower-client-1:
    build:
      context: .
      dockerfile: Dockerfile.client
    container_name: fl-client-1
    depends_on:
      - flower-server
    networks:
      - fl-network
    environment:
      - CLIENT_ID=0
      - SERVER_ADDRESS=flower-server:8080
      - DATA_PARTITION=0
    volumes:
      - ./client_data:/app/data
    command: python client.py --client-id 0 --server flower-server:8080

  flower-client-2:
    build:
      context: .
      dockerfile: Dockerfile.client
    container_name: fl-client-2
    depends_on:
      - flower-server
    networks:
      - fl-network
    environment:
      - CLIENT_ID=1
      - SERVER_ADDRESS=flower-server:8080
      - DATA_PARTITION=1
    volumes:
      - ./client_data:/app/data
    command: python client.py --client-id 1 --server flower-server:8080

  flower-client-3:
    build:
      context: .
      dockerfile: Dockerfile.client
    container_name: fl-client-3
    depends_on:
      - flower-server
    networks:
      - fl-network
    environment:
      - CLIENT_ID=2
      - SERVER_ADDRESS=flower-server:8080
      - DATA_PARTITION=2
    volumes:
      - ./client_data:/app/data
    command: python client.py --client-id 2 --server flower-server:8080

networks:
  fl-network:
    driver: bridge
    ipam:
      config:
        - subnet: 172.25.0.0/16
```

```dockerfile
# Dockerfile.server
FROM python:3.9-slim

WORKDIR /app

# Install dependencies
RUN pip install --no-cache-dir flwr tensorflow numpy

# Copy server code
COPY server_robust.py server.py
COPY common.py .

# Expose FL port
EXPOSE 8080

CMD ["python", "server.py"]
```

```dockerfile
# Dockerfile.client
FROM python:3.9-slim

WORKDIR /app

# Install dependencies
RUN pip install --no-cache-dir flwr tensorflow numpy

# Copy client code
COPY client_robust.py client.py
COPY common.py .
COPY data_partitioner.py .

CMD ["python", "client.py"]
```

```python
# server_robust.py
import flwr as fl
from typing import List, Tuple, Dict, Optional
from flwr.common import Metrics, FitRes, Parameters
from flwr.server.client_proxy import ClientProxy
import logging
import time
import os

# Setup logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('/app/logs/server.log'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger("FlowerServer")

class RobustFedAvg(fl.server.strategy.FedAvg):
    """FedAvg with robust client handling and monitoring."""

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.client_history = {}
        self.round_start_time = None

    def initialize_parameters(self, client_manager):
        """Initialize global model parameters."""
        logger.info("Initializing global model parameters")
        return super().initialize_parameters(client_manager)

    def configure_fit(self, server_round: int, parameters: Parameters,
                     client_manager) -> List[Tuple[ClientProxy, FitRes]]:
        """Configure clients for training round."""
        self.round_start_time = time.time()
        logger.info(f"Round {server_round}: Configuring fit")

        # Sample clients
        sample_size, min_num_clients = self.num_fit_clients(
            client_manager.num_available()
        )

        logger.info(f"  Available clients: {client_manager.num_available()}")
        logger.info(f"  Sampling {sample_size} clients (min: {min_num_clients})")

        clients = client_manager.sample(
            num_clients=sample_size,
            min_num_clients=min_num_clients
        )

        # Log client selection
        for client in clients:
            cid = client.cid
            if cid not in self.client_history:
                self.client_history[cid] = {
                    'rounds_participated': 0,
                    'total_samples': 0,
                    'failures': 0
                }
            self.client_history[cid]['rounds_participated'] += 1

        return super().configure_fit(server_round, parameters, client_manager)

    def aggregate_fit(self, server_round, results, failures):
        """Aggregate training results from clients."""
        round_duration = time.time() - self.round_start_time

        logger.info(f"Round {server_round}: Aggregating results")
        logger.info(f"  Successful clients: {len(results)}")
        logger.info(f"  Failed clients: {len(failures)}")
        logger.info(f"  Round duration: {round_duration:.2f}s")

        # Log failures
        for failure in failures:
            logger.warning(f"  Client failure: {failure}")

        # Track client statistics
        for client, fit_res in results:
            cid = client.cid
            self.client_history[cid]['total_samples'] += fit_res.num_examples

        # Perform aggregation
        aggregated = super().aggregate_fit(server_round, results, failures)

        return aggregated

    def evaluate(self, server_round: int, parameters: Parameters):
        """Evaluate global model."""
        logger.info(f"Round {server_round}: Server-side evaluation")
        return super().evaluate(server_round, parameters)

    def aggregate_evaluate(self, server_round, results, failures):
        """Aggregate evaluation results."""
        logger.info(f"Round {server_round}: Aggregating evaluation")
        logger.info(f"  Evaluated on {len(results)} clients")

        aggregated = super().aggregate_evaluate(server_round, results, failures)

        if aggregated:
            loss, metrics = aggregated
            logger.info(f"  Global loss: {loss:.4f}")
            if 'accuracy' in metrics:
                logger.info(f"  Global accuracy: {metrics['accuracy']:.2%}")

        return aggregated

def weighted_average(metrics: List[Tuple[int, Metrics]]) -> Metrics:
    """Aggregate evaluation metrics."""
    accuracies = [num_examples * m["accuracy"] for num_examples, m in metrics]
    examples = [num_examples for num_examples, m in metrics]
    return {"accuracy": sum(accuracies) / sum(examples)}

def main():
    """Start Flower server with robust configuration."""
    # Configuration from environment
    num_rounds = int(os.getenv("NUM_ROUNDS", 10))
    min_clients = int(os.getenv("MIN_CLIENTS", 3))

    logger.info("="*60)
    logger.info("STARTING FEDERATED LEARNING SERVER")
    logger.info("="*60)
    logger.info(f"Configuration:")
    logger.info(f"  Rounds: {num_rounds}")
    logger.info(f"  Minimum clients: {min_clients}")
    logger.info(f"  Server address: 0.0.0.0:8080")

    # Define strategy
    strategy = RobustFedAvg(
        fraction_fit=1.0,
        fraction_evaluate=1.0,
        min_fit_clients=min_clients,
        min_evaluate_clients=min_clients,
        min_available_clients=min_clients,
        evaluate_metrics_aggregation_fn=weighted_average,
    )

    # Start server
    try:
        fl.server.start_server(
            server_address="0.0.0.0:8080",
            config=fl.server.ServerConfig(num_rounds=num_rounds),
            strategy=strategy,
        )

        logger.info("\n" + "="*60)
        logger.info("TRAINING COMPLETED SUCCESSFULLY")
        logger.info("="*60)

        # Log client statistics
        logger.info("\nClient Statistics:")
        for cid, stats in strategy.client_history.items():
            logger.info(f"  Client {cid}:")
            logger.info(f"    Rounds participated: {stats['rounds_participated']}")
            logger.info(f"    Total samples: {stats['total_samples']}")

    except Exception as e:
        logger.error(f"Server error: {e}", exc_info=True)
        raise

if __name__ == "__main__":
    main()
```

```python
# client_robust.py
import flwr as fl
import tensorflow as tf
import numpy as np
import argparse
import logging
import time
import os
from typing import Optional

# Setup logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)

class RobustFlowerClient(fl.client.NumPyClient):
    """Flower client with connection handling and retry logic."""

    def __init__(self, client_id, model, x_train, y_train, x_test, y_test):
        self.client_id = client_id
        self.model = model
        self.x_train = x_train
        self.y_train = y_train
        self.x_test = x_test
        self.y_test = y_test
        self.logger = logging.getLogger(f"Client-{client_id}")
        self.round_num = 0

    def get_parameters(self, config):
        return self.model.get_weights()

    def fit(self, parameters, config):
        """Train with error handling."""
        self.round_num += 1
        self.logger.info(f"Round {self.round_num}: Starting training")

        try:
            self.model.set_weights(parameters)

            # Simulate occasional client issues
            if np.random.random() < 0.05:  # 5% chance
                self.logger.warning("Simulating temporary compute issue...")
                time.sleep(2)

            # Train
            history = self.model.fit(
                self.x_train, self.y_train,
                batch_size=32,
                epochs=config.get("local_epochs", 5),
                verbose=0
            )

            self.logger.info(f"  Training complete - "
                           f"Acc: {history.history['accuracy'][-1]:.2%}")

            return (
                self.model.get_weights(),
                len(self.x_train),
                {"accuracy": float(history.history['accuracy'][-1])}
            )

        except Exception as e:
            self.logger.error(f"Training error: {e}")
            raise

    def evaluate(self, parameters, config):
        """Evaluate with error handling."""
        try:
            self.model.set_weights(parameters)
            loss, accuracy = self.model.evaluate(
                self.x_test, self.y_test, verbose=0
            )

            self.logger.info(f"  Evaluation - Acc: {accuracy:.2%}, Loss: {loss:.4f}")

            return float(loss), len(self.x_test), {"accuracy": float(accuracy)}

        except Exception as e:
            self.logger.error(f"Evaluation error: {e}")
            raise

def connect_with_retry(server_address: str, client, max_retries: int = 5):
    """Connect to server with retry logic."""
    logger = logging.getLogger("ClientConnector")

    for attempt in range(max_retries):
        try:
            logger.info(f"Connection attempt {attempt + 1}/{max_retries} to {server_address}")
            fl.client.start_numpy_client(
                server_address=server_address,
                client=client
            )
            return  # Success

        except Exception as e:
            logger.warning(f"Connection failed: {e}")
            if attempt < max_retries - 1:
                wait_time = 2 ** attempt  # Exponential backoff
                logger.info(f"Retrying in {wait_time}s...")
                time.sleep(wait_time)
            else:
                logger.error("Max retries reached. Giving up.")
                raise

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--client-id", type=int, required=True)
    parser.add_argument("--server", type=str, default="localhost:8080")
    args = parser.parse_args()

    logger = logging.getLogger(f"Client-{args.client_id}")
    logger.info("="*60)
    logger.info(f"STARTING FEDERATED LEARNING CLIENT {args.client_id}")
    logger.info("="*60)

    # Load data
    from data_partitioner import DataPartitioner
    (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
    x_train = x_train.astype('float32') / 255.0
    x_test = x_test.astype('float32') / 255.0
    x_train = x_train.reshape(-1, 28, 28, 1)
    x_test = x_test.reshape(-1, 28, 28, 1)

    # Partition
    partitions = DataPartitioner.partition_iid(x_train, y_train, num_clients=5)
    x_train, y_train = partitions[args.client_id]

    logger.info(f"Data loaded: {len(x_train)} training samples")

    # Create model
    model = tf.keras.Sequential([
        tf.keras.layers.Conv2D(32, (3, 3), activation='relu', input_shape=(28, 28, 1)),
        tf.keras.layers.MaxPooling2D((2, 2)),
        tf.keras.layers.Flatten(),
        tf.keras.layers.Dense(128, activation='relu'),
        tf.keras.layers.Dense(10, activation='softmax')
    ])
    model.compile(optimizer='adam', loss='sparse_categorical_crossentropy',
                  metrics=['accuracy'])

    # Start client with retry logic
    client = RobustFlowerClient(args.client_id, model, x_train, y_train,
                                x_test[:1000], y_test[:1000])
    connect_with_retry(args.server, client, max_retries=5)

if __name__ == "__main__":
    main()
```

Run the deployment:
```bash
# Build and start all containers
docker-compose up --build

# Scale clients dynamically
docker-compose up --scale flower-client-2=3

# Simulate client dropout
docker-compose stop fl-client-2
# ... training continues with remaining clients

# Restart dropped client
docker-compose start fl-client-2

# View logs
docker-compose logs -f flower-server
docker-compose logs -f flower-client-1

# Cleanup
docker-compose down
```

**Network testing with tc (traffic control):**
```bash
# Add latency to simulate slow network
docker exec fl-client-1 tc qdisc add dev eth0 root netem delay 100ms

# Limit bandwidth
docker exec fl-client-1 tc qdisc add dev eth0 root tbf rate 1mbit burst 32kbit latency 400ms

# Add packet loss
docker exec fl-client-1 tc qdisc add dev eth0 root netem loss 5%

# Remove limitations
docker exec fl-client-1 tc qdisc del dev eth0 root
```

**Expected Outcomes:**
- Federated learning runs across isolated containers
- Server handles client dropouts gracefully (continues with minimum clients)
- Clients reconnect automatically with exponential backoff
- Network conditions affect round duration but not correctness
- Logs provide full visibility into distributed training process
- System scales horizontally by adding more client containers
:::

## Summary

These Level 2 exercises provide hands-on experience with production-ready federated learning deployments. You've learned to:

1. Configure Flower servers with proper client selection and aggregation strategies
2. Implement robust FL clients with comprehensive data partitioning
3. Analyze the impact of non-IID data on convergence and accuracy
4. Add privacy-preserving techniques with quantifiable privacy budgets
5. Deploy FL systems across multiple machines/containers with fault tolerance

These skills prepare you for Level 3 deployment on real edge devices like Raspberry Pi clusters.
