flwr-nightly 1.16.0.dev20250304__py3-none-any.whl → 1.16.0.dev20250305__py3-none-any.whl

flwr/server/superlink/linkstate/sqlite_linkstate.py

@@ -26,7 +26,7 @@ from logging import DEBUG, ERROR, WARNING
 from typing import Any, Optional, Union, cast
 from uuid import UUID, uuid4
 
-from flwr.common import Context, log, now
+from flwr.common import Context, Message, Metadata, log, now
 from flwr.common.constant import (
     MESSAGE_TTL_TOLERANCE,
     NODE_ID_NUM_BYTES,
@@ -35,15 +35,22 @@ from flwr.common.constant import (
     Status,
 )
 from flwr.common.record import ConfigsRecord
+from flwr.common.serde import (
+    error_from_proto,
+    error_to_proto,
+    recordset_from_proto,
+    recordset_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.node_pb2 import Node
 from flwr.proto.recordset_pb2 import RecordSet as ProtoRecordSet
 from flwr.proto.task_pb2 import Task, TaskIns, TaskRes
 
 # pylint: enable=E0611
-from flwr.server.utils.validator import validate_task_ins_or_res
+from flwr.server.utils.validator import validate_message, validate_task_ins_or_res
 
 from .linkstate import LinkState
 from .utils import (
@@ -58,7 +65,9 @@ from .utils import (
     generate_rand_int_from_bytes,
     has_valid_sub_status,
     is_valid_transition,
+    verify_found_message_replies,
     verify_found_taskres,
+    verify_message_ids,
     verify_taskins_ids,
 )
 
@@ -134,6 +143,24 @@ CREATE TABLE IF NOT EXISTS task_ins(
 );
 """
 
+SQL_CREATE_TABLE_MESSAGE_INS = """
+CREATE TABLE IF NOT EXISTS message_ins(
+    message_id              TEXT UNIQUE,
+    group_id                TEXT,
+    run_id                  INTEGER,
+    src_node_id             INTEGER,
+    dst_node_id             INTEGER,
+    reply_to_message        TEXT,
+    created_at              REAL,
+    delivered_at            TEXT,
+    ttl                     REAL,
+    message_type            TEXT,
+    content                 BLOB NULL,
+    error                   BLOB NULL,
+    FOREIGN KEY(run_id) REFERENCES run(run_id)
+);
+"""
+
 SQL_CREATE_TABLE_TASK_RES = """
 CREATE TABLE IF NOT EXISTS task_res(
     task_id                 TEXT UNIQUE,
@@ -151,6 +178,25 @@ CREATE TABLE IF NOT EXISTS task_res(
 );
 """
 
+
+SQL_CREATE_TABLE_MESSAGE_RES = """
+CREATE TABLE IF NOT EXISTS message_res(
+    message_id              TEXT UNIQUE,
+    group_id                TEXT,
+    run_id                  INTEGER,
+    src_node_id             INTEGER,
+    dst_node_id             INTEGER,
+    reply_to_message        TEXT,
+    created_at              REAL,
+    delivered_at            TEXT,
+    ttl                     REAL,
+    message_type            TEXT,
+    content                 BLOB NULL,
+    error                   BLOB NULL,
+    FOREIGN KEY(run_id) REFERENCES run(run_id)
+);
+"""
+
 DictOrTuple = Union[tuple[Any, ...], dict[str, Any]]
 
 
@@ -198,6 +244,8 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         cur.execute(SQL_CREATE_TABLE_CONTEXT)
         cur.execute(SQL_CREATE_TABLE_TASK_INS)
         cur.execute(SQL_CREATE_TABLE_TASK_RES)
+        cur.execute(SQL_CREATE_TABLE_MESSAGE_INS)
+        cur.execute(SQL_CREATE_TABLE_MESSAGE_RES)
         cur.execute(SQL_CREATE_TABLE_NODE)
         cur.execute(SQL_CREATE_TABLE_PUBLIC_KEY)
         cur.execute(SQL_CREATE_INDEX_ONLINE_UNTIL)
@@ -302,6 +350,60 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
 
         return task_id
 
+    def store_message_ins(self, message: Message) -> Optional[UUID]:
+        """Store one Message."""
+        # Validate message
+        errors = validate_message(message=message, is_reply_message=False)
+        if any(errors):
+            log(ERROR, errors)
+            return None
+        # Create message_id
+        message_id = uuid4()
+
+        # Store Message
+        # pylint: disable-next=W0212
+        message.metadata._message_id = str(message_id)  # type: ignore
+        data = (message_to_dict(message),)
+
+        # Convert values from uint64 to sint64 for SQLite
+        convert_uint64_values_in_dict_to_sint64(
+            data[0], ["run_id", "src_node_id", "dst_node_id"]
+        )
+
+        # Validate run_id
+        query = "SELECT run_id FROM run WHERE run_id = ?;"
+        if not self.query(query, (data[0]["run_id"],)):
+            log(ERROR, "Invalid run ID for Message: %s", message.metadata.run_id)
+            return None
+
+        # Validate source node ID
+        if message.metadata.src_node_id != SUPERLINK_NODE_ID:
+            log(
+                ERROR,
+                "Invalid source node ID for Message: %s",
+                message.metadata.src_node_id,
+            )
+            return None
+
+        # Validate destination node ID
+        query = "SELECT node_id FROM node WHERE node_id = ?;"
+        if not self.query(query, (data[0]["dst_node_id"],)):
+            log(
+                ERROR,
+                "Invalid destination node ID for Message: %s",
+                message.metadata.dst_node_id,
+            )
+            return None
+
+        columns = ", ".join([f":{key}" for key in data[0]])
+        query = f"INSERT INTO message_ins VALUES({columns});"
+
+        # Only invalid run_id can trigger IntegrityError.
+        # This may need to be changed in the future version with more integrity checks.
+        self.query(query, data)
+
+        return message_id
+
     def get_task_ins(self, node_id: int, limit: Optional[int]) -> list[TaskIns]:
         """Get undelivered TaskIns for one node.
 
@@ -380,6 +482,67 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
 
         return result
 
+    def get_message_ins(self, node_id: int, limit: Optional[int]) -> list[Message]:
+        """Get all Messages that have not been delivered yet."""
+        if limit is not None and limit < 1:
+            raise AssertionError("`limit` must be >= 1")
+
+        if node_id == SUPERLINK_NODE_ID:
+            msg = f"`node_id` must be != {SUPERLINK_NODE_ID}"
+            raise AssertionError(msg)
+
+        data: dict[str, Union[str, int]] = {}
+
+        # Convert the uint64 value to sint64 for SQLite
+        data["node_id"] = convert_uint64_to_sint64(node_id)
+
+        # Retrieve all Messages for node_id
+        query = """
+            SELECT message_id
+            FROM message_ins
+            WHERE   dst_node_id == :node_id
+            AND   delivered_at = ""
+            AND   (created_at + ttl) > CAST(strftime('%s', 'now') AS REAL)
+        """
+
+        if limit is not None:
+            query += " LIMIT :limit"
+            data["limit"] = limit
+
+        query += ";"
+
+        rows = self.query(query, data)
+
+        if rows:
+            # Prepare query
+            message_ids = [row["message_id"] for row in rows]
+            placeholders: str = ",".join([f":id_{i}" for i in range(len(message_ids))])
+            query = f"""
+                UPDATE message_ins
+                SET delivered_at = :delivered_at
+                WHERE message_id IN ({placeholders})
+                RETURNING *;
+            """
+
+            # Prepare data for query
+            delivered_at = now().isoformat()
+            data = {"delivered_at": delivered_at}
+            for index, msg_id in enumerate(message_ids):
+                data[f"id_{index}"] = str(msg_id)
+
+            # Run query
+            rows = self.query(query, data)
+
+        for row in rows:
+            # Convert values from sint64 to uint64
+            convert_sint64_values_in_dict_to_uint64(
+                row, ["run_id", "src_node_id", "dst_node_id"]
+            )
+
+        result = [dict_to_message(row) for row in rows]
+
+        return result
+
     def store_task_res(self, task_res: TaskRes) -> Optional[UUID]:
         """Store one TaskRes.
 
@@ -464,6 +627,84 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
 
         return task_id
 
+    def store_message_res(self, message: Message) -> Optional[UUID]:
+        """Store one Message."""
+        # Validate message
+        errors = validate_message(message=message, is_reply_message=True)
+        if any(errors):
+            log(ERROR, errors)
+            return None
+
+        res_metadata = message.metadata
+        msg_ins_id = res_metadata.reply_to_message
+        msg_ins = self.get_valid_message_ins(msg_ins_id)
+        if msg_ins is None:
+            log(
+                ERROR,
+                "Failed to store Message reply: "
+                "The message it replies to with message_id %s does not exist or "
+                "has expired.",
+                msg_ins_id,
+            )
+            return None
+
+        # Ensure that the dst_node_id of the original message matches the src_node_id of
+        # reply being processed.
+        if (
+            msg_ins
+            and message
+            and convert_sint64_to_uint64(msg_ins["dst_node_id"])
+            != res_metadata.src_node_id
+        ):
+            return None
+
+        # Fail if the Message TTL exceeds the
+        # expiration time of the Message it replies to.
+        # Condition: ins_metadata.created_at + ins_metadata.ttl ≥
+        #            res_metadata.created_at + res_metadata.ttl
+        # A small tolerance is introduced to account
+        # for floating-point precision issues.
+        max_allowed_ttl = (
+            msg_ins["created_at"] + msg_ins["ttl"] - res_metadata.created_at
+        )
+        if res_metadata.ttl and (
+            res_metadata.ttl - max_allowed_ttl > MESSAGE_TTL_TOLERANCE
+        ):
+            log(
+                WARNING,
+                "Received Message with TTL %.2f exceeding the allowed maximum "
+                "TTL %.2f.",
+                res_metadata.ttl,
+                max_allowed_ttl,
+            )
+            return None
+
+        # Create message_id
+        message_id = uuid4()
+
+        # Store Message
+        # pylint: disable-next=W0212
+        message.metadata._message_id = str(message_id)  # type: ignore
+        data = (message_to_dict(message),)
+
+        # Convert values from uint64 to sint64 for SQLite
+        convert_uint64_values_in_dict_to_sint64(
+            data[0], ["run_id", "src_node_id", "dst_node_id"]
+        )
+
+        columns = ", ".join([f":{key}" for key in data[0]])
+        query = f"INSERT INTO message_res VALUES({columns});"
+
+        # Only invalid run_id can trigger IntegrityError.
+        # This may need to be changed in the future version with more integrity checks.
+        try:
+            self.query(query, data)
+        except sqlite3.IntegrityError:
+            log(ERROR, "`run` is invalid")
+            return None
+
+        return message_id
+
     # pylint: disable-next=R0912,R0915,R0914
     def get_task_res(self, task_ids: set[UUID]) -> list[TaskRes]:
         """Get TaskRes for the given TaskIns IDs."""
@@ -525,6 +766,68 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
 
         return list(ret.values())
 
+    def get_message_res(self, message_ids: set[UUID]) -> list[Message]:
+        """Get reply Messages for the given Message IDs."""
+        ret: dict[UUID, Message] = {}
+
+        # Verify Message IDs
+        current = time.time()
+        query = f"""
+            SELECT *
+            FROM message_ins
+            WHERE message_id IN ({",".join(["?"] * len(message_ids))});
+        """
+        rows = self.query(query, tuple(str(message_id) for message_id in message_ids))
+        found_message_ins_dict: dict[UUID, Message] = {}
+        for row in rows:
+            convert_sint64_values_in_dict_to_uint64(
+                row, ["run_id", "src_node_id", "dst_node_id"]
+            )
+            found_message_ins_dict[UUID(row["message_id"])] = dict_to_message(row)
+
+        ret = verify_message_ids(
+            inquired_message_ids=message_ids,
+            found_message_ins_dict=found_message_ins_dict,
+            current_time=current,
+        )
+
+        # Find all reply Messages
+        query = f"""
+            SELECT *
+            FROM message_res
+            WHERE reply_to_message IN ({",".join(["?"] * len(message_ids))})
+            AND delivered_at = "";
+        """
+        rows = self.query(query, tuple(str(message_id) for message_id in message_ids))
+        for row in rows:
+            convert_sint64_values_in_dict_to_uint64(
+                row, ["run_id", "src_node_id", "dst_node_id"]
+            )
+        tmp_ret_dict = verify_found_message_replies(
+            inquired_message_ids=message_ids,
+            found_message_ins_dict=found_message_ins_dict,
+            found_message_res_list=[dict_to_message(row) for row in rows],
+            current_time=current,
+        )
+        ret.update(tmp_ret_dict)
+
+        # Mark existing reply Messages to be returned as delivered
+        delivered_at = now().isoformat()
+        for message_res in ret.values():
+            message_res.metadata.delivered_at = delivered_at
+        message_res_ids = [
+            message_res.metadata.message_id for message_res in ret.values()
+        ]
+        query = f"""
+            UPDATE message_res
+            SET delivered_at = ?
+            WHERE message_id IN ({",".join(["?"] * len(message_res_ids))});
+        """
+        data: list[Any] = [delivered_at] + message_res_ids
+        self.query(query, data)
+
+        return list(ret.values())
+
     def num_task_ins(self) -> int:
         """Calculate the number of task_ins in store.
 
@@ -536,6 +839,17 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         num = cast(int, result["num"])
         return num
 
+    def num_message_ins(self) -> int:
+        """Calculate the number of instruction Messages in store.
+
+        This includes delivered but not yet deleted.
+        """
+        query = "SELECT count(*) AS num FROM message_ins;"
+        rows = self.query(query)
+        result = rows[0]
+        num = cast(int, result["num"])
+        return num
+
     def num_task_res(self) -> int:
         """Calculate the number of task_res in store.
 
@@ -546,6 +860,16 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         result: dict[str, int] = rows[0]
         return result["num"]
 
+    def num_message_res(self) -> int:
+        """Calculate the number of reply Messages in store.
+
+        This includes delivered but not yet deleted.
+        """
+        query = "SELECT count(*) AS num FROM message_res;"
+        rows = self.query(query)
+        result: dict[str, int] = rows[0]
+        return result["num"]
+
     def delete_tasks(self, task_ins_ids: set[UUID]) -> None:
         """Delete TaskIns/TaskRes pairs based on provided TaskIns IDs."""
         if not task_ins_ids:
@@ -572,6 +896,32 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
             self.conn.execute(query_1, data)
             self.conn.execute(query_2, data)
 
+    def delete_messages(self, message_ins_ids: set[UUID]) -> None:
+        """Delete a Message and its reply based on provided Message IDs."""
+        if not message_ins_ids:
+            return
+        if self.conn is None:
+            raise AttributeError("LinkState not initialized")
+
+        placeholders = ",".join(["?"] * len(message_ins_ids))
+        data = tuple(str(message_id) for message_id in message_ins_ids)
+
+        # Delete Message
+        query_1 = f"""
+            DELETE FROM message_ins
+            WHERE message_id IN ({placeholders});
+        """
+
+        # Delete reply Message
+        query_2 = f"""
+            DELETE FROM message_res
+            WHERE reply_to_message IN ({placeholders});
+        """
+
+        with self.conn:
+            self.conn.execute(query_1, data)
+            self.conn.execute(query_2, data)
+
     def get_task_ids_from_run_id(self, run_id: int) -> set[UUID]:
         """Get all TaskIns IDs for the given run_id."""
         if self.conn is None:
@@ -591,6 +941,25 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
 
         return {UUID(row["task_id"]) for row in rows}
 
+    def get_message_ids_from_run_id(self, run_id: int) -> set[UUID]:
+        """Get all instruction Message IDs for the given run_id."""
+        if self.conn is None:
+            raise AttributeError("LinkState not initialized")
+
+        query = """
+            SELECT message_id
+            FROM message_ins
+            WHERE run_id = :run_id;
+        """
+
+        sint64_run_id = convert_uint64_to_sint64(run_id)
+        data = {"run_id": sint64_run_id}
+
+        with self.conn:
+            rows = self.conn.execute(query, data).fetchall()
+
+        return {UUID(row["message_id"]) for row in rows}
+
     def create_node(self, ping_interval: float) -> int:
         """Create, store in the link state, and return `node_id`."""
         # Sample a random uint64 as node_id
@@ -1028,6 +1397,33 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
 
         return task_ins
 
+    def get_valid_message_ins(self, message_id: str) -> Optional[dict[str, Any]]:
+        """Check if the Message exists and is valid (not expired).
+
+        Return Message if valid.
+        """
+        query = """
+            SELECT *
+            FROM message_ins
+            WHERE message_id = :message_id
+        """
+        data = {"message_id": message_id}
+        rows = self.query(query, data)
+        if not rows:
+            # Message does not exist
+            return None
+
+        message_ins = rows[0]
+        created_at = message_ins["created_at"]
+        ttl = message_ins["ttl"]
+        current_time = time.time()
+
+        # Check if TaskIns is expired
+        if ttl is not None and created_at + ttl <= current_time:
+            return None
+
+        return message_ins
+
 
 def dict_factory(
     cursor: sqlite3.Cursor,
@@ -1077,6 +1473,31 @@ def task_res_to_dict(task_msg: TaskRes) -> dict[str, Any]:
     return result
 
 
+def message_to_dict(message: Message) -> dict[str, Any]:
+    """Transform Message to dict."""
+    result = {
+        "message_id": message.metadata.message_id,
+        "group_id": message.metadata.group_id,
+        "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,
+        "created_at": message.metadata.created_at,
+        "delivered_at": message.metadata.delivered_at,
+        "ttl": message.metadata.ttl,
+        "message_type": message.metadata.message_type,
+        "content": None,
+        "error": None,
+    }
+
+    if message.has_content():
+        result["content"] = recordset_to_proto(message.content).SerializeToString()
+    else:
+        result["error"] = error_to_proto(message.error).SerializeToString()
+
+    return result
+
+
 def dict_to_task_ins(task_dict: dict[str, Any]) -> TaskIns:
     """Turn task_dict into protobuf message."""
     recordset = ProtoRecordSet()
@@ -1131,6 +1552,28 @@ def dict_to_task_res(task_dict: dict[str, Any]) -> TaskRes:
     return result
 
 
+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))
+    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"]
+        }
+    )
+    msg = Message(metadata=metadata, content=content, error=error)
+    msg.metadata.__dict__["_created_at"] = message_dict["created_at"]
+    msg.metadata.delivered_at = message_dict["delivered_at"]
+    return msg
+
+
 def determine_run_status(row: dict[str, Any]) -> str:
     """Determine the status of the run based on timestamp fields."""
     if row["pending_at"]: