federated/client.py

"""Federated Learning Client for Myanmar Ghost project.

Enables training on distributed data (e.g., hospitals, restaurants)
without centralizing sensitive data.
"""

import flwr as fl
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple
import yaml

logger = __import__("loguru").logger


class FederatedClient(fl.client.NumPyClient):
    """Flower client for federated learning."""
    
    def __init__(
        self,
        model: nn.Module,
        trainloader,
        valloader,
        device: str = "cuda" if torch.cuda.is_available() else "cpu",
        client_id: str = "client_1",
        output_dir: str = "outputs/federated",
    ):
        self.model = model.to(device)
        self.trainloader = trainloader
        self.valloader = valloader
        self.device = device
        self.client_id = client_id
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)
        
        self.criterion = nn.CrossEntropyLoss()
        self.optimizer = optim.AdamW(self.model.parameters(), lr=1e-4)
    
    def get_parameters(self) -> List[np.ndarray]:
        """Get model parameters as numpy arrays."""
        return [val.cpu().numpy() for _, val in self.model.state_dict().items()]
    
    def set_parameters(self, parameters: List[np.ndarray]) -> None:
        """Set model parameters from numpy arrays."""
        state_dict = dict(self.model.state_dict())
        for i, (key, _) in enumerate(state_dict.items()):
            state_dict[key] = torch.from_numpy(parameters[i])
        self.model.load_state_dict(state_dict)
    
    def fit(
        self,
        parameters: List[np.ndarray],
        config: Dict[str, Any],
    ) -> Tuple[List[np.ndarray], int, Dict]:
        """Train model on local data.
        
        Args:
            parameters: Global model parameters
            config: Training configuration
        
        Returns:
            Updated parameters, number of samples, metrics
        """
        # Set global parameters
        self.set_parameters(parameters)
        
        # Training configuration
        epochs = config.get("local_epochs", 1)
        batch_size = config.get("batch_size", 32)
        learning_rate = config.get("learning_rate", 1e-4)
        
        self.optimizer = optim.AdamW(
            self.model.parameters(),
            lr=learning_rate,
        )
        
        # Local training
        self.model.train()
        total_loss = 0.0
        total_samples = 0
        
        for epoch in range(epochs):
            epoch_loss = 0.0
            epoch_samples = 0
            
            for batch_idx, (inputs, labels) in enumerate(self.trainloader):
                inputs = inputs.to(self.device)
                labels = labels.to(self.device)
                
                self.optimizer.zero_grad()
                outputs = self.model(inputs)
                loss = self.criterion(outputs, labels)
                loss.backward()
                self.optimizer.step()
                
                epoch_loss += loss.item() * inputs.size(0)
                epoch_samples += inputs.size(0)
            
            total_loss += epoch_loss
            total_samples += epoch_samples
            
            logger.info(
                f"Client {self.client_id} - Epoch {epoch+1}/{epochs}: "
                f"Loss={epoch_loss/epoch_samples:.4f}"
            )
        
        # Save checkpoint
        self._save_checkpoint(epochs)
        
        metrics = {
            "loss": total_loss / total_samples,
            "samples": total_samples,
            "epochs": epochs,
        }
        
        return self.get_parameters(), total_samples, metrics
    
    def evaluate(
        self,
        parameters: List[np.ndarray],
        config: Dict[str, Any],
    ) -> Tuple[float, int, Dict]:
        """Evaluate model on local validation data.
        
        Args:
            parameters: Model parameters
            config: Evaluation configuration
        
        Returns:
            Loss, number of samples, metrics
        """
        self.set_parameters(parameters)
        
        self.model.eval()
        total_loss = 0.0
        total_correct = 0
        total_samples = 0
        
        with torch.no_grad():
            for inputs, labels in self.valloader:
                inputs = inputs.to(self.device)
                labels = labels.to(self.device)
                
                outputs = self.model(inputs)
                loss = self.criterion(outputs, labels)
                
                total_loss += loss.item() * inputs.size(0)
                _, predicted = outputs.max(1)
                total_correct += predicted.eq(labels).sum().item()
                total_samples += inputs.size(0)
        
        accuracy = total_correct / total_samples if total_samples > 0 else 0.0
        
        logger.info(
            f"Client {self.client_id} - Evaluation: "
            f"Loss={total_loss/total_samples:.4f}, Accuracy={accuracy:.4f}"
        )
        
        metrics = {
            "loss": total_loss / total_samples,
            "accuracy": accuracy,
            "samples": total_samples,
        }
        
        return total_loss / total_samples, total_samples, metrics
    
    def _save_checkpoint(self, epochs: int) -> None:
        """Save model checkpoint."""
        path = self.output_dir / f"{self.client_id}_checkpoint.pt"
        torch.save({
            "model_state_dict": self.model.state_dict(),
            "optimizer_state_dict": self.optimizer.state_dict(),
            "epochs": epochs,
            "client_id": self.client_id,
        }, path)
        logger.info(f"Checkpoint saved to {path}")


def load_client_config(config_path: str) -> Dict:
    """Load client configuration from YAML."""
    with open(config_path, "r", encoding="utf-8") as f:
        return yaml.safe_load(f)


class ClientFactory:
    """Factory for creating federated clients."""
    
    def __init__(
        self,
        model_fn,
        data_dir: str,
        output_dir: str = "outputs/federated",
    ):
        self.model_fn = model_fn
        self.data_dir = Path(data_dir)
        self.output_dir = Path(output_dir)
    
    def create_client(
        self,
        client_id: str,
        config: Dict,
    ) -> FederatedClient:
        """Create a client for the given configuration."""
        from torch.utils.data import DataLoader
        
        # Load data
        train_data = self._load_partition(
            client_id,
            config.get("train_file", f"{client_id}_train.pt"),
        )
        val_data = self._load_partition(
            client_id,
            config.get("val_file", f"{client_id}_val.pt"),
        )
        
        trainloader = DataLoader(
            train_data,
            batch_size=config.get("batch_size", 32),
            shuffle=True,
        )
        valloader = DataLoader(
            val_data,
            batch_size=config.get("batch_size", 32),
            shuffle=False,
        )
        
        model = self.model_fn()
        
        return FederatedClient(
            model=model,
            trainloader=trainloader,
            valloader=valloader,
            device=config.get("device", "cuda"),
            client_id=client_id,
            output_dir=str(self.output_dir),
        )
    
    def _load_partition(self, client_id: str, filename: str):
        """Load data partition for a client."""
        path = self.data_dir / client_id / filename
        if path.exists():
            return torch.load(path)
        raise FileNotFoundError(f"Data partition not found: {path}")


def start_client(
    model_fn,
    data_dir: str,
    client_id: str,
    server_address: str = "localhost:8080",
    config_path: Optional[str] = None,
) -> None:
    """Start a federated learning client."""
    config = {}
    if config_path:
        config = load_client_config(config_path)
    
    factory = ClientFactory(model_fn, data_dir)
    client = factory.create_client(client_id, config)
    
    app = fl.client.start_numpy_client(
        server_address=server_address,
        client=client,
    )


if __name__ == "__main__":
    print("FederatedClient module loaded")
    print("Use start_client() to start a federated learning client")