flwr-nightly 1.17.0.dev20250318__py3-none-any.whl → 1.17.0.dev20250319__py3-none-any.whl

flwr/server/serverapp/app.py

@@ -60,7 +60,7 @@ from flwr.proto.serverappio_pb2 import (  # pylint: disable=E0611
     PullServerAppInputsResponse,
     PushServerAppOutputsRequest,
 )
-from flwr.server.grid.grpc_grid import GrpcDriver
+from flwr.server.grid.grpc_grid import GrpcGrid
 from flwr.server.run_serverapp import run as run_
 
 
@@ -106,7 +106,7 @@ def run_serverapp(  # pylint: disable=R0914, disable=W0212, disable=R0915
     certificates: Optional[bytes] = None,
 ) -> None:
     """Run Flower ServerApp process."""
-    driver = GrpcDriver(
+    grid = GrpcGrid(
         serverappio_service_address=serverappio_api_address,
         root_certificates=certificates,
     )
@@ -123,7 +123,7 @@ def run_serverapp(  # pylint: disable=R0914, disable=W0212, disable=R0915
             # Pull ServerAppInputs from LinkState
             req = PullServerAppInputsRequest()
             log(DEBUG, "[flwr-serverapp] Pull ServerAppInputs")
-            res: PullServerAppInputsResponse = driver._stub.PullServerAppInputs(req)
+            res: PullServerAppInputsResponse = grid._stub.PullServerAppInputs(req)
             if not res.HasField("run"):
                 sleep(3)
                 run_status = None
@@ -135,14 +135,14 @@ def run_serverapp(  # pylint: disable=R0914, disable=W0212, disable=R0915
 
             hash_run_id = get_sha256_hash(run.run_id)
 
-            driver.set_run(run.run_id)
+            grid.set_run(run.run_id)
 
             # Start log uploader for this run
             log_uploader = start_log_uploader(
                 log_queue=log_queue,
                 node_id=0,
                 run_id=run.run_id,
-                stub=driver._stub,
+                stub=grid._stub,
             )
 
             log(DEBUG, "[flwr-serverapp] Start FAB installation.")
@@ -173,7 +173,7 @@ def run_serverapp(  # pylint: disable=R0914, disable=W0212, disable=R0915
 
             # Change status to Running
             run_status_proto = run_status_to_proto(RunStatus(Status.RUNNING, "", ""))
-            driver._stub.UpdateRunStatus(
+            grid._stub.UpdateRunStatus(
                 UpdateRunStatusRequest(run_id=run.run_id, run_status=run_status_proto)
             )
 
@@ -182,9 +182,9 @@ def run_serverapp(  # pylint: disable=R0914, disable=W0212, disable=R0915
                 event_details={"run-id-hash": hash_run_id},
             )
 
-            # Load and run the ServerApp with the Driver
+            # Load and run the ServerApp with the Grid
             updated_context = run_(
-                driver=driver,
+                grid=grid,
                 server_app_dir=app_path,
                 server_app_attr=server_app_attr,
                 context=context,
@@ -196,7 +196,7 @@ def run_serverapp(  # pylint: disable=R0914, disable=W0212, disable=R0915
             out_req = PushServerAppOutputsRequest(
                 run_id=run.run_id, context=context_proto
             )
-            _ = driver._stub.PushServerAppOutputs(out_req)
+            _ = grid._stub.PushServerAppOutputs(out_req)
 
             run_status = RunStatus(Status.FINISHED, SubStatus.COMPLETED, "")
         except RunNotRunningException:
@@ -221,7 +221,7 @@ def run_serverapp(  # pylint: disable=R0914, disable=W0212, disable=R0915
             # Update run status
             if run_status:
                 run_status_proto = run_status_to_proto(run_status)
-                driver._stub.UpdateRunStatus(
+                grid._stub.UpdateRunStatus(
                     UpdateRunStatusRequest(
                         run_id=run.run_id, run_status=run_status_proto
                     )

flwr/server/superlink/linkstate/in_memory_linkstate.py

@@ -27,6 +27,7 @@ from flwr.common import Context, Message, log, now
 from flwr.common.constant import (
     MESSAGE_TTL_TOLERANCE,
     NODE_ID_NUM_BYTES,
+    PING_PATIENCE,
     RUN_ID_NUM_BYTES,
     SUPERLINK_NODE_ID,
     Status,
@@ -37,6 +38,7 @@ from flwr.server.superlink.linkstate.linkstate import LinkState
 from flwr.server.utils import validate_message
 
 from .utils import (
+    check_node_availability_for_in_message,
     generate_rand_int_from_bytes,
     has_valid_sub_status,
     is_valid_transition,
@@ -232,13 +234,28 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         with self.lock:
             current = time.time()
 
-            # Verify Messge IDs
+            # Verify Message IDs
             ret = verify_message_ids(
                 inquired_message_ids=message_ids,
                 found_message_ins_dict=self.message_ins_store,
                 current_time=current,
             )
 
+            # Check node availability
+            dst_node_ids = {
+                self.message_ins_store[message_id].metadata.dst_node_id
+                for message_id in message_ids
+            }
+            tmp_ret_dict = check_node_availability_for_in_message(
+                inquired_in_message_ids=message_ids,
+                found_in_message_dict=self.message_ins_store,
+                node_id_to_online_until={
+                    node_id: self.node_ids[node_id][0] for node_id in dst_node_ids
+                },
+                current_time=current,
+            )
+            ret.update(tmp_ret_dict)
+
             # Find all reply Messages
             message_res_found: list[Message] = []
             for message_id in message_ids:
@@ -317,6 +334,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
                 log(ERROR, "Unexpected node registration failure.")
                 return 0
 
+            # Mark the node online util time.time() + ping_interval
             self.node_ids[node_id] = (time.time() + ping_interval, ping_interval)
             return node_id
 
@@ -519,10 +537,17 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
             return self.federation_options[run_id]
 
     def acknowledge_ping(self, node_id: int, ping_interval: float) -> bool:
-        """Acknowledge a ping received from a node, serving as a heartbeat."""
+        """Acknowledge a ping received from a node, serving as a heartbeat.
+
+        It allows for one missed ping (in a PING_PATIENCE * ping_interval) before
+        marking the node as offline, where PING_PATIENCE = 2 in default.
+        """
         with self.lock:
             if node_id in self.node_ids:
-                self.node_ids[node_id] = (time.time() + ping_interval, ping_interval)
+                self.node_ids[node_id] = (
+                    time.time() + PING_PATIENCE * ping_interval,
+                    ping_interval,
+                )
                 return True
         return False
 

flwr/server/superlink/linkstate/sqlite_linkstate.py

@@ -30,6 +30,7 @@ from flwr.common import Context, Message, Metadata, log, now
 from flwr.common.constant import (
     MESSAGE_TTL_TOLERANCE,
     NODE_ID_NUM_BYTES,
+    PING_PATIENCE,
     RUN_ID_NUM_BYTES,
     SUPERLINK_NODE_ID,
     Status,
@@ -52,6 +53,7 @@ from flwr.server.utils.validator import validate_message
 
 from .linkstate import LinkState
 from .utils import (
+    check_node_availability_for_in_message,
     configsrecord_from_bytes,
     configsrecord_to_bytes,
     context_from_bytes,
@@ -442,6 +444,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
 
     def get_message_res(self, message_ids: set[UUID]) -> list[Message]:
         """Get reply Messages for the given Message IDs."""
+        # pylint: disable-msg=too-many-locals
         ret: dict[UUID, Message] = {}
 
         # Verify Message IDs
@@ -465,6 +468,29 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
             current_time=current,
         )
 
+        # Check node availability
+        dst_node_ids: set[int] = set()
+        for message_id in message_ids:
+            in_message = found_message_ins_dict[message_id]
+            sint_node_id = convert_uint64_to_sint64(in_message.metadata.dst_node_id)
+            dst_node_ids.add(sint_node_id)
+        query = f"""
+                    SELECT node_id, online_until
+                    FROM node
+                    WHERE node_id IN ({",".join(["?"] * len(dst_node_ids))});
+                """
+        rows = self.query(query, tuple(dst_node_ids))
+        tmp_ret_dict = check_node_availability_for_in_message(
+            inquired_in_message_ids=message_ids,
+            found_in_message_dict=found_message_ins_dict,
+            node_id_to_online_until={
+                convert_sint64_to_uint64(row["node_id"]): row["online_until"]
+                for row in rows
+            },
+            current_time=current,
+        )
+        ret.update(tmp_ret_dict)
+
         # Find all reply Messages
         query = f"""
             SELECT *
@@ -584,6 +610,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
             "VALUES (?, ?, ?, ?)"
         )
 
+        # Mark the node online util time.time() + ping_interval
         try:
             self.query(
                 query,
@@ -899,7 +926,11 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         return configsrecord_from_bytes(row["federation_options"])
 
     def acknowledge_ping(self, node_id: int, ping_interval: float) -> bool:
-        """Acknowledge a ping received from a node, serving as a heartbeat."""
+        """Acknowledge a ping received from a node, serving as a heartbeat.
+
+        It allows for one missed ping (in a PING_PATIENCE * ping_interval) before
+        marking the node as offline, where PING_PATIENCE = 2 in default.
+        """
         sint64_node_id = convert_uint64_to_sint64(node_id)
 
         # Check if the node exists in the `node` table
@@ -909,7 +940,14 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
 
         # Update `online_until` and `ping_interval` for the given `node_id`
         query = "UPDATE node SET online_until = ?, ping_interval = ? WHERE node_id = ?"
-        self.query(query, (time.time() + ping_interval, ping_interval, sint64_node_id))
+        self.query(
+            query,
+            (
+                time.time() + PING_PATIENCE * ping_interval,
+                ping_interval,
+                sint64_node_id,
+            ),
+        )
         return True
 
     def get_serverapp_context(self, run_id: int) -> Optional[Context]:

flwr/server/superlink/linkstate/utils.py

@@ -34,12 +34,6 @@ from flwr.proto.message_pb2 import Context as ProtoContext
 from flwr.proto.recordset_pb2 import ConfigsRecord as ProtoConfigsRecord
 
 # pylint: enable=E0611
-
-NODE_UNAVAILABLE_ERROR_REASON = (
-    "Error: Node Unavailable - The destination node is currently unavailable. "
-    "It exceeds the time limit specified in its last ping."
-)
-
 VALID_RUN_STATUS_TRANSITIONS = {
     (Status.PENDING, Status.STARTING),
     (Status.STARTING, Status.RUNNING),
@@ -60,6 +54,10 @@ MESSAGE_UNAVAILABLE_ERROR_REASON = (
 REPLY_MESSAGE_UNAVAILABLE_ERROR_REASON = (
     "Error: Reply Message Unavailable - The reply message has expired."
 )
+NODE_UNAVAILABLE_ERROR_REASON = (
+    "Error: Node Unavailable - The destination node is currently unavailable. "
+    "It exceeds twice the time limit specified in its last ping."
+)
 
 
 def generate_rand_int_from_bytes(
@@ -237,7 +235,9 @@ def has_valid_sub_status(status: RunStatus) -> bool:
     return status.sub_status == ""
 
 
-def create_message_error_unavailable_res_message(ins_metadata: Metadata) -> Message:
+def create_message_error_unavailable_res_message(
+    ins_metadata: Metadata, error_type: str
+) -> Message:
     """Generate an error Message that the SuperLink returns carrying the specified
     error."""
     current_time = now().timestamp()
@@ -256,8 +256,16 @@ def create_message_error_unavailable_res_message(ins_metadata: Metadata) -> Mess
     return Message(
         metadata=metadata,
         error=Error(
-            code=ErrorCode.REPLY_MESSAGE_UNAVAILABLE,
-            reason=REPLY_MESSAGE_UNAVAILABLE_ERROR_REASON,
+            code=(
+                ErrorCode.REPLY_MESSAGE_UNAVAILABLE
+                if error_type == "msg_unavail"
+                else ErrorCode.NODE_UNAVAILABLE
+            ),
+            reason=(
+                REPLY_MESSAGE_UNAVAILABLE_ERROR_REASON
+                if error_type == "msg_unavail"
+                else NODE_UNAVAILABLE_ERROR_REASON
+            ),
         ),
     )
 
@@ -371,7 +379,56 @@ def verify_found_message_replies(
         if message_ttl_has_expired(message_res.metadata, current):
             # No need to insert the error Message
             message_res = create_message_error_unavailable_res_message(
-                found_message_ins_dict[message_ins_id].metadata
+                found_message_ins_dict[message_ins_id].metadata, "msg_unavail"
             )
         ret_dict[message_ins_id] = message_res
     return ret_dict
+
+
+def check_node_availability_for_in_message(
+    inquired_in_message_ids: set[UUID],
+    found_in_message_dict: dict[UUID, Message],
+    node_id_to_online_until: dict[int, float],
+    current_time: Optional[float] = None,
+    update_set: bool = True,
+) -> dict[UUID, Message]:
+    """Check node availability for given Message and generate error reply Message if
+    unavailable. A Message error indicating node unavailability will be generated for
+    each given Message whose destination node is offline or non-existent.
+
+    Parameters
+    ----------
+    inquired_in_message_ids : set[UUID]
+        Set of Message IDs for which to check destination node availability.
+    found_in_message_dict : dict[UUID, Message]
+        Dictionary containing all found Message indexed by their IDs.
+    node_id_to_online_until : dict[int, float]
+        Dictionary mapping node IDs to their online-until timestamps.
+    current_time : Optional[float] (default: None)
+        The current time to check for expiration. If set to `None`, the current time
+        will automatically be set to the current timestamp using `now().timestamp()`.
+    update_set : bool (default: True)
+        If True, the `inquired_in_message_ids` will be updated to remove invalid ones,
+        by default True.
+
+    Returns
+    -------
+    dict[UUID, Message]
+        A dictionary of error Message indexed by the corresponding Message ID.
+    """
+    ret_dict = {}
+    current = current_time if current_time else now().timestamp()
+    for in_message_id in list(inquired_in_message_ids):
+        in_message = found_in_message_dict[in_message_id]
+        node_id = in_message.metadata.dst_node_id
+        online_until = node_id_to_online_until.get(node_id)
+        # Generate a reply message containing an error reply
+        # if the node is offline or doesn't exist.
+        if online_until is None or online_until < current:
+            if update_set:
+                inquired_in_message_ids.remove(in_message_id)
+            reply_message = create_message_error_unavailable_res_message(
+                in_message.metadata, "node_unavail"
+            )
+            ret_dict[in_message_id] = reply_message
+    return ret_dict

flwr/server/typing.py

@@ -19,9 +19,9 @@ from typing import Callable
 
 from flwr.common import Context
 
-from .grid import Driver
+from .grid import Grid
 from .serverapp_components import ServerAppComponents
 
-ServerAppCallable = Callable[[Driver, Context], None]
-Workflow = Callable[[Driver, Context], None]
+ServerAppCallable = Callable[[Grid, Context], None]
+Workflow = Callable[[Grid, Context], None]
 ServerFn = Callable[[Context], ServerAppComponents]

flwr/server/workflow/default_workflows.py

@@ -36,7 +36,7 @@ from flwr.common.constant import MessageType, MessageTypeLegacy
 from ..client_proxy import ClientProxy
 from ..compat.app_utils import start_update_client_manager_thread
 from ..compat.legacy_context import LegacyContext
-from ..grid import Driver
+from ..grid import Grid
 from ..typing import Workflow
 from .constant import MAIN_CONFIGS_RECORD, MAIN_PARAMS_RECORD, Key
 
@@ -56,7 +56,7 @@ class DefaultWorkflow:
         self.fit_workflow: Workflow = fit_workflow
         self.evaluate_workflow: Workflow = evaluate_workflow
 
-    def __call__(self, driver: Driver, context: Context) -> None:
+    def __call__(self, grid: Grid, context: Context) -> None:
         """Execute the workflow."""
         if not isinstance(context, LegacyContext):
             raise TypeError(
@@ -65,7 +65,7 @@ class DefaultWorkflow:
 
         # Start the thread updating nodes
         thread, f_stop, c_done = start_update_client_manager_thread(
-            driver, context.client_manager
+            grid, context.client_manager
         )
 
         # Wait until the node registration done
@@ -73,7 +73,7 @@ class DefaultWorkflow:
 
         # Initialize parameters
         log(INFO, "[INIT]")
-        default_init_params_workflow(driver, context)
+        default_init_params_workflow(grid, context)
 
         # Run federated learning for num_rounds
         start_time = timeit.default_timer()
@@ -87,13 +87,13 @@ class DefaultWorkflow:
             cfg[Key.CURRENT_ROUND] = current_round
 
             # Fit round
-            self.fit_workflow(driver, context)
+            self.fit_workflow(grid, context)
 
             # Centralized evaluation
-            default_centralized_evaluation_workflow(driver, context)
+            default_centralized_evaluation_workflow(grid, context)
 
             # Evaluate round
-            self.evaluate_workflow(driver, context)
+            self.evaluate_workflow(grid, context)
 
         # Bookkeeping and log results
         end_time = timeit.default_timer()
@@ -119,7 +119,7 @@ class DefaultWorkflow:
         thread.join()
 
 
-def default_init_params_workflow(driver: Driver, context: Context) -> None:
+def default_init_params_workflow(grid: Grid, context: Context) -> None:
     """Execute the default workflow for parameters initialization."""
     if not isinstance(context, LegacyContext):
         raise TypeError(f"Expect a LegacyContext, but get {type(context).__name__}.")
@@ -138,9 +138,9 @@ def default_init_params_workflow(driver: Driver, context: Context) -> None:
         random_client = context.client_manager.sample(1)[0]
         # Send GetParametersIns and get the response
         content = compat.getparametersins_to_recordset(GetParametersIns({}))
-        messages = driver.send_and_receive(
+        messages = grid.send_and_receive(
             [
-                driver.create_message(
+                grid.create_message(
                     content=content,
                     message_type=MessageTypeLegacy.GET_PARAMETERS,
                     dst_node_id=random_client.node_id,
@@ -186,7 +186,7 @@ def default_init_params_workflow(driver: Driver, context: Context) -> None:
         log(INFO, "Evaluation returned no results (`None`)")
 
 
-def default_centralized_evaluation_workflow(_: Driver, context: Context) -> None:
+def default_centralized_evaluation_workflow(_: Grid, context: Context) -> None:
     """Execute the default workflow for centralized evaluation."""
     if not isinstance(context, LegacyContext):
         raise TypeError(f"Expect a LegacyContext, but get {type(context).__name__}.")
@@ -218,9 +218,7 @@ def default_centralized_evaluation_workflow(_: Driver, context: Context) -> None
         )
 
 
-def default_fit_workflow(  # pylint: disable=R0914
-    driver: Driver, context: Context
-) -> None:
+def default_fit_workflow(grid: Grid, context: Context) -> None:  # pylint: disable=R0914
     """Execute the default workflow for a single fit round."""
     if not isinstance(context, LegacyContext):
         raise TypeError(f"Expect a LegacyContext, but get {type(context).__name__}.")
@@ -255,7 +253,7 @@ def default_fit_workflow(  # pylint: disable=R0914
 
     # Build out messages
     out_messages = [
-        driver.create_message(
+        grid.create_message(
             content=compat.fitins_to_recordset(fitins, True),
             message_type=MessageType.TRAIN,
             dst_node_id=proxy.node_id,
@@ -266,7 +264,7 @@ def default_fit_workflow(  # pylint: disable=R0914
 
     # Send instructions to clients and
     # collect `fit` results from all clients participating in this round
-    messages = list(driver.send_and_receive(out_messages))
+    messages = list(grid.send_and_receive(out_messages))
     del out_messages
     num_failures = len([msg for msg in messages if msg.has_error()])
 
@@ -307,7 +305,7 @@ def default_fit_workflow(  # pylint: disable=R0914
 
 
 # pylint: disable-next=R0914
-def default_evaluate_workflow(driver: Driver, context: Context) -> None:
+def default_evaluate_workflow(grid: Grid, context: Context) -> None:
     """Execute the default workflow for a single evaluate round."""
     if not isinstance(context, LegacyContext):
         raise TypeError(f"Expect a LegacyContext, but get {type(context).__name__}.")
@@ -341,7 +339,7 @@ def default_evaluate_workflow(driver: Driver, context: Context) -> None:
 
     # Build out messages
     out_messages = [
-        driver.create_message(
+        grid.create_message(
             content=compat.evaluateins_to_recordset(evalins, True),
             message_type=MessageType.EVALUATE,
             dst_node_id=proxy.node_id,
@@ -352,7 +350,7 @@ def default_evaluate_workflow(driver: Driver, context: Context) -> None:
 
     # Send instructions to clients and
     # collect `evaluate` results from all clients participating in this round
-    messages = list(driver.send_and_receive(out_messages))
+    messages = list(grid.send_and_receive(out_messages))
     del out_messages
     num_failures = len([msg for msg in messages if msg.has_error()])
 

flwr/server/workflow/secure_aggregation/secaggplus_workflow.py

@@ -55,7 +55,7 @@ from flwr.common.secure_aggregation.secaggplus_constants import (
 from flwr.common.secure_aggregation.secaggplus_utils import pseudo_rand_gen
 from flwr.server.client_proxy import ClientProxy
 from flwr.server.compat.legacy_context import LegacyContext
-from flwr.server.grid import Driver
+from flwr.server.grid import Grid
 
 from ..constant import MAIN_CONFIGS_RECORD, MAIN_PARAMS_RECORD
 from ..constant import Key as WorkflowKey
@@ -186,7 +186,7 @@ class SecAggPlusWorkflow:
 
         self._check_init_params()
 
-    def __call__(self, driver: Driver, context: Context) -> None:
+    def __call__(self, grid: Grid, context: Context) -> None:
         """Run the SecAgg+ protocol."""
         if not isinstance(context, LegacyContext):
             raise TypeError(
@@ -202,7 +202,7 @@ class SecAggPlusWorkflow:
         )
         log(INFO, "Secure aggregation commencing.")
         for step in steps:
-            if not step(driver, context, state):
+            if not step(grid, context, state):
                 log(INFO, "Secure aggregation halted.")
                 return
         log(INFO, "Secure aggregation completed.")
@@ -279,7 +279,7 @@ class SecAggPlusWorkflow:
         return True
 
     def setup_stage(  # pylint: disable=R0912, R0914, R0915
-        self, driver: Driver, context: LegacyContext, state: WorkflowState
+        self, grid: Grid, context: LegacyContext, state: WorkflowState
     ) -> bool:
         """Execute the 'setup' stage."""
         # Obtain fit instructions
@@ -370,7 +370,7 @@ class SecAggPlusWorkflow:
         content = RecordSet({RECORD_KEY_CONFIGS: cfgs_record})
 
         def make(nid: int) -> Message:
-            return driver.create_message(
+            return grid.create_message(
                 content=content,
                 message_type=MessageType.TRAIN,
                 dst_node_id=nid,
@@ -382,7 +382,7 @@ class SecAggPlusWorkflow:
             "[Stage 0] Sending configurations to %s clients.",
             len(state.active_node_ids),
         )
-        msgs = driver.send_and_receive(
+        msgs = grid.send_and_receive(
             [make(node_id) for node_id in state.active_node_ids], timeout=self.timeout
         )
         state.active_node_ids = {
@@ -406,7 +406,7 @@ class SecAggPlusWorkflow:
         return self._check_threshold(state)
 
     def share_keys_stage(  # pylint: disable=R0914
-        self, driver: Driver, context: LegacyContext, state: WorkflowState
+        self, grid: Grid, context: LegacyContext, state: WorkflowState
     ) -> bool:
         """Execute the 'share keys' stage."""
         cfg = context.state.configs_records[MAIN_CONFIGS_RECORD]
@@ -418,7 +418,7 @@ class SecAggPlusWorkflow:
             )
             cfgs_record[Key.STAGE] = Stage.SHARE_KEYS
             content = RecordSet({RECORD_KEY_CONFIGS: cfgs_record})
-            return driver.create_message(
+            return grid.create_message(
                 content=content,
                 message_type=MessageType.TRAIN,
                 dst_node_id=nid,
@@ -431,7 +431,7 @@ class SecAggPlusWorkflow:
             "[Stage 1] Forwarding public keys to %s clients.",
             len(state.active_node_ids),
         )
-        msgs = driver.send_and_receive(
+        msgs = grid.send_and_receive(
             [make(node_id) for node_id in state.active_node_ids], timeout=self.timeout
         )
         state.active_node_ids = {
@@ -476,7 +476,7 @@ class SecAggPlusWorkflow:
         return self._check_threshold(state)
 
     def collect_masked_vectors_stage(
-        self, driver: Driver, context: LegacyContext, state: WorkflowState
+        self, grid: Grid, context: LegacyContext, state: WorkflowState
     ) -> bool:
         """Execute the 'collect masked vectors' stage."""
         cfg = context.state.configs_records[MAIN_CONFIGS_RECORD]
@@ -491,7 +491,7 @@ class SecAggPlusWorkflow:
             cfgs_record = ConfigsRecord(cfgs_dict)  # type: ignore
             content = state.nid_to_fitins[nid]
             content.configs_records[RECORD_KEY_CONFIGS] = cfgs_record
-            return driver.create_message(
+            return grid.create_message(
                 content=content,
                 message_type=MessageType.TRAIN,
                 dst_node_id=nid,
@@ -503,7 +503,7 @@ class SecAggPlusWorkflow:
             "[Stage 2] Forwarding encrypted key shares to %s clients.",
             len(state.active_node_ids),
         )
-        msgs = driver.send_and_receive(
+        msgs = grid.send_and_receive(
             [make(node_id) for node_id in state.active_node_ids], timeout=self.timeout
         )
         state.active_node_ids = {
@@ -547,7 +547,7 @@ class SecAggPlusWorkflow:
         return self._check_threshold(state)
 
     def unmask_stage(  # pylint: disable=R0912, R0914, R0915
-        self, driver: Driver, context: LegacyContext, state: WorkflowState
+        self, grid: Grid, context: LegacyContext, state: WorkflowState
     ) -> bool:
         """Execute the 'unmask' stage."""
         cfg = context.state.configs_records[MAIN_CONFIGS_RECORD]
@@ -567,7 +567,7 @@ class SecAggPlusWorkflow:
             }
             cfgs_record = ConfigsRecord(cfgs_dict)  # type: ignore
             content = RecordSet({RECORD_KEY_CONFIGS: cfgs_record})
-            return driver.create_message(
+            return grid.create_message(
                 content=content,
                 message_type=MessageType.TRAIN,
                 dst_node_id=nid,
@@ -579,7 +579,7 @@ class SecAggPlusWorkflow:
             "[Stage 3] Requesting key shares from %s clients to remove masks.",
             len(state.active_node_ids),
         )
-        msgs = driver.send_and_receive(
+        msgs = grid.send_and_receive(
             [make(node_id) for node_id in state.active_node_ids], timeout=self.timeout
         )
         state.active_node_ids = {