[pre-commit.ci] Add auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Author: pre-commit-ci[bot]

fl4health/clients/flexible_client.py

@@ -6,9 +6,9 @@
 from typing import Any
 
 import torch
-import torch.nn as nn
 from flwr.common.logger import log
 from flwr.common.typing import Config, NDArrays, Scalar
+from torch import nn
 from torch.nn.modules.loss import _Loss
 from torch.optim import Optimizer
 from torch.optim.lr_scheduler import LRScheduler
@@ -74,7 +74,6 @@ def __init__(
                 If not passed, a hash is randomly generated. Client state will use this as part of its state file
                 name. Defaults to None.
         """
-
         self.data_path = data_path
         self.device = device
         self.metrics = metrics
@@ -202,12 +201,11 @@ def get_parameters(self, config: Config) -> NDArrays:
 
             # Need all parameters even if normally exchanging partial
             return FullParameterExchanger().push_parameters(self.model, config=config)
-        else:
-            assert self.model is not None and self.parameter_exchanger is not None
-            # If the client has early stopping module and the patience is None, we load the best saved state
-            # to send the best checkpointed local model's parameters to the server
-            self._maybe_load_saved_best_local_model_state()
-            return self.parameter_exchanger.push_parameters(self.model, config=config)
+        assert self.model is not None and self.parameter_exchanger is not None
+        # If the client has early stopping module and the patience is None, we load the best saved state
+        # to send the best checkpointed local model's parameters to the server
+        self._maybe_load_saved_best_local_model_state()
+        return self.parameter_exchanger.push_parameters(self.model, config=config)
 
     def _maybe_load_saved_best_local_model_state(self) -> None:
         if self.early_stopper is not None and self.early_stopper.patience is None:
@@ -280,7 +278,7 @@ def process_config(self, config: Config) -> tuple[int | None, int | None, int, b
         # Parse config to determine train by steps or train by epochs
         if ("local_epochs" in config) and ("local_steps" in config):
             raise ValueError("Config cannot contain both local_epochs and local_steps. Please specify only one.")
-        elif "local_epochs" in config:
+        if "local_epochs" in config:
             local_epochs = narrow_dict_type(config, "local_epochs", int)
             local_steps = None
         elif "local_steps" in config:
@@ -479,7 +477,6 @@ def _log_header_str(
             current_epoch (int | None, optional): The current epoch of local
                 training. Defaults to None.
         """
-
         log_str = f"Current FL Round: {int(current_round)} " if current_round is not None else ""
         log_str += f"Current Epoch: {int(current_epoch)} " if current_epoch is not None else ""
 
@@ -591,7 +588,8 @@ def update_metric_manager(
     def _compute_preds_and_losses(
         self, model: nn.Module, optimizer: Optimizer, input: TorchInputType, target: TorchTargetType
     ) -> tuple[TrainingLosses, TorchPredType]:
-        """Helper method within the train step for computing preds and losses.
+        """
+        Helper method within the train step for computing preds and losses.
 
         NOTE: Subclasses should implement this helper method if there is a need
         to specialize this part of the overall train step.
@@ -620,7 +618,8 @@ def _compute_preds_and_losses(
     def _apply_backwards_on_losses_and_take_step(
         self, model: nn.Module, optimizer: Optimizer, losses: TrainingLosses
     ) -> TrainingLosses:
-        """Helper method within the train step for applying backwards on losses and taking step with optimizer.
+        """
+        Helper method within the train step for applying backwards on losses and taking step with optimizer.
 
         NOTE: Subclasses should implement this helper method if there is a need
         to specialize this part of the overall train step.
@@ -643,7 +642,8 @@ def _apply_backwards_on_losses_and_take_step(
     def _train_step_with_model_and_optimizer(
         self, model: torch.nn.Module, optimizer: Optimizer, input: TorchInputType, target: TorchTargetType
     ) -> tuple[TrainingLosses, TorchPredType]:
-        """Helper train step method that allows for injection of model and optimizer.
+        """
+        Helper train step method that allows for injection of model and optimizer.
 
         NOTE: Subclasses should implement this method if there is a need to specialize
         the train_step logic.
@@ -656,7 +656,6 @@ def _train_step_with_model_and_optimizer(
             tuple[TrainingLosses, TorchPredType]: The losses object from the train step along with
             a dictionary of any predictions produced by the model.
         """
-
         losses, preds = self._compute_preds_and_losses(model, optimizer, input, target)
         losses = self._apply_backwards_on_losses_and_take_step(model, optimizer, losses)
 
@@ -681,7 +680,8 @@ def train_step(self, input: TorchInputType, target: TorchTargetType) -> tuple[Tr
     def _val_step_with_model(
         self, model: nn.Module, input: TorchInputType, target: TorchTargetType
     ) -> tuple[EvaluationLosses, TorchPredType]:
-        """Helper method for val_step that allows for injection of model.
+        """
+        Helper method for val_step that allows for injection of model.
 
         NOTE: Subclasses should implement this method if there is a need to
         specialize the val_step logic.
@@ -694,7 +694,6 @@ def _val_step_with_model(
             tuple[EvaluationLosses, TorchPredType]: The losses object from the val step along with a dictionary of the
             predictions produced by the model.
         """
-
         # Get preds and compute loss
         with torch.no_grad():
             preds, features = self.predict_with_model(model, input)
@@ -715,7 +714,6 @@ def val_step(self, input: TorchInputType, target: TorchTargetType) -> tuple[Eval
             tuple[EvaluationLosses, TorchPredType]: The losses object from the val step along with a dictionary of the
             predictions produced by the model.
         """
-
         return self._val_step_with_model(self.model, input, target)
 
     def train_by_epochs(
@@ -874,7 +872,6 @@ def _validate_by_steps(
         Returns:
             tuple[float, dict[str, Scalar]]: The loss and a dictionary of metrics from evaluation.
         """
-
         assert self.num_validation_steps is not None, "num_validation_steps must be defined to use this function"
 
         self.model.eval()
@@ -1047,9 +1044,9 @@ def setup_client(self, config: Config) -> None:
         # batch_size * num_validation_steps
         self.num_val_samples = len(self.val_loader.dataset)  # type: ignore
         if self.num_validation_steps is not None:
-            assert (
-                self.val_loader.batch_size is not None
-            ), "Validation batch size must be defined if we want to limit the number of validation steps"
+            assert self.val_loader.batch_size is not None, (
+                "Validation batch size must be defined if we want to limit the number of validation steps"
+            )
             self.num_val_samples = self.num_validation_steps * self.val_loader.batch_size
 
         if self.test_loader:
@@ -1087,7 +1084,8 @@ def get_parameter_exchanger(self, config: Config) -> ParameterExchanger:
     def predict_with_model(
         self, model: torch.nn.Module, input: TorchInputType
     ) -> tuple[TorchPredType, TorchFeatureType]:
-        """Helper predict method that allows for injection of model.
+        """
+        Helper predict method that allows for injection of model.
 
         NOTE: Subclasses should implement this method if there is need to specialize
         the predict logic of the client.
@@ -1110,7 +1108,6 @@ def predict_with_model(
             ValueError: Occurs when something other than a tensor or dict of tensors is returned by the model
                 forward.
         """
-
         if isinstance(input, torch.Tensor):
             output = model(input)
         elif isinstance(input, dict):
@@ -1123,15 +1120,14 @@ def predict_with_model(
 
         if isinstance(output, dict):
             return output, {}
-        elif isinstance(output, torch.Tensor):
+        if isinstance(output, torch.Tensor):
             return {"prediction": output}, {}
-        elif isinstance(output, tuple):
+        if isinstance(output, tuple):
             if len(output) != 2:
                 raise ValueError(f"Output tuple should have length 2 but has length {len(output)}")
             preds, features = output
             return preds, features
-        else:
-            raise ValueError("Model forward did not return a tensor, dictionary of tensors, or tuple of tensors")
+        raise ValueError("Model forward did not return a tensor, dictionary of tensors, or tuple of tensors")
 
     def compute_loss_and_additional_losses(
         self, preds: TorchPredType, features: TorchFeatureType, target: TorchTargetType
@@ -1328,7 +1324,6 @@ def update_lr_schedulers(self, step: int | None = None, epoch: int | None = None
             step (int | None): If using ``local_steps``, current step of this round. Otherwise None.
             epoch (int | None): If using ``local_epochs`` current epoch of this round. Otherwise None.
         """
-
         assert (step is None) ^ (epoch is None)
 
         for lr_scheduler in self.lr_schedulers.values():
@@ -1396,7 +1391,8 @@ def update_before_epoch(self, epoch: int) -> None:
         pass
 
     def _transform_gradients_with_model(self, model: torch.nn.Module, losses: TrainingLosses) -> None:
-        """Helper transform gradients method that allows for injection of model.
+        """
+        Helper transform gradients method that allows for injection of model.
 
         NOTE: Subclasses should implement this helper should there be a need to specialize the logic
         for transforming gradients.
@@ -1405,7 +1401,6 @@ def _transform_gradients_with_model(self, model: torch.nn.Module, losses: Traini
             model (torch.nn.Module): the model used to generate predictions to compute losses
             losses (TrainingLosses): The losses object from the train step
         """
-
         pass
 
     def transform_gradients(self, losses: TrainingLosses) -> None:

fl4health/mixins/adaptive_drift_constrained.py

@@ -143,7 +143,6 @@ def set_parameters(
     def train_step(
         self: AdaptiveDriftConstrainedProtocol, input: TorchInputType, target: TorchTargetType
     ) -> tuple[TrainingLosses, TorchPredType]:
-
         losses, preds = self._compute_preds_and_losses(self.model, self.optimizers["global"], input, target)
         loss_clone = losses.backward["backward"].clone()
 
@@ -207,7 +206,8 @@ def compute_penalty_loss(self: AdaptiveDriftConstrainedProtocol) -> torch.Tensor
 
 
 def apply_adaptive_drift_to_client(client_base_type: type[FlexibleClient]) -> type[FlexibleClient]:
-    """Dynamically create an adapted client class.
+    """
+    Dynamically create an adapted client class.
 
     Args:
         client_base_type (type[FlexibleClient]): The class to be mixed.

fl4health/mixins/personalized/utils.py

@@ -8,7 +8,8 @@
 
 @wrapt.decorator
 def ensure_protocol_compliance(func: Callable, instance: Any | None, args: Any, kwargs: Any) -> Any:
-    """Wrapper to ensure that a the instance is of `FlexibleClient` type.
+    """
+    Wrapper to ensure that a the instance is of `FlexibleClient` type.
 
     NOTE: this should only be used within a `FlexibleClient`.
 

tests/clients/test_flexible_client.py

@@ -24,6 +24,7 @@
 from tests.test_utils.assert_metrics_dict import assert_metrics_dict
 from tests.test_utils.models_for_test import LinearModel
 
+
 freezegun.configure(extend_ignore_list=["transformers"])  # type: ignore
 
 
@@ -249,8 +250,7 @@ def mock_validate_or_test(  # type: ignore
             fold_loss_dict_into_metrics(self.mock_metrics, self.mock_loss_dict, logging_mode)
         if logging_mode == LoggingMode.VALIDATION:
             return self.mock_loss, self.mock_metrics
-        else:
-            return self.mock_loss, self.mock_metrics_test
+        return self.mock_loss, self.mock_metrics_test
 
 
 def test_subclass_raises_warning_if_override_val_step() -> None:

tests/mixins/personalized/test_ditto.py

@@ -49,7 +49,6 @@ class _TestDittoedClient(DittoPersonalizedMixin, _TestFlexibleClient):
 
 
 class _DummyParent:
-
     def __init__(self) -> None:
         pass
 
@@ -81,7 +80,6 @@ class _InvalidTestDittoClient(DittoPersonalizedMixin):
         pass
 
     with pytest.raises(RuntimeError, match="This object needs to satisfy `FlexibleClientProtocolPreSetup`."):
-
         _InvalidTestDittoClient(data_path=Path(""), metrics=[Accuracy()])
 
 
@@ -370,7 +368,6 @@ def test_val_step(
 
 def test_raise_runtime_error_not_flexible_client() -> None:
     """Test that an invalid parent raises RuntimeError."""
-
     with pytest.raises(
         RuntimeError, match=re.escape("This object needs to satisfy `FlexibleClientProtocolPreSetup`.")
     ):

tests/mixins/test_adaptive_drift_constrained.py

@@ -67,7 +67,6 @@ class _InvalidTestAdaptedClient(AdaptiveDriftConstrainedMixin):
         pass
 
     with pytest.raises(RuntimeError, match="This object needs to satisfy `FlexibleClientProtocolPreSetup`."):
-
         _InvalidTestAdaptedClient(data_path=Path(""), metrics=[Accuracy()])