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

flwr/server/superlink/linkstate/sqlite_linkstate.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.
@@ -24,25 +24,23 @@ import time
 from collections.abc import Sequence
 from logging import DEBUG, ERROR, WARNING
 from typing import Any, Optional, Union, cast
-from uuid import UUID, uuid4
 
 from flwr.common import Context, Message, Metadata, log, now
 from flwr.common.constant import (
+    HEARTBEAT_MAX_INTERVAL,
+    HEARTBEAT_PATIENCE,
     MESSAGE_TTL_TOLERANCE,
     NODE_ID_NUM_BYTES,
-    PING_PATIENCE,
+    RUN_FAILURE_DETAILS_NO_HEARTBEAT,
     RUN_ID_NUM_BYTES,
     SUPERLINK_NODE_ID,
     Status,
+    SubStatus,
 )
 from flwr.common.message import make_message
 from flwr.common.record import ConfigRecord
-from flwr.common.serde import (
-    error_from_proto,
-    error_to_proto,
-    recorddict_from_proto,
-    recorddict_to_proto,
-)
+from flwr.common.serde import recorddict_from_proto, recorddict_to_proto
+from flwr.common.serde_utils import error_from_proto, error_to_proto
 from flwr.common.typing import Run, RunStatus, UserConfig
 
 # pylint: disable=E0611
@@ -74,7 +72,7 @@ SQL_CREATE_TABLE_NODE = """
 CREATE TABLE IF NOT EXISTS node(
     node_id         INTEGER UNIQUE,
     online_until    REAL,
-    ping_interval   REAL,
+    heartbeat_interval   REAL,
     public_key      BLOB
 );
 """
@@ -92,6 +90,8 @@ CREATE INDEX IF NOT EXISTS idx_online_until ON node (online_until);
 SQL_CREATE_TABLE_RUN = """
 CREATE TABLE IF NOT EXISTS run(
     run_id                INTEGER UNIQUE,
+    active_until          REAL,
+    heartbeat_interval    REAL,
     fab_id                TEXT,
     fab_version           TEXT,
     fab_hash              TEXT,
@@ -102,7 +102,8 @@ CREATE TABLE IF NOT EXISTS run(
     finished_at           TEXT,
     sub_status            TEXT,
     details               TEXT,
-    federation_options    BLOB
+    federation_options    BLOB,
+    flwr_aid              TEXT
 );
 """
 
@@ -250,19 +251,15 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
 
         return result
 
-    def store_message_ins(self, message: Message) -> Optional[UUID]:
+    def store_message_ins(self, message: Message) -> Optional[str]:
         """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
@@ -302,7 +299,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         # This may need to be changed in the future version with more integrity checks.
         self.query(query, data)
 
-        return message_id
+        return message.metadata.message_id
 
     def get_message_ins(self, node_id: int, limit: Optional[int]) -> list[Message]:
         """Get all Messages that have not been delivered yet."""
@@ -365,7 +362,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
 
         return result
 
-    def store_message_res(self, message: Message) -> Optional[UUID]:
+    def store_message_res(self, message: Message) -> Optional[str]:
         """Store one Message."""
         # Validate message
         errors = validate_message(message=message, is_reply_message=True)
@@ -417,12 +414,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
             )
             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
@@ -441,12 +433,12 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
             log(ERROR, "`run` is invalid")
             return None
 
-        return message_id
+        return message.metadata.message_id
 
-    def get_message_res(self, message_ids: set[UUID]) -> list[Message]:
+    def get_message_res(self, message_ids: set[str]) -> list[Message]:
         """Get reply Messages for the given Message IDs."""
         # pylint: disable-msg=too-many-locals
-        ret: dict[UUID, Message] = {}
+        ret: dict[str, Message] = {}
 
         # Verify Message IDs
         current = time.time()
@@ -456,12 +448,12 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
             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] = {}
+        found_message_ins_dict: dict[str, 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)
+            found_message_ins_dict[row["message_id"]] = dict_to_message(row)
 
         ret = verify_message_ids(
             inquired_message_ids=message_ids,
@@ -550,7 +542,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         result: dict[str, int] = rows[0]
         return result["num"]
 
-    def delete_messages(self, message_ins_ids: set[UUID]) -> None:
+    def delete_messages(self, message_ins_ids: set[str]) -> None:
         """Delete a Message and its reply based on provided Message IDs."""
         if not message_ins_ids:
             return
@@ -576,7 +568,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
             self.conn.execute(query_1, data)
             self.conn.execute(query_2, data)
 
-    def get_message_ids_from_run_id(self, run_id: int) -> set[UUID]:
+    def get_message_ids_from_run_id(self, run_id: int) -> set[str]:
         """Get all instruction Message IDs for the given run_id."""
         if self.conn is None:
             raise AttributeError("LinkState not initialized")
@@ -593,9 +585,9 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         with self.conn:
             rows = self.conn.execute(query, data).fetchall()
 
-        return {UUID(row["message_id"]) for row in rows}
+        return {row["message_id"] for row in rows}
 
-    def create_node(self, ping_interval: float) -> int:
+    def create_node(self, heartbeat_interval: float) -> int:
         """Create, store in the link state, and return `node_id`."""
         # Sample a random uint64 as node_id
         uint64_node_id = generate_rand_int_from_bytes(
@@ -607,18 +599,18 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
 
         query = (
             "INSERT INTO node "
-            "(node_id, online_until, ping_interval, public_key) "
+            "(node_id, online_until, heartbeat_interval, public_key) "
             "VALUES (?, ?, ?, ?)"
         )
 
-        # Mark the node online util time.time() + ping_interval
+        # Mark the node online util time.time() + heartbeat_interval
         try:
             self.query(
                 query,
                 (
                     sint64_node_id,
-                    time.time() + ping_interval,
-                    ping_interval,
+                    time.time() + heartbeat_interval,
+                    heartbeat_interval,
                     b"",  # Initialize with an empty public key
                 ),
             )
@@ -728,6 +720,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         fab_hash: Optional[str],
         override_config: UserConfig,
         federation_options: ConfigRecord,
+        flwr_aid: Optional[str],
     ) -> int:
         """Create a new run for the specified `fab_id` and `fab_version`."""
         # Sample a random int64 as run_id
@@ -742,26 +735,28 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         if self.query(query, (sint64_run_id,))[0]["COUNT(*)"] == 0:
             query = (
                 "INSERT INTO run "
-                "(run_id, fab_id, fab_version, fab_hash, override_config, "
-                "federation_options, pending_at, starting_at, running_at, finished_at, "
-                "sub_status, details) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?);"
+                "(run_id, active_until, heartbeat_interval, fab_id, fab_version, "
+                "fab_hash, override_config, federation_options, pending_at, "
+                "starting_at, running_at, finished_at, sub_status, details, flwr_aid) "
+                "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?);"
             )
             override_config_json = json.dumps(override_config)
             data = [
                 sint64_run_id,
+                0,  # The `active_until` is not used until the run is started
+                0,  # This `heartbeat_interval` is not used until the run is started
                 fab_id,
                 fab_version,
                 fab_hash,
                 override_config_json,
                 configrecord_to_bytes(federation_options),
-            ]
-            data += [
                 now().isoformat(),
                 "",
                 "",
                 "",
                 "",
                 "",
+                flwr_aid or "",
             ]
             self.query(query, tuple(data))
             return uint64_run_id
@@ -790,14 +785,47 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         result: set[bytes] = {row["public_key"] for row in rows}
         return result
 
-    def get_run_ids(self) -> set[int]:
-        """Retrieve all run IDs."""
-        query = "SELECT run_id FROM run;"
-        rows = self.query(query)
+    def get_run_ids(self, flwr_aid: Optional[str]) -> set[int]:
+        """Retrieve all run IDs if `flwr_aid` is not specified.
+
+        Otherwise, retrieve all run IDs for the specified `flwr_aid`.
+        """
+        if flwr_aid:
+            rows = self.query(
+                "SELECT run_id FROM run WHERE flwr_aid = ?;",
+                (flwr_aid,),
+            )
+        else:
+            rows = self.query("SELECT run_id FROM run;", ())
         return {convert_sint64_to_uint64(row["run_id"]) for row in rows}
 
+    def _check_and_tag_inactive_run(self, run_ids: set[int]) -> None:
+        """Check if any runs are no longer active.
+
+        Marks runs with status 'starting' or 'running' as failed
+        if they have not sent a heartbeat before `active_until`.
+        """
+        sint_run_ids = [convert_uint64_to_sint64(run_id) for run_id in run_ids]
+        query = "UPDATE run SET finished_at = ?, sub_status = ?, details = ? "
+        query += "WHERE starting_at != '' AND finished_at = '' AND active_until < ?"
+        query += f" AND run_id IN ({','.join(['?'] * len(run_ids))});"
+        current = now()
+        self.query(
+            query,
+            (
+                current.isoformat(),
+                SubStatus.FAILED,
+                RUN_FAILURE_DETAILS_NO_HEARTBEAT,
+                current.timestamp(),
+                *sint_run_ids,
+            ),
+        )
+
     def get_run(self, run_id: int) -> Optional[Run]:
         """Retrieve information about the run with the specified `run_id`."""
+        # Check if runs are still active
+        self._check_and_tag_inactive_run(run_ids={run_id})
+
         # Convert the uint64 value to sint64 for SQLite
         sint64_run_id = convert_uint64_to_sint64(run_id)
         query = "SELECT * FROM run WHERE run_id = ?;"
@@ -819,12 +847,16 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
                     sub_status=row["sub_status"],
                     details=row["details"],
                 ),
+                flwr_aid=row["flwr_aid"],
             )
         log(ERROR, "`run_id` does not exist.")
         return None
 
     def get_run_status(self, run_ids: set[int]) -> dict[int, RunStatus]:
         """Retrieve the statuses for the specified runs."""
+        # Check if runs are still active
+        self._check_and_tag_inactive_run(run_ids=run_ids)
+
         # Convert the uint64 value to sint64 for SQLite
         sint64_run_ids = (convert_uint64_to_sint64(run_id) for run_id in set(run_ids))
         query = f"SELECT * FROM run WHERE run_id IN ({','.join(['?'] * len(run_ids))});"
@@ -842,6 +874,9 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
 
     def update_run_status(self, run_id: int, new_status: RunStatus) -> bool:
         """Update the status of the run with the specified `run_id`."""
+        # Check if runs are still active
+        self._check_and_tag_inactive_run(run_ids={run_id})
+
         # Convert the uint64 value to sint64 for SQLite
         sint64_run_id = convert_uint64_to_sint64(run_id)
         query = "SELECT * FROM run WHERE run_id = ?;"
@@ -879,9 +914,22 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
             return False
 
         # Update the status
-        query = "UPDATE run SET %s= ?, sub_status = ?, details = ? "
+        query = "UPDATE run SET %s= ?, sub_status = ?, details = ?, "
+        query += "active_until = ?, heartbeat_interval = ? "
         query += "WHERE run_id = ?;"
 
+        # Prepare data for query
+        # Initialize heartbeat_interval and active_until
+        # when switching to starting or running
+        current = now()
+        if new_status.status in (Status.STARTING, Status.RUNNING):
+            heartbeat_interval = HEARTBEAT_MAX_INTERVAL
+            active_until = current.timestamp() + heartbeat_interval
+        else:
+            heartbeat_interval = 0
+            active_until = 0
+
+        # Determine the timestamp field based on the new status
         timestamp_fld = ""
         if new_status.status == Status.STARTING:
             timestamp_fld = "starting_at"
@@ -891,10 +939,12 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
             timestamp_fld = "finished_at"
 
         data = (
-            now().isoformat(),
+            current.isoformat(),
             new_status.sub_status,
             new_status.details,
-            sint64_run_id,
+            active_until,
+            heartbeat_interval,
+            convert_uint64_to_sint64(run_id),
         )
         self.query(query % timestamp_fld, data)
         return True
@@ -926,11 +976,15 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         row = rows[0]
         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.
+    def acknowledge_node_heartbeat(
+        self, node_id: int, heartbeat_interval: float
+    ) -> bool:
+        """Acknowledge a heartbeat 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.
+        A node is considered online as long as it sends heartbeats within
+        the tolerated interval: HEARTBEAT_PATIENCE × heartbeat_interval.
+        HEARTBEAT_PATIENCE = N allows for N-1 missed heartbeat before
+        the node is marked as offline.
         """
         sint64_node_id = convert_uint64_to_sint64(node_id)
 
@@ -939,18 +993,58 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         if not self.query(query, (sint64_node_id,)):
             return False
 
-        # Update `online_until` and `ping_interval` for the given `node_id`
-        query = "UPDATE node SET online_until = ?, ping_interval = ? WHERE node_id = ?"
+        # Update `online_until` and `heartbeat_interval` for the given `node_id`
+        query = (
+            "UPDATE node SET online_until = ?, heartbeat_interval = ? WHERE node_id = ?"
+        )
         self.query(
             query,
             (
-                time.time() + PING_PATIENCE * ping_interval,
-                ping_interval,
+                time.time() + HEARTBEAT_PATIENCE * heartbeat_interval,
+                heartbeat_interval,
                 sint64_node_id,
             ),
         )
         return True
 
+    def acknowledge_app_heartbeat(self, run_id: int, heartbeat_interval: float) -> bool:
+        """Acknowledge a heartbeat received from a ServerApp for a given run.
+
+        A run with status `"running"` is considered alive as long as it sends heartbeats
+        within the tolerated interval: HEARTBEAT_PATIENCE × heartbeat_interval.
+        HEARTBEAT_PATIENCE = N allows for N-1 missed heartbeat before the run is
+        marked as `"completed:failed"`.
+        """
+        # Check if runs are still active
+        self._check_and_tag_inactive_run(run_ids={run_id})
+
+        # Search for the run
+        sint_run_id = convert_uint64_to_sint64(run_id)
+        query = "SELECT * FROM run WHERE run_id = ?;"
+        rows = self.query(query, (sint_run_id,))
+
+        if not rows:
+            log(ERROR, "`run_id` is invalid")
+            return False
+
+        # Check if the run is of status "running"/"starting"
+        row = rows[0]
+        status = determine_run_status(row)
+        if status not in (Status.RUNNING, Status.STARTING):
+            log(
+                ERROR,
+                'Cannot acknowledge heartbeat for run with status "%s"',
+                status,
+            )
+            return False
+
+        # Update the `active_until` and `heartbeat_interval` for the given run
+        active_until = now().timestamp() + HEARTBEAT_PATIENCE * heartbeat_interval
+        query = "UPDATE run SET active_until = ?, heartbeat_interval = ? "
+        query += "WHERE run_id = ?"
+        self.query(query, (active_until, heartbeat_interval, sint_run_id))
+        return True
+
     def get_serverapp_context(self, run_id: int) -> Optional[Context]:
         """Get the context for the specified `run_id`."""
         # Retrieve context if any

flwr/server/superlink/linkstate/utils.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.
@@ -17,10 +17,10 @@
 
 from os import urandom
 from typing import Optional
-from uuid import UUID, uuid4
 
 from flwr.common import ConfigRecord, Context, Error, Message, Metadata, now, serde
 from flwr.common.constant import (
+    HEARTBEAT_PATIENCE,
     SUPERLINK_NODE_ID,
     ErrorCode,
     MessageType,
@@ -56,8 +56,8 @@ 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."
+    "Error: Node Unavailable — The destination node failed to report a heartbeat "
+    f"within {HEARTBEAT_PATIENCE} × its expected interval."
 )
 
 
@@ -245,7 +245,7 @@ def create_message_error_unavailable_res_message(
     ttl = max(ins_metadata.ttl - (current_time - ins_metadata.created_at), 0)
     metadata = Metadata(
         run_id=ins_metadata.run_id,
-        message_id=str(uuid4()),
+        message_id="",
         src_node_id=SUPERLINK_NODE_ID,
         dst_node_id=SUPERLINK_NODE_ID,
         reply_to_message_id=ins_metadata.message_id,
@@ -255,7 +255,7 @@ def create_message_error_unavailable_res_message(
         ttl=ttl,
     )
 
-    return make_message(
+    msg = make_message(
         metadata=metadata,
         error=Error(
             code=(
@@ -270,30 +270,34 @@ def create_message_error_unavailable_res_message(
             ),
         ),
     )
+    msg.metadata.__dict__["_message_id"] = msg.object_id
+    return msg
 
 
-def create_message_error_unavailable_ins_message(reply_to_message_id: UUID) -> Message:
+def create_message_error_unavailable_ins_message(reply_to_message_id: str) -> Message:
     """Error to indicate that the enquired Message had expired before reply arrived or
     that it isn't found."""
     metadata = Metadata(
         run_id=0,  # Unknown
-        message_id=str(uuid4()),
+        message_id="",
         src_node_id=SUPERLINK_NODE_ID,
         dst_node_id=SUPERLINK_NODE_ID,
-        reply_to_message_id=str(reply_to_message_id),
+        reply_to_message_id=reply_to_message_id,
         group_id="",  # Unknown
         message_type=MessageType.SYSTEM,
         created_at=now().timestamp(),
         ttl=0,
     )
 
-    return make_message(
+    msg = make_message(
         metadata=metadata,
         error=Error(
             code=ErrorCode.MESSAGE_UNAVAILABLE,
             reason=MESSAGE_UNAVAILABLE_ERROR_REASON,
         ),
     )
+    msg.metadata.__dict__["_message_id"] = msg.object_id
+    return msg
 
 
 def message_ttl_has_expired(message_metadata: Metadata, current_time: float) -> bool:
@@ -302,18 +306,18 @@ def message_ttl_has_expired(message_metadata: Metadata, current_time: float) ->
 
 
 def verify_message_ids(
-    inquired_message_ids: set[UUID],
-    found_message_ins_dict: dict[UUID, Message],
+    inquired_message_ids: set[str],
+    found_message_ins_dict: dict[str, Message],
     current_time: Optional[float] = None,
     update_set: bool = True,
-) -> dict[UUID, Message]:
+) -> dict[str, Message]:
     """Verify found Messages and generate error Messages for invalid ones.
 
     Parameters
     ----------
-    inquired_message_ids : set[UUID]
+    inquired_message_ids : set[str]
         Set of Message IDs for which to generate error Message if invalid.
-    found_message_ins_dict : dict[UUID, Message]
+    found_message_ins_dict : dict[str, Message]
         Dictionary containing all found Message indexed by their IDs.
     current_time : Optional[float] (default: None)
         The current time to check for expiration. If set to `None`, the current time
@@ -324,7 +328,7 @@ def verify_message_ids(
 
     Returns
     -------
-    dict[UUID, Message]
+    dict[str, Message]
         A dictionary of error Message indexed by the corresponding ID of the message
         they are a reply of.
     """
@@ -344,19 +348,19 @@ def verify_message_ids(
 
 
 def verify_found_message_replies(
-    inquired_message_ids: set[UUID],
-    found_message_ins_dict: dict[UUID, Message],
+    inquired_message_ids: set[str],
+    found_message_ins_dict: dict[str, Message],
     found_message_res_list: list[Message],
     current_time: Optional[float] = None,
     update_set: bool = True,
-) -> dict[UUID, Message]:
+) -> dict[str, Message]:
     """Verify found Message replies and generate error Message for invalid ones.
 
     Parameters
     ----------
-    inquired_message_ids : set[UUID]
+    inquired_message_ids : set[str]
         Set of Message IDs for which to generate error Message if invalid.
-    found_message_ins_dict : dict[UUID, Message]
+    found_message_ins_dict : dict[str, Message]
         Dictionary containing all found instruction Messages indexed by their IDs.
     found_message_res_list : dict[Message, Message]
         List of found Message to be verified.
@@ -369,13 +373,13 @@ def verify_found_message_replies(
 
     Returns
     -------
-    dict[UUID, Message]
+    dict[str, Message]
         A dictionary of Message indexed by the corresponding Message ID.
     """
-    ret_dict: dict[UUID, Message] = {}
+    ret_dict: dict[str, 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_id)
+        message_ins_id = message_res.metadata.reply_to_message_id
         if update_set:
             inquired_message_ids.remove(message_ins_id)
         # Check if the reply Message has expired
@@ -389,21 +393,21 @@ def verify_found_message_replies(
 
 
 def check_node_availability_for_in_message(
-    inquired_in_message_ids: set[UUID],
-    found_in_message_dict: dict[UUID, Message],
+    inquired_in_message_ids: set[str],
+    found_in_message_dict: dict[str, Message],
     node_id_to_online_until: dict[int, float],
     current_time: Optional[float] = None,
     update_set: bool = True,
-) -> dict[UUID, Message]:
+) -> dict[str, 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]
+    inquired_in_message_ids : set[str]
         Set of Message IDs for which to check destination node availability.
-    found_in_message_dict : dict[UUID, Message]
+    found_in_message_dict : dict[str, 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.
@@ -416,7 +420,7 @@ def check_node_availability_for_in_message(
 
     Returns
     -------
-    dict[UUID, Message]
+    dict[str, Message]
         A dictionary of error Message indexed by the corresponding Message ID.
     """
     ret_dict = {}

flwr/server/superlink/serverappio/serverappio_grpc.py

@@ -28,6 +28,7 @@ from flwr.proto.serverappio_pb2_grpc import (  # pylint: disable=E0611
 )
 from flwr.server.superlink.ffs.ffs_factory import FfsFactory
 from flwr.server.superlink.linkstate import LinkStateFactory
+from flwr.supercore.object_store import ObjectStoreFactory
 
 from .serverappio_servicer import ServerAppIoServicer
 
@@ -36,6 +37,7 @@ def run_serverappio_api_grpc(
     address: str,
     state_factory: LinkStateFactory,
     ffs_factory: FfsFactory,
+    objectstore_factory: ObjectStoreFactory,
     certificates: Optional[tuple[bytes, bytes, bytes]],
 ) -> grpc.Server:
     """Run ServerAppIo API (gRPC, request-response)."""
@@ -43,6 +45,7 @@ def run_serverappio_api_grpc(
     serverappio_servicer: grpc.Server = ServerAppIoServicer(
         state_factory=state_factory,
         ffs_factory=ffs_factory,
+        objectstore_factory=objectstore_factory,
     )
     serverappio_add_servicer_to_server_fn = add_ServerAppIoServicer_to_server
     serverappio_grpc_server = generic_create_grpc_server(