flwr 1.16.0__py3-none-any.whl → 1.17.0__py3-none-any.whl

flwr/server/superlink/linkstate/sqlite_linkstate.py

@@ -30,30 +30,33 @@ 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,
 )
-from flwr.common.record import ConfigsRecord
+from flwr.common.message import make_message
+from flwr.common.record import ConfigRecord
 from flwr.common.serde import (
     error_from_proto,
     error_to_proto,
-    recordset_from_proto,
-    recordset_to_proto,
+    recorddict_from_proto,
+    recorddict_to_proto,
 )
 from flwr.common.typing import Run, RunStatus, UserConfig
 
 # pylint: disable=E0611
 from flwr.proto.error_pb2 import Error as ProtoError
-from flwr.proto.recordset_pb2 import RecordSet as ProtoRecordSet
+from flwr.proto.recorddict_pb2 import RecordDict as ProtoRecordDict
 
 # pylint: enable=E0611
 from flwr.server.utils.validator import validate_message
 
 from .linkstate import LinkState
 from .utils import (
-    configsrecord_from_bytes,
-    configsrecord_to_bytes,
+    check_node_availability_for_in_message,
+    configrecord_from_bytes,
+    configrecord_to_bytes,
     context_from_bytes,
     context_to_bytes,
     convert_sint64_to_uint64,
@@ -129,7 +132,7 @@ CREATE TABLE IF NOT EXISTS message_ins(
     run_id                  INTEGER,
     src_node_id             INTEGER,
     dst_node_id             INTEGER,
-    reply_to_message        TEXT,
+    reply_to_message_id     TEXT,
     created_at              REAL,
     delivered_at            TEXT,
     ttl                     REAL,
@@ -148,7 +151,7 @@ CREATE TABLE IF NOT EXISTS message_res(
     run_id                  INTEGER,
     src_node_id             INTEGER,
     dst_node_id             INTEGER,
-    reply_to_message        TEXT,
+    reply_to_message_id     TEXT,
     created_at              REAL,
     delivered_at            TEXT,
     ttl                     REAL,
@@ -371,7 +374,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
             return None
 
         res_metadata = message.metadata
-        msg_ins_id = res_metadata.reply_to_message
+        msg_ins_id = res_metadata.reply_to_message_id
         msg_ins = self.get_valid_message_ins(msg_ins_id)
         if msg_ins is None:
             log(
@@ -442,6 +445,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,11 +469,34 @@ 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 *
             FROM message_res
-            WHERE reply_to_message IN ({",".join(["?"] * len(message_ids))})
+            WHERE reply_to_message_id IN ({",".join(["?"] * len(message_ids))})
             AND delivered_at = "";
         """
         rows = self.query(query, tuple(str(message_id) for message_id in message_ids))
@@ -542,7 +569,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         # Delete reply Message
         query_2 = f"""
             DELETE FROM message_res
-            WHERE reply_to_message IN ({placeholders});
+            WHERE reply_to_message_id IN ({placeholders});
         """
 
         with self.conn:
@@ -584,6 +611,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
             "VALUES (?, ?, ?, ?)"
         )
 
+        # Mark the node online util time.time() + ping_interval
         try:
             self.query(
                 query,
@@ -699,7 +727,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         fab_version: Optional[str],
         fab_hash: Optional[str],
         override_config: UserConfig,
-        federation_options: ConfigsRecord,
+        federation_options: ConfigRecord,
     ) -> int:
         """Create a new run for the specified `fab_id` and `fab_version`."""
         # Sample a random int64 as run_id
@@ -725,7 +753,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
                 fab_version,
                 fab_hash,
                 override_config_json,
-                configsrecord_to_bytes(federation_options),
+                configrecord_to_bytes(federation_options),
             ]
             data += [
                 now().isoformat(),
@@ -883,7 +911,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
 
         return pending_run_id
 
-    def get_federation_options(self, run_id: int) -> Optional[ConfigsRecord]:
+    def get_federation_options(self, run_id: int) -> Optional[ConfigRecord]:
         """Retrieve the federation options for the specified `run_id`."""
         # Convert the uint64 value to sint64 for SQLite
         sint64_run_id = convert_uint64_to_sint64(run_id)
@@ -896,10 +924,14 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
             return None
 
         row = rows[0]
-        return configsrecord_from_bytes(row["federation_options"])
+        return configrecord_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 +941,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]:
@@ -1026,7 +1065,7 @@ def message_to_dict(message: Message) -> dict[str, Any]:
         "run_id": message.metadata.run_id,
         "src_node_id": message.metadata.src_node_id,
         "dst_node_id": message.metadata.dst_node_id,
-        "reply_to_message": message.metadata.reply_to_message,
+        "reply_to_message_id": message.metadata.reply_to_message_id,
         "created_at": message.metadata.created_at,
         "delivered_at": message.metadata.delivered_at,
         "ttl": message.metadata.ttl,
@@ -1036,7 +1075,7 @@ def message_to_dict(message: Message) -> dict[str, Any]:
     }
 
     if message.has_content():
-        result["content"] = recordset_to_proto(message.content).SerializeToString()
+        result["content"] = recorddict_to_proto(message.content).SerializeToString()
     else:
         result["error"] = error_to_proto(message.error).SerializeToString()
 
@@ -1047,20 +1086,15 @@ def dict_to_message(message_dict: dict[str, Any]) -> Message:
     """Transform dict to Message."""
     content, error = None, None
     if (b_content := message_dict.pop("content")) is not None:
-        content = recordset_from_proto(ProtoRecordSet.FromString(b_content))
+        content = recorddict_from_proto(ProtoRecordDict.FromString(b_content))
     if (b_error := message_dict.pop("error")) is not None:
         error = error_from_proto(ProtoError.FromString(b_error))
 
     # Metadata constructor doesn't allow passing created_at. We set it later
     metadata = Metadata(
-        **{
-            k: v
-            for k, v in message_dict.items()
-            if k not in ["created_at", "delivered_at"]
-        }
+        **{k: v for k, v in message_dict.items() if k not in ["delivered_at"]}
     )
-    msg = Message(metadata=metadata, content=content, error=error)
-    msg.metadata.__dict__["_created_at"] = message_dict["created_at"]
+    msg = make_message(metadata=metadata, content=content, error=error)
     msg.metadata.delivered_at = message_dict["delivered_at"]
     return msg
 

flwr/server/superlink/linkstate/utils.py

@@ -19,21 +19,22 @@ from os import urandom
 from typing import Optional
 from uuid import UUID, uuid4
 
-from flwr.common import ConfigsRecord, Context, Error, Message, Metadata, now, serde
-from flwr.common.constant import SUPERLINK_NODE_ID, ErrorCode, Status, SubStatus
+from flwr.common import ConfigRecord, Context, Error, Message, Metadata, now, serde
+from flwr.common.constant import (
+    SUPERLINK_NODE_ID,
+    ErrorCode,
+    MessageType,
+    Status,
+    SubStatus,
+)
+from flwr.common.message import make_message
 from flwr.common.typing import RunStatus
 
 # pylint: disable=E0611
 from flwr.proto.message_pb2 import Context as ProtoContext
-from flwr.proto.recordset_pb2 import ConfigsRecord as ProtoConfigsRecord
+from flwr.proto.recorddict_pb2 import ConfigRecord as ProtoConfigRecord
 
 # 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),
@@ -54,6 +55,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(
@@ -167,15 +172,15 @@ def context_from_bytes(context_bytes: bytes) -> Context:
     return serde.context_from_proto(ProtoContext.FromString(context_bytes))
 
 
-def configsrecord_to_bytes(configs_record: ConfigsRecord) -> bytes:
-    """Serialize a `ConfigsRecord` to bytes."""
-    return serde.configs_record_to_proto(configs_record).SerializeToString()
+def configrecord_to_bytes(config_record: ConfigRecord) -> bytes:
+    """Serialize a `ConfigRecord` to bytes."""
+    return serde.config_record_to_proto(config_record).SerializeToString()
 
 
-def configsrecord_from_bytes(configsrecord_bytes: bytes) -> ConfigsRecord:
-    """Deserialize `ConfigsRecord` from bytes."""
-    return serde.configs_record_from_proto(
-        ProtoConfigsRecord.FromString(configsrecord_bytes)
+def configrecord_from_bytes(configrecord_bytes: bytes) -> ConfigRecord:
+    """Deserialize `ConfigRecord` from bytes."""
+    return serde.config_record_from_proto(
+        ProtoConfigRecord.FromString(configrecord_bytes)
     )
 
 
@@ -231,7 +236,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()
@@ -241,22 +248,31 @@ def create_message_error_unavailable_res_message(ins_metadata: Metadata) -> Mess
         message_id=str(uuid4()),
         src_node_id=SUPERLINK_NODE_ID,
         dst_node_id=SUPERLINK_NODE_ID,
-        reply_to_message=ins_metadata.message_id,
+        reply_to_message_id=ins_metadata.message_id,
         group_id=ins_metadata.group_id,
         message_type=ins_metadata.message_type,
+        created_at=current_time,
         ttl=ttl,
     )
 
-    return Message(
+    return make_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
+            ),
         ),
     )
 
 
-def create_message_error_unavailable_ins_message(reply_to_message: UUID) -> Message:
+def create_message_error_unavailable_ins_message(reply_to_message_id: UUID) -> Message:
     """Error to indicate that the enquired Message had expired before reply arrived or
     that it isn't found."""
     metadata = Metadata(
@@ -264,13 +280,14 @@ def create_message_error_unavailable_ins_message(reply_to_message: UUID) -> Mess
         message_id=str(uuid4()),
         src_node_id=SUPERLINK_NODE_ID,
         dst_node_id=SUPERLINK_NODE_ID,
-        reply_to_message=str(reply_to_message),
+        reply_to_message_id=str(reply_to_message_id),
         group_id="",  # Unknown
-        message_type="",  # Unknown
+        message_type=MessageType.SYSTEM,
+        created_at=now().timestamp(),
         ttl=0,
     )
 
-    return Message(
+    return make_message(
         metadata=metadata,
         error=Error(
             code=ErrorCode.MESSAGE_UNAVAILABLE,
@@ -358,14 +375,63 @@ def verify_found_message_replies(
     ret_dict: dict[UUID, Message] = {}
     current = current_time if current_time else now().timestamp()
     for message_res in found_message_res_list:
-        message_ins_id = UUID(message_res.metadata.reply_to_message)
+        message_ins_id = UUID(message_res.metadata.reply_to_message_id)
         if update_set:
             inquired_message_ids.remove(message_ins_id)
         # Check if the reply Message has expired
         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/superlink/{driver → serverappio}/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2024 Flower Labs GmbH. All Rights Reserved.
+# Copyright 2025 Flower Labs GmbH. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

flwr/server/superlink/{driver → serverappio}/serverappio_grpc.py

@@ -1,4 +1,4 @@
-# Copyright 2024 Flower Labs GmbH. All Rights Reserved.
+# Copyright 2025 Flower Labs GmbH. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

flwr/server/superlink/{driver → serverappio}/serverappio_servicer.py

@@ -1,4 +1,4 @@
-# Copyright 2024 Flower Labs GmbH. All Rights Reserved.
+# Copyright 2025 Flower Labs GmbH. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -22,7 +22,7 @@ from uuid import UUID
 
 import grpc
 
-from flwr.common import ConfigsRecord, Message
+from flwr.common import ConfigRecord, Message
 from flwr.common.constant import SUPERLINK_NODE_ID, Status
 from flwr.common.logger import log
 from flwr.common.serde import (
@@ -127,7 +127,7 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
             request.fab_version,
             fab_hash,
             user_config_from_proto(request.override_config),
-            ConfigsRecord(),
+            ConfigRecord(),
         )
         return CreateRunResponse(run_id=run_id)
 
@@ -206,7 +206,7 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
 
         # Delete the instruction Messages and their replies if found
         message_ins_ids_to_delete = {
-            UUID(msg_res.metadata.reply_to_message) for msg_res in messages_res
+            UUID(msg_res.metadata.reply_to_message_id) for msg_res in messages_res
         }
 
         state.delete_messages(message_ins_ids=message_ins_ids_to_delete)

flwr/server/superlink/simulation/simulationio_servicer.py

@@ -24,7 +24,7 @@ from grpc import ServicerContext
 from flwr.common.constant import Status
 from flwr.common.logger import log
 from flwr.common.serde import (
-    configs_record_to_proto,
+    config_record_to_proto,
     context_from_proto,
     context_to_proto,
     fab_to_proto,
@@ -182,5 +182,5 @@ class SimulationIoServicer(simulationio_pb2_grpc.SimulationIoServicer):
             )
             return GetFederationOptionsResponse()
         return GetFederationOptionsResponse(
-            federation_options=configs_record_to_proto(federation_options)
+            federation_options=config_record_to_proto(federation_options)
         )

flwr/server/typing.py

@@ -19,9 +19,9 @@ from typing import Callable
 
 from flwr.common import Context
 
-from .driver 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/utils/validator.py

@@ -68,8 +68,8 @@ def validate_message(message: Message, is_reply_message: bool) -> list[str]:
 
     # Link respose to original message
     if not is_reply_message:
-        if metadata.reply_to_message != "":
-            validation_errors.append("`metadata.reply_to_message` MUST not be set.")
+        if metadata.reply_to_message_id != "":
+            validation_errors.append("`metadata.reply_to_message_id` MUST not be set.")
         if metadata.src_node_id != SUPERLINK_NODE_ID:
             validation_errors.append(
                 f"`metadata.src_node_id` is not {SUPERLINK_NODE_ID} (SuperLink node ID)"
@@ -79,8 +79,8 @@ def validate_message(message: Message, is_reply_message: bool) -> list[str]:
                 f"`metadata.dst_node_id` is {SUPERLINK_NODE_ID} (SuperLink node ID)"
             )
     else:
-        if metadata.reply_to_message == "":
-            validation_errors.append("`metadata.reply_to_message` MUST be set.")
+        if metadata.reply_to_message_id == "":
+            validation_errors.append("`metadata.reply_to_message_id` MUST be set.")
         if metadata.src_node_id == SUPERLINK_NODE_ID:
             validation_errors.append(
                 f"`metadata.src_node_id` is {SUPERLINK_NODE_ID} (SuperLink node ID)"