diff --git a/MEDfl/LearningManager/client.py b/MEDfl/LearningManager/client.py index 4a18ea8..4b6d631 100644 --- a/MEDfl/LearningManager/client.py +++ b/MEDfl/LearningManager/client.py @@ -1,4 +1,7 @@ #!/usr/bin/env python3 +import logging +import sys + import flwr as fl from opacus import PrivacyEngine from torch.utils.data import DataLoader @@ -6,52 +9,62 @@ from .model import Model from .utils import params import torch +import numpy as np + + +# --------------------------------------------------------------------------- +# Per-worker logger setup +# --------------------------------------------------------------------------- +# Ray workers are separate processes. We configure this logger inside the +# worker so that messages go to stdout — Ray then forwards stdout to the +# driver when log_to_driver=True is set in ray_init_args. +# --------------------------------------------------------------------------- +def _get_logger() -> logging.Logger: + logger = logging.getLogger("MEDfl.client") + if not logger.handlers: + handler = logging.StreamHandler(sys.stdout) + handler.setFormatter( + logging.Formatter( + "[%(asctime)s][%(name)s][%(levelname)s] %(message)s", + datefmt="%H:%M:%S", + ) + ) + logger.addHandler(handler) + logger.setLevel(logging.DEBUG) + logger.propagate = False # avoid double-printing on driver + return logger + class FlowerClient(fl.client.NumPyClient): """ FlowerClient class for creating MEDfl clients. - - Attributes: - cid (str): Client ID. - local_model (Model): Local model of the federated learning network. - trainloader (DataLoader): DataLoader for training data. - valloader (DataLoader): DataLoader for validation data. - diff_priv (bool): Flag indicating whether to use differential privacy. """ - def __init__(self, cid: str, local_model: Model, trainloader: DataLoader, valloader: DataLoader, diff_priv: bool = params["diff_privacy"]): - """ - Initializes the FlowerClient instance. - - Args: - cid (str): Client ID. - local_model (Model): Local model of the federated learning network. - trainloader (DataLoader): DataLoader for training data. - valloader (DataLoader): DataLoader for validation data. - diff_priv (bool): Flag indicating whether to use differential privacy. - """ + + def __init__( + self, + cid: str, + local_model: Model, + trainloader: DataLoader, + valloader: DataLoader, + diff_priv: bool = params["diff_privacy"], + ): self.cid = cid self.local_model = local_model self.trainloader = trainloader self.valloader = valloader - if torch.cuda.is_available(): - num_cuda_devices = torch.cuda.device_count() - if num_cuda_devices > 0: - device_idx = int(self.cid) % num_cuda_devices - self.device = torch.device(f"cuda:{device_idx}") - self.local_model.model.to(self.device) - else: - # Handle case where CUDA is available but no CUDA devices are found - raise RuntimeError("CUDA is available, but no CUDA devices are found.") - else: - # Handle case where CUDA is not available - self.device = torch.device("cpu") - self.local_model.model.to(self.device) + self.log = _get_logger() + + self.device = torch.device("cpu") + self.local_model.model.to(self.device) self.privacy_engine = PrivacyEngine(secure_mode=False) self.diff_priv = diff_priv self.epsilons = [] self.accuracies = [] self.losses = [] + self.train_round_losses = [] + self.train_round_metrics = [] + if self.diff_priv: model, optimizer, self.trainloader = self.privacy_engine.make_private_with_epsilon( module=self.local_model.model.train(), @@ -59,92 +72,98 @@ def __init__(self, cid: str, local_model: Model, trainloader: DataLoader, valloa data_loader=self.trainloader, epochs=params["train_epochs"], target_epsilon=float(params["EPSILON"]), - target_delta= float(params["DELTA"]), + target_delta=float(params["DELTA"]), max_grad_norm=params["MAX_GRAD_NORM"], ) setattr(self.local_model, "model", model) setattr(self.local_model, "optimizer", optimizer) self.validate() + # ------------------------------------------------------------------ def validate(self): """Validates cid, local_model, trainloader, valloader.""" if not isinstance(self.cid, str): raise TypeError("cid argument must be a string") - if not isinstance(self.local_model, Model): raise TypeError("local_model argument must be a MEDfl.LearningManager.model.Model") - if not isinstance(self.trainloader, DataLoader): raise TypeError("trainloader argument must be a torch.utils.data.dataloader") - if not isinstance(self.valloader, DataLoader): raise TypeError("valloader argument must be a torch.utils.data.dataloader") - if not isinstance(self.diff_priv, bool): raise TypeError("diff_priv argument must be a bool") + # ------------------------------------------------------------------ def get_parameters(self, config): - """ - Returns the current parameters of the local model. - - Args: - config: Configuration information. - - Returns: - Numpy array: Parameters of the local model. - """ - print(f"[Client {self.cid}] get_parameters") + self.log.info(f"[Client {self.cid}] get_parameters") return self.local_model.get_parameters() + # ------------------------------------------------------------------ def fit(self, parameters, config): - """ - Fits the local model to the received parameters using federated learning. - - Args: - parameters: Parameters received from the server. - config: Configuration information. - - Returns: - Tuple: Parameters of the local model, number of training examples, and privacy information. - """ - print('\n -------------------------------- \n this is the config of the client') - print(f"[Client {self.cid}] fit, config: {config}") - # print(config['epochs']) - print('\n -------------------------------- \n ') + self.log.info(f"[Client {self.cid}] fit | config: {config}") + self.local_model.set_parameters(parameters) - for _ in range(params["train_epochs"]): - epsilon = self.local_model.train( + + epoch_losses, epoch_metrics = [], [] + + for ep in range(params["train_epochs"]): + epsilon, tr_loss, tr_metric = self.local_model.train( self.trainloader, - epoch=_, + epoch=ep, device=self.device, privacy_engine=self.privacy_engine, diff_priv=self.diff_priv, ) self.epsilons.append(epsilon) - print(f"epsilon of client {self.cid} : eps = {epsilon}") + epoch_losses.append(tr_loss) + epoch_metrics.append(tr_metric) + + avg_grad_norm = float(np.nanmean(self.local_model.grad_norms)) \ + if hasattr(self.local_model, "grad_norms") and self.local_model.grad_norms \ + else float("nan") + self.local_model.grad_norms = [] # reset for next round + + + self.log.info( + f"[Client {self.cid}] fit done | " + f"epsilon={epsilon:.4f} | " + f"avg_loss={float(np.nanmean(epoch_losses)):.4f}" + ) + + round_train_loss = float(np.nanmean(epoch_losses)) if len(epoch_losses) else float("nan") + round_train_metric = float(np.nanmean(epoch_metrics)) if len(epoch_metrics) else float("nan") + + self.train_round_losses.append(round_train_loss) + self.train_round_metrics.append(round_train_metric) + + metric_key = "train_accuracy" if self.local_model.task_type == "binary" else "train_rmse" + return ( self.local_model.get_parameters(), - len(self.trainloader), - {"epsilon": epsilon}, + len(self.trainloader.dataset), + { + "epsilon": float(epsilon), + "train_loss": float(round_train_loss), + metric_key: float(round_train_metric), + "grad_norm": avg_grad_norm, # <-- NEW + }, ) + # ------------------------------------------------------------------ def evaluate(self, parameters, config): - """ - Evaluates the local model on the validation data and returns the loss and accuracy. + self.log.info(f"[Client {self.cid}] evaluate | config: {config}") + self.local_model.set_parameters(parameters) - Args: - parameters: Parameters received from the server. - config: Configuration information. + loss, metrics = self.local_model.evaluate(self.valloader, device=self.device) - Returns: - Tuple: Loss, number of validation examples, and accuracy information. - """ - print(f"[Client {self.cid}] evaluate, config: {config}") - self.local_model.set_parameters(parameters) - loss, accuracy , auc = self.local_model.evaluate( - self.valloader, device=self.device - ) self.losses.append(loss) - self.accuracies.append(accuracy) - - return float(loss), len(self.valloader), {"accuracy": float(accuracy)} + + if isinstance(metrics, dict) and "accuracy" in metrics: + self.accuracies.append(metrics["accuracy"]) + + self.log.info( + f"[Client {self.cid}] evaluate done | " + f"loss={loss:.4f} | metrics={metrics}" + ) + + return float(loss), len(self.valloader.dataset), {k: float(v) for k, v in (metrics or {}).items()} \ No newline at end of file diff --git a/MEDfl/LearningManager/federated_dataset.py b/MEDfl/LearningManager/federated_dataset.py index 4a451cc..c0a8acd 100644 --- a/MEDfl/LearningManager/federated_dataset.py +++ b/MEDfl/LearningManager/federated_dataset.py @@ -1,6 +1,8 @@ from MEDfl.NetManager.net_helper import * from MEDfl.NetManager.net_manager_queries import * from MEDfl.NetManager.database_connector import DatabaseManager +from typing import List, Optional + class FederatedDataset: def __init__( @@ -11,6 +13,7 @@ def __init__( trainloaders: list, valloaders: list, testloaders: list, + feature_names: Optional[List[str]] = None, ): """ Represents a Federated Dataset. @@ -30,6 +33,25 @@ def __init__( self.testloaders = testloaders self.size = len(self.trainloaders[0].dataset[0][0]) + if feature_names is None: + feature_names = [ + f"feature_{index}" + for index in range(self.size) + ] + + if len(feature_names) != self.size: + raise ValueError( + "feature_names length must match the " + f"federated dataset size. Expected {self.size}, " + f"received {len(feature_names)}." + ) + + self.feature_names = list( + feature_names + ) + + + db_manager = DatabaseManager() db_manager.connect() self.eng = db_manager.get_connection() diff --git a/MEDfl/LearningManager/model.py b/MEDfl/LearningManager/model.py index f68dc73..206873c 100644 --- a/MEDfl/LearningManager/model.py +++ b/MEDfl/LearningManager/model.py @@ -1,6 +1,7 @@ #!/usr/bin/env python3 -# froked from https://github.com/pythonlessons/mltu/blob/main/mltu/torch/model.py +import logging +import sys import typing from collections import OrderedDict from typing import Dict, List, Optional, Tuple @@ -8,19 +9,39 @@ import numpy as np import torch import torch.nn as nn -from sklearn.metrics import accuracy_score,roc_auc_score +from sklearn.metrics import ( + accuracy_score, roc_auc_score, + mean_squared_error, mean_absolute_error, r2_score, +) from .utils import params +# --------------------------------------------------------------------------- +# Per-worker logger — same pattern as client.py so Ray forwards to notebook +# --------------------------------------------------------------------------- +def _get_logger() -> logging.Logger: + logger = logging.getLogger("MEDfl.model") + if not logger.handlers: + handler = logging.StreamHandler(sys.stdout) + handler.setFormatter( + logging.Formatter( + "[%(asctime)s][%(name)s][%(levelname)s] %(message)s", + datefmt="%H:%M:%S", + ) + ) + logger.addHandler(handler) + logger.setLevel(logging.DEBUG) + logger.propagate = False + return logger + + +_log = _get_logger() + + class Model: """ Model class for training and testing PyTorch neural networks. - - Attributes: - model (torch.nn.Module): PyTorch neural network. - optimizer (torch.optim.Optimizer): PyTorch optimizer. - criterion (typing.Callable): Loss function. """ def __init__( @@ -28,197 +49,232 @@ def __init__( model: torch.nn.Module, optimizer: torch.optim.Optimizer, criterion: typing.Callable, + task_type: str = "binary", ) -> None: - """ - Initialize Model class with the specified model, optimizer, and criterion. - - Args: - model (torch.nn.Module): PyTorch neural network. - optimizer (torch.optim.Optimizer): PyTorch optimizer. - criterion (typing.Callable): Loss function. - """ self.model = model self.optimizer = optimizer self.criterion = criterion - # Get device on which model is running + self.task_type = task_type self.validate() def validate(self) -> None: - """ - Validate model and optimizer. - """ if not isinstance(self.model, torch.nn.Module): raise TypeError("model argument must be a torch.nn.Module") - if not isinstance(self.optimizer, torch.optim.Optimizer): - raise TypeError( - "optimizer argument must be a torch.optim.Optimizer" - ) + raise TypeError("optimizer argument must be a torch.optim.Optimizer") def get_parameters(self) -> List[np.ndarray]: - """ - Get the parameters of the model as a list of NumPy arrays. - - Returns: - List[np.ndarray]: The parameters of the model as a list of NumPy arrays. - """ - return [ - val.cpu().numpy() for _, val in self.model.state_dict().items() - ] + return [val.detach().cpu().numpy() for val in self.model.state_dict().values()] def set_parameters(self, parameters: List[np.ndarray]) -> None: - """ - Set the parameters of the model from a list of NumPy arrays. - - Args: - parameters (List[np.ndarray]): The parameters to be set. - """ - params_dict = zip(self.model.state_dict().keys(), parameters) - state_dict = OrderedDict({k: torch.Tensor(v) for k, v in params_dict}) - self.model.load_state_dict(state_dict, strict=True) - - def train( - self, train_loader, epoch, device, privacy_engine, diff_priv=False - ) -> float: - """ - Train the model on the given train_loader for one epoch. - - Args: - train_loader: The data loader for training data. - epoch (int): The current epoch number. - device: The device on which to perform the training. - privacy_engine: The privacy engine used for differential privacy (if enabled). - diff_priv (bool, optional): Whether differential privacy is used. Default is False. - - Returns: - float: The value of epsilon used in differential privacy. - """ + state_dict = self.model.state_dict() + keys = list(state_dict.keys()) + device = next(self.model.parameters()).device + new_state_dict = OrderedDict() + for k, v in zip(keys, parameters): + ref = state_dict[k] + new_state_dict[k] = torch.as_tensor(v, device=device, dtype=ref.dtype) + self.model.load_state_dict(new_state_dict, strict=True) + + def train(self, train_loader, epoch, device, privacy_engine, diff_priv=False) -> float: self.model.train() - epsilon = 0 + epsilon = 0.0 losses = [] - top1_acc = [] + metrics = [] for i, (X_train, y_train) in enumerate(train_loader): - X_train, y_train = X_train.to(device), y_train.to(device) - - self.optimizer.zero_grad() + X_train = X_train.to(device) + y_train = y_train.to(device).view(-1) + + if torch.isnan(X_train).any() or torch.isinf(X_train).any(): + _log.warning(f"Train batch {i}: X_train has NaN/Inf -> skipping batch") + continue + if torch.isnan(y_train).any() or torch.isinf(y_train).any(): + _log.warning(f"Train batch {i}: y_train has NaN/Inf -> skipping batch") + continue + + if i == 0 and epoch == 0: + try: + _log.debug( + f"Train scale check | " + f"X min/max: {X_train.min().item():.4f} / {X_train.max().item():.4f} | " + f"y min/max: {y_train.min().item():.4f} / {y_train.max().item():.4f}" + ) + except Exception: + pass - # compute output - y_hat = torch.squeeze(self.model(X_train), 1) - loss = self.criterion(y_hat, y_train) + self.optimizer.zero_grad(set_to_none=True) - preds = np.argmax(y_hat.detach().cpu().numpy(), axis=0) - labels = y_train.detach().cpu().numpy() + y_hat = self.model(X_train).view(-1) - # measure accuracy and record loss - acc = (preds == labels).mean() + if torch.isnan(y_hat).any() or torch.isinf(y_hat).any(): + _log.warning(f"Train batch {i}: y_hat has NaN/Inf -> skipping batch") + continue + + loss = self.criterion(y_hat, y_train) - losses.append(loss.item()) - top1_acc.append(acc) + if torch.isnan(loss) or torch.isinf(loss): + _log.warning(f"Train batch {i}: loss is NaN/Inf -> skipping batch") + continue + losses.append(float(loss.item())) loss.backward() + + bad_grad = False + for name, p in self.model.named_parameters(): + if p.grad is None: + continue + if torch.isnan(p.grad).any() or torch.isinf(p.grad).any(): + _log.warning( + f"Train batch {i}: grad NaN/Inf in {name} " + f"(NaN={torch.isnan(p.grad).sum().item()}, " + f"Inf={torch.isinf(p.grad).sum().item()}) -> skipping update" + ) + bad_grad = True + break + + if bad_grad: + self.optimizer.zero_grad(set_to_none=True) + continue + + torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0) self.optimizer.step() + if self.task_type == "binary": + y_prob = y_hat.detach().cpu().numpy() + y_true = y_train.detach().cpu().numpy() + acc = accuracy_score(y_true, (y_prob >= 0.5).astype(int)) + metrics.append(float(acc)) + elif self.task_type == "regression": + rmse = torch.sqrt(torch.mean((y_hat - y_train) ** 2)).item() + metrics.append(float(rmse)) + if diff_priv: - epsilon = privacy_engine.get_epsilon(float(params["DELTA"])) - - if (i + 1) % 10 == 0: - if diff_priv: - epsilon = privacy_engine.get_epsilon(float(params["DELTA"])) - print( - f"\tTrain Epoch: {epoch} \t" - f"Loss: {np.mean(losses):.6f} " - f"Acc@1: {np.mean(top1_acc) * 100:.6f} " - f"(ε = {epsilon:.2f}, δ = {params['DELTA']})" + epsilon = float(privacy_engine.get_epsilon(float(params["DELTA"]))) + + if (i + 1) % 10 == 0 and len(losses) > 0: + if self.task_type == "binary" and len(metrics) > 0: + _log.info( + f"Train Epoch {epoch} | batch {i+1} | " + f"Loss {np.mean(losses):.6f} | " + f"Acc {np.mean(metrics) * 100:.2f}%" + + (f" | ε={epsilon:.2f}" if diff_priv else "") ) - else: - print( - f"\tTrain Epoch: {epoch} \t" - f"Loss: {np.mean(losses):.6f} " - f"Acc@1: {np.mean(top1_acc) * 100:.6f}" + elif self.task_type == "regression" and len(metrics) > 0: + _log.info( + f"Train Epoch {epoch} | batch {i+1} | " + f"Loss {np.mean(losses):.6f} | " + f"RMSE {np.mean(metrics):.6f}" + + (f" | ε={epsilon:.2f}" if diff_priv else "") ) - return epsilon + mean_train_loss = float(np.mean(losses)) if len(losses) else float("nan") + mean_train_metric = float(np.mean(metrics)) if len(metrics) else float("nan") - def evaluate(self, val_loader, device=torch.device("cpu")) -> Tuple[float, float]: - """ - Evaluate the model on the given validation data. + total_grad_norm = 0.0 + for p in self.model.parameters(): + if p.grad is not None: + total_grad_norm += p.grad.detach().norm(2).item() ** 2 + total_grad_norm = total_grad_norm ** 0.5 + # store it + if not hasattr(self, "grad_norms"): + self.grad_norms = [] + self.grad_norms.append(total_grad_norm) - Args: - val_loader: The data loader for validation data. - device: The device on which to perform the evaluation. Default is 'cpu'. + return epsilon, mean_train_loss, mean_train_metric - Returns: - Tuple[float, float]: The evaluation loss and accuracy. - """ - correct, total, loss, accuracy, auc = 0, 0, 0.0, [], [] + def evaluate(self, val_loader, device=torch.device("cpu")): self.model.eval() + losses = [] + y_true_all, y_pred_all = [], [] with torch.no_grad(): - for X_test, y_test in val_loader: - X_test, y_test = X_test.to(device), y_test.to(device) # Move data to device - - y_hat = torch.squeeze(self.model(X_test), 1) - - - criterion = self.criterion.to(y_hat.device) - loss += criterion(y_hat, y_test).item() - - - # Move y_hat to CPU for accuracy computation - y_hat_cpu = y_hat.cpu().detach().numpy() - accuracy.append(accuracy_score(y_test.cpu().numpy(), y_hat_cpu.round())) - - # Move y_test to CPU for AUC computation - y_test_cpu = y_test.cpu().numpy() - y_prob_cpu = y_hat.cpu().detach().numpy() - if (len(np.unique(y_test_cpu)) != 1): - auc.append(roc_auc_score(y_test_cpu, y_prob_cpu)) - - total += y_test.size(0) - correct += np.sum(y_hat_cpu.round() == y_test_cpu) - - loss /= len(val_loader.dataset) - return loss, np.mean(accuracy), np.mean(auc) - + for i, (X_test, y_test) in enumerate(val_loader): + X_test = X_test.to(device) + y_test_t = y_test.to(device).view(-1) + y_hat_t = self.model(X_test).view(-1) + + if torch.isnan(y_hat_t).any() or torch.isinf(y_hat_t).any(): + _log.warning(f"Eval batch {i}: y_hat NaN/Inf -> skipping") + continue + if torch.isnan(y_test_t).any() or torch.isinf(y_test_t).any(): + _log.warning(f"Eval batch {i}: y_test NaN/Inf -> skipping") + continue + + crit = self.criterion.to(y_hat_t.device) if hasattr(self.criterion, "to") else self.criterion + loss_t = crit(y_hat_t, y_test_t) + + if torch.isnan(loss_t) or torch.isinf(loss_t): + _log.warning(f"Eval batch {i}: loss NaN/Inf -> skipping") + continue + + losses.append(float(loss_t.item())) + y_true_all.append(np.atleast_1d(y_test_t.detach().cpu().numpy())) + y_pred_all.append(np.atleast_1d(y_hat_t.detach().cpu().numpy())) + + if len(losses) == 0 or len(y_true_all) == 0: + return float("nan"), {} + + y_true = np.concatenate(y_true_all, axis=0) + y_pred = np.concatenate(y_pred_all, axis=0) + mean_loss = float(np.mean(losses)) + + if self.task_type == "binary": + acc = accuracy_score(y_true, (y_pred >= 0.5).astype(int)) + auc = roc_auc_score(y_true, y_pred) if len(np.unique(y_true)) > 1 else float("nan") + return mean_loss, {"accuracy": float(acc), "auc": float(auc)} + + if self.task_type == "regression": + mask = np.isfinite(y_true) & np.isfinite(y_pred) + if mask.sum() == 0: + return mean_loss, {"rmse": float("nan"), "mae": float("nan"), + "r2": float("nan"), "adj_r2": float("nan")} + + y_true_c = y_true[mask] + y_pred_c = y_pred[mask] + mse = mean_squared_error(y_true_c, y_pred_c) + mae = mean_absolute_error(y_true_c, y_pred_c) + r2 = r2_score(y_true_c, y_pred_c) + rmse = float(np.sqrt(mse)) + n = int(len(y_true_c)) + + p = None + try: + for m in self.model.modules(): + if isinstance(m, torch.nn.Linear): + p = int(m.in_features) + break + except Exception: + p = None + + if p is None or n <= p + 1: + adj_r2 = float("nan") + else: + adj_r2 = 1.0 - (1.0 - float(r2)) * (n - 1) / (n - p - 1) + + _log.debug(f"adj_r2 check | n={n} | p={p} | condition n>p+1={n > p + 1}") + + return mean_loss, { + "rmse": float(rmse), + "mae": float(mae), + "r2": float(r2), + "adj_r2": float(adj_r2), + } + + return mean_loss, {} @staticmethod - def save_model(model , model_name:str): - """ - Saves a PyTorch model to a file. - - Args: - model (torch.nn.Module): PyTorch model to be saved. - model_name (str): Name of the model file. - - Raises: - Exception: If there is an issue during the saving process. - - Returns: - None - """ + def save_model(model, model_name: str): try: torch.save(model, '../../notebooks/.ipynb_checkpoints/trainedModels/' + model_name + ".pth") except Exception as e: raise Exception(f"Error saving the model: {str(e)}") - + @staticmethod def load_model(model_path: str): - """ - Loads a PyTorch model from a file. - - Args: - model_path (str): Path to the model file to be loaded. - - Returns: - torch.nn.Module: Loaded PyTorch model. - """ - # Ensure models are loaded onto the CPU when CUDA is not available torch_kwargs = {"weights_only": False} if torch.cuda.is_available(): - loaded_model = torch.load(model_path , **torch_kwargs) + loaded_model = torch.load(model_path, **torch_kwargs) else: - loaded_model = torch.load(model_path, map_location=torch.device('cpu') , **torch_kwargs) - return loaded_model - - + loaded_model = torch.load(model_path, map_location=torch.device('cpu'), **torch_kwargs) + return loaded_model \ No newline at end of file diff --git a/MEDfl/LearningManager/params.yaml b/MEDfl/LearningManager/params.yaml index 83fc6fa..95f0440 100644 --- a/MEDfl/LearningManager/params.yaml +++ b/MEDfl/LearningManager/params.yaml @@ -1,8 +1,8 @@ DELTA: 1.0e-05 EPSILON: 5.0 MAX_GRAD_NORM: 1.0 -diff_privacy: true -lr: 0.001 +diff_privacy: false +lr: 0.01 min_evalclient: 2 num_rounds: 20 optimizer: SGD @@ -10,5 +10,5 @@ path_to_master_csv: /home/local/USHERBROOKE/saho6810/MEDfl/code/MEDfl/notebooks/ path_to_test_csv: /home/local/USHERBROOKE/saho6810/MEDfl/code/MEDfl/notebooks/data/masterDataSet/Mimic_train.csv task: BinaryClassification test_batch_size: 1 -train_batch_size: 32 -train_epochs: 20 +train_batch_size: 6 +train_epochs: 1 diff --git a/MEDfl/LearningManager/plot.py b/MEDfl/LearningManager/plot.py index c5f6c6e..23aa596 100644 --- a/MEDfl/LearningManager/plot.py +++ b/MEDfl/LearningManager/plot.py @@ -227,3 +227,108 @@ def plot_classification_report(pipeline_name: str): plt.show() return + + + + +def plot_client_regression_weights(strategy_obj, feature_names=None, include_bias=True): + """ + Plot per-client linear-regression weights over FL rounds. + + Assumes each regression entry in: + strategy_obj.client_fit_history[cid] + contains: + - "round" + - "weights" + + For a linear regressor: + weights[0] -> weight matrix, shape (1, input_dim) or similar + weights[1] -> bias, shape (1,) or similar + """ + + history = strategy_obj.client_fit_history + + if not history: + print("No client_fit_history found.") + return + + for cid, entries in history.items(): + # keep only rounds where weights were stored + reg_entries = [e for e in entries if "weights" in e] + + if not reg_entries: + print(f"Client {cid}: no stored regression weights.") + continue + + rounds = [] + coeff_series = [] + bias_series = [] + + for e in reg_entries: + rounds.append(int(e["round"])) + + weights = e["weights"] + + # linear regression assumption + w = np.array(weights[0]).reshape(-1) + + coeff_series.append(w) + + if include_bias and len(weights) > 1: + b = float(np.array(weights[1]).reshape(-1)[0]) + bias_series.append(b) + + coeff_series = np.array(coeff_series) # shape: (num_rounds, num_features) + + plt.figure(figsize=(10, 5)) + + n_features = coeff_series.shape[1] + for j in range(n_features): + label = feature_names[j] if feature_names is not None and j < len(feature_names) else f"w{j}" + plt.plot(rounds, coeff_series[:, j], marker="o", label=label) + + if include_bias and len(bias_series) == len(rounds): + plt.plot(rounds, bias_series, marker="x", linestyle="--", label="bias") + + plt.title(f"Client {cid} - weights over rounds") + plt.xlabel("Round") + plt.ylabel("Weight value") + plt.grid(True) + plt.legend() + plt.tight_layout() + plt.show() + + + + +def plot_one_weight_across_clients(strategy_obj, weight_index=0, feature_name=None): + history = strategy_obj.client_fit_history + + plt.figure(figsize=(10, 5)) + + for cid, entries in history.items(): + reg_entries = [e for e in entries if "weights" in e] + if not reg_entries: + continue + + rounds = [] + vals = [] + + for e in reg_entries: + w = np.array(e["weights"][0]).reshape(-1) + if weight_index >= len(w): + continue + rounds.append(int(e["round"])) + vals.append(float(w[weight_index])) + + if rounds: + plt.plot(rounds, vals, marker=",", label=f"Client {cid}") + + label = feature_name if feature_name is not None else f"w{weight_index}" + plt.title(f"{label} across clients over rounds") + plt.xlabel("Round") + plt.ylabel("Weight value") + plt.grid(True) + plt.legend() + plt.tight_layout() + plt.show() \ No newline at end of file diff --git a/MEDfl/LearningManager/server.py b/MEDfl/LearningManager/server.py index fa88c10..8b54913 100644 --- a/MEDfl/LearningManager/server.py +++ b/MEDfl/LearningManager/server.py @@ -1,32 +1,65 @@ #!/usr/bin/env python3 import copy +import logging +import sys from typing import Dict, Optional, Tuple import flwr as fl import torch +from torch.utils.data import ConcatDataset, DataLoader + from .client import FlowerClient from .federated_dataset import FederatedDataset from .model import Model from .strategy import Strategy +from .shap import ( + SHAPConfig, + LocalPyTorchSHAPExplainer, + aggregate_federated_shap, +) + + +# --------------------------------------------------------------------------- +# Logging setup — call this ONCE before starting the simulation. +# It forces Flower's logger (and the root logger) to write to stdout so that +# Jupyter captures every message from the driver process. +# --------------------------------------------------------------------------- +def setup_fl_logging(level: int = logging.INFO) -> None: + """ + Configure logging so that all FL-related messages appear in the notebook. + + Call this at the top of your notebook cell before creating FlowerServer: + + from MEDfl.LearningManager.server import setup_fl_logging + setup_fl_logging() + """ + fmt = logging.Formatter( + "[%(asctime)s][%(name)s][%(levelname)s] %(message)s", + datefmt="%H:%M:%S", + ) + handler = logging.StreamHandler(sys.stdout) + handler.setFormatter(fmt) + + # Root logger — catches everything + root = logging.getLogger() + root.setLevel(level) + if not any(isinstance(h, logging.StreamHandler) and h.stream is sys.stdout + for h in root.handlers): + root.addHandler(handler) + + # Flower logger specifically + for name in ("flwr", "flwr.server", "flwr.simulation", "flwr.client"): + lg = logging.getLogger(name) + lg.setLevel(level) + lg.propagate = True # bubble up to root so our handler fires + class FlowerServer: """ A class representing the central server for Federated Learning using Flower. - - Attributes: - global_model (Model): The global model that will be federated among clients. - strategy (Strategy): The strategy used for federated learning, specifying communication and aggregation methods. - num_rounds (int): The number of federated learning rounds to perform. - num_clients (int): The number of clients participating in the federated learning process. - fed_dataset (FederatedDataset): The federated dataset used for training and evaluation. - diff_priv (bool): Whether differential privacy is used during the federated learning process. - accuracies (List[float]): A list to store the accuracy of the global model during each round. - losses (List[float]): A list to store the loss of the global model during each round. - flower_clients (List[FlowerClient]): A list to store the FlowerClient objects representing individual clients. - """ def __init__( @@ -37,23 +70,13 @@ def __init__( num_clients: int, fed_dataset: FederatedDataset, diff_privacy: bool = False, - client_resources: Optional[Dict[str, float]] = {'num_cpus': 1, 'num_gpus': 0.0} + client_resources: Optional[Dict[str, float]] = {'num_cpus': 1, 'num_gpus': 0.0} , + shap_config: Optional[ + SHAPConfig + ] = None, ) -> None: - """ - Initialize a FlowerServer object with the specified parameters. - - Args: - global_model (Model): The global model that will be federated among clients. - strategy (Strategy): The strategy used for federated learning, specifying communication and aggregation methods. - num_rounds (int): The number of federated learning rounds to perform. - num_clients (int): The number of clients participating in the federated learning process. - fed_dataset (FederatedDataset): The federated dataset used for training and evaluation. - diff_privacy (bool, optional): Whether differential privacy is used during the federated learning process. - Default is False. - """ - self.device = torch.device( - f"cuda" if torch.cuda.is_available() else "cpu" - ) + + self.device = torch.device("cpu") self.global_model = global_model self.params = global_model.get_parameters() self.global_model.model = global_model.model.to(self.device) @@ -62,120 +85,531 @@ def __init__( self.fed_dataset = fed_dataset self.strategy = strategy self.client_resources = client_resources - setattr( - self.strategy.strategy_object, - "min_available_clients", - self.num_clients, + + self.shap_config = ( + shap_config + if shap_config is not None + else SHAPConfig(enabled=False) ) - setattr( - self.strategy.strategy_object, - "initial_parameters", - fl.common.ndarrays_to_parameters(self.params), + + self.shap_config.validate( + input_size=self.fed_dataset.size ) + + self.federated_shap_result = None + self.local_shap_results = [] + + setattr(self.strategy.strategy_object, "min_available_clients", self.num_clients) + setattr(self.strategy.strategy_object, "initial_parameters", fl.common.ndarrays_to_parameters(self.params)) setattr(self.strategy.strategy_object, "evaluate_fn", self.evaluate) - self.fed_dataset = fed_dataset + self.diff_priv = diff_privacy + + # tracking self.accuracies = [] self.losses = [] self.auc = [] + self.rmse = [] + self.mae = [] + self.r2 = [] + self.adj_r2 = [] + + self.train_losses = [] + self.train_accuracies = [] + self.train_auc = [] + self.train_rmse = [] + self.train_mae = [] + self.train_r2 = [] + self.train_adj_r2 = [] + self.flower_clients = [] + self.device = torch.device("cpu") + self.global_model = global_model + + # move model to CPU + self.global_model.model = self.global_model.model.to(self.device) + + # also move optimizer state to CPU + for state in self.global_model.optimizer.state.values(): + for k, v in state.items(): + if torch.is_tensor(v): + state[k] = v.detach().cpu() + + # if criterion has tensor buffers (e.g., pos_weight), ensure CPU too + if hasattr(self.global_model.criterion, "pos_weight") and self.global_model.criterion.pos_weight is not None: + self.global_model.criterion.pos_weight = self.global_model.criterion.pos_weight.detach().cpu() self.validate() + + def synchronize_final_global_model( + self, + ) -> None: + """ + Copy the last Flower-aggregated parameters into + the MEDfl global model wrapper. + """ + + strategy_object = ( + self.strategy.strategy_object + ) + + if strategy_object is None: + raise RuntimeError( + "The Flower strategy has not been created" + ) + + final_parameters = getattr( + strategy_object, + "latest_aggregated_parameters", + None, + ) + + if final_parameters is None: + # Your evaluate() method already records the latest + # centralized-evaluation parameter arrays. + final_arrays = getattr( + self, + "last_parameters", + None, + ) + + if final_arrays is None: + raise RuntimeError( + "The final global model parameters " + "could not be found" + ) + else: + final_arrays = ( + fl.common.parameters_to_ndarrays( + final_parameters + ) + ) + + self.global_model.set_parameters( + final_arrays + ) + + self.global_model.model.to( + self.device + ) + self.global_model.model.eval() + + logging.getLogger( + "MEDfl.server" + ).info( + "Final global parameters synchronized " + "for post-training operations" + ) + + def _get_shap_loaders( + self, + ) -> list[DataLoader]: + """ + Return one loader per client according to SHAPConfig. + """ + + split = self.shap_config.data_split + + if split == "train": + return self.fed_dataset.trainloaders + + if split == "validation": + return self.fed_dataset.valloaders + + if split == "test": + return self.fed_dataset.testloaders + + raise ValueError( + f"Unsupported SHAP data split: {split}" + ) + + def calculate_local_shap_results( + self, + ) -> list[dict]: + """ + Calculate local SHAP summaries for every client partition. + + This method runs only after federated training has finished. + """ + + if not self.shap_config.enabled: + return [] + + shap_logger = logging.getLogger( + "MEDfl.shap" + ) + + loaders = self._get_shap_loaders() + + feature_names = ( + self.shap_config.feature_names + if self.shap_config.feature_names + is not None + else self.fed_dataset.feature_names + ) + + if len(feature_names) != self.fed_dataset.size: + raise ValueError( + "SHAP feature names do not match the " + "federated dataset input size" + ) + + local_explainer = ( + LocalPyTorchSHAPExplainer( + explainer_type=( + self.shap_config.explainer + ), + background_size=( + self.shap_config.background_size + ), + explanation_size=( + self.shap_config.explanation_size + ), + random_seed=( + self.shap_config.random_seed + ), + clipping_value=( + self.shap_config.clipping_value + ), + ) + ) + + local_results: list[dict] = [] + + for client_index, loader in enumerate( + loaders + ): + client_id = str(client_index) + + if loader is None: + shap_logger.warning( + "Skipping SHAP client %s: loader is None", + client_id, + ) + continue + + available_samples = len( + loader.dataset + ) + + if available_samples == 0: + shap_logger.warning( + "Skipping SHAP client %s: empty dataset", + client_id, + ) + continue + + shap_logger.info( + "Starting local SHAP | client=%s | " + "available_samples=%d", + client_id, + available_samples, + ) + + try: + # The model is used for inference only. + # No optimizer or training method is called. + local_result = ( + local_explainer.calculate( + model=self.global_model.model, + loader=loader, + device=self.device, + ) + ) + + explained_samples = int( + local_result["sample_count"] + ) + + if ( + explained_samples + < self.shap_config.minimum_samples + ): + shap_logger.warning( + "Skipping client %s because only %d " + "samples were explained; minimum=%d", + client_id, + explained_samples, + self.shap_config.minimum_samples, + ) + continue + + local_result[ + "client_id" + ] = client_id + + local_result[ + "feature_names" + ] = list(feature_names) + + local_result[ + "available_samples" + ] = available_samples + + local_results.append( + local_result + ) + + shap_logger.info( + "Local SHAP completed | client=%s | " + "explained_samples=%d", + client_id, + explained_samples, + ) + + except Exception as exc: + shap_logger.exception( + "Local SHAP failed for client %s: %s", + client_id, + exc, + ) + + if not local_results: + raise RuntimeError( + "SHAP was enabled, but no client produced " + "a valid explanation result" + ) + + self.local_shap_results = ( + local_results + ) + + return local_results + + def calculate_federated_shap( + self, + ) -> dict: + """ + Run the complete post-training SHAP phase. + """ + + shap_logger = logging.getLogger( + "MEDfl.shap" + ) + + if not self.shap_config.enabled: + return { + "status": "disabled", + } + + shap_logger.info( + "========================================" + ) + shap_logger.info( + "Starting post-training federated SHAP" + ) + shap_logger.info( + "Explainer: %s", + self.shap_config.explainer, + ) + shap_logger.info( + "Data split: %s", + self.shap_config.data_split, + ) + shap_logger.info( + "========================================" + ) + + local_results = ( + self.calculate_local_shap_results() + ) + + federated_result = ( + aggregate_federated_shap( + client_results=local_results, + include_client_results=( + self.shap_config + .include_client_results + ), + ) + ) + + federated_result["explainer"] = ( + self.shap_config.explainer + ) + federated_result["data_split"] = ( + self.shap_config.data_split + ) + federated_result["background_size"] = ( + self.shap_config.background_size + ) + federated_result["requested_explanation_size"] = ( + self.shap_config.explanation_size + ) + + self.federated_shap_result = ( + federated_result + ) + + shap_logger.info( + "Federated SHAP completed | clients=%d | " + "samples=%d", + federated_result[ + "participating_clients" + ], + federated_result[ + "explained_samples" + ], + ) + + return federated_result def validate(self) -> None: - """Validate global_model, strategy, num_clients, num_rounds, fed_dataset, diff_privacy""" if not isinstance(self.global_model, Model): raise TypeError("global_model argument must be a Model instance") - - # if not isinstance(self.strategy, Strategy): - # print(self.strategy) - # print(isinstance(self.strategy, Strategy)) - # raise TypeError("strategy argument must be a Strategy instance") - if not isinstance(self.num_clients, int): raise TypeError("num_clients argument must be an int") - if not isinstance(self.num_rounds, int): raise TypeError("num_rounds argument must be an int") - if not isinstance(self.diff_priv, bool): raise TypeError("diff_priv argument must be a bool") def client_fn(self, cid) -> FlowerClient: - """ - Return a FlowerClient object for a specific client ID. - - Args: - cid: The client ID. + device = torch.device(f"cuda:{int(cid) % 4}" if torch.cuda.is_available() else "cpu") - Returns: - FlowerClient: A FlowerClient object representing the individual client. - """ - - device = torch.device( - f"cuda:{int(cid) % 4}" if torch.cuda.is_available() else "cpu" - ) client_model = copy.deepcopy(self.global_model) - + trainloader = self.fed_dataset.trainloaders[int(cid)] - valloader = self.fed_dataset.valloaders[int(cid)] - # this helps in making plots + + # If no validation → reuse trainloader - client = FlowerClient( - cid, client_model, trainloader, valloader, self.diff_priv - ) + if len(self.fed_dataset.valloaders[int(cid)].dataset) == 0: + valloader = trainloader + else: + valloader = self.fed_dataset.valloaders[int(cid)] + + client_model = copy.deepcopy(self.global_model) + client = FlowerClient(cid, client_model, trainloader, valloader, self.diff_priv) self.flower_clients.append(client) return client + def _concat_loader(self, loaders, shuffle: bool = False) -> DataLoader: + ds = ConcatDataset([ldr.dataset for ldr in loaders]) + batch_size = getattr(loaders[0], "batch_size", None) or 64 + num_workers = getattr(loaders[0], "num_workers", 0) + pin_memory = getattr(loaders[0], "pin_memory", False) + + return DataLoader( + ds, + batch_size=batch_size, + shuffle=shuffle, + num_workers=num_workers, + pin_memory=pin_memory, + ) + def evaluate( self, server_round: int, parameters: fl.common.NDArrays, config: Dict[str, fl.common.Scalar], ) -> Optional[Tuple[float, Dict[str, fl.common.Scalar]]]: - """ - Evaluate the global model on the validation dataset and update the accuracies and losses. - Args: - server_round (int): The current round of the federated learning process. - parameters (fl.common.NDArrays): The global model parameters. - config (Dict[str, fl.common.Scalar]): Configuration dictionary. + self.global_model.set_parameters(parameters) + self.last_parameters = parameters - Returns: - Optional[Tuple[float, Dict[str, fl.common.Scalar]]]: The evaluation loss and accuracy. - """ - testloader = self.fed_dataset.valloaders[0] - - self.global_model.set_parameters( - parameters - ) # Update model with the latest parameters - loss, accuracy ,auc = self.global_model.evaluate(testloader, self.device) - self.auc.append(auc) - self.losses.append(loss) - self.accuracies.append(accuracy) - - return loss, {"accuracy": accuracy} + global_valloader = self._concat_loader(self.fed_dataset.valloaders, shuffle=False) + global_trainloader = self._concat_loader(self.fed_dataset.trainloaders, shuffle=False) + + val_loss, val_metrics = self.global_model.evaluate(global_valloader, self.device) + self.losses.append(val_loss) + + if isinstance(val_metrics, dict): + if "accuracy" in val_metrics: self.accuracies.append(val_metrics["accuracy"]) + if "auc" in val_metrics: self.auc.append(val_metrics["auc"]) + if "rmse" in val_metrics: self.rmse.append(val_metrics["rmse"]) + if "mae" in val_metrics: self.mae.append(val_metrics["mae"]) + if "r2" in val_metrics: self.r2.append(val_metrics["r2"]) + if "adj_r2" in val_metrics: self.adj_r2.append(val_metrics["adj_r2"]) + + train_loss, train_metrics = self.global_model.evaluate(global_trainloader, self.device) + self.train_losses.append(train_loss) + + if isinstance(train_metrics, dict): + if "accuracy" in train_metrics: self.train_accuracies.append(train_metrics["accuracy"]) + if "auc" in train_metrics: self.train_auc.append(train_metrics["auc"]) + if "rmse" in train_metrics: self.train_rmse.append(train_metrics["rmse"]) + if "mae" in train_metrics: self.train_mae.append(train_metrics["mae"]) + if "r2" in train_metrics: self.train_r2.append(train_metrics["r2"]) + if "adj_r2" in train_metrics: self.train_adj_r2.append(train_metrics["adj_r2"]) + + out_metrics = dict(val_metrics or {}) + out_metrics["train_loss"] = float(train_loss) if train_loss is not None else float("nan") + + if isinstance(train_metrics, dict): + for k, v in train_metrics.items(): + out_metrics[f"train_{k}"] = float(v) if v is not None else float("nan") + + return val_loss, out_metrics + + def score_each_client(self, parameters=None, split: str = "val"): + if parameters is not None: + self.global_model.set_parameters(parameters) + + loaders = self.fed_dataset.valloaders if split == "val" else self.fed_dataset.trainloaders + + results = {} + for cid, loader in enumerate(loaders): + loss, metrics = self.global_model.evaluate(loader, self.device) + results[str(cid)] = { + "num_examples": len(loader.dataset), + "loss": float(loss), + **{k: float(v) for k, v in (metrics or {}).items()}, + } + return results def run(self) -> None: - """ - Run the federated learning process using Flower simulation. + # --------------------------------------------------------------- + # KEY FIX 1: log_to_driver=True is required, AND we must redirect + # Ray worker stdout/stderr back to the driver. + # --------------------------------------------------------------- + ray_init_args = { + "include_dashboard": False, + "object_store_memory": 78643200, + "log_to_driver": True, # sends worker logs → driver + "logging_level": logging.INFO, # Ray's own log level + } + + # --------------------------------------------------------------- + # KEY FIX 2: make sure the driver-side logger is set up before + # simulation starts so Flower's internal messages reach the notebook. + # --------------------------------------------------------------- + setup_fl_logging() - Returns: - History: The history of the accuracies and losses during the training of each node - """ - # Increase the object store memory to the minimum allowed value or higher - ray_init_args = {"include_dashboard": False - , "object_store_memory": 78643200 - } self.fed_dataset.eng = None - + history = fl.simulation.start_simulation( client_fn=self.client_fn, num_clients=self.num_clients, config=fl.server.ServerConfig(self.num_rounds), strategy=self.strategy.strategy_object, ray_init_args=ray_init_args, - client_resources = self.client_resources + client_resources=self.client_resources ) + # Training has now completely finished. + self.synchronize_final_global_model() + + if self.shap_config.enabled: + try: + self.federated_shap_result = ( + self.calculate_federated_shap() + ) + + except Exception as exc: + logging.getLogger( + "MEDfl.shap" + ).exception( + "Post-training federated SHAP failed: %s", + exc, + ) - return history + # Important: a SHAP failure does not mean + # federated training failed. + self.federated_shap_result = { + "status": "failed", + "error": str(exc), + } + else: + self.federated_shap_result = { + "status": "disabled", + } + return history \ No newline at end of file diff --git a/MEDfl/LearningManager/shap/__init__.py b/MEDfl/LearningManager/shap/__init__.py new file mode 100644 index 0000000..8805c35 --- /dev/null +++ b/MEDfl/LearningManager/shap/__init__.py @@ -0,0 +1,14 @@ +from .config import SHAPConfig +from .explainer import LocalPyTorchSHAPExplainer +from .xgboost_explainer import ( + LocalXGBoostSHAPExplainer, +) +from .aggregation import aggregate_federated_shap + + +__all__ = [ + "SHAPConfig", + "LocalPyTorchSHAPExplainer", + "LocalXGBoostSHAPExplainer", + "aggregate_federated_shap", +] \ No newline at end of file diff --git a/MEDfl/LearningManager/shap/aggregation.py b/MEDfl/LearningManager/shap/aggregation.py new file mode 100644 index 0000000..0e70e69 --- /dev/null +++ b/MEDfl/LearningManager/shap/aggregation.py @@ -0,0 +1,249 @@ +from __future__ import annotations + +from typing import Any + +import numpy as np + + +def _validate_client_results( + client_results: list[dict], +) -> tuple[list[str], int]: + if not client_results: + raise ValueError( + "No client SHAP results were provided" + ) + + reference_features = client_results[0][ + "feature_names" + ] + + feature_count = len(reference_features) + + required_array_keys = { + "abs_sum", + "signed_sum", + "square_sum", + "positive_count", + "negative_count", + "zero_count", + } + + for result in client_results: + if result["feature_names"] != reference_features: + raise ValueError( + "Clients used different feature names or orders" + ) + + if int(result["sample_count"]) <= 0: + raise ValueError( + f"Client {result.get('client_id')} returned an " + "invalid SHAP sample count" + ) + + for key in required_array_keys: + if key not in result: + raise ValueError( + f"Client result is missing '{key}'" + ) + + values = np.asarray(result[key]) + + if values.shape != (feature_count,): + raise ValueError( + f"Client {result.get('client_id')} returned " + f"invalid shape {values.shape} for {key}. " + f"Expected ({feature_count},)." + ) + + return reference_features, feature_count + + +def aggregate_federated_shap( + client_results: list[dict], + include_client_results: bool = True, +) -> dict[str, Any]: + """ + Aggregate additive local SHAP statistics. + + The global mean absolute SHAP is equivalent to calculating: + + sum of all |SHAP values| / number of all explained samples + """ + + feature_names, feature_count = ( + _validate_client_results(client_results) + ) + + total_samples = sum( + int(result["sample_count"]) + for result in client_results + ) + + abs_sum = np.zeros( + feature_count, + dtype=np.float64, + ) + signed_sum = np.zeros( + feature_count, + dtype=np.float64, + ) + square_sum = np.zeros( + feature_count, + dtype=np.float64, + ) + positive_count = np.zeros( + feature_count, + dtype=np.float64, + ) + negative_count = np.zeros( + feature_count, + dtype=np.float64, + ) + zero_count = np.zeros( + feature_count, + dtype=np.float64, + ) + + for result in client_results: + abs_sum += np.asarray( + result["abs_sum"], + dtype=np.float64, + ) + signed_sum += np.asarray( + result["signed_sum"], + dtype=np.float64, + ) + square_sum += np.asarray( + result["square_sum"], + dtype=np.float64, + ) + positive_count += np.asarray( + result["positive_count"], + dtype=np.float64, + ) + negative_count += np.asarray( + result["negative_count"], + dtype=np.float64, + ) + zero_count += np.asarray( + result["zero_count"], + dtype=np.float64, + ) + + mean_abs = abs_sum / total_samples + mean_signed = signed_sum / total_samples + mean_square = square_sum / total_samples + + # Var(X) = E(X²) - E(X)² + variance = np.maximum( + mean_square - np.square(mean_signed), + 0.0, + ) + + standard_deviation = np.sqrt(variance) + + features = [] + + for index, feature_name in enumerate( + feature_names + ): + features.append( + { + "feature": feature_name, + "feature_index": index, + "mean_abs_shap": float( + mean_abs[index] + ), + "mean_signed_shap": float( + mean_signed[index] + ), + "shap_std": float( + standard_deviation[index] + ), + "positive_rate": float( + positive_count[index] + / total_samples + ), + "negative_rate": float( + negative_count[index] + / total_samples + ), + "zero_rate": float( + zero_count[index] + / total_samples + ), + } + ) + + # Highest global importance first. + features.sort( + key=lambda item: item["mean_abs_shap"], + reverse=True, + ) + + final_result: dict[str, Any] = { + "status": "completed", + "participating_clients": len( + client_results + ), + "explained_samples": int( + total_samples + ), + "feature_importance": features, + } + + if include_client_results: + client_summaries = [] + + for result in client_results: + sample_count = int( + result["sample_count"] + ) + + client_mean_abs = ( + np.asarray( + result["abs_sum"], + dtype=np.float64, + ) + / sample_count + ) + + local_features = [ + { + "feature": feature_name, + "feature_index": index, + "mean_abs_shap": float( + client_mean_abs[index] + ), + } + for index, feature_name in enumerate( + feature_names + ) + ] + + local_features.sort( + key=lambda item: item[ + "mean_abs_shap" + ], + reverse=True, + ) + + client_summaries.append( + { + "client_id": result[ + "client_id" + ], + "explained_samples": ( + sample_count + ), + "feature_importance": ( + local_features + ), + } + ) + + final_result["client_results"] = ( + client_summaries + ) + + return final_result \ No newline at end of file diff --git a/MEDfl/LearningManager/shap/config.py b/MEDfl/LearningManager/shap/config.py new file mode 100644 index 0000000..7eb810c --- /dev/null +++ b/MEDfl/LearningManager/shap/config.py @@ -0,0 +1,196 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import List, Literal, Optional + + +@dataclass +class SHAPConfig: + """ + Configuration for the optional post-training federated SHAP phase. + """ + + enabled: bool = False + + # Neural networks: + # gradient + # deep + # + # XGBoost: + # tree + explainer: Literal[ + "gradient", + "deep", + "tree", + ] = "gradient" + + data_split: Literal[ + "train", + "validation", + "test", + ] = "validation" + + background_size: int = 100 + explanation_size: int = 500 + random_seed: int = 42 + + feature_names: Optional[List[str]] = None + + minimum_samples: int = 10 + clipping_value: Optional[float] = None + + include_client_results: bool = True + + # TreeExplainer output: + # + # raw: + # Binary classification → log-odds margin + # Regression → raw prediction + # + # probability: + # Binary classification probability when supported by the + # installed SHAP/XGBoost versions. + model_output: Literal[ + "raw", + "probability", + ] = "raw" + + def validate( + self, + input_size: int, + ) -> None: + if not isinstance(self.enabled, bool): + raise TypeError( + "SHAP enabled must be a boolean" + ) + + if self.explainer not in { + "gradient", + "deep", + "tree", + }: + raise ValueError( + "SHAP explainer must be 'gradient', " + "'deep', or 'tree'" + ) + + if self.data_split not in { + "train", + "validation", + "test", + }: + raise ValueError( + "SHAP data_split must be train, " + "validation, or test" + ) + + if self.background_size <= 0: + raise ValueError( + "SHAP background_size must be greater than zero" + ) + + if self.explanation_size <= 0: + raise ValueError( + "SHAP explanation_size must be greater than zero" + ) + + if self.minimum_samples <= 0: + raise ValueError( + "SHAP minimum_samples must be greater than zero" + ) + + if self.clipping_value is not None: + if self.clipping_value <= 0: + raise ValueError( + "SHAP clipping_value must be greater than zero" + ) + + if self.model_output not in { + "raw", + "probability", + }: + raise ValueError( + "SHAP model_output must be 'raw' or " + "'probability'" + ) + + if self.feature_names is not None: + if len(self.feature_names) != input_size: + raise ValueError( + "The number of SHAP feature names must " + "match the model input size. " + f"Expected {input_size}, " + f"received {len(self.feature_names)}." + ) + + if len(set(self.feature_names)) != len( + self.feature_names + ): + raise ValueError( + "SHAP feature names must be unique" + ) + + def validate_backend( + self, + input_size: int, + backend: str, + task: Optional[str] = None, + ) -> None: + """ + Validate both the general configuration and the selected + machine-learning backend. + """ + + self.validate(input_size) + + backend = str(backend).lower() + + if backend == "pytorch": + if self.explainer not in { + "gradient", + "deep", + }: + raise ValueError( + "PyTorch models require explainer='gradient' " + "or explainer='deep'" + ) + + elif backend == "xgboost": + if self.explainer != "tree": + raise ValueError( + "XGBoost models require explainer='tree'" + ) + + if task == "multiclass": + raise ValueError( + "The current federated XGBoost SHAP " + "implementation supports binary classification " + "and regression. Multiclass requires " + "class-specific SHAP aggregation." + ) + + if ( + task == "regression" + and self.model_output == "probability" + ): + raise ValueError( + "model_output='probability' is not valid for " + "XGBoost regression" + ) + + else: + raise ValueError( + f"Unsupported SHAP backend: {backend}" + ) + + def resolved_feature_names( + self, + input_size: int, + ) -> List[str]: + if self.feature_names is not None: + return list(self.feature_names) + + return [ + f"feature_{index}" + for index in range(input_size) + ] \ No newline at end of file diff --git a/MEDfl/LearningManager/shap/explainer.py b/MEDfl/LearningManager/shap/explainer.py new file mode 100644 index 0000000..b884709 --- /dev/null +++ b/MEDfl/LearningManager/shap/explainer.py @@ -0,0 +1,324 @@ +from __future__ import annotations + +import logging +from typing import Any + +import numpy as np +import torch +from torch.utils.data import DataLoader + + +logger = logging.getLogger("MEDfl.shap") + + +class LocalPyTorchSHAPExplainer: + """ + Calculate SHAP values for one client partition. + + Raw samples and per-sample SHAP values are not returned. + Only additive summary statistics are returned. + """ + + def __init__( + self, + explainer_type: str = "gradient", + background_size: int = 100, + explanation_size: int = 500, + random_seed: int = 42, + clipping_value: float | None = None, + ) -> None: + self.explainer_type = explainer_type + self.background_size = background_size + self.explanation_size = explanation_size + self.random_seed = random_seed + self.clipping_value = clipping_value + + @staticmethod + def _collect_features( + loader: DataLoader, + ) -> torch.Tensor: + """ + Collect only X values from the selected client loader. + Labels are not required to calculate SHAP values. + """ + feature_batches: list[torch.Tensor] = [] + + for batch_index, batch in enumerate(loader): + if not isinstance(batch, (tuple, list)): + raise TypeError( + "Each DataLoader batch must be a tuple or list" + ) + + if len(batch) == 0: + raise ValueError( + f"Empty batch received at index {batch_index}" + ) + + X = batch[0] + + if not isinstance(X, torch.Tensor): + X = torch.as_tensor(X) + + X = X.detach().cpu().float() + + if torch.isnan(X).any(): + raise ValueError( + f"SHAP input contains NaN values in batch " + f"{batch_index}" + ) + + if torch.isinf(X).any(): + raise ValueError( + f"SHAP input contains infinite values in batch " + f"{batch_index}" + ) + + feature_batches.append(X) + + if not feature_batches: + raise ValueError( + "Cannot calculate SHAP using an empty DataLoader" + ) + + features = torch.cat( + feature_batches, + dim=0, + ) + + # First implementation: tabular data. + if features.ndim != 2: + raise ValueError( + "The current SHAP implementation supports tabular " + "input with shape (samples, features). " + f"Received shape {tuple(features.shape)}." + ) + + return features + + @staticmethod + def _sample_rows( + data: torch.Tensor, + requested_size: int, + generator: torch.Generator, + ) -> torch.Tensor: + selected_size = min( + requested_size, + len(data), + ) + + indices = torch.randperm( + len(data), + generator=generator, + )[:selected_size] + + return data[indices] + + @staticmethod + def _normalize_shap_values( + raw_values: Any, + sample_count: int, + feature_count: int, + ) -> np.ndarray: + """ + Normalize the different shapes returned by SHAP versions. + + Expected final shape: + (number_of_samples, number_of_features) + """ + + if hasattr(raw_values, "values"): + raw_values = raw_values.values + + if isinstance(raw_values, list): + if len(raw_values) != 1: + raise ValueError( + "The current implementation supports binary or " + "single-output models only. Multiclass SHAP " + "requires class-specific aggregation." + ) + + raw_values = raw_values[0] + + values = np.asarray(raw_values) + + # Possible shape: (samples, features, 1) + if ( + values.ndim == 3 + and values.shape[-1] == 1 + ): + values = values[..., 0] + + # Possible shape: (1, samples, features) + if ( + values.ndim == 3 + and values.shape[0] == 1 + ): + values = values[0] + + expected_shape = ( + sample_count, + feature_count, + ) + + if values.shape != expected_shape: + raise ValueError( + "Unexpected SHAP output shape. " + f"Expected {expected_shape}, received {values.shape}." + ) + + values = values.astype( + np.float64, + copy=False, + ) + + if np.isnan(values).any(): + raise ValueError( + "Calculated SHAP values contain NaN values" + ) + + if np.isinf(values).any(): + raise ValueError( + "Calculated SHAP values contain infinite values" + ) + + return values + + def _create_explainer( + self, + model: torch.nn.Module, + background: torch.Tensor, + ): + try: + import shap + except ImportError as exc: + raise ImportError( + "SHAP is not installed. Install it with: " + "pip install shap" + ) from exc + + if self.explainer_type == "gradient": + return shap.GradientExplainer( + model, + background, + ) + + if self.explainer_type == "deep": + return shap.DeepExplainer( + model, + background, + ) + + raise ValueError( + f"Unsupported SHAP explainer: {self.explainer_type}" + ) + + def calculate( + self, + model: torch.nn.Module, + loader: DataLoader, + device: torch.device, + ) -> dict: + """ + Calculate SHAP for one client partition. + + The returned values are additive summaries. They can be + combined across clients without transmitting raw SHAP rows. + """ + + features = self._collect_features(loader) + + if len(features) < 2: + raise ValueError( + "At least two samples are required for SHAP" + ) + + generator = torch.Generator() + generator.manual_seed(self.random_seed) + + background = self._sample_rows( + data=features, + requested_size=self.background_size, + generator=generator, + ) + + explanation_data = self._sample_rows( + data=features, + requested_size=self.explanation_size, + generator=generator, + ) + + background = background.to(device) + explanation_data = explanation_data.to(device) + + model = model.to(device) + model.eval() + + logger.info( + "Calculating local SHAP | " + "explainer=%s | available=%d | background=%d | " + "explained=%d", + self.explainer_type, + len(features), + len(background), + len(explanation_data), + ) + + explainer = self._create_explainer( + model=model, + background=background, + ) + + raw_values = explainer.shap_values( + explanation_data + ) + + shap_values = self._normalize_shap_values( + raw_values=raw_values, + sample_count=len(explanation_data), + feature_count=explanation_data.shape[1], + ) + + if self.clipping_value is not None: + shap_values = np.clip( + shap_values, + -self.clipping_value, + self.clipping_value, + ) + + # Additive summaries are preferable because they can later + # be used with distributed or secure aggregation. + result = { + "sample_count": int(shap_values.shape[0]), + + "abs_sum": np.sum( + np.abs(shap_values), + axis=0, + ), + + "signed_sum": np.sum( + shap_values, + axis=0, + ), + + "square_sum": np.sum( + np.square(shap_values), + axis=0, + ), + + "positive_count": np.sum( + shap_values > 0, + axis=0, + ).astype(np.int64), + + "negative_count": np.sum( + shap_values < 0, + axis=0, + ).astype(np.int64), + + "zero_count": np.sum( + shap_values == 0, + axis=0, + ).astype(np.int64), + } + + return result \ No newline at end of file diff --git a/MEDfl/LearningManager/shap/xgboost_explainer.py b/MEDfl/LearningManager/shap/xgboost_explainer.py new file mode 100644 index 0000000..00f251f --- /dev/null +++ b/MEDfl/LearningManager/shap/xgboost_explainer.py @@ -0,0 +1,436 @@ +from __future__ import annotations + +import logging +from typing import Any, Dict, List, Optional + +import numpy as np +import torch +import xgboost as xgb +from torch.utils.data import DataLoader + + +logger = logging.getLogger("MEDfl.shap.xgboost") + + +class LocalXGBoostSHAPExplainer: + """ + Calculate local SHAP values for a federated XGBoost client. + + The final federated XGBoost booster is explained separately on + every client's selected local partition. + + Raw samples and raw per-sample SHAP values are never returned. + Only additive summary statistics are returned so the existing + federated SHAP aggregation function can be reused. + """ + + def __init__( + self, + background_size: int = 100, + explanation_size: int = 500, + random_seed: int = 42, + clipping_value: Optional[float] = None, + model_output: str = "raw", + ) -> None: + self.background_size = int(background_size) + self.explanation_size = int(explanation_size) + self.random_seed = int(random_seed) + self.clipping_value = clipping_value + self.model_output = str(model_output) + + self._validate_configuration() + + def _validate_configuration(self) -> None: + if self.background_size <= 0: + raise ValueError( + "background_size must be greater than zero" + ) + + if self.explanation_size <= 0: + raise ValueError( + "explanation_size must be greater than zero" + ) + + if self.clipping_value is not None: + if self.clipping_value <= 0: + raise ValueError( + "clipping_value must be greater than zero" + ) + + if self.model_output not in { + "raw", + "probability", + }: + raise ValueError( + "model_output must be 'raw' or 'probability'" + ) + + @staticmethod + def _collect_features( + loader: DataLoader, + ) -> np.ndarray: + """ + Collect only the input features from a client DataLoader. + + Labels are not needed for SHAP calculation. + """ + + if loader is None: + raise ValueError( + "Cannot calculate XGBoost SHAP using a None loader" + ) + + feature_batches: List[np.ndarray] = [] + + for batch_index, batch in enumerate(loader): + if not isinstance(batch, (tuple, list)): + raise TypeError( + "Each DataLoader batch must be a tuple or list" + ) + + if len(batch) == 0: + raise ValueError( + f"Empty batch received at index {batch_index}" + ) + + features = batch[0] + + if isinstance(features, torch.Tensor): + features = ( + features + .detach() + .cpu() + .numpy() + ) + + features = np.asarray(features) + + if features.ndim == 1: + features = features.reshape(1, -1) + + if features.ndim != 2: + raise ValueError( + "The XGBoost SHAP implementation supports " + "two-dimensional tabular inputs. " + f"Received shape {features.shape}." + ) + + if not np.issubdtype( + features.dtype, + np.number, + ): + try: + features = features.astype( + np.float32 + ) + except Exception as exc: + raise TypeError( + "XGBoost SHAP features must be numeric" + ) from exc + + features = features.astype( + np.float32, + copy=False, + ) + + if np.isnan(features).any(): + raise ValueError( + "XGBoost SHAP input contains NaN values " + f"in batch {batch_index}" + ) + + if np.isinf(features).any(): + raise ValueError( + "XGBoost SHAP input contains infinite values " + f"in batch {batch_index}" + ) + + feature_batches.append(features) + + if not feature_batches: + raise ValueError( + "Cannot calculate XGBoost SHAP using an empty " + "DataLoader" + ) + + collected_features = np.concatenate( + feature_batches, + axis=0, + ) + + if collected_features.ndim != 2: + raise ValueError( + "Unexpected collected XGBoost feature shape: " + f"{collected_features.shape}" + ) + + return collected_features + + @staticmethod + def _sample_rows( + data: np.ndarray, + requested_size: int, + random_generator: np.random.Generator, + ) -> np.ndarray: + selected_size = min( + int(requested_size), + int(data.shape[0]), + ) + + indices = random_generator.choice( + data.shape[0], + size=selected_size, + replace=False, + ) + + return data[indices] + + @staticmethod + def _normalize_shap_values( + raw_values: Any, + sample_count: int, + feature_count: int, + ) -> np.ndarray: + """ + Normalize SHAP output to: + + (number_of_samples, number_of_features) + + This first implementation supports binary classification and + regression. Multiclass SHAP requires class-specific aggregation. + """ + + if hasattr(raw_values, "values"): + raw_values = raw_values.values + + if isinstance(raw_values, list): + if len(raw_values) != 1: + raise ValueError( + "Multiclass XGBoost SHAP is not yet supported. " + "The returned SHAP result contains several " + "class-specific arrays." + ) + + raw_values = raw_values[0] + + shap_values = np.asarray(raw_values) + + # Possible binary/single-output shape: + # (samples, features, 1) + if ( + shap_values.ndim == 3 + and shap_values.shape[-1] == 1 + ): + shap_values = shap_values[..., 0] + + # Possible shape: + # (1, samples, features) + if ( + shap_values.ndim == 3 + and shap_values.shape[0] == 1 + ): + shap_values = shap_values[0] + + if shap_values.ndim == 3: + raise ValueError( + "Multiclass XGBoost SHAP is not yet supported. " + "Class-specific SHAP values were returned with " + f"shape {shap_values.shape}." + ) + + expected_shape = ( + sample_count, + feature_count, + ) + + if shap_values.shape != expected_shape: + raise ValueError( + "Unexpected XGBoost SHAP output shape. " + f"Expected {expected_shape}, " + f"received {shap_values.shape}." + ) + + shap_values = shap_values.astype( + np.float64, + copy=False, + ) + + if np.isnan(shap_values).any(): + raise ValueError( + "Calculated XGBoost SHAP values contain NaN values" + ) + + if np.isinf(shap_values).any(): + raise ValueError( + "Calculated XGBoost SHAP values contain " + "infinite values" + ) + + return shap_values + + def _create_explainer( + self, + booster: xgb.Booster, + background: np.ndarray, + ): + try: + import shap + except ImportError as exc: + raise ImportError( + "SHAP is enabled for XGBoost, but the shap " + "package is not installed." + ) from exc + + if not isinstance(booster, xgb.Booster): + raise TypeError( + "LocalXGBoostSHAPExplainer expects an " + "xgboost.Booster" + ) + + try: + return shap.TreeExplainer( + model=booster, + data=background, + feature_perturbation="interventional", + model_output=self.model_output, + ) + + except Exception as exc: + # Some SHAP/XGBoost combinations may not support every + # TreeExplainer argument. Keep a conservative fallback. + logger.warning( + "Could not initialize TreeExplainer with the " + "configured background and model output. " + "Falling back to the default TreeExplainer: %s", + exc, + ) + + return shap.TreeExplainer( + model=booster, + ) + + @staticmethod + def _calculate_raw_values( + explainer, + explanation_data: np.ndarray, + ): + """ + Support different SHAP versions. + """ + + try: + return explainer.shap_values( + explanation_data, + check_additivity=False, + ) + + except TypeError: + # Older SHAP versions may not accept + # check_additivity here. + return explainer.shap_values( + explanation_data + ) + + def calculate( + self, + booster: xgb.Booster, + loader: DataLoader, + ) -> Dict[str, Any]: + """ + Calculate local XGBoost SHAP additive summaries. + """ + + features = self._collect_features(loader) + + if len(features) < 2: + raise ValueError( + "At least two samples are required for " + "XGBoost SHAP" + ) + + random_generator = np.random.default_rng( + self.random_seed + ) + + background = self._sample_rows( + data=features, + requested_size=self.background_size, + random_generator=random_generator, + ) + + explanation_data = self._sample_rows( + data=features, + requested_size=self.explanation_size, + random_generator=random_generator, + ) + + logger.info( + "Calculating local XGBoost SHAP | " + "available=%d | background=%d | explained=%d | " + "model_output=%s", + len(features), + len(background), + len(explanation_data), + self.model_output, + ) + + explainer = self._create_explainer( + booster=booster, + background=background, + ) + + raw_values = self._calculate_raw_values( + explainer=explainer, + explanation_data=explanation_data, + ) + + shap_values = self._normalize_shap_values( + raw_values=raw_values, + sample_count=len(explanation_data), + feature_count=explanation_data.shape[1], + ) + + if self.clipping_value is not None: + shap_values = np.clip( + shap_values, + -self.clipping_value, + self.clipping_value, + ) + + result: Dict[str, Any] = { + "sample_count": int( + shap_values.shape[0] + ), + + "abs_sum": np.sum( + np.abs(shap_values), + axis=0, + ), + + "signed_sum": np.sum( + shap_values, + axis=0, + ), + + "square_sum": np.sum( + np.square(shap_values), + axis=0, + ), + + "positive_count": np.sum( + shap_values > 0, + axis=0, + ).astype(np.int64), + + "negative_count": np.sum( + shap_values < 0, + axis=0, + ).astype(np.int64), + + "zero_count": np.sum( + shap_values == 0, + axis=0, + ).astype(np.int64), + } + + return result \ No newline at end of file diff --git a/MEDfl/LearningManager/strategy.py b/MEDfl/LearningManager/strategy.py index 038d8e6..f746f19 100644 --- a/MEDfl/LearningManager/strategy.py +++ b/MEDfl/LearningManager/strategy.py @@ -1,31 +1,83 @@ +""" +MEDfl Strategy wrapper for Flower simulation (Ray) -from collections import OrderedDict -from typing import Dict, List, Optional, Tuple +This file provides a Strategy factory that wraps any Flower server strategy +(e.g., FedAvg, FedProx, FedAdam...) with extra debugging: -import flwr as fl -import numpy as np +✅ What it adds: +- Per-client fit/eval history with round + num_examples (+ loss for eval) +- Per-round server summary history (results/failures, cids, example stats) +- Global-weights fingerprints after aggregation: + * w_norm, delta_norm, cos_sim_prev, nan_or_inf +- Uses Flower logging (NOT print) so logs show up reliably in simulation +- Optional metrics aggregation functions to remove: + "No fit_metrics_aggregation_fn provided" warning + +Usage (server side): + strat_factory = Strategy(name="FedAvg", debug_print=True) + strategy = strat_factory.create_strategy() + + fl.simulation.start_simulation(..., strategy=strategy, ...) +After run: + strategy.round_history + strategy.client_fit_history + strategy.client_eval_history +""" + +from __future__ import annotations + +from typing import Any, Dict, Optional, Type, List, Tuple, Callable + +import numpy as np +import flwr as fl import optuna +import logging +Metrics = Dict[str, float] +MetricsAggFn = Callable[[List[Tuple[int, Metrics]]], Metrics] -class Strategy: +def weighted_avg(metrics: List[Tuple[int, Metrics]]) -> Metrics: """ - A class representing a strategy for Federated Learning. - - Attributes: - name (str): The name of the strategy. Default is "FedAvg". - fraction_fit (float): Fraction of clients to use for training during each round. Default is 1.0. - fraction_evaluate (float): Fraction of clients to use for evaluation during each round. Default is 1.0. - min_fit_clients (int): Minimum number of clients to use for training during each round. Default is 2. - min_evaluate_clients (int): Minimum number of clients to use for evaluation during each round. Default is 2. - min_available_clients (int): Minimum number of available clients required to start a round. Default is 2. - initial_parameters (Optional[]): The initial parameters of the server model - Methods: - + Weighted average of client metrics by num_examples. + + Flower passes: + metrics = [(num_examples, {"rmse": 1.2, "mae": 0.5}), ...] """ + if not metrics: + return {} + + total = sum(int(n) for n, _ in metrics) + if total <= 0: + return {} + + # union of all keys + keys = set() + for _, m in metrics: + keys |= set(m.keys()) + out: Dict[str, float] = {} + for k in keys: + acc = 0.0 + used = 0 + for n, m in metrics: + if k not in m: + continue + try: + v = float(m[k]) + except Exception: + continue + acc += float(n) * v + used += int(n) + if used > 0: + out[k] = acc / float(total) + + return out + + +class Strategy: def __init__( self, name: str = "FedAvg", @@ -34,49 +86,267 @@ def __init__( min_fit_clients: int = 2, min_evaluate_clients: int = 2, min_available_clients: int = 2, - initial_parameters = [], - evaluation_methode = "centralized" + initial_parameters=None, + evaluation_methode: str = "centralized", + strategy_kwargs: Optional[Dict[str, Any]] = None, + debug_print: bool = True, + enable_metrics_aggregation: bool = True, # ✅ remove "No fit_metrics_aggregation_fn" warnings ) -> None: - """ - Initialize a Strategy object with the specified parameters. - - Args: - name (str): The name of the strategy. Default is "FedAvg". - fraction_fit (float): Fraction of clients to use for training during each round. Default is 1.0. - fraction_evaluate (float): Fraction of clients to use for evaluation during each round. Default is 1.0. - min_fit_clients (int): Minimum number of clients to use for training during each round. Default is 2. - min_evaluate_clients (int): Minimum number of clients to use for evaluation during each round. Default is 2. - min_available_clients (int): Minimum number of available clients required to start a round. Default is 2. - initial_parameters (Optional[]): The initial parametres of the server model - evaluation_methode ( "centralized" | "distributed") - """ self.fraction_fit = fraction_fit self.fraction_evaluate = fraction_evaluate self.min_fit_clients = min_fit_clients self.min_evaluate_clients = min_evaluate_clients self.min_available_clients = min_available_clients - self.initial_parameters = initial_parameters + self.initial_parameters = initial_parameters if initial_parameters is not None else [] self.evaluate_fn = None self.name = name - - def optuna_fed_optimization(self, direction:str , hpo_rate:int , params_config): + self.evaluation_methode = evaluation_methode + self.strategy_kwargs = strategy_kwargs or {} + self.debug_print = debug_print + self.enable_metrics_aggregation = enable_metrics_aggregation + + # Exposed after create_strategy() + self.strategy_object: Optional[fl.server.strategy.Strategy] = None + + # Optuna + self.study: Optional[optuna.Study] = None + self.hpo_rate: Optional[int] = None + self.params_config: Optional[dict] = None + + def optuna_fed_optimization(self, direction: str, hpo_rate: int, params_config): self.study = optuna.create_study(direction=direction) self.hpo_rate = hpo_rate self.params_config = params_config - - + + def get_strategy_by_name(self): + try: + return getattr(fl.server.strategy, self.name) + except AttributeError as e: + available = [n for n in dir(fl.server.strategy) if n and n[0].isupper()] + raise ValueError(f"Unknown strategy '{self.name}'. Available: {available}") from e + def create_strategy(self): - self.strategy_object = self.get_strategy_by_name()( + BaseStrategy: Type[fl.server.strategy.Strategy] = self.get_strategy_by_name() + + debug_print = self.debug_print + + class TrackingStrategy(BaseStrategy): # type: ignore[misc] + """ + Wraps a Flower Strategy and adds: + - per-client fit/eval history + - per-round summary history + - global weights fingerprints after aggregation + - robust logging (shows in simulation logs) + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + # cid -> list of dicts per round + self.client_fit_history: Dict[str, list] = {} + self.client_eval_history: Dict[str, list] = {} + + # list of dicts (round summaries) + self.round_history: list = [] + + # previous global model (flattened) for delta/cos sim + self._prev_global_flat: Optional[np.ndarray] = None + + self.latest_aggregated_parameters = None + + # logger (Flower config will handle output) + self._log = logging.getLogger("flwr") + self._debug_print = debug_print + + # prove that THIS strategy is active + if self._debug_print: + self._log.warning("✅ TrackingStrategy is ACTIVE (debug enabled)") + + # ---------- helpers ---------- + def _params_to_flat(self, parameters: fl.common.Parameters) -> np.ndarray: + nds = fl.common.parameters_to_ndarrays(parameters) + if not nds: + return np.array([], dtype=np.float64) + return np.concatenate([x.reshape(-1).astype(np.float64, copy=False) for x in nds]) + + def _fingerprint(self, flat: np.ndarray) -> Dict[str, Any]: + if flat.size == 0: + return { + "w_norm": 0.0, + "nan_or_inf": False, + "delta_norm": None, + "cos_sim_prev": None, + } + + w_norm = float(np.linalg.norm(flat)) + nan_or_inf = bool(np.isnan(flat).any() or np.isinf(flat).any()) + + delta_norm = None + cos_sim = None + if self._prev_global_flat is not None and self._prev_global_flat.size == flat.size: + diff = flat - self._prev_global_flat + delta_norm = float(np.linalg.norm(diff)) + + denom = np.linalg.norm(flat) * np.linalg.norm(self._prev_global_flat) + if denom > 0: + cos_sim = float(np.dot(flat, self._prev_global_flat) / denom) + + return { + "w_norm": w_norm, + "nan_or_inf": nan_or_inf, + "delta_norm": delta_norm, + "cos_sim_prev": cos_sim, + } + + def _safe_failure_summaries(self, failures: list, cap: int = 5) -> List[str]: + return [str(f) for f in failures][:cap] + + # ---------- overrides ---------- + def aggregate_fit(self, server_round, results, failures): + # Store per-client fit history WITH round + num_examples + cids: List[str] = [] + num_examples_list: List[int] = [] + + for client_proxy, fit_res in results: + cid = getattr(client_proxy, "cid", "unknown") + cids.append(cid) + + ne = int(getattr(fit_res, "num_examples", 0)) + num_examples_list.append(ne) + + entry = { + "round": int(server_round), + "num_examples": ne, + **dict(fit_res.metrics or {}), + } + # store client weights only for regression runs + if fit_res.metrics is not None and "train_rmse" in fit_res.metrics: + client_nds = fl.common.parameters_to_ndarrays(fit_res.parameters) + entry["weights"] = [arr.copy() for arr in client_nds] + self.client_fit_history.setdefault(cid, []).append(entry) + + # Call base aggregation + agg = super().aggregate_fit(server_round, results, failures) + + + # Round summary + global fingerprint (only if aggregation produced parameters) + if agg is not None: + parameters_agg, metrics_agg = agg # type: ignore[misc] + self.latest_aggregated_parameters = (parameters_agg) + flat = self._params_to_flat(parameters_agg) + fp = self._fingerprint(flat) + self._prev_global_flat = flat + + nds = fl.common.parameters_to_ndarrays(parameters_agg) + + round_row = { + "round": int(server_round), + "phase": "fit", + "n_results": int(len(results)), + "n_failures": int(len(failures)), + "cids": cids, + "total_examples": int(sum(num_examples_list)), + "min_examples": int(min(num_examples_list)) if num_examples_list else 0, + "max_examples": int(max(num_examples_list)) if num_examples_list else 0, + "failures": self._safe_failure_summaries(failures), + "weights": [arr.copy() for arr in nds], # <- add this + **fp, + **(dict(metrics_agg or {})), + } + + self.round_history.append(round_row) + + if self._debug_print: + self._log.warning( + f"[fit r={server_round}] " + f"results={round_row['n_results']} failures={round_row['n_failures']} " + f"total_ex={round_row['total_examples']} " + f"w_norm={round_row['w_norm']:.4g} " + f"delta={round_row['delta_norm']} " + f"cos={round_row['cos_sim_prev']} " + f"nan/inf={round_row['nan_or_inf']}" + ) + + return agg + + def aggregate_evaluate(self, server_round, results, failures): + cids: List[str] = [] + num_examples_list: List[int] = [] + + for client_proxy, eval_res in results: + cid = getattr(client_proxy, "cid", "unknown") + cids.append(cid) + + ne = int(getattr(eval_res, "num_examples", 0)) + num_examples_list.append(ne) + + entry = { + "round": int(server_round), + "num_examples": ne, + "loss": float(getattr(eval_res, "loss", float("nan"))), + **dict(eval_res.metrics or {}), + } + self.client_eval_history.setdefault(cid, []).append(entry) + + agg = super().aggregate_evaluate(server_round, results, failures) + + if agg is not None: + loss_agg, metrics_agg = agg # type: ignore[misc] + round_row = { + "round": int(server_round), + "phase": "eval", + "n_results": int(len(results)), + "n_failures": int(len(failures)), + "cids": cids, + "total_examples": int(sum(num_examples_list)), + "min_examples": int(min(num_examples_list)) if num_examples_list else 0, + "max_examples": int(max(num_examples_list)) if num_examples_list else 0, + "failures": self._safe_failure_summaries(failures), + "loss": float(loss_agg) if loss_agg is not None else None, + **(dict(metrics_agg or {})), + } + self.round_history.append(round_row) + + if self._debug_print: + self._log.warning( + f"[eval r={server_round}] " + f"results={round_row['n_results']} failures={round_row['n_failures']} " + f"loss={round_row['loss']}" + ) + + return agg + + def get_debug_state(self) -> Dict[str, Any]: + return { + "round_history": self.round_history, + "client_fit_history": self.client_fit_history, + "client_eval_history": self.client_eval_history, + } + + # ---- base kwargs ---- + kwargs: Dict[str, Any] = dict( fraction_fit=self.fraction_fit, fraction_evaluate=self.fraction_evaluate, min_fit_clients=self.min_fit_clients, min_evaluate_clients=self.min_evaluate_clients, min_available_clients=self.min_available_clients, - initial_parameters=fl.common.ndarrays_to_parameters(self.initial_parameters), - evaluate_fn=self.evaluate_fn ) - def get_strategy_by_name(self): - return eval(f"fl.server.strategy.{self.name}") - - + # Initial parameters + if self.initial_parameters is not None: + kwargs["initial_parameters"] = fl.common.ndarrays_to_parameters(self.initial_parameters) + + # Optional evaluate function + if self.evaluate_fn is not None: + kwargs["evaluate_fn"] = self.evaluate_fn + + # ✅ Metrics aggregation (removes warnings and gives aggregated metrics) + if self.enable_metrics_aggregation: + kwargs["fit_metrics_aggregation_fn"] = weighted_avg + kwargs["evaluate_metrics_aggregation_fn"] = weighted_avg + + # User strategy overrides + kwargs.update(self.strategy_kwargs) + + self.strategy_object = TrackingStrategy(**kwargs) + return self.strategy_object \ No newline at end of file diff --git a/MEDfl/LearningManager/utils.py b/MEDfl/LearningManager/utils.py index a5bd4dd..62ff32e 100644 --- a/MEDfl/LearningManager/utils.py +++ b/MEDfl/LearningManager/utils.py @@ -1,6 +1,5 @@ #!/usr/bin/env python3 -import pkg_resources import torch import yaml from sklearn.metrics import * diff --git a/MEDfl/LearningManager/xgboost/__init__.py b/MEDfl/LearningManager/xgboost/__init__.py new file mode 100644 index 0000000..91a6300 --- /dev/null +++ b/MEDfl/LearningManager/xgboost/__init__.py @@ -0,0 +1,15 @@ +""" +Federated XGBoost backend for MEDfl simulation mode. +""" + +from .client import XGBoostSimulationClient +from .config import XGBoostConfig +from .server import XGBoostSimulationServer +from .strategy import XGBoostSimulationStrategy + +__all__ = [ + "XGBoostSimulationClient", + "XGBoostConfig", + "XGBoostSimulationServer", + "XGBoostSimulationStrategy", +] \ No newline at end of file diff --git a/MEDfl/LearningManager/xgboost/client.py b/MEDfl/LearningManager/xgboost/client.py new file mode 100644 index 0000000..0975248 --- /dev/null +++ b/MEDfl/LearningManager/xgboost/client.py @@ -0,0 +1,328 @@ +#!/usr/bin/env python3 + +import logging +import sys +from typing import Any, Dict, List, Optional + +import flwr as fl +import xgboost as xgb +from flwr.common import ( + Code, + EvaluateIns, + EvaluateRes, + FitIns, + FitRes, + GetParametersIns, + GetParametersRes, + GetPropertiesIns, + GetPropertiesRes, + Status, +) +from torch.utils.data import DataLoader + +from .data import SimulationXGBoostDataModule +from .metrics import evaluate_predictions +from .serialization import ( + booster_to_parameters, + parameters_to_booster, +) + + +def _get_logger() -> logging.Logger: + logger = logging.getLogger( + "MEDfl.simulation.xgboost.client" + ) + + if not logger.handlers: + handler = logging.StreamHandler(sys.stdout) + handler.setFormatter( + logging.Formatter( + "[%(asctime)s][%(name)s][%(levelname)s] " + "%(message)s", + datefmt="%H:%M:%S", + ) + ) + + logger.addHandler(handler) + logger.setLevel(logging.INFO) + logger.propagate = False + + return logger + + +class XGBoostSimulationClient(fl.client.Client): + def __init__( + self, + cid: str, + trainloader: DataLoader, + valloader: Optional[DataLoader] = None, + testloader: Optional[DataLoader] = None, + task: str = "binary", + xgb_params: Optional[Dict[str, Any]] = None, + local_num_boost_round: int = 10, + threshold: float = 0.5, + feature_names: Optional[List[str]] = None, + ): + self.cid = str(cid) + self.task = task + self.xgb_params = xgb_params or {} + self.local_num_boost_round = int( + local_num_boost_round + ) + self.threshold = float(threshold) + self.feature_names = feature_names + + self.log = _get_logger() + self.booster: Optional[xgb.Booster] = None + + data_module = SimulationXGBoostDataModule( + trainloader=trainloader, + valloader=valloader, + testloader=testloader, + feature_names=feature_names, + ) + + self.dtrain = data_module.dtrain + self.dval = data_module.dval + self.dtest = data_module.dtest + self.metadata = data_module.metadata + + self._validate() + + self.log.info( + "[Client %s] initialized: " + "train=%s, val=%s, test=%s, features=%s", + self.cid, + self.metadata["num_train"], + self.metadata["num_val"], + self.metadata["num_test"], + self.metadata["num_features"], + ) + + def _validate(self) -> None: + if self.task not in { + "binary", + "regression", + "multiclass", + }: + raise ValueError( + f"Unsupported task: {self.task}" + ) + + if self.dtrain is None: + raise ValueError( + f"Client {self.cid} has no training data." + ) + + if self.local_num_boost_round <= 0: + raise ValueError( + "local_num_boost_round must be greater than zero." + ) + + def get_parameters( + self, + ins: GetParametersIns, + ) -> GetParametersRes: + return GetParametersRes( + status=Status( + code=Code.OK, + message="Success", + ), + parameters=booster_to_parameters( + self.booster + ), + ) + + def fit(self, ins: FitIns) -> FitRes: + self.log.info( + "[Client %s] starting local XGBoost fit", + self.cid, + ) + + received_booster = parameters_to_booster( + ins.parameters + ) + + if received_booster is not None: + self.booster = received_booster + + evals = [] + + if self.dval is not None: + evals.append( + (self.dval, "validation") + ) + + self.booster = xgb.train( + params=self.xgb_params, + dtrain=self.dtrain, + num_boost_round=self.local_num_boost_round, + xgb_model=self.booster, + evals=evals, + verbose_eval=False, + ) + + train_predictions = self.booster.predict( + self.dtrain + ) + + train_metrics = evaluate_predictions( + task=self.task, + y_true=self.dtrain.get_label(), + y_pred_or_prob=train_predictions, + threshold=self.threshold, + num_features=int( + self.metadata["num_features"] + ), + ) + + metrics = { + "backend": "xgboost", + "task": self.task, + "cid": self.cid, + "num_train": int( + self.metadata["num_train"] + ), + } + + for key, value in train_metrics.items(): + metrics[f"train_{key}"] = value + + self.log.info( + "[Client %s] fit completed: %s", + self.cid, + metrics, + ) + + return FitRes( + status=Status( + code=Code.OK, + message="Success", + ), + parameters=booster_to_parameters( + self.booster + ), + num_examples=int( + self.metadata["num_train"] + ), + metrics=metrics, + ) + + def evaluate( + self, + ins: EvaluateIns, + ) -> EvaluateRes: + received_booster = parameters_to_booster( + ins.parameters + ) + + if received_booster is not None: + self.booster = received_booster + + evaluation_matrix = ( + self.dtest + if self.dtest is not None + else self.dval + ) + + if ( + self.booster is None + or evaluation_matrix is None + ): + return EvaluateRes( + status=Status( + code=Code.EVALUATE_NOT_IMPLEMENTED, + message=( + "No booster or evaluation data." + ), + ), + loss=0.0, + num_examples=0, + metrics={}, + ) + + predictions = self.booster.predict( + evaluation_matrix + ) + + eval_metrics = evaluate_predictions( + task=self.task, + y_true=evaluation_matrix.get_label(), + y_pred_or_prob=predictions, + threshold=self.threshold, + num_features=int( + self.metadata["num_features"] + ), + ) + + metrics = { + "backend": "xgboost", + "task": self.task, + "cid": self.cid, + } + + for key, value in eval_metrics.items(): + metrics[f"eval_{key}"] = value + + if self.task == "binary": + loss = float( + eval_metrics.get("logloss", 0.0) + ) + elif self.task == "regression": + loss = float( + eval_metrics.get("rmse", 0.0) + ) + else: + loss = float( + eval_metrics.get("mlogloss", 0.0) + ) + + self.log.info( + "[Client %s] evaluation completed: " + "loss=%s metrics=%s", + self.cid, + loss, + metrics, + ) + + return EvaluateRes( + status=Status( + code=Code.OK, + message="Success", + ), + loss=loss, + num_examples=int( + evaluation_matrix.num_row() + ), + metrics=metrics, + ) + + def get_properties( + self, + ins: GetPropertiesIns, + ) -> GetPropertiesRes: + properties = { + "backend": "xgboost", + "cid": self.cid, + "task": self.task, + "num_train": int( + self.metadata["num_train"] + ), + "num_val": int( + self.metadata["num_val"] + ), + "num_test": int( + self.metadata["num_test"] + ), + "num_features": int( + self.metadata["num_features"] + ), + } + + return GetPropertiesRes( + status=Status( + code=Code.OK, + message="Success", + ), + properties=properties, + ) \ No newline at end of file diff --git a/MEDfl/LearningManager/xgboost/config.py b/MEDfl/LearningManager/xgboost/config.py new file mode 100644 index 0000000..9fd6ba8 --- /dev/null +++ b/MEDfl/LearningManager/xgboost/config.py @@ -0,0 +1,63 @@ +from dataclasses import dataclass, field +from typing import Any, Dict, Optional + + +@dataclass +class XGBoostConfig: + task: str = "binary" + mode: str = "bagging" + local_num_boost_round: int = 10 + params: Dict[str, Any] = field(default_factory=dict) + num_classes: Optional[int] = None + + def build_params(self) -> Dict[str, Any]: + if self.task == "binary": + defaults = { + "objective": "binary:logistic", + "eval_metric": "auc", + "tree_method": "hist", + "max_depth": 6, + "eta": 0.1, + "subsample": 1.0, + "colsample_bytree": 1.0, + } + + elif self.task == "regression": + defaults = { + "objective": "reg:squarederror", + "eval_metric": "rmse", + "tree_method": "hist", + "max_depth": 6, + "eta": 0.1, + "subsample": 1.0, + "colsample_bytree": 1.0, + } + + elif self.task == "multiclass": + if self.num_classes is None or self.num_classes < 2: + raise ValueError( + "num_classes must be provided and >= 2 " + "for multiclass XGBoost." + ) + + defaults = { + "objective": "multi:softprob", + "num_class": self.num_classes, + "eval_metric": "mlogloss", + "tree_method": "hist", + "max_depth": 6, + "eta": 0.1, + "subsample": 1.0, + "colsample_bytree": 1.0, + } + + else: + raise ValueError( + f"Unsupported XGBoost task: {self.task}. " + "Expected binary, regression, or multiclass." + ) + + merged = defaults.copy() + merged.update(self.params or {}) + + return merged \ No newline at end of file diff --git a/MEDfl/LearningManager/xgboost/data.py b/MEDfl/LearningManager/xgboost/data.py new file mode 100644 index 0000000..9333c75 --- /dev/null +++ b/MEDfl/LearningManager/xgboost/data.py @@ -0,0 +1,122 @@ +from typing import List, Optional, Tuple + +import numpy as np +import torch +import xgboost as xgb +from torch.utils.data import DataLoader + + +def dataloader_to_numpy( + loader: DataLoader, +) -> Tuple[np.ndarray, np.ndarray]: + if loader is None: + raise ValueError("DataLoader cannot be None.") + + x_batches = [] + y_batches = [] + + for x_batch, y_batch in loader: + if isinstance(x_batch, torch.Tensor): + x_batch = x_batch.detach().cpu().numpy() + + if isinstance(y_batch, torch.Tensor): + y_batch = y_batch.detach().cpu().numpy() + + x_batch = np.asarray(x_batch) + y_batch = np.asarray(y_batch).reshape(-1) + + x_batches.append(x_batch) + y_batches.append(y_batch) + + if not x_batches: + raise ValueError( + "Cannot convert an empty DataLoader to DMatrix." + ) + + x = np.concatenate(x_batches, axis=0) + y = np.concatenate(y_batches, axis=0) + + if x.ndim != 2: + raise ValueError( + "XGBoost tabular simulation expects a two-dimensional " + f"feature matrix. Received shape={x.shape}." + ) + + return x, y + + +def dataloader_to_dmatrix( + loader: Optional[DataLoader], + feature_names: Optional[List[str]] = None, +) -> Optional[xgb.DMatrix]: + if loader is None: + return None + + try: + if len(loader.dataset) == 0: + return None + except TypeError: + pass + + x, y = dataloader_to_numpy(loader) + + if feature_names is not None: + if len(feature_names) != x.shape[1]: + raise ValueError( + "feature_names length does not match the number " + f"of features: {len(feature_names)} != {x.shape[1]}" + ) + + return xgb.DMatrix( + data=x, + label=y, + feature_names=feature_names, + ) + + +class SimulationXGBoostDataModule: + def __init__( + self, + trainloader: DataLoader, + valloader: Optional[DataLoader] = None, + testloader: Optional[DataLoader] = None, + feature_names: Optional[List[str]] = None, + ): + self.dtrain = dataloader_to_dmatrix( + trainloader, + feature_names, + ) + + self.dval = dataloader_to_dmatrix( + valloader, + feature_names, + ) + + self.dtest = dataloader_to_dmatrix( + testloader, + feature_names, + ) + + self.metadata = { + "num_train": ( + self.dtrain.num_row() + if self.dtrain is not None + else 0 + ), + "num_val": ( + self.dval.num_row() + if self.dval is not None + else 0 + ), + "num_test": ( + self.dtest.num_row() + if self.dtest is not None + else 0 + ), + "num_features": ( + self.dtrain.num_col() + if self.dtrain is not None + else 0 + ), + "feature_names": feature_names or [], + } \ No newline at end of file diff --git a/MEDfl/LearningManager/xgboost/metrics.py b/MEDfl/LearningManager/xgboost/metrics.py new file mode 100644 index 0000000..0df8827 --- /dev/null +++ b/MEDfl/LearningManager/xgboost/metrics.py @@ -0,0 +1,192 @@ +import math +from typing import Any, Dict + +import numpy as np +from sklearn.metrics import ( + accuracy_score, + f1_score, + log_loss, + mean_absolute_error, + mean_squared_error, + r2_score, + roc_auc_score, +) + + +def _safe_float(value: Any) -> float: + try: + value = float(value) + + if math.isnan(value) or math.isinf(value): + return float("nan") + + return value + except Exception: + return float("nan") + + +def binary_metrics( + y_true, + y_prob, + threshold: float = 0.5, +) -> Dict[str, float]: + y_true = np.asarray(y_true) + y_prob = np.asarray(y_prob) + y_pred = (y_prob >= threshold).astype(int) + + metrics = { + "accuracy": _safe_float( + accuracy_score(y_true, y_pred) + ), + } + + try: + metrics["auc"] = _safe_float( + roc_auc_score(y_true, y_prob) + ) + except Exception: + metrics["auc"] = float("nan") + + try: + metrics["logloss"] = _safe_float( + log_loss(y_true, y_prob) + ) + except Exception: + metrics["logloss"] = float("nan") + + return metrics + + +def regression_metrics( + y_true, + y_pred, + num_features: int = 0, +) -> Dict[str, float]: + y_true = np.asarray(y_true) + y_pred = np.asarray(y_pred) + + rmse = mean_squared_error( + y_true, + y_pred, + squared=False, + ) + mae = mean_absolute_error(y_true, y_pred) + r2 = r2_score(y_true, y_pred) + + n = len(y_true) + p = int(num_features) + + if n > p + 1: + adjusted_r2 = ( + 1.0 + - (1.0 - r2) + * ((n - 1.0) / (n - p - 1.0)) + ) + else: + adjusted_r2 = float("nan") + + return { + "rmse": _safe_float(rmse), + "mae": _safe_float(mae), + "r2": _safe_float(r2), + "adjusted_r2": _safe_float(adjusted_r2), + } + + +def multiclass_metrics( + y_true, + y_prob, +) -> Dict[str, float]: + y_true = np.asarray(y_true) + y_prob = np.asarray(y_prob) + y_pred = np.argmax(y_prob, axis=1) + + metrics = { + "accuracy": _safe_float( + accuracy_score(y_true, y_pred) + ), + } + + try: + metrics["macro_f1"] = _safe_float( + f1_score( + y_true, + y_pred, + average="macro", + ) + ) + except Exception: + metrics["macro_f1"] = float("nan") + + try: + metrics["mlogloss"] = _safe_float( + log_loss(y_true, y_prob) + ) + except Exception: + metrics["mlogloss"] = float("nan") + + return metrics + + +def evaluate_predictions( + task: str, + y_true, + y_pred_or_prob, + threshold: float = 0.5, + num_features: int = 0, +) -> Dict[str, float]: + if task == "binary": + return binary_metrics( + y_true, + y_pred_or_prob, + threshold=threshold, + ) + + if task == "regression": + return regression_metrics( + y_true, + y_pred_or_prob, + num_features=num_features, + ) + + if task == "multiclass": + return multiclass_metrics( + y_true, + y_pred_or_prob, + ) + + raise ValueError(f"Unsupported task: {task}") + + +def aggregate_weighted_metrics(results): + aggregated = {} + + keys = set() + for _, metrics in results: + keys.update(metrics.keys()) + + for key in keys: + weighted_sum = 0.0 + contributing_examples = 0 + + for num_examples, metrics in results: + if key not in metrics: + continue + + try: + value = float(metrics[key]) + except Exception: + continue + + if math.isnan(value) or math.isinf(value): + continue + + weighted_sum += num_examples * value + contributing_examples += num_examples + + if contributing_examples > 0: + aggregated[key] = ( + weighted_sum / contributing_examples + ) + + return aggregated \ No newline at end of file diff --git a/MEDfl/LearningManager/xgboost/serialization.py b/MEDfl/LearningManager/xgboost/serialization.py new file mode 100644 index 0000000..58bbacf --- /dev/null +++ b/MEDfl/LearningManager/xgboost/serialization.py @@ -0,0 +1,52 @@ +from typing import Optional + +import xgboost as xgb +from flwr.common import Parameters + + +XGBOOST_TENSOR_TYPE = "xgboost.ubj" + + +def booster_to_parameters( + booster: Optional[xgb.Booster], +) -> Parameters: + if booster is None: + return Parameters( + tensors=[], + tensor_type=XGBOOST_TENSOR_TYPE, + ) + + raw = booster.save_raw(raw_format="ubj") + + return Parameters( + tensors=[bytes(raw)], + tensor_type=XGBOOST_TENSOR_TYPE, + ) + + +def parameters_to_booster( + parameters: Optional[Parameters], +) -> Optional[xgb.Booster]: + if parameters is None: + return None + + if not parameters.tensors: + return None + + booster = xgb.Booster() + booster.load_model(bytearray(parameters.tensors[0])) + + return booster + + +def save_booster( + booster: xgb.Booster, + path: str, +) -> None: + booster.save_model(path) + + +def load_booster(path: str) -> xgb.Booster: + booster = xgb.Booster() + booster.load_model(path) + return booster \ No newline at end of file diff --git a/MEDfl/LearningManager/xgboost/server.py b/MEDfl/LearningManager/xgboost/server.py new file mode 100644 index 0000000..93e4d1e --- /dev/null +++ b/MEDfl/LearningManager/xgboost/server.py @@ -0,0 +1,746 @@ +#!/usr/bin/env python3 + +import logging +from typing import Any, Dict, List, Optional + +import flwr as fl + +from MEDfl.LearningManager.federated_dataset import FederatedDataset +from MEDfl.LearningManager.server import setup_fl_logging + +from .client import XGBoostSimulationClient +from .strategy import XGBoostSimulationStrategy + +from MEDfl.LearningManager.shap import ( + SHAPConfig, + LocalXGBoostSHAPExplainer, + aggregate_federated_shap, +) + + +class XGBoostSimulationServer: + """ + Flower/Ray simulation server for federated XGBoost. + + Important: + Ray must serialize client_fn and everything captured by it. + Therefore client_fn must not capture the server instance, a SQLAlchemy + connection, a DatabaseManager, or another non-serializable object. + """ + + backend = "xgboost" + + def __init__( + self, + strategy: XGBoostSimulationStrategy, + num_rounds: int, + fed_dataset: FederatedDataset, + task: str = "binary", + xgb_params: Optional[Dict[str, Any]] = None, + local_num_boost_round: int = 10, + threshold: float = 0.5, + feature_names: Optional[List[str]] = None, + client_resources: Optional[Dict[str, float]] = None, + shap_config: Optional[SHAPConfig] = None, + ): + self.strategy = strategy + self.num_rounds = int(num_rounds) + self.fed_dataset = fed_dataset + + self.task = str(task) + self.xgb_params = dict(xgb_params or {}) + self.local_num_boost_round = int( + local_num_boost_round + ) + self.threshold = float(threshold) + + self.feature_names = ( + list(feature_names) + if feature_names is not None + else getattr( + fed_dataset, + "feature_names", + None, + ) + ) + + self.client_resources = ( + dict(client_resources) + if client_resources is not None + else { + "num_cpus": 1.0, + "num_gpus": 0.0, + } + ) + + self.shap_config = ( + shap_config + if shap_config is not None + else SHAPConfig(enabled=False) + ) + + self.num_clients = len( + self.fed_dataset.trainloaders + ) + + self.history = None + self.shap_results = None + + self._validate() + + self.strategy.strategy_object.min_available_clients = ( + self.num_clients + ) + + def _validate(self) -> None: + if self.strategy.strategy_object is None: + raise ValueError( + "strategy.create_strategy() must be called before " + "creating XGBoostSimulationServer." + ) + + if self.num_rounds <= 0: + raise ValueError("num_rounds must be greater than zero.") + + if self.num_clients <= 0: + raise ValueError( + "The federated dataset does not contain training clients." + ) + + if len(self.fed_dataset.valloaders) < self.num_clients: + raise ValueError( + "A validation-loader entry is required for every client." + ) + + if self.task not in { + "binary", + "regression", + "multiclass", + }: + raise ValueError( + f"Unsupported task '{self.task}'." + ) + + @staticmethod + def _clean_loader(loader): + """ + Return None for an unavailable or empty loader. + """ + + if loader is None: + return None + + try: + if len(loader.dataset) == 0: + return None + except Exception: + pass + + return loader + + def _detach_database_state(self) -> None: + """ + Remove live database objects before Ray serializes simulation data. + + The loaders are already in memory, so the XGBoost simulation no longer + requires the database connection. + """ + + if hasattr(self.fed_dataset, "eng"): + try: + connection = self.fed_dataset.eng + + if connection is not None: + try: + connection.close() + except Exception: + pass + finally: + self.fed_dataset.eng = None + + # Defensive cleanup in case these attributes are added later. + for attribute_name in ( + "engine", + "connection", + "db_manager", + "database_manager", + ): + if hasattr(self.fed_dataset, attribute_name): + setattr( + self.fed_dataset, + attribute_name, + None, + ) + + @staticmethod + def _get_loader_input_size(loader) -> int: + if loader is None: + raise ValueError( + "Cannot determine the input size from a None loader" + ) + + for batch in loader: + if not isinstance(batch, (tuple, list)): + raise TypeError( + "Expected each DataLoader batch to be a " + "tuple or list" + ) + + if len(batch) == 0: + continue + + features = batch[0] + + if hasattr(features, "shape"): + shape = features.shape + else: + features = np.asarray(features) + shape = features.shape + + if len(shape) != 2: + raise ValueError( + "XGBoost SHAP expects tabular input with shape " + "(samples, features). " + f"Received shape {shape}." + ) + + return int(shape[1]) + + raise ValueError( + "Cannot determine input size from an empty DataLoader" + ) + + def _get_shap_loaders(self): + """ + Return the selected local partition for every client. + """ + + data_split = self.shap_config.data_split + + if data_split == "train": + source_loaders = self.fed_dataset.trainloaders + + elif data_split == "validation": + source_loaders = self.fed_dataset.valloaders + + elif data_split == "test": + source_loaders = self.fed_dataset.testloaders + + else: + raise ValueError( + f"Unsupported SHAP data split: {data_split}" + ) + + selected_loaders = [] + + for client_index in range(self.num_clients): + loader = None + + if client_index < len(source_loaders): + loader = self._clean_loader( + source_loaders[client_index] + ) + + selected_loaders.append(loader) + + return selected_loaders + + def _resolve_shap_feature_names( + self, + input_size: int, + ) -> List[str]: + """ + Resolve feature names in this order: + + 1. SHAPConfig.feature_names + 2. XGBoost server feature_names + 3. FederatedDataset feature_names + 4. Automatically generated feature_0, feature_1, ... + """ + + if self.shap_config.feature_names is not None: + feature_names = list( + self.shap_config.feature_names + ) + + elif self.feature_names is not None: + feature_names = list( + self.feature_names + ) + + else: + dataset_feature_names = getattr( + self.fed_dataset, + "feature_names", + None, + ) + + if dataset_feature_names is not None: + feature_names = list( + dataset_feature_names + ) + else: + feature_names = ( + self.shap_config.resolved_feature_names( + input_size + ) + ) + + if len(feature_names) != input_size: + raise ValueError( + "The number of XGBoost SHAP feature names must " + "match the input size. " + f"Expected {input_size}, " + f"received {len(feature_names)}." + ) + + return feature_names + + def _run_federated_shap(self): + """ + Run the optional post-training federated SHAP phase. + + The same final global XGBoost booster is explained on each + participating client's selected local data partition. + """ + + if not self.shap_config.enabled: + print( + "[XGBoost/SHAP] SHAP is disabled.", + flush=True, + ) + + self.shap_results = None + return None + + final_booster = self.strategy.current_booster + + if final_booster is None: + raise RuntimeError( + "Cannot calculate XGBoost SHAP because the " + "federated strategy does not contain a final booster." + ) + + selected_loaders = self._get_shap_loaders() + + first_available_loader = None + + for loader in selected_loaders: + if loader is not None: + first_available_loader = loader + break + + if first_available_loader is None: + raise ValueError( + "No client has data for the selected SHAP split: " + f"{self.shap_config.data_split}" + ) + + input_size = self._get_loader_input_size( + first_available_loader + ) + + feature_names = self._resolve_shap_feature_names( + input_size + ) + + self.shap_config.validate_backend( + input_size=input_size, + backend="xgboost", + task=self.task, + ) + + local_explainer = LocalXGBoostSHAPExplainer( + background_size=( + self.shap_config.background_size + ), + explanation_size=( + self.shap_config.explanation_size + ), + random_seed=( + self.shap_config.random_seed + ), + clipping_value=( + self.shap_config.clipping_value + ), + model_output=( + self.shap_config.model_output + ), + ) + + client_results = [] + skipped_clients = [] + + print( + "[XGBoost/SHAP] Starting federated SHAP:", + { + "explainer": self.shap_config.explainer, + "data_split": self.shap_config.data_split, + "background_size": ( + self.shap_config.background_size + ), + "explanation_size": ( + self.shap_config.explanation_size + ), + "minimum_samples": ( + self.shap_config.minimum_samples + ), + "model_output": ( + self.shap_config.model_output + ), + "num_clients": self.num_clients, + }, + flush=True, + ) + + for client_index, loader in enumerate( + selected_loaders + ): + client_id = str(client_index) + + if loader is None: + skipped_clients.append( + { + "client_id": client_id, + "reason": ( + "Selected DataLoader is unavailable or empty" + ), + } + ) + + print( + f"[XGBoost/SHAP] Skipping client " + f"{client_id}: no data.", + flush=True, + ) + + continue + + try: + available_samples = len(loader.dataset) + except Exception: + available_samples = None + + if ( + available_samples is not None + and available_samples + < self.shap_config.minimum_samples + ): + skipped_clients.append( + { + "client_id": client_id, + "reason": ( + "Insufficient samples" + ), + "available_samples": int( + available_samples + ), + "minimum_samples": int( + self.shap_config.minimum_samples + ), + } + ) + + print( + f"[XGBoost/SHAP] Skipping client " + f"{client_id}: available={available_samples}, " + f"minimum={self.shap_config.minimum_samples}.", + flush=True, + ) + + continue + + print( + f"[XGBoost/SHAP] Calculating SHAP for " + f"client {client_id}.", + flush=True, + ) + + try: + local_result = local_explainer.calculate( + booster=final_booster, + loader=loader, + ) + + if ( + local_result["sample_count"] + < self.shap_config.minimum_samples + ): + skipped_clients.append( + { + "client_id": client_id, + "reason": ( + "Explained sample count is below " + "minimum_samples" + ), + "explained_samples": int( + local_result[ + "sample_count" + ] + ), + "minimum_samples": int( + self.shap_config.minimum_samples + ), + } + ) + + continue + + local_result["client_id"] = client_id + local_result["feature_names"] = list( + feature_names + ) + + client_results.append(local_result) + + print( + f"[XGBoost/SHAP] Client {client_id} " + f"completed: explained_samples=" + f"{local_result['sample_count']}.", + flush=True, + ) + + except Exception as exc: + skipped_clients.append( + { + "client_id": client_id, + "reason": str(exc), + } + ) + + print( + f"[XGBoost/SHAP] Client {client_id} failed: " + f"{exc}", + flush=True, + ) + + if not client_results: + self.shap_results = { + "status": "failed", + "backend": "xgboost", + "explainer": "tree", + "data_split": ( + self.shap_config.data_split + ), + "participating_clients": 0, + "explained_samples": 0, + "feature_importance": [], + "skipped_clients": skipped_clients, + "error": ( + "No XGBoost client produced a valid " + "local SHAP result" + ), + } + + print( + "[XGBoost/SHAP] Federated SHAP failed: " + "no valid client results.", + flush=True, + ) + + return self.shap_results + + federated_result = aggregate_federated_shap( + client_results=client_results, + include_client_results=( + self.shap_config.include_client_results + ), + ) + + federated_result.update( + { + "backend": "xgboost", + "task": self.task, + "explainer": "tree", + "data_split": ( + self.shap_config.data_split + ), + "background_size": int( + self.shap_config.background_size + ), + "requested_explanation_size": int( + self.shap_config.explanation_size + ), + "minimum_samples": int( + self.shap_config.minimum_samples + ), + "random_seed": int( + self.shap_config.random_seed + ), + "model_output": ( + self.shap_config.model_output + ), + "feature_names": list( + feature_names + ), + "skipped_clients": skipped_clients, + } + ) + + self.shap_results = federated_result + + print( + "[XGBoost/SHAP] Federated SHAP completed:", + { + "participating_clients": ( + federated_result[ + "participating_clients" + ] + ), + "explained_samples": ( + federated_result[ + "explained_samples" + ] + ), + "skipped_clients": len( + skipped_clients + ), + }, + flush=True, + ) + + print( + "[XGBoost/SHAP] Results:", + federated_result, + flush=True, + ) + + return federated_result + + def run(self): + setup_fl_logging() + + self._detach_database_state() + + # ----------------------------------------------------------- + # Extract only values that Ray can serialize. + # + # Do not pass self.client_fn because that bound method captures + # the complete server object, including fed_dataset and any DB state. + # ----------------------------------------------------------- + + trainloaders = list( + self.fed_dataset.trainloaders + ) + valloaders = list( + self.fed_dataset.valloaders + ) + testloaders = list( + self.fed_dataset.testloaders + ) + + task = self.task + xgb_params = dict(self.xgb_params) + local_num_boost_round = int( + self.local_num_boost_round + ) + threshold = float(self.threshold) + + feature_names = ( + list(self.feature_names) + if self.feature_names is not None + else None + ) + + num_clients = int(self.num_clients) + num_rounds = int(self.num_rounds) + + def client_fn(cid: str): + """ + Ray-serializable client factory. + + This function captures only DataLoaders, primitive values, + dictionaries, and lists. It does not capture the server or + SQLAlchemy objects. + """ + + client_index = int(cid) + + if client_index < 0 or client_index >= num_clients: + raise IndexError( + f"Invalid simulated client ID {cid}. " + f"Expected an index in [0, {num_clients - 1}]." + ) + + trainloader = trainloaders[client_index] + + valloader = None + if client_index < len(valloaders): + valloader = XGBoostSimulationServer._clean_loader( + valloaders[client_index] + ) + + testloader = None + if client_index < len(testloaders): + testloader = XGBoostSimulationServer._clean_loader( + testloaders[client_index] + ) + + client = XGBoostSimulationClient( + cid=str(cid), + trainloader=trainloader, + valloader=valloader, + testloader=testloader, + task=task, + xgb_params=xgb_params, + local_num_boost_round=local_num_boost_round, + threshold=threshold, + feature_names=feature_names, + ) + + # XGBoostSimulationClient already subclasses fl.client.Client. + return client + + print( + "[XGBoostSimulationServer] Starting simulation:", + { + "num_clients": num_clients, + "num_rounds": num_rounds, + "task": task, + "local_num_boost_round": local_num_boost_round, + "client_resources": self.client_resources, + }, + flush=True, + ) + + self.history = fl.simulation.start_simulation( + client_fn=client_fn, + num_clients=num_clients, + config=fl.server.ServerConfig( + num_rounds=num_rounds, + ), + strategy=self.strategy.strategy_object, + client_resources=self.client_resources, + ray_init_args={ + "include_dashboard": False, + "log_to_driver": True, + "logging_level": logging.INFO, + }, + ) + + # ----------------------------------------------------------- + # Optional post-training federated SHAP + # ----------------------------------------------------------- + + if self.shap_config.enabled: + try: + self._run_federated_shap() + + except Exception as exc: + self.shap_results = { + "status": "failed", + "backend": "xgboost", + "explainer": "tree", + "data_split": ( + self.shap_config.data_split + ), + "error": str(exc), + } + + print( + "[XGBoost/SHAP] Federated SHAP failed:", + exc, + flush=True, + ) + + return self.history \ No newline at end of file diff --git a/MEDfl/LearningManager/xgboost/strategy.py b/MEDfl/LearningManager/xgboost/strategy.py new file mode 100644 index 0000000..c5e7cb5 --- /dev/null +++ b/MEDfl/LearningManager/xgboost/strategy.py @@ -0,0 +1,576 @@ +import json +import os +import tempfile +from typing import Any, Dict, Optional + +import flwr as fl +import xgboost as xgb +from flwr.common import Parameters + +from .config import XGBoostConfig +from .metrics import aggregate_weighted_metrics +from .serialization import ( + XGBOOST_TENSOR_TYPE, + booster_to_parameters, + parameters_to_booster, + save_booster, +) + +class XGBoostSimulationStrategy: + """ + Strategy wrapper for MEDfl real-world federated XGBoost. + + This class follows the same high-level contract as the existing MEDfl + Strategy wrapper: + - it exposes create_strategy(); + - it stores the Flower strategy in self.strategy_object; + - it can be passed directly into FederatedServer. + + First supported mode: + - bagging-style federated XGBoost using Flower strategy hooks. + + Notes + ----- + This implementation is conservative and non-invasive: + - It does not modify the existing PyTorch FedAvg/FedAdam/FedYogi path. + - It uses Flower Parameters containing serialized XGBoost boosters. + - It saves XGBoost models as .ubj files instead of .npz files. + """ + + def __init__( + self, + mode="bagging", + fraction_fit=1.0, + fraction_evaluate=1.0, + min_fit_clients=2, + min_evaluate_clients=2, + min_available_clients=2, + task="binary", + num_classes=None, + xgb_params=None, + local_num_boost_round=10, + threshold=0.5, + savingPath="", + save_on_rounds=1, + total_rounds=3, + ): + self.name = "FedXgbBagging" + self.mode = mode + + if self.mode != "bagging": + raise ValueError( + "Only mode='bagging' is implemented in this first integration. " + "Cyclic mode should be added after bagging is stable." + ) + + self.fraction_fit = fraction_fit + self.fraction_evaluate = fraction_evaluate + self.min_fit_clients = min_fit_clients + self.min_evaluate_clients = min_evaluate_clients + self.min_available_clients = min_available_clients + + self.task = task + self.num_classes = num_classes + self.threshold = threshold + self.savingPath = savingPath + self.save_on_rounds = save_on_rounds + + config = XGBoostConfig( + task=task, + mode=mode, + local_num_boost_round=local_num_boost_round, + params=xgb_params or {}, + num_classes=num_classes, + ) + + self.xgb_params = config.build_params() + self.local_num_boost_round = local_num_boost_round + + self.strategy_object = None + self.current_booster = None + self.current_num_trees = 0 + + self.round_history = [] + self.client_fit_history = {} + self.client_evaluation_history = {} + + + self.total_rounds = total_rounds + + self.strategy_object = None + self.current_booster: Optional[xgb.Booster] = None + self.current_num_trees: int = 0 + + def create_strategy(self) -> None: + def fit_config_fn(server_round: int): + return { + "backend": "xgboost", + "server_round": int(server_round), + "mode": self.mode, + "task": self.task, + "threshold": float(self.threshold), + "local_num_boost_round": int( + self.local_num_boost_round + ), + "xgb_params": json.dumps( + self.xgb_params + ), + } + + def evaluate_config_fn(server_round: int): + return { + "backend": "xgboost", + "server_round": int(server_round), + "task": self.task, + "threshold": float(self.threshold), + } + + strategy = fl.server.strategy.FedAvg( + fraction_fit=self.fraction_fit, + fraction_evaluate=self.fraction_evaluate, + min_fit_clients=self.min_fit_clients, + min_evaluate_clients=( + self.min_evaluate_clients + ), + min_available_clients=( + self.min_available_clients + ), + on_fit_config_fn=fit_config_fn, + on_evaluate_config_fn=( + evaluate_config_fn + ), + fit_metrics_aggregation_fn=( + aggregate_weighted_metrics + ), + evaluate_metrics_aggregation_fn=( + aggregate_weighted_metrics + ), + initial_parameters=Parameters( + tensors=[], + tensor_type=XGBOOST_TENSOR_TYPE, + ), + ) + + self._wrap_aggregate_fit(strategy) + self._wrap_aggregate_evaluate(strategy) + + self.strategy_object = strategy + def _wrap_aggregate_fit(self, strategy) -> None: + """ + Wrap aggregate_fit for XGBoost. + + IMPORTANT: + XGBoost boosters are NOT NumPy tensors. + Therefore we MUST NOT call FedAvg tensor aggregation. + + Current strategy: + - select the booster from the client with the largest dataset + - aggregate only metrics manually + - save the selected booster + """ + + def wrapped_aggregate_fit(server_round, results, failures): + print(f"\n[Server/XGBoost] 🔄 Round {server_round} - Client Training Metrics:") + + for client, fit_res in results: + print( + f" XGB CTM Round {server_round} " + f"Client:{client.cid}: {fit_res.metrics}" + ) + + if failures: + print(f"[Server/XGBoost] ⚠️ Round {server_round} failures: {failures}") + + if not results: + print("[Server/XGBoost] No fit results received.") + return None, {} + + # Aggregate metrics manually, same as before. + metric_results = [] + + for _, fit_res in results: + metric_results.append( + ( + fit_res.num_examples, + fit_res.metrics, + ) + ) + + metrics = aggregate_weighted_metrics(metric_results) + + previous_num_trees = int(self.current_num_trees) + new_local_boosters = [] + + print( + f"[Server/XGBoost] Previous global tree count before round " + f"{server_round}: {previous_num_trees}" + ) + + for client, fit_res in results: + client_booster = parameters_to_booster(fit_res.parameters) + + if client_booster is None: + print( + f"[Server/XGBoost] Client {client.cid} returned no booster. Skipping." + ) + continue + + client_total_trees = self._num_boosted_rounds(client_booster) + + print( + f"[Server/XGBoost] Client {client.cid} returned booster with " + f"{client_total_trees} trees." + ) + + # In round 1, previous_num_trees is 0, so this keeps all local trees. + # In later rounds, this keeps only the new trees added by the client. + new_trees_booster = self._slice_new_trees( + client_booster=client_booster, + previous_num_trees=previous_num_trees, + ) + + if new_trees_booster is None: + continue + + new_tree_count = self._num_boosted_rounds(new_trees_booster) + + print( + f"[Server/XGBoost] Client {client.cid} contributed " + f"{new_tree_count} new trees." + ) + + new_local_boosters.append(new_trees_booster) + + if not new_local_boosters: + print( + "[Server/XGBoost] No new local trees received. " + "Keeping previous global booster." + ) + + if self.current_booster is None: + return None, metrics + + agg_parameters = booster_to_parameters(self.current_booster) + return agg_parameters, metrics + + boosters_to_combine = [] + + if self.current_booster is not None: + boosters_to_combine.append(self.current_booster) + + boosters_to_combine.extend(new_local_boosters) + + combined_booster = self._combine_boosters_by_json(boosters_to_combine) + + if combined_booster is None: + print("[Server/XGBoost] Failed to combine boosters.") + return None, metrics + + self.current_booster = combined_booster + self.current_num_trees = self._num_boosted_rounds(combined_booster) + + print( + f"[Server/XGBoost] ✅ Round {server_round} - Flower-style bagging aggregation done." + ) + print( + f"[Server/XGBoost] Global booster now has {self.current_num_trees} trees." + ) + print(f"[Server/XGBoost] Aggregated metrics: {metrics}\n") + + # Optional: print importance directly, no saving required. + if server_round == self.total_rounds: + self._print_feature_importance(combined_booster, server_round) + + agg_parameters = booster_to_parameters(combined_booster) + + self.round_history.append( + { + "round": int(server_round), + "num_fit_clients": len(results), + "num_failures": len(failures), + "global_num_trees": int( + self.current_num_trees + ), + **metrics, + } +) + + # Optional: keep your existing save behavior if savingPath is provided. + self._save_if_needed(server_round, agg_parameters) + + return agg_parameters, metrics + strategy.aggregate_fit = wrapped_aggregate_fit + + + def _wrap_aggregate_evaluate(self, strategy) -> None: + """ + Wrap aggregate_evaluate for logging. + """ + + original_aggregate_evaluate = strategy.aggregate_evaluate + + def wrapped_aggregate_evaluate(server_round, results, failures): + print(f"\n[Server/XGBoost] 📊 Round {server_round} - Client Evaluation Metrics:") + + for client, eval_res in results: + print( + f" XGB CEM Round {server_round} " + f"Client:{client.cid}: {eval_res.metrics}" + ) + + loss, metrics = original_aggregate_evaluate(server_round, results, failures) + + print(f"[Server/XGBoost] ✅ Round {server_round} - Aggregated Evaluation:") + print(f" Loss: {loss}, Metrics: {metrics}\n") + + return loss, metrics + + strategy.aggregate_evaluate = wrapped_aggregate_evaluate + + def _save_if_needed(self, server_round: int, parameters: Parameters) -> None: + """ + Save XGBoost model as .ubj. + """ + + if not self.savingPath: + return + + should_save = ( + server_round % self.save_on_rounds == 0 + or server_round == self.total_rounds + ) + + if not should_save: + return + + os.makedirs(self.savingPath, exist_ok=True) + + booster = parameters_to_booster(parameters) + + if booster is None: + print("[Server/XGBoost] No booster to save.") + return + + filename = ( + f"round_{server_round}_final_model.ubj" + if server_round == self.total_rounds + else f"round_{server_round}_model.ubj" + ) + + filepath = os.path.join(self.savingPath, filename) + + save_booster(booster, filepath) + + print(f"[Server/XGBoost] 💾 Saved booster to: {filepath}") + def _num_boosted_rounds(self, booster: xgb.Booster) -> int: + """ + Return the number of boosted rounds/trees in the booster. + """ + try: + return int(booster.num_boosted_rounds()) + except Exception: + try: + return len(booster.get_dump()) + except Exception: + return 0 + + + def _slice_new_trees( + self, + client_booster: xgb.Booster, + previous_num_trees: int, + ) -> Optional[xgb.Booster]: + """ + Extract only the trees added by the client in this round. + + Client booster = previous global booster + new local trees. + We keep only new local trees. + """ + + total_trees = self._num_boosted_rounds(client_booster) + + if total_trees <= previous_num_trees: + print( + "[Server/XGBoost] Client booster has no new trees. " + f"total_trees={total_trees}, previous_num_trees={previous_num_trees}" + ) + return None + + try: + return client_booster[previous_num_trees:total_trees] + except Exception as exc: + print(f"[Server/XGBoost] Failed to slice new trees: {exc}") + return None + + + def _booster_to_raw_bytes(self, booster: xgb.Booster) -> bytes: + """ + Serialize booster to raw UBJ bytes. + """ + with tempfile.NamedTemporaryFile(suffix=".ubj", delete=False) as tmp: + tmp_path = tmp.name + + try: + booster.save_model(tmp_path) + with open(tmp_path, "rb") as f: + return f.read() + finally: + try: + os.remove(tmp_path) + except Exception: + pass + + + def _booster_from_raw_bytes(self, raw: bytes) -> Optional[xgb.Booster]: + """ + Deserialize booster from raw UBJ bytes. + """ + if not raw: + return None + + with tempfile.NamedTemporaryFile(suffix=".ubj", delete=False) as tmp: + tmp_path = tmp.name + tmp.write(raw) + + try: + booster = xgb.Booster() + booster.load_model(tmp_path) + return booster + except Exception as exc: + print(f"[Server/XGBoost] Failed to load booster from bytes: {exc}") + return None + finally: + try: + os.remove(tmp_path) + except Exception: + pass + + def _combine_boosters_by_json( + self, + boosters: list[xgb.Booster], + ) -> Optional[xgb.Booster]: + """ + Combine multiple XGBoost boosters by concatenating their trees. + + Fixes: + - updates trees + - updates tree_info + - updates gbtree_model_param["num_trees"] + - updates iteration_indptr + + This is the missing part that caused: + Check failed: model.iteration_indptr.back() == model.param.num_trees + """ + + boosters = [b for b in boosters if b is not None] + + if not boosters: + return None + + if len(boosters) == 1: + return boosters[0] + + model_jsons = [] + + for booster in boosters: + raw_json = booster.save_raw(raw_format="json") + if isinstance(raw_json, bytes): + raw_json = raw_json.decode("utf-8") + model_jsons.append(json.loads(raw_json)) + + base = model_jsons[0] + + base_model = base["learner"]["gradient_booster"]["model"] + base_trees = base_model["trees"] + base_tree_info = base_model["tree_info"] + + # Start from the base model iteration structure + base_iteration_indptr = base_model.get("iteration_indptr", None) + + if base_iteration_indptr is None: + # Safe fallback for binary/regression: one tree per iteration + combined_iteration_indptr = list(range(0, len(base_trees) + 1)) + else: + combined_iteration_indptr = list(base_iteration_indptr) + + next_tree_id = len(base_trees) + + for extra in model_jsons[1:]: + extra_model = extra["learner"]["gradient_booster"]["model"] + + extra_trees = extra_model["trees"] + extra_tree_info = extra_model["tree_info"] + + # Append trees and reassign IDs + for tree in extra_trees: + tree["id"] = next_tree_id + base_trees.append(tree) + next_tree_id += 1 + + # Append tree_info + base_tree_info.extend(extra_tree_info) + + # Append iteration_indptr correctly + extra_iteration_indptr = extra_model.get("iteration_indptr", None) + + if extra_iteration_indptr is None: + # Fallback: one tree per boosting iteration + extra_iteration_indptr = list(range(0, len(extra_trees) + 1)) + + offset = combined_iteration_indptr[-1] + + # Skip the first 0 from the extra booster + for ptr in extra_iteration_indptr[1:]: + combined_iteration_indptr.append(offset + ptr) + + # Write back merged model fields + base_model["trees"] = base_trees + base_model["tree_info"] = base_tree_info + base_model["iteration_indptr"] = combined_iteration_indptr + + # Update num_trees + if "gbtree_model_param" in base_model: + base_model["gbtree_model_param"]["num_trees"] = str(len(base_trees)) + + # Save/load to validate the merged model + with tempfile.NamedTemporaryFile( + mode="w", + suffix=".json", + delete=False, + encoding="utf-8", + ) as tmp: + tmp_path = tmp.name + json.dump(base, tmp) + + try: + combined = xgb.Booster() + combined.load_model(tmp_path) + return combined + finally: + try: + os.remove(tmp_path) + except Exception: + pass + def _print_feature_importance( + self, + booster: xgb.Booster, + server_round: int, + ) -> None: + """ + Print feature importance for the current global booster. + """ + if booster is None: + return + + importance_gain = booster.get_score(importance_type="gain") + importance_weight = booster.get_score(importance_type="weight") + importance_cover = booster.get_score(importance_type="cover") + + print(f"\n[Server/XGBoost] 📊 Round {server_round} Global Feature Importance - Gain:") + print(importance_gain) + + print(f"[Server/XGBoost] 📊 Round {server_round} Global Feature Importance - Weight:") + print(importance_weight) + + print(f"[Server/XGBoost] 📊 Round {server_round} Global Feature Importance - Cover:") + print(importance_cover) \ No newline at end of file diff --git a/MEDfl/NetManager/flsetup.py b/MEDfl/NetManager/flsetup.py index 10567ad..79457bb 100644 --- a/MEDfl/NetManager/flsetup.py +++ b/MEDfl/NetManager/flsetup.py @@ -1,7 +1,7 @@ from datetime import datetime -from torch.utils.data import random_split, DataLoader, Dataset +from torch.utils.data import random_split, DataLoader, Dataset , TensorDataset from MEDfl.LearningManager.federated_dataset import FederatedDataset from .net_helper import * @@ -12,6 +12,11 @@ from MEDfl.NetManager.database_connector import DatabaseManager +import numpy as np +import torch +from sklearn.impute import SimpleImputer +from sklearn.preprocessing import LabelEncoder, StandardScaler +from torch.utils.data import random_split, DataLoader, Dataset, TensorDataset class FLsetup: def __init__(self, name: str, description: str, network: Network): @@ -176,7 +181,7 @@ def create_dataloader_from_node( fill_strategy="mean", fit_encode=[], to_drop=[], train_batch_size: int = 32, test_batch_size: int = 1, - split_frac: float = 0.2, + split_frac: float = 0, dataset: Dataset = None, ): @@ -212,84 +217,205 @@ def create_dataloader_from_node( return trainloader, testloader def create_federated_dataset( - self, output, fill_strategy="mean", fit_encode=[], to_drop=[], val_frac=0.1, test_frac=0.2 + self, + output, + fill_strategy="mean", + fit_encode=[], + to_drop=[], + val_frac=0.2, + test_frac=0.15, + train_batch_size=8, + task_type="binary", + scale_y=None, ) -> FederatedDataset: - """Create a federated dataset. - Args: - output (string): the output feature of the dataset - val_frac (float): The fraction of data to be used for validation. - test_frac (float): The fraction of data to be used for testing. + to_drop = list(to_drop) - Returns: - FederatedDataset: The FederatedDataset instance containing train, validation, and test data. - """ - if not self.column_name: - to_drop.extend(["DataSetName" , "NodeId" , "DataSetId"]) - else : - to_drop.extend(["PatientId"]) - + to_drop.extend(["DataSetName", "NodeId", "DataSetId"]) + else: + to_drop.extend(["PatientId"]) + netid = self.network.id + train_nodes = pd.read_sql( - text( - f"SELECT Nodes.NodeName FROM Nodes WHERE Nodes.NetId = {netid} AND Nodes.train = 1 " - ), + text(f"SELECT Nodes.NodeName FROM Nodes WHERE Nodes.NetId = {netid} AND Nodes.train = 1"), self.eng, ) test_nodes = pd.read_sql( - text( - f"SELECT Nodes.NodeName FROM Nodes WHERE Nodes.NetId = {netid} AND Nodes.train = 0 " - ), + text(f"SELECT Nodes.NodeName FROM Nodes WHERE Nodes.NetId = {netid} AND Nodes.train = 0"), self.eng, ) - train_nodes = [ - Node(val[0], 1, test_frac) for val in train_nodes.values.tolist() - ] + train_nodes = [Node(val[0], 1, test_frac) for val in train_nodes.values.tolist()] test_nodes = [Node(val[0], 0) for val in test_nodes.values.tolist()] trainloaders, valloaders, testloaders = [], [], [] - # if len(test_nodes) == 0: - # raise "test node empty" - if test_nodes is None: - _, testloader = self.create_dataloader_from_node( - train_nodes[0], output, fill_strategy=fill_strategy, fit_encode=fit_encode, to_drop=to_drop) - testloaders.append(testloader) - else: - for train_node in train_nodes: - train_valloader, testloader = self.create_dataloader_from_node( - train_node, output, fill_strategy=fill_strategy, - fit_encode=fit_encode, to_drop=to_drop,) - trainloader, valloader = self.create_dataloader_from_node( - train_node, - output, fill_strategy=fill_strategy, fit_encode=fit_encode, to_drop=to_drop, - test_batch_size=32, - split_frac=val_frac, - dataset=train_valloader.dataset, - ) - trainloaders.append(trainloader) - valloaders.append(valloader) - testloaders.append(testloader) + self.scalers = {} - for test_node in test_nodes: - _, testloader = self.create_dataloader_from_node( - test_node, output, fill_strategy=fill_strategy, fit_encode=fit_encode, to_drop=to_drop, split_frac=1.0 - ) - testloaders.append(testloader) - train_nodes_names = [node.name for node in train_nodes] - test_nodes_names = train_nodes_names + [ - node.name for node in test_nodes - ] + # ----------------------------- + # Pass 1: collect all training raw data + # ----------------------------- + cached_splits = {} + all_train_X = [] + all_train_y = [] + + for train_node in train_nodes: + if self.column_name is not None: + raw_data = train_node.get_dataset(self.column_name) + else: + raw_data = train_node.get_dataset() + + n = len(raw_data) + n_test = max(0, int(n * test_frac)) + n_trainval = n - n_test + + raw_data = raw_data.sample(frac=1, random_state=42).reset_index(drop=True) + trainval_raw = raw_data.iloc[:n_trainval].copy() + test_raw = raw_data.iloc[n_trainval:].copy() + + cached_splits[train_node.name] = { + "trainval_raw": trainval_raw, + "test_raw": test_raw, + } + + tmp = trainval_raw.copy() + + for s in fit_encode: + tmp[s] = LabelEncoder().fit_transform(tmp[s]) + + y_series = pd.to_numeric(tmp[output], errors="coerce") + X_df = tmp.copy() + + for column in to_drop: + if column in X_df.columns: + X_df = X_df.drop([column], axis=1) + + features = [col for col in X_df.columns if col != output] + + imputer = SimpleImputer(strategy=fill_strategy) + X_df[features] = imputer.fit_transform(X_df[features]) + + X_vals = X_df[features].values + y_vals = y_series.values.reshape(-1, 1) + + valid_mask = ~np.isnan(y_vals).ravel() + X_vals = X_vals[valid_mask] + y_vals = y_vals[valid_mask] + + if len(X_vals) > 0: + all_train_X.append(X_vals) + all_train_y.append(y_vals) + + if len(all_train_X) == 0: + raise ValueError("No training data available to fit global scalers") + + # ----------------------------- + # Pass 2: fit one global X scaler and one global y scaler + # ----------------------------- + global_scaler_X = StandardScaler() + global_scaler_X.fit(np.vstack(all_train_X)) + + if scale_y is None: + scale_y = (task_type == "regression") - # test_nodes_names = [ - # node.name for node in test_nodes - # ] + global_scaler_y = None + if task_type == "regression" and scale_y: + if len(all_train_y) == 0: + raise ValueError("No training targets available to fit global y scaler") + global_scaler_y = StandardScaler() + global_scaler_y.fit(np.vstack(all_train_y)) + + # ----------------------------- + # Pass 3: build datasets using the same global scalers + # ----------------------------- + for train_node in train_nodes: + trainval_raw = cached_splits[train_node.name]["trainval_raw"] + test_raw = cached_splits[train_node.name]["test_raw"] + + trainval_dataset, _, _ = process_data_after_reading( + trainval_raw.copy(), + output, + fill_strategy=fill_strategy, + fit_encode=fit_encode, + to_drop=to_drop, + scaler_X=global_scaler_X, + scaler_y=global_scaler_y, + fit_scaler=False, + task_type=task_type, + scale_y=scale_y, + ) + + self.scalers[train_node.name] = { + "X": global_scaler_X, + "y": global_scaler_y, + } + + n_tv = len(trainval_dataset) + n_val = max(0, int(n_tv * val_frac)) + n_tr = n_tv - n_val + + train_ds, val_ds = random_split(trainval_dataset, [n_tr, n_val]) - # Add FlSetup on to the DataBase - # self.create() + if len(test_raw) == 0: + if len(trainval_dataset) > 0: + n_features = len(trainval_dataset[0][0]) + else: + n_features = len(features) + + empty_X = torch.empty((0, n_features), dtype=torch.float32) + empty_y = torch.empty((0,), dtype=torch.float32) + test_dataset = TensorDataset(empty_X, empty_y) + else: + test_dataset, _, _ = process_data_after_reading( + test_raw.copy(), + output, + fill_strategy=fill_strategy, + fit_encode=fit_encode, + to_drop=to_drop, + scaler_X=global_scaler_X, + scaler_y=global_scaler_y, + fit_scaler=False, + task_type=task_type, + scale_y=scale_y, + ) + + trainloaders.append(DataLoader(train_ds, batch_size=train_batch_size, shuffle=True)) + valloaders.append(DataLoader(val_ds, batch_size=8)) + testloaders.append(DataLoader(test_dataset, batch_size=8)) + + print(f"{train_node.name}: train={n_tr}, val={n_val}, test={len(test_dataset)}") + + # test-only nodes + for test_node in test_nodes: + raw_data = test_node.get_dataset() + + if len(raw_data) == 0: + continue + + test_dataset, _, _ = process_data_after_reading( + raw_data.copy(), + output, + fill_strategy=fill_strategy, + fit_encode=fit_encode, + to_drop=to_drop, + scaler_X=global_scaler_X, + scaler_y=global_scaler_y, + fit_scaler=False, + task_type=task_type, + scale_y=scale_y, + ) + + self.scalers[test_node.name] = { + "X": global_scaler_X, + "y": global_scaler_y, + } + + testloaders.append(DataLoader(test_dataset, batch_size=8)) + + train_nodes_names = [node.name for node in train_nodes] + test_nodes_names = train_nodes_names + [node.name for node in test_nodes] - # self.network.update_network(FLsetupId=self.id) fed_dataset = FederatedDataset( self.name + "_Feddataset", train_nodes_names, @@ -298,13 +424,11 @@ def create_federated_dataset( valloaders, testloaders, ) + self.fed_dataset = fed_dataset self.fed_dataset.create(self.id) - return self.fed_dataset - - - + return self.fed_dataset def get_flDataSet(self): """ Retrieve the federated dataset associated with the FL setup using the FL setup's name. diff --git a/MEDfl/NetManager/net_helper.py b/MEDfl/NetManager/net_helper.py index 2d9aff0..3988345 100644 --- a/MEDfl/NetManager/net_helper.py +++ b/MEDfl/NetManager/net_helper.py @@ -9,7 +9,8 @@ import numpy as np from MEDfl.NetManager.database_connector import DatabaseManager - +from sklearn.preprocessing import LabelEncoder, StandardScaler +from sklearn.impute import SimpleImputer def is_str(data_df, row, x): """ @@ -61,63 +62,89 @@ def process_eicu(data_df): # remove indiserd columns after reading from the DB -def process_data_after_reading(data, output, fill_strategy="mean", fit_encode=[], to_drop=[]): - """ - Process data after reading from the database, including encoding, dropping columns, and creating a PyTorch TensorDataset. - - Args: - data (pandas.DataFrame): Input data. - output (str): Output column name. - fill_strategy (str, optional): Imputation strategy for missing values. Default is "mean". - fit_encode (list, optional): List of columns to be label-encoded. Default is an empty list. - to_drop (list, optional): List of columns to be dropped from the DataFrame. Default is an empty list. - - Returns: - torch.utils.data.TensorDataset: Processed data as a PyTorch TensorDataset. - """ - - # Check if there is a DataSet assigned to the node - if (len(data) == 0): - raise "Node doesn't Have dataSet" +def process_data_after_reading( + data, + output, + fill_strategy="mean", + fit_encode=[], + to_drop=[], + scaler_X=None, + scaler_y=None, + fit_scaler=True, + task_type="binary", # NEW + scale_y=None, # NEW +): + if len(data) == 0: + raise ValueError("Node doesn't have a dataset") encoder = LabelEncoder() - # En Code some columns for s in fit_encode: - try: - data[s] = encoder.fit_transform(data[s]) - except: - raise print(s) - - # The output of the DATA - y = data[output] + data[s] = encoder.fit_transform(data[s]) - X = data + y = data[output].copy() + X = data.copy() - # remove indisered columns when reading the dataframe from the DB for column in to_drop: - try: - X = X.drop( - [column], axis=1 - ) - except Exception as e: - raise e + if column in X.columns: + X = X.drop([column], axis=1) - # Get the DATAset Features features = [col for col in X.columns if col != output] - # Impute missing values using the mean strategy - try: - imputer = SimpleImputer(strategy=fill_strategy) - X[features] = imputer.fit_transform(X[features]) - except: - print() + imputer = SimpleImputer(strategy=fill_strategy) + X[features] = imputer.fit_transform(X[features]) - X = torch.tensor(X.values, dtype=torch.float32) - y = torch.tensor(y.values, dtype=torch.float32) - data = TensorDataset(X, y) + if scaler_X is None: + scaler_X = StandardScaler() + + X_vals = X[features].values + y_vals = pd.to_numeric(y, errors="coerce").values.reshape(-1, 1) + + valid_mask = ~np.isnan(y_vals).ravel() + X_vals = X_vals[valid_mask] + y_vals = y_vals[valid_mask] + + if fit_scaler: + X_vals = scaler_X.fit_transform(X_vals) + else: + X_vals = scaler_X.transform(X_vals) + + # default behavior depends on task + if scale_y is None: + scale_y = (task_type == "regression") + + if task_type == "binary": + y_vals = y_vals.ravel().astype(np.float32) + + uniq = set(np.unique(y_vals)) + if uniq == {-1.0, 1.0}: + y_vals = ((y_vals + 1.0) / 2.0).astype(np.float32) + elif uniq == {1.0, 2.0}: + y_vals = (y_vals - 1.0).astype(np.float32) + + uniq = set(np.unique(y_vals)) + if not uniq.issubset({0.0, 1.0}): + raise ValueError( + f"Binary classification requires targets in {{0,1}}. Found: {sorted(uniq)}" + ) + + elif task_type == "regression": + if scale_y: + if scaler_y is None: + scaler_y = StandardScaler() + if fit_scaler: + y_vals = scaler_y.fit_transform(y_vals).ravel() + else: + y_vals = scaler_y.transform(y_vals).ravel() + else: + y_vals = y_vals.ravel().astype(np.float32) + + else: + raise ValueError(f"Unsupported task_type: {task_type}") - return data + X_tensor = torch.tensor(X_vals, dtype=torch.float32) + y_tensor = torch.tensor(y_vals, dtype=torch.float32) + return TensorDataset(X_tensor, y_tensor), scaler_X, scaler_y def get_nodeid_from_name(name): """ diff --git a/MEDfl/rw/client.py b/MEDfl/rw/client.py index 5c60d2f..7e2d226 100644 --- a/MEDfl/rw/client.py +++ b/MEDfl/rw/client.py @@ -1,6 +1,7 @@ # File: client.py import argparse +import json import pandas as pd import flwr as fl import torch @@ -10,6 +11,7 @@ from sklearn.metrics import accuracy_score, roc_auc_score from sklearn.model_selection import train_test_split from MEDfl.rw.model import Net # votre définition de modèle +from MEDfl.LearningManager.shap import LocalPyTorchSHAPExplainer import socket import platform import psutil @@ -293,6 +295,155 @@ def _lazy_init_from_server_config(self, config): self._initialized = True print(f"[Client] Initialized with features={feat_cols}, target={target_col}, val={val_frac}, test={self.effective_test_frac}") + + # ---------- Federated SHAP helpers ---------- + + def _get_shap_loader(self, data_split): + """ + Return the local partition selected for the post-training SHAP phase. + """ + split = str(data_split).lower() + + if split == "train": + loader = self.train_loader + elif split == "validation": + loader = self.val_loader + elif split == "test": + loader = self.test_loader + else: + raise ValueError( + f"Unsupported SHAP data split '{data_split}'. " + "Expected train, validation, or test." + ) + + if loader is None: + raise ValueError( + f"The selected SHAP split '{split}' is not available " + "on this client." + ) + + if len(loader.dataset) == 0: + raise ValueError( + f"The selected SHAP split '{split}' is empty " + "on this client." + ) + + return loader + + def _calculate_local_shap(self, config): + """ + Explain the final global model on this client's local data. + + Only additive SHAP summary statistics are returned. Raw samples and + per-sample SHAP values never leave the client. + """ + if not self._initialized: + self._lazy_init_from_server_config(config) + + data_split = str( + config.get("shap_data_split", "validation") + ) + explainer_type = str( + config.get("shap_explainer", "gradient") + ) + background_size = int( + config.get("shap_background_size", 100) + ) + explanation_size = int( + config.get("shap_explanation_size", 500) + ) + random_seed = int( + config.get("shap_random_seed", 42) + ) + minimum_samples = int( + config.get("shap_minimum_samples", 10) + ) + + clipping_raw = float( + config.get("shap_clipping_value", -1.0) + ) + clipping_value = ( + None if clipping_raw <= 0 else clipping_raw + ) + + feature_names_raw = str( + config.get("shap_feature_names_json", "[]") + ) + try: + configured_feature_names = json.loads( + feature_names_raw + ) + except Exception: + configured_feature_names = [] + + feature_names = ( + list(configured_feature_names) + if configured_feature_names + else list(self.effective_features) + ) + + if len(feature_names) != len(self.effective_features): + raise ValueError( + "SHAP feature names do not match the local model input size. " + f"Expected {len(self.effective_features)}, " + f"received {len(feature_names)}." + ) + + loader = self._get_shap_loader(data_split) + + # Opacus can wrap the original torch.nn.Module in GradSampleModule. + # SHAP should explain the underlying neural network. + shap_model = getattr(self.model, "_module", self.model) + + local_explainer = LocalPyTorchSHAPExplainer( + explainer_type=explainer_type, + background_size=background_size, + explanation_size=explanation_size, + random_seed=random_seed, + clipping_value=clipping_value, + ) + + result = local_explainer.calculate( + model=shap_model, + loader=loader, + device=torch.device("cpu"), + ) + + explained_samples = int(result["sample_count"]) + if explained_samples < minimum_samples: + raise ValueError( + f"Only {explained_samples} samples were explained, " + f"but minimum_samples={minimum_samples}." + ) + + # Convert NumPy arrays to lists before JSON serialization. + serializable_result = { + "client_id": socket.gethostname(), + "feature_names": feature_names, + "available_samples": int(len(loader.dataset)), + "sample_count": explained_samples, + "abs_sum": np.asarray( + result["abs_sum"], dtype=np.float64 + ).tolist(), + "signed_sum": np.asarray( + result["signed_sum"], dtype=np.float64 + ).tolist(), + "square_sum": np.asarray( + result["square_sum"], dtype=np.float64 + ).tolist(), + "positive_count": np.asarray( + result["positive_count"], dtype=np.int64 + ).tolist(), + "negative_count": np.asarray( + result["negative_count"], dtype=np.int64 + ).tolist(), + "zero_count": np.asarray( + result["zero_count"], dtype=np.int64 + ).tolist(), + } + + return serializable_result + # ---------- FL API (unchanged behavior) ---------- def get_parameters(self, config): @@ -390,6 +541,50 @@ def evaluate(self, parameters, config): "eval_accuracy": acc, "eval_auc": auc, } + + # --------------------------------------------------------------- + # Final-round federated SHAP + # + # Flower calls evaluate() with the parameters produced by the + # current round's aggregation. On the last round, these are the + # final global parameters. Therefore every participating client + # explains exactly the same final global model. + # --------------------------------------------------------------- + if bool(config.get("shap_enabled", False)): + print( + "[Client/SHAP] Starting local SHAP on the final " + "global model...", + flush=True, + ) + + try: + local_shap_result = self._calculate_local_shap( + config + ) + + metrics["shap_status"] = "completed" + metrics["shap_result"] = json.dumps( + local_shap_result, + allow_nan=False, + ) + + print( + "[Client/SHAP] Local SHAP completed | " + f"client={local_shap_result['client_id']} | " + f"samples={local_shap_result['sample_count']}", + flush=True, + ) + + except Exception as exc: + metrics["shap_status"] = "failed" + metrics["shap_error"] = str(exc) + + print( + "[Client/SHAP] Local SHAP failed: " + f"{exc}", + flush=True, + ) + print(f"Evaluation metrics: {metrics}") return float(avg_loss), len(self.test_loader.dataset), metrics diff --git a/MEDfl/rw/server.py b/MEDfl/rw/server.py index 178658a..501a089 100644 --- a/MEDfl/rw/server.py +++ b/MEDfl/rw/server.py @@ -2,6 +2,7 @@ from flwr.server.strategy import FedAvg from flwr.server.server import ServerConfig from MEDfl.rw.strategy import Strategy +from MEDfl.LearningManager.shap import SHAPConfig import asyncio from flwr.server.client_manager import ClientManager from flwr.server.client_proxy import ClientProxy @@ -33,6 +34,7 @@ def __init__( num_rounds=3, strategy=None, certificates=None, + shap_config=None, ): """ Initialize the FederatedServer. @@ -50,14 +52,62 @@ def __init__( # Use custom or default strategy self.strategy_wrapper = strategy or Strategy() + + self.shap_config = ( + shap_config + if shap_config is not None + else SHAPConfig(enabled=False) + ) + + # The Strategy owns the final-round SHAP orchestration. + self.strategy_wrapper.shap_config = self.shap_config + + # Keep the strategy's final-round marker synchronized with + # the actual Flower ServerConfig. + self.strategy_wrapper.total_rounds = int(num_rounds) + + # Validate the PyTorch SHAP configuration before the server starts. + if self.shap_config.enabled: + feature_columns = [ + column.strip() + for column in ( + self.strategy_wrapper._features or "" + ).split(",") + if column.strip() + ] + + if not feature_columns: + raise ValueError( + "Federated SHAP requires the Strategy 'features' " + "configuration so the neural-network input size " + "can be validated." + ) + + self.shap_config.validate_backend( + input_size=len(feature_columns), + backend="pytorch", + ) + self.strategy_wrapper.create_strategy() if self.strategy_wrapper.strategy_object is None: - raise ValueError("Strategy object not initialized. Call create_strategy() first.") + raise ValueError( + "Strategy object not initialized. " + "Call create_strategy() first." + ) self.strategy = self.strategy_wrapper.strategy_object self.certificates = certificates self.connected_clients = [] # Track connected client IDs + self.federated_shap_result = { + "status": ( + "pending" + if self.shap_config.enabled + else "disabled" + ) + } + self.local_shap_results = [] + def start(self): """ @@ -78,6 +128,21 @@ def start(self): client_manager=client_manager, ) + # The strategy aggregates SHAP during the final evaluation. + self.federated_shap_result = ( + self.strategy_wrapper.federated_shap_result + ) + self.local_shap_results = ( + self.strategy_wrapper.local_shap_results + ) + + if self.shap_config.enabled: + print( + "[Server/SHAP] Final federated SHAP status: " + f"{self.federated_shap_result.get('status')}", + flush=True, + ) + class TrackingClientManager(fl.server.client_manager.SimpleClientManager): """ diff --git a/MEDfl/rw/strategy.py b/MEDfl/rw/strategy.py index ba1ce42..96d966a 100644 --- a/MEDfl/rw/strategy.py +++ b/MEDfl/rw/strategy.py @@ -1,14 +1,19 @@ # File: MEDfl/rw/strategy.py import os +import json import numpy as np import flwr as fl -from flwr.common import GetPropertiesIns +from flwr.common import GetPropertiesIns, EvaluateIns from flwr.server.client_manager import ClientManager from flwr.server.client_proxy import ClientProxy import time from MEDfl.rw.model import Net import torch +from MEDfl.LearningManager.shap import ( + SHAPConfig, + aggregate_federated_shap, +) # =================================================== # Custom metric aggregation functions @@ -92,6 +97,8 @@ def __init__( client_fractions=None, # NEW: id column for test_ids mapping id_col="id", + # Optional post-training federated SHAP configuration + shap_config=None, ): self.name = name self.fraction_fit = fraction_fit @@ -123,6 +130,16 @@ def __init__( self.client_fractions = client_fractions or {} self.id_col = id_col + self.shap_config = ( + shap_config + if shap_config is not None + else SHAPConfig(enabled=False) + ) + self.federated_shap_result = { + "status": "disabled" + } + self.local_shap_results = [] + self.strategy_object = None def create_strategy(self): @@ -144,6 +161,66 @@ def fit_config_fn(server_round): "id_col": self.id_col, } + def evaluate_config_fn(server_round): + """ + Send the ordinary dataset/schema configuration on every + evaluation. Enable SHAP only on the last federated round. + """ + is_final_round = ( + int(server_round) == int(self.total_rounds) + ) + + cfg = { + "features": self._features, + "target": self._target, + "val_fraction": float(self._val_fraction), + "test_fraction": float(self._test_fraction), + "id_col": self.id_col, + "shap_enabled": bool( + self.shap_config.enabled + and is_final_round + ), + } + + if cfg["shap_enabled"]: + cfg.update( + { + "shap_explainer": ( + self.shap_config.explainer + ), + "shap_data_split": ( + self.shap_config.data_split + ), + "shap_background_size": int( + self.shap_config.background_size + ), + "shap_explanation_size": int( + self.shap_config.explanation_size + ), + "shap_random_seed": int( + self.shap_config.random_seed + ), + "shap_minimum_samples": int( + self.shap_config.minimum_samples + ), + # Flower configs are Scalar-only, so None is + # represented by a negative sentinel. + "shap_clipping_value": ( + -1.0 + if self.shap_config.clipping_value + is None + else float( + self.shap_config.clipping_value + ) + ), + "shap_feature_names_json": json.dumps( + self.shap_config.feature_names or [] + ), + } + ) + + return cfg + # 3) Build params including on_fit_config_fn params = { "fraction_fit": self.fraction_fit, @@ -153,6 +230,7 @@ def fit_config_fn(server_round): "min_available_clients": self.min_available_clients, "evaluate_fn": self.evaluate_fn, "on_fit_config_fn": fit_config_fn, + "on_evaluate_config_fn": evaluate_config_fn, "fit_metrics_aggregation_fn": self.fit_metrics_aggregation_fn, "evaluate_metrics_aggregation_fn": self.evaluate_metrics_aggregation_fn, } @@ -205,16 +283,195 @@ def logged_agg_fit(server_round, results, failures): original_agg_eval = strat.aggregate_evaluate def logged_agg_eval(server_round, results, failures): - print(f"\n[Server] 📊 Round {server_round} - Client Evaluation Metrics:") + print( + f"\n[Server] 📊 Round {server_round} - " + "Client Evaluation Metrics:" + ) for i, (client_id, eval_res) in enumerate(results): - print(f" CEM Round {server_round} Client:{client_id.cid}: {eval_res.metrics}") - loss, metrics = original_agg_eval(server_round, results, failures) - print(f"[Server] ✅ Round {server_round} - Aggregated Evaluation Metrics:") + print( + f" CEM Round {server_round} " + f"Client:{client_id.cid}: {eval_res.metrics}" + ) + + loss, metrics = original_agg_eval( + server_round, + results, + failures, + ) + + print( + f"[Server] ✅ Round {server_round} - " + "Aggregated Evaluation Metrics:" + ) print(f" Loss: {loss}, Metrics: {metrics}\n") + + is_final_round = ( + int(server_round) == int(self.total_rounds) + ) + + if self.shap_config.enabled and is_final_round: + local_results = [] + shap_failures = [] + + for client_proxy, eval_res in results: + client_metrics = eval_res.metrics or {} + + if ( + client_metrics.get("shap_status") + == "completed" + and client_metrics.get("shap_result") + ): + try: + local_result = json.loads( + client_metrics["shap_result"] + ) + local_results.append(local_result) + except Exception as exc: + shap_failures.append( + { + "client_id": client_proxy.cid, + "error": ( + "Invalid SHAP JSON: " + f"{exc}" + ), + } + ) + else: + shap_failures.append( + { + "client_id": client_proxy.cid, + "error": client_metrics.get( + "shap_error", + "No SHAP result returned", + ), + } + ) + + self.local_shap_results = local_results + + if local_results: + try: + federated_result = ( + aggregate_federated_shap( + client_results=local_results, + include_client_results=( + self.shap_config + .include_client_results + ), + ) + ) + + federated_result["explainer"] = ( + self.shap_config.explainer + ) + federated_result["data_split"] = ( + self.shap_config.data_split + ) + federated_result["background_size"] = ( + self.shap_config.background_size + ) + federated_result[ + "requested_explanation_size" + ] = ( + self.shap_config.explanation_size + ) + + if shap_failures: + federated_result[ + "failed_clients" + ] = shap_failures + + self.federated_shap_result = ( + federated_result + ) + + print( + "[Server/SHAP] Federated SHAP " + "completed:" + ) + print( + json.dumps( + self.federated_shap_result, + indent=2, + ), + flush=True, + ) + + except Exception as exc: + self.federated_shap_result = { + "status": "failed", + "error": str(exc), + "failed_clients": shap_failures, + } + print( + "[Server/SHAP] Federated SHAP " + f"aggregation failed: {exc}", + flush=True, + ) + + else: + self.federated_shap_result = { + "status": "failed", + "error": ( + "SHAP was enabled but no client " + "returned a valid SHAP result." + ), + "failed_clients": shap_failures, + } + print( + "[Server/SHAP] No valid local SHAP " + "results were returned.", + flush=True, + ) + return loss, metrics strat.aggregate_evaluate = logged_agg_eval + # 6b) On the last round, federated SHAP should run on every + # currently connected client, regardless of fraction_evaluate. + original_conf_eval = strat.configure_evaluate + + def wrapped_conf_eval( + server_round, + parameters, + client_manager, + ): + is_final_shap_round = ( + self.shap_config.enabled + and int(server_round) == int(self.total_rounds) + ) + + if not is_final_shap_round: + return original_conf_eval( + server_round=server_round, + parameters=parameters, + client_manager=client_manager, + ) + + evaluate_ins = EvaluateIns( + parameters=parameters, + config=evaluate_config_fn(server_round), + ) + + available_clients = list( + client_manager.all().values() + ) + + print( + "[Server/SHAP] Final round: requesting local " + f"SHAP from {len(available_clients)} connected " + "client(s).", + flush=True, + ) + + return [ + (client, evaluate_ins) + for client in available_clients + ] + + strat.configure_evaluate = wrapped_conf_eval + # 7) Wrap configure_fit to: # - log client properties (unchanged) # - apply split_mode/client_fractions to fit_ins.config (NEW) diff --git a/MEDfl/rw/xgboost/__init__.py b/MEDfl/rw/xgboost/__init__.py new file mode 100644 index 0000000..726124b --- /dev/null +++ b/MEDfl/rw/xgboost/__init__.py @@ -0,0 +1,16 @@ +""" +Federated XGBoost backend for MEDfl real-world mode. + +This module adds XGBoost support without modifying the existing PyTorch +real-world client, model, or strategy implementation. +""" + +from MEDfl.rw.xgboost.client import XGBoostFlowerClient +from MEDfl.rw.xgboost.strategy import XGBoostStrategy +from MEDfl.rw.xgboost.config import XGBoostConfig + +__all__ = [ + "XGBoostFlowerClient", + "XGBoostStrategy", + "XGBoostConfig", +] \ No newline at end of file diff --git a/MEDfl/rw/xgboost/client.py b/MEDfl/rw/xgboost/client.py new file mode 100644 index 0000000..0d7844e --- /dev/null +++ b/MEDfl/rw/xgboost/client.py @@ -0,0 +1,379 @@ +import argparse +import json +import socket +import platform +import shutil +from typing import Dict, Any, Optional + +import flwr as fl +import numpy as np +import pandas as pd +import psutil +import xgboost as xgb +from flwr.common import ( + Code, + EvaluateIns, + EvaluateRes, + FitIns, + FitRes, + GetParametersIns, + GetParametersRes, + GetPropertiesIns, + GetPropertiesRes, + Parameters, + Status, +) + +try: + import GPUtil +except ImportError: + GPUtil = None + +from MEDfl.rw.xgboost.data import XGBoostDataModule +from MEDfl.rw.xgboost.metrics import evaluate_predictions +from MEDfl.rw.xgboost.serialization import ( + booster_to_parameters, + parameters_to_booster, +) + + +class XGBoostFlowerClient(fl.client.Client): + """ + Real-world Flower client for federated XGBoost in MEDfl. + + This client is intentionally separate from the existing PyTorch NumPyClient. + The PyTorch client exchanges neural network weights. This client exchanges + serialized XGBoost boosters. + + Parameters + ---------- + server_address : str + Flower server address, e.g. "127.0.0.1:8080". + data_path : str + Local CSV dataset path. + val_frac : Optional[float] + Optional client-side validation fraction override. + test_frac : Optional[float] + Optional client-side test fraction override. + id_col : Optional[str] + Column used for ID-based test split. + test_ids : Optional[str] + Comma-separated test IDs. + seed : int + Random seed. + """ + + def __init__( + self, + server_address: str, + data_path: str = "data/data.csv", + val_frac: Optional[float] = None, + test_frac: Optional[float] = None, + id_col: Optional[str] = None, + test_ids: Optional[str] = None, + seed: int = 42, + ): + self.server_address = server_address + self.data_path = data_path + + self.data_module = XGBoostDataModule( + data_path=data_path, + val_frac=val_frac, + test_frac=test_frac, + id_col=id_col, + test_ids=test_ids, + seed=seed, + ) + + self.booster: Optional[xgb.Booster] = None + + self.dtrain = None + self.dval = None + self.dtest = None + self.metadata: Dict[str, Any] = {} + + self.task = "binary" + self.threshold = 0.5 + self.xgb_params: Dict[str, Any] = {} + self.local_num_boost_round = 10 + + self._initialized = False + + self._df_preview = pd.read_csv(data_path, nrows=5) + + def _initialize_from_config(self, config: Dict[str, Any]) -> None: + """ + Lazily initialize DMatrix objects from server config. + """ + + self.task = str(config.get("task", self.task) or self.task) + self.threshold = float(config.get("threshold", self.threshold)) + self.local_num_boost_round = int( + config.get("local_num_boost_round", self.local_num_boost_round) + ) + + xgb_params_raw = config.get("xgb_params", "{}") + + if isinstance(xgb_params_raw, str): + try: + self.xgb_params = json.loads(xgb_params_raw) + except json.JSONDecodeError as exc: + raise ValueError( + f"Invalid xgb_params JSON received from server: {xgb_params_raw}" + ) from exc + elif isinstance(xgb_params_raw, dict): + self.xgb_params = xgb_params_raw + else: + raise ValueError( + "xgb_params must be either a JSON string or a Python dictionary." + ) + + self.dtrain, self.dval, self.dtest, self.metadata = ( + self.data_module.prepare_from_config(config) + ) + + self._initialized = True + + print( + "[XGBoostClient] Initialized with " + f"task={self.task}, " + f"features={self.metadata.get('features')}, " + f"target={self.metadata.get('target')}, " + f"train={self.metadata.get('num_train')}, " + f"val={self.metadata.get('num_val')}, " + f"test={self.metadata.get('num_test')}" + ) + + def get_parameters(self, ins: GetParametersIns) -> GetParametersRes: + """ + Return current booster parameters. + + Before training, this can be empty. + """ + + return GetParametersRes( + status=Status(code=Code.OK, message="Success"), + parameters=booster_to_parameters(self.booster), + ) + + def fit(self, ins: FitIns) -> FitRes: + """ + Receive global booster, continue local XGBoost training, return updated booster. + """ + + config = dict(ins.config or {}) + + if not self._initialized: + self._initialize_from_config(config) + + received_booster = parameters_to_booster(ins.parameters) + + if received_booster is not None: + self.booster = received_booster + + evals = [] + if self.dval is not None and self.metadata.get("num_val", 0) > 0: + evals.append((self.dval, "validation")) + + self.booster = xgb.train( + params=self.xgb_params, + dtrain=self.dtrain, + num_boost_round=self.local_num_boost_round, + xgb_model=self.booster, + evals=evals, + verbose_eval=False, + ) + + train_pred = self.booster.predict(self.dtrain) + + train_metrics = evaluate_predictions( + task=self.task, + y_true=self.dtrain.get_label(), + y_pred_or_prob=train_pred, + threshold=self.threshold, + num_features=int(self.metadata.get("num_features", 0)), + ) + + metrics = { + "backend": "xgboost", + "task": self.task, + "hostname": socket.gethostname(), + "os_type": platform.system(), + "num_train": int(self.metadata.get("num_train", 0)), + "features": ",".join(self.metadata.get("features", [])), + "target": str(self.metadata.get("target", "")), + "val_fraction": float(self.metadata.get("val_fraction", 0.0)), + "test_fraction": float(self.metadata.get("test_fraction", 0.0)), + } + + for key, value in train_metrics.items(): + metrics[f"train_{key}"] = value + + return FitRes( + status=Status(code=Code.OK, message="Success"), + parameters=booster_to_parameters(self.booster), + num_examples=int(self.metadata.get("num_train", 0)), + metrics=metrics, + ) + + def evaluate(self, ins: EvaluateIns) -> EvaluateRes: + """ + Evaluate received global booster on local test set. + """ + + config = dict(ins.config or {}) + + if not self._initialized: + self._initialize_from_config(config) + + received_booster = parameters_to_booster(ins.parameters) + + if received_booster is not None: + self.booster = received_booster + + if self.booster is None: + return EvaluateRes( + status=Status(code=Code.EVALUATE_NOT_IMPLEMENTED, message="No booster."), + loss=0.0, + num_examples=0, + metrics={}, + ) + + y_true = self.dtest.get_label() + y_pred = self.booster.predict(self.dtest) + + eval_metrics = evaluate_predictions( + task=self.task, + y_true=y_true, + y_pred_or_prob=y_pred, + threshold=self.threshold, + num_features=int(self.metadata.get("num_features", 0)), + ) + + metrics = { + "backend": "xgboost", + "task": self.task, + } + + for key, value in eval_metrics.items(): + metrics[f"eval_{key}"] = value + + if self.task == "regression": + loss = float(eval_metrics.get("rmse", 0.0)) + elif self.task == "binary": + loss = float(eval_metrics.get("logloss", 0.0)) + else: + loss = float(eval_metrics.get("mlogloss", 0.0)) + + return EvaluateRes( + status=Status(code=Code.OK, message="Success"), + loss=loss, + num_examples=int(self.metadata.get("num_test", 0)), + metrics=metrics, + ) + + def get_properties(self, ins: GetPropertiesIns) -> GetPropertiesRes: + """ + Return machine and dataset properties. + """ + + hostname = socket.gethostname() + os_type = platform.system() + + df = self.data_module.df + + feature_names = df.columns[:-1].tolist() + target_name = df.columns[-1] + + try: + label_counts = df[target_name].value_counts().to_dict() + dist_str = ",".join(f"{k}:{v}" for k, v in label_counts.items()) + classes_str = ",".join(map(str, sorted(label_counts.keys()))) + except Exception: + dist_str = "" + classes_str = "" + + cpu_physical = psutil.cpu_count(logical=False) + cpu_logical = psutil.cpu_count(logical=True) + total_mem_gb = round(psutil.virtual_memory().total / (1024**3), 2) + + driver_present = shutil.which("nvidia-smi") is not None + gpu_count = 0 + + if GPUtil and driver_present: + try: + gpu_count = len(GPUtil.getGPUs()) + except Exception: + gpu_count = 0 + + properties = { + "backend": "xgboost", + "hostname": hostname, + "os_type": os_type, + "num_samples": int(len(df)), + "num_features": int(len(feature_names)), + "features": ",".join(feature_names), + "target": target_name, + "classes": classes_str, + "label_distribution": dist_str, + "cpu_physical_cores": int(cpu_physical or 0), + "cpu_logical_cores": int(cpu_logical or 0), + "total_memory_gb": float(total_mem_gb), + "gpu_driver_present": str(driver_present), + "gpu_count": int(gpu_count), + } + + return GetPropertiesRes( + status=Status(code=Code.OK, message="Success"), + properties=properties, + ) + + def start(self) -> None: + """ + Start the Flower client. + """ + + fl.client.start_client( + server_address=self.server_address, + client=self, + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description="MEDfl real-world XGBoost client") + + parser.add_argument( + "--server_address", + type=str, + required=True, + help="Flower server address, e.g. 127.0.0.1:8080", + ) + parser.add_argument( + "--data_path", + type=str, + default="data/data.csv", + help="Path to local CSV dataset.", + ) + parser.add_argument("--val_frac", type=float, default=None) + parser.add_argument("--test_frac", type=float, default=None) + parser.add_argument("--id_col", type=str, default=None) + parser.add_argument("--test_ids", type=str, default=None) + parser.add_argument("--seed", type=int, default=42) + + args = parser.parse_args() + + client = XGBoostFlowerClient( + server_address=args.server_address, + data_path=args.data_path, + val_frac=args.val_frac, + test_frac=args.test_frac, + id_col=args.id_col, + test_ids=args.test_ids, + seed=args.seed, + ) + + client.start() + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/MEDfl/rw/xgboost/config.py b/MEDfl/rw/xgboost/config.py new file mode 100644 index 0000000..2e26b70 --- /dev/null +++ b/MEDfl/rw/xgboost/config.py @@ -0,0 +1,84 @@ +from dataclasses import dataclass, field +from typing import Dict, Any, Optional + + +@dataclass +class XGBoostConfig: + """ + Configuration object for MEDfl federated XGBoost. + + Parameters + ---------- + task : str + One of: "binary", "regression", "multiclass". + mode : str + Federated XGBoost mode. First supported mode: "bagging". + local_num_boost_round : int + Number of local boosting rounds per federated round. + params : dict + Native XGBoost training parameters. + num_classes : Optional[int] + Required for multiclass classification. + """ + + task: str = "binary" + mode: str = "bagging" + local_num_boost_round: int = 10 + params: Dict[str, Any] = field(default_factory=dict) + num_classes: Optional[int] = None + + def build_params(self) -> Dict[str, Any]: + """ + Build a safe default XGBoost parameter dictionary. + + User-provided params override the defaults. + """ + + if self.task == "binary": + defaults = { + "objective": "binary:logistic", + "eval_metric": "auc", + "tree_method": "hist", + "max_depth": 6, + "eta": 0.1, + "subsample": 1.0, + "colsample_bytree": 1.0, + } + + elif self.task == "regression": + defaults = { + "objective": "reg:squarederror", + "eval_metric": "rmse", + "tree_method": "hist", + "max_depth": 6, + "eta": 0.1, + "subsample": 1.0, + "colsample_bytree": 1.0, + } + + elif self.task == "multiclass": + if self.num_classes is None or self.num_classes < 2: + raise ValueError( + "num_classes must be provided and >= 2 for multiclass XGBoost." + ) + + defaults = { + "objective": "multi:softprob", + "num_class": self.num_classes, + "eval_metric": "mlogloss", + "tree_method": "hist", + "max_depth": 6, + "eta": 0.1, + "subsample": 1.0, + "colsample_bytree": 1.0, + } + + else: + raise ValueError( + f"Unsupported XGBoost task '{self.task}'. " + "Choose from: 'binary', 'regression', 'multiclass'." + ) + + merged = defaults.copy() + merged.update(self.params or {}) + return merged \ No newline at end of file diff --git a/MEDfl/rw/xgboost/data.py b/MEDfl/rw/xgboost/data.py new file mode 100644 index 0000000..73adb07 --- /dev/null +++ b/MEDfl/rw/xgboost/data.py @@ -0,0 +1,253 @@ +from typing import Dict, Any, Optional, Tuple, List + +import numpy as np +import pandas as pd +import xgboost as xgb +from sklearn.model_selection import train_test_split + + +class XGBoostDataModule: + """ + Data preparation module for MEDfl federated XGBoost clients. + + This mirrors the behavior of the existing PyTorch real-world client: + - server-enforced features; + - server-enforced target; + - global or per-client split fractions; + - optional ID-based test split. + """ + + def __init__( + self, + data_path: str, + val_frac: Optional[float] = None, + test_frac: Optional[float] = None, + id_col: Optional[str] = None, + test_ids: Optional[str] = None, + seed: int = 42, + ): + self.data_path = data_path + self.df = pd.read_csv(data_path) + + self.client_val_frac = val_frac + self.client_test_frac = test_frac + self.id_col = id_col + self.test_ids = test_ids + self.seed = seed + + self.feature_names = self.df.columns[:-1].tolist() + self.target_name = self.df.columns[-1] + + self.effective_features = self.feature_names[:] + self.effective_target = self.target_name + self.effective_val_frac = 0.0 + self.effective_test_frac = 0.0 + + def prepare_from_config( + self, + config: Dict[str, Any], + ) -> Tuple[xgb.DMatrix, Optional[xgb.DMatrix], xgb.DMatrix, Dict[str, Any]]: + """ + Prepare XGBoost DMatrix objects from server config. + """ + + srv_features = str(config.get("features", "") or "").strip() + srv_target = str(config.get("target", "") or "").strip() + + if srv_target: + if srv_target not in self.df.columns: + raise ValueError( + f"Server-specified target '{srv_target}' not found in CSV columns: " + f"{list(self.df.columns)}" + ) + target_col = srv_target + else: + target_col = self.df.columns[-1] + + if srv_features: + feat_cols = [c.strip() for c in srv_features.split(",") if c.strip()] + missing = [c for c in feat_cols if c not in self.df.columns] + if missing: + raise ValueError( + f"Server-specified feature columns not found in CSV: {missing}" + ) + else: + feat_cols = [c for c in self.df.columns if c != target_col] + + srv_val = config.get("val_fraction", None) + srv_test = config.get("test_fraction", None) + + val_frac = ( + float(self.client_val_frac) + if self.client_val_frac is not None + else float(srv_val) + if srv_val is not None + else 0.10 + ) + + test_frac = ( + float(self.client_test_frac) + if self.client_test_frac is not None + else float(srv_test) + if srv_test is not None + else 0.10 + ) + + if not (0.0 <= val_frac < 1.0): + raise ValueError(f"Invalid val_fraction={val_frac}. Must be in [0, 1).") + + if config.get("id_col") and not self.id_col: + self.id_col = str(config.get("id_col")) + + if config.get("test_ids") and not self.test_ids: + self.test_ids = str(config.get("test_ids")) + + X_all = self.df[feat_cols].values + y_all = self.df[target_col].values + + if self.test_ids and str(self.test_ids).strip(): + X_train, X_val, X_test, y_train, y_val, y_test, actual_test_frac = ( + self._split_with_test_ids( + X_all=X_all, + y_all=y_all, + val_frac=val_frac, + ) + ) + self.effective_test_frac = actual_test_frac + + else: + if not (0.0 <= test_frac < 1.0 and val_frac + test_frac < 1.0): + raise ValueError( + f"Invalid fractions: val={val_frac}, test={test_frac}. " + "Require val >= 0, test >= 0, val + test < 1." + ) + + X_train, X_val, X_test, y_train, y_val, y_test = self._split_with_fractions( + X_all=X_all, + y_all=y_all, + val_frac=val_frac, + test_frac=test_frac, + ) + self.effective_test_frac = test_frac + + self.effective_features = feat_cols + self.effective_target = target_col + self.effective_val_frac = val_frac + + dtrain = xgb.DMatrix(X_train, label=y_train, feature_names=feat_cols) + + dval = None + if X_val is not None and len(y_val) > 0: + dval = xgb.DMatrix(X_val, label=y_val, feature_names=feat_cols) + + dtest = xgb.DMatrix(X_test, label=y_test, feature_names=feat_cols) + + metadata = { + "features": feat_cols, + "target": target_col, + "num_features": len(feat_cols), + "num_train": int(len(y_train)), + "num_val": int(len(y_val)) if y_val is not None else 0, + "num_test": int(len(y_test)), + "val_fraction": float(self.effective_val_frac), + "test_fraction": float(self.effective_test_frac), + } + + return dtrain, dval, dtest, metadata + + def _split_with_fractions( + self, + X_all, + y_all, + val_frac: float, + test_frac: float, + ): + strat_all = y_all if len(np.unique(y_all)) > 1 else None + + X_trval, X_test, y_trval, y_test = train_test_split( + X_all, + y_all, + test_size=test_frac, + random_state=self.seed, + stratify=strat_all, + ) + + if val_frac > 0 and len(y_trval) > 0: + actual_test_frac = len(y_test) / len(y_all) + rel_val = val_frac / (1.0 - actual_test_frac) + + strat_tr = y_trval if len(np.unique(y_trval)) > 1 else None + + X_train, X_val, y_train, y_val = train_test_split( + X_trval, + y_trval, + test_size=rel_val, + random_state=self.seed, + stratify=strat_tr, + ) + else: + X_train, y_train = X_trval, y_trval + X_val = np.empty((0, X_all.shape[1])) + y_val = np.empty((0,)) + + return X_train, X_val, X_test, y_train, y_val, y_test + + def _split_with_test_ids( + self, + X_all, + y_all, + val_frac: float, + ): + test_ids_list = [i.strip() for i in str(self.test_ids).split(",") if i.strip()] + + if self.id_col and self.id_col in self.df.columns: + id_series = self.df[self.id_col] + + if np.issubdtype(id_series.dtype, np.number): + test_ids_list = [int(i) for i in test_ids_list] + else: + test_ids_list = [str(i) for i in test_ids_list] + + else: + id_series = self.df.index + + try: + test_ids_list = [int(i) for i in test_ids_list] + except ValueError as exc: + raise ValueError( + "test_ids must be integers when id_col is not provided " + "or not found in the CSV." + ) from exc + + test_mask = id_series.isin(test_ids_list) + + X_test = X_all[test_mask] + y_test = y_all[test_mask] + X_trval = X_all[~test_mask] + y_trval = y_all[~test_mask] + + actual_test_frac = len(y_test) / len(y_all) if len(y_all) > 0 else 0.0 + + if val_frac + actual_test_frac >= 1.0: + raise ValueError( + f"Validation fraction {val_frac} + actual test fraction " + f"{actual_test_frac} >= 1." + ) + + if val_frac > 0 and len(y_trval) > 0: + rel_val = val_frac / (1.0 - actual_test_frac) + strat_tr = y_trval if len(np.unique(y_trval)) > 1 else None + + X_train, X_val, y_train, y_val = train_test_split( + X_trval, + y_trval, + test_size=rel_val, + random_state=self.seed, + stratify=strat_tr, + ) + else: + X_train, y_train = X_trval, y_trval + X_val = np.empty((0, X_all.shape[1])) + y_val = np.empty((0,)) + + return X_train, X_val, X_test, y_train, y_val, y_test, actual_test_frac \ No newline at end of file diff --git a/MEDfl/rw/xgboost/metrics.py b/MEDfl/rw/xgboost/metrics.py new file mode 100644 index 0000000..ea08866 --- /dev/null +++ b/MEDfl/rw/xgboost/metrics.py @@ -0,0 +1,176 @@ +import math +from typing import Dict, Any + +import numpy as np +from sklearn.metrics import ( + accuracy_score, + roc_auc_score, + log_loss, + mean_squared_error, + mean_absolute_error, + r2_score, + f1_score, +) + + +def _safe_float(value: Any) -> float: + try: + value = float(value) + if math.isnan(value) or math.isinf(value): + return float("nan") + return value + except Exception: + return float("nan") + + +def binary_metrics(y_true, y_prob, threshold: float = 0.5) -> Dict[str, float]: + """ + Metrics for binary classification. + """ + + y_true = np.asarray(y_true) + y_prob = np.asarray(y_prob) + + y_pred = (y_prob >= threshold).astype(int) + + metrics = { + "accuracy": _safe_float(accuracy_score(y_true, y_pred)), + } + + try: + metrics["auc"] = _safe_float(roc_auc_score(y_true, y_prob)) + except Exception: + metrics["auc"] = float("nan") + + try: + metrics["logloss"] = _safe_float(log_loss(y_true, y_prob)) + except Exception: + metrics["logloss"] = float("nan") + + return metrics + + +def regression_metrics(y_true, y_pred, num_features: int = 0) -> Dict[str, float]: + """ + Metrics for regression. + """ + + y_true = np.asarray(y_true) + y_pred = np.asarray(y_pred) + + rmse = mean_squared_error(y_true, y_pred, squared=False) + mae = mean_absolute_error(y_true, y_pred) + r2 = r2_score(y_true, y_pred) + + n = len(y_true) + p = int(num_features) + + if n > p + 1: + adjusted_r2 = 1.0 - (1.0 - r2) * ((n - 1.0) / (n - p - 1.0)) + else: + adjusted_r2 = float("nan") + + return { + "rmse": _safe_float(rmse), + "mae": _safe_float(mae), + "r2": _safe_float(r2), + "adjusted_r2": _safe_float(adjusted_r2), + } + + +def multiclass_metrics(y_true, y_prob) -> Dict[str, float]: + """ + Metrics for multiclass classification. + """ + + y_true = np.asarray(y_true) + y_prob = np.asarray(y_prob) + + y_pred = np.argmax(y_prob, axis=1) + + metrics = { + "accuracy": _safe_float(accuracy_score(y_true, y_pred)), + } + + try: + metrics["macro_f1"] = _safe_float(f1_score(y_true, y_pred, average="macro")) + except Exception: + metrics["macro_f1"] = float("nan") + + try: + metrics["mlogloss"] = _safe_float(log_loss(y_true, y_prob)) + except Exception: + metrics["mlogloss"] = float("nan") + + return metrics + + +def evaluate_predictions( + task: str, + y_true, + y_pred_or_prob, + threshold: float = 0.5, + num_features: int = 0, +) -> Dict[str, float]: + """ + Task-aware metric dispatcher. + """ + + if task == "binary": + return binary_metrics(y_true, y_pred_or_prob, threshold=threshold) + + if task == "regression": + return regression_metrics(y_true, y_pred_or_prob, num_features=num_features) + + if task == "multiclass": + return multiclass_metrics(y_true, y_pred_or_prob) + + raise ValueError( + f"Unsupported task '{task}'. Choose from 'binary', 'regression', 'multiclass'." + ) + + +def aggregate_weighted_metrics(results): + """ + Weighted aggregation for Flower metric aggregation callbacks. + + Expected Flower format: + List[Tuple[int, Dict[str, Scalar]]] + """ + + total = sum(num_examples for num_examples, _ in results) + + if total == 0: + return {} + + keys = set() + for _, metrics in results: + keys.update(metrics.keys()) + + aggregated = {} + + for key in keys: + values = [] + weights = [] + + for num_examples, metrics in results: + if key not in metrics: + continue + + value = metrics[key] + + try: + value = float(value) + except Exception: + continue + + if math.isnan(value) or math.isinf(value): + continue + + values.append(value) + weights.append(num_examples) + + if values: + aggregated[key] = float(np.average(values, weights=weights)) + + return aggregated \ No newline at end of file diff --git a/MEDfl/rw/xgboost/serialization.py b/MEDfl/rw/xgboost/serialization.py new file mode 100644 index 0000000..e316c05 --- /dev/null +++ b/MEDfl/rw/xgboost/serialization.py @@ -0,0 +1,56 @@ +from typing import Optional + +import xgboost as xgb +from flwr.common import Parameters + + +XGBOOST_TENSOR_TYPE = "xgboost.ubj" + + +def booster_to_parameters(booster: Optional[xgb.Booster]) -> Parameters: + """ + Serialize an XGBoost Booster into Flower Parameters. + + XGBoost does not exchange PyTorch-style NumPy tensors. Instead, we serialize + the Booster as raw UBJSON bytes and store them in Flower Parameters. + """ + + if booster is None: + return Parameters(tensors=[], tensor_type=XGBOOST_TENSOR_TYPE) + + raw = booster.save_raw(raw_format="ubj") + return Parameters(tensors=[bytes(raw)], tensor_type=XGBOOST_TENSOR_TYPE) + + +def parameters_to_booster(parameters: Optional[Parameters]) -> Optional[xgb.Booster]: + """ + Deserialize Flower Parameters into an XGBoost Booster. + """ + + if parameters is None: + return None + + if not parameters.tensors: + return None + + booster = xgb.Booster() + booster.load_model(bytearray(parameters.tensors[0])) + return booster + + +def save_booster(booster: xgb.Booster, path: str) -> None: + """ + Save an XGBoost Booster to disk. + """ + + booster.save_model(path) + + +def load_booster(path: str) -> xgb.Booster: + """ + Load an XGBoost Booster from disk. + """ + + booster = xgb.Booster() + booster.load_model(path) + return booster \ No newline at end of file diff --git a/MEDfl/rw/xgboost/strategy.py b/MEDfl/rw/xgboost/strategy.py new file mode 100644 index 0000000..029ebc3 --- /dev/null +++ b/MEDfl/rw/xgboost/strategy.py @@ -0,0 +1,668 @@ +import json +import os +from typing import Dict, Any, Optional +import tempfile +import xgboost as xgb + +import flwr as fl +from flwr.common import ( + Parameters, + FitRes, + parameters_to_ndarrays, +) +from flwr.common import GetPropertiesIns + +from MEDfl.rw.xgboost.config import XGBoostConfig +from MEDfl.rw.xgboost.metrics import aggregate_weighted_metrics +from MEDfl.rw.xgboost.serialization import ( + parameters_to_booster,booster_to_parameters, + save_booster, + XGBOOST_TENSOR_TYPE, +) + + +class XGBoostStrategy: + """ + Strategy wrapper for MEDfl real-world federated XGBoost. + + This class follows the same high-level contract as the existing MEDfl + Strategy wrapper: + - it exposes create_strategy(); + - it stores the Flower strategy in self.strategy_object; + - it can be passed directly into FederatedServer. + + First supported mode: + - bagging-style federated XGBoost using Flower strategy hooks. + + Notes + ----- + This implementation is conservative and non-invasive: + - It does not modify the existing PyTorch FedAvg/FedAdam/FedYogi path. + - It uses Flower Parameters containing serialized XGBoost boosters. + - It saves XGBoost models as .ubj files instead of .npz files. + """ + + def __init__( + self, + mode: str = "bagging", + fraction_fit: float = 1.0, + fraction_evaluate: float = 1.0, + min_fit_clients: int = 2, + min_evaluate_clients: int = 2, + min_available_clients: int = 2, + task: str = "binary", + num_classes: Optional[int] = None, + xgb_params: Optional[Dict[str, Any]] = None, + local_num_boost_round: int = 10, + threshold: float = 0.5, + features: str = "", + target: str = "", + val_fraction: float = 0.10, + test_fraction: float = 0.10, + split_mode: str = "global", + client_fractions: Optional[Dict[str, Dict[str, Any]]] = None, + id_col: str = "id", + savingPath: str = "", + saveOnRounds: int = 1, + total_rounds: int = 3, + ): + self.name = "FedXgbBagging" + self.mode = mode + + if self.mode != "bagging": + raise ValueError( + "Only mode='bagging' is implemented in this first integration. " + "Cyclic mode should be added after bagging is stable." + ) + + self.fraction_fit = fraction_fit + self.fraction_evaluate = fraction_evaluate + self.min_fit_clients = min_fit_clients + self.min_evaluate_clients = min_evaluate_clients + self.min_available_clients = min_available_clients + + self.task = task + self.num_classes = num_classes + self.threshold = threshold + + config = XGBoostConfig( + task=task, + mode=mode, + local_num_boost_round=local_num_boost_round, + params=xgb_params or {}, + num_classes=num_classes, + ) + + self.xgb_params = config.build_params() + self.local_num_boost_round = local_num_boost_round + + self._features = features + self._target = target + self._val_fraction = val_fraction + self._test_fraction = test_fraction + + self.split_mode = split_mode + self.client_fractions = client_fractions or {} + self.id_col = id_col + + self.savingPath = savingPath + self.saveOnRounds = saveOnRounds + self.total_rounds = total_rounds + + self.strategy_object = None + self.current_booster: Optional[xgb.Booster] = None + self.current_num_trees: int = 0 + + def create_strategy(self) -> None: + """ + Create the real Flower strategy object. + + We use FedAvg as the orchestration base but override aggregation behavior + to keep XGBoost booster exchange explicit. + """ + + def fit_config_fn(server_round: int) -> Dict[str, Any]: + return { + "backend": "xgboost", + "mode": self.mode, + "task": self.task, + "threshold": float(self.threshold), + "features": self._features, + "target": self._target, + "val_fraction": float(self._val_fraction), + "test_fraction": float(self._test_fraction), + "id_col": self.id_col, + "local_num_boost_round": int(self.local_num_boost_round), + "xgb_params": json.dumps(self.xgb_params), + } + + def evaluate_config_fn(server_round: int) -> Dict[str, Any]: + return { + "backend": "xgboost", + "mode": self.mode, + "task": self.task, + "threshold": float(self.threshold), + "features": self._features, + "target": self._target, + "val_fraction": float(self._val_fraction), + "test_fraction": float(self._test_fraction), + "id_col": self.id_col, + "local_num_boost_round": int(self.local_num_boost_round), + "xgb_params": json.dumps(self.xgb_params), + } + + strategy = fl.server.strategy.FedAvg( + fraction_fit=self.fraction_fit, + fraction_evaluate=self.fraction_evaluate, + min_fit_clients=self.min_fit_clients, + min_evaluate_clients=self.min_evaluate_clients, + min_available_clients=self.min_available_clients, + on_fit_config_fn=fit_config_fn, + on_evaluate_config_fn=evaluate_config_fn, + fit_metrics_aggregation_fn=aggregate_weighted_metrics, + evaluate_metrics_aggregation_fn=aggregate_weighted_metrics, + initial_parameters=Parameters( + tensors=[], + tensor_type=XGBOOST_TENSOR_TYPE, + ), + ) + + self._wrap_configure_fit(strategy) + self._wrap_aggregate_fit(strategy) + self._wrap_aggregate_evaluate(strategy) + + self.strategy_object = strategy + + def _wrap_configure_fit(self, strategy) -> None: + """ + Inject per-client test fractions / test IDs exactly like the existing MEDfl strategy. + """ + + original_configure_fit = strategy.configure_fit + + def wrapped_configure_fit(server_round, parameters, client_manager): + selected = original_configure_fit( + server_round=server_round, + parameters=parameters, + client_manager=client_manager, + ) + + ins = GetPropertiesIns(config={}) + + for client, fit_ins in selected: + hostname = None + + try: + props = client.get_properties( + ins=ins, + timeout=10.0, + group_id=0, + ) + print( + f"\n📋 [XGBoost Round {server_round}] " + f"Client {client.cid} Properties: {props.properties}" + ) + hostname = props.properties.get("hostname", None) + + except Exception as exc: + print(f"⚠️ Failed to get properties from {client.cid}: {exc}") + + if not hostname: + hostname = client.cid + + cfg = fit_ins.config + + if self.split_mode == "per_client": + per_cfg = ( + self.client_fractions.get(hostname) + or self.client_fractions.get(client.cid) + or {} + ) + + if "val_fraction" in per_cfg: + try: + cfg["val_fraction"] = float(per_cfg["val_fraction"]) + except Exception: + pass + + if "test_ids" in per_cfg and per_cfg["test_ids"]: + test_ids_val = per_cfg["test_ids"] + + if isinstance(test_ids_val, (list, tuple, set)): + test_ids_str = ",".join(str(x) for x in test_ids_val) + else: + test_ids_str = str(test_ids_val) + + cfg["test_ids"] = test_ids_str + + if "test_fraction" in cfg: + del cfg["test_fraction"] + + cfg["id_col"] = self.id_col + + else: + if "test_fraction" in per_cfg: + try: + cfg["test_fraction"] = float(per_cfg["test_fraction"]) + except Exception: + pass + + else: + if "test_ids" in cfg: + del cfg["test_ids"] + + cfg["val_fraction"] = float(self._val_fraction) + cfg["test_fraction"] = float(self._test_fraction) + cfg["id_col"] = self.id_col + + return selected + + strategy.configure_fit = wrapped_configure_fit + + def _wrap_aggregate_fit(self, strategy) -> None: + """ + Wrap aggregate_fit for XGBoost. + + IMPORTANT: + XGBoost boosters are NOT NumPy tensors. + Therefore we MUST NOT call FedAvg tensor aggregation. + + Current strategy: + - select the booster from the client with the largest dataset + - aggregate only metrics manually + - save the selected booster + """ + + def wrapped_aggregate_fit(server_round, results, failures): + print(f"\n[Server/XGBoost] 🔄 Round {server_round} - Client Training Metrics:") + + for client, fit_res in results: + print( + f" XGB CTM Round {server_round} " + f"Client:{client.cid}: {fit_res.metrics}" + ) + + if failures: + print(f"[Server/XGBoost] ⚠️ Round {server_round} failures: {failures}") + + if not results: + print("[Server/XGBoost] No fit results received.") + return None, {} + + # Aggregate metrics manually, same as before. + metric_results = [] + + for _, fit_res in results: + metric_results.append( + ( + fit_res.num_examples, + fit_res.metrics, + ) + ) + + metrics = aggregate_weighted_metrics(metric_results) + + previous_num_trees = int(self.current_num_trees) + new_local_boosters = [] + + print( + f"[Server/XGBoost] Previous global tree count before round " + f"{server_round}: {previous_num_trees}" + ) + + for client, fit_res in results: + client_booster = parameters_to_booster(fit_res.parameters) + + if client_booster is None: + print( + f"[Server/XGBoost] Client {client.cid} returned no booster. Skipping." + ) + continue + + client_total_trees = self._num_boosted_rounds(client_booster) + + print( + f"[Server/XGBoost] Client {client.cid} returned booster with " + f"{client_total_trees} trees." + ) + + # In round 1, previous_num_trees is 0, so this keeps all local trees. + # In later rounds, this keeps only the new trees added by the client. + new_trees_booster = self._slice_new_trees( + client_booster=client_booster, + previous_num_trees=previous_num_trees, + ) + + if new_trees_booster is None: + continue + + new_tree_count = self._num_boosted_rounds(new_trees_booster) + + print( + f"[Server/XGBoost] Client {client.cid} contributed " + f"{new_tree_count} new trees." + ) + + new_local_boosters.append(new_trees_booster) + + if not new_local_boosters: + print( + "[Server/XGBoost] No new local trees received. " + "Keeping previous global booster." + ) + + if self.current_booster is None: + return None, metrics + + agg_parameters = booster_to_parameters(self.current_booster) + return agg_parameters, metrics + + boosters_to_combine = [] + + if self.current_booster is not None: + boosters_to_combine.append(self.current_booster) + + boosters_to_combine.extend(new_local_boosters) + + combined_booster = self._combine_boosters_by_json(boosters_to_combine) + + if combined_booster is None: + print("[Server/XGBoost] Failed to combine boosters.") + return None, metrics + + self.current_booster = combined_booster + self.current_num_trees = self._num_boosted_rounds(combined_booster) + + print( + f"[Server/XGBoost] ✅ Round {server_round} - Flower-style bagging aggregation done." + ) + print( + f"[Server/XGBoost] Global booster now has {self.current_num_trees} trees." + ) + print(f"[Server/XGBoost] Aggregated metrics: {metrics}\n") + + # Optional: print importance directly, no saving required. + if server_round == self.total_rounds: + self._print_feature_importance(combined_booster, server_round) + + agg_parameters = booster_to_parameters(combined_booster) + + # Optional: keep your existing save behavior if savingPath is provided. + self._save_if_needed(server_round, agg_parameters) + + return agg_parameters, metrics + strategy.aggregate_fit = wrapped_aggregate_fit + + + def _wrap_aggregate_evaluate(self, strategy) -> None: + """ + Wrap aggregate_evaluate for logging. + """ + + original_aggregate_evaluate = strategy.aggregate_evaluate + + def wrapped_aggregate_evaluate(server_round, results, failures): + print(f"\n[Server/XGBoost] 📊 Round {server_round} - Client Evaluation Metrics:") + + for client, eval_res in results: + print( + f" XGB CEM Round {server_round} " + f"Client:{client.cid}: {eval_res.metrics}" + ) + + loss, metrics = original_aggregate_evaluate(server_round, results, failures) + + print(f"[Server/XGBoost] ✅ Round {server_round} - Aggregated Evaluation:") + print(f" Loss: {loss}, Metrics: {metrics}\n") + + return loss, metrics + + strategy.aggregate_evaluate = wrapped_aggregate_evaluate + + def _save_if_needed(self, server_round: int, parameters: Parameters) -> None: + """ + Save XGBoost model as .ubj. + """ + + if not self.savingPath: + return + + should_save = ( + server_round % self.saveOnRounds == 0 + or server_round == self.total_rounds + ) + + if not should_save: + return + + os.makedirs(self.savingPath, exist_ok=True) + + booster = parameters_to_booster(parameters) + + if booster is None: + print("[Server/XGBoost] No booster to save.") + return + + filename = ( + f"round_{server_round}_final_model.ubj" + if server_round == self.total_rounds + else f"round_{server_round}_model.ubj" + ) + + filepath = os.path.join(self.savingPath, filename) + + save_booster(booster, filepath) + + print(f"[Server/XGBoost] 💾 Saved booster to: {filepath}") + def _num_boosted_rounds(self, booster: xgb.Booster) -> int: + """ + Return the number of boosted rounds/trees in the booster. + """ + try: + return int(booster.num_boosted_rounds()) + except Exception: + try: + return len(booster.get_dump()) + except Exception: + return 0 + + + def _slice_new_trees( + self, + client_booster: xgb.Booster, + previous_num_trees: int, + ) -> Optional[xgb.Booster]: + """ + Extract only the trees added by the client in this round. + + Client booster = previous global booster + new local trees. + We keep only new local trees. + """ + + total_trees = self._num_boosted_rounds(client_booster) + + if total_trees <= previous_num_trees: + print( + "[Server/XGBoost] Client booster has no new trees. " + f"total_trees={total_trees}, previous_num_trees={previous_num_trees}" + ) + return None + + try: + return client_booster[previous_num_trees:total_trees] + except Exception as exc: + print(f"[Server/XGBoost] Failed to slice new trees: {exc}") + return None + + + def _booster_to_raw_bytes(self, booster: xgb.Booster) -> bytes: + """ + Serialize booster to raw UBJ bytes. + """ + with tempfile.NamedTemporaryFile(suffix=".ubj", delete=False) as tmp: + tmp_path = tmp.name + + try: + booster.save_model(tmp_path) + with open(tmp_path, "rb") as f: + return f.read() + finally: + try: + os.remove(tmp_path) + except Exception: + pass + + + def _booster_from_raw_bytes(self, raw: bytes) -> Optional[xgb.Booster]: + """ + Deserialize booster from raw UBJ bytes. + """ + if not raw: + return None + + with tempfile.NamedTemporaryFile(suffix=".ubj", delete=False) as tmp: + tmp_path = tmp.name + tmp.write(raw) + + try: + booster = xgb.Booster() + booster.load_model(tmp_path) + return booster + except Exception as exc: + print(f"[Server/XGBoost] Failed to load booster from bytes: {exc}") + return None + finally: + try: + os.remove(tmp_path) + except Exception: + pass + + def _combine_boosters_by_json( + self, + boosters: list[xgb.Booster], + ) -> Optional[xgb.Booster]: + """ + Combine multiple XGBoost boosters by concatenating their trees. + + Fixes: + - updates trees + - updates tree_info + - updates gbtree_model_param["num_trees"] + - updates iteration_indptr + + This is the missing part that caused: + Check failed: model.iteration_indptr.back() == model.param.num_trees + """ + + boosters = [b for b in boosters if b is not None] + + if not boosters: + return None + + if len(boosters) == 1: + return boosters[0] + + model_jsons = [] + + for booster in boosters: + raw_json = booster.save_raw(raw_format="json") + if isinstance(raw_json, bytes): + raw_json = raw_json.decode("utf-8") + model_jsons.append(json.loads(raw_json)) + + base = model_jsons[0] + + base_model = base["learner"]["gradient_booster"]["model"] + base_trees = base_model["trees"] + base_tree_info = base_model["tree_info"] + + # Start from the base model iteration structure + base_iteration_indptr = base_model.get("iteration_indptr", None) + + if base_iteration_indptr is None: + # Safe fallback for binary/regression: one tree per iteration + combined_iteration_indptr = list(range(0, len(base_trees) + 1)) + else: + combined_iteration_indptr = list(base_iteration_indptr) + + next_tree_id = len(base_trees) + + for extra in model_jsons[1:]: + extra_model = extra["learner"]["gradient_booster"]["model"] + + extra_trees = extra_model["trees"] + extra_tree_info = extra_model["tree_info"] + + # Append trees and reassign IDs + for tree in extra_trees: + tree["id"] = next_tree_id + base_trees.append(tree) + next_tree_id += 1 + + # Append tree_info + base_tree_info.extend(extra_tree_info) + + # Append iteration_indptr correctly + extra_iteration_indptr = extra_model.get("iteration_indptr", None) + + if extra_iteration_indptr is None: + # Fallback: one tree per boosting iteration + extra_iteration_indptr = list(range(0, len(extra_trees) + 1)) + + offset = combined_iteration_indptr[-1] + + # Skip the first 0 from the extra booster + for ptr in extra_iteration_indptr[1:]: + combined_iteration_indptr.append(offset + ptr) + + # Write back merged model fields + base_model["trees"] = base_trees + base_model["tree_info"] = base_tree_info + base_model["iteration_indptr"] = combined_iteration_indptr + + # Update num_trees + if "gbtree_model_param" in base_model: + base_model["gbtree_model_param"]["num_trees"] = str(len(base_trees)) + + # Save/load to validate the merged model + with tempfile.NamedTemporaryFile( + mode="w", + suffix=".json", + delete=False, + encoding="utf-8", + ) as tmp: + tmp_path = tmp.name + json.dump(base, tmp) + + try: + combined = xgb.Booster() + combined.load_model(tmp_path) + return combined + finally: + try: + os.remove(tmp_path) + except Exception: + pass + def _print_feature_importance( + self, + booster: xgb.Booster, + server_round: int, + ) -> None: + """ + Print feature importance for the current global booster. + """ + if booster is None: + return + + importance_gain = booster.get_score(importance_type="gain") + importance_weight = booster.get_score(importance_type="weight") + importance_cover = booster.get_score(importance_type="cover") + + print(f"\n[Server/XGBoost] 📊 Round {server_round} Global Feature Importance - Gain:") + print(importance_gain) + + print(f"[Server/XGBoost] 📊 Round {server_round} Global Feature Importance - Weight:") + print(importance_weight) + + print(f"[Server/XGBoost] 📊 Round {server_round} Global Feature Importance - Cover:") + print(importance_cover) \ No newline at end of file diff --git a/build/lib/MEDfl/LearningManager/client.py b/build/lib/MEDfl/LearningManager/client.py index 4a18ea8..4b6d631 100644 --- a/build/lib/MEDfl/LearningManager/client.py +++ b/build/lib/MEDfl/LearningManager/client.py @@ -1,4 +1,7 @@ #!/usr/bin/env python3 +import logging +import sys + import flwr as fl from opacus import PrivacyEngine from torch.utils.data import DataLoader @@ -6,52 +9,62 @@ from .model import Model from .utils import params import torch +import numpy as np + + +# --------------------------------------------------------------------------- +# Per-worker logger setup +# --------------------------------------------------------------------------- +# Ray workers are separate processes. We configure this logger inside the +# worker so that messages go to stdout — Ray then forwards stdout to the +# driver when log_to_driver=True is set in ray_init_args. +# --------------------------------------------------------------------------- +def _get_logger() -> logging.Logger: + logger = logging.getLogger("MEDfl.client") + if not logger.handlers: + handler = logging.StreamHandler(sys.stdout) + handler.setFormatter( + logging.Formatter( + "[%(asctime)s][%(name)s][%(levelname)s] %(message)s", + datefmt="%H:%M:%S", + ) + ) + logger.addHandler(handler) + logger.setLevel(logging.DEBUG) + logger.propagate = False # avoid double-printing on driver + return logger + class FlowerClient(fl.client.NumPyClient): """ FlowerClient class for creating MEDfl clients. - - Attributes: - cid (str): Client ID. - local_model (Model): Local model of the federated learning network. - trainloader (DataLoader): DataLoader for training data. - valloader (DataLoader): DataLoader for validation data. - diff_priv (bool): Flag indicating whether to use differential privacy. """ - def __init__(self, cid: str, local_model: Model, trainloader: DataLoader, valloader: DataLoader, diff_priv: bool = params["diff_privacy"]): - """ - Initializes the FlowerClient instance. - - Args: - cid (str): Client ID. - local_model (Model): Local model of the federated learning network. - trainloader (DataLoader): DataLoader for training data. - valloader (DataLoader): DataLoader for validation data. - diff_priv (bool): Flag indicating whether to use differential privacy. - """ + + def __init__( + self, + cid: str, + local_model: Model, + trainloader: DataLoader, + valloader: DataLoader, + diff_priv: bool = params["diff_privacy"], + ): self.cid = cid self.local_model = local_model self.trainloader = trainloader self.valloader = valloader - if torch.cuda.is_available(): - num_cuda_devices = torch.cuda.device_count() - if num_cuda_devices > 0: - device_idx = int(self.cid) % num_cuda_devices - self.device = torch.device(f"cuda:{device_idx}") - self.local_model.model.to(self.device) - else: - # Handle case where CUDA is available but no CUDA devices are found - raise RuntimeError("CUDA is available, but no CUDA devices are found.") - else: - # Handle case where CUDA is not available - self.device = torch.device("cpu") - self.local_model.model.to(self.device) + self.log = _get_logger() + + self.device = torch.device("cpu") + self.local_model.model.to(self.device) self.privacy_engine = PrivacyEngine(secure_mode=False) self.diff_priv = diff_priv self.epsilons = [] self.accuracies = [] self.losses = [] + self.train_round_losses = [] + self.train_round_metrics = [] + if self.diff_priv: model, optimizer, self.trainloader = self.privacy_engine.make_private_with_epsilon( module=self.local_model.model.train(), @@ -59,92 +72,98 @@ def __init__(self, cid: str, local_model: Model, trainloader: DataLoader, valloa data_loader=self.trainloader, epochs=params["train_epochs"], target_epsilon=float(params["EPSILON"]), - target_delta= float(params["DELTA"]), + target_delta=float(params["DELTA"]), max_grad_norm=params["MAX_GRAD_NORM"], ) setattr(self.local_model, "model", model) setattr(self.local_model, "optimizer", optimizer) self.validate() + # ------------------------------------------------------------------ def validate(self): """Validates cid, local_model, trainloader, valloader.""" if not isinstance(self.cid, str): raise TypeError("cid argument must be a string") - if not isinstance(self.local_model, Model): raise TypeError("local_model argument must be a MEDfl.LearningManager.model.Model") - if not isinstance(self.trainloader, DataLoader): raise TypeError("trainloader argument must be a torch.utils.data.dataloader") - if not isinstance(self.valloader, DataLoader): raise TypeError("valloader argument must be a torch.utils.data.dataloader") - if not isinstance(self.diff_priv, bool): raise TypeError("diff_priv argument must be a bool") + # ------------------------------------------------------------------ def get_parameters(self, config): - """ - Returns the current parameters of the local model. - - Args: - config: Configuration information. - - Returns: - Numpy array: Parameters of the local model. - """ - print(f"[Client {self.cid}] get_parameters") + self.log.info(f"[Client {self.cid}] get_parameters") return self.local_model.get_parameters() + # ------------------------------------------------------------------ def fit(self, parameters, config): - """ - Fits the local model to the received parameters using federated learning. - - Args: - parameters: Parameters received from the server. - config: Configuration information. - - Returns: - Tuple: Parameters of the local model, number of training examples, and privacy information. - """ - print('\n -------------------------------- \n this is the config of the client') - print(f"[Client {self.cid}] fit, config: {config}") - # print(config['epochs']) - print('\n -------------------------------- \n ') + self.log.info(f"[Client {self.cid}] fit | config: {config}") + self.local_model.set_parameters(parameters) - for _ in range(params["train_epochs"]): - epsilon = self.local_model.train( + + epoch_losses, epoch_metrics = [], [] + + for ep in range(params["train_epochs"]): + epsilon, tr_loss, tr_metric = self.local_model.train( self.trainloader, - epoch=_, + epoch=ep, device=self.device, privacy_engine=self.privacy_engine, diff_priv=self.diff_priv, ) self.epsilons.append(epsilon) - print(f"epsilon of client {self.cid} : eps = {epsilon}") + epoch_losses.append(tr_loss) + epoch_metrics.append(tr_metric) + + avg_grad_norm = float(np.nanmean(self.local_model.grad_norms)) \ + if hasattr(self.local_model, "grad_norms") and self.local_model.grad_norms \ + else float("nan") + self.local_model.grad_norms = [] # reset for next round + + + self.log.info( + f"[Client {self.cid}] fit done | " + f"epsilon={epsilon:.4f} | " + f"avg_loss={float(np.nanmean(epoch_losses)):.4f}" + ) + + round_train_loss = float(np.nanmean(epoch_losses)) if len(epoch_losses) else float("nan") + round_train_metric = float(np.nanmean(epoch_metrics)) if len(epoch_metrics) else float("nan") + + self.train_round_losses.append(round_train_loss) + self.train_round_metrics.append(round_train_metric) + + metric_key = "train_accuracy" if self.local_model.task_type == "binary" else "train_rmse" + return ( self.local_model.get_parameters(), - len(self.trainloader), - {"epsilon": epsilon}, + len(self.trainloader.dataset), + { + "epsilon": float(epsilon), + "train_loss": float(round_train_loss), + metric_key: float(round_train_metric), + "grad_norm": avg_grad_norm, # <-- NEW + }, ) + # ------------------------------------------------------------------ def evaluate(self, parameters, config): - """ - Evaluates the local model on the validation data and returns the loss and accuracy. + self.log.info(f"[Client {self.cid}] evaluate | config: {config}") + self.local_model.set_parameters(parameters) - Args: - parameters: Parameters received from the server. - config: Configuration information. + loss, metrics = self.local_model.evaluate(self.valloader, device=self.device) - Returns: - Tuple: Loss, number of validation examples, and accuracy information. - """ - print(f"[Client {self.cid}] evaluate, config: {config}") - self.local_model.set_parameters(parameters) - loss, accuracy , auc = self.local_model.evaluate( - self.valloader, device=self.device - ) self.losses.append(loss) - self.accuracies.append(accuracy) - - return float(loss), len(self.valloader), {"accuracy": float(accuracy)} + + if isinstance(metrics, dict) and "accuracy" in metrics: + self.accuracies.append(metrics["accuracy"]) + + self.log.info( + f"[Client {self.cid}] evaluate done | " + f"loss={loss:.4f} | metrics={metrics}" + ) + + return float(loss), len(self.valloader.dataset), {k: float(v) for k, v in (metrics or {}).items()} \ No newline at end of file diff --git a/build/lib/MEDfl/LearningManager/federated_dataset.py b/build/lib/MEDfl/LearningManager/federated_dataset.py index 4a451cc..c0a8acd 100644 --- a/build/lib/MEDfl/LearningManager/federated_dataset.py +++ b/build/lib/MEDfl/LearningManager/federated_dataset.py @@ -1,6 +1,8 @@ from MEDfl.NetManager.net_helper import * from MEDfl.NetManager.net_manager_queries import * from MEDfl.NetManager.database_connector import DatabaseManager +from typing import List, Optional + class FederatedDataset: def __init__( @@ -11,6 +13,7 @@ def __init__( trainloaders: list, valloaders: list, testloaders: list, + feature_names: Optional[List[str]] = None, ): """ Represents a Federated Dataset. @@ -30,6 +33,25 @@ def __init__( self.testloaders = testloaders self.size = len(self.trainloaders[0].dataset[0][0]) + if feature_names is None: + feature_names = [ + f"feature_{index}" + for index in range(self.size) + ] + + if len(feature_names) != self.size: + raise ValueError( + "feature_names length must match the " + f"federated dataset size. Expected {self.size}, " + f"received {len(feature_names)}." + ) + + self.feature_names = list( + feature_names + ) + + + db_manager = DatabaseManager() db_manager.connect() self.eng = db_manager.get_connection() diff --git a/build/lib/MEDfl/LearningManager/model.py b/build/lib/MEDfl/LearningManager/model.py index f68dc73..206873c 100644 --- a/build/lib/MEDfl/LearningManager/model.py +++ b/build/lib/MEDfl/LearningManager/model.py @@ -1,6 +1,7 @@ #!/usr/bin/env python3 -# froked from https://github.com/pythonlessons/mltu/blob/main/mltu/torch/model.py +import logging +import sys import typing from collections import OrderedDict from typing import Dict, List, Optional, Tuple @@ -8,19 +9,39 @@ import numpy as np import torch import torch.nn as nn -from sklearn.metrics import accuracy_score,roc_auc_score +from sklearn.metrics import ( + accuracy_score, roc_auc_score, + mean_squared_error, mean_absolute_error, r2_score, +) from .utils import params +# --------------------------------------------------------------------------- +# Per-worker logger — same pattern as client.py so Ray forwards to notebook +# --------------------------------------------------------------------------- +def _get_logger() -> logging.Logger: + logger = logging.getLogger("MEDfl.model") + if not logger.handlers: + handler = logging.StreamHandler(sys.stdout) + handler.setFormatter( + logging.Formatter( + "[%(asctime)s][%(name)s][%(levelname)s] %(message)s", + datefmt="%H:%M:%S", + ) + ) + logger.addHandler(handler) + logger.setLevel(logging.DEBUG) + logger.propagate = False + return logger + + +_log = _get_logger() + + class Model: """ Model class for training and testing PyTorch neural networks. - - Attributes: - model (torch.nn.Module): PyTorch neural network. - optimizer (torch.optim.Optimizer): PyTorch optimizer. - criterion (typing.Callable): Loss function. """ def __init__( @@ -28,197 +49,232 @@ def __init__( model: torch.nn.Module, optimizer: torch.optim.Optimizer, criterion: typing.Callable, + task_type: str = "binary", ) -> None: - """ - Initialize Model class with the specified model, optimizer, and criterion. - - Args: - model (torch.nn.Module): PyTorch neural network. - optimizer (torch.optim.Optimizer): PyTorch optimizer. - criterion (typing.Callable): Loss function. - """ self.model = model self.optimizer = optimizer self.criterion = criterion - # Get device on which model is running + self.task_type = task_type self.validate() def validate(self) -> None: - """ - Validate model and optimizer. - """ if not isinstance(self.model, torch.nn.Module): raise TypeError("model argument must be a torch.nn.Module") - if not isinstance(self.optimizer, torch.optim.Optimizer): - raise TypeError( - "optimizer argument must be a torch.optim.Optimizer" - ) + raise TypeError("optimizer argument must be a torch.optim.Optimizer") def get_parameters(self) -> List[np.ndarray]: - """ - Get the parameters of the model as a list of NumPy arrays. - - Returns: - List[np.ndarray]: The parameters of the model as a list of NumPy arrays. - """ - return [ - val.cpu().numpy() for _, val in self.model.state_dict().items() - ] + return [val.detach().cpu().numpy() for val in self.model.state_dict().values()] def set_parameters(self, parameters: List[np.ndarray]) -> None: - """ - Set the parameters of the model from a list of NumPy arrays. - - Args: - parameters (List[np.ndarray]): The parameters to be set. - """ - params_dict = zip(self.model.state_dict().keys(), parameters) - state_dict = OrderedDict({k: torch.Tensor(v) for k, v in params_dict}) - self.model.load_state_dict(state_dict, strict=True) - - def train( - self, train_loader, epoch, device, privacy_engine, diff_priv=False - ) -> float: - """ - Train the model on the given train_loader for one epoch. - - Args: - train_loader: The data loader for training data. - epoch (int): The current epoch number. - device: The device on which to perform the training. - privacy_engine: The privacy engine used for differential privacy (if enabled). - diff_priv (bool, optional): Whether differential privacy is used. Default is False. - - Returns: - float: The value of epsilon used in differential privacy. - """ + state_dict = self.model.state_dict() + keys = list(state_dict.keys()) + device = next(self.model.parameters()).device + new_state_dict = OrderedDict() + for k, v in zip(keys, parameters): + ref = state_dict[k] + new_state_dict[k] = torch.as_tensor(v, device=device, dtype=ref.dtype) + self.model.load_state_dict(new_state_dict, strict=True) + + def train(self, train_loader, epoch, device, privacy_engine, diff_priv=False) -> float: self.model.train() - epsilon = 0 + epsilon = 0.0 losses = [] - top1_acc = [] + metrics = [] for i, (X_train, y_train) in enumerate(train_loader): - X_train, y_train = X_train.to(device), y_train.to(device) - - self.optimizer.zero_grad() + X_train = X_train.to(device) + y_train = y_train.to(device).view(-1) + + if torch.isnan(X_train).any() or torch.isinf(X_train).any(): + _log.warning(f"Train batch {i}: X_train has NaN/Inf -> skipping batch") + continue + if torch.isnan(y_train).any() or torch.isinf(y_train).any(): + _log.warning(f"Train batch {i}: y_train has NaN/Inf -> skipping batch") + continue + + if i == 0 and epoch == 0: + try: + _log.debug( + f"Train scale check | " + f"X min/max: {X_train.min().item():.4f} / {X_train.max().item():.4f} | " + f"y min/max: {y_train.min().item():.4f} / {y_train.max().item():.4f}" + ) + except Exception: + pass - # compute output - y_hat = torch.squeeze(self.model(X_train), 1) - loss = self.criterion(y_hat, y_train) + self.optimizer.zero_grad(set_to_none=True) - preds = np.argmax(y_hat.detach().cpu().numpy(), axis=0) - labels = y_train.detach().cpu().numpy() + y_hat = self.model(X_train).view(-1) - # measure accuracy and record loss - acc = (preds == labels).mean() + if torch.isnan(y_hat).any() or torch.isinf(y_hat).any(): + _log.warning(f"Train batch {i}: y_hat has NaN/Inf -> skipping batch") + continue + + loss = self.criterion(y_hat, y_train) - losses.append(loss.item()) - top1_acc.append(acc) + if torch.isnan(loss) or torch.isinf(loss): + _log.warning(f"Train batch {i}: loss is NaN/Inf -> skipping batch") + continue + losses.append(float(loss.item())) loss.backward() + + bad_grad = False + for name, p in self.model.named_parameters(): + if p.grad is None: + continue + if torch.isnan(p.grad).any() or torch.isinf(p.grad).any(): + _log.warning( + f"Train batch {i}: grad NaN/Inf in {name} " + f"(NaN={torch.isnan(p.grad).sum().item()}, " + f"Inf={torch.isinf(p.grad).sum().item()}) -> skipping update" + ) + bad_grad = True + break + + if bad_grad: + self.optimizer.zero_grad(set_to_none=True) + continue + + torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0) self.optimizer.step() + if self.task_type == "binary": + y_prob = y_hat.detach().cpu().numpy() + y_true = y_train.detach().cpu().numpy() + acc = accuracy_score(y_true, (y_prob >= 0.5).astype(int)) + metrics.append(float(acc)) + elif self.task_type == "regression": + rmse = torch.sqrt(torch.mean((y_hat - y_train) ** 2)).item() + metrics.append(float(rmse)) + if diff_priv: - epsilon = privacy_engine.get_epsilon(float(params["DELTA"])) - - if (i + 1) % 10 == 0: - if diff_priv: - epsilon = privacy_engine.get_epsilon(float(params["DELTA"])) - print( - f"\tTrain Epoch: {epoch} \t" - f"Loss: {np.mean(losses):.6f} " - f"Acc@1: {np.mean(top1_acc) * 100:.6f} " - f"(ε = {epsilon:.2f}, δ = {params['DELTA']})" + epsilon = float(privacy_engine.get_epsilon(float(params["DELTA"]))) + + if (i + 1) % 10 == 0 and len(losses) > 0: + if self.task_type == "binary" and len(metrics) > 0: + _log.info( + f"Train Epoch {epoch} | batch {i+1} | " + f"Loss {np.mean(losses):.6f} | " + f"Acc {np.mean(metrics) * 100:.2f}%" + + (f" | ε={epsilon:.2f}" if diff_priv else "") ) - else: - print( - f"\tTrain Epoch: {epoch} \t" - f"Loss: {np.mean(losses):.6f} " - f"Acc@1: {np.mean(top1_acc) * 100:.6f}" + elif self.task_type == "regression" and len(metrics) > 0: + _log.info( + f"Train Epoch {epoch} | batch {i+1} | " + f"Loss {np.mean(losses):.6f} | " + f"RMSE {np.mean(metrics):.6f}" + + (f" | ε={epsilon:.2f}" if diff_priv else "") ) - return epsilon + mean_train_loss = float(np.mean(losses)) if len(losses) else float("nan") + mean_train_metric = float(np.mean(metrics)) if len(metrics) else float("nan") - def evaluate(self, val_loader, device=torch.device("cpu")) -> Tuple[float, float]: - """ - Evaluate the model on the given validation data. + total_grad_norm = 0.0 + for p in self.model.parameters(): + if p.grad is not None: + total_grad_norm += p.grad.detach().norm(2).item() ** 2 + total_grad_norm = total_grad_norm ** 0.5 + # store it + if not hasattr(self, "grad_norms"): + self.grad_norms = [] + self.grad_norms.append(total_grad_norm) - Args: - val_loader: The data loader for validation data. - device: The device on which to perform the evaluation. Default is 'cpu'. + return epsilon, mean_train_loss, mean_train_metric - Returns: - Tuple[float, float]: The evaluation loss and accuracy. - """ - correct, total, loss, accuracy, auc = 0, 0, 0.0, [], [] + def evaluate(self, val_loader, device=torch.device("cpu")): self.model.eval() + losses = [] + y_true_all, y_pred_all = [], [] with torch.no_grad(): - for X_test, y_test in val_loader: - X_test, y_test = X_test.to(device), y_test.to(device) # Move data to device - - y_hat = torch.squeeze(self.model(X_test), 1) - - - criterion = self.criterion.to(y_hat.device) - loss += criterion(y_hat, y_test).item() - - - # Move y_hat to CPU for accuracy computation - y_hat_cpu = y_hat.cpu().detach().numpy() - accuracy.append(accuracy_score(y_test.cpu().numpy(), y_hat_cpu.round())) - - # Move y_test to CPU for AUC computation - y_test_cpu = y_test.cpu().numpy() - y_prob_cpu = y_hat.cpu().detach().numpy() - if (len(np.unique(y_test_cpu)) != 1): - auc.append(roc_auc_score(y_test_cpu, y_prob_cpu)) - - total += y_test.size(0) - correct += np.sum(y_hat_cpu.round() == y_test_cpu) - - loss /= len(val_loader.dataset) - return loss, np.mean(accuracy), np.mean(auc) - + for i, (X_test, y_test) in enumerate(val_loader): + X_test = X_test.to(device) + y_test_t = y_test.to(device).view(-1) + y_hat_t = self.model(X_test).view(-1) + + if torch.isnan(y_hat_t).any() or torch.isinf(y_hat_t).any(): + _log.warning(f"Eval batch {i}: y_hat NaN/Inf -> skipping") + continue + if torch.isnan(y_test_t).any() or torch.isinf(y_test_t).any(): + _log.warning(f"Eval batch {i}: y_test NaN/Inf -> skipping") + continue + + crit = self.criterion.to(y_hat_t.device) if hasattr(self.criterion, "to") else self.criterion + loss_t = crit(y_hat_t, y_test_t) + + if torch.isnan(loss_t) or torch.isinf(loss_t): + _log.warning(f"Eval batch {i}: loss NaN/Inf -> skipping") + continue + + losses.append(float(loss_t.item())) + y_true_all.append(np.atleast_1d(y_test_t.detach().cpu().numpy())) + y_pred_all.append(np.atleast_1d(y_hat_t.detach().cpu().numpy())) + + if len(losses) == 0 or len(y_true_all) == 0: + return float("nan"), {} + + y_true = np.concatenate(y_true_all, axis=0) + y_pred = np.concatenate(y_pred_all, axis=0) + mean_loss = float(np.mean(losses)) + + if self.task_type == "binary": + acc = accuracy_score(y_true, (y_pred >= 0.5).astype(int)) + auc = roc_auc_score(y_true, y_pred) if len(np.unique(y_true)) > 1 else float("nan") + return mean_loss, {"accuracy": float(acc), "auc": float(auc)} + + if self.task_type == "regression": + mask = np.isfinite(y_true) & np.isfinite(y_pred) + if mask.sum() == 0: + return mean_loss, {"rmse": float("nan"), "mae": float("nan"), + "r2": float("nan"), "adj_r2": float("nan")} + + y_true_c = y_true[mask] + y_pred_c = y_pred[mask] + mse = mean_squared_error(y_true_c, y_pred_c) + mae = mean_absolute_error(y_true_c, y_pred_c) + r2 = r2_score(y_true_c, y_pred_c) + rmse = float(np.sqrt(mse)) + n = int(len(y_true_c)) + + p = None + try: + for m in self.model.modules(): + if isinstance(m, torch.nn.Linear): + p = int(m.in_features) + break + except Exception: + p = None + + if p is None or n <= p + 1: + adj_r2 = float("nan") + else: + adj_r2 = 1.0 - (1.0 - float(r2)) * (n - 1) / (n - p - 1) + + _log.debug(f"adj_r2 check | n={n} | p={p} | condition n>p+1={n > p + 1}") + + return mean_loss, { + "rmse": float(rmse), + "mae": float(mae), + "r2": float(r2), + "adj_r2": float(adj_r2), + } + + return mean_loss, {} @staticmethod - def save_model(model , model_name:str): - """ - Saves a PyTorch model to a file. - - Args: - model (torch.nn.Module): PyTorch model to be saved. - model_name (str): Name of the model file. - - Raises: - Exception: If there is an issue during the saving process. - - Returns: - None - """ + def save_model(model, model_name: str): try: torch.save(model, '../../notebooks/.ipynb_checkpoints/trainedModels/' + model_name + ".pth") except Exception as e: raise Exception(f"Error saving the model: {str(e)}") - + @staticmethod def load_model(model_path: str): - """ - Loads a PyTorch model from a file. - - Args: - model_path (str): Path to the model file to be loaded. - - Returns: - torch.nn.Module: Loaded PyTorch model. - """ - # Ensure models are loaded onto the CPU when CUDA is not available torch_kwargs = {"weights_only": False} if torch.cuda.is_available(): - loaded_model = torch.load(model_path , **torch_kwargs) + loaded_model = torch.load(model_path, **torch_kwargs) else: - loaded_model = torch.load(model_path, map_location=torch.device('cpu') , **torch_kwargs) - return loaded_model - - + loaded_model = torch.load(model_path, map_location=torch.device('cpu'), **torch_kwargs) + return loaded_model \ No newline at end of file diff --git a/build/lib/MEDfl/LearningManager/params.yaml b/build/lib/MEDfl/LearningManager/params.yaml index 83fc6fa..95f0440 100644 --- a/build/lib/MEDfl/LearningManager/params.yaml +++ b/build/lib/MEDfl/LearningManager/params.yaml @@ -1,8 +1,8 @@ DELTA: 1.0e-05 EPSILON: 5.0 MAX_GRAD_NORM: 1.0 -diff_privacy: true -lr: 0.001 +diff_privacy: false +lr: 0.01 min_evalclient: 2 num_rounds: 20 optimizer: SGD @@ -10,5 +10,5 @@ path_to_master_csv: /home/local/USHERBROOKE/saho6810/MEDfl/code/MEDfl/notebooks/ path_to_test_csv: /home/local/USHERBROOKE/saho6810/MEDfl/code/MEDfl/notebooks/data/masterDataSet/Mimic_train.csv task: BinaryClassification test_batch_size: 1 -train_batch_size: 32 -train_epochs: 20 +train_batch_size: 6 +train_epochs: 1 diff --git a/build/lib/MEDfl/LearningManager/plot.py b/build/lib/MEDfl/LearningManager/plot.py index c5f6c6e..23aa596 100644 --- a/build/lib/MEDfl/LearningManager/plot.py +++ b/build/lib/MEDfl/LearningManager/plot.py @@ -227,3 +227,108 @@ def plot_classification_report(pipeline_name: str): plt.show() return + + + + +def plot_client_regression_weights(strategy_obj, feature_names=None, include_bias=True): + """ + Plot per-client linear-regression weights over FL rounds. + + Assumes each regression entry in: + strategy_obj.client_fit_history[cid] + contains: + - "round" + - "weights" + + For a linear regressor: + weights[0] -> weight matrix, shape (1, input_dim) or similar + weights[1] -> bias, shape (1,) or similar + """ + + history = strategy_obj.client_fit_history + + if not history: + print("No client_fit_history found.") + return + + for cid, entries in history.items(): + # keep only rounds where weights were stored + reg_entries = [e for e in entries if "weights" in e] + + if not reg_entries: + print(f"Client {cid}: no stored regression weights.") + continue + + rounds = [] + coeff_series = [] + bias_series = [] + + for e in reg_entries: + rounds.append(int(e["round"])) + + weights = e["weights"] + + # linear regression assumption + w = np.array(weights[0]).reshape(-1) + + coeff_series.append(w) + + if include_bias and len(weights) > 1: + b = float(np.array(weights[1]).reshape(-1)[0]) + bias_series.append(b) + + coeff_series = np.array(coeff_series) # shape: (num_rounds, num_features) + + plt.figure(figsize=(10, 5)) + + n_features = coeff_series.shape[1] + for j in range(n_features): + label = feature_names[j] if feature_names is not None and j < len(feature_names) else f"w{j}" + plt.plot(rounds, coeff_series[:, j], marker="o", label=label) + + if include_bias and len(bias_series) == len(rounds): + plt.plot(rounds, bias_series, marker="x", linestyle="--", label="bias") + + plt.title(f"Client {cid} - weights over rounds") + plt.xlabel("Round") + plt.ylabel("Weight value") + plt.grid(True) + plt.legend() + plt.tight_layout() + plt.show() + + + + +def plot_one_weight_across_clients(strategy_obj, weight_index=0, feature_name=None): + history = strategy_obj.client_fit_history + + plt.figure(figsize=(10, 5)) + + for cid, entries in history.items(): + reg_entries = [e for e in entries if "weights" in e] + if not reg_entries: + continue + + rounds = [] + vals = [] + + for e in reg_entries: + w = np.array(e["weights"][0]).reshape(-1) + if weight_index >= len(w): + continue + rounds.append(int(e["round"])) + vals.append(float(w[weight_index])) + + if rounds: + plt.plot(rounds, vals, marker=",", label=f"Client {cid}") + + label = feature_name if feature_name is not None else f"w{weight_index}" + plt.title(f"{label} across clients over rounds") + plt.xlabel("Round") + plt.ylabel("Weight value") + plt.grid(True) + plt.legend() + plt.tight_layout() + plt.show() \ No newline at end of file diff --git a/build/lib/MEDfl/LearningManager/server.py b/build/lib/MEDfl/LearningManager/server.py index fa88c10..8b54913 100644 --- a/build/lib/MEDfl/LearningManager/server.py +++ b/build/lib/MEDfl/LearningManager/server.py @@ -1,32 +1,65 @@ #!/usr/bin/env python3 import copy +import logging +import sys from typing import Dict, Optional, Tuple import flwr as fl import torch +from torch.utils.data import ConcatDataset, DataLoader + from .client import FlowerClient from .federated_dataset import FederatedDataset from .model import Model from .strategy import Strategy +from .shap import ( + SHAPConfig, + LocalPyTorchSHAPExplainer, + aggregate_federated_shap, +) + + +# --------------------------------------------------------------------------- +# Logging setup — call this ONCE before starting the simulation. +# It forces Flower's logger (and the root logger) to write to stdout so that +# Jupyter captures every message from the driver process. +# --------------------------------------------------------------------------- +def setup_fl_logging(level: int = logging.INFO) -> None: + """ + Configure logging so that all FL-related messages appear in the notebook. + + Call this at the top of your notebook cell before creating FlowerServer: + + from MEDfl.LearningManager.server import setup_fl_logging + setup_fl_logging() + """ + fmt = logging.Formatter( + "[%(asctime)s][%(name)s][%(levelname)s] %(message)s", + datefmt="%H:%M:%S", + ) + handler = logging.StreamHandler(sys.stdout) + handler.setFormatter(fmt) + + # Root logger — catches everything + root = logging.getLogger() + root.setLevel(level) + if not any(isinstance(h, logging.StreamHandler) and h.stream is sys.stdout + for h in root.handlers): + root.addHandler(handler) + + # Flower logger specifically + for name in ("flwr", "flwr.server", "flwr.simulation", "flwr.client"): + lg = logging.getLogger(name) + lg.setLevel(level) + lg.propagate = True # bubble up to root so our handler fires + class FlowerServer: """ A class representing the central server for Federated Learning using Flower. - - Attributes: - global_model (Model): The global model that will be federated among clients. - strategy (Strategy): The strategy used for federated learning, specifying communication and aggregation methods. - num_rounds (int): The number of federated learning rounds to perform. - num_clients (int): The number of clients participating in the federated learning process. - fed_dataset (FederatedDataset): The federated dataset used for training and evaluation. - diff_priv (bool): Whether differential privacy is used during the federated learning process. - accuracies (List[float]): A list to store the accuracy of the global model during each round. - losses (List[float]): A list to store the loss of the global model during each round. - flower_clients (List[FlowerClient]): A list to store the FlowerClient objects representing individual clients. - """ def __init__( @@ -37,23 +70,13 @@ def __init__( num_clients: int, fed_dataset: FederatedDataset, diff_privacy: bool = False, - client_resources: Optional[Dict[str, float]] = {'num_cpus': 1, 'num_gpus': 0.0} + client_resources: Optional[Dict[str, float]] = {'num_cpus': 1, 'num_gpus': 0.0} , + shap_config: Optional[ + SHAPConfig + ] = None, ) -> None: - """ - Initialize a FlowerServer object with the specified parameters. - - Args: - global_model (Model): The global model that will be federated among clients. - strategy (Strategy): The strategy used for federated learning, specifying communication and aggregation methods. - num_rounds (int): The number of federated learning rounds to perform. - num_clients (int): The number of clients participating in the federated learning process. - fed_dataset (FederatedDataset): The federated dataset used for training and evaluation. - diff_privacy (bool, optional): Whether differential privacy is used during the federated learning process. - Default is False. - """ - self.device = torch.device( - f"cuda" if torch.cuda.is_available() else "cpu" - ) + + self.device = torch.device("cpu") self.global_model = global_model self.params = global_model.get_parameters() self.global_model.model = global_model.model.to(self.device) @@ -62,120 +85,531 @@ def __init__( self.fed_dataset = fed_dataset self.strategy = strategy self.client_resources = client_resources - setattr( - self.strategy.strategy_object, - "min_available_clients", - self.num_clients, + + self.shap_config = ( + shap_config + if shap_config is not None + else SHAPConfig(enabled=False) ) - setattr( - self.strategy.strategy_object, - "initial_parameters", - fl.common.ndarrays_to_parameters(self.params), + + self.shap_config.validate( + input_size=self.fed_dataset.size ) + + self.federated_shap_result = None + self.local_shap_results = [] + + setattr(self.strategy.strategy_object, "min_available_clients", self.num_clients) + setattr(self.strategy.strategy_object, "initial_parameters", fl.common.ndarrays_to_parameters(self.params)) setattr(self.strategy.strategy_object, "evaluate_fn", self.evaluate) - self.fed_dataset = fed_dataset + self.diff_priv = diff_privacy + + # tracking self.accuracies = [] self.losses = [] self.auc = [] + self.rmse = [] + self.mae = [] + self.r2 = [] + self.adj_r2 = [] + + self.train_losses = [] + self.train_accuracies = [] + self.train_auc = [] + self.train_rmse = [] + self.train_mae = [] + self.train_r2 = [] + self.train_adj_r2 = [] + self.flower_clients = [] + self.device = torch.device("cpu") + self.global_model = global_model + + # move model to CPU + self.global_model.model = self.global_model.model.to(self.device) + + # also move optimizer state to CPU + for state in self.global_model.optimizer.state.values(): + for k, v in state.items(): + if torch.is_tensor(v): + state[k] = v.detach().cpu() + + # if criterion has tensor buffers (e.g., pos_weight), ensure CPU too + if hasattr(self.global_model.criterion, "pos_weight") and self.global_model.criterion.pos_weight is not None: + self.global_model.criterion.pos_weight = self.global_model.criterion.pos_weight.detach().cpu() self.validate() + + def synchronize_final_global_model( + self, + ) -> None: + """ + Copy the last Flower-aggregated parameters into + the MEDfl global model wrapper. + """ + + strategy_object = ( + self.strategy.strategy_object + ) + + if strategy_object is None: + raise RuntimeError( + "The Flower strategy has not been created" + ) + + final_parameters = getattr( + strategy_object, + "latest_aggregated_parameters", + None, + ) + + if final_parameters is None: + # Your evaluate() method already records the latest + # centralized-evaluation parameter arrays. + final_arrays = getattr( + self, + "last_parameters", + None, + ) + + if final_arrays is None: + raise RuntimeError( + "The final global model parameters " + "could not be found" + ) + else: + final_arrays = ( + fl.common.parameters_to_ndarrays( + final_parameters + ) + ) + + self.global_model.set_parameters( + final_arrays + ) + + self.global_model.model.to( + self.device + ) + self.global_model.model.eval() + + logging.getLogger( + "MEDfl.server" + ).info( + "Final global parameters synchronized " + "for post-training operations" + ) + + def _get_shap_loaders( + self, + ) -> list[DataLoader]: + """ + Return one loader per client according to SHAPConfig. + """ + + split = self.shap_config.data_split + + if split == "train": + return self.fed_dataset.trainloaders + + if split == "validation": + return self.fed_dataset.valloaders + + if split == "test": + return self.fed_dataset.testloaders + + raise ValueError( + f"Unsupported SHAP data split: {split}" + ) + + def calculate_local_shap_results( + self, + ) -> list[dict]: + """ + Calculate local SHAP summaries for every client partition. + + This method runs only after federated training has finished. + """ + + if not self.shap_config.enabled: + return [] + + shap_logger = logging.getLogger( + "MEDfl.shap" + ) + + loaders = self._get_shap_loaders() + + feature_names = ( + self.shap_config.feature_names + if self.shap_config.feature_names + is not None + else self.fed_dataset.feature_names + ) + + if len(feature_names) != self.fed_dataset.size: + raise ValueError( + "SHAP feature names do not match the " + "federated dataset input size" + ) + + local_explainer = ( + LocalPyTorchSHAPExplainer( + explainer_type=( + self.shap_config.explainer + ), + background_size=( + self.shap_config.background_size + ), + explanation_size=( + self.shap_config.explanation_size + ), + random_seed=( + self.shap_config.random_seed + ), + clipping_value=( + self.shap_config.clipping_value + ), + ) + ) + + local_results: list[dict] = [] + + for client_index, loader in enumerate( + loaders + ): + client_id = str(client_index) + + if loader is None: + shap_logger.warning( + "Skipping SHAP client %s: loader is None", + client_id, + ) + continue + + available_samples = len( + loader.dataset + ) + + if available_samples == 0: + shap_logger.warning( + "Skipping SHAP client %s: empty dataset", + client_id, + ) + continue + + shap_logger.info( + "Starting local SHAP | client=%s | " + "available_samples=%d", + client_id, + available_samples, + ) + + try: + # The model is used for inference only. + # No optimizer or training method is called. + local_result = ( + local_explainer.calculate( + model=self.global_model.model, + loader=loader, + device=self.device, + ) + ) + + explained_samples = int( + local_result["sample_count"] + ) + + if ( + explained_samples + < self.shap_config.minimum_samples + ): + shap_logger.warning( + "Skipping client %s because only %d " + "samples were explained; minimum=%d", + client_id, + explained_samples, + self.shap_config.minimum_samples, + ) + continue + + local_result[ + "client_id" + ] = client_id + + local_result[ + "feature_names" + ] = list(feature_names) + + local_result[ + "available_samples" + ] = available_samples + + local_results.append( + local_result + ) + + shap_logger.info( + "Local SHAP completed | client=%s | " + "explained_samples=%d", + client_id, + explained_samples, + ) + + except Exception as exc: + shap_logger.exception( + "Local SHAP failed for client %s: %s", + client_id, + exc, + ) + + if not local_results: + raise RuntimeError( + "SHAP was enabled, but no client produced " + "a valid explanation result" + ) + + self.local_shap_results = ( + local_results + ) + + return local_results + + def calculate_federated_shap( + self, + ) -> dict: + """ + Run the complete post-training SHAP phase. + """ + + shap_logger = logging.getLogger( + "MEDfl.shap" + ) + + if not self.shap_config.enabled: + return { + "status": "disabled", + } + + shap_logger.info( + "========================================" + ) + shap_logger.info( + "Starting post-training federated SHAP" + ) + shap_logger.info( + "Explainer: %s", + self.shap_config.explainer, + ) + shap_logger.info( + "Data split: %s", + self.shap_config.data_split, + ) + shap_logger.info( + "========================================" + ) + + local_results = ( + self.calculate_local_shap_results() + ) + + federated_result = ( + aggregate_federated_shap( + client_results=local_results, + include_client_results=( + self.shap_config + .include_client_results + ), + ) + ) + + federated_result["explainer"] = ( + self.shap_config.explainer + ) + federated_result["data_split"] = ( + self.shap_config.data_split + ) + federated_result["background_size"] = ( + self.shap_config.background_size + ) + federated_result["requested_explanation_size"] = ( + self.shap_config.explanation_size + ) + + self.federated_shap_result = ( + federated_result + ) + + shap_logger.info( + "Federated SHAP completed | clients=%d | " + "samples=%d", + federated_result[ + "participating_clients" + ], + federated_result[ + "explained_samples" + ], + ) + + return federated_result def validate(self) -> None: - """Validate global_model, strategy, num_clients, num_rounds, fed_dataset, diff_privacy""" if not isinstance(self.global_model, Model): raise TypeError("global_model argument must be a Model instance") - - # if not isinstance(self.strategy, Strategy): - # print(self.strategy) - # print(isinstance(self.strategy, Strategy)) - # raise TypeError("strategy argument must be a Strategy instance") - if not isinstance(self.num_clients, int): raise TypeError("num_clients argument must be an int") - if not isinstance(self.num_rounds, int): raise TypeError("num_rounds argument must be an int") - if not isinstance(self.diff_priv, bool): raise TypeError("diff_priv argument must be a bool") def client_fn(self, cid) -> FlowerClient: - """ - Return a FlowerClient object for a specific client ID. - - Args: - cid: The client ID. + device = torch.device(f"cuda:{int(cid) % 4}" if torch.cuda.is_available() else "cpu") - Returns: - FlowerClient: A FlowerClient object representing the individual client. - """ - - device = torch.device( - f"cuda:{int(cid) % 4}" if torch.cuda.is_available() else "cpu" - ) client_model = copy.deepcopy(self.global_model) - + trainloader = self.fed_dataset.trainloaders[int(cid)] - valloader = self.fed_dataset.valloaders[int(cid)] - # this helps in making plots + + # If no validation → reuse trainloader - client = FlowerClient( - cid, client_model, trainloader, valloader, self.diff_priv - ) + if len(self.fed_dataset.valloaders[int(cid)].dataset) == 0: + valloader = trainloader + else: + valloader = self.fed_dataset.valloaders[int(cid)] + + client_model = copy.deepcopy(self.global_model) + client = FlowerClient(cid, client_model, trainloader, valloader, self.diff_priv) self.flower_clients.append(client) return client + def _concat_loader(self, loaders, shuffle: bool = False) -> DataLoader: + ds = ConcatDataset([ldr.dataset for ldr in loaders]) + batch_size = getattr(loaders[0], "batch_size", None) or 64 + num_workers = getattr(loaders[0], "num_workers", 0) + pin_memory = getattr(loaders[0], "pin_memory", False) + + return DataLoader( + ds, + batch_size=batch_size, + shuffle=shuffle, + num_workers=num_workers, + pin_memory=pin_memory, + ) + def evaluate( self, server_round: int, parameters: fl.common.NDArrays, config: Dict[str, fl.common.Scalar], ) -> Optional[Tuple[float, Dict[str, fl.common.Scalar]]]: - """ - Evaluate the global model on the validation dataset and update the accuracies and losses. - Args: - server_round (int): The current round of the federated learning process. - parameters (fl.common.NDArrays): The global model parameters. - config (Dict[str, fl.common.Scalar]): Configuration dictionary. + self.global_model.set_parameters(parameters) + self.last_parameters = parameters - Returns: - Optional[Tuple[float, Dict[str, fl.common.Scalar]]]: The evaluation loss and accuracy. - """ - testloader = self.fed_dataset.valloaders[0] - - self.global_model.set_parameters( - parameters - ) # Update model with the latest parameters - loss, accuracy ,auc = self.global_model.evaluate(testloader, self.device) - self.auc.append(auc) - self.losses.append(loss) - self.accuracies.append(accuracy) - - return loss, {"accuracy": accuracy} + global_valloader = self._concat_loader(self.fed_dataset.valloaders, shuffle=False) + global_trainloader = self._concat_loader(self.fed_dataset.trainloaders, shuffle=False) + + val_loss, val_metrics = self.global_model.evaluate(global_valloader, self.device) + self.losses.append(val_loss) + + if isinstance(val_metrics, dict): + if "accuracy" in val_metrics: self.accuracies.append(val_metrics["accuracy"]) + if "auc" in val_metrics: self.auc.append(val_metrics["auc"]) + if "rmse" in val_metrics: self.rmse.append(val_metrics["rmse"]) + if "mae" in val_metrics: self.mae.append(val_metrics["mae"]) + if "r2" in val_metrics: self.r2.append(val_metrics["r2"]) + if "adj_r2" in val_metrics: self.adj_r2.append(val_metrics["adj_r2"]) + + train_loss, train_metrics = self.global_model.evaluate(global_trainloader, self.device) + self.train_losses.append(train_loss) + + if isinstance(train_metrics, dict): + if "accuracy" in train_metrics: self.train_accuracies.append(train_metrics["accuracy"]) + if "auc" in train_metrics: self.train_auc.append(train_metrics["auc"]) + if "rmse" in train_metrics: self.train_rmse.append(train_metrics["rmse"]) + if "mae" in train_metrics: self.train_mae.append(train_metrics["mae"]) + if "r2" in train_metrics: self.train_r2.append(train_metrics["r2"]) + if "adj_r2" in train_metrics: self.train_adj_r2.append(train_metrics["adj_r2"]) + + out_metrics = dict(val_metrics or {}) + out_metrics["train_loss"] = float(train_loss) if train_loss is not None else float("nan") + + if isinstance(train_metrics, dict): + for k, v in train_metrics.items(): + out_metrics[f"train_{k}"] = float(v) if v is not None else float("nan") + + return val_loss, out_metrics + + def score_each_client(self, parameters=None, split: str = "val"): + if parameters is not None: + self.global_model.set_parameters(parameters) + + loaders = self.fed_dataset.valloaders if split == "val" else self.fed_dataset.trainloaders + + results = {} + for cid, loader in enumerate(loaders): + loss, metrics = self.global_model.evaluate(loader, self.device) + results[str(cid)] = { + "num_examples": len(loader.dataset), + "loss": float(loss), + **{k: float(v) for k, v in (metrics or {}).items()}, + } + return results def run(self) -> None: - """ - Run the federated learning process using Flower simulation. + # --------------------------------------------------------------- + # KEY FIX 1: log_to_driver=True is required, AND we must redirect + # Ray worker stdout/stderr back to the driver. + # --------------------------------------------------------------- + ray_init_args = { + "include_dashboard": False, + "object_store_memory": 78643200, + "log_to_driver": True, # sends worker logs → driver + "logging_level": logging.INFO, # Ray's own log level + } + + # --------------------------------------------------------------- + # KEY FIX 2: make sure the driver-side logger is set up before + # simulation starts so Flower's internal messages reach the notebook. + # --------------------------------------------------------------- + setup_fl_logging() - Returns: - History: The history of the accuracies and losses during the training of each node - """ - # Increase the object store memory to the minimum allowed value or higher - ray_init_args = {"include_dashboard": False - , "object_store_memory": 78643200 - } self.fed_dataset.eng = None - + history = fl.simulation.start_simulation( client_fn=self.client_fn, num_clients=self.num_clients, config=fl.server.ServerConfig(self.num_rounds), strategy=self.strategy.strategy_object, ray_init_args=ray_init_args, - client_resources = self.client_resources + client_resources=self.client_resources ) + # Training has now completely finished. + self.synchronize_final_global_model() + + if self.shap_config.enabled: + try: + self.federated_shap_result = ( + self.calculate_federated_shap() + ) + + except Exception as exc: + logging.getLogger( + "MEDfl.shap" + ).exception( + "Post-training federated SHAP failed: %s", + exc, + ) - return history + # Important: a SHAP failure does not mean + # federated training failed. + self.federated_shap_result = { + "status": "failed", + "error": str(exc), + } + else: + self.federated_shap_result = { + "status": "disabled", + } + return history \ No newline at end of file diff --git a/build/lib/MEDfl/LearningManager/shap/__init__.py b/build/lib/MEDfl/LearningManager/shap/__init__.py new file mode 100644 index 0000000..8805c35 --- /dev/null +++ b/build/lib/MEDfl/LearningManager/shap/__init__.py @@ -0,0 +1,14 @@ +from .config import SHAPConfig +from .explainer import LocalPyTorchSHAPExplainer +from .xgboost_explainer import ( + LocalXGBoostSHAPExplainer, +) +from .aggregation import aggregate_federated_shap + + +__all__ = [ + "SHAPConfig", + "LocalPyTorchSHAPExplainer", + "LocalXGBoostSHAPExplainer", + "aggregate_federated_shap", +] \ No newline at end of file diff --git a/build/lib/MEDfl/LearningManager/shap/aggregation.py b/build/lib/MEDfl/LearningManager/shap/aggregation.py new file mode 100644 index 0000000..0e70e69 --- /dev/null +++ b/build/lib/MEDfl/LearningManager/shap/aggregation.py @@ -0,0 +1,249 @@ +from __future__ import annotations + +from typing import Any + +import numpy as np + + +def _validate_client_results( + client_results: list[dict], +) -> tuple[list[str], int]: + if not client_results: + raise ValueError( + "No client SHAP results were provided" + ) + + reference_features = client_results[0][ + "feature_names" + ] + + feature_count = len(reference_features) + + required_array_keys = { + "abs_sum", + "signed_sum", + "square_sum", + "positive_count", + "negative_count", + "zero_count", + } + + for result in client_results: + if result["feature_names"] != reference_features: + raise ValueError( + "Clients used different feature names or orders" + ) + + if int(result["sample_count"]) <= 0: + raise ValueError( + f"Client {result.get('client_id')} returned an " + "invalid SHAP sample count" + ) + + for key in required_array_keys: + if key not in result: + raise ValueError( + f"Client result is missing '{key}'" + ) + + values = np.asarray(result[key]) + + if values.shape != (feature_count,): + raise ValueError( + f"Client {result.get('client_id')} returned " + f"invalid shape {values.shape} for {key}. " + f"Expected ({feature_count},)." + ) + + return reference_features, feature_count + + +def aggregate_federated_shap( + client_results: list[dict], + include_client_results: bool = True, +) -> dict[str, Any]: + """ + Aggregate additive local SHAP statistics. + + The global mean absolute SHAP is equivalent to calculating: + + sum of all |SHAP values| / number of all explained samples + """ + + feature_names, feature_count = ( + _validate_client_results(client_results) + ) + + total_samples = sum( + int(result["sample_count"]) + for result in client_results + ) + + abs_sum = np.zeros( + feature_count, + dtype=np.float64, + ) + signed_sum = np.zeros( + feature_count, + dtype=np.float64, + ) + square_sum = np.zeros( + feature_count, + dtype=np.float64, + ) + positive_count = np.zeros( + feature_count, + dtype=np.float64, + ) + negative_count = np.zeros( + feature_count, + dtype=np.float64, + ) + zero_count = np.zeros( + feature_count, + dtype=np.float64, + ) + + for result in client_results: + abs_sum += np.asarray( + result["abs_sum"], + dtype=np.float64, + ) + signed_sum += np.asarray( + result["signed_sum"], + dtype=np.float64, + ) + square_sum += np.asarray( + result["square_sum"], + dtype=np.float64, + ) + positive_count += np.asarray( + result["positive_count"], + dtype=np.float64, + ) + negative_count += np.asarray( + result["negative_count"], + dtype=np.float64, + ) + zero_count += np.asarray( + result["zero_count"], + dtype=np.float64, + ) + + mean_abs = abs_sum / total_samples + mean_signed = signed_sum / total_samples + mean_square = square_sum / total_samples + + # Var(X) = E(X²) - E(X)² + variance = np.maximum( + mean_square - np.square(mean_signed), + 0.0, + ) + + standard_deviation = np.sqrt(variance) + + features = [] + + for index, feature_name in enumerate( + feature_names + ): + features.append( + { + "feature": feature_name, + "feature_index": index, + "mean_abs_shap": float( + mean_abs[index] + ), + "mean_signed_shap": float( + mean_signed[index] + ), + "shap_std": float( + standard_deviation[index] + ), + "positive_rate": float( + positive_count[index] + / total_samples + ), + "negative_rate": float( + negative_count[index] + / total_samples + ), + "zero_rate": float( + zero_count[index] + / total_samples + ), + } + ) + + # Highest global importance first. + features.sort( + key=lambda item: item["mean_abs_shap"], + reverse=True, + ) + + final_result: dict[str, Any] = { + "status": "completed", + "participating_clients": len( + client_results + ), + "explained_samples": int( + total_samples + ), + "feature_importance": features, + } + + if include_client_results: + client_summaries = [] + + for result in client_results: + sample_count = int( + result["sample_count"] + ) + + client_mean_abs = ( + np.asarray( + result["abs_sum"], + dtype=np.float64, + ) + / sample_count + ) + + local_features = [ + { + "feature": feature_name, + "feature_index": index, + "mean_abs_shap": float( + client_mean_abs[index] + ), + } + for index, feature_name in enumerate( + feature_names + ) + ] + + local_features.sort( + key=lambda item: item[ + "mean_abs_shap" + ], + reverse=True, + ) + + client_summaries.append( + { + "client_id": result[ + "client_id" + ], + "explained_samples": ( + sample_count + ), + "feature_importance": ( + local_features + ), + } + ) + + final_result["client_results"] = ( + client_summaries + ) + + return final_result \ No newline at end of file diff --git a/build/lib/MEDfl/LearningManager/shap/config.py b/build/lib/MEDfl/LearningManager/shap/config.py new file mode 100644 index 0000000..7eb810c --- /dev/null +++ b/build/lib/MEDfl/LearningManager/shap/config.py @@ -0,0 +1,196 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import List, Literal, Optional + + +@dataclass +class SHAPConfig: + """ + Configuration for the optional post-training federated SHAP phase. + """ + + enabled: bool = False + + # Neural networks: + # gradient + # deep + # + # XGBoost: + # tree + explainer: Literal[ + "gradient", + "deep", + "tree", + ] = "gradient" + + data_split: Literal[ + "train", + "validation", + "test", + ] = "validation" + + background_size: int = 100 + explanation_size: int = 500 + random_seed: int = 42 + + feature_names: Optional[List[str]] = None + + minimum_samples: int = 10 + clipping_value: Optional[float] = None + + include_client_results: bool = True + + # TreeExplainer output: + # + # raw: + # Binary classification → log-odds margin + # Regression → raw prediction + # + # probability: + # Binary classification probability when supported by the + # installed SHAP/XGBoost versions. + model_output: Literal[ + "raw", + "probability", + ] = "raw" + + def validate( + self, + input_size: int, + ) -> None: + if not isinstance(self.enabled, bool): + raise TypeError( + "SHAP enabled must be a boolean" + ) + + if self.explainer not in { + "gradient", + "deep", + "tree", + }: + raise ValueError( + "SHAP explainer must be 'gradient', " + "'deep', or 'tree'" + ) + + if self.data_split not in { + "train", + "validation", + "test", + }: + raise ValueError( + "SHAP data_split must be train, " + "validation, or test" + ) + + if self.background_size <= 0: + raise ValueError( + "SHAP background_size must be greater than zero" + ) + + if self.explanation_size <= 0: + raise ValueError( + "SHAP explanation_size must be greater than zero" + ) + + if self.minimum_samples <= 0: + raise ValueError( + "SHAP minimum_samples must be greater than zero" + ) + + if self.clipping_value is not None: + if self.clipping_value <= 0: + raise ValueError( + "SHAP clipping_value must be greater than zero" + ) + + if self.model_output not in { + "raw", + "probability", + }: + raise ValueError( + "SHAP model_output must be 'raw' or " + "'probability'" + ) + + if self.feature_names is not None: + if len(self.feature_names) != input_size: + raise ValueError( + "The number of SHAP feature names must " + "match the model input size. " + f"Expected {input_size}, " + f"received {len(self.feature_names)}." + ) + + if len(set(self.feature_names)) != len( + self.feature_names + ): + raise ValueError( + "SHAP feature names must be unique" + ) + + def validate_backend( + self, + input_size: int, + backend: str, + task: Optional[str] = None, + ) -> None: + """ + Validate both the general configuration and the selected + machine-learning backend. + """ + + self.validate(input_size) + + backend = str(backend).lower() + + if backend == "pytorch": + if self.explainer not in { + "gradient", + "deep", + }: + raise ValueError( + "PyTorch models require explainer='gradient' " + "or explainer='deep'" + ) + + elif backend == "xgboost": + if self.explainer != "tree": + raise ValueError( + "XGBoost models require explainer='tree'" + ) + + if task == "multiclass": + raise ValueError( + "The current federated XGBoost SHAP " + "implementation supports binary classification " + "and regression. Multiclass requires " + "class-specific SHAP aggregation." + ) + + if ( + task == "regression" + and self.model_output == "probability" + ): + raise ValueError( + "model_output='probability' is not valid for " + "XGBoost regression" + ) + + else: + raise ValueError( + f"Unsupported SHAP backend: {backend}" + ) + + def resolved_feature_names( + self, + input_size: int, + ) -> List[str]: + if self.feature_names is not None: + return list(self.feature_names) + + return [ + f"feature_{index}" + for index in range(input_size) + ] \ No newline at end of file diff --git a/build/lib/MEDfl/LearningManager/shap/explainer.py b/build/lib/MEDfl/LearningManager/shap/explainer.py new file mode 100644 index 0000000..b884709 --- /dev/null +++ b/build/lib/MEDfl/LearningManager/shap/explainer.py @@ -0,0 +1,324 @@ +from __future__ import annotations + +import logging +from typing import Any + +import numpy as np +import torch +from torch.utils.data import DataLoader + + +logger = logging.getLogger("MEDfl.shap") + + +class LocalPyTorchSHAPExplainer: + """ + Calculate SHAP values for one client partition. + + Raw samples and per-sample SHAP values are not returned. + Only additive summary statistics are returned. + """ + + def __init__( + self, + explainer_type: str = "gradient", + background_size: int = 100, + explanation_size: int = 500, + random_seed: int = 42, + clipping_value: float | None = None, + ) -> None: + self.explainer_type = explainer_type + self.background_size = background_size + self.explanation_size = explanation_size + self.random_seed = random_seed + self.clipping_value = clipping_value + + @staticmethod + def _collect_features( + loader: DataLoader, + ) -> torch.Tensor: + """ + Collect only X values from the selected client loader. + Labels are not required to calculate SHAP values. + """ + feature_batches: list[torch.Tensor] = [] + + for batch_index, batch in enumerate(loader): + if not isinstance(batch, (tuple, list)): + raise TypeError( + "Each DataLoader batch must be a tuple or list" + ) + + if len(batch) == 0: + raise ValueError( + f"Empty batch received at index {batch_index}" + ) + + X = batch[0] + + if not isinstance(X, torch.Tensor): + X = torch.as_tensor(X) + + X = X.detach().cpu().float() + + if torch.isnan(X).any(): + raise ValueError( + f"SHAP input contains NaN values in batch " + f"{batch_index}" + ) + + if torch.isinf(X).any(): + raise ValueError( + f"SHAP input contains infinite values in batch " + f"{batch_index}" + ) + + feature_batches.append(X) + + if not feature_batches: + raise ValueError( + "Cannot calculate SHAP using an empty DataLoader" + ) + + features = torch.cat( + feature_batches, + dim=0, + ) + + # First implementation: tabular data. + if features.ndim != 2: + raise ValueError( + "The current SHAP implementation supports tabular " + "input with shape (samples, features). " + f"Received shape {tuple(features.shape)}." + ) + + return features + + @staticmethod + def _sample_rows( + data: torch.Tensor, + requested_size: int, + generator: torch.Generator, + ) -> torch.Tensor: + selected_size = min( + requested_size, + len(data), + ) + + indices = torch.randperm( + len(data), + generator=generator, + )[:selected_size] + + return data[indices] + + @staticmethod + def _normalize_shap_values( + raw_values: Any, + sample_count: int, + feature_count: int, + ) -> np.ndarray: + """ + Normalize the different shapes returned by SHAP versions. + + Expected final shape: + (number_of_samples, number_of_features) + """ + + if hasattr(raw_values, "values"): + raw_values = raw_values.values + + if isinstance(raw_values, list): + if len(raw_values) != 1: + raise ValueError( + "The current implementation supports binary or " + "single-output models only. Multiclass SHAP " + "requires class-specific aggregation." + ) + + raw_values = raw_values[0] + + values = np.asarray(raw_values) + + # Possible shape: (samples, features, 1) + if ( + values.ndim == 3 + and values.shape[-1] == 1 + ): + values = values[..., 0] + + # Possible shape: (1, samples, features) + if ( + values.ndim == 3 + and values.shape[0] == 1 + ): + values = values[0] + + expected_shape = ( + sample_count, + feature_count, + ) + + if values.shape != expected_shape: + raise ValueError( + "Unexpected SHAP output shape. " + f"Expected {expected_shape}, received {values.shape}." + ) + + values = values.astype( + np.float64, + copy=False, + ) + + if np.isnan(values).any(): + raise ValueError( + "Calculated SHAP values contain NaN values" + ) + + if np.isinf(values).any(): + raise ValueError( + "Calculated SHAP values contain infinite values" + ) + + return values + + def _create_explainer( + self, + model: torch.nn.Module, + background: torch.Tensor, + ): + try: + import shap + except ImportError as exc: + raise ImportError( + "SHAP is not installed. Install it with: " + "pip install shap" + ) from exc + + if self.explainer_type == "gradient": + return shap.GradientExplainer( + model, + background, + ) + + if self.explainer_type == "deep": + return shap.DeepExplainer( + model, + background, + ) + + raise ValueError( + f"Unsupported SHAP explainer: {self.explainer_type}" + ) + + def calculate( + self, + model: torch.nn.Module, + loader: DataLoader, + device: torch.device, + ) -> dict: + """ + Calculate SHAP for one client partition. + + The returned values are additive summaries. They can be + combined across clients without transmitting raw SHAP rows. + """ + + features = self._collect_features(loader) + + if len(features) < 2: + raise ValueError( + "At least two samples are required for SHAP" + ) + + generator = torch.Generator() + generator.manual_seed(self.random_seed) + + background = self._sample_rows( + data=features, + requested_size=self.background_size, + generator=generator, + ) + + explanation_data = self._sample_rows( + data=features, + requested_size=self.explanation_size, + generator=generator, + ) + + background = background.to(device) + explanation_data = explanation_data.to(device) + + model = model.to(device) + model.eval() + + logger.info( + "Calculating local SHAP | " + "explainer=%s | available=%d | background=%d | " + "explained=%d", + self.explainer_type, + len(features), + len(background), + len(explanation_data), + ) + + explainer = self._create_explainer( + model=model, + background=background, + ) + + raw_values = explainer.shap_values( + explanation_data + ) + + shap_values = self._normalize_shap_values( + raw_values=raw_values, + sample_count=len(explanation_data), + feature_count=explanation_data.shape[1], + ) + + if self.clipping_value is not None: + shap_values = np.clip( + shap_values, + -self.clipping_value, + self.clipping_value, + ) + + # Additive summaries are preferable because they can later + # be used with distributed or secure aggregation. + result = { + "sample_count": int(shap_values.shape[0]), + + "abs_sum": np.sum( + np.abs(shap_values), + axis=0, + ), + + "signed_sum": np.sum( + shap_values, + axis=0, + ), + + "square_sum": np.sum( + np.square(shap_values), + axis=0, + ), + + "positive_count": np.sum( + shap_values > 0, + axis=0, + ).astype(np.int64), + + "negative_count": np.sum( + shap_values < 0, + axis=0, + ).astype(np.int64), + + "zero_count": np.sum( + shap_values == 0, + axis=0, + ).astype(np.int64), + } + + return result \ No newline at end of file diff --git a/build/lib/MEDfl/LearningManager/shap/xgboost_explainer.py b/build/lib/MEDfl/LearningManager/shap/xgboost_explainer.py new file mode 100644 index 0000000..00f251f --- /dev/null +++ b/build/lib/MEDfl/LearningManager/shap/xgboost_explainer.py @@ -0,0 +1,436 @@ +from __future__ import annotations + +import logging +from typing import Any, Dict, List, Optional + +import numpy as np +import torch +import xgboost as xgb +from torch.utils.data import DataLoader + + +logger = logging.getLogger("MEDfl.shap.xgboost") + + +class LocalXGBoostSHAPExplainer: + """ + Calculate local SHAP values for a federated XGBoost client. + + The final federated XGBoost booster is explained separately on + every client's selected local partition. + + Raw samples and raw per-sample SHAP values are never returned. + Only additive summary statistics are returned so the existing + federated SHAP aggregation function can be reused. + """ + + def __init__( + self, + background_size: int = 100, + explanation_size: int = 500, + random_seed: int = 42, + clipping_value: Optional[float] = None, + model_output: str = "raw", + ) -> None: + self.background_size = int(background_size) + self.explanation_size = int(explanation_size) + self.random_seed = int(random_seed) + self.clipping_value = clipping_value + self.model_output = str(model_output) + + self._validate_configuration() + + def _validate_configuration(self) -> None: + if self.background_size <= 0: + raise ValueError( + "background_size must be greater than zero" + ) + + if self.explanation_size <= 0: + raise ValueError( + "explanation_size must be greater than zero" + ) + + if self.clipping_value is not None: + if self.clipping_value <= 0: + raise ValueError( + "clipping_value must be greater than zero" + ) + + if self.model_output not in { + "raw", + "probability", + }: + raise ValueError( + "model_output must be 'raw' or 'probability'" + ) + + @staticmethod + def _collect_features( + loader: DataLoader, + ) -> np.ndarray: + """ + Collect only the input features from a client DataLoader. + + Labels are not needed for SHAP calculation. + """ + + if loader is None: + raise ValueError( + "Cannot calculate XGBoost SHAP using a None loader" + ) + + feature_batches: List[np.ndarray] = [] + + for batch_index, batch in enumerate(loader): + if not isinstance(batch, (tuple, list)): + raise TypeError( + "Each DataLoader batch must be a tuple or list" + ) + + if len(batch) == 0: + raise ValueError( + f"Empty batch received at index {batch_index}" + ) + + features = batch[0] + + if isinstance(features, torch.Tensor): + features = ( + features + .detach() + .cpu() + .numpy() + ) + + features = np.asarray(features) + + if features.ndim == 1: + features = features.reshape(1, -1) + + if features.ndim != 2: + raise ValueError( + "The XGBoost SHAP implementation supports " + "two-dimensional tabular inputs. " + f"Received shape {features.shape}." + ) + + if not np.issubdtype( + features.dtype, + np.number, + ): + try: + features = features.astype( + np.float32 + ) + except Exception as exc: + raise TypeError( + "XGBoost SHAP features must be numeric" + ) from exc + + features = features.astype( + np.float32, + copy=False, + ) + + if np.isnan(features).any(): + raise ValueError( + "XGBoost SHAP input contains NaN values " + f"in batch {batch_index}" + ) + + if np.isinf(features).any(): + raise ValueError( + "XGBoost SHAP input contains infinite values " + f"in batch {batch_index}" + ) + + feature_batches.append(features) + + if not feature_batches: + raise ValueError( + "Cannot calculate XGBoost SHAP using an empty " + "DataLoader" + ) + + collected_features = np.concatenate( + feature_batches, + axis=0, + ) + + if collected_features.ndim != 2: + raise ValueError( + "Unexpected collected XGBoost feature shape: " + f"{collected_features.shape}" + ) + + return collected_features + + @staticmethod + def _sample_rows( + data: np.ndarray, + requested_size: int, + random_generator: np.random.Generator, + ) -> np.ndarray: + selected_size = min( + int(requested_size), + int(data.shape[0]), + ) + + indices = random_generator.choice( + data.shape[0], + size=selected_size, + replace=False, + ) + + return data[indices] + + @staticmethod + def _normalize_shap_values( + raw_values: Any, + sample_count: int, + feature_count: int, + ) -> np.ndarray: + """ + Normalize SHAP output to: + + (number_of_samples, number_of_features) + + This first implementation supports binary classification and + regression. Multiclass SHAP requires class-specific aggregation. + """ + + if hasattr(raw_values, "values"): + raw_values = raw_values.values + + if isinstance(raw_values, list): + if len(raw_values) != 1: + raise ValueError( + "Multiclass XGBoost SHAP is not yet supported. " + "The returned SHAP result contains several " + "class-specific arrays." + ) + + raw_values = raw_values[0] + + shap_values = np.asarray(raw_values) + + # Possible binary/single-output shape: + # (samples, features, 1) + if ( + shap_values.ndim == 3 + and shap_values.shape[-1] == 1 + ): + shap_values = shap_values[..., 0] + + # Possible shape: + # (1, samples, features) + if ( + shap_values.ndim == 3 + and shap_values.shape[0] == 1 + ): + shap_values = shap_values[0] + + if shap_values.ndim == 3: + raise ValueError( + "Multiclass XGBoost SHAP is not yet supported. " + "Class-specific SHAP values were returned with " + f"shape {shap_values.shape}." + ) + + expected_shape = ( + sample_count, + feature_count, + ) + + if shap_values.shape != expected_shape: + raise ValueError( + "Unexpected XGBoost SHAP output shape. " + f"Expected {expected_shape}, " + f"received {shap_values.shape}." + ) + + shap_values = shap_values.astype( + np.float64, + copy=False, + ) + + if np.isnan(shap_values).any(): + raise ValueError( + "Calculated XGBoost SHAP values contain NaN values" + ) + + if np.isinf(shap_values).any(): + raise ValueError( + "Calculated XGBoost SHAP values contain " + "infinite values" + ) + + return shap_values + + def _create_explainer( + self, + booster: xgb.Booster, + background: np.ndarray, + ): + try: + import shap + except ImportError as exc: + raise ImportError( + "SHAP is enabled for XGBoost, but the shap " + "package is not installed." + ) from exc + + if not isinstance(booster, xgb.Booster): + raise TypeError( + "LocalXGBoostSHAPExplainer expects an " + "xgboost.Booster" + ) + + try: + return shap.TreeExplainer( + model=booster, + data=background, + feature_perturbation="interventional", + model_output=self.model_output, + ) + + except Exception as exc: + # Some SHAP/XGBoost combinations may not support every + # TreeExplainer argument. Keep a conservative fallback. + logger.warning( + "Could not initialize TreeExplainer with the " + "configured background and model output. " + "Falling back to the default TreeExplainer: %s", + exc, + ) + + return shap.TreeExplainer( + model=booster, + ) + + @staticmethod + def _calculate_raw_values( + explainer, + explanation_data: np.ndarray, + ): + """ + Support different SHAP versions. + """ + + try: + return explainer.shap_values( + explanation_data, + check_additivity=False, + ) + + except TypeError: + # Older SHAP versions may not accept + # check_additivity here. + return explainer.shap_values( + explanation_data + ) + + def calculate( + self, + booster: xgb.Booster, + loader: DataLoader, + ) -> Dict[str, Any]: + """ + Calculate local XGBoost SHAP additive summaries. + """ + + features = self._collect_features(loader) + + if len(features) < 2: + raise ValueError( + "At least two samples are required for " + "XGBoost SHAP" + ) + + random_generator = np.random.default_rng( + self.random_seed + ) + + background = self._sample_rows( + data=features, + requested_size=self.background_size, + random_generator=random_generator, + ) + + explanation_data = self._sample_rows( + data=features, + requested_size=self.explanation_size, + random_generator=random_generator, + ) + + logger.info( + "Calculating local XGBoost SHAP | " + "available=%d | background=%d | explained=%d | " + "model_output=%s", + len(features), + len(background), + len(explanation_data), + self.model_output, + ) + + explainer = self._create_explainer( + booster=booster, + background=background, + ) + + raw_values = self._calculate_raw_values( + explainer=explainer, + explanation_data=explanation_data, + ) + + shap_values = self._normalize_shap_values( + raw_values=raw_values, + sample_count=len(explanation_data), + feature_count=explanation_data.shape[1], + ) + + if self.clipping_value is not None: + shap_values = np.clip( + shap_values, + -self.clipping_value, + self.clipping_value, + ) + + result: Dict[str, Any] = { + "sample_count": int( + shap_values.shape[0] + ), + + "abs_sum": np.sum( + np.abs(shap_values), + axis=0, + ), + + "signed_sum": np.sum( + shap_values, + axis=0, + ), + + "square_sum": np.sum( + np.square(shap_values), + axis=0, + ), + + "positive_count": np.sum( + shap_values > 0, + axis=0, + ).astype(np.int64), + + "negative_count": np.sum( + shap_values < 0, + axis=0, + ).astype(np.int64), + + "zero_count": np.sum( + shap_values == 0, + axis=0, + ).astype(np.int64), + } + + return result \ No newline at end of file diff --git a/build/lib/MEDfl/LearningManager/strategy.py b/build/lib/MEDfl/LearningManager/strategy.py index 038d8e6..f746f19 100644 --- a/build/lib/MEDfl/LearningManager/strategy.py +++ b/build/lib/MEDfl/LearningManager/strategy.py @@ -1,31 +1,83 @@ +""" +MEDfl Strategy wrapper for Flower simulation (Ray) -from collections import OrderedDict -from typing import Dict, List, Optional, Tuple +This file provides a Strategy factory that wraps any Flower server strategy +(e.g., FedAvg, FedProx, FedAdam...) with extra debugging: -import flwr as fl -import numpy as np +✅ What it adds: +- Per-client fit/eval history with round + num_examples (+ loss for eval) +- Per-round server summary history (results/failures, cids, example stats) +- Global-weights fingerprints after aggregation: + * w_norm, delta_norm, cos_sim_prev, nan_or_inf +- Uses Flower logging (NOT print) so logs show up reliably in simulation +- Optional metrics aggregation functions to remove: + "No fit_metrics_aggregation_fn provided" warning + +Usage (server side): + strat_factory = Strategy(name="FedAvg", debug_print=True) + strategy = strat_factory.create_strategy() + + fl.simulation.start_simulation(..., strategy=strategy, ...) +After run: + strategy.round_history + strategy.client_fit_history + strategy.client_eval_history +""" + +from __future__ import annotations + +from typing import Any, Dict, Optional, Type, List, Tuple, Callable + +import numpy as np +import flwr as fl import optuna +import logging +Metrics = Dict[str, float] +MetricsAggFn = Callable[[List[Tuple[int, Metrics]]], Metrics] -class Strategy: +def weighted_avg(metrics: List[Tuple[int, Metrics]]) -> Metrics: """ - A class representing a strategy for Federated Learning. - - Attributes: - name (str): The name of the strategy. Default is "FedAvg". - fraction_fit (float): Fraction of clients to use for training during each round. Default is 1.0. - fraction_evaluate (float): Fraction of clients to use for evaluation during each round. Default is 1.0. - min_fit_clients (int): Minimum number of clients to use for training during each round. Default is 2. - min_evaluate_clients (int): Minimum number of clients to use for evaluation during each round. Default is 2. - min_available_clients (int): Minimum number of available clients required to start a round. Default is 2. - initial_parameters (Optional[]): The initial parameters of the server model - Methods: - + Weighted average of client metrics by num_examples. + + Flower passes: + metrics = [(num_examples, {"rmse": 1.2, "mae": 0.5}), ...] """ + if not metrics: + return {} + + total = sum(int(n) for n, _ in metrics) + if total <= 0: + return {} + + # union of all keys + keys = set() + for _, m in metrics: + keys |= set(m.keys()) + out: Dict[str, float] = {} + for k in keys: + acc = 0.0 + used = 0 + for n, m in metrics: + if k not in m: + continue + try: + v = float(m[k]) + except Exception: + continue + acc += float(n) * v + used += int(n) + if used > 0: + out[k] = acc / float(total) + + return out + + +class Strategy: def __init__( self, name: str = "FedAvg", @@ -34,49 +86,267 @@ def __init__( min_fit_clients: int = 2, min_evaluate_clients: int = 2, min_available_clients: int = 2, - initial_parameters = [], - evaluation_methode = "centralized" + initial_parameters=None, + evaluation_methode: str = "centralized", + strategy_kwargs: Optional[Dict[str, Any]] = None, + debug_print: bool = True, + enable_metrics_aggregation: bool = True, # ✅ remove "No fit_metrics_aggregation_fn" warnings ) -> None: - """ - Initialize a Strategy object with the specified parameters. - - Args: - name (str): The name of the strategy. Default is "FedAvg". - fraction_fit (float): Fraction of clients to use for training during each round. Default is 1.0. - fraction_evaluate (float): Fraction of clients to use for evaluation during each round. Default is 1.0. - min_fit_clients (int): Minimum number of clients to use for training during each round. Default is 2. - min_evaluate_clients (int): Minimum number of clients to use for evaluation during each round. Default is 2. - min_available_clients (int): Minimum number of available clients required to start a round. Default is 2. - initial_parameters (Optional[]): The initial parametres of the server model - evaluation_methode ( "centralized" | "distributed") - """ self.fraction_fit = fraction_fit self.fraction_evaluate = fraction_evaluate self.min_fit_clients = min_fit_clients self.min_evaluate_clients = min_evaluate_clients self.min_available_clients = min_available_clients - self.initial_parameters = initial_parameters + self.initial_parameters = initial_parameters if initial_parameters is not None else [] self.evaluate_fn = None self.name = name - - def optuna_fed_optimization(self, direction:str , hpo_rate:int , params_config): + self.evaluation_methode = evaluation_methode + self.strategy_kwargs = strategy_kwargs or {} + self.debug_print = debug_print + self.enable_metrics_aggregation = enable_metrics_aggregation + + # Exposed after create_strategy() + self.strategy_object: Optional[fl.server.strategy.Strategy] = None + + # Optuna + self.study: Optional[optuna.Study] = None + self.hpo_rate: Optional[int] = None + self.params_config: Optional[dict] = None + + def optuna_fed_optimization(self, direction: str, hpo_rate: int, params_config): self.study = optuna.create_study(direction=direction) self.hpo_rate = hpo_rate self.params_config = params_config - - + + def get_strategy_by_name(self): + try: + return getattr(fl.server.strategy, self.name) + except AttributeError as e: + available = [n for n in dir(fl.server.strategy) if n and n[0].isupper()] + raise ValueError(f"Unknown strategy '{self.name}'. Available: {available}") from e + def create_strategy(self): - self.strategy_object = self.get_strategy_by_name()( + BaseStrategy: Type[fl.server.strategy.Strategy] = self.get_strategy_by_name() + + debug_print = self.debug_print + + class TrackingStrategy(BaseStrategy): # type: ignore[misc] + """ + Wraps a Flower Strategy and adds: + - per-client fit/eval history + - per-round summary history + - global weights fingerprints after aggregation + - robust logging (shows in simulation logs) + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + # cid -> list of dicts per round + self.client_fit_history: Dict[str, list] = {} + self.client_eval_history: Dict[str, list] = {} + + # list of dicts (round summaries) + self.round_history: list = [] + + # previous global model (flattened) for delta/cos sim + self._prev_global_flat: Optional[np.ndarray] = None + + self.latest_aggregated_parameters = None + + # logger (Flower config will handle output) + self._log = logging.getLogger("flwr") + self._debug_print = debug_print + + # prove that THIS strategy is active + if self._debug_print: + self._log.warning("✅ TrackingStrategy is ACTIVE (debug enabled)") + + # ---------- helpers ---------- + def _params_to_flat(self, parameters: fl.common.Parameters) -> np.ndarray: + nds = fl.common.parameters_to_ndarrays(parameters) + if not nds: + return np.array([], dtype=np.float64) + return np.concatenate([x.reshape(-1).astype(np.float64, copy=False) for x in nds]) + + def _fingerprint(self, flat: np.ndarray) -> Dict[str, Any]: + if flat.size == 0: + return { + "w_norm": 0.0, + "nan_or_inf": False, + "delta_norm": None, + "cos_sim_prev": None, + } + + w_norm = float(np.linalg.norm(flat)) + nan_or_inf = bool(np.isnan(flat).any() or np.isinf(flat).any()) + + delta_norm = None + cos_sim = None + if self._prev_global_flat is not None and self._prev_global_flat.size == flat.size: + diff = flat - self._prev_global_flat + delta_norm = float(np.linalg.norm(diff)) + + denom = np.linalg.norm(flat) * np.linalg.norm(self._prev_global_flat) + if denom > 0: + cos_sim = float(np.dot(flat, self._prev_global_flat) / denom) + + return { + "w_norm": w_norm, + "nan_or_inf": nan_or_inf, + "delta_norm": delta_norm, + "cos_sim_prev": cos_sim, + } + + def _safe_failure_summaries(self, failures: list, cap: int = 5) -> List[str]: + return [str(f) for f in failures][:cap] + + # ---------- overrides ---------- + def aggregate_fit(self, server_round, results, failures): + # Store per-client fit history WITH round + num_examples + cids: List[str] = [] + num_examples_list: List[int] = [] + + for client_proxy, fit_res in results: + cid = getattr(client_proxy, "cid", "unknown") + cids.append(cid) + + ne = int(getattr(fit_res, "num_examples", 0)) + num_examples_list.append(ne) + + entry = { + "round": int(server_round), + "num_examples": ne, + **dict(fit_res.metrics or {}), + } + # store client weights only for regression runs + if fit_res.metrics is not None and "train_rmse" in fit_res.metrics: + client_nds = fl.common.parameters_to_ndarrays(fit_res.parameters) + entry["weights"] = [arr.copy() for arr in client_nds] + self.client_fit_history.setdefault(cid, []).append(entry) + + # Call base aggregation + agg = super().aggregate_fit(server_round, results, failures) + + + # Round summary + global fingerprint (only if aggregation produced parameters) + if agg is not None: + parameters_agg, metrics_agg = agg # type: ignore[misc] + self.latest_aggregated_parameters = (parameters_agg) + flat = self._params_to_flat(parameters_agg) + fp = self._fingerprint(flat) + self._prev_global_flat = flat + + nds = fl.common.parameters_to_ndarrays(parameters_agg) + + round_row = { + "round": int(server_round), + "phase": "fit", + "n_results": int(len(results)), + "n_failures": int(len(failures)), + "cids": cids, + "total_examples": int(sum(num_examples_list)), + "min_examples": int(min(num_examples_list)) if num_examples_list else 0, + "max_examples": int(max(num_examples_list)) if num_examples_list else 0, + "failures": self._safe_failure_summaries(failures), + "weights": [arr.copy() for arr in nds], # <- add this + **fp, + **(dict(metrics_agg or {})), + } + + self.round_history.append(round_row) + + if self._debug_print: + self._log.warning( + f"[fit r={server_round}] " + f"results={round_row['n_results']} failures={round_row['n_failures']} " + f"total_ex={round_row['total_examples']} " + f"w_norm={round_row['w_norm']:.4g} " + f"delta={round_row['delta_norm']} " + f"cos={round_row['cos_sim_prev']} " + f"nan/inf={round_row['nan_or_inf']}" + ) + + return agg + + def aggregate_evaluate(self, server_round, results, failures): + cids: List[str] = [] + num_examples_list: List[int] = [] + + for client_proxy, eval_res in results: + cid = getattr(client_proxy, "cid", "unknown") + cids.append(cid) + + ne = int(getattr(eval_res, "num_examples", 0)) + num_examples_list.append(ne) + + entry = { + "round": int(server_round), + "num_examples": ne, + "loss": float(getattr(eval_res, "loss", float("nan"))), + **dict(eval_res.metrics or {}), + } + self.client_eval_history.setdefault(cid, []).append(entry) + + agg = super().aggregate_evaluate(server_round, results, failures) + + if agg is not None: + loss_agg, metrics_agg = agg # type: ignore[misc] + round_row = { + "round": int(server_round), + "phase": "eval", + "n_results": int(len(results)), + "n_failures": int(len(failures)), + "cids": cids, + "total_examples": int(sum(num_examples_list)), + "min_examples": int(min(num_examples_list)) if num_examples_list else 0, + "max_examples": int(max(num_examples_list)) if num_examples_list else 0, + "failures": self._safe_failure_summaries(failures), + "loss": float(loss_agg) if loss_agg is not None else None, + **(dict(metrics_agg or {})), + } + self.round_history.append(round_row) + + if self._debug_print: + self._log.warning( + f"[eval r={server_round}] " + f"results={round_row['n_results']} failures={round_row['n_failures']} " + f"loss={round_row['loss']}" + ) + + return agg + + def get_debug_state(self) -> Dict[str, Any]: + return { + "round_history": self.round_history, + "client_fit_history": self.client_fit_history, + "client_eval_history": self.client_eval_history, + } + + # ---- base kwargs ---- + kwargs: Dict[str, Any] = dict( fraction_fit=self.fraction_fit, fraction_evaluate=self.fraction_evaluate, min_fit_clients=self.min_fit_clients, min_evaluate_clients=self.min_evaluate_clients, min_available_clients=self.min_available_clients, - initial_parameters=fl.common.ndarrays_to_parameters(self.initial_parameters), - evaluate_fn=self.evaluate_fn ) - def get_strategy_by_name(self): - return eval(f"fl.server.strategy.{self.name}") - - + # Initial parameters + if self.initial_parameters is not None: + kwargs["initial_parameters"] = fl.common.ndarrays_to_parameters(self.initial_parameters) + + # Optional evaluate function + if self.evaluate_fn is not None: + kwargs["evaluate_fn"] = self.evaluate_fn + + # ✅ Metrics aggregation (removes warnings and gives aggregated metrics) + if self.enable_metrics_aggregation: + kwargs["fit_metrics_aggregation_fn"] = weighted_avg + kwargs["evaluate_metrics_aggregation_fn"] = weighted_avg + + # User strategy overrides + kwargs.update(self.strategy_kwargs) + + self.strategy_object = TrackingStrategy(**kwargs) + return self.strategy_object \ No newline at end of file diff --git a/build/lib/MEDfl/LearningManager/utils.py b/build/lib/MEDfl/LearningManager/utils.py index a5bd4dd..62ff32e 100644 --- a/build/lib/MEDfl/LearningManager/utils.py +++ b/build/lib/MEDfl/LearningManager/utils.py @@ -1,6 +1,5 @@ #!/usr/bin/env python3 -import pkg_resources import torch import yaml from sklearn.metrics import * diff --git a/build/lib/MEDfl/LearningManager/xgboost/__init__.py b/build/lib/MEDfl/LearningManager/xgboost/__init__.py new file mode 100644 index 0000000..91a6300 --- /dev/null +++ b/build/lib/MEDfl/LearningManager/xgboost/__init__.py @@ -0,0 +1,15 @@ +""" +Federated XGBoost backend for MEDfl simulation mode. +""" + +from .client import XGBoostSimulationClient +from .config import XGBoostConfig +from .server import XGBoostSimulationServer +from .strategy import XGBoostSimulationStrategy + +__all__ = [ + "XGBoostSimulationClient", + "XGBoostConfig", + "XGBoostSimulationServer", + "XGBoostSimulationStrategy", +] \ No newline at end of file diff --git a/build/lib/MEDfl/LearningManager/xgboost/client.py b/build/lib/MEDfl/LearningManager/xgboost/client.py new file mode 100644 index 0000000..0975248 --- /dev/null +++ b/build/lib/MEDfl/LearningManager/xgboost/client.py @@ -0,0 +1,328 @@ +#!/usr/bin/env python3 + +import logging +import sys +from typing import Any, Dict, List, Optional + +import flwr as fl +import xgboost as xgb +from flwr.common import ( + Code, + EvaluateIns, + EvaluateRes, + FitIns, + FitRes, + GetParametersIns, + GetParametersRes, + GetPropertiesIns, + GetPropertiesRes, + Status, +) +from torch.utils.data import DataLoader + +from .data import SimulationXGBoostDataModule +from .metrics import evaluate_predictions +from .serialization import ( + booster_to_parameters, + parameters_to_booster, +) + + +def _get_logger() -> logging.Logger: + logger = logging.getLogger( + "MEDfl.simulation.xgboost.client" + ) + + if not logger.handlers: + handler = logging.StreamHandler(sys.stdout) + handler.setFormatter( + logging.Formatter( + "[%(asctime)s][%(name)s][%(levelname)s] " + "%(message)s", + datefmt="%H:%M:%S", + ) + ) + + logger.addHandler(handler) + logger.setLevel(logging.INFO) + logger.propagate = False + + return logger + + +class XGBoostSimulationClient(fl.client.Client): + def __init__( + self, + cid: str, + trainloader: DataLoader, + valloader: Optional[DataLoader] = None, + testloader: Optional[DataLoader] = None, + task: str = "binary", + xgb_params: Optional[Dict[str, Any]] = None, + local_num_boost_round: int = 10, + threshold: float = 0.5, + feature_names: Optional[List[str]] = None, + ): + self.cid = str(cid) + self.task = task + self.xgb_params = xgb_params or {} + self.local_num_boost_round = int( + local_num_boost_round + ) + self.threshold = float(threshold) + self.feature_names = feature_names + + self.log = _get_logger() + self.booster: Optional[xgb.Booster] = None + + data_module = SimulationXGBoostDataModule( + trainloader=trainloader, + valloader=valloader, + testloader=testloader, + feature_names=feature_names, + ) + + self.dtrain = data_module.dtrain + self.dval = data_module.dval + self.dtest = data_module.dtest + self.metadata = data_module.metadata + + self._validate() + + self.log.info( + "[Client %s] initialized: " + "train=%s, val=%s, test=%s, features=%s", + self.cid, + self.metadata["num_train"], + self.metadata["num_val"], + self.metadata["num_test"], + self.metadata["num_features"], + ) + + def _validate(self) -> None: + if self.task not in { + "binary", + "regression", + "multiclass", + }: + raise ValueError( + f"Unsupported task: {self.task}" + ) + + if self.dtrain is None: + raise ValueError( + f"Client {self.cid} has no training data." + ) + + if self.local_num_boost_round <= 0: + raise ValueError( + "local_num_boost_round must be greater than zero." + ) + + def get_parameters( + self, + ins: GetParametersIns, + ) -> GetParametersRes: + return GetParametersRes( + status=Status( + code=Code.OK, + message="Success", + ), + parameters=booster_to_parameters( + self.booster + ), + ) + + def fit(self, ins: FitIns) -> FitRes: + self.log.info( + "[Client %s] starting local XGBoost fit", + self.cid, + ) + + received_booster = parameters_to_booster( + ins.parameters + ) + + if received_booster is not None: + self.booster = received_booster + + evals = [] + + if self.dval is not None: + evals.append( + (self.dval, "validation") + ) + + self.booster = xgb.train( + params=self.xgb_params, + dtrain=self.dtrain, + num_boost_round=self.local_num_boost_round, + xgb_model=self.booster, + evals=evals, + verbose_eval=False, + ) + + train_predictions = self.booster.predict( + self.dtrain + ) + + train_metrics = evaluate_predictions( + task=self.task, + y_true=self.dtrain.get_label(), + y_pred_or_prob=train_predictions, + threshold=self.threshold, + num_features=int( + self.metadata["num_features"] + ), + ) + + metrics = { + "backend": "xgboost", + "task": self.task, + "cid": self.cid, + "num_train": int( + self.metadata["num_train"] + ), + } + + for key, value in train_metrics.items(): + metrics[f"train_{key}"] = value + + self.log.info( + "[Client %s] fit completed: %s", + self.cid, + metrics, + ) + + return FitRes( + status=Status( + code=Code.OK, + message="Success", + ), + parameters=booster_to_parameters( + self.booster + ), + num_examples=int( + self.metadata["num_train"] + ), + metrics=metrics, + ) + + def evaluate( + self, + ins: EvaluateIns, + ) -> EvaluateRes: + received_booster = parameters_to_booster( + ins.parameters + ) + + if received_booster is not None: + self.booster = received_booster + + evaluation_matrix = ( + self.dtest + if self.dtest is not None + else self.dval + ) + + if ( + self.booster is None + or evaluation_matrix is None + ): + return EvaluateRes( + status=Status( + code=Code.EVALUATE_NOT_IMPLEMENTED, + message=( + "No booster or evaluation data." + ), + ), + loss=0.0, + num_examples=0, + metrics={}, + ) + + predictions = self.booster.predict( + evaluation_matrix + ) + + eval_metrics = evaluate_predictions( + task=self.task, + y_true=evaluation_matrix.get_label(), + y_pred_or_prob=predictions, + threshold=self.threshold, + num_features=int( + self.metadata["num_features"] + ), + ) + + metrics = { + "backend": "xgboost", + "task": self.task, + "cid": self.cid, + } + + for key, value in eval_metrics.items(): + metrics[f"eval_{key}"] = value + + if self.task == "binary": + loss = float( + eval_metrics.get("logloss", 0.0) + ) + elif self.task == "regression": + loss = float( + eval_metrics.get("rmse", 0.0) + ) + else: + loss = float( + eval_metrics.get("mlogloss", 0.0) + ) + + self.log.info( + "[Client %s] evaluation completed: " + "loss=%s metrics=%s", + self.cid, + loss, + metrics, + ) + + return EvaluateRes( + status=Status( + code=Code.OK, + message="Success", + ), + loss=loss, + num_examples=int( + evaluation_matrix.num_row() + ), + metrics=metrics, + ) + + def get_properties( + self, + ins: GetPropertiesIns, + ) -> GetPropertiesRes: + properties = { + "backend": "xgboost", + "cid": self.cid, + "task": self.task, + "num_train": int( + self.metadata["num_train"] + ), + "num_val": int( + self.metadata["num_val"] + ), + "num_test": int( + self.metadata["num_test"] + ), + "num_features": int( + self.metadata["num_features"] + ), + } + + return GetPropertiesRes( + status=Status( + code=Code.OK, + message="Success", + ), + properties=properties, + ) \ No newline at end of file diff --git a/build/lib/MEDfl/LearningManager/xgboost/config.py b/build/lib/MEDfl/LearningManager/xgboost/config.py new file mode 100644 index 0000000..9fd6ba8 --- /dev/null +++ b/build/lib/MEDfl/LearningManager/xgboost/config.py @@ -0,0 +1,63 @@ +from dataclasses import dataclass, field +from typing import Any, Dict, Optional + + +@dataclass +class XGBoostConfig: + task: str = "binary" + mode: str = "bagging" + local_num_boost_round: int = 10 + params: Dict[str, Any] = field(default_factory=dict) + num_classes: Optional[int] = None + + def build_params(self) -> Dict[str, Any]: + if self.task == "binary": + defaults = { + "objective": "binary:logistic", + "eval_metric": "auc", + "tree_method": "hist", + "max_depth": 6, + "eta": 0.1, + "subsample": 1.0, + "colsample_bytree": 1.0, + } + + elif self.task == "regression": + defaults = { + "objective": "reg:squarederror", + "eval_metric": "rmse", + "tree_method": "hist", + "max_depth": 6, + "eta": 0.1, + "subsample": 1.0, + "colsample_bytree": 1.0, + } + + elif self.task == "multiclass": + if self.num_classes is None or self.num_classes < 2: + raise ValueError( + "num_classes must be provided and >= 2 " + "for multiclass XGBoost." + ) + + defaults = { + "objective": "multi:softprob", + "num_class": self.num_classes, + "eval_metric": "mlogloss", + "tree_method": "hist", + "max_depth": 6, + "eta": 0.1, + "subsample": 1.0, + "colsample_bytree": 1.0, + } + + else: + raise ValueError( + f"Unsupported XGBoost task: {self.task}. " + "Expected binary, regression, or multiclass." + ) + + merged = defaults.copy() + merged.update(self.params or {}) + + return merged \ No newline at end of file diff --git a/build/lib/MEDfl/LearningManager/xgboost/data.py b/build/lib/MEDfl/LearningManager/xgboost/data.py new file mode 100644 index 0000000..9333c75 --- /dev/null +++ b/build/lib/MEDfl/LearningManager/xgboost/data.py @@ -0,0 +1,122 @@ +from typing import List, Optional, Tuple + +import numpy as np +import torch +import xgboost as xgb +from torch.utils.data import DataLoader + + +def dataloader_to_numpy( + loader: DataLoader, +) -> Tuple[np.ndarray, np.ndarray]: + if loader is None: + raise ValueError("DataLoader cannot be None.") + + x_batches = [] + y_batches = [] + + for x_batch, y_batch in loader: + if isinstance(x_batch, torch.Tensor): + x_batch = x_batch.detach().cpu().numpy() + + if isinstance(y_batch, torch.Tensor): + y_batch = y_batch.detach().cpu().numpy() + + x_batch = np.asarray(x_batch) + y_batch = np.asarray(y_batch).reshape(-1) + + x_batches.append(x_batch) + y_batches.append(y_batch) + + if not x_batches: + raise ValueError( + "Cannot convert an empty DataLoader to DMatrix." + ) + + x = np.concatenate(x_batches, axis=0) + y = np.concatenate(y_batches, axis=0) + + if x.ndim != 2: + raise ValueError( + "XGBoost tabular simulation expects a two-dimensional " + f"feature matrix. Received shape={x.shape}." + ) + + return x, y + + +def dataloader_to_dmatrix( + loader: Optional[DataLoader], + feature_names: Optional[List[str]] = None, +) -> Optional[xgb.DMatrix]: + if loader is None: + return None + + try: + if len(loader.dataset) == 0: + return None + except TypeError: + pass + + x, y = dataloader_to_numpy(loader) + + if feature_names is not None: + if len(feature_names) != x.shape[1]: + raise ValueError( + "feature_names length does not match the number " + f"of features: {len(feature_names)} != {x.shape[1]}" + ) + + return xgb.DMatrix( + data=x, + label=y, + feature_names=feature_names, + ) + + +class SimulationXGBoostDataModule: + def __init__( + self, + trainloader: DataLoader, + valloader: Optional[DataLoader] = None, + testloader: Optional[DataLoader] = None, + feature_names: Optional[List[str]] = None, + ): + self.dtrain = dataloader_to_dmatrix( + trainloader, + feature_names, + ) + + self.dval = dataloader_to_dmatrix( + valloader, + feature_names, + ) + + self.dtest = dataloader_to_dmatrix( + testloader, + feature_names, + ) + + self.metadata = { + "num_train": ( + self.dtrain.num_row() + if self.dtrain is not None + else 0 + ), + "num_val": ( + self.dval.num_row() + if self.dval is not None + else 0 + ), + "num_test": ( + self.dtest.num_row() + if self.dtest is not None + else 0 + ), + "num_features": ( + self.dtrain.num_col() + if self.dtrain is not None + else 0 + ), + "feature_names": feature_names or [], + } \ No newline at end of file diff --git a/build/lib/MEDfl/LearningManager/xgboost/metrics.py b/build/lib/MEDfl/LearningManager/xgboost/metrics.py new file mode 100644 index 0000000..0df8827 --- /dev/null +++ b/build/lib/MEDfl/LearningManager/xgboost/metrics.py @@ -0,0 +1,192 @@ +import math +from typing import Any, Dict + +import numpy as np +from sklearn.metrics import ( + accuracy_score, + f1_score, + log_loss, + mean_absolute_error, + mean_squared_error, + r2_score, + roc_auc_score, +) + + +def _safe_float(value: Any) -> float: + try: + value = float(value) + + if math.isnan(value) or math.isinf(value): + return float("nan") + + return value + except Exception: + return float("nan") + + +def binary_metrics( + y_true, + y_prob, + threshold: float = 0.5, +) -> Dict[str, float]: + y_true = np.asarray(y_true) + y_prob = np.asarray(y_prob) + y_pred = (y_prob >= threshold).astype(int) + + metrics = { + "accuracy": _safe_float( + accuracy_score(y_true, y_pred) + ), + } + + try: + metrics["auc"] = _safe_float( + roc_auc_score(y_true, y_prob) + ) + except Exception: + metrics["auc"] = float("nan") + + try: + metrics["logloss"] = _safe_float( + log_loss(y_true, y_prob) + ) + except Exception: + metrics["logloss"] = float("nan") + + return metrics + + +def regression_metrics( + y_true, + y_pred, + num_features: int = 0, +) -> Dict[str, float]: + y_true = np.asarray(y_true) + y_pred = np.asarray(y_pred) + + rmse = mean_squared_error( + y_true, + y_pred, + squared=False, + ) + mae = mean_absolute_error(y_true, y_pred) + r2 = r2_score(y_true, y_pred) + + n = len(y_true) + p = int(num_features) + + if n > p + 1: + adjusted_r2 = ( + 1.0 + - (1.0 - r2) + * ((n - 1.0) / (n - p - 1.0)) + ) + else: + adjusted_r2 = float("nan") + + return { + "rmse": _safe_float(rmse), + "mae": _safe_float(mae), + "r2": _safe_float(r2), + "adjusted_r2": _safe_float(adjusted_r2), + } + + +def multiclass_metrics( + y_true, + y_prob, +) -> Dict[str, float]: + y_true = np.asarray(y_true) + y_prob = np.asarray(y_prob) + y_pred = np.argmax(y_prob, axis=1) + + metrics = { + "accuracy": _safe_float( + accuracy_score(y_true, y_pred) + ), + } + + try: + metrics["macro_f1"] = _safe_float( + f1_score( + y_true, + y_pred, + average="macro", + ) + ) + except Exception: + metrics["macro_f1"] = float("nan") + + try: + metrics["mlogloss"] = _safe_float( + log_loss(y_true, y_prob) + ) + except Exception: + metrics["mlogloss"] = float("nan") + + return metrics + + +def evaluate_predictions( + task: str, + y_true, + y_pred_or_prob, + threshold: float = 0.5, + num_features: int = 0, +) -> Dict[str, float]: + if task == "binary": + return binary_metrics( + y_true, + y_pred_or_prob, + threshold=threshold, + ) + + if task == "regression": + return regression_metrics( + y_true, + y_pred_or_prob, + num_features=num_features, + ) + + if task == "multiclass": + return multiclass_metrics( + y_true, + y_pred_or_prob, + ) + + raise ValueError(f"Unsupported task: {task}") + + +def aggregate_weighted_metrics(results): + aggregated = {} + + keys = set() + for _, metrics in results: + keys.update(metrics.keys()) + + for key in keys: + weighted_sum = 0.0 + contributing_examples = 0 + + for num_examples, metrics in results: + if key not in metrics: + continue + + try: + value = float(metrics[key]) + except Exception: + continue + + if math.isnan(value) or math.isinf(value): + continue + + weighted_sum += num_examples * value + contributing_examples += num_examples + + if contributing_examples > 0: + aggregated[key] = ( + weighted_sum / contributing_examples + ) + + return aggregated \ No newline at end of file diff --git a/build/lib/MEDfl/LearningManager/xgboost/serialization.py b/build/lib/MEDfl/LearningManager/xgboost/serialization.py new file mode 100644 index 0000000..58bbacf --- /dev/null +++ b/build/lib/MEDfl/LearningManager/xgboost/serialization.py @@ -0,0 +1,52 @@ +from typing import Optional + +import xgboost as xgb +from flwr.common import Parameters + + +XGBOOST_TENSOR_TYPE = "xgboost.ubj" + + +def booster_to_parameters( + booster: Optional[xgb.Booster], +) -> Parameters: + if booster is None: + return Parameters( + tensors=[], + tensor_type=XGBOOST_TENSOR_TYPE, + ) + + raw = booster.save_raw(raw_format="ubj") + + return Parameters( + tensors=[bytes(raw)], + tensor_type=XGBOOST_TENSOR_TYPE, + ) + + +def parameters_to_booster( + parameters: Optional[Parameters], +) -> Optional[xgb.Booster]: + if parameters is None: + return None + + if not parameters.tensors: + return None + + booster = xgb.Booster() + booster.load_model(bytearray(parameters.tensors[0])) + + return booster + + +def save_booster( + booster: xgb.Booster, + path: str, +) -> None: + booster.save_model(path) + + +def load_booster(path: str) -> xgb.Booster: + booster = xgb.Booster() + booster.load_model(path) + return booster \ No newline at end of file diff --git a/build/lib/MEDfl/LearningManager/xgboost/server.py b/build/lib/MEDfl/LearningManager/xgboost/server.py new file mode 100644 index 0000000..93e4d1e --- /dev/null +++ b/build/lib/MEDfl/LearningManager/xgboost/server.py @@ -0,0 +1,746 @@ +#!/usr/bin/env python3 + +import logging +from typing import Any, Dict, List, Optional + +import flwr as fl + +from MEDfl.LearningManager.federated_dataset import FederatedDataset +from MEDfl.LearningManager.server import setup_fl_logging + +from .client import XGBoostSimulationClient +from .strategy import XGBoostSimulationStrategy + +from MEDfl.LearningManager.shap import ( + SHAPConfig, + LocalXGBoostSHAPExplainer, + aggregate_federated_shap, +) + + +class XGBoostSimulationServer: + """ + Flower/Ray simulation server for federated XGBoost. + + Important: + Ray must serialize client_fn and everything captured by it. + Therefore client_fn must not capture the server instance, a SQLAlchemy + connection, a DatabaseManager, or another non-serializable object. + """ + + backend = "xgboost" + + def __init__( + self, + strategy: XGBoostSimulationStrategy, + num_rounds: int, + fed_dataset: FederatedDataset, + task: str = "binary", + xgb_params: Optional[Dict[str, Any]] = None, + local_num_boost_round: int = 10, + threshold: float = 0.5, + feature_names: Optional[List[str]] = None, + client_resources: Optional[Dict[str, float]] = None, + shap_config: Optional[SHAPConfig] = None, + ): + self.strategy = strategy + self.num_rounds = int(num_rounds) + self.fed_dataset = fed_dataset + + self.task = str(task) + self.xgb_params = dict(xgb_params or {}) + self.local_num_boost_round = int( + local_num_boost_round + ) + self.threshold = float(threshold) + + self.feature_names = ( + list(feature_names) + if feature_names is not None + else getattr( + fed_dataset, + "feature_names", + None, + ) + ) + + self.client_resources = ( + dict(client_resources) + if client_resources is not None + else { + "num_cpus": 1.0, + "num_gpus": 0.0, + } + ) + + self.shap_config = ( + shap_config + if shap_config is not None + else SHAPConfig(enabled=False) + ) + + self.num_clients = len( + self.fed_dataset.trainloaders + ) + + self.history = None + self.shap_results = None + + self._validate() + + self.strategy.strategy_object.min_available_clients = ( + self.num_clients + ) + + def _validate(self) -> None: + if self.strategy.strategy_object is None: + raise ValueError( + "strategy.create_strategy() must be called before " + "creating XGBoostSimulationServer." + ) + + if self.num_rounds <= 0: + raise ValueError("num_rounds must be greater than zero.") + + if self.num_clients <= 0: + raise ValueError( + "The federated dataset does not contain training clients." + ) + + if len(self.fed_dataset.valloaders) < self.num_clients: + raise ValueError( + "A validation-loader entry is required for every client." + ) + + if self.task not in { + "binary", + "regression", + "multiclass", + }: + raise ValueError( + f"Unsupported task '{self.task}'." + ) + + @staticmethod + def _clean_loader(loader): + """ + Return None for an unavailable or empty loader. + """ + + if loader is None: + return None + + try: + if len(loader.dataset) == 0: + return None + except Exception: + pass + + return loader + + def _detach_database_state(self) -> None: + """ + Remove live database objects before Ray serializes simulation data. + + The loaders are already in memory, so the XGBoost simulation no longer + requires the database connection. + """ + + if hasattr(self.fed_dataset, "eng"): + try: + connection = self.fed_dataset.eng + + if connection is not None: + try: + connection.close() + except Exception: + pass + finally: + self.fed_dataset.eng = None + + # Defensive cleanup in case these attributes are added later. + for attribute_name in ( + "engine", + "connection", + "db_manager", + "database_manager", + ): + if hasattr(self.fed_dataset, attribute_name): + setattr( + self.fed_dataset, + attribute_name, + None, + ) + + @staticmethod + def _get_loader_input_size(loader) -> int: + if loader is None: + raise ValueError( + "Cannot determine the input size from a None loader" + ) + + for batch in loader: + if not isinstance(batch, (tuple, list)): + raise TypeError( + "Expected each DataLoader batch to be a " + "tuple or list" + ) + + if len(batch) == 0: + continue + + features = batch[0] + + if hasattr(features, "shape"): + shape = features.shape + else: + features = np.asarray(features) + shape = features.shape + + if len(shape) != 2: + raise ValueError( + "XGBoost SHAP expects tabular input with shape " + "(samples, features). " + f"Received shape {shape}." + ) + + return int(shape[1]) + + raise ValueError( + "Cannot determine input size from an empty DataLoader" + ) + + def _get_shap_loaders(self): + """ + Return the selected local partition for every client. + """ + + data_split = self.shap_config.data_split + + if data_split == "train": + source_loaders = self.fed_dataset.trainloaders + + elif data_split == "validation": + source_loaders = self.fed_dataset.valloaders + + elif data_split == "test": + source_loaders = self.fed_dataset.testloaders + + else: + raise ValueError( + f"Unsupported SHAP data split: {data_split}" + ) + + selected_loaders = [] + + for client_index in range(self.num_clients): + loader = None + + if client_index < len(source_loaders): + loader = self._clean_loader( + source_loaders[client_index] + ) + + selected_loaders.append(loader) + + return selected_loaders + + def _resolve_shap_feature_names( + self, + input_size: int, + ) -> List[str]: + """ + Resolve feature names in this order: + + 1. SHAPConfig.feature_names + 2. XGBoost server feature_names + 3. FederatedDataset feature_names + 4. Automatically generated feature_0, feature_1, ... + """ + + if self.shap_config.feature_names is not None: + feature_names = list( + self.shap_config.feature_names + ) + + elif self.feature_names is not None: + feature_names = list( + self.feature_names + ) + + else: + dataset_feature_names = getattr( + self.fed_dataset, + "feature_names", + None, + ) + + if dataset_feature_names is not None: + feature_names = list( + dataset_feature_names + ) + else: + feature_names = ( + self.shap_config.resolved_feature_names( + input_size + ) + ) + + if len(feature_names) != input_size: + raise ValueError( + "The number of XGBoost SHAP feature names must " + "match the input size. " + f"Expected {input_size}, " + f"received {len(feature_names)}." + ) + + return feature_names + + def _run_federated_shap(self): + """ + Run the optional post-training federated SHAP phase. + + The same final global XGBoost booster is explained on each + participating client's selected local data partition. + """ + + if not self.shap_config.enabled: + print( + "[XGBoost/SHAP] SHAP is disabled.", + flush=True, + ) + + self.shap_results = None + return None + + final_booster = self.strategy.current_booster + + if final_booster is None: + raise RuntimeError( + "Cannot calculate XGBoost SHAP because the " + "federated strategy does not contain a final booster." + ) + + selected_loaders = self._get_shap_loaders() + + first_available_loader = None + + for loader in selected_loaders: + if loader is not None: + first_available_loader = loader + break + + if first_available_loader is None: + raise ValueError( + "No client has data for the selected SHAP split: " + f"{self.shap_config.data_split}" + ) + + input_size = self._get_loader_input_size( + first_available_loader + ) + + feature_names = self._resolve_shap_feature_names( + input_size + ) + + self.shap_config.validate_backend( + input_size=input_size, + backend="xgboost", + task=self.task, + ) + + local_explainer = LocalXGBoostSHAPExplainer( + background_size=( + self.shap_config.background_size + ), + explanation_size=( + self.shap_config.explanation_size + ), + random_seed=( + self.shap_config.random_seed + ), + clipping_value=( + self.shap_config.clipping_value + ), + model_output=( + self.shap_config.model_output + ), + ) + + client_results = [] + skipped_clients = [] + + print( + "[XGBoost/SHAP] Starting federated SHAP:", + { + "explainer": self.shap_config.explainer, + "data_split": self.shap_config.data_split, + "background_size": ( + self.shap_config.background_size + ), + "explanation_size": ( + self.shap_config.explanation_size + ), + "minimum_samples": ( + self.shap_config.minimum_samples + ), + "model_output": ( + self.shap_config.model_output + ), + "num_clients": self.num_clients, + }, + flush=True, + ) + + for client_index, loader in enumerate( + selected_loaders + ): + client_id = str(client_index) + + if loader is None: + skipped_clients.append( + { + "client_id": client_id, + "reason": ( + "Selected DataLoader is unavailable or empty" + ), + } + ) + + print( + f"[XGBoost/SHAP] Skipping client " + f"{client_id}: no data.", + flush=True, + ) + + continue + + try: + available_samples = len(loader.dataset) + except Exception: + available_samples = None + + if ( + available_samples is not None + and available_samples + < self.shap_config.minimum_samples + ): + skipped_clients.append( + { + "client_id": client_id, + "reason": ( + "Insufficient samples" + ), + "available_samples": int( + available_samples + ), + "minimum_samples": int( + self.shap_config.minimum_samples + ), + } + ) + + print( + f"[XGBoost/SHAP] Skipping client " + f"{client_id}: available={available_samples}, " + f"minimum={self.shap_config.minimum_samples}.", + flush=True, + ) + + continue + + print( + f"[XGBoost/SHAP] Calculating SHAP for " + f"client {client_id}.", + flush=True, + ) + + try: + local_result = local_explainer.calculate( + booster=final_booster, + loader=loader, + ) + + if ( + local_result["sample_count"] + < self.shap_config.minimum_samples + ): + skipped_clients.append( + { + "client_id": client_id, + "reason": ( + "Explained sample count is below " + "minimum_samples" + ), + "explained_samples": int( + local_result[ + "sample_count" + ] + ), + "minimum_samples": int( + self.shap_config.minimum_samples + ), + } + ) + + continue + + local_result["client_id"] = client_id + local_result["feature_names"] = list( + feature_names + ) + + client_results.append(local_result) + + print( + f"[XGBoost/SHAP] Client {client_id} " + f"completed: explained_samples=" + f"{local_result['sample_count']}.", + flush=True, + ) + + except Exception as exc: + skipped_clients.append( + { + "client_id": client_id, + "reason": str(exc), + } + ) + + print( + f"[XGBoost/SHAP] Client {client_id} failed: " + f"{exc}", + flush=True, + ) + + if not client_results: + self.shap_results = { + "status": "failed", + "backend": "xgboost", + "explainer": "tree", + "data_split": ( + self.shap_config.data_split + ), + "participating_clients": 0, + "explained_samples": 0, + "feature_importance": [], + "skipped_clients": skipped_clients, + "error": ( + "No XGBoost client produced a valid " + "local SHAP result" + ), + } + + print( + "[XGBoost/SHAP] Federated SHAP failed: " + "no valid client results.", + flush=True, + ) + + return self.shap_results + + federated_result = aggregate_federated_shap( + client_results=client_results, + include_client_results=( + self.shap_config.include_client_results + ), + ) + + federated_result.update( + { + "backend": "xgboost", + "task": self.task, + "explainer": "tree", + "data_split": ( + self.shap_config.data_split + ), + "background_size": int( + self.shap_config.background_size + ), + "requested_explanation_size": int( + self.shap_config.explanation_size + ), + "minimum_samples": int( + self.shap_config.minimum_samples + ), + "random_seed": int( + self.shap_config.random_seed + ), + "model_output": ( + self.shap_config.model_output + ), + "feature_names": list( + feature_names + ), + "skipped_clients": skipped_clients, + } + ) + + self.shap_results = federated_result + + print( + "[XGBoost/SHAP] Federated SHAP completed:", + { + "participating_clients": ( + federated_result[ + "participating_clients" + ] + ), + "explained_samples": ( + federated_result[ + "explained_samples" + ] + ), + "skipped_clients": len( + skipped_clients + ), + }, + flush=True, + ) + + print( + "[XGBoost/SHAP] Results:", + federated_result, + flush=True, + ) + + return federated_result + + def run(self): + setup_fl_logging() + + self._detach_database_state() + + # ----------------------------------------------------------- + # Extract only values that Ray can serialize. + # + # Do not pass self.client_fn because that bound method captures + # the complete server object, including fed_dataset and any DB state. + # ----------------------------------------------------------- + + trainloaders = list( + self.fed_dataset.trainloaders + ) + valloaders = list( + self.fed_dataset.valloaders + ) + testloaders = list( + self.fed_dataset.testloaders + ) + + task = self.task + xgb_params = dict(self.xgb_params) + local_num_boost_round = int( + self.local_num_boost_round + ) + threshold = float(self.threshold) + + feature_names = ( + list(self.feature_names) + if self.feature_names is not None + else None + ) + + num_clients = int(self.num_clients) + num_rounds = int(self.num_rounds) + + def client_fn(cid: str): + """ + Ray-serializable client factory. + + This function captures only DataLoaders, primitive values, + dictionaries, and lists. It does not capture the server or + SQLAlchemy objects. + """ + + client_index = int(cid) + + if client_index < 0 or client_index >= num_clients: + raise IndexError( + f"Invalid simulated client ID {cid}. " + f"Expected an index in [0, {num_clients - 1}]." + ) + + trainloader = trainloaders[client_index] + + valloader = None + if client_index < len(valloaders): + valloader = XGBoostSimulationServer._clean_loader( + valloaders[client_index] + ) + + testloader = None + if client_index < len(testloaders): + testloader = XGBoostSimulationServer._clean_loader( + testloaders[client_index] + ) + + client = XGBoostSimulationClient( + cid=str(cid), + trainloader=trainloader, + valloader=valloader, + testloader=testloader, + task=task, + xgb_params=xgb_params, + local_num_boost_round=local_num_boost_round, + threshold=threshold, + feature_names=feature_names, + ) + + # XGBoostSimulationClient already subclasses fl.client.Client. + return client + + print( + "[XGBoostSimulationServer] Starting simulation:", + { + "num_clients": num_clients, + "num_rounds": num_rounds, + "task": task, + "local_num_boost_round": local_num_boost_round, + "client_resources": self.client_resources, + }, + flush=True, + ) + + self.history = fl.simulation.start_simulation( + client_fn=client_fn, + num_clients=num_clients, + config=fl.server.ServerConfig( + num_rounds=num_rounds, + ), + strategy=self.strategy.strategy_object, + client_resources=self.client_resources, + ray_init_args={ + "include_dashboard": False, + "log_to_driver": True, + "logging_level": logging.INFO, + }, + ) + + # ----------------------------------------------------------- + # Optional post-training federated SHAP + # ----------------------------------------------------------- + + if self.shap_config.enabled: + try: + self._run_federated_shap() + + except Exception as exc: + self.shap_results = { + "status": "failed", + "backend": "xgboost", + "explainer": "tree", + "data_split": ( + self.shap_config.data_split + ), + "error": str(exc), + } + + print( + "[XGBoost/SHAP] Federated SHAP failed:", + exc, + flush=True, + ) + + return self.history \ No newline at end of file diff --git a/build/lib/MEDfl/LearningManager/xgboost/strategy.py b/build/lib/MEDfl/LearningManager/xgboost/strategy.py new file mode 100644 index 0000000..c5e7cb5 --- /dev/null +++ b/build/lib/MEDfl/LearningManager/xgboost/strategy.py @@ -0,0 +1,576 @@ +import json +import os +import tempfile +from typing import Any, Dict, Optional + +import flwr as fl +import xgboost as xgb +from flwr.common import Parameters + +from .config import XGBoostConfig +from .metrics import aggregate_weighted_metrics +from .serialization import ( + XGBOOST_TENSOR_TYPE, + booster_to_parameters, + parameters_to_booster, + save_booster, +) + +class XGBoostSimulationStrategy: + """ + Strategy wrapper for MEDfl real-world federated XGBoost. + + This class follows the same high-level contract as the existing MEDfl + Strategy wrapper: + - it exposes create_strategy(); + - it stores the Flower strategy in self.strategy_object; + - it can be passed directly into FederatedServer. + + First supported mode: + - bagging-style federated XGBoost using Flower strategy hooks. + + Notes + ----- + This implementation is conservative and non-invasive: + - It does not modify the existing PyTorch FedAvg/FedAdam/FedYogi path. + - It uses Flower Parameters containing serialized XGBoost boosters. + - It saves XGBoost models as .ubj files instead of .npz files. + """ + + def __init__( + self, + mode="bagging", + fraction_fit=1.0, + fraction_evaluate=1.0, + min_fit_clients=2, + min_evaluate_clients=2, + min_available_clients=2, + task="binary", + num_classes=None, + xgb_params=None, + local_num_boost_round=10, + threshold=0.5, + savingPath="", + save_on_rounds=1, + total_rounds=3, + ): + self.name = "FedXgbBagging" + self.mode = mode + + if self.mode != "bagging": + raise ValueError( + "Only mode='bagging' is implemented in this first integration. " + "Cyclic mode should be added after bagging is stable." + ) + + self.fraction_fit = fraction_fit + self.fraction_evaluate = fraction_evaluate + self.min_fit_clients = min_fit_clients + self.min_evaluate_clients = min_evaluate_clients + self.min_available_clients = min_available_clients + + self.task = task + self.num_classes = num_classes + self.threshold = threshold + self.savingPath = savingPath + self.save_on_rounds = save_on_rounds + + config = XGBoostConfig( + task=task, + mode=mode, + local_num_boost_round=local_num_boost_round, + params=xgb_params or {}, + num_classes=num_classes, + ) + + self.xgb_params = config.build_params() + self.local_num_boost_round = local_num_boost_round + + self.strategy_object = None + self.current_booster = None + self.current_num_trees = 0 + + self.round_history = [] + self.client_fit_history = {} + self.client_evaluation_history = {} + + + self.total_rounds = total_rounds + + self.strategy_object = None + self.current_booster: Optional[xgb.Booster] = None + self.current_num_trees: int = 0 + + def create_strategy(self) -> None: + def fit_config_fn(server_round: int): + return { + "backend": "xgboost", + "server_round": int(server_round), + "mode": self.mode, + "task": self.task, + "threshold": float(self.threshold), + "local_num_boost_round": int( + self.local_num_boost_round + ), + "xgb_params": json.dumps( + self.xgb_params + ), + } + + def evaluate_config_fn(server_round: int): + return { + "backend": "xgboost", + "server_round": int(server_round), + "task": self.task, + "threshold": float(self.threshold), + } + + strategy = fl.server.strategy.FedAvg( + fraction_fit=self.fraction_fit, + fraction_evaluate=self.fraction_evaluate, + min_fit_clients=self.min_fit_clients, + min_evaluate_clients=( + self.min_evaluate_clients + ), + min_available_clients=( + self.min_available_clients + ), + on_fit_config_fn=fit_config_fn, + on_evaluate_config_fn=( + evaluate_config_fn + ), + fit_metrics_aggregation_fn=( + aggregate_weighted_metrics + ), + evaluate_metrics_aggregation_fn=( + aggregate_weighted_metrics + ), + initial_parameters=Parameters( + tensors=[], + tensor_type=XGBOOST_TENSOR_TYPE, + ), + ) + + self._wrap_aggregate_fit(strategy) + self._wrap_aggregate_evaluate(strategy) + + self.strategy_object = strategy + def _wrap_aggregate_fit(self, strategy) -> None: + """ + Wrap aggregate_fit for XGBoost. + + IMPORTANT: + XGBoost boosters are NOT NumPy tensors. + Therefore we MUST NOT call FedAvg tensor aggregation. + + Current strategy: + - select the booster from the client with the largest dataset + - aggregate only metrics manually + - save the selected booster + """ + + def wrapped_aggregate_fit(server_round, results, failures): + print(f"\n[Server/XGBoost] 🔄 Round {server_round} - Client Training Metrics:") + + for client, fit_res in results: + print( + f" XGB CTM Round {server_round} " + f"Client:{client.cid}: {fit_res.metrics}" + ) + + if failures: + print(f"[Server/XGBoost] ⚠️ Round {server_round} failures: {failures}") + + if not results: + print("[Server/XGBoost] No fit results received.") + return None, {} + + # Aggregate metrics manually, same as before. + metric_results = [] + + for _, fit_res in results: + metric_results.append( + ( + fit_res.num_examples, + fit_res.metrics, + ) + ) + + metrics = aggregate_weighted_metrics(metric_results) + + previous_num_trees = int(self.current_num_trees) + new_local_boosters = [] + + print( + f"[Server/XGBoost] Previous global tree count before round " + f"{server_round}: {previous_num_trees}" + ) + + for client, fit_res in results: + client_booster = parameters_to_booster(fit_res.parameters) + + if client_booster is None: + print( + f"[Server/XGBoost] Client {client.cid} returned no booster. Skipping." + ) + continue + + client_total_trees = self._num_boosted_rounds(client_booster) + + print( + f"[Server/XGBoost] Client {client.cid} returned booster with " + f"{client_total_trees} trees." + ) + + # In round 1, previous_num_trees is 0, so this keeps all local trees. + # In later rounds, this keeps only the new trees added by the client. + new_trees_booster = self._slice_new_trees( + client_booster=client_booster, + previous_num_trees=previous_num_trees, + ) + + if new_trees_booster is None: + continue + + new_tree_count = self._num_boosted_rounds(new_trees_booster) + + print( + f"[Server/XGBoost] Client {client.cid} contributed " + f"{new_tree_count} new trees." + ) + + new_local_boosters.append(new_trees_booster) + + if not new_local_boosters: + print( + "[Server/XGBoost] No new local trees received. " + "Keeping previous global booster." + ) + + if self.current_booster is None: + return None, metrics + + agg_parameters = booster_to_parameters(self.current_booster) + return agg_parameters, metrics + + boosters_to_combine = [] + + if self.current_booster is not None: + boosters_to_combine.append(self.current_booster) + + boosters_to_combine.extend(new_local_boosters) + + combined_booster = self._combine_boosters_by_json(boosters_to_combine) + + if combined_booster is None: + print("[Server/XGBoost] Failed to combine boosters.") + return None, metrics + + self.current_booster = combined_booster + self.current_num_trees = self._num_boosted_rounds(combined_booster) + + print( + f"[Server/XGBoost] ✅ Round {server_round} - Flower-style bagging aggregation done." + ) + print( + f"[Server/XGBoost] Global booster now has {self.current_num_trees} trees." + ) + print(f"[Server/XGBoost] Aggregated metrics: {metrics}\n") + + # Optional: print importance directly, no saving required. + if server_round == self.total_rounds: + self._print_feature_importance(combined_booster, server_round) + + agg_parameters = booster_to_parameters(combined_booster) + + self.round_history.append( + { + "round": int(server_round), + "num_fit_clients": len(results), + "num_failures": len(failures), + "global_num_trees": int( + self.current_num_trees + ), + **metrics, + } +) + + # Optional: keep your existing save behavior if savingPath is provided. + self._save_if_needed(server_round, agg_parameters) + + return agg_parameters, metrics + strategy.aggregate_fit = wrapped_aggregate_fit + + + def _wrap_aggregate_evaluate(self, strategy) -> None: + """ + Wrap aggregate_evaluate for logging. + """ + + original_aggregate_evaluate = strategy.aggregate_evaluate + + def wrapped_aggregate_evaluate(server_round, results, failures): + print(f"\n[Server/XGBoost] 📊 Round {server_round} - Client Evaluation Metrics:") + + for client, eval_res in results: + print( + f" XGB CEM Round {server_round} " + f"Client:{client.cid}: {eval_res.metrics}" + ) + + loss, metrics = original_aggregate_evaluate(server_round, results, failures) + + print(f"[Server/XGBoost] ✅ Round {server_round} - Aggregated Evaluation:") + print(f" Loss: {loss}, Metrics: {metrics}\n") + + return loss, metrics + + strategy.aggregate_evaluate = wrapped_aggregate_evaluate + + def _save_if_needed(self, server_round: int, parameters: Parameters) -> None: + """ + Save XGBoost model as .ubj. + """ + + if not self.savingPath: + return + + should_save = ( + server_round % self.save_on_rounds == 0 + or server_round == self.total_rounds + ) + + if not should_save: + return + + os.makedirs(self.savingPath, exist_ok=True) + + booster = parameters_to_booster(parameters) + + if booster is None: + print("[Server/XGBoost] No booster to save.") + return + + filename = ( + f"round_{server_round}_final_model.ubj" + if server_round == self.total_rounds + else f"round_{server_round}_model.ubj" + ) + + filepath = os.path.join(self.savingPath, filename) + + save_booster(booster, filepath) + + print(f"[Server/XGBoost] 💾 Saved booster to: {filepath}") + def _num_boosted_rounds(self, booster: xgb.Booster) -> int: + """ + Return the number of boosted rounds/trees in the booster. + """ + try: + return int(booster.num_boosted_rounds()) + except Exception: + try: + return len(booster.get_dump()) + except Exception: + return 0 + + + def _slice_new_trees( + self, + client_booster: xgb.Booster, + previous_num_trees: int, + ) -> Optional[xgb.Booster]: + """ + Extract only the trees added by the client in this round. + + Client booster = previous global booster + new local trees. + We keep only new local trees. + """ + + total_trees = self._num_boosted_rounds(client_booster) + + if total_trees <= previous_num_trees: + print( + "[Server/XGBoost] Client booster has no new trees. " + f"total_trees={total_trees}, previous_num_trees={previous_num_trees}" + ) + return None + + try: + return client_booster[previous_num_trees:total_trees] + except Exception as exc: + print(f"[Server/XGBoost] Failed to slice new trees: {exc}") + return None + + + def _booster_to_raw_bytes(self, booster: xgb.Booster) -> bytes: + """ + Serialize booster to raw UBJ bytes. + """ + with tempfile.NamedTemporaryFile(suffix=".ubj", delete=False) as tmp: + tmp_path = tmp.name + + try: + booster.save_model(tmp_path) + with open(tmp_path, "rb") as f: + return f.read() + finally: + try: + os.remove(tmp_path) + except Exception: + pass + + + def _booster_from_raw_bytes(self, raw: bytes) -> Optional[xgb.Booster]: + """ + Deserialize booster from raw UBJ bytes. + """ + if not raw: + return None + + with tempfile.NamedTemporaryFile(suffix=".ubj", delete=False) as tmp: + tmp_path = tmp.name + tmp.write(raw) + + try: + booster = xgb.Booster() + booster.load_model(tmp_path) + return booster + except Exception as exc: + print(f"[Server/XGBoost] Failed to load booster from bytes: {exc}") + return None + finally: + try: + os.remove(tmp_path) + except Exception: + pass + + def _combine_boosters_by_json( + self, + boosters: list[xgb.Booster], + ) -> Optional[xgb.Booster]: + """ + Combine multiple XGBoost boosters by concatenating their trees. + + Fixes: + - updates trees + - updates tree_info + - updates gbtree_model_param["num_trees"] + - updates iteration_indptr + + This is the missing part that caused: + Check failed: model.iteration_indptr.back() == model.param.num_trees + """ + + boosters = [b for b in boosters if b is not None] + + if not boosters: + return None + + if len(boosters) == 1: + return boosters[0] + + model_jsons = [] + + for booster in boosters: + raw_json = booster.save_raw(raw_format="json") + if isinstance(raw_json, bytes): + raw_json = raw_json.decode("utf-8") + model_jsons.append(json.loads(raw_json)) + + base = model_jsons[0] + + base_model = base["learner"]["gradient_booster"]["model"] + base_trees = base_model["trees"] + base_tree_info = base_model["tree_info"] + + # Start from the base model iteration structure + base_iteration_indptr = base_model.get("iteration_indptr", None) + + if base_iteration_indptr is None: + # Safe fallback for binary/regression: one tree per iteration + combined_iteration_indptr = list(range(0, len(base_trees) + 1)) + else: + combined_iteration_indptr = list(base_iteration_indptr) + + next_tree_id = len(base_trees) + + for extra in model_jsons[1:]: + extra_model = extra["learner"]["gradient_booster"]["model"] + + extra_trees = extra_model["trees"] + extra_tree_info = extra_model["tree_info"] + + # Append trees and reassign IDs + for tree in extra_trees: + tree["id"] = next_tree_id + base_trees.append(tree) + next_tree_id += 1 + + # Append tree_info + base_tree_info.extend(extra_tree_info) + + # Append iteration_indptr correctly + extra_iteration_indptr = extra_model.get("iteration_indptr", None) + + if extra_iteration_indptr is None: + # Fallback: one tree per boosting iteration + extra_iteration_indptr = list(range(0, len(extra_trees) + 1)) + + offset = combined_iteration_indptr[-1] + + # Skip the first 0 from the extra booster + for ptr in extra_iteration_indptr[1:]: + combined_iteration_indptr.append(offset + ptr) + + # Write back merged model fields + base_model["trees"] = base_trees + base_model["tree_info"] = base_tree_info + base_model["iteration_indptr"] = combined_iteration_indptr + + # Update num_trees + if "gbtree_model_param" in base_model: + base_model["gbtree_model_param"]["num_trees"] = str(len(base_trees)) + + # Save/load to validate the merged model + with tempfile.NamedTemporaryFile( + mode="w", + suffix=".json", + delete=False, + encoding="utf-8", + ) as tmp: + tmp_path = tmp.name + json.dump(base, tmp) + + try: + combined = xgb.Booster() + combined.load_model(tmp_path) + return combined + finally: + try: + os.remove(tmp_path) + except Exception: + pass + def _print_feature_importance( + self, + booster: xgb.Booster, + server_round: int, + ) -> None: + """ + Print feature importance for the current global booster. + """ + if booster is None: + return + + importance_gain = booster.get_score(importance_type="gain") + importance_weight = booster.get_score(importance_type="weight") + importance_cover = booster.get_score(importance_type="cover") + + print(f"\n[Server/XGBoost] 📊 Round {server_round} Global Feature Importance - Gain:") + print(importance_gain) + + print(f"[Server/XGBoost] 📊 Round {server_round} Global Feature Importance - Weight:") + print(importance_weight) + + print(f"[Server/XGBoost] 📊 Round {server_round} Global Feature Importance - Cover:") + print(importance_cover) \ No newline at end of file diff --git a/build/lib/MEDfl/NetManager/dataset.py b/build/lib/MEDfl/NetManager/dataset.py index 2023495..d50b847 100644 --- a/build/lib/MEDfl/NetManager/dataset.py +++ b/build/lib/MEDfl/NetManager/dataset.py @@ -48,10 +48,12 @@ def upload_dataset(self, NodeId=-1): data_df = pd.read_csv(self.path) nodeId = NodeId - columns = data_df.columns.tolist() + data_df = process_eicu(data_df) + columns = data_df.columns.tolist() + for index, row in data_df.iterrows(): query_1 = "INSERT INTO DataSets(DataSetName,nodeId," + "".join( f"{x}," for x in columns diff --git a/build/lib/MEDfl/NetManager/flsetup.py b/build/lib/MEDfl/NetManager/flsetup.py index 10567ad..79457bb 100644 --- a/build/lib/MEDfl/NetManager/flsetup.py +++ b/build/lib/MEDfl/NetManager/flsetup.py @@ -1,7 +1,7 @@ from datetime import datetime -from torch.utils.data import random_split, DataLoader, Dataset +from torch.utils.data import random_split, DataLoader, Dataset , TensorDataset from MEDfl.LearningManager.federated_dataset import FederatedDataset from .net_helper import * @@ -12,6 +12,11 @@ from MEDfl.NetManager.database_connector import DatabaseManager +import numpy as np +import torch +from sklearn.impute import SimpleImputer +from sklearn.preprocessing import LabelEncoder, StandardScaler +from torch.utils.data import random_split, DataLoader, Dataset, TensorDataset class FLsetup: def __init__(self, name: str, description: str, network: Network): @@ -176,7 +181,7 @@ def create_dataloader_from_node( fill_strategy="mean", fit_encode=[], to_drop=[], train_batch_size: int = 32, test_batch_size: int = 1, - split_frac: float = 0.2, + split_frac: float = 0, dataset: Dataset = None, ): @@ -212,84 +217,205 @@ def create_dataloader_from_node( return trainloader, testloader def create_federated_dataset( - self, output, fill_strategy="mean", fit_encode=[], to_drop=[], val_frac=0.1, test_frac=0.2 + self, + output, + fill_strategy="mean", + fit_encode=[], + to_drop=[], + val_frac=0.2, + test_frac=0.15, + train_batch_size=8, + task_type="binary", + scale_y=None, ) -> FederatedDataset: - """Create a federated dataset. - Args: - output (string): the output feature of the dataset - val_frac (float): The fraction of data to be used for validation. - test_frac (float): The fraction of data to be used for testing. + to_drop = list(to_drop) - Returns: - FederatedDataset: The FederatedDataset instance containing train, validation, and test data. - """ - if not self.column_name: - to_drop.extend(["DataSetName" , "NodeId" , "DataSetId"]) - else : - to_drop.extend(["PatientId"]) - + to_drop.extend(["DataSetName", "NodeId", "DataSetId"]) + else: + to_drop.extend(["PatientId"]) + netid = self.network.id + train_nodes = pd.read_sql( - text( - f"SELECT Nodes.NodeName FROM Nodes WHERE Nodes.NetId = {netid} AND Nodes.train = 1 " - ), + text(f"SELECT Nodes.NodeName FROM Nodes WHERE Nodes.NetId = {netid} AND Nodes.train = 1"), self.eng, ) test_nodes = pd.read_sql( - text( - f"SELECT Nodes.NodeName FROM Nodes WHERE Nodes.NetId = {netid} AND Nodes.train = 0 " - ), + text(f"SELECT Nodes.NodeName FROM Nodes WHERE Nodes.NetId = {netid} AND Nodes.train = 0"), self.eng, ) - train_nodes = [ - Node(val[0], 1, test_frac) for val in train_nodes.values.tolist() - ] + train_nodes = [Node(val[0], 1, test_frac) for val in train_nodes.values.tolist()] test_nodes = [Node(val[0], 0) for val in test_nodes.values.tolist()] trainloaders, valloaders, testloaders = [], [], [] - # if len(test_nodes) == 0: - # raise "test node empty" - if test_nodes is None: - _, testloader = self.create_dataloader_from_node( - train_nodes[0], output, fill_strategy=fill_strategy, fit_encode=fit_encode, to_drop=to_drop) - testloaders.append(testloader) - else: - for train_node in train_nodes: - train_valloader, testloader = self.create_dataloader_from_node( - train_node, output, fill_strategy=fill_strategy, - fit_encode=fit_encode, to_drop=to_drop,) - trainloader, valloader = self.create_dataloader_from_node( - train_node, - output, fill_strategy=fill_strategy, fit_encode=fit_encode, to_drop=to_drop, - test_batch_size=32, - split_frac=val_frac, - dataset=train_valloader.dataset, - ) - trainloaders.append(trainloader) - valloaders.append(valloader) - testloaders.append(testloader) + self.scalers = {} - for test_node in test_nodes: - _, testloader = self.create_dataloader_from_node( - test_node, output, fill_strategy=fill_strategy, fit_encode=fit_encode, to_drop=to_drop, split_frac=1.0 - ) - testloaders.append(testloader) - train_nodes_names = [node.name for node in train_nodes] - test_nodes_names = train_nodes_names + [ - node.name for node in test_nodes - ] + # ----------------------------- + # Pass 1: collect all training raw data + # ----------------------------- + cached_splits = {} + all_train_X = [] + all_train_y = [] + + for train_node in train_nodes: + if self.column_name is not None: + raw_data = train_node.get_dataset(self.column_name) + else: + raw_data = train_node.get_dataset() + + n = len(raw_data) + n_test = max(0, int(n * test_frac)) + n_trainval = n - n_test + + raw_data = raw_data.sample(frac=1, random_state=42).reset_index(drop=True) + trainval_raw = raw_data.iloc[:n_trainval].copy() + test_raw = raw_data.iloc[n_trainval:].copy() + + cached_splits[train_node.name] = { + "trainval_raw": trainval_raw, + "test_raw": test_raw, + } + + tmp = trainval_raw.copy() + + for s in fit_encode: + tmp[s] = LabelEncoder().fit_transform(tmp[s]) + + y_series = pd.to_numeric(tmp[output], errors="coerce") + X_df = tmp.copy() + + for column in to_drop: + if column in X_df.columns: + X_df = X_df.drop([column], axis=1) + + features = [col for col in X_df.columns if col != output] + + imputer = SimpleImputer(strategy=fill_strategy) + X_df[features] = imputer.fit_transform(X_df[features]) + + X_vals = X_df[features].values + y_vals = y_series.values.reshape(-1, 1) + + valid_mask = ~np.isnan(y_vals).ravel() + X_vals = X_vals[valid_mask] + y_vals = y_vals[valid_mask] + + if len(X_vals) > 0: + all_train_X.append(X_vals) + all_train_y.append(y_vals) + + if len(all_train_X) == 0: + raise ValueError("No training data available to fit global scalers") + + # ----------------------------- + # Pass 2: fit one global X scaler and one global y scaler + # ----------------------------- + global_scaler_X = StandardScaler() + global_scaler_X.fit(np.vstack(all_train_X)) + + if scale_y is None: + scale_y = (task_type == "regression") - # test_nodes_names = [ - # node.name for node in test_nodes - # ] + global_scaler_y = None + if task_type == "regression" and scale_y: + if len(all_train_y) == 0: + raise ValueError("No training targets available to fit global y scaler") + global_scaler_y = StandardScaler() + global_scaler_y.fit(np.vstack(all_train_y)) + + # ----------------------------- + # Pass 3: build datasets using the same global scalers + # ----------------------------- + for train_node in train_nodes: + trainval_raw = cached_splits[train_node.name]["trainval_raw"] + test_raw = cached_splits[train_node.name]["test_raw"] + + trainval_dataset, _, _ = process_data_after_reading( + trainval_raw.copy(), + output, + fill_strategy=fill_strategy, + fit_encode=fit_encode, + to_drop=to_drop, + scaler_X=global_scaler_X, + scaler_y=global_scaler_y, + fit_scaler=False, + task_type=task_type, + scale_y=scale_y, + ) + + self.scalers[train_node.name] = { + "X": global_scaler_X, + "y": global_scaler_y, + } + + n_tv = len(trainval_dataset) + n_val = max(0, int(n_tv * val_frac)) + n_tr = n_tv - n_val + + train_ds, val_ds = random_split(trainval_dataset, [n_tr, n_val]) - # Add FlSetup on to the DataBase - # self.create() + if len(test_raw) == 0: + if len(trainval_dataset) > 0: + n_features = len(trainval_dataset[0][0]) + else: + n_features = len(features) + + empty_X = torch.empty((0, n_features), dtype=torch.float32) + empty_y = torch.empty((0,), dtype=torch.float32) + test_dataset = TensorDataset(empty_X, empty_y) + else: + test_dataset, _, _ = process_data_after_reading( + test_raw.copy(), + output, + fill_strategy=fill_strategy, + fit_encode=fit_encode, + to_drop=to_drop, + scaler_X=global_scaler_X, + scaler_y=global_scaler_y, + fit_scaler=False, + task_type=task_type, + scale_y=scale_y, + ) + + trainloaders.append(DataLoader(train_ds, batch_size=train_batch_size, shuffle=True)) + valloaders.append(DataLoader(val_ds, batch_size=8)) + testloaders.append(DataLoader(test_dataset, batch_size=8)) + + print(f"{train_node.name}: train={n_tr}, val={n_val}, test={len(test_dataset)}") + + # test-only nodes + for test_node in test_nodes: + raw_data = test_node.get_dataset() + + if len(raw_data) == 0: + continue + + test_dataset, _, _ = process_data_after_reading( + raw_data.copy(), + output, + fill_strategy=fill_strategy, + fit_encode=fit_encode, + to_drop=to_drop, + scaler_X=global_scaler_X, + scaler_y=global_scaler_y, + fit_scaler=False, + task_type=task_type, + scale_y=scale_y, + ) + + self.scalers[test_node.name] = { + "X": global_scaler_X, + "y": global_scaler_y, + } + + testloaders.append(DataLoader(test_dataset, batch_size=8)) + + train_nodes_names = [node.name for node in train_nodes] + test_nodes_names = train_nodes_names + [node.name for node in test_nodes] - # self.network.update_network(FLsetupId=self.id) fed_dataset = FederatedDataset( self.name + "_Feddataset", train_nodes_names, @@ -298,13 +424,11 @@ def create_federated_dataset( valloaders, testloaders, ) + self.fed_dataset = fed_dataset self.fed_dataset.create(self.id) - return self.fed_dataset - - - + return self.fed_dataset def get_flDataSet(self): """ Retrieve the federated dataset associated with the FL setup using the FL setup's name. diff --git a/build/lib/MEDfl/NetManager/net_helper.py b/build/lib/MEDfl/NetManager/net_helper.py index 2d9aff0..3988345 100644 --- a/build/lib/MEDfl/NetManager/net_helper.py +++ b/build/lib/MEDfl/NetManager/net_helper.py @@ -9,7 +9,8 @@ import numpy as np from MEDfl.NetManager.database_connector import DatabaseManager - +from sklearn.preprocessing import LabelEncoder, StandardScaler +from sklearn.impute import SimpleImputer def is_str(data_df, row, x): """ @@ -61,63 +62,89 @@ def process_eicu(data_df): # remove indiserd columns after reading from the DB -def process_data_after_reading(data, output, fill_strategy="mean", fit_encode=[], to_drop=[]): - """ - Process data after reading from the database, including encoding, dropping columns, and creating a PyTorch TensorDataset. - - Args: - data (pandas.DataFrame): Input data. - output (str): Output column name. - fill_strategy (str, optional): Imputation strategy for missing values. Default is "mean". - fit_encode (list, optional): List of columns to be label-encoded. Default is an empty list. - to_drop (list, optional): List of columns to be dropped from the DataFrame. Default is an empty list. - - Returns: - torch.utils.data.TensorDataset: Processed data as a PyTorch TensorDataset. - """ - - # Check if there is a DataSet assigned to the node - if (len(data) == 0): - raise "Node doesn't Have dataSet" +def process_data_after_reading( + data, + output, + fill_strategy="mean", + fit_encode=[], + to_drop=[], + scaler_X=None, + scaler_y=None, + fit_scaler=True, + task_type="binary", # NEW + scale_y=None, # NEW +): + if len(data) == 0: + raise ValueError("Node doesn't have a dataset") encoder = LabelEncoder() - # En Code some columns for s in fit_encode: - try: - data[s] = encoder.fit_transform(data[s]) - except: - raise print(s) - - # The output of the DATA - y = data[output] + data[s] = encoder.fit_transform(data[s]) - X = data + y = data[output].copy() + X = data.copy() - # remove indisered columns when reading the dataframe from the DB for column in to_drop: - try: - X = X.drop( - [column], axis=1 - ) - except Exception as e: - raise e + if column in X.columns: + X = X.drop([column], axis=1) - # Get the DATAset Features features = [col for col in X.columns if col != output] - # Impute missing values using the mean strategy - try: - imputer = SimpleImputer(strategy=fill_strategy) - X[features] = imputer.fit_transform(X[features]) - except: - print() + imputer = SimpleImputer(strategy=fill_strategy) + X[features] = imputer.fit_transform(X[features]) - X = torch.tensor(X.values, dtype=torch.float32) - y = torch.tensor(y.values, dtype=torch.float32) - data = TensorDataset(X, y) + if scaler_X is None: + scaler_X = StandardScaler() + + X_vals = X[features].values + y_vals = pd.to_numeric(y, errors="coerce").values.reshape(-1, 1) + + valid_mask = ~np.isnan(y_vals).ravel() + X_vals = X_vals[valid_mask] + y_vals = y_vals[valid_mask] + + if fit_scaler: + X_vals = scaler_X.fit_transform(X_vals) + else: + X_vals = scaler_X.transform(X_vals) + + # default behavior depends on task + if scale_y is None: + scale_y = (task_type == "regression") + + if task_type == "binary": + y_vals = y_vals.ravel().astype(np.float32) + + uniq = set(np.unique(y_vals)) + if uniq == {-1.0, 1.0}: + y_vals = ((y_vals + 1.0) / 2.0).astype(np.float32) + elif uniq == {1.0, 2.0}: + y_vals = (y_vals - 1.0).astype(np.float32) + + uniq = set(np.unique(y_vals)) + if not uniq.issubset({0.0, 1.0}): + raise ValueError( + f"Binary classification requires targets in {{0,1}}. Found: {sorted(uniq)}" + ) + + elif task_type == "regression": + if scale_y: + if scaler_y is None: + scaler_y = StandardScaler() + if fit_scaler: + y_vals = scaler_y.fit_transform(y_vals).ravel() + else: + y_vals = scaler_y.transform(y_vals).ravel() + else: + y_vals = y_vals.ravel().astype(np.float32) + + else: + raise ValueError(f"Unsupported task_type: {task_type}") - return data + X_tensor = torch.tensor(X_vals, dtype=torch.float32) + y_tensor = torch.tensor(y_vals, dtype=torch.float32) + return TensorDataset(X_tensor, y_tensor), scaler_X, scaler_y def get_nodeid_from_name(name): """ diff --git a/build/lib/MEDfl/rw/xgboost/__init__.py b/build/lib/MEDfl/rw/xgboost/__init__.py new file mode 100644 index 0000000..726124b --- /dev/null +++ b/build/lib/MEDfl/rw/xgboost/__init__.py @@ -0,0 +1,16 @@ +""" +Federated XGBoost backend for MEDfl real-world mode. + +This module adds XGBoost support without modifying the existing PyTorch +real-world client, model, or strategy implementation. +""" + +from MEDfl.rw.xgboost.client import XGBoostFlowerClient +from MEDfl.rw.xgboost.strategy import XGBoostStrategy +from MEDfl.rw.xgboost.config import XGBoostConfig + +__all__ = [ + "XGBoostFlowerClient", + "XGBoostStrategy", + "XGBoostConfig", +] \ No newline at end of file diff --git a/build/lib/MEDfl/rw/xgboost/client.py b/build/lib/MEDfl/rw/xgboost/client.py new file mode 100644 index 0000000..d02c1ce --- /dev/null +++ b/build/lib/MEDfl/rw/xgboost/client.py @@ -0,0 +1,379 @@ +import argparse +import json +import socket +import platform +import shutil +from typing import Dict, Any, Optional + +import flwr as fl +import numpy as np +import pandas as pd +import psutil +import xgboost as xgb +from flwr.common import ( + Code, + EvaluateIns, + EvaluateRes, + FitIns, + FitRes, + GetParametersIns, + GetParametersRes, + GetPropertiesIns, + GetPropertiesRes, + Parameters, + Status, +) + +try: + import GPUtil +except ImportError: + GPUtil = None + +from MEDfl.rw.xgboost.data import XGBoostDataModule +from MEDfl.rw.xgboost.metrics import evaluate_predictions +from MEDfl.rw.xgboost.serialization import ( + booster_to_parameters, + parameters_to_booster, +) + + +class XGBoostFlowerClient(fl.client.Client): + """ + Real-world Flower client for federated XGBoost in MEDfl. + + This client is intentionally separate from the existing PyTorch NumPyClient. + The PyTorch client exchanges neural network weights. This client exchanges + serialized XGBoost boosters. + + Parameters + ---------- + server_address : str + Flower server address, e.g. "127.0.0.1:8080". + data_path : str + Local CSV dataset path. + val_frac : Optional[float] + Optional client-side validation fraction override. + test_frac : Optional[float] + Optional client-side test fraction override. + id_col : Optional[str] + Column used for ID-based test split. + test_ids : Optional[str] + Comma-separated test IDs. + seed : int + Random seed. + """ + + def __init__( + self, + server_address: str, + data_path: str = "data/data.csv", + val_frac: Optional[float] = None, + test_frac: Optional[float] = None, + id_col: Optional[str] = None, + test_ids: Optional[str] = None, + seed: int = 42, + ): + self.server_address = server_address + self.data_path = data_path + + self.data_module = XGBoostDataModule( + data_path=data_path, + val_frac=val_frac, + test_frac=test_frac, + id_col=id_col, + test_ids=test_ids, + seed=seed, + ) + + self.booster: Optional[xgb.Booster] = None + + self.dtrain = None + self.dval = None + self.dtest = None + self.metadata: Dict[str, Any] = {} + + self.task = "binary" + self.threshold = 0.5 + self.xgb_params: Dict[str, Any] = {} + self.local_num_boost_round = 1 + + self._initialized = False + + self._df_preview = pd.read_csv(data_path, nrows=5) + + def _initialize_from_config(self, config: Dict[str, Any]) -> None: + """ + Lazily initialize DMatrix objects from server config. + """ + + self.task = str(config.get("task", self.task) or self.task) + self.threshold = float(config.get("threshold", self.threshold)) + self.local_num_boost_round = int( + config.get("local_num_boost_round", self.local_num_boost_round) + ) + + xgb_params_raw = config.get("xgb_params", "{}") + + if isinstance(xgb_params_raw, str): + try: + self.xgb_params = json.loads(xgb_params_raw) + except json.JSONDecodeError as exc: + raise ValueError( + f"Invalid xgb_params JSON received from server: {xgb_params_raw}" + ) from exc + elif isinstance(xgb_params_raw, dict): + self.xgb_params = xgb_params_raw + else: + raise ValueError( + "xgb_params must be either a JSON string or a Python dictionary." + ) + + self.dtrain, self.dval, self.dtest, self.metadata = ( + self.data_module.prepare_from_config(config) + ) + + self._initialized = True + + print( + "[XGBoostClient] Initialized with " + f"task={self.task}, " + f"features={self.metadata.get('features')}, " + f"target={self.metadata.get('target')}, " + f"train={self.metadata.get('num_train')}, " + f"val={self.metadata.get('num_val')}, " + f"test={self.metadata.get('num_test')}" + ) + + def get_parameters(self, ins: GetParametersIns) -> GetParametersRes: + """ + Return current booster parameters. + + Before training, this can be empty. + """ + + return GetParametersRes( + status=Status(code=Code.OK, message="Success"), + parameters=booster_to_parameters(self.booster), + ) + + def fit(self, ins: FitIns) -> FitRes: + """ + Receive global booster, continue local XGBoost training, return updated booster. + """ + + config = dict(ins.config or {}) + + if not self._initialized: + self._initialize_from_config(config) + + received_booster = parameters_to_booster(ins.parameters) + + if received_booster is not None: + self.booster = received_booster + + evals = [] + if self.dval is not None and self.metadata.get("num_val", 0) > 0: + evals.append((self.dval, "validation")) + + self.booster = xgb.train( + params=self.xgb_params, + dtrain=self.dtrain, + num_boost_round=self.local_num_boost_round, + xgb_model=self.booster, + evals=evals, + verbose_eval=False, + ) + + train_pred = self.booster.predict(self.dtrain) + + train_metrics = evaluate_predictions( + task=self.task, + y_true=self.dtrain.get_label(), + y_pred_or_prob=train_pred, + threshold=self.threshold, + num_features=int(self.metadata.get("num_features", 0)), + ) + + metrics = { + "backend": "xgboost", + "task": self.task, + "hostname": socket.gethostname(), + "os_type": platform.system(), + "num_train": int(self.metadata.get("num_train", 0)), + "features": ",".join(self.metadata.get("features", [])), + "target": str(self.metadata.get("target", "")), + "val_fraction": float(self.metadata.get("val_fraction", 0.0)), + "test_fraction": float(self.metadata.get("test_fraction", 0.0)), + } + + for key, value in train_metrics.items(): + metrics[f"train_{key}"] = value + + return FitRes( + status=Status(code=Code.OK, message="Success"), + parameters=booster_to_parameters(self.booster), + num_examples=int(self.metadata.get("num_train", 0)), + metrics=metrics, + ) + + def evaluate(self, ins: EvaluateIns) -> EvaluateRes: + """ + Evaluate received global booster on local test set. + """ + + config = dict(ins.config or {}) + + if not self._initialized: + self._initialize_from_config(config) + + received_booster = parameters_to_booster(ins.parameters) + + if received_booster is not None: + self.booster = received_booster + + if self.booster is None: + return EvaluateRes( + status=Status(code=Code.EVALUATE_NOT_IMPLEMENTED, message="No booster."), + loss=0.0, + num_examples=0, + metrics={}, + ) + + y_true = self.dtest.get_label() + y_pred = self.booster.predict(self.dtest) + + eval_metrics = evaluate_predictions( + task=self.task, + y_true=y_true, + y_pred_or_prob=y_pred, + threshold=self.threshold, + num_features=int(self.metadata.get("num_features", 0)), + ) + + metrics = { + "backend": "xgboost", + "task": self.task, + } + + for key, value in eval_metrics.items(): + metrics[f"eval_{key}"] = value + + if self.task == "regression": + loss = float(eval_metrics.get("rmse", 0.0)) + elif self.task == "binary": + loss = float(eval_metrics.get("logloss", 0.0)) + else: + loss = float(eval_metrics.get("mlogloss", 0.0)) + + return EvaluateRes( + status=Status(code=Code.OK, message="Success"), + loss=loss, + num_examples=int(self.metadata.get("num_test", 0)), + metrics=metrics, + ) + + def get_properties(self, ins: GetPropertiesIns) -> GetPropertiesRes: + """ + Return machine and dataset properties. + """ + + hostname = socket.gethostname() + os_type = platform.system() + + df = self.data_module.df + + feature_names = df.columns[:-1].tolist() + target_name = df.columns[-1] + + try: + label_counts = df[target_name].value_counts().to_dict() + dist_str = ",".join(f"{k}:{v}" for k, v in label_counts.items()) + classes_str = ",".join(map(str, sorted(label_counts.keys()))) + except Exception: + dist_str = "" + classes_str = "" + + cpu_physical = psutil.cpu_count(logical=False) + cpu_logical = psutil.cpu_count(logical=True) + total_mem_gb = round(psutil.virtual_memory().total / (1024**3), 2) + + driver_present = shutil.which("nvidia-smi") is not None + gpu_count = 0 + + if GPUtil and driver_present: + try: + gpu_count = len(GPUtil.getGPUs()) + except Exception: + gpu_count = 0 + + properties = { + "backend": "xgboost", + "hostname": hostname, + "os_type": os_type, + "num_samples": int(len(df)), + "num_features": int(len(feature_names)), + "features": ",".join(feature_names), + "target": target_name, + "classes": classes_str, + "label_distribution": dist_str, + "cpu_physical_cores": int(cpu_physical or 0), + "cpu_logical_cores": int(cpu_logical or 0), + "total_memory_gb": float(total_mem_gb), + "gpu_driver_present": str(driver_present), + "gpu_count": int(gpu_count), + } + + return GetPropertiesRes( + status=Status(code=Code.OK, message="Success"), + properties=properties, + ) + + def start(self) -> None: + """ + Start the Flower client. + """ + + fl.client.start_client( + server_address=self.server_address, + client=self, + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description="MEDfl real-world XGBoost client") + + parser.add_argument( + "--server_address", + type=str, + required=True, + help="Flower server address, e.g. 127.0.0.1:8080", + ) + parser.add_argument( + "--data_path", + type=str, + default="data/data.csv", + help="Path to local CSV dataset.", + ) + parser.add_argument("--val_frac", type=float, default=None) + parser.add_argument("--test_frac", type=float, default=None) + parser.add_argument("--id_col", type=str, default=None) + parser.add_argument("--test_ids", type=str, default=None) + parser.add_argument("--seed", type=int, default=42) + + args = parser.parse_args() + + client = XGBoostFlowerClient( + server_address=args.server_address, + data_path=args.data_path, + val_frac=args.val_frac, + test_frac=args.test_frac, + id_col=args.id_col, + test_ids=args.test_ids, + seed=args.seed, + ) + + client.start() + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/build/lib/MEDfl/rw/xgboost/config.py b/build/lib/MEDfl/rw/xgboost/config.py new file mode 100644 index 0000000..2e26b70 --- /dev/null +++ b/build/lib/MEDfl/rw/xgboost/config.py @@ -0,0 +1,84 @@ +from dataclasses import dataclass, field +from typing import Dict, Any, Optional + + +@dataclass +class XGBoostConfig: + """ + Configuration object for MEDfl federated XGBoost. + + Parameters + ---------- + task : str + One of: "binary", "regression", "multiclass". + mode : str + Federated XGBoost mode. First supported mode: "bagging". + local_num_boost_round : int + Number of local boosting rounds per federated round. + params : dict + Native XGBoost training parameters. + num_classes : Optional[int] + Required for multiclass classification. + """ + + task: str = "binary" + mode: str = "bagging" + local_num_boost_round: int = 10 + params: Dict[str, Any] = field(default_factory=dict) + num_classes: Optional[int] = None + + def build_params(self) -> Dict[str, Any]: + """ + Build a safe default XGBoost parameter dictionary. + + User-provided params override the defaults. + """ + + if self.task == "binary": + defaults = { + "objective": "binary:logistic", + "eval_metric": "auc", + "tree_method": "hist", + "max_depth": 6, + "eta": 0.1, + "subsample": 1.0, + "colsample_bytree": 1.0, + } + + elif self.task == "regression": + defaults = { + "objective": "reg:squarederror", + "eval_metric": "rmse", + "tree_method": "hist", + "max_depth": 6, + "eta": 0.1, + "subsample": 1.0, + "colsample_bytree": 1.0, + } + + elif self.task == "multiclass": + if self.num_classes is None or self.num_classes < 2: + raise ValueError( + "num_classes must be provided and >= 2 for multiclass XGBoost." + ) + + defaults = { + "objective": "multi:softprob", + "num_class": self.num_classes, + "eval_metric": "mlogloss", + "tree_method": "hist", + "max_depth": 6, + "eta": 0.1, + "subsample": 1.0, + "colsample_bytree": 1.0, + } + + else: + raise ValueError( + f"Unsupported XGBoost task '{self.task}'. " + "Choose from: 'binary', 'regression', 'multiclass'." + ) + + merged = defaults.copy() + merged.update(self.params or {}) + return merged \ No newline at end of file diff --git a/build/lib/MEDfl/rw/xgboost/data.py b/build/lib/MEDfl/rw/xgboost/data.py new file mode 100644 index 0000000..73adb07 --- /dev/null +++ b/build/lib/MEDfl/rw/xgboost/data.py @@ -0,0 +1,253 @@ +from typing import Dict, Any, Optional, Tuple, List + +import numpy as np +import pandas as pd +import xgboost as xgb +from sklearn.model_selection import train_test_split + + +class XGBoostDataModule: + """ + Data preparation module for MEDfl federated XGBoost clients. + + This mirrors the behavior of the existing PyTorch real-world client: + - server-enforced features; + - server-enforced target; + - global or per-client split fractions; + - optional ID-based test split. + """ + + def __init__( + self, + data_path: str, + val_frac: Optional[float] = None, + test_frac: Optional[float] = None, + id_col: Optional[str] = None, + test_ids: Optional[str] = None, + seed: int = 42, + ): + self.data_path = data_path + self.df = pd.read_csv(data_path) + + self.client_val_frac = val_frac + self.client_test_frac = test_frac + self.id_col = id_col + self.test_ids = test_ids + self.seed = seed + + self.feature_names = self.df.columns[:-1].tolist() + self.target_name = self.df.columns[-1] + + self.effective_features = self.feature_names[:] + self.effective_target = self.target_name + self.effective_val_frac = 0.0 + self.effective_test_frac = 0.0 + + def prepare_from_config( + self, + config: Dict[str, Any], + ) -> Tuple[xgb.DMatrix, Optional[xgb.DMatrix], xgb.DMatrix, Dict[str, Any]]: + """ + Prepare XGBoost DMatrix objects from server config. + """ + + srv_features = str(config.get("features", "") or "").strip() + srv_target = str(config.get("target", "") or "").strip() + + if srv_target: + if srv_target not in self.df.columns: + raise ValueError( + f"Server-specified target '{srv_target}' not found in CSV columns: " + f"{list(self.df.columns)}" + ) + target_col = srv_target + else: + target_col = self.df.columns[-1] + + if srv_features: + feat_cols = [c.strip() for c in srv_features.split(",") if c.strip()] + missing = [c for c in feat_cols if c not in self.df.columns] + if missing: + raise ValueError( + f"Server-specified feature columns not found in CSV: {missing}" + ) + else: + feat_cols = [c for c in self.df.columns if c != target_col] + + srv_val = config.get("val_fraction", None) + srv_test = config.get("test_fraction", None) + + val_frac = ( + float(self.client_val_frac) + if self.client_val_frac is not None + else float(srv_val) + if srv_val is not None + else 0.10 + ) + + test_frac = ( + float(self.client_test_frac) + if self.client_test_frac is not None + else float(srv_test) + if srv_test is not None + else 0.10 + ) + + if not (0.0 <= val_frac < 1.0): + raise ValueError(f"Invalid val_fraction={val_frac}. Must be in [0, 1).") + + if config.get("id_col") and not self.id_col: + self.id_col = str(config.get("id_col")) + + if config.get("test_ids") and not self.test_ids: + self.test_ids = str(config.get("test_ids")) + + X_all = self.df[feat_cols].values + y_all = self.df[target_col].values + + if self.test_ids and str(self.test_ids).strip(): + X_train, X_val, X_test, y_train, y_val, y_test, actual_test_frac = ( + self._split_with_test_ids( + X_all=X_all, + y_all=y_all, + val_frac=val_frac, + ) + ) + self.effective_test_frac = actual_test_frac + + else: + if not (0.0 <= test_frac < 1.0 and val_frac + test_frac < 1.0): + raise ValueError( + f"Invalid fractions: val={val_frac}, test={test_frac}. " + "Require val >= 0, test >= 0, val + test < 1." + ) + + X_train, X_val, X_test, y_train, y_val, y_test = self._split_with_fractions( + X_all=X_all, + y_all=y_all, + val_frac=val_frac, + test_frac=test_frac, + ) + self.effective_test_frac = test_frac + + self.effective_features = feat_cols + self.effective_target = target_col + self.effective_val_frac = val_frac + + dtrain = xgb.DMatrix(X_train, label=y_train, feature_names=feat_cols) + + dval = None + if X_val is not None and len(y_val) > 0: + dval = xgb.DMatrix(X_val, label=y_val, feature_names=feat_cols) + + dtest = xgb.DMatrix(X_test, label=y_test, feature_names=feat_cols) + + metadata = { + "features": feat_cols, + "target": target_col, + "num_features": len(feat_cols), + "num_train": int(len(y_train)), + "num_val": int(len(y_val)) if y_val is not None else 0, + "num_test": int(len(y_test)), + "val_fraction": float(self.effective_val_frac), + "test_fraction": float(self.effective_test_frac), + } + + return dtrain, dval, dtest, metadata + + def _split_with_fractions( + self, + X_all, + y_all, + val_frac: float, + test_frac: float, + ): + strat_all = y_all if len(np.unique(y_all)) > 1 else None + + X_trval, X_test, y_trval, y_test = train_test_split( + X_all, + y_all, + test_size=test_frac, + random_state=self.seed, + stratify=strat_all, + ) + + if val_frac > 0 and len(y_trval) > 0: + actual_test_frac = len(y_test) / len(y_all) + rel_val = val_frac / (1.0 - actual_test_frac) + + strat_tr = y_trval if len(np.unique(y_trval)) > 1 else None + + X_train, X_val, y_train, y_val = train_test_split( + X_trval, + y_trval, + test_size=rel_val, + random_state=self.seed, + stratify=strat_tr, + ) + else: + X_train, y_train = X_trval, y_trval + X_val = np.empty((0, X_all.shape[1])) + y_val = np.empty((0,)) + + return X_train, X_val, X_test, y_train, y_val, y_test + + def _split_with_test_ids( + self, + X_all, + y_all, + val_frac: float, + ): + test_ids_list = [i.strip() for i in str(self.test_ids).split(",") if i.strip()] + + if self.id_col and self.id_col in self.df.columns: + id_series = self.df[self.id_col] + + if np.issubdtype(id_series.dtype, np.number): + test_ids_list = [int(i) for i in test_ids_list] + else: + test_ids_list = [str(i) for i in test_ids_list] + + else: + id_series = self.df.index + + try: + test_ids_list = [int(i) for i in test_ids_list] + except ValueError as exc: + raise ValueError( + "test_ids must be integers when id_col is not provided " + "or not found in the CSV." + ) from exc + + test_mask = id_series.isin(test_ids_list) + + X_test = X_all[test_mask] + y_test = y_all[test_mask] + X_trval = X_all[~test_mask] + y_trval = y_all[~test_mask] + + actual_test_frac = len(y_test) / len(y_all) if len(y_all) > 0 else 0.0 + + if val_frac + actual_test_frac >= 1.0: + raise ValueError( + f"Validation fraction {val_frac} + actual test fraction " + f"{actual_test_frac} >= 1." + ) + + if val_frac > 0 and len(y_trval) > 0: + rel_val = val_frac / (1.0 - actual_test_frac) + strat_tr = y_trval if len(np.unique(y_trval)) > 1 else None + + X_train, X_val, y_train, y_val = train_test_split( + X_trval, + y_trval, + test_size=rel_val, + random_state=self.seed, + stratify=strat_tr, + ) + else: + X_train, y_train = X_trval, y_trval + X_val = np.empty((0, X_all.shape[1])) + y_val = np.empty((0,)) + + return X_train, X_val, X_test, y_train, y_val, y_test, actual_test_frac \ No newline at end of file diff --git a/build/lib/MEDfl/rw/xgboost/metrics.py b/build/lib/MEDfl/rw/xgboost/metrics.py new file mode 100644 index 0000000..ea08866 --- /dev/null +++ b/build/lib/MEDfl/rw/xgboost/metrics.py @@ -0,0 +1,176 @@ +import math +from typing import Dict, Any + +import numpy as np +from sklearn.metrics import ( + accuracy_score, + roc_auc_score, + log_loss, + mean_squared_error, + mean_absolute_error, + r2_score, + f1_score, +) + + +def _safe_float(value: Any) -> float: + try: + value = float(value) + if math.isnan(value) or math.isinf(value): + return float("nan") + return value + except Exception: + return float("nan") + + +def binary_metrics(y_true, y_prob, threshold: float = 0.5) -> Dict[str, float]: + """ + Metrics for binary classification. + """ + + y_true = np.asarray(y_true) + y_prob = np.asarray(y_prob) + + y_pred = (y_prob >= threshold).astype(int) + + metrics = { + "accuracy": _safe_float(accuracy_score(y_true, y_pred)), + } + + try: + metrics["auc"] = _safe_float(roc_auc_score(y_true, y_prob)) + except Exception: + metrics["auc"] = float("nan") + + try: + metrics["logloss"] = _safe_float(log_loss(y_true, y_prob)) + except Exception: + metrics["logloss"] = float("nan") + + return metrics + + +def regression_metrics(y_true, y_pred, num_features: int = 0) -> Dict[str, float]: + """ + Metrics for regression. + """ + + y_true = np.asarray(y_true) + y_pred = np.asarray(y_pred) + + rmse = mean_squared_error(y_true, y_pred, squared=False) + mae = mean_absolute_error(y_true, y_pred) + r2 = r2_score(y_true, y_pred) + + n = len(y_true) + p = int(num_features) + + if n > p + 1: + adjusted_r2 = 1.0 - (1.0 - r2) * ((n - 1.0) / (n - p - 1.0)) + else: + adjusted_r2 = float("nan") + + return { + "rmse": _safe_float(rmse), + "mae": _safe_float(mae), + "r2": _safe_float(r2), + "adjusted_r2": _safe_float(adjusted_r2), + } + + +def multiclass_metrics(y_true, y_prob) -> Dict[str, float]: + """ + Metrics for multiclass classification. + """ + + y_true = np.asarray(y_true) + y_prob = np.asarray(y_prob) + + y_pred = np.argmax(y_prob, axis=1) + + metrics = { + "accuracy": _safe_float(accuracy_score(y_true, y_pred)), + } + + try: + metrics["macro_f1"] = _safe_float(f1_score(y_true, y_pred, average="macro")) + except Exception: + metrics["macro_f1"] = float("nan") + + try: + metrics["mlogloss"] = _safe_float(log_loss(y_true, y_prob)) + except Exception: + metrics["mlogloss"] = float("nan") + + return metrics + + +def evaluate_predictions( + task: str, + y_true, + y_pred_or_prob, + threshold: float = 0.5, + num_features: int = 0, +) -> Dict[str, float]: + """ + Task-aware metric dispatcher. + """ + + if task == "binary": + return binary_metrics(y_true, y_pred_or_prob, threshold=threshold) + + if task == "regression": + return regression_metrics(y_true, y_pred_or_prob, num_features=num_features) + + if task == "multiclass": + return multiclass_metrics(y_true, y_pred_or_prob) + + raise ValueError( + f"Unsupported task '{task}'. Choose from 'binary', 'regression', 'multiclass'." + ) + + +def aggregate_weighted_metrics(results): + """ + Weighted aggregation for Flower metric aggregation callbacks. + + Expected Flower format: + List[Tuple[int, Dict[str, Scalar]]] + """ + + total = sum(num_examples for num_examples, _ in results) + + if total == 0: + return {} + + keys = set() + for _, metrics in results: + keys.update(metrics.keys()) + + aggregated = {} + + for key in keys: + values = [] + weights = [] + + for num_examples, metrics in results: + if key not in metrics: + continue + + value = metrics[key] + + try: + value = float(value) + except Exception: + continue + + if math.isnan(value) or math.isinf(value): + continue + + values.append(value) + weights.append(num_examples) + + if values: + aggregated[key] = float(np.average(values, weights=weights)) + + return aggregated \ No newline at end of file diff --git a/build/lib/MEDfl/rw/xgboost/serialization.py b/build/lib/MEDfl/rw/xgboost/serialization.py new file mode 100644 index 0000000..e316c05 --- /dev/null +++ b/build/lib/MEDfl/rw/xgboost/serialization.py @@ -0,0 +1,56 @@ +from typing import Optional + +import xgboost as xgb +from flwr.common import Parameters + + +XGBOOST_TENSOR_TYPE = "xgboost.ubj" + + +def booster_to_parameters(booster: Optional[xgb.Booster]) -> Parameters: + """ + Serialize an XGBoost Booster into Flower Parameters. + + XGBoost does not exchange PyTorch-style NumPy tensors. Instead, we serialize + the Booster as raw UBJSON bytes and store them in Flower Parameters. + """ + + if booster is None: + return Parameters(tensors=[], tensor_type=XGBOOST_TENSOR_TYPE) + + raw = booster.save_raw(raw_format="ubj") + return Parameters(tensors=[bytes(raw)], tensor_type=XGBOOST_TENSOR_TYPE) + + +def parameters_to_booster(parameters: Optional[Parameters]) -> Optional[xgb.Booster]: + """ + Deserialize Flower Parameters into an XGBoost Booster. + """ + + if parameters is None: + return None + + if not parameters.tensors: + return None + + booster = xgb.Booster() + booster.load_model(bytearray(parameters.tensors[0])) + return booster + + +def save_booster(booster: xgb.Booster, path: str) -> None: + """ + Save an XGBoost Booster to disk. + """ + + booster.save_model(path) + + +def load_booster(path: str) -> xgb.Booster: + """ + Load an XGBoost Booster from disk. + """ + + booster = xgb.Booster() + booster.load_model(path) + return booster \ No newline at end of file diff --git a/build/lib/MEDfl/rw/xgboost/strategy.py b/build/lib/MEDfl/rw/xgboost/strategy.py new file mode 100644 index 0000000..029ebc3 --- /dev/null +++ b/build/lib/MEDfl/rw/xgboost/strategy.py @@ -0,0 +1,668 @@ +import json +import os +from typing import Dict, Any, Optional +import tempfile +import xgboost as xgb + +import flwr as fl +from flwr.common import ( + Parameters, + FitRes, + parameters_to_ndarrays, +) +from flwr.common import GetPropertiesIns + +from MEDfl.rw.xgboost.config import XGBoostConfig +from MEDfl.rw.xgboost.metrics import aggregate_weighted_metrics +from MEDfl.rw.xgboost.serialization import ( + parameters_to_booster,booster_to_parameters, + save_booster, + XGBOOST_TENSOR_TYPE, +) + + +class XGBoostStrategy: + """ + Strategy wrapper for MEDfl real-world federated XGBoost. + + This class follows the same high-level contract as the existing MEDfl + Strategy wrapper: + - it exposes create_strategy(); + - it stores the Flower strategy in self.strategy_object; + - it can be passed directly into FederatedServer. + + First supported mode: + - bagging-style federated XGBoost using Flower strategy hooks. + + Notes + ----- + This implementation is conservative and non-invasive: + - It does not modify the existing PyTorch FedAvg/FedAdam/FedYogi path. + - It uses Flower Parameters containing serialized XGBoost boosters. + - It saves XGBoost models as .ubj files instead of .npz files. + """ + + def __init__( + self, + mode: str = "bagging", + fraction_fit: float = 1.0, + fraction_evaluate: float = 1.0, + min_fit_clients: int = 2, + min_evaluate_clients: int = 2, + min_available_clients: int = 2, + task: str = "binary", + num_classes: Optional[int] = None, + xgb_params: Optional[Dict[str, Any]] = None, + local_num_boost_round: int = 10, + threshold: float = 0.5, + features: str = "", + target: str = "", + val_fraction: float = 0.10, + test_fraction: float = 0.10, + split_mode: str = "global", + client_fractions: Optional[Dict[str, Dict[str, Any]]] = None, + id_col: str = "id", + savingPath: str = "", + saveOnRounds: int = 1, + total_rounds: int = 3, + ): + self.name = "FedXgbBagging" + self.mode = mode + + if self.mode != "bagging": + raise ValueError( + "Only mode='bagging' is implemented in this first integration. " + "Cyclic mode should be added after bagging is stable." + ) + + self.fraction_fit = fraction_fit + self.fraction_evaluate = fraction_evaluate + self.min_fit_clients = min_fit_clients + self.min_evaluate_clients = min_evaluate_clients + self.min_available_clients = min_available_clients + + self.task = task + self.num_classes = num_classes + self.threshold = threshold + + config = XGBoostConfig( + task=task, + mode=mode, + local_num_boost_round=local_num_boost_round, + params=xgb_params or {}, + num_classes=num_classes, + ) + + self.xgb_params = config.build_params() + self.local_num_boost_round = local_num_boost_round + + self._features = features + self._target = target + self._val_fraction = val_fraction + self._test_fraction = test_fraction + + self.split_mode = split_mode + self.client_fractions = client_fractions or {} + self.id_col = id_col + + self.savingPath = savingPath + self.saveOnRounds = saveOnRounds + self.total_rounds = total_rounds + + self.strategy_object = None + self.current_booster: Optional[xgb.Booster] = None + self.current_num_trees: int = 0 + + def create_strategy(self) -> None: + """ + Create the real Flower strategy object. + + We use FedAvg as the orchestration base but override aggregation behavior + to keep XGBoost booster exchange explicit. + """ + + def fit_config_fn(server_round: int) -> Dict[str, Any]: + return { + "backend": "xgboost", + "mode": self.mode, + "task": self.task, + "threshold": float(self.threshold), + "features": self._features, + "target": self._target, + "val_fraction": float(self._val_fraction), + "test_fraction": float(self._test_fraction), + "id_col": self.id_col, + "local_num_boost_round": int(self.local_num_boost_round), + "xgb_params": json.dumps(self.xgb_params), + } + + def evaluate_config_fn(server_round: int) -> Dict[str, Any]: + return { + "backend": "xgboost", + "mode": self.mode, + "task": self.task, + "threshold": float(self.threshold), + "features": self._features, + "target": self._target, + "val_fraction": float(self._val_fraction), + "test_fraction": float(self._test_fraction), + "id_col": self.id_col, + "local_num_boost_round": int(self.local_num_boost_round), + "xgb_params": json.dumps(self.xgb_params), + } + + strategy = fl.server.strategy.FedAvg( + fraction_fit=self.fraction_fit, + fraction_evaluate=self.fraction_evaluate, + min_fit_clients=self.min_fit_clients, + min_evaluate_clients=self.min_evaluate_clients, + min_available_clients=self.min_available_clients, + on_fit_config_fn=fit_config_fn, + on_evaluate_config_fn=evaluate_config_fn, + fit_metrics_aggregation_fn=aggregate_weighted_metrics, + evaluate_metrics_aggregation_fn=aggregate_weighted_metrics, + initial_parameters=Parameters( + tensors=[], + tensor_type=XGBOOST_TENSOR_TYPE, + ), + ) + + self._wrap_configure_fit(strategy) + self._wrap_aggregate_fit(strategy) + self._wrap_aggregate_evaluate(strategy) + + self.strategy_object = strategy + + def _wrap_configure_fit(self, strategy) -> None: + """ + Inject per-client test fractions / test IDs exactly like the existing MEDfl strategy. + """ + + original_configure_fit = strategy.configure_fit + + def wrapped_configure_fit(server_round, parameters, client_manager): + selected = original_configure_fit( + server_round=server_round, + parameters=parameters, + client_manager=client_manager, + ) + + ins = GetPropertiesIns(config={}) + + for client, fit_ins in selected: + hostname = None + + try: + props = client.get_properties( + ins=ins, + timeout=10.0, + group_id=0, + ) + print( + f"\n📋 [XGBoost Round {server_round}] " + f"Client {client.cid} Properties: {props.properties}" + ) + hostname = props.properties.get("hostname", None) + + except Exception as exc: + print(f"⚠️ Failed to get properties from {client.cid}: {exc}") + + if not hostname: + hostname = client.cid + + cfg = fit_ins.config + + if self.split_mode == "per_client": + per_cfg = ( + self.client_fractions.get(hostname) + or self.client_fractions.get(client.cid) + or {} + ) + + if "val_fraction" in per_cfg: + try: + cfg["val_fraction"] = float(per_cfg["val_fraction"]) + except Exception: + pass + + if "test_ids" in per_cfg and per_cfg["test_ids"]: + test_ids_val = per_cfg["test_ids"] + + if isinstance(test_ids_val, (list, tuple, set)): + test_ids_str = ",".join(str(x) for x in test_ids_val) + else: + test_ids_str = str(test_ids_val) + + cfg["test_ids"] = test_ids_str + + if "test_fraction" in cfg: + del cfg["test_fraction"] + + cfg["id_col"] = self.id_col + + else: + if "test_fraction" in per_cfg: + try: + cfg["test_fraction"] = float(per_cfg["test_fraction"]) + except Exception: + pass + + else: + if "test_ids" in cfg: + del cfg["test_ids"] + + cfg["val_fraction"] = float(self._val_fraction) + cfg["test_fraction"] = float(self._test_fraction) + cfg["id_col"] = self.id_col + + return selected + + strategy.configure_fit = wrapped_configure_fit + + def _wrap_aggregate_fit(self, strategy) -> None: + """ + Wrap aggregate_fit for XGBoost. + + IMPORTANT: + XGBoost boosters are NOT NumPy tensors. + Therefore we MUST NOT call FedAvg tensor aggregation. + + Current strategy: + - select the booster from the client with the largest dataset + - aggregate only metrics manually + - save the selected booster + """ + + def wrapped_aggregate_fit(server_round, results, failures): + print(f"\n[Server/XGBoost] 🔄 Round {server_round} - Client Training Metrics:") + + for client, fit_res in results: + print( + f" XGB CTM Round {server_round} " + f"Client:{client.cid}: {fit_res.metrics}" + ) + + if failures: + print(f"[Server/XGBoost] ⚠️ Round {server_round} failures: {failures}") + + if not results: + print("[Server/XGBoost] No fit results received.") + return None, {} + + # Aggregate metrics manually, same as before. + metric_results = [] + + for _, fit_res in results: + metric_results.append( + ( + fit_res.num_examples, + fit_res.metrics, + ) + ) + + metrics = aggregate_weighted_metrics(metric_results) + + previous_num_trees = int(self.current_num_trees) + new_local_boosters = [] + + print( + f"[Server/XGBoost] Previous global tree count before round " + f"{server_round}: {previous_num_trees}" + ) + + for client, fit_res in results: + client_booster = parameters_to_booster(fit_res.parameters) + + if client_booster is None: + print( + f"[Server/XGBoost] Client {client.cid} returned no booster. Skipping." + ) + continue + + client_total_trees = self._num_boosted_rounds(client_booster) + + print( + f"[Server/XGBoost] Client {client.cid} returned booster with " + f"{client_total_trees} trees." + ) + + # In round 1, previous_num_trees is 0, so this keeps all local trees. + # In later rounds, this keeps only the new trees added by the client. + new_trees_booster = self._slice_new_trees( + client_booster=client_booster, + previous_num_trees=previous_num_trees, + ) + + if new_trees_booster is None: + continue + + new_tree_count = self._num_boosted_rounds(new_trees_booster) + + print( + f"[Server/XGBoost] Client {client.cid} contributed " + f"{new_tree_count} new trees." + ) + + new_local_boosters.append(new_trees_booster) + + if not new_local_boosters: + print( + "[Server/XGBoost] No new local trees received. " + "Keeping previous global booster." + ) + + if self.current_booster is None: + return None, metrics + + agg_parameters = booster_to_parameters(self.current_booster) + return agg_parameters, metrics + + boosters_to_combine = [] + + if self.current_booster is not None: + boosters_to_combine.append(self.current_booster) + + boosters_to_combine.extend(new_local_boosters) + + combined_booster = self._combine_boosters_by_json(boosters_to_combine) + + if combined_booster is None: + print("[Server/XGBoost] Failed to combine boosters.") + return None, metrics + + self.current_booster = combined_booster + self.current_num_trees = self._num_boosted_rounds(combined_booster) + + print( + f"[Server/XGBoost] ✅ Round {server_round} - Flower-style bagging aggregation done." + ) + print( + f"[Server/XGBoost] Global booster now has {self.current_num_trees} trees." + ) + print(f"[Server/XGBoost] Aggregated metrics: {metrics}\n") + + # Optional: print importance directly, no saving required. + if server_round == self.total_rounds: + self._print_feature_importance(combined_booster, server_round) + + agg_parameters = booster_to_parameters(combined_booster) + + # Optional: keep your existing save behavior if savingPath is provided. + self._save_if_needed(server_round, agg_parameters) + + return agg_parameters, metrics + strategy.aggregate_fit = wrapped_aggregate_fit + + + def _wrap_aggregate_evaluate(self, strategy) -> None: + """ + Wrap aggregate_evaluate for logging. + """ + + original_aggregate_evaluate = strategy.aggregate_evaluate + + def wrapped_aggregate_evaluate(server_round, results, failures): + print(f"\n[Server/XGBoost] 📊 Round {server_round} - Client Evaluation Metrics:") + + for client, eval_res in results: + print( + f" XGB CEM Round {server_round} " + f"Client:{client.cid}: {eval_res.metrics}" + ) + + loss, metrics = original_aggregate_evaluate(server_round, results, failures) + + print(f"[Server/XGBoost] ✅ Round {server_round} - Aggregated Evaluation:") + print(f" Loss: {loss}, Metrics: {metrics}\n") + + return loss, metrics + + strategy.aggregate_evaluate = wrapped_aggregate_evaluate + + def _save_if_needed(self, server_round: int, parameters: Parameters) -> None: + """ + Save XGBoost model as .ubj. + """ + + if not self.savingPath: + return + + should_save = ( + server_round % self.saveOnRounds == 0 + or server_round == self.total_rounds + ) + + if not should_save: + return + + os.makedirs(self.savingPath, exist_ok=True) + + booster = parameters_to_booster(parameters) + + if booster is None: + print("[Server/XGBoost] No booster to save.") + return + + filename = ( + f"round_{server_round}_final_model.ubj" + if server_round == self.total_rounds + else f"round_{server_round}_model.ubj" + ) + + filepath = os.path.join(self.savingPath, filename) + + save_booster(booster, filepath) + + print(f"[Server/XGBoost] 💾 Saved booster to: {filepath}") + def _num_boosted_rounds(self, booster: xgb.Booster) -> int: + """ + Return the number of boosted rounds/trees in the booster. + """ + try: + return int(booster.num_boosted_rounds()) + except Exception: + try: + return len(booster.get_dump()) + except Exception: + return 0 + + + def _slice_new_trees( + self, + client_booster: xgb.Booster, + previous_num_trees: int, + ) -> Optional[xgb.Booster]: + """ + Extract only the trees added by the client in this round. + + Client booster = previous global booster + new local trees. + We keep only new local trees. + """ + + total_trees = self._num_boosted_rounds(client_booster) + + if total_trees <= previous_num_trees: + print( + "[Server/XGBoost] Client booster has no new trees. " + f"total_trees={total_trees}, previous_num_trees={previous_num_trees}" + ) + return None + + try: + return client_booster[previous_num_trees:total_trees] + except Exception as exc: + print(f"[Server/XGBoost] Failed to slice new trees: {exc}") + return None + + + def _booster_to_raw_bytes(self, booster: xgb.Booster) -> bytes: + """ + Serialize booster to raw UBJ bytes. + """ + with tempfile.NamedTemporaryFile(suffix=".ubj", delete=False) as tmp: + tmp_path = tmp.name + + try: + booster.save_model(tmp_path) + with open(tmp_path, "rb") as f: + return f.read() + finally: + try: + os.remove(tmp_path) + except Exception: + pass + + + def _booster_from_raw_bytes(self, raw: bytes) -> Optional[xgb.Booster]: + """ + Deserialize booster from raw UBJ bytes. + """ + if not raw: + return None + + with tempfile.NamedTemporaryFile(suffix=".ubj", delete=False) as tmp: + tmp_path = tmp.name + tmp.write(raw) + + try: + booster = xgb.Booster() + booster.load_model(tmp_path) + return booster + except Exception as exc: + print(f"[Server/XGBoost] Failed to load booster from bytes: {exc}") + return None + finally: + try: + os.remove(tmp_path) + except Exception: + pass + + def _combine_boosters_by_json( + self, + boosters: list[xgb.Booster], + ) -> Optional[xgb.Booster]: + """ + Combine multiple XGBoost boosters by concatenating their trees. + + Fixes: + - updates trees + - updates tree_info + - updates gbtree_model_param["num_trees"] + - updates iteration_indptr + + This is the missing part that caused: + Check failed: model.iteration_indptr.back() == model.param.num_trees + """ + + boosters = [b for b in boosters if b is not None] + + if not boosters: + return None + + if len(boosters) == 1: + return boosters[0] + + model_jsons = [] + + for booster in boosters: + raw_json = booster.save_raw(raw_format="json") + if isinstance(raw_json, bytes): + raw_json = raw_json.decode("utf-8") + model_jsons.append(json.loads(raw_json)) + + base = model_jsons[0] + + base_model = base["learner"]["gradient_booster"]["model"] + base_trees = base_model["trees"] + base_tree_info = base_model["tree_info"] + + # Start from the base model iteration structure + base_iteration_indptr = base_model.get("iteration_indptr", None) + + if base_iteration_indptr is None: + # Safe fallback for binary/regression: one tree per iteration + combined_iteration_indptr = list(range(0, len(base_trees) + 1)) + else: + combined_iteration_indptr = list(base_iteration_indptr) + + next_tree_id = len(base_trees) + + for extra in model_jsons[1:]: + extra_model = extra["learner"]["gradient_booster"]["model"] + + extra_trees = extra_model["trees"] + extra_tree_info = extra_model["tree_info"] + + # Append trees and reassign IDs + for tree in extra_trees: + tree["id"] = next_tree_id + base_trees.append(tree) + next_tree_id += 1 + + # Append tree_info + base_tree_info.extend(extra_tree_info) + + # Append iteration_indptr correctly + extra_iteration_indptr = extra_model.get("iteration_indptr", None) + + if extra_iteration_indptr is None: + # Fallback: one tree per boosting iteration + extra_iteration_indptr = list(range(0, len(extra_trees) + 1)) + + offset = combined_iteration_indptr[-1] + + # Skip the first 0 from the extra booster + for ptr in extra_iteration_indptr[1:]: + combined_iteration_indptr.append(offset + ptr) + + # Write back merged model fields + base_model["trees"] = base_trees + base_model["tree_info"] = base_tree_info + base_model["iteration_indptr"] = combined_iteration_indptr + + # Update num_trees + if "gbtree_model_param" in base_model: + base_model["gbtree_model_param"]["num_trees"] = str(len(base_trees)) + + # Save/load to validate the merged model + with tempfile.NamedTemporaryFile( + mode="w", + suffix=".json", + delete=False, + encoding="utf-8", + ) as tmp: + tmp_path = tmp.name + json.dump(base, tmp) + + try: + combined = xgb.Booster() + combined.load_model(tmp_path) + return combined + finally: + try: + os.remove(tmp_path) + except Exception: + pass + def _print_feature_importance( + self, + booster: xgb.Booster, + server_round: int, + ) -> None: + """ + Print feature importance for the current global booster. + """ + if booster is None: + return + + importance_gain = booster.get_score(importance_type="gain") + importance_weight = booster.get_score(importance_type="weight") + importance_cover = booster.get_score(importance_type="cover") + + print(f"\n[Server/XGBoost] 📊 Round {server_round} Global Feature Importance - Gain:") + print(importance_gain) + + print(f"[Server/XGBoost] 📊 Round {server_round} Global Feature Importance - Weight:") + print(importance_weight) + + print(f"[Server/XGBoost] 📊 Round {server_round} Global Feature Importance - Cover:") + print(importance_cover) \ No newline at end of file diff --git a/medfl.egg-info/PKG-INFO b/medfl.egg-info/PKG-INFO index 76956df..f8a2c44 100644 --- a/medfl.egg-info/PKG-INFO +++ b/medfl.egg-info/PKG-INFO @@ -1,6 +1,6 @@ -Metadata-Version: 2.4 +Metadata-Version: 2.1 Name: medfl -Version: 2.0.5.dev9 +Version: 2.0.10.dev5 Summary: Python Open-source package for simulating federated learning and differential privacy Home-page: https://github.com/MEDomics-UdeS/MEDfl Author: MEDomics consortium @@ -17,37 +17,6 @@ Classifier: Programming Language :: Python :: 3.9 Requires-Python: >=3.8,<3.15 Description-Content-Type: text/markdown License-File: LICENSE -Requires-Dist: flwr~=1.18.0 -Requires-Dist: matplotlib~=3.10.8 -Requires-Dist: numpy~=1.26.4 -Requires-Dist: opacus~=1.5.3 -Requires-Dist: pandas~=2.3.1 -Requires-Dist: PyYAML~=6.0 -Requires-Dist: setuptools~=68.0.0 -Requires-Dist: Sphinx~=5.3.0 -Requires-Dist: SQLAlchemy~=1.4.47 -Requires-Dist: torch>=2.7.0 -Requires-Dist: datetime~=5.1 -Requires-Dist: scikit-learn~=1.8.0 -Requires-Dist: sphinx-jsonschema==1.19.1 -Requires-Dist: sphinx-rtd-dark-mode==1.2.4 -Requires-Dist: plotly==5.19.0 -Requires-Dist: optuna==3.5.0 -Requires-Dist: mysql-connector-python~=9.3.0 -Requires-Dist: seaborn~=0.13.2 -Requires-Dist: flwr[simulation] -Dynamic: author -Dynamic: author-email -Dynamic: classifier -Dynamic: description -Dynamic: description-content-type -Dynamic: home-page -Dynamic: keywords -Dynamic: license-file -Dynamic: project-url -Dynamic: requires-dist -Dynamic: requires-python -Dynamic: summary # MEDfl: A Collaborative Framework for Federated Learning in Medicine diff --git a/medfl.egg-info/SOURCES.txt b/medfl.egg-info/SOURCES.txt index 19fd8da..24e708b 100644 --- a/medfl.egg-info/SOURCES.txt +++ b/medfl.egg-info/SOURCES.txt @@ -14,6 +14,19 @@ MEDfl/LearningManager/plot.py MEDfl/LearningManager/server.py MEDfl/LearningManager/strategy.py MEDfl/LearningManager/utils.py +MEDfl/LearningManager/shap/__init__.py +MEDfl/LearningManager/shap/aggregation.py +MEDfl/LearningManager/shap/config.py +MEDfl/LearningManager/shap/explainer.py +MEDfl/LearningManager/shap/xgboost_explainer.py +MEDfl/LearningManager/xgboost/__init__.py +MEDfl/LearningManager/xgboost/client.py +MEDfl/LearningManager/xgboost/config.py +MEDfl/LearningManager/xgboost/data.py +MEDfl/LearningManager/xgboost/metrics.py +MEDfl/LearningManager/xgboost/serialization.py +MEDfl/LearningManager/xgboost/server.py +MEDfl/LearningManager/xgboost/strategy.py MEDfl/NetManager/__init__.py MEDfl/NetManager/database_connector.py MEDfl/NetManager/dataset.py @@ -27,6 +40,13 @@ MEDfl/rw/client.py MEDfl/rw/model.py MEDfl/rw/server.py MEDfl/rw/strategy.py +MEDfl/rw/xgboost/__init__.py +MEDfl/rw/xgboost/client.py +MEDfl/rw/xgboost/config.py +MEDfl/rw/xgboost/data.py +MEDfl/rw/xgboost/metrics.py +MEDfl/rw/xgboost/serialization.py +MEDfl/rw/xgboost/strategy.py MEDfl/scripts/__init__.py MEDfl/scripts/base.py MEDfl/scripts/create_db.py diff --git a/medfl.egg-info/requires.txt b/medfl.egg-info/requires.txt index 6041958..cddb882 100644 --- a/medfl.egg-info/requires.txt +++ b/medfl.egg-info/requires.txt @@ -1,19 +1,18 @@ -flwr~=1.18.0 -matplotlib~=3.10.8 -numpy~=1.26.4 -opacus~=1.5.3 -pandas~=2.3.1 -PyYAML~=6.0 -setuptools~=68.0.0 -Sphinx~=5.3.0 -SQLAlchemy~=1.4.47 -torch>=2.7.0 -datetime~=5.1 -scikit-learn~=1.8.0 -sphinx-jsonschema==1.19.1 -sphinx-rtd-dark-mode==1.2.4 +flwr[simulation]~=1.18.0 +numpy<2.1,>=1.26 +pandas<2.2,>=1.5 +scikit-learn<1.6,>=1.1 +matplotlib<3.10,>=3.6 +seaborn<0.14,>=0.13 plotly==5.19.0 +torch>=2.0 +opacus<2.0,>=1.5 optuna==3.5.0 -mysql-connector-python~=9.3.0 -seaborn~=0.13.2 -flwr[simulation] +SQLAlchemy<2.1,>=1.4.47 +mysql-connector-python<10.0,>=9.3 +PyYAML<7.0,>=6.0 +Sphinx<6.0,>=5.3 +sphinx-jsonschema==1.19.1 +sphinx-rtd-dark-mode==1.2.4 +xgboost +shap diff --git a/requirements.txt b/requirements.txt index 2250260..7949df5 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,20 +1,42 @@ -flwr~=1.18.0 -matplotlib~=3.10.8 -numpy~=1.26.4 -opacus~=1.5.3 -pandas~=2.3.1 -PyYAML~=6.0 -setuptools~=68.0.0 -Sphinx~=5.3.0 -SQLAlchemy~=1.4.47 -torch>=2.7.0 -datetime~=5.1 -scikit-learn~=1.8.0 -sphinx-jsonschema==1.19.1 -sphinx-rtd-dark-mode==1.2.4 +# ========================= +# MEDfl FULL requirements +# (runtime + docs) +# Compatible with Python 3.9 +# ========================= + +# ---- Flower (keep only one; simulation includes extra deps) +flwr[simulation]~=1.18.0 + +# ---- Core scientific +numpy>=1.26,<2.1 +pandas>=1.5,<2.2 +scikit-learn>=1.1,<1.6 + +# ---- Visualization +matplotlib>=3.6,<3.10 +seaborn>=0.13,<0.14 plotly==5.19.0 + +# ---- Deep learning + DP +torch>=2.0 +opacus>=1.5,<2.0 + +# ---- Optimization optuna==3.5.0 -mysql-connector-python~=9.3.0 -seaborn~=0.13.2 -flwr[simulation] +# ---- DB +SQLAlchemy>=1.4.47,<2.1 +mysql-connector-python>=9.3,<10.0 + +# ---- Utils +PyYAML>=6.0,<7.0 + +# ========================= +# Docs dependencies (normally should be optional) +# ========================= +Sphinx>=5.3,<6.0 +sphinx-jsonschema==1.19.1 +sphinx-rtd-dark-mode==1.2.4 + +xgboost +shap \ No newline at end of file diff --git a/setup.py b/setup.py index 5b3f523..acfbc4b 100644 --- a/setup.py +++ b/setup.py @@ -14,7 +14,7 @@ setup( name="medfl", - version="2.0.5.dev9", + version="2.0.10.dev5", author="MEDomics consortium", author_email="medomics.info@gmail.com", description="Python Open-source package for simulating federated learning and differential privacy",