Fix strategy API for aggregate_fit()

baselines/flanders/flanders/attacks.py

@@ -218,16 +218,18 @@ def fang_attack(
     # Set corrupted clients' updates to w_1
     results = [
         (
-            proxy,
-            FitRes(
-                fitres.status,
-                parameters=ndarrays_to_parameters(corrupted_params),
-                num_examples=fitres.num_examples,
-                metrics=fitres.metrics,
-            ),
+            (
+                proxy,
+                FitRes(
+                    fitres.status,
+                    parameters=ndarrays_to_parameters(corrupted_params),
+                    num_examples=fitres.num_examples,
+                    metrics=fitres.metrics,
+                ),
+            )
+            if states[fitres.metrics["cid"]]
+            else (proxy, fitres)
         )
-        if states[fitres.metrics["cid"]]
-        else (proxy, fitres)
         for proxy, fitres in ordered_results
     ]
 

baselines/flanders/flanders/client.py

@@ -1,4 +1,5 @@
 """Clients implementation for Flanders."""
+
 from collections import OrderedDict
 from pathlib import Path
 from typing import Tuple
@@ -80,7 +81,7 @@ def fit(self, parameters, config):
         return (
             get_params(self.net),
             len(trainloader.dataset),
-            {"cid": self.cid},
+            {"cid": self.cid, "malicious": config["malicious"]},
         )
 
     def evaluate(self, parameters, config):

baselines/flanders/flanders/dataset.py

@@ -151,8 +151,8 @@ def mnist_transformation(img):
 class TorchVisionFL(VisionDataset):
     """TorchVision FL class.
 
-    Use this class by either passing a path to a torch file (.pt) containing
-    (data, targets) or pass the data, targets directly instead.
+    Use this class by either passing a path to a torch file (.pt) containing (data,
+    targets) or pass the data, targets directly instead.
 
     This is just a trimmed down version of torchvision.datasets.MNIST.
     """

baselines/flanders/flanders/main.py

@@ -22,7 +22,7 @@
 from .utils import fmnist_evaluate, l2_norm, mnist_evaluate
 
 
-# pylint: disable=too-many-locals
+# pylint: disable=too-many-locals, too-many-branches, too-many-statements
 @hydra.main(config_path="conf", config_name="base", version_base=None)
 def main(cfg: DictConfig) -> None:
     """Run the baseline.
@@ -36,8 +36,8 @@ def main(cfg: DictConfig) -> None:
     seed = cfg.seed
     np.random.seed(seed)
     np.random.set_state(
-        np.random.RandomState(seed).get_state()
-    )  # pylint: disable=no-member
+        np.random.RandomState(seed).get_state()  # pylint: disable=no-member
+    )
     random.seed(seed)
     torch.manual_seed(seed)
     torch.cuda.manual_seed(seed)

baselines/flanders/flanders/models.py

@@ -1,4 +1,5 @@
 """Models for FLANDERS experiments."""
+
 import itertools
 
 import torch

baselines/flanders/flanders/server.py

@@ -221,11 +221,15 @@ def fit_round(
         )
 
         # Create dict clients_state to keep track of malicious clients
+        # and send the information to the clients
         clients_state = {}
-        for _, (proxy, _) in enumerate(client_instructions):
+        for _, (proxy, ins) in enumerate(client_instructions):
             clients_state[proxy.cid] = False
+            ins.config["malicious"] = False
             if proxy.cid in self.malicious_lst:
                 clients_state[proxy.cid] = True
+                ins.config["malicious"] = True
+
         # Sort clients states
         clients_state = {k: clients_state[k] for k in sorted(clients_state)}
         log(
@@ -326,20 +330,18 @@ def fit_round(
         else:
             results = ordered_results
 
+        # Aggregate training results
+        log(INFO, "fit_round - Aggregating training results")
         good_clients_idx = []
         malicious_clients_idx = []
+        aggregated_result = self.strategy.aggregate_fit(
+            server_round, results, failures
+        )
         if isinstance(self.strategy, Flanders):
-            # Aggregate training results
-            log(INFO, "fit_round - Aggregating training results")
-            aggregated_result = self.strategy.aggregate_fit(
-                server_round, results, failures, clients_state
-            )
-            (
-                parameters_aggregated,
-                metrics_aggregated,
-                good_clients_idx,
-                malicious_clients_idx,
-            ) = aggregated_result
+            parameters_aggregated, metrics_aggregated = aggregated_result
+            malicious_clients_idx = metrics_aggregated["malicious_clients_idx"]
+            good_clients_idx = metrics_aggregated["good_clients_idx"]
+
             log(INFO, "Malicious clients: %s", malicious_clients_idx)
 
             log(INFO, "clients_state: %s", clients_state)
@@ -376,9 +378,6 @@ def fit_round(
         else:
             # Aggregate training results
             log(INFO, "fit_round - Aggregating training results")
-            aggregated_result = self.strategy.aggregate_fit(
-                server_round, results, failures
-            )
             parameters_aggregated, metrics_aggregated = aggregated_result
 
         self.clients_state = clients_state

baselines/flanders/flanders/strategies/dnc.py

@@ -101,9 +101,11 @@ def configure_fit(
         fit_ins_list = [
             FitIns(
                 parameters,
-                {}
-                if not self.on_fit_config_fn
-                else self.on_fit_config_fn(server_round),
+                (
+                    {}
+                    if not self.on_fit_config_fn
+                    else self.on_fit_config_fn(server_round)
+                ),
             )
             for _ in range(sample_size)
         ]

baselines/flanders/flanders/strategy.py

@@ -162,9 +162,11 @@ def configure_fit(
         fit_ins_list = [
             FitIns(
                 parameters,
-                {}
-                if not self.on_fit_config_fn
-                else self.on_fit_config_fn(server_round),
+                (
+                    {}
+                    if not self.on_fit_config_fn
+                    else self.on_fit_config_fn(server_round)
+                ),
             )
             for _ in range(sample_size)
         ]
@@ -186,7 +188,6 @@ def aggregate_fit(
         server_round: int,
         results: List[Tuple[ClientProxy, FitRes]],
         failures: List[Union[Tuple[ClientProxy, FitRes], BaseException]],
-        clients_state: Dict[str, bool],
     ) -> Tuple[Optional[Parameters], Dict[str, Scalar]]:
         """Apply MAR forecasting to exclude malicious clients from FedAvg.
 
@@ -269,7 +270,11 @@ def aggregate_fit(
         # Check that self.aggregate_fn has num_malicious parameter
         if "num_malicious" in self.aggregate_fn.__code__.co_varnames:
             # Count the number of malicious clients in
-            # good_clients_idx by checking clients_state
+            # good_clients_idx by checking FitRes
+            clients_state = {
+                str(fit_res.metrics["cid"]): fit_res.metrics["malicious"]
+                for _, fit_res in results
+            }
             num_malicious = sum([clients_state[str(cid)] for cid in good_clients_idx])
             log(
                 INFO,
@@ -313,12 +318,11 @@ def aggregate_fit(
         elif server_round == 1:  # Only log this warning once
             log(WARNING, "No fit_metrics_aggregation_fn provided")
 
-        return (
-            parameters_aggregated,
-            metrics_aggregated,
-            good_clients_idx,
-            malicious_clients_idx,
-        )
+        # Add good_clients_idx and malicious_clients_idx to metrics_aggregated
+        metrics_aggregated["good_clients_idx"] = good_clients_idx
+        metrics_aggregated["malicious_clients_idx"] = malicious_clients_idx
+
+        return parameters_aggregated, metrics_aggregated
 
 
 # pylint: disable=too-many-locals, too-many-arguments, invalid-name