Fixed type checking issues in examples/three_layer_fl, and in

examples/unlearning.

Author: baochunli

examples/three_layer_fl/fedsaw/fedsaw_algorithm.py

@@ -47,7 +47,11 @@ def prune_weight_updates(
         )
 
         # Clone the reference model to host update tensors for pruning.
-        delta_model = copy.deepcopy(self.model.cpu())
+        model = self.model
+        if model is None:
+            raise RuntimeError("Model must be initialised before pruning updates.")
+
+        delta_model = copy.deepcopy(model).cpu()
         cpu_updates: MutableMapping[str, torch.Tensor] = OrderedDict(
             (name, tensor.detach().cpu()) for name, tensor in updates.items()
         )

examples/three_layer_fl/fedsaw/fedsaw_client.py

@@ -2,8 +2,12 @@
 A federated learning client using pruning.
 """
 
+from __future__ import annotations
+
 import copy
 import logging
+from collections.abc import Mapping
+from typing import Any
 
 from fedsaw_algorithm import Algorithm as FedSawAlgorithm
 
@@ -20,10 +24,12 @@ class FedSawClientLifecycleStrategy(DefaultLifecycleStrategy):
     _STATE_KEY = "fedsaw_client"
 
     @staticmethod
-    def _state(context):
+    def _state(context: ClientContext) -> dict[str, Any]:
         return context.state.setdefault(FedSawClientLifecycleStrategy._STATE_KEY, {})
 
-    def process_server_response(self, context, server_response):
+    def process_server_response(
+        self, context: ClientContext, server_response: dict[str, Any]
+    ) -> None:
         super().process_server_response(context, server_response)
         amount = server_response.get("pruning_amount")
         if amount is None:
@@ -33,17 +39,17 @@ def process_server_response(self, context, server_response):
         state["pruning_amount"] = amount
 
         owner = context.owner
-        if owner is not None:
-            owner.pruning_amount = amount
+        if isinstance(owner, FedSawClient) and isinstance(amount, (int, float)):
+            owner.pruning_amount = float(amount)
 
 
 class FedSawTrainingStrategy(DefaultTrainingStrategy):
     """Training strategy that prunes local updates before transmission."""
 
     async def train(self, context: ClientContext):
         algorithm = context.algorithm
-        if algorithm is None:
-            raise RuntimeError("Algorithm is required for FedSaw training.")
+        if not isinstance(algorithm, FedSawAlgorithm):
+            raise RuntimeError("FedSaw training requires a FedSaw algorithm instance.")
 
         previous_weights = copy.deepcopy(algorithm.extract_weights())
         report, new_weights = await super().train(context)
@@ -53,25 +59,69 @@ async def train(self, context: ClientContext):
 
         return report, weight_updates
 
-    def _prune_updates(self, context, previous_weights, new_weights):
+    def _prune_updates(
+        self,
+        context: ClientContext,
+        previous_weights: Mapping[str, Any],
+        new_weights: Mapping[str, Any],
+    ):
         algorithm = context.algorithm
+        if not isinstance(algorithm, FedSawAlgorithm):
+            raise RuntimeError("FedSaw algorithm required to prune weight updates.")
+
         updates = algorithm.compute_weight_updates(previous_weights, new_weights)
 
         pruning_method = (
             "random"
             if getattr(Config().clients, "pruning_method", None) == "random"
             else "l1"
         )
-        pruning_amount = getattr(context.owner, "pruning_amount", None)
+        owner = context.owner
+        pruning_amount: float | int | None = None
+        if isinstance(owner, FedSawClient):
+            pruning_amount = owner.pruning_amount
+
         if pruning_amount is None:
             state = FedSawClientLifecycleStrategy._state(context)
-            pruning_amount = state.get("pruning_amount", 0)
+            stored_amount = state.get("pruning_amount", 0)
+            pruning_amount = stored_amount if isinstance(stored_amount, (int, float)) else 0
 
         return algorithm.prune_weight_updates(
             updates, amount=pruning_amount, method=pruning_method
         )
 
 
+class FedSawClient(simple.Client):
+    """Client implementation that tracks pruning metadata for FedSaw."""
+
+    def __init__(
+        self,
+        model=None,
+        datasource=None,
+        algorithm=None,
+        trainer=None,
+        callbacks=None,
+        trainer_callbacks=None,
+    ):
+        super().__init__(
+            model=model,
+            datasource=datasource,
+            algorithm=algorithm or FedSawAlgorithm,
+            trainer=trainer,
+            callbacks=callbacks,
+            trainer_callbacks=trainer_callbacks,
+        )
+        self.pruning_amount: float = 0.0
+
+        self._configure_composable(
+            lifecycle_strategy=FedSawClientLifecycleStrategy(),
+            payload_strategy=self.payload_strategy,
+            training_strategy=FedSawTrainingStrategy(),
+            reporting_strategy=self.reporting_strategy,
+            communication_strategy=self.communication_strategy,
+        )
+
+
 def create_client(
     *,
     model=None,
@@ -80,28 +130,17 @@ def create_client(
     trainer=None,
     callbacks=None,
     trainer_callbacks=None,
-):
+) -> FedSawClient:
     """Build a FedSaw client that prunes its updates before reporting."""
-    client = simple.Client(
+    return FedSawClient(
         model=model,
         datasource=datasource,
-        algorithm=algorithm or FedSawAlgorithm,
+        algorithm=algorithm,
         trainer=trainer,
         callbacks=callbacks,
         trainer_callbacks=trainer_callbacks,
     )
-    client.pruning_amount = 0
-
-    client._configure_composable(
-        lifecycle_strategy=FedSawClientLifecycleStrategy(),
-        payload_strategy=client.payload_strategy,
-        training_strategy=FedSawTrainingStrategy(),
-        reporting_strategy=client.reporting_strategy,
-        communication_strategy=client.communication_strategy,
-    )
-
-    return client
 
 
 # Maintain compatibility for imports expecting a Client callable/class.
-Client = create_client
+Client = FedSawClient

examples/three_layer_fl/fedsaw/fedsaw_server.py

@@ -3,6 +3,8 @@
 as either central or edge servers.
 """
 
+from __future__ import annotations
+
 import math
 import statistics
 
@@ -20,7 +22,7 @@ def __init__(self, model=None, algorithm=None, trainer=None):
         super().__init__(model=model, algorithm=selected_algorithm, trainer=trainer)
 
         # The central server uses a list to store each edge server's clients' pruning amount
-        self.pruning_amount_list = None
+        self.pruning_amount_list: dict[int, float] | None = None
 
         if Config().is_central_server():
             init_pruning_amount = (
@@ -36,8 +38,9 @@ def __init__(self, model=None, algorithm=None, trainer=None):
                 )
             }
 
+        self.edge_pruning_amount: float = 0.0
         if Config().is_edge_server():
-            self.edge_pruning_amount = 0
+            self.edge_pruning_amount = 0.0
 
     def customize_server_response(self, server_response: dict, client_id) -> dict:
         """Wraps up generating the server response with any additional information."""
@@ -55,7 +58,8 @@ def customize_server_response(self, server_response: dict, client_id) -> dict:
     async def aggregate_weights(self, updates, baseline_weights, weights_received):
         """Aggregates the reported weight updates from the selected clients."""
         deltas = await self.aggregate_deltas(updates, weights_received)
-        updated_weights = self.algorithm.update_weights(deltas)
+        algorithm = self.require_algorithm()
+        updated_weights = algorithm.update_weights(deltas)
         return updated_weights
 
     def update_pruning_amount_list(self):
@@ -64,6 +68,11 @@ def update_pruning_amount_list(self):
 
         median = statistics.median(weights_diff_list)
 
+        if self.pruning_amount_list is None:
+            raise RuntimeError(
+                "Pruning amount list is unavailable on this server instance."
+            )
+
         for client_id in weights_diff_dict:
             if weights_diff_dict[client_id]:
                 self.pruning_amount_list[client_id] = 1 / (
@@ -84,13 +93,14 @@ def get_weights_differences(self):
         }
 
         weights_diff_list = []
+        algorithm = self.require_algorithm()
 
         for update in self.updates:
             client_id = update.report.client_id
             num_samples = update.report.num_samples
             received_updates = update.payload
 
-            weights_diff = self.algorithm.compute_weight_difference(
+            weights_diff = algorithm.compute_weight_difference(
                 received_updates,
                 num_samples=num_samples,
                 total_samples=self.total_samples,

examples/three_layer_fl/tempo/tempo_server.py

@@ -3,6 +3,8 @@
 clients' local epoch numbers of each edge server (institution).
 """
 
+from __future__ import annotations
+
 import math
 
 import torch
@@ -21,7 +23,7 @@ def __init__(self):
         super().__init__()
 
         # The central server uses a list to store each edge server's clients' local epoch numbers
-        self.local_epoch_list = None
+        self.local_epoch_list: list[int] | None = None
         if Config().is_central_server():
             self.local_epoch_list = [
                 Config().trainer.epochs for i in range(Config().algorithm.total_silos)
@@ -111,10 +113,11 @@ def compute_weights_difference(self, local_weights, num_samples):
         Computes the weight difference of an edge server's aggregated model
         and the global model.
         """
-        weights_diff = 0
+        weights_diff = 0.0
 
         # Extract global model weights
-        global_weights = self.algorithm.extract_weights()
+        algorithm = self.require_algorithm()
+        global_weights = algorithm.extract_weights()
 
         for name, local_weight in local_weights.items():
             global_weight = global_weights[name]

examples/unlearning/knot/solver.py

@@ -6,10 +6,12 @@
 the number of papers
 """
 
+from typing import Any, cast
+
 try:
     import mosek  # type: ignore
 except ImportError:  # pragma: no cover - optional dependency
-    mosek = None
+    mosek = cast(Any, None)
 
 from cvxopt import matrix, solvers, sparse, spmatrix
 

plato/config.py

@@ -473,14 +473,19 @@ def simulate_client_speed() -> np.ndarray:
     @staticmethod
     def is_edge_server() -> bool:
         """Returns whether the current instance is an edge server in cross-silo FL."""
-        return Config().args.port is not None and bool(
-            getattr(Config().algorithm, "cross_silo", False)
-        )
+        if not bool(getattr(Config().algorithm, "cross_silo", False)):
+            return False
+
+        args = Config().args
+        return args.port is not None and args.id is not None
 
     @staticmethod
     def is_central_server() -> bool:
         """Returns whether the current instance is a central server in cross-silo FL."""
-        return hasattr(Config().algorithm, "cross_silo") and Config().args.port is None
+        if not bool(getattr(Config().algorithm, "cross_silo", False)):
+            return False
+
+        return Config().args.id is None
 
     @staticmethod
     def gpu_count() -> int: