Rw version

MEDfl/LearningManager/client.py

@@ -1,150 +1,169 @@
 #!/usr/bin/env python3
+import logging
+import sys
+
 import flwr as fl
 from opacus import PrivacyEngine
 from torch.utils.data import DataLoader
 
 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(),
                 optimizer=self.local_model.optimizer,
                 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()}

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()