Merge pull request #405 from VectorInstitute/dbe/loosening_some_ruff_ignores

Loosening Some Ruff Ignores

Author: emersodb

examples/fedllm_example/model.py

@@ -7,6 +7,10 @@
 from transformers import AutoModelForCausalLM, BitsAndBytesConfig
 
 
+EIGHT_BIT = 8
+FOUR_BIT = 4
+
+
 def cosine_annealing(
     total_round: int,
     current_round: int = 0,
@@ -48,9 +52,9 @@ def get_model(model_cfg: dict[str, Any]) -> torch.nn.Module:
         torch.nn.Module: The model.
     """
     quantization_config = model_cfg["quantization"]
-    if quantization_config == 4:
+    if quantization_config == FOUR_BIT:
         quantization_config = BitsAndBytesConfig(load_in_4bit=True)
-    elif quantization_config == 8:
+    elif quantization_config == EIGHT_BIT:
         quantization_config = BitsAndBytesConfig(load_in_8bit=True)
     else:
         raise ValueError(f"Use 4-bit or 8-bit quantization. You passed: {quantization_config}/")

fl4health/checkpointing/state_checkpointer.py

@@ -488,13 +488,13 @@ def maybe_load_server_state(
         Load the state of the server from checkpoint.
 
         Args:
-            server (FlServer): server into which the attributes will be loaded
-            model nn.Module: The model structure to be loaded as part of the server state.
-            attributes (list[str] | None): List of attributes to load from the checkpoint. If None, all attributes
-                specified in ``snapshot_attrs`` are loaded. Defaults to None.
+            server (FlServer): Server into which the attributes will be loaded.
+            model (nn.Module): The model structure to be loaded as part of the server state.
+            attributes (list[str] | None, optional): List of attributes to load from the checkpoint. If None, all
+                attributes specified in ``snapshot_attrs`` are loaded. Defaults to None.
 
         Returns:
-            nn.Module | None: Returns a model if a checkpoint exists to load from. Otherwise returns None
+            nn.Module | None: Returns a model if a checkpoint exists to load from. Otherwise returns None.
         """
         # Store server for access in functions
         self.server = server

fl4health/clients/basic_client.py

@@ -37,6 +37,9 @@
 from fl4health.utils.typing import LogLevel, TorchFeatureType, TorchInputType, TorchPredType, TorchTargetType
 
 
+EXPECTED_OUTPUT_TUPLE_SIZE = 2
+
+
 class BasicClient(NumPyClient):
     def __init__(
         self,
@@ -140,8 +143,10 @@ def _maybe_checkpoint(self, loss: float, metrics: dict[str, Scalar], checkpoint_
         If checkpointer exists, maybe checkpoint model based on the provided metric values.
 
         Args:
-            loss (float): validation loss to potentially be used for checkpointing
-            metrics (dict[str, float]): validation metrics to potentially be used for checkpointing
+            loss (float): Validation loss to potentially be used for checkpointing.
+            metrics (dict[str, Scalar]): Validation metrics to potentially be used for checkpointing
+            checkpoint_mode (CheckpointMode): Whether we're doing checkpointing pre- or post-aggregation on the server
+                side.
         """
         self.checkpoint_and_state_module.maybe_checkpoint(self.model, loss, metrics, checkpoint_mode)
 
@@ -439,10 +444,10 @@ def _log_header_str(
         the round if training by steps.
 
         Args:
-            current_round (int | None, optional): The current FL round. (Ie current
-                server round). Defaults to None.
-            current_epoch (int | None, optional): The current epoch of local
-                training. Defaults to None.
+            current_round (int | None, optional): The current FL round. (Ie current server round). Defaults to None.
+            current_epoch (int | None, optional): The current epoch of local training. Defaults to None.
+            logging_mode (LoggingMode, optional): The logging mode to be used in logging. This mainly changes the
+                way in which logging is decorated. Defaults to LoggingMode.TRAIN.
         """
         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 ""
@@ -1002,7 +1007,7 @@ def predict(self, input: TorchInputType) -> tuple[TorchPredType, TorchFeatureTyp
         if isinstance(output, torch.Tensor):
             return {"prediction": output}, {}
         if isinstance(output, tuple):
-            if len(output) != 2:
+            if len(output) != EXPECTED_OUTPUT_TUPLE_SIZE:
                 raise ValueError(f"Output tuple should have length 2 but has length {len(output)}")
             preds, features = output
             return preds, features

fl4health/clients/constrained_fenda_client.py

@@ -152,6 +152,7 @@ def update_after_train(self, local_steps: int, loss_dict: dict[str, float], conf
         Args:
             local_steps (int): Number of steps performed during training
             loss_dict (dict[str, float]): Losses computed during training.
+            config (Config): FL training configuration object.
         """
         # Save the parameters of the old model
         assert isinstance(self.model, FendaModelWithFeatureState)

fl4health/clients/evaluate_client.py

@@ -45,7 +45,8 @@ def __init__(
                 "cuda"
             loss_meter_type (LossMeterType, optional): Type of meter used to track and compute the losses over
                 each batch. Defaults to ``LossMeterType.AVERAGE``.
-            model_checkpoint_path (Path | None, optional): _description_. Defaults to None.
+            model_checkpoint_path (Path | None, optional): Path to which the model should be checkpointed. Defaults to
+                None.
             reporters (Sequence[BaseReporter] | None, optional): A sequence of FL4Health reporters which the client
                 should send data to. Defaults to None.
             client_name (str | None, optional): An optional client name that uniquely identifies a client.

fl4health/clients/fedrep_client.py

@@ -122,7 +122,7 @@ def _prefix_loss_and_metrics_dictionaries(
         Args:
             prefix (str): Prefix to be attached to all keys of the provided dictionaries.
             loss_dict (dict[str, float]): Dictionary of loss values obtained during training.
-            metrics (dict[str, Scalar]): Dictionary of metrics values measured during training
+            metrics_dict (dict[str, Scalar]): Dictionary of metrics values measured during training.
         """
         for loss_key in list(loss_dict):
             loss_dict[f"{prefix}_{loss_key}"] = loss_dict.pop(loss_key)
@@ -317,12 +317,13 @@ def train_fedrep_by_epochs(
         Train locally for the specified number of epochs.
 
         Args:
-            epochs (int): The number of epochs for local training.
-            current_round (int | None): The current FL round.
+            head_epochs (int): The number of epochs for local training of the head module.
+            rep_epochs (int): The number of epochs for local training of the representation module
+            current_round (int | None, optional): The current FL round. Defaults to None.
 
         Returns:
             tuple[dict[str, float], dict[str, Scalar]]: The loss and metrics dictionary from the local training.
-            Loss is a dictionary of one or more losses that represent the different components of the loss.
+                Loss is a dictionary of one or more losses that represent the different components of the loss.
         """
         # First we train the head module for head_epochs with the representations frozen in place
         self._prepare_train_head()
@@ -356,11 +357,13 @@ def train_fedrep_by_steps(
         Train locally for the specified number of steps.
 
         Args:
-            steps (int): The number of steps to train locally.
+            head_steps (int): The number of steps to train locally for the head model.
+            rep_steps (int): The number of steps to train locally for the representation model
+            current_round (int | None, optional): What round of FL training we're currently on. Defaults to None.
 
         Returns:
             tuple[dict[str, float], dict[str, Scalar]]: The loss and metrics dictionary from the local training.
-            Loss is a dictionary of one or more losses that represent the different components of the loss.
+                Loss is a dictionary of one or more losses that represent the different components of the loss.
         """
         assert isinstance(self.model, FedRepModel)
         # First we train the head module for head_steps with the representations frozen in place

fl4health/clients/flexible_client.py

@@ -17,6 +17,9 @@
 from fl4health.utils.typing import TorchFeatureType, TorchInputType, TorchPredType, TorchTargetType
 
 
+EXPECTED_OUTPUT_TUPLE_SIZE = 2
+
+
 class FlexibleClient(BasicClient):
     def __init__(
         self,
@@ -296,7 +299,7 @@ def predict_with_model(
         if isinstance(output, torch.Tensor):
             return {"prediction": output}, {}
         if isinstance(output, tuple):
-            if len(output) != 2:
+            if len(output) != EXPECTED_OUTPUT_TUPLE_SIZE:
                 raise ValueError(f"Output tuple should have length 2 but has length {len(output)}")
             preds, features = output
             return preds, features

fl4health/clients/nnunet_client.py

@@ -138,6 +138,8 @@ def __init__(
                 No checkpointing (state or model) is done if not provided. Defaults to None.
             reporters (Sequence[BaseReporter], optional): A sequence of FL4Health reporters which the client should
                 send data to.
+            client_name (str | None, optional): Name of the client. If None the class will set a default and random
+                name. Defaults to None.
             nnunet_trainer_class (type[nnUNetTrainer]): A ``nnUNetTrainer`` constructor. Useful for passing custom
                 ``nnUNetTrainer``. Defaults to the standard nnUNetTrainer class. Must match the
                 ``nnunet_trainer_class`` passed to the ``NnunetServer``.
@@ -340,7 +342,13 @@ def get_lr_scheduler(self, optimizer_key: str, config: Config) -> _LRScheduler:
         this method to set your own LR scheduler.
 
         Args:
-            config (Config): The server config. This method will look for the
+            optimizer_key (str): Key of the optimizer to which the scheduler will be applied.
+            config (Config): The server config. This method will determine the total number of steps that will be
+                applied across all server rounds on FL training.
+
+        Raises:
+            ValueError: Thrown if the configuration does not contain either a steps or epochs specification.
+
         Returns:
             _LRScheduler: The default nnunet LR Scheduler for nnunetv2 2.5.1
         """
@@ -904,7 +912,7 @@ def update_before_train(self, current_server_round: int) -> None:
         gc.collect()  # Cleans up unused variables
         # As the linked issue above points out, calling gc.freeze() greatly reduces the
         # overhead of garbage collection. (from 1.5s to 0.005s)
-        if current_server_round == 2:
+        if current_server_round == 2:  # noqa: PLR2004
             # Collect runs even faster if we freeze after the end of the first iteration
             # Likely because a lot of variables are created in the first pass. If we
             # freeze before the first pass, gc.collect has to check all those variables

fl4health/clients/scaffold_client.py

@@ -123,6 +123,7 @@ def set_parameters(self, parameters: NDArrays, config: Config, fitting_round: bo
             parameters (NDArrays): Parameters have information about model state to be added to the relevant client
                 model and also the server control variates (initial or after aggregation)
             config (Config): The config is sent by the FL server to allow for customization in the function if desired.
+            fitting_round (bool): Which fitting round (i.e. server round of fitting) that we're on.
         """
         assert self.model is not None and self.parameter_exchanger is not None
 
@@ -309,12 +310,6 @@ def setup_client(self, config: Config) -> None:
 
 
 class DPScaffoldClient(ScaffoldClient, InstanceLevelDpClient):
-    """
-    Federated Learning client for Instance Level Differentially Private Scaffold strategy.
-
-    Implemented as specified in https://arxiv.org/abs/2111.09278
-    """
-
     def __init__(
         self,
         data_path: Path,
@@ -326,6 +321,30 @@ def __init__(
         progress_bar: bool = False,
         client_name: str | None = None,
     ) -> None:
+        """
+        Federated Learning client for Instance Level Differentially Private Scaffold strategy.
+
+        Implemented as specified in https://arxiv.org/abs/2111.09278
+
+        Args:
+            data_path (Path): path to the data to be used to load the data for client-side training
+            metrics (Sequence[Metric]): Metrics to be computed based on the labels and predictions of the client model
+            device (torch.device): Device indicator for where to send the model, batches, labels etc. Often "cpu" or
+                "cuda"
+            loss_meter_type (LossMeterType, optional): Type of meter used to track and compute the losses over
+                each batch. Defaults to ``LossMeterType.AVERAGE``.
+            checkpoint_and_state_module (ClientCheckpointAndStateModule | None, optional): A module meant to handle
+                both checkpointing and state saving. The module, and its underlying model and state checkpointing
+                components will determine when and how to do checkpointing during client-side training.
+                No checkpointing (state or model) is done if not provided. Defaults to None.
+            reporters (Sequence[BaseReporter] | None, optional): A sequence of FL4Health reporters which the client
+                should send data to. Defaults to None.
+            progress_bar (bool, optional): Whether or not to display a progress bar during client training and
+                validation. Uses ``tqdm``. Defaults to False.
+            client_name (str | None, optional): n optional client name that uniquely identifies a client.
+                If not passed, a hash is randomly generated. Client state will use this as part of its state file
+                name. Defaults to None.
+        """
         ScaffoldClient.__init__(
             self,
             data_path=data_path,

fl4health/feature_alignment/string_columns_transformer.py

@@ -20,10 +20,10 @@ def __init__(self, transformer: TextFeatureTransformer):
     def fit(self, x: pd.DataFrame, y: pd.DataFrame | None = None) -> TextMulticolumnTransformer:
         """
         Fit the transformer to the provided dataframe. The dataframe should have multiple string columns
-        The transformer is fit on the appended text from all columns in the ``X`` dataframe.
+        The transformer is fit on the appended text from all columns in the ``x`` dataframe.
 
         Args:
-            X (pd.DataFrame): Columns on which to fit the transformer
+            x (pd.DataFrame): Columns on which to fit the transformer
             y (pd.DataFrame | None, optional): Not used. Defaults to None.
 
         Returns:
@@ -35,10 +35,10 @@ def fit(self, x: pd.DataFrame, y: pd.DataFrame | None = None) -> TextMulticolumn
 
     def transform(self, x: pd.DataFrame) -> pd.DataFrame:
         """
-        Transforms the concatenation of all columns of text in the ``X`` dataframe.
+        Transforms the concatenation of all columns of text in the ``x`` dataframe.
 
         Args:
-            X (pd.DataFrame): Dataframe of text-based columns to be transformed
+            x (pd.DataFrame): Dataframe of text-based columns to be transformed
 
         Returns:
             pd.DataFrame: Transformed dataframe.
@@ -61,10 +61,10 @@ def __init__(self, transformer: TextFeatureTransformer):
     def fit(self, x: pd.DataFrame, y: pd.DataFrame | None = None) -> TextColumnTransformer:
         """
         Fit the transformer to the provided dataframe. The dataframe should have a single string column
-        The transformer is fit on the text from the single columns in the ``X`` dataframe.
+        The transformer is fit on the text from the single columns in the ``x`` dataframe.
 
         Args:
-            X (pd.DataFrame): Column on which to fit the transformer
+            x (pd.DataFrame): Column on which to fit the transformer
             y (pd.DataFrame | None, optional): Not used. Defaults to None.
 
         Returns:
@@ -76,10 +76,10 @@ def fit(self, x: pd.DataFrame, y: pd.DataFrame | None = None) -> TextColumnTrans
 
     def transform(self, x: pd.DataFrame) -> pd.DataFrame:
         """
-        Transforms the concatenation of a single column of text in the ``X`` dataframe.
+        Transforms the concatenation of a single column of text in the ``x`` dataframe.
 
         Args:
-            X (pd.DataFrame): Dataframe of text-based column to be transformed
+            x (pd.DataFrame): Dataframe of text-based column to be transformed
 
         Returns:
             pd.DataFrame: Transformed dataframe.