flwr 1.23.0__py3-none-any.whl → 1.24.0__py3-none-any.whl

flwr/server/superlink/linkstate/utils.py

@@ -16,7 +16,6 @@
 
 
 from os import urandom
-from typing import Optional
 
 from flwr.common import ConfigRecord, Context, Error, Message, Metadata, now, serde
 from flwr.common.constant import (
@@ -51,7 +50,8 @@ VALID_RUN_SUB_STATUSES = {
 }
 MESSAGE_UNAVAILABLE_ERROR_REASON = (
     "Error: Message Unavailable - The requested message could not be found in the "
-    "database. It may have expired due to its TTL or never existed."
+    "database. It may have expired due to its TTL, been deleted because the "
+    "destination SuperNode was removed from the federation, or never existed."
 )
 REPLY_MESSAGE_UNAVAILABLE_ERROR_REASON = (
     "Error: Reply Message Unavailable - The reply message has expired."
@@ -63,7 +63,7 @@ NODE_UNAVAILABLE_ERROR_REASON = (
 
 
 def generate_rand_int_from_bytes(
-    num_bytes: int, exclude: Optional[list[int]] = None
+    num_bytes: int, exclude: list[int] | None = None
 ) -> int:
     """Generate a random unsigned integer from `num_bytes` bytes.
 
@@ -257,7 +257,7 @@ def message_ttl_has_expired(message_metadata: Metadata, current_time: float) ->
 def verify_message_ids(
     inquired_message_ids: set[str],
     found_message_ins_dict: dict[str, Message],
-    current_time: Optional[float] = None,
+    current_time: float | None = None,
     update_set: bool = True,
 ) -> dict[str, Message]:
     """Verify found Messages and generate error Messages for invalid ones.
@@ -300,7 +300,7 @@ def verify_found_message_replies(
     inquired_message_ids: set[str],
     found_message_ins_dict: dict[str, Message],
     found_message_res_list: list[Message],
-    current_time: Optional[float] = None,
+    current_time: float | None = None,
     update_set: bool = True,
 ) -> dict[str, Message]:
     """Verify found Message replies and generate error Message for invalid ones.
@@ -345,7 +345,7 @@ def check_node_availability_for_in_message(
     inquired_in_message_ids: set[str],
     found_in_message_dict: dict[str, Message],
     node_id_to_online_until: dict[int, float],
-    current_time: Optional[float] = None,
+    current_time: float | None = None,
     update_set: bool = True,
 ) -> dict[str, Message]:
     """Check node availability for given Message and generate error reply Message if

flwr/server/superlink/serverappio/serverappio_grpc.py

@@ -16,7 +16,6 @@
 
 
 from logging import INFO
-from typing import Optional
 
 import grpc
 
@@ -38,7 +37,7 @@ def run_serverappio_api_grpc(
     state_factory: LinkStateFactory,
     ffs_factory: FfsFactory,
     objectstore_factory: ObjectStoreFactory,
-    certificates: Optional[tuple[bytes, bytes, bytes]],
+    certificates: tuple[bytes, bytes, bytes] | None,
 ) -> grpc.Server:
     """Run ServerAppIo API (gRPC, request-response)."""
     # Create ServerAppIo API gRPC server

flwr/server/superlink/serverappio/serverappio_servicer.py

@@ -17,7 +17,6 @@
 
 import threading
 from logging import DEBUG, ERROR, INFO
-from typing import Optional
 
 import grpc
 
@@ -91,7 +90,6 @@ from flwr.server.superlink.utils import abort_if
 from flwr.server.utils.validator import validate_message
 from flwr.supercore.ffs import Ffs, FfsFactory
 from flwr.supercore.object_store import NoObjectInStoreError, ObjectStoreFactory
-from flwr.supercore.object_store.utils import store_mapping_and_register_objects
 
 
 class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
@@ -141,6 +139,13 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
         # Attempt to create a token for the provided run ID
         token = state.create_token(request.run_id)
 
+        # Transition the run to STARTING if token creation was successful
+        if token:
+            state.update_run_status(
+                run_id=request.run_id,
+                new_status=RunStatus(Status.STARTING, "", ""),
+            )
+
         # Return the token
         return RequestTokenResponse(token=token or "")
 
@@ -192,8 +197,11 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
             request_name="PushMessages",
             detail="`messages_list` must not be empty",
         )
-        message_ids: list[Optional[str]] = []
-        for message_proto in request.messages_list:
+        message_ids: list[str | None] = []
+        objects_to_push: set[str] = set()
+        for message_proto, object_tree in zip(
+            request.messages_list, request.message_object_trees, strict=True
+        ):
             message = message_from_proto(message_proto=message_proto)
             validation_errors = validate_message(message, is_reply_message=False)
             _raise_if(
@@ -206,13 +214,12 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
                 request_name="PushMessages",
                 detail="`Message.metadata` has mismatched `run_id`",
             )
-            # Store
-            message_id: Optional[str] = state.store_message_ins(message=message)
+            # Store objects
+            objects_to_push |= set(store.preregister(request.run_id, object_tree))
+            # Store message
+            message_id: str | None = state.store_message_ins(message=message)
             message_ids.append(message_id)
 
-        # Store Message object to descendants mapping and preregister objects
-        objects_to_push = store_mapping_and_register_objects(store, request=request)
-
         return PushAppMessagesResponse(
             message_ids=[
                 str(message_id) if message_id else "" for message_id in message_ids
@@ -345,8 +352,8 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
                 if result := ffs.get(run.fab_hash):
                     fab = Fab(run.fab_hash, result[0], result[1])
             if run and fab and serverapp_ctxt:
-                # Update run status to STARTING
-                if state.update_run_status(run_id, RunStatus(Status.STARTING, "", "")):
+                # Update run status to RUNNING
+                if state.update_run_status(run_id, RunStatus(Status.RUNNING, "", "")):
                     log(INFO, "Starting run %d", run_id)
                     return PullAppInputsResponse(
                         context=context_to_proto(serverapp_ctxt),
@@ -355,8 +362,12 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
                     )
 
         # Raise an exception if the Run or Fab is not found,
-        # or if the status cannot be updated to STARTING
-        raise RuntimeError(f"Failed to start run {run_id}")
+        # or if the status cannot be updated to RUNNING
+        context.abort(
+            grpc.StatusCode.FAILED_PRECONDITION,
+            f"Failed to start run {run_id}",
+        )
+        raise RuntimeError("Unreachable code")  # for mypy
 
     def PushAppOutputs(
         self, request: PushAppOutputsRequest, context: grpc.ServicerContext
@@ -441,20 +452,14 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
     def SendAppHeartbeat(
         self, request: SendAppHeartbeatRequest, context: grpc.ServicerContext
     ) -> SendAppHeartbeatResponse:
-        """Handle a heartbeat from the ServerApp."""
+        """Handle a heartbeat from an app process."""
         log(DEBUG, "ServerAppIoServicer.SendAppHeartbeat")
 
         # Init state
         state = self.state_factory.state()
 
         # Acknowledge the heartbeat
-        # The app heartbeat can only be acknowledged if the run is in
-        # starting or running status.
-        success = state.acknowledge_app_heartbeat(
-            run_id=request.run_id,
-            heartbeat_interval=request.heartbeat_interval,
-        )
-
+        success = state.acknowledge_app_heartbeat(request.token)
         return SendAppHeartbeatResponse(success=success)
 
     def PushObject(

flwr/server/superlink/simulation/simulationio_grpc.py

@@ -16,7 +16,6 @@
 
 
 from logging import INFO
-from typing import Optional
 
 import grpc
 
@@ -36,7 +35,7 @@ def run_simulationio_api_grpc(
     address: str,
     state_factory: LinkStateFactory,
     ffs_factory: FfsFactory,
-    certificates: Optional[tuple[bytes, bytes, bytes]],
+    certificates: tuple[bytes, bytes, bytes] | None,
 ) -> grpc.Server:
     """Run SimulationIo API (gRPC, request-response)."""
     # Create SimulationIo API gRPC server

flwr/server/superlink/simulation/simulationio_servicer.py

@@ -110,6 +110,13 @@ class SimulationIoServicer(simulationio_pb2_grpc.SimulationIoServicer):
         # Attempt to create a token for the provided run ID
         token = state.create_token(request.run_id)
 
+        # Transition the run to STARTING if token creation was successful
+        if token:
+            state.update_run_status(
+                run_id=request.run_id,
+                new_status=RunStatus(Status.STARTING, "", ""),
+            )
+
         # Return the token
         return RequestTokenResponse(token=token or "")
 
@@ -152,8 +159,8 @@ class SimulationIoServicer(simulationio_pb2_grpc.SimulationIoServicer):
                 if result := ffs.get(run.fab_hash):
                     fab = Fab(run.fab_hash, result[0], result[1])
             if run and fab and serverapp_ctxt:
-                # Update run status to STARTING
-                if state.update_run_status(run_id, RunStatus(Status.STARTING, "", "")):
+                # Update run status to RUNNING
+                if state.update_run_status(run_id, RunStatus(Status.RUNNING, "", "")):
                     log(INFO, "Starting run %d", run_id)
                     return PullAppInputsResponse(
                         context=context_to_proto(serverapp_ctxt),
@@ -162,8 +169,12 @@ class SimulationIoServicer(simulationio_pb2_grpc.SimulationIoServicer):
                     )
 
         # Raise an exception if the Run or Fab is not found,
-        # or if the status cannot be updated to STARTING
-        raise RuntimeError(f"Failed to start run {run_id}")
+        # or if the status cannot be updated to RUNNING
+        context.abort(
+            grpc.StatusCode.FAILED_PRECONDITION,
+            f"Failed to start run {run_id}",
+        )
+        raise RuntimeError("Unreachable code")  # for mypy
 
     def PushAppOutputs(
         self, request: PushAppOutputsRequest, context: ServicerContext
@@ -257,20 +268,14 @@ class SimulationIoServicer(simulationio_pb2_grpc.SimulationIoServicer):
     def SendAppHeartbeat(
         self, request: SendAppHeartbeatRequest, context: grpc.ServicerContext
     ) -> SendAppHeartbeatResponse:
-        """Handle a heartbeat from the ServerApp in simulation."""
-        log(DEBUG, "SimultionIoServicer.SendAppHeartbeat")
+        """Handle a heartbeat from an app process."""
+        log(DEBUG, "SimulationIoServicer.SendAppHeartbeat")
 
         # Init state
         state = self.state_factory.state()
 
         # Acknowledge the heartbeat
-        # The app heartbeat can only be acknowledged if the run is in
-        # starting or running status.
-        success = state.acknowledge_app_heartbeat(
-            run_id=request.run_id,
-            heartbeat_interval=request.heartbeat_interval,
-        )
-
+        success = state.acknowledge_app_heartbeat(request.token)
         return SendAppHeartbeatResponse(success=success)
 
     def _verify_token(self, token: str, context: grpc.ServicerContext) -> int:

flwr/server/superlink/utils.py

@@ -15,8 +15,6 @@
 """SuperLink utilities."""
 
 
-from typing import Optional, Union
-
 import grpc
 
 from flwr.common.constant import Status, SubStatus
@@ -36,8 +34,8 @@ def check_abort(
     run_id: int,
     abort_status_list: list[str],
     state: LinkState,
-    store: Optional[ObjectStore] = None,
-) -> Union[str, None]:
+    store: ObjectStore | None = None,
+) -> str | None:
     """Check if the status of the provided `run_id` is in `abort_status_list`."""
     run_status: RunStatus = state.get_run_status({run_id})[run_id]
 
@@ -54,7 +52,7 @@ def check_abort(
     return None
 
 
-def abort_grpc_context(msg: Union[str, None], context: grpc.ServicerContext) -> None:
+def abort_grpc_context(msg: str | None, context: grpc.ServicerContext) -> None:
     """Abort context with statuscode PERMISSION_DENIED if `msg` is not None."""
     if msg is not None:
         context.abort(grpc.StatusCode.PERMISSION_DENIED, msg)
@@ -64,7 +62,7 @@ def abort_if(
     run_id: int,
     abort_status_list: list[str],
     state: LinkState,
-    store: Optional[ObjectStore],
+    store: ObjectStore | None,
     context: grpc.ServicerContext,
 ) -> None:
     """Abort context if status of the provided `run_id` is in `abort_status_list`."""

flwr/server/typing.py

@@ -15,7 +15,7 @@
 """Custom types for Flower servers."""
 
 
-from typing import Callable
+from collections.abc import Callable
 
 from flwr.common import Context
 

flwr/server/utils/tensorboard.py

@@ -16,20 +16,16 @@
 
 
 import os
+from collections.abc import Callable
 from datetime import datetime
 from logging import WARN
-from typing import Callable, Optional, Union, cast
+from typing import cast
 
 from flwr.common import EvaluateRes, Scalar
 from flwr.common.logger import log
 from flwr.server.client_proxy import ClientProxy
 from flwr.server.strategy import Strategy
 
-try:
-    import tensorflow as TF
-except ModuleNotFoundError:
-    TF = None
-
 MISSING_EXTRA_TF = """
 Extra dependency required for using tensorboard are missing.
 The program will continue without tensorboard.
@@ -59,6 +55,17 @@ def tensorboard(logdir: str) -> Callable[[Strategy], Strategy]:
         # Variant 2
         strategy = tensorboard(logdir=LOGDIR)(FedAvg)()
     """
+    log(
+        WARN,
+        "The `tensorboard` function is deprecated and will be removed "
+        "in a future release.",
+    )
+    # Lazy import of TensorFlow to avoid slow import times
+    try:
+        import tensorflow as TF  # pylint: disable=import-outside-toplevel
+    except ModuleNotFoundError:
+        TF = None  # pylint: disable=invalid-name
+
     print(
         "\n\t\033[32mStart TensorBoard with the following parameters"
         f"\n\t$ tensorboard --logdir {logdir}\033[39m\n"
@@ -93,8 +100,8 @@ def tensorboard(logdir: str) -> Callable[[Strategy], Strategy]:
                 self,
                 server_round: int,
                 results: list[tuple[ClientProxy, EvaluateRes]],
-                failures: list[Union[tuple[ClientProxy, EvaluateRes], BaseException]],
-            ) -> tuple[Optional[float], dict[str, Scalar]]:
+                failures: list[tuple[ClientProxy, EvaluateRes] | BaseException],
+            ) -> tuple[float | None, dict[str, Scalar]]:
                 """Hooks into aggregate_evaluate for TensorBoard logging purpose."""
                 # Execute decorated function and extract results for logging
                 # They will be returned at the end of this function but also

flwr/server/workflow/default_workflows.py

@@ -18,7 +18,7 @@
 import io
 import timeit
 from logging import INFO, WARN
-from typing import Optional, Union, cast
+from typing import cast
 
 import flwr.common.recorddict_compat as compat
 from flwr.common import (
@@ -47,8 +47,8 @@ class DefaultWorkflow:
 
     def __init__(
         self,
-        fit_workflow: Optional[Workflow] = None,
-        evaluate_workflow: Optional[Workflow] = None,
+        fit_workflow: Workflow | None = None,
+        evaluate_workflow: Workflow | None = None,
     ) -> None:
         if fit_workflow is None:
             fit_workflow = default_fit_workflow
@@ -275,7 +275,7 @@ def default_fit_workflow(grid: Grid, context: Context) -> None:  # pylint: disab
 
     # Aggregate training results
     results: list[tuple[ClientProxy, FitRes]] = []
-    failures: list[Union[tuple[ClientProxy, FitRes], BaseException]] = []
+    failures: list[tuple[ClientProxy, FitRes] | BaseException] = []
     for msg in messages:
         if msg.has_content():
             proxy = node_id_to_proxy[msg.metadata.src_node_id]
@@ -357,7 +357,7 @@ def default_evaluate_workflow(grid: Grid, context: Context) -> None:
 
     # Aggregate the evaluation results
     results: list[tuple[ClientProxy, EvaluateRes]] = []
-    failures: list[Union[tuple[ClientProxy, EvaluateRes], BaseException]] = []
+    failures: list[tuple[ClientProxy, EvaluateRes] | BaseException] = []
     for msg in messages:
         if msg.has_content():
             proxy = node_id_to_proxy[msg.metadata.src_node_id]

flwr/server/workflow/secure_aggregation/secagg_workflow.py

@@ -15,8 +15,6 @@
 """Workflow for the SecAgg protocol."""
 
 
-from typing import Optional, Union
-
 from .secaggplus_workflow import SecAggPlusWorkflow
 
 
@@ -94,13 +92,13 @@ class SecAggWorkflow(SecAggPlusWorkflow):
 
     def __init__(  # pylint: disable=R0913
         self,
-        reconstruction_threshold: Union[int, float],
+        reconstruction_threshold: int | float,
         *,
         max_weight: float = 1000.0,
         clipping_range: float = 8.0,
         quantization_range: int = 4194304,
         modulus_range: int = 4294967296,
-        timeout: Optional[float] = None,
+        timeout: float | None = None,
     ) -> None:
         super().__init__(
             num_shares=1.0,

flwr/server/workflow/secure_aggregation/secaggplus_workflow.py

@@ -18,7 +18,7 @@
 import random
 from dataclasses import dataclass, field
 from logging import DEBUG, ERROR, INFO, WARN
-from typing import Optional, Union, cast
+from typing import cast
 
 import flwr.common.recorddict_compat as compat
 from flwr.common import (
@@ -169,14 +169,14 @@ class SecAggPlusWorkflow:
 
     def __init__(  # pylint: disable=R0913
         self,
-        num_shares: Union[int, float],
-        reconstruction_threshold: Union[int, float],
+        num_shares: int | float,
+        reconstruction_threshold: int | float,
         *,
         max_weight: float = 1000.0,
         clipping_range: float = 8.0,
         quantization_range: int = 4194304,
         modulus_range: int = 4294967296,
-        timeout: Optional[float] = None,
+        timeout: float | None = None,
     ) -> None:
         self.num_shares = num_shares
         self.reconstruction_threshold = reconstruction_threshold
@@ -211,7 +211,7 @@ class SecAggPlusWorkflow:
 
     def _check_init_params(self) -> None:  # pylint: disable=R0912
         # Check `num_shares`
-        if not isinstance(self.num_shares, (int, float)):
+        if not isinstance(self.num_shares, (int | float)):
             raise TypeError("`num_shares` must be of type int or float.")
         if isinstance(self.num_shares, int):
             if self.num_shares == 1:
@@ -229,7 +229,7 @@ class SecAggPlusWorkflow:
             raise ValueError("`num_shares` as a float must be greater than 0.")
 
         # Check `reconstruction_threshold`
-        if not isinstance(self.reconstruction_threshold, (int, float)):
+        if not isinstance(self.reconstruction_threshold, (int | float)):
             raise TypeError("`reconstruction_threshold` must be of type int or float.")
         if isinstance(self.reconstruction_threshold, int):
             if self.reconstruction_threshold == 1:
@@ -467,7 +467,7 @@ class SecAggPlusWorkflow:
             dsts += dst_lst
             ciphertexts += ctxt_lst
 
-        for src, dst, ciphertext in zip(srcs, dsts, ciphertexts):
+        for src, dst, ciphertext in zip(srcs, dsts, ciphertexts, strict=True):
             if dst in fwd_ciphertexts:
                 fwd_ciphertexts[dst].append(ciphertext)
                 fwd_srcs[dst].append(src)
@@ -604,7 +604,7 @@ class SecAggPlusWorkflow:
             res_dict = msg.content.config_records[RECORD_KEY_CONFIGS]
             nids = cast(list[int], res_dict[Key.NODE_ID_LIST])
             shares = cast(list[bytes], res_dict[Key.SHARE_LIST])
-            for owner_nid, share in zip(nids, shares):
+            for owner_nid, share in zip(nids, shares, strict=True):
                 collected_shares_dict[owner_nid].append(share)
 
         # Remove masks for every active client after collect_masked_vectors stage

flwr/serverapp/strategy/bulyan.py

@@ -18,10 +18,9 @@ Paper: arxiv.org/abs/1802.07927
 """
 
 
-from collections import OrderedDict
-from collections.abc import Iterable
+from collections.abc import Callable, Iterable
 from logging import INFO, WARN
-from typing import Callable, Optional, cast
+from typing import cast
 
 import numpy as np
 
@@ -104,15 +103,15 @@ class Bulyan(FedAvg):
         weighted_by_key: str = "num-examples",
         arrayrecord_key: str = "arrays",
         configrecord_key: str = "config",
-        train_metrics_aggr_fn: Optional[
-            Callable[[list[RecordDict], str], MetricRecord]
-        ] = None,
-        evaluate_metrics_aggr_fn: Optional[
-            Callable[[list[RecordDict], str], MetricRecord]
-        ] = None,
-        selection_rule: Optional[
-            Callable[[list[RecordDict], int, int], list[RecordDict]]
-        ] = None,
+        train_metrics_aggr_fn: (
+            Callable[[list[RecordDict], str], MetricRecord] | None
+        ) = None,
+        evaluate_metrics_aggr_fn: (
+            Callable[[list[RecordDict], str], MetricRecord] | None
+        ) = None,
+        selection_rule: (
+            Callable[[list[RecordDict], int, int], list[RecordDict]] | None
+        ) = None,
     ) -> None:
         super().__init__(
             fraction_train=fraction_train,
@@ -140,7 +139,7 @@ class Bulyan(FedAvg):
         self,
         server_round: int,
         replies: Iterable[Message],
-    ) -> tuple[Optional[ArrayRecord], Optional[MetricRecord]]:
+    ) -> tuple[ArrayRecord | None, MetricRecord | None]:
         """Aggregate ArrayRecords and MetricRecords in the received Messages."""
         valid_replies, _ = self._check_and_log_replies(replies, is_train=True)
 
@@ -175,7 +174,9 @@ class Bulyan(FedAvg):
         ]
 
         # Compute median
-        median_ndarrays = [np.median(arr, axis=0) for arr in zip(*selected_ndarrays)]
+        median_ndarrays = [
+            np.median(arr, axis=0) for arr in zip(*selected_ndarrays, strict=True)
+        ]
 
         # Aggregate the beta closest weights element-wise
         aggregated_ndarrays = aggregate_n_closest_weights(
@@ -184,7 +185,7 @@ class Bulyan(FedAvg):
 
         # Convert to ArrayRecord
         arrays = ArrayRecord(
-            OrderedDict(zip(array_keys, map(Array, aggregated_ndarrays)))
+            dict(zip(array_keys, map(Array, aggregated_ndarrays), strict=True))
         )
 
         # Aggregate MetricRecords

flwr/serverapp/strategy/dp_adaptive_clipping.py

@@ -19,10 +19,8 @@ Paper (Andrew et al.): https://arxiv.org/abs/1905.03871
 
 import math
 from abc import ABC
-from collections import OrderedDict
 from collections.abc import Iterable
 from logging import INFO
-from typing import Optional
 
 import numpy as np
 
@@ -53,7 +51,7 @@ class DifferentialPrivacyAdaptiveBase(Strategy, ABC):
         initial_clipping_norm: float = 0.1,
         target_clipped_quantile: float = 0.5,
         clip_norm_lr: float = 0.2,
-        clipped_count_stddev: Optional[float] = None,
+        clipped_count_stddev: float | None = None,
     ) -> None:
         super().__init__()
 
@@ -96,7 +94,7 @@ class DifferentialPrivacyAdaptiveBase(Strategy, ABC):
         add_gaussian_noise_inplace(nds, stdv)
         log(INFO, "aggregate_fit: central DP noise with %.4f stdev added", stdv)
         return ArrayRecord(
-            OrderedDict({k: Array(v) for k, v in zip(aggregated.keys(), nds)})
+            {k: Array(v) for k, v in zip(aggregated.keys(), nds, strict=True)}
         )
 
     def _noisy_fraction(self, count: int, total: int) -> float:
@@ -115,7 +113,7 @@ class DifferentialPrivacyAdaptiveBase(Strategy, ABC):
 
     def aggregate_evaluate(
         self, server_round: int, replies: Iterable[Message]
-    ) -> Optional[MetricRecord]:
+    ) -> MetricRecord | None:
         """Aggregate MetricRecords in the received Messages."""
         return self.strategy.aggregate_evaluate(server_round, replies)
 
@@ -136,7 +134,7 @@ class DifferentialPrivacyServerSideAdaptiveClipping(DifferentialPrivacyAdaptiveB
         initial_clipping_norm: float = 0.1,
         target_clipped_quantile: float = 0.5,
         clip_norm_lr: float = 0.2,
-        clipped_count_stddev: Optional[float] = None,
+        clipped_count_stddev: float | None = None,
     ) -> None:
         super().__init__(
             strategy,
@@ -171,7 +169,7 @@ class DifferentialPrivacyServerSideAdaptiveClipping(DifferentialPrivacyAdaptiveB
 
     def aggregate_train(
         self, server_round: int, replies: Iterable[Message]
-    ) -> tuple[Optional[ArrayRecord], Optional[MetricRecord]]:
+    ) -> tuple[ArrayRecord | None, MetricRecord | None]:
         """Aggregate ArrayRecords and MetricRecords in the received Messages."""
         if not validate_replies(replies, self.num_sampled_clients):
             return None, None
@@ -184,16 +182,19 @@ class DifferentialPrivacyServerSideAdaptiveClipping(DifferentialPrivacyAdaptiveB
             for arr_name, record in reply.content.array_records.items():
                 reply_nd = record.to_numpy_ndarrays()
                 model_update = [
-                    np.subtract(x, y) for (x, y) in zip(reply_nd, current_nd)
+                    np.subtract(x, y)
+                    for (x, y) in zip(reply_nd, current_nd, strict=True)
                 ]
                 norm_bit = adaptive_clip_inputs_inplace(
                     model_update, self.clipping_norm
                 )
                 clipped_indicator_count += int(norm_bit)
                 # reconstruct array using clipped contribution from current round
-                restored = [c + u for c, u in zip(current_nd, model_update)]
+                restored = [
+                    c + u for c, u in zip(current_nd, model_update, strict=True)
+                ]
                 reply.content[arr_name] = ArrayRecord(
-                    OrderedDict({k: Array(v) for k, v in zip(record.keys(), restored)})
+                    {k: Array(v) for k, v in zip(record.keys(), restored, strict=True)}
                 )
             log(
                 INFO,
@@ -287,7 +288,7 @@ class DifferentialPrivacyClientSideAdaptiveClipping(DifferentialPrivacyAdaptiveB
 
     def aggregate_train(
         self, server_round: int, replies: Iterable[Message]
-    ) -> tuple[Optional[ArrayRecord], Optional[MetricRecord]]:
+    ) -> tuple[ArrayRecord | None, MetricRecord | None]:
         """Aggregate ArrayRecords and MetricRecords in the received Messages."""
         if not validate_replies(replies, self.num_sampled_clients):
             return None, None