wip

Author: nerdai

examples/nnunet_example/client.py

@@ -28,7 +28,6 @@
 from fl4health.utils.msd_dataset_sources import get_msd_dataset_enum, msd_num_labels
 from fl4health.utils.nnunet_utils import get_segs_from_probs, set_nnunet_env
 
-
 personalized_client_classes = {"ditto": make_it_personal(NnunetClient, "ditto")}
 
 
@@ -107,7 +106,7 @@ def main(
             log(INFO, f"Setting up client without personalization")
             client = NnunetClient(**client_kwargs)
         log(INFO, f"Using client: {type(client).__name__}")
-        log(INFO, f"Parameter exchanger: {type(client.parameter_exchanger).__name__}")
+        # log(INFO, f"Parameter exchanger: {type(client.parameter_exchanger).__name__}")
 
         start_client(server_address=server_address, client=client.to_client())
 

fl4health/mixins/adaptive_drift_contrained.py

@@ -3,7 +3,7 @@
 import warnings
 from collections.abc import Sequence
 from logging import INFO
-from typing import cast
+from typing import Protocol, runtime_checkable
 
 import torch
 from flwr.common.logger import log
@@ -19,13 +19,16 @@
 from fl4health.utils.typing import TorchFeatureType, TorchPredType, TorchTargetType
 
 
-class AdaptiveProtocol(BasicClientProtocol):
+@runtime_checkable
+class AdaptiveProtocol(BasicClientProtocol, Protocol):
     loss_for_adaptation: float | None
     drift_penalty_tensors: list[torch.Tensor] | None
     drift_penalty_weight: float | None
 
     def compute_penalty_loss(self) -> torch.Tensor: ...
 
+    def ensure_protocol_compliance(self) -> None: ...
+
 
 class AdaptiveDriftConstrainedMixin:
     def __init_subclass__(cls, **kwargs):
@@ -44,15 +47,10 @@ def __init_subclass__(cls, **kwargs):
             RuntimeWarning,
         )
 
-    def __init__(self, *args, **kwargs):
-        # Verify at instance creation time
+    def ensure_protocol_compliance(self) -> None:
+        """Call this after the object is fully initialized"""
         if not isinstance(self, BasicClientProtocol):
-            raise TypeError(
-                f"Class {self.__class__.__name__} uses AdaptiveMixin but does not "
-                f"implement BasicClientProtocol. Make sure a parent class implements "
-                f"the required methods and attributes."
-            )
-        super().__init__(*args, **kwargs)
+            raise TypeError(f"Protocol requirements not met.")
 
     def penalty_loss_function(self: AdaptiveProtocol) -> WeightDriftLoss:
         """Function to compute the penalty loss."""

fl4health/mixins/personalized/ditto.py

@@ -1,26 +1,26 @@
 """Ditto Personalized Mixin"""
 
+import warnings
 from logging import DEBUG, INFO
-from typing import cast, runtime_checkable
+from typing import Protocol, cast, runtime_checkable
 
 import torch
 import torch.nn as nn
-import warnings
 from flwr.common.logger import log
 from flwr.common.typing import Config, NDArrays, Scalar
 from torch.optim import Optimizer
 
+from fl4health.clients.basic_client import BasicClientProtocol
 from fl4health.mixins.adaptive_drift_contrained import AdaptiveDriftConstrainedMixin, AdaptiveProtocol
 from fl4health.mixins.personalized.base import BasePersonalizedMixin
 from fl4health.parameter_exchange.full_exchanger import FullParameterExchanger
 from fl4health.utils.config import narrow_dict_type
 from fl4health.utils.losses import EvaluationLosses, TrainingLosses
 from fl4health.utils.typing import TorchFeatureType, TorchInputType, TorchPredType, TorchTargetType
-from fl4health.clients.basic_client import BasicClientProtocol
 
 
 @runtime_checkable
-class DittoProtocol(AdaptiveProtocol):
+class DittoProtocol(AdaptiveProtocol, Protocol):
     global_model: torch.nn.Module | None
 
     def get_global_model(self, config: Config) -> nn.Module: ...
@@ -50,16 +50,6 @@ def __init_subclass__(cls, **kwargs):
             RuntimeWarning,
         )
 
-    def __init__(self, *args, **kwargs):
-        # Verify at instance creation time
-        if not isinstance(self, BasicClientProtocol):
-            raise TypeError(
-                f"Class {self.__class__.__name__} uses AdaptiveMixin but does not "
-                f"implement BasicClientProtocol. Make sure a parent class implements "
-                f"the required methods and attributes."
-            )
-        super().__init__(*args, **kwargs)
-
     @property
     def optimizer_keys(self: DittoProtocol) -> list[str]:
         """Returns the optimizer keys."""
@@ -143,6 +133,8 @@ def setup_client(self: DittoProtocol, config: Config) -> None:
         Args:
             config (Config): The config from the server.
         """
+        self.ensure_protocol_compliance()
+
         try:
             self.global_model = self.get_global_model(config)
             log(INFO, f"global model set: {type(self.global_model).__name__}")
@@ -254,6 +246,7 @@ def update_before_train(self: DittoProtocol, current_server_round: int) -> None:
         Args:
             current_server_round (int): Indicates which server round we are currently executing.
         """
+        self.ensure_protocol_compliance()
         self.set_initial_global_tensors()
 
         # Need to also set the global model to train mode before any training begins.
@@ -479,6 +472,7 @@ def compute_evaluation_loss(
             EvaluationLosses: An instance of ``EvaluationLosses`` containing checkpoint loss and additional losses
             indexed by name.
         """
+        self.ensure_protocol_compliance()
         # Check that both models are in eval mode
         assert not self.global_model is None and not self.global_model.training and not self.model.training
         return super().compute_evaluation_loss(preds, features, target)