Add docstrings, type hints, and modernize code

Author: ArthurDanjou

.python-version

@@ -1 +1 @@
->=3.10
+>=3.10, <3.15

histogram.py

@@ -69,7 +69,7 @@ def gtk_run(closure):
         Args:
           closure Ignored parameter
         """
-        tools.warning("GTK 3.0 is unavailable: %s" % (err,))
+        tools.warning(f"GTK 3.0 is unavailable: {err}")
 
 # ---------------------------------------------------------------------------- #
 # Data frame columns selection helper
@@ -91,7 +91,7 @@ def select(data, *only_columns):
     if len(only_columns) == 0:
         return data
     # Intelligent selection
-    columns = list()
+    columns = []
     for only_column in only_columns:
         only_column = only_column.lower()
         for column in data.columns:
@@ -136,7 +136,7 @@ def to_string(x):
           Converted data to string
         """
         if type(x) is float:
-            return "%e" % x
+            return f"{x:e}"
         return str(x).strip()
 
     def __init__(self, data, title="Display data"):
@@ -149,7 +149,7 @@ def __init__(self, data, title="Display data"):
         # Make and fill list store
         store = Gtk.ListStore(*([str] * (len(data.columns) + 1)))
         for row in data.itertuples():
-            store.append(list(self.to_string(x) for x in row))
+            store.append([self.to_string(x) for x in row])
         # Make the associated tree view
         view = Gtk.TreeView(store)
         columns = list(data.columns)
@@ -196,17 +196,15 @@ def __init__(self, path_results):
         # Ensure directory exist
         if not path_results.exists():
             raise tools.UserException(
-                "Result directory %r cannot be accessed or does not exist"
-                % str(path_results)
+                f"Result directory {str(path_results)!r} cannot be accessed or does not exist"
             )
         # Load configuration string
         path_config = path_results / "config"
         try:
             data_config = path_config.read_text().strip()
         except Exception as err:
             tools.warning(
-                "Result directory %r: unable to read configuration (%s)"
-                % (str(path_results), err)
+                f"Result directory {str(path_results)!r}: unable to read configuration ({err})"
             )
             data_config = None
         # Load configuration json
@@ -216,8 +214,7 @@ def __init__(self, path_results):
                 data_json = json.load(fd)
         except Exception as err:
             tools.warning(
-                "Result directory %r: unable to read JSON configuration (%s)"
-                % (str(path_results), err)
+                f"Result directory {str(path_results)!r}: unable to read JSON configuration ({err})"
             )
             data_json = None
         # Load training data
@@ -229,8 +226,7 @@ def __init__(self, path_results):
             data_study.index.name = "Step number"
         except Exception as err:
             tools.warning(
-                "Result directory %r: unable to read training data (%s)"
-                % (str(path_results), err)
+                f"Result directory {str(path_results)!r}: unable to read training data ({err})"
             )
             data_study = None
         # Load evaluation data
@@ -240,8 +236,7 @@ def __init__(self, path_results):
             data_eval.index.name = "Step number"
         except Exception as err:
             tools.warning(
-                "Result directory %r: unable to read evaluation data (%s)"
-                % (str(path_results), err)
+                f"Result directory {str(path_results)!r}: unable to read evaluation data ({err})"
             )
             data_eval = None
         # Merge data frames
@@ -284,8 +279,7 @@ def display(self, *only_columns, name=None):
         display(
             self.get(*only_columns),
             title=(
-                "Session data%s for %r"
-                % (" (subset)" if len(only_columns) > 0 else "", self.name)
+                "Session data{} for {!r}".format(" (subset)" if len(only_columns) > 0 else "", self.name)
             ),
         )
         # Return self to enable chaining
@@ -343,7 +337,7 @@ def compute_epoch(self):
         }.get(dataset_name)
         if training_size is None:
             tools.warning(
-                "Unknown dataset %r, cannot compute the epoch number" % dataset_name
+                f"Unknown dataset {dataset_name!r}, cannot compute the epoch number"
             )
             return self
         self.data[column_name] = self.data["Training point count"] / training_size
@@ -465,17 +459,15 @@ def include(self, data, *cols, errs=None, lalp=1.0, ccnt=None):
             data = data.data
         elif not isinstance(data, pandas.DataFrame):
             raise RuntimeError(
-                "Expected a Session or DataFrame for 'data', got a %r"
-                % tools.fullqual(type(data))
+                f"Expected a Session or DataFrame for 'data', got a {tools.fullqual(type(data))!r}"
             )
         # Get the x-axis values
         if self._idx is None:
             x = data.index.to_numpy()
         else:
             if self._idx not in data:
                 raise RuntimeError(
-                    "No column named %r to use as index in the given session/dataframe"
-                    % (self._idx,)
+                    f"No column named {self._idx!r} to use as index in the given session/dataframe"
                 )
             x = data[self._idx].to_numpy()
         # Select semantic: empty list = select all
@@ -541,17 +533,15 @@ def include_single(self, data, key, col, err=None, lalp=1.0, ccnt=None):
             data = data.data
         elif not isinstance(data, pandas.DataFrame):
             raise RuntimeError(
-                "Expected a Session or DataFrame for 'data', got a %r"
-                % tools.fullqual(type(data))
+                f"Expected a Session or DataFrame for 'data', got a {tools.fullqual(type(data))!r}"
             )
         # Get the x-axis values
         if self._idx is None:
             x = data.index.to_numpy()
         else:
             if self._idx not in data:
                 raise RuntimeError(
-                    "No column named %r to use as index in the given session/dataframe"
-                    % (self._idx,)
+                    f"No column named {self._idx!r} to use as index in the given session/dataframe"
                 )
             x = data[self._idx].to_numpy()
         # Pick a new line style and color
@@ -641,15 +631,13 @@ def generator_sum(gen):
         if zlabel is not None:
             if self._tax is None:
                 tools.warning(
-                    "No secondary y-axis found, but its label %r was provided"
-                    % (zlabel,)
+                    f"No secondary y-axis found, but its label {zlabel!r} was provided"
                 )
             else:
                 self._tax.set_ylabel(zlabel)
         elif self._tax is not None:
             tools.warning(
-                "No label provided for the secondary y-axis; using label %r from the primary"
-                % (ylabel,)
+                f"No label provided for the secondary y-axis; using label {ylabel!r} from the primary"
             )
             self._tax.set_ylabel(ylabel)
         self._ax.set_xlim(left=xmin, right=xmax)

krum/aggregators/__init__.py

@@ -10,51 +10,91 @@
 # @section DESCRIPTION
 #
 # Loading of the local modules.
-#
-# Each rule MUST support taking any named arguments, possibly ignoring them.
-# The parameters MUST all be passed as their keyword arguments.
-# The reserved argument names, and their interface, are the following:
-# · gradients: Non-empty list of gradients to aggregate
-# · f        : Number of Byzantine gradients to support
-# · model    : Model (duck-typing 'experiments.Model') with valid default dataset and loss set
-# The rule, given "valid" parameter(s), MUST NOT return a tensor that is a reference to any tensor given as parameter.
-#
-# Each rule MUST provide a "check" member function, taking the same arguments as the rule itself.
-# The "check" member function returns 'None' when the parameters are valid,
-# or an explanatory string when the parameters are not valid.
-# The check member function MUST NOT modify the given parameters.
-#
-# Once registered, the check member function will be available as member "check".
-# The raw function and a wrapped checking the input/output of the raw function
-# will respectively be available as members "unchecked" and "checked".
-# Which of these two functions is called by default depends whether debug mode is enabled.
 ###
 
+"""
+Gradient aggregation rules (GARs) for Byzantine-resilient distributed learning.
+
+Each rule combines a keyword-only aggregation function with a validation
+function and optional metadata used by the training and experiment scripts.
+
+Contract
+--------
+
+Each aggregation rule MUST:
+
+1. Accept keyword-only arguments
+2. Accept the reserved parameter ``gradients`` (non-empty list of gradients)
+3. Accept the reserved parameter ``f`` (number of Byzantine gradients to tolerate)
+4. Accept the reserved parameter ``model`` (model with configured defaults)
+5. NOT return a tensor that aliases any input tensor
+
+Each rule MUST provide a ``check`` function that validates parameters and
+returns ``None`` when valid, or a user-facing error message otherwise.
+
+The module exposes three variants for each rule:
+
+- ``rule``: The default version (checked in debug mode, unchecked in release)
+- ``rule.checked``: Always validates parameters
+- ``rule.unchecked``: Skips validation (faster in production)
+
+Additional metadata available on each rule:
+
+- ``rule.check``: The validation function
+- ``rule.upper_bound``: Theoretical bound on stddev/norm ratio (if available)
+- ``rule.influence``: Attack acceptance ratio (if available)
+"""
+
 import pathlib
+from collections.abc import Callable
 
-from krum import tools
+import tools
+import torch
 
 # ---------------------------------------------------------------------------- #
 # Automated GAR loader
 
 
-def make_gar(unchecked, check, upper_bound=None, influence=None):
-    """GAR wrapper helper.
-    Args:
-      unchecked   Associated function (see module description)
-      check       Parameter validity check function
-      upper_bound Compute the theoretical upper bound on the ratio non-Byzantine standard deviation / norm to use this aggregation rule: (n, f, d) -> float
-      influence   Attack acceptation ratio function
-    Returns:
-      Wrapped GAR
+def make_gar(
+    unchecked: Callable,
+    check: Callable,
+    upper_bound: Callable | None = None,
+    influence: Callable | None = None,
+) -> Callable:
+    """
+    Wrap an unchecked GAR with validation and metadata.
+
+    Parameters
+    ----------
+    unchecked : callable
+        Aggregation function implementing the rule without parameter checks.
+    check : callable
+        Validation function. It must return ``None`` when parameters are valid,
+        or an error message otherwise.
+    upper_bound : callable, optional
+        Function computing the theoretical upper bound on the ratio between
+        non-Byzantine standard deviation and gradient norm. The expected
+        signature is ``(n, f, d) -> float``.
+    influence : callable, optional
+        Function computing the accepted Byzantine-gradient ratio for a given
+        set of honest and attack gradients.
+
+    Returns
+    -------
+    callable
+        Checked or unchecked GAR selected according to ``__debug__``. The
+        returned callable is annotated with ``check``, ``checked``,
+        ``unchecked``, ``upper_bound``, and ``influence`` attributes.
     """
 
     # Closure wrapping the call with checks
     def checked(**kwargs):
         # Check parameter validity
         message = check(**kwargs)
         if message is not None:
-            raise tools.UserException(f"Aggregation rule {name!r} cannot be used with the given parameters: {message}")
+            raise tools.UserException(
+                f"Aggregation rule {name!r} cannot be used with the given parameters: {message}"
+            )
         # Aggregation (hard to assert return value, duck-typing is allowed...)
         return unchecked(**kwargs)
 
@@ -70,26 +110,42 @@ def checked(**kwargs):
     return func
 
 
-def register(name, unchecked, check, upper_bound=None, influence=None):
-    """Simple registration-wrapper helper.
-    Args:
-      name        GAR name
-      unchecked   Associated function (see module description)
-      check       Parameter validity check function
-      upper_bound Compute the theoretical upper bound on the ratio non-Byzantine standard deviation / norm to use this aggregation rule: (n, f, d) -> float
-      influence   Attack acceptation ratio function
+def register(
+    name: str,
+    unchecked: Callable,
+    check: Callable,
+    upper_bound: Callable | None = None,
+    influence: Callable | None = None,
+) -> None:
+    """
+    Register a gradient aggregation rule.
+
+    Parameters
+    ----------
+    name : str
+        User-visible GAR name.
+    unchecked : callable
+        Aggregation function implementing the rule without parameter checks.
+    check : callable
+        Validation function associated with ``unchecked``.
+    upper_bound : callable, optional
+        Function computing the rule's theoretical upper bound.
+    influence : callable, optional
+        Function computing the accepted Byzantine-gradient ratio.
     """
     global gars
     # Check if name already in use
     if name in gars:
         tools.warning(f"Unable to register {name!r} GAR: name already in use")
         return
     # Export the selected function with the associated name
-    gars[name] = make_gar(unchecked, check, upper_bound=upper_bound, influence=influence)
+    gars[name] = make_gar(
+        unchecked, check, upper_bound=upper_bound, influence=influence
+    )
 
 
 # Registered rules (mapping name -> aggregation rule)
-gars = dict()
+gars = {}
 
 # Load all local modules
 with tools.Context("aggregators", None):