MEDfl 2.0.10.dev5__tar.gz → 2.0.10.dev7__tar.gz

{medfl-2.0.10.dev5 → medfl-2.0.10.dev7}/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 torch.utils.data import TensorDataset, DataLoader
 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 @@ class FlowerClient(fl.client.NumPyClient):
         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 @@ class FlowerClient(fl.client.NumPyClient):
             "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

{medfl-2.0.10.dev5 → medfl-2.0.10.dev7}/MEDfl/rw/server.py

@@ -2,6 +2,7 @@ import flwr as fl
 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 @@ class FederatedServer:
         num_rounds=3,
         strategy=None,
         certificates=None,
+        shap_config=None,
     ):
         """
         Initialize the FederatedServer.
@@ -50,14 +52,88 @@ class FederatedServer:
 
         # 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 SHAP against the active real-world backend.
+        if self.shap_config.enabled:
+            feature_columns = [
+                column.strip()
+                for column in (
+                    getattr(
+                        self.strategy_wrapper,
+                        "_features",
+                        "",
+                    )
+                    or ""
+                ).split(",")
+                if column.strip()
+            ]
+
+            if not feature_columns:
+                raise ValueError(
+                    "Federated SHAP requires the Strategy 'features' "
+                    "configuration so the model input size can be "
+                    "validated."
+                )
+
+            backend = (
+                "xgboost"
+                if getattr(
+                    self.strategy_wrapper,
+                    "name",
+                    "",
+                ) == "FedXgbBagging"
+                else "pytorch"
+            )
+
+            task = (
+                getattr(
+                    self.strategy_wrapper,
+                    "task",
+                    None,
+                )
+                if backend == "xgboost"
+                else None
+            )
+
+            self.shap_config.validate_backend(
+                input_size=len(feature_columns),
+                backend=backend,
+                task=task,
+            )
+
         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 +154,21 @@ class FederatedServer:
             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):
     """
@@ -131,4 +222,4 @@ class TrackingClientManager(fl.server.client_manager.SimpleClientManager):
         #     hostname = f"Error: {e}"
 
         print(f"✅ Client connected - CID: {client.cid}")
-        self.server.connected_clients.append(client.cid)
+        self.server.connected_clients.append(client.cid)

{medfl-2.0.10.dev5 → medfl-2.0.10.dev7}/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 @@ class Strategy:
         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 @@ class Strategy:
         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 @@ class Strategy:
                 "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 @@ class Strategy:
             "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 @@ class Strategy:
         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)