flwr-nightly 1.12.0.dev20240906__py3-none-any.whl → 1.12.0.dev20240913__py3-none-any.whl

flwr/server/superlink/fleet/vce/backend/raybackend.py

@@ -16,7 +16,7 @@
 
 import sys
 from logging import DEBUG, ERROR
-from typing import Callable, Dict, Optional, Tuple, Union
+from typing import Callable, Optional, Union
 
 import ray
 
@@ -31,8 +31,8 @@ from flwr.simulation.ray_transport.utils import enable_tf_gpu_growth
 
 from .backend import Backend, BackendConfig
 
-ClientResourcesDict = Dict[str, Union[int, float]]
-ActorArgsDict = Dict[str, Union[int, float, Callable[[], None]]]
+ClientResourcesDict = dict[str, Union[int, float]]
+ActorArgsDict = dict[str, Union[int, float, Callable[[], None]]]
 
 
 class RayBackend(Backend):
@@ -52,16 +52,11 @@ class RayBackend(Backend):
 
         # Validate client resources
         self.client_resources_key = "client_resources"
-        client_resources = self._validate_client_resources(config=backend_config)
+        self.client_resources = self._validate_client_resources(config=backend_config)
 
-        # Create actor pool
-        actor_kwargs = self._validate_actor_arguments(config=backend_config)
-
-        self.pool = BasicActorPool(
-            actor_type=ClientAppActor,
-            client_resources=client_resources,
-            actor_kwargs=actor_kwargs,
-        )
+        # Valide actor resources
+        self.actor_kwargs = self._validate_actor_arguments(config=backend_config)
+        self.pool: Optional[BasicActorPool] = None
 
         self.app_fn: Optional[Callable[[], ClientApp]] = None
 
@@ -106,7 +101,7 @@ class RayBackend(Backend):
     def init_ray(self, backend_config: BackendConfig) -> None:
         """Intialises Ray if not already initialised."""
         if not ray.is_initialized():
-            ray_init_args: Dict[
+            ray_init_args: dict[
                 str,
                 ConfigsRecordValues,
             ] = {}
@@ -122,14 +117,24 @@ class RayBackend(Backend):
     @property
     def num_workers(self) -> int:
         """Return number of actors in pool."""
-        return self.pool.num_actors
+        return self.pool.num_actors if self.pool else 0
 
     def is_worker_idle(self) -> bool:
         """Report whether the pool has idle actors."""
-        return self.pool.is_actor_available()
+        return self.pool.is_actor_available() if self.pool else False
 
     def build(self, app_fn: Callable[[], ClientApp]) -> None:
         """Build pool of Ray actors that this backend will submit jobs to."""
+        # Create Actor Pool
+        try:
+            self.pool = BasicActorPool(
+                actor_type=ClientAppActor,
+                client_resources=self.client_resources,
+                actor_kwargs=self.actor_kwargs,
+            )
+        except Exception as ex:
+            raise ex
+
         self.pool.add_actors_to_pool(self.pool.actors_capacity)
         # Set ClientApp callable that ray actors will use
         self.app_fn = app_fn
@@ -139,13 +144,16 @@ class RayBackend(Backend):
         self,
         message: Message,
         context: Context,
-    ) -> Tuple[Message, Context]:
+    ) -> tuple[Message, Context]:
         """Run ClientApp that process a given message.
 
         Return output message and updated context.
         """
         partition_id = context.node_config[PARTITION_ID_KEY]
 
+        if self.pool is None:
+            raise ValueError("The actor pool is empty, unfit to process messages.")
+
         if self.app_fn is None:
             raise ValueError(
                 "Unspecified function to load a `ClientApp`. "
@@ -179,6 +187,7 @@ class RayBackend(Backend):
 
     def terminate(self) -> None:
         """Terminate all actors in actor pool."""
-        self.pool.terminate_all_actors()
+        if self.pool:
+            self.pool.terminate_all_actors()
         ray.shutdown()
         log(DEBUG, "Terminated %s", self.__class__.__name__)

flwr/server/superlink/fleet/vce/vce_api.py

@@ -24,7 +24,7 @@ from logging import DEBUG, ERROR, INFO, WARN
 from pathlib import Path
 from queue import Empty, Queue
 from time import sleep
-from typing import Callable, Dict, Optional
+from typing import Callable, Optional
 
 from flwr.client.client_app import ClientApp, ClientAppException, LoadClientAppError
 from flwr.client.clientapp.utils import get_load_client_app_fn
@@ -44,7 +44,7 @@ from flwr.server.superlink.state import State, StateFactory
 
 from .backend import Backend, error_messages_backends, supported_backends
 
-NodeToPartitionMapping = Dict[int, int]
+NodeToPartitionMapping = dict[int, int]
 
 
 def _register_nodes(
@@ -64,9 +64,9 @@ def _register_node_states(
     nodes_mapping: NodeToPartitionMapping,
     run: Run,
     app_dir: Optional[str] = None,
-) -> Dict[int, NodeState]:
+) -> dict[int, NodeState]:
     """Create NodeState objects and pre-register the context for the run."""
-    node_states: Dict[int, NodeState] = {}
+    node_states: dict[int, NodeState] = {}
     num_partitions = len(set(nodes_mapping.values()))
     for node_id, partition_id in nodes_mapping.items():
         node_states[node_id] = NodeState(
@@ -89,7 +89,7 @@ def _register_node_states(
 def worker(
     taskins_queue: "Queue[TaskIns]",
     taskres_queue: "Queue[TaskRes]",
-    node_states: Dict[int, NodeState],
+    node_states: dict[int, NodeState],
     backend: Backend,
     f_stop: threading.Event,
 ) -> None:
@@ -177,7 +177,7 @@ def run_api(
     backend_fn: Callable[[], Backend],
     nodes_mapping: NodeToPartitionMapping,
     state_factory: StateFactory,
-    node_states: Dict[int, NodeState],
+    node_states: dict[int, NodeState],
     f_stop: threading.Event,
 ) -> None:
     """Run the VCE."""

flwr/server/superlink/state/in_memory_state.py

@@ -18,7 +18,7 @@
 import threading
 import time
 from logging import ERROR
-from typing import Dict, List, Optional, Set, Tuple
+from typing import Optional
 from uuid import UUID, uuid4
 
 from flwr.common import log, now
@@ -37,15 +37,15 @@ class InMemoryState(State):  # pylint: disable=R0902,R0904
     def __init__(self) -> None:
 
         # Map node_id to (online_until, ping_interval)
-        self.node_ids: Dict[int, Tuple[float, float]] = {}
-        self.public_key_to_node_id: Dict[bytes, int] = {}
+        self.node_ids: dict[int, tuple[float, float]] = {}
+        self.public_key_to_node_id: dict[bytes, int] = {}
 
         # Map run_id to (fab_id, fab_version)
-        self.run_ids: Dict[int, Run] = {}
-        self.task_ins_store: Dict[UUID, TaskIns] = {}
-        self.task_res_store: Dict[UUID, TaskRes] = {}
+        self.run_ids: dict[int, Run] = {}
+        self.task_ins_store: dict[UUID, TaskIns] = {}
+        self.task_res_store: dict[UUID, TaskRes] = {}
 
-        self.node_public_keys: Set[bytes] = set()
+        self.node_public_keys: set[bytes] = set()
         self.server_public_key: Optional[bytes] = None
         self.server_private_key: Optional[bytes] = None
 
@@ -76,13 +76,13 @@ class InMemoryState(State):  # pylint: disable=R0902,R0904
 
     def get_task_ins(
         self, node_id: Optional[int], limit: Optional[int]
-    ) -> List[TaskIns]:
+    ) -> list[TaskIns]:
         """Get all TaskIns that have not been delivered yet."""
         if limit is not None and limit < 1:
             raise AssertionError("`limit` must be >= 1")
 
         # Find TaskIns for node_id that were not delivered yet
-        task_ins_list: List[TaskIns] = []
+        task_ins_list: list[TaskIns] = []
         with self.lock:
             for _, task_ins in self.task_ins_store.items():
                 # pylint: disable=too-many-boolean-expressions
@@ -133,15 +133,15 @@ class InMemoryState(State):  # pylint: disable=R0902,R0904
         # Return the new task_id
         return task_id
 
-    def get_task_res(self, task_ids: Set[UUID], limit: Optional[int]) -> List[TaskRes]:
+    def get_task_res(self, task_ids: set[UUID], limit: Optional[int]) -> list[TaskRes]:
         """Get all TaskRes that have not been delivered yet."""
         if limit is not None and limit < 1:
             raise AssertionError("`limit` must be >= 1")
 
         with self.lock:
             # Find TaskRes that were not delivered yet
-            task_res_list: List[TaskRes] = []
-            replied_task_ids: Set[UUID] = set()
+            task_res_list: list[TaskRes] = []
+            replied_task_ids: set[UUID] = set()
             for _, task_res in self.task_res_store.items():
                 reply_to = UUID(task_res.task.ancestry[0])
                 if reply_to in task_ids and task_res.task.delivered_at == "":
@@ -175,10 +175,10 @@ class InMemoryState(State):  # pylint: disable=R0902,R0904
             # Return TaskRes
             return task_res_list
 
-    def delete_tasks(self, task_ids: Set[UUID]) -> None:
+    def delete_tasks(self, task_ids: set[UUID]) -> None:
         """Delete all delivered TaskIns/TaskRes pairs."""
-        task_ins_to_be_deleted: Set[UUID] = set()
-        task_res_to_be_deleted: Set[UUID] = set()
+        task_ins_to_be_deleted: set[UUID] = set()
+        task_res_to_be_deleted: set[UUID] = set()
 
         with self.lock:
             for task_ins_id in task_ids:
@@ -253,7 +253,7 @@ class InMemoryState(State):  # pylint: disable=R0902,R0904
 
             del self.node_ids[node_id]
 
-    def get_nodes(self, run_id: int) -> Set[int]:
+    def get_nodes(self, run_id: int) -> set[int]:
         """Return all available nodes.
 
         Constraints
@@ -318,7 +318,7 @@ class InMemoryState(State):  # pylint: disable=R0902,R0904
         """Retrieve `server_public_key` in urlsafe bytes."""
         return self.server_public_key
 
-    def store_node_public_keys(self, public_keys: Set[bytes]) -> None:
+    def store_node_public_keys(self, public_keys: set[bytes]) -> None:
         """Store a set of `node_public_keys` in state."""
         with self.lock:
             self.node_public_keys = public_keys
@@ -328,7 +328,7 @@ class InMemoryState(State):  # pylint: disable=R0902,R0904
         with self.lock:
             self.node_public_keys.add(public_key)
 
-    def get_node_public_keys(self) -> Set[bytes]:
+    def get_node_public_keys(self) -> set[bytes]:
         """Retrieve all currently stored `node_public_keys` as a set."""
         return self.node_public_keys
 

flwr/server/superlink/state/sqlite_state.py

@@ -19,8 +19,9 @@ import json
 import re
 import sqlite3
 import time
+from collections.abc import Sequence
 from logging import DEBUG, ERROR
-from typing import Any, Dict, List, Optional, Sequence, Set, Tuple, Union, cast
+from typing import Any, Optional, Union, cast
 from uuid import UUID, uuid4
 
 from flwr.common import log, now
@@ -110,7 +111,7 @@ CREATE TABLE IF NOT EXISTS task_res(
 );
 """
 
-DictOrTuple = Union[Tuple[Any, ...], Dict[str, Any]]
+DictOrTuple = Union[tuple[Any, ...], dict[str, Any]]
 
 
 class SqliteState(State):  # pylint: disable=R0904
@@ -131,7 +132,7 @@ class SqliteState(State):  # pylint: disable=R0904
         self.database_path = database_path
         self.conn: Optional[sqlite3.Connection] = None
 
-    def initialize(self, log_queries: bool = False) -> List[Tuple[str]]:
+    def initialize(self, log_queries: bool = False) -> list[tuple[str]]:
         """Create tables if they don't exist yet.
 
         Parameters
@@ -162,7 +163,7 @@ class SqliteState(State):  # pylint: disable=R0904
         self,
         query: str,
         data: Optional[Union[Sequence[DictOrTuple], DictOrTuple]] = None,
-    ) -> List[Dict[str, Any]]:
+    ) -> list[dict[str, Any]]:
         """Execute a SQL query."""
         if self.conn is None:
             raise AttributeError("State is not initialized.")
@@ -237,7 +238,7 @@ class SqliteState(State):  # pylint: disable=R0904
 
     def get_task_ins(
         self, node_id: Optional[int], limit: Optional[int]
-    ) -> List[TaskIns]:
+    ) -> list[TaskIns]:
         """Get undelivered TaskIns for one node (either anonymous or with ID).
 
         Usually, the Fleet API calls this for Nodes planning to work on one or more
@@ -271,7 +272,7 @@ class SqliteState(State):  # pylint: disable=R0904
             )
             raise AssertionError(msg)
 
-        data: Dict[str, Union[str, int]] = {}
+        data: dict[str, Union[str, int]] = {}
 
         if node_id is None:
             # Retrieve all anonymous Tasks
@@ -367,7 +368,7 @@ class SqliteState(State):  # pylint: disable=R0904
         return task_id
 
     # pylint: disable-next=R0914
-    def get_task_res(self, task_ids: Set[UUID], limit: Optional[int]) -> List[TaskRes]:
+    def get_task_res(self, task_ids: set[UUID], limit: Optional[int]) -> list[TaskRes]:
         """Get TaskRes for task_ids.
 
         Usually, the Driver API calls this method to get results for instructions it has
@@ -397,7 +398,7 @@ class SqliteState(State):  # pylint: disable=R0904
             AND delivered_at = ""
         """
 
-        data: Dict[str, Union[str, float, int]] = {}
+        data: dict[str, Union[str, float, int]] = {}
 
         if limit is not None:
             query += " LIMIT :limit"
@@ -435,7 +436,7 @@ class SqliteState(State):  # pylint: disable=R0904
         # 1. Query: Fetch consumer_node_id of remaining task_ids
         # Assume the ancestry field only contains one element
         data.clear()
-        replied_task_ids: Set[UUID] = {UUID(str(row["ancestry"])) for row in rows}
+        replied_task_ids: set[UUID] = {UUID(str(row["ancestry"])) for row in rows}
         remaining_task_ids = task_ids - replied_task_ids
         placeholders = ",".join([f":id_{i}" for i in range(len(remaining_task_ids))])
         query = f"""
@@ -499,10 +500,10 @@ class SqliteState(State):  # pylint: disable=R0904
         """
         query = "SELECT count(*) AS num FROM task_res;"
         rows = self.query(query)
-        result: Dict[str, int] = rows[0]
+        result: dict[str, int] = rows[0]
         return result["num"]
 
-    def delete_tasks(self, task_ids: Set[UUID]) -> None:
+    def delete_tasks(self, task_ids: set[UUID]) -> None:
         """Delete all delivered TaskIns/TaskRes pairs."""
         ids = list(task_ids)
         if len(ids) == 0:
@@ -588,7 +589,7 @@ class SqliteState(State):  # pylint: disable=R0904
         except KeyError as exc:
             log(ERROR, {"query": query, "data": params, "exception": exc})
 
-    def get_nodes(self, run_id: int) -> Set[int]:
+    def get_nodes(self, run_id: int) -> set[int]:
         """Retrieve all currently stored node IDs as a set.
 
         Constraints
@@ -604,7 +605,7 @@ class SqliteState(State):  # pylint: disable=R0904
         # Get nodes
         query = "SELECT node_id FROM node WHERE online_until > ?;"
         rows = self.query(query, (time.time(),))
-        result: Set[int] = {row["node_id"] for row in rows}
+        result: set[int] = {row["node_id"] for row in rows}
         return result
 
     def get_node_id(self, node_public_key: bytes) -> Optional[int]:
@@ -684,7 +685,7 @@ class SqliteState(State):  # pylint: disable=R0904
             public_key = None
         return public_key
 
-    def store_node_public_keys(self, public_keys: Set[bytes]) -> None:
+    def store_node_public_keys(self, public_keys: set[bytes]) -> None:
         """Store a set of `node_public_keys` in state."""
         query = "INSERT INTO public_key (public_key) VALUES (?)"
         data = [(key,) for key in public_keys]
@@ -695,11 +696,11 @@ class SqliteState(State):  # pylint: disable=R0904
         query = "INSERT INTO public_key (public_key) VALUES (:public_key)"
         self.query(query, {"public_key": public_key})
 
-    def get_node_public_keys(self) -> Set[bytes]:
+    def get_node_public_keys(self) -> set[bytes]:
         """Retrieve all currently stored `node_public_keys` as a set."""
         query = "SELECT public_key FROM public_key"
         rows = self.query(query)
-        result: Set[bytes] = {row["public_key"] for row in rows}
+        result: set[bytes] = {row["public_key"] for row in rows}
         return result
 
     def get_run(self, run_id: int) -> Optional[Run]:
@@ -733,7 +734,7 @@ class SqliteState(State):  # pylint: disable=R0904
 def dict_factory(
     cursor: sqlite3.Cursor,
     row: sqlite3.Row,
-) -> Dict[str, Any]:
+) -> dict[str, Any]:
     """Turn SQLite results into dicts.
 
     Less efficent for retrival of large amounts of data but easier to use.
@@ -742,7 +743,7 @@ def dict_factory(
     return dict(zip(fields, row))
 
 
-def task_ins_to_dict(task_msg: TaskIns) -> Dict[str, Any]:
+def task_ins_to_dict(task_msg: TaskIns) -> dict[str, Any]:
     """Transform TaskIns to dict."""
     result = {
         "task_id": task_msg.task_id,
@@ -763,7 +764,7 @@ def task_ins_to_dict(task_msg: TaskIns) -> Dict[str, Any]:
     return result
 
 
-def task_res_to_dict(task_msg: TaskRes) -> Dict[str, Any]:
+def task_res_to_dict(task_msg: TaskRes) -> dict[str, Any]:
     """Transform TaskRes to dict."""
     result = {
         "task_id": task_msg.task_id,
@@ -784,7 +785,7 @@ def task_res_to_dict(task_msg: TaskRes) -> Dict[str, Any]:
     return result
 
 
-def dict_to_task_ins(task_dict: Dict[str, Any]) -> TaskIns:
+def dict_to_task_ins(task_dict: dict[str, Any]) -> TaskIns:
     """Turn task_dict into protobuf message."""
     recordset = RecordSet()
     recordset.ParseFromString(task_dict["recordset"])
@@ -814,7 +815,7 @@ def dict_to_task_ins(task_dict: Dict[str, Any]) -> TaskIns:
     return result
 
 
-def dict_to_task_res(task_dict: Dict[str, Any]) -> TaskRes:
+def dict_to_task_res(task_dict: dict[str, Any]) -> TaskRes:
     """Turn task_dict into protobuf message."""
     recordset = RecordSet()
     recordset.ParseFromString(task_dict["recordset"])

flwr/server/superlink/state/state.py

@@ -16,7 +16,7 @@
 
 
 import abc
-from typing import List, Optional, Set
+from typing import Optional
 from uuid import UUID
 
 from flwr.common.typing import Run, UserConfig
@@ -51,7 +51,7 @@ class State(abc.ABC):  # pylint: disable=R0904
     @abc.abstractmethod
     def get_task_ins(
         self, node_id: Optional[int], limit: Optional[int]
-    ) -> List[TaskIns]:
+    ) -> list[TaskIns]:
         """Get TaskIns optionally filtered by node_id.
 
         Usually, the Fleet API calls this for Nodes planning to work on one or more
@@ -98,7 +98,7 @@ class State(abc.ABC):  # pylint: disable=R0904
         """
 
     @abc.abstractmethod
-    def get_task_res(self, task_ids: Set[UUID], limit: Optional[int]) -> List[TaskRes]:
+    def get_task_res(self, task_ids: set[UUID], limit: Optional[int]) -> list[TaskRes]:
         """Get TaskRes for task_ids.
 
         Usually, the Driver API calls this method to get results for instructions it has
@@ -129,7 +129,7 @@ class State(abc.ABC):  # pylint: disable=R0904
         """
 
     @abc.abstractmethod
-    def delete_tasks(self, task_ids: Set[UUID]) -> None:
+    def delete_tasks(self, task_ids: set[UUID]) -> None:
         """Delete all delivered TaskIns/TaskRes pairs."""
 
     @abc.abstractmethod
@@ -143,7 +143,7 @@ class State(abc.ABC):  # pylint: disable=R0904
         """Remove `node_id` from state."""
 
     @abc.abstractmethod
-    def get_nodes(self, run_id: int) -> Set[int]:
+    def get_nodes(self, run_id: int) -> set[int]:
         """Retrieve all currently stored node IDs as a set.
 
         Constraints
@@ -199,7 +199,7 @@ class State(abc.ABC):  # pylint: disable=R0904
         """Retrieve `server_public_key` in urlsafe bytes."""
 
     @abc.abstractmethod
-    def store_node_public_keys(self, public_keys: Set[bytes]) -> None:
+    def store_node_public_keys(self, public_keys: set[bytes]) -> None:
         """Store a set of `node_public_keys` in state."""
 
     @abc.abstractmethod
@@ -207,7 +207,7 @@ class State(abc.ABC):  # pylint: disable=R0904
         """Store a `node_public_key` in state."""
 
     @abc.abstractmethod
-    def get_node_public_keys(self) -> Set[bytes]:
+    def get_node_public_keys(self) -> set[bytes]:
         """Retrieve all currently stored `node_public_keys` as a set."""
 
     @abc.abstractmethod

flwr/server/utils/tensorboard.py

@@ -18,7 +18,7 @@
 import os
 from datetime import datetime
 from logging import WARN
-from typing import Callable, Dict, List, Optional, Tuple, Union, cast
+from typing import Callable, Optional, Union, cast
 
 from flwr.common import EvaluateRes, Scalar
 from flwr.common.logger import log
@@ -92,9 +92,9 @@ def tensorboard(logdir: str) -> Callable[[Strategy], Strategy]:
             def aggregate_evaluate(
                 self,
                 server_round: int,
-                results: List[Tuple[ClientProxy, EvaluateRes]],
-                failures: List[Union[Tuple[ClientProxy, EvaluateRes], BaseException]],
-            ) -> Tuple[Optional[float], Dict[str, Scalar]]:
+                results: list[tuple[ClientProxy, EvaluateRes]],
+                failures: list[Union[tuple[ClientProxy, EvaluateRes], BaseException]],
+            ) -> tuple[Optional[float], 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/utils/validator.py

@@ -15,13 +15,13 @@
 """Validators."""
 
 
-from typing import List, Union
+from typing import Union
 
 from flwr.proto.task_pb2 import TaskIns, TaskRes  # pylint: disable=E0611
 
 
 # pylint: disable-next=too-many-branches,too-many-statements
-def validate_task_ins_or_res(tasks_ins_res: Union[TaskIns, TaskRes]) -> List[str]:
+def validate_task_ins_or_res(tasks_ins_res: Union[TaskIns, TaskRes]) -> list[str]:
     """Validate a TaskIns or TaskRes."""
     validation_errors = []
 

flwr/server/workflow/default_workflows.py

@@ -18,7 +18,7 @@
 import io
 import timeit
 from logging import INFO, WARN
-from typing import List, Optional, Tuple, Union, cast
+from typing import Optional, Union, cast
 
 import flwr.common.recordset_compat as compat
 from flwr.common import (
@@ -276,8 +276,8 @@ def default_fit_workflow(  # pylint: disable=R0914
     )
 
     # Aggregate training results
-    results: List[Tuple[ClientProxy, FitRes]] = []
-    failures: List[Union[Tuple[ClientProxy, FitRes], BaseException]] = []
+    results: list[tuple[ClientProxy, FitRes]] = []
+    failures: list[Union[tuple[ClientProxy, FitRes], BaseException]] = []
     for msg in messages:
         if msg.has_content():
             proxy = node_id_to_proxy[msg.metadata.src_node_id]
@@ -362,8 +362,8 @@ def default_evaluate_workflow(driver: Driver, context: Context) -> None:
     )
 
     # Aggregate the evaluation results
-    results: List[Tuple[ClientProxy, EvaluateRes]] = []
-    failures: List[Union[Tuple[ClientProxy, EvaluateRes], BaseException]] = []
+    results: list[tuple[ClientProxy, EvaluateRes]] = []
+    failures: list[Union[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/secaggplus_workflow.py

@@ -18,7 +18,7 @@
 import random
 from dataclasses import dataclass, field
 from logging import DEBUG, ERROR, INFO, WARN
-from typing import Dict, List, Optional, Set, Tuple, Union, cast
+from typing import Optional, Union, cast
 
 import flwr.common.recordset_compat as compat
 from flwr.common import (
@@ -65,23 +65,23 @@ from ..constant import Key as WorkflowKey
 class WorkflowState:  # pylint: disable=R0902
     """The state of the SecAgg+ protocol."""
 
-    nid_to_proxies: Dict[int, ClientProxy] = field(default_factory=dict)
-    nid_to_fitins: Dict[int, RecordSet] = field(default_factory=dict)
-    sampled_node_ids: Set[int] = field(default_factory=set)
-    active_node_ids: Set[int] = field(default_factory=set)
+    nid_to_proxies: dict[int, ClientProxy] = field(default_factory=dict)
+    nid_to_fitins: dict[int, RecordSet] = field(default_factory=dict)
+    sampled_node_ids: set[int] = field(default_factory=set)
+    active_node_ids: set[int] = field(default_factory=set)
     num_shares: int = 0
     threshold: int = 0
     clipping_range: float = 0.0
     quantization_range: int = 0
     mod_range: int = 0
     max_weight: float = 0.0
-    nid_to_neighbours: Dict[int, Set[int]] = field(default_factory=dict)
-    nid_to_publickeys: Dict[int, List[bytes]] = field(default_factory=dict)
-    forward_srcs: Dict[int, List[int]] = field(default_factory=dict)
-    forward_ciphertexts: Dict[int, List[bytes]] = field(default_factory=dict)
+    nid_to_neighbours: dict[int, set[int]] = field(default_factory=dict)
+    nid_to_publickeys: dict[int, list[bytes]] = field(default_factory=dict)
+    forward_srcs: dict[int, list[int]] = field(default_factory=dict)
+    forward_ciphertexts: dict[int, list[bytes]] = field(default_factory=dict)
     aggregate_ndarrays: NDArrays = field(default_factory=list)
-    legacy_results: List[Tuple[ClientProxy, FitRes]] = field(default_factory=list)
-    failures: List[Exception] = field(default_factory=list)
+    legacy_results: list[tuple[ClientProxy, FitRes]] = field(default_factory=list)
+    failures: list[Exception] = field(default_factory=list)
 
 
 class SecAggPlusWorkflow:
@@ -444,13 +444,13 @@ class SecAggPlusWorkflow:
         )
 
         # Build forward packet list dictionary
-        srcs: List[int] = []
-        dsts: List[int] = []
-        ciphertexts: List[bytes] = []
-        fwd_ciphertexts: Dict[int, List[bytes]] = {
+        srcs: list[int] = []
+        dsts: list[int] = []
+        ciphertexts: list[bytes] = []
+        fwd_ciphertexts: dict[int, list[bytes]] = {
             nid: [] for nid in state.active_node_ids
         }  # dest node ID -> list of ciphertexts
-        fwd_srcs: Dict[int, List[int]] = {
+        fwd_srcs: dict[int, list[int]] = {
             nid: [] for nid in state.active_node_ids
         }  # dest node ID -> list of src node IDs
         for msg in msgs:
@@ -459,8 +459,8 @@ class SecAggPlusWorkflow:
                 continue
             node_id = msg.metadata.src_node_id
             res_dict = msg.content.configs_records[RECORD_KEY_CONFIGS]
-            dst_lst = cast(List[int], res_dict[Key.DESTINATION_LIST])
-            ctxt_lst = cast(List[bytes], res_dict[Key.CIPHERTEXT_LIST])
+            dst_lst = cast(list[int], res_dict[Key.DESTINATION_LIST])
+            ctxt_lst = cast(list[bytes], res_dict[Key.CIPHERTEXT_LIST])
             srcs += [node_id] * len(dst_lst)
             dsts += dst_lst
             ciphertexts += ctxt_lst
@@ -525,7 +525,7 @@ class SecAggPlusWorkflow:
                 state.failures.append(Exception(msg.error))
                 continue
             res_dict = msg.content.configs_records[RECORD_KEY_CONFIGS]
-            bytes_list = cast(List[bytes], res_dict[Key.MASKED_PARAMETERS])
+            bytes_list = cast(list[bytes], res_dict[Key.MASKED_PARAMETERS])
             client_masked_vec = [bytes_to_ndarray(b) for b in bytes_list]
             if masked_vector is None:
                 masked_vector = client_masked_vec
@@ -592,7 +592,7 @@ class SecAggPlusWorkflow:
         )
 
         # Build collected shares dict
-        collected_shares_dict: Dict[int, List[bytes]] = {}
+        collected_shares_dict: dict[int, list[bytes]] = {}
         for nid in state.sampled_node_ids:
             collected_shares_dict[nid] = []
         for msg in msgs:
@@ -600,8 +600,8 @@ class SecAggPlusWorkflow:
                 state.failures.append(Exception(msg.error))
                 continue
             res_dict = msg.content.configs_records[RECORD_KEY_CONFIGS]
-            nids = cast(List[int], res_dict[Key.NODE_ID_LIST])
-            shares = cast(List[bytes], res_dict[Key.SHARE_LIST])
+            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):
                 collected_shares_dict[owner_nid].append(share)