Merge branch 'main' into update-changelog

src/py/flwr/common/constant.py

@@ -83,6 +83,7 @@ class ErrorCode:
     UNKNOWN = 0
     LOAD_CLIENT_APP_EXCEPTION = 1
     CLIENT_APP_RAISED_EXCEPTION = 2
+    NODE_UNAVAILABLE = 3
 
     def __new__(cls) -> ErrorCode:
         """Prevent instantiation."""

src/py/flwr/server/superlink/state/in_memory_state.py

@@ -27,6 +27,8 @@
 from flwr.server.superlink.state.state import State
 from flwr.server.utils import validate_task_ins_or_res
 
+from .utils import make_node_unavailable_taskres
+
 
 class InMemoryState(State):
     """In-memory State implementation."""
@@ -129,14 +131,31 @@ def get_task_res(self, task_ids: Set[UUID], limit: Optional[int]) -> List[TaskRe
         with self.lock:
             # Find TaskRes that were not delivered yet
             task_res_list: List[TaskRes] = []
+            replied_task_ids: Set[UUID] = set()
             for _, task_res in self.task_res_store.items():
-                if (
-                    UUID(task_res.task.ancestry[0]) in task_ids
-                    and task_res.task.delivered_at == ""
-                ):
+                reply_to = UUID(task_res.task.ancestry[0])
+                if reply_to in task_ids and task_res.task.delivered_at == "":
                     task_res_list.append(task_res)
+                    replied_task_ids.add(reply_to)
+                if limit and len(task_res_list) == limit:
+                    break
+
+            # Check if the node is offline
+            for task_id in task_ids - replied_task_ids:
                 if limit and len(task_res_list) == limit:
                     break
+                task_ins = self.task_ins_store.get(task_id)
+                if task_ins is None:
+                    continue
+                node_id = task_ins.task.consumer.node_id
+                online_until, _ = self.node_ids[node_id]
+                # Generate a TaskRes containing an error reply if the node is offline.
+                if online_until < time.time():
+                    err_taskres = make_node_unavailable_taskres(
+                        ref_taskins=task_ins,
+                    )
+                    self.task_res_store[UUID(err_taskres.task_id)] = err_taskres
+                    task_res_list.append(err_taskres)
 
             # Mark all of them as delivered
             delivered_at = now().isoformat()

src/py/flwr/server/superlink/state/sqlite_state.py

@@ -30,6 +30,7 @@
 from flwr.server.utils.validator import validate_task_ins_or_res
 
 from .state import State
+from .utils import make_node_unavailable_taskres
 
 SQL_CREATE_TABLE_NODE = """
 CREATE TABLE IF NOT EXISTS node(
@@ -344,6 +345,7 @@ def store_task_res(self, task_res: TaskRes) -> Optional[UUID]:
 
         return task_id
 
+    # pylint: disable-next=R0914
     def get_task_res(self, task_ids: Set[UUID], limit: Optional[int]) -> List[TaskRes]:
         """Get TaskRes for task_ids.
 
@@ -374,7 +376,7 @@ def get_task_res(self, task_ids: Set[UUID], limit: Optional[int]) -> List[TaskRe
             AND delivered_at = ""
         """
 
-        data: Dict[str, Union[str, int]] = {}
+        data: Dict[str, Union[str, float, int]] = {}
 
         if limit is not None:
             query += " LIMIT :limit"
@@ -408,6 +410,54 @@ def get_task_res(self, task_ids: Set[UUID], limit: Optional[int]) -> List[TaskRe
             rows = self.query(query, data)
 
         result = [dict_to_task_res(row) for row in rows]
+
+        # 1. Query: Fetch consumer_node_id of remaining task_ids
+        # Assume the ancestry field only contains one element
+        data.clear()
+        replied_task_ids: Set[UUID] = {UUID(str(row["ancestry"])) for row in rows}
+        remaining_task_ids = task_ids - replied_task_ids
+        placeholders = ",".join([f":id_{i}" for i in range(len(remaining_task_ids))])
+        query = f"""
+            SELECT consumer_node_id
+            FROM task_ins
+            WHERE task_id IN ({placeholders});
+        """
+        for index, task_id in enumerate(remaining_task_ids):
+            data[f"id_{index}"] = str(task_id)
+        node_ids = [int(row["consumer_node_id"]) for row in self.query(query, data)]
+
+        # 2. Query: Select offline nodes
+        placeholders = ",".join([f":id_{i}" for i in range(len(node_ids))])
+        query = f"""
+            SELECT node_id
+            FROM node
+            WHERE node_id IN ({placeholders})
+            AND online_until < :time;
+        """
+        data = {f"id_{i}": str(node_id) for i, node_id in enumerate(node_ids)}
+        data["time"] = time.time()
+        offline_node_ids = [int(row["node_id"]) for row in self.query(query, data)]
+
+        # 3. Query: Select TaskIns for offline nodes
+        placeholders = ",".join([f":id_{i}" for i in range(len(offline_node_ids))])
+        query = f"""
+            SELECT *
+            FROM task_ins
+            WHERE consumer_node_id IN ({placeholders});
+        """
+        data = {f"id_{i}": str(node_id) for i, node_id in enumerate(offline_node_ids)}
+        task_ins_rows = self.query(query, data)
+
+        # Make TaskRes containing node unavailabe error
+        for row in task_ins_rows:
+            if limit and len(result) == limit:
+                break
+            task_ins = dict_to_task_ins(row)
+            err_taskres = make_node_unavailable_taskres(
+                ref_taskins=task_ins,
+            )
+            result.append(err_taskres)
+
         return result
 
     def num_task_ins(self) -> int:

src/py/flwr/server/superlink/state/state_test.py

@@ -25,6 +25,7 @@
 from uuid import uuid4
 
 from flwr.common import DEFAULT_TTL
+from flwr.common.constant import ErrorCode
 from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
 from flwr.proto.recordset_pb2 import RecordSet  # pylint: disable=E0611
 from flwr.proto.task_pb2 import Task, TaskIns, TaskRes  # pylint: disable=E0611
@@ -413,6 +414,48 @@ def test_acknowledge_ping(self) -> None:
         # Assert
         self.assertSetEqual(actual_node_ids, set(node_ids[70:]))
 
+    def test_node_unavailable_error(self) -> None:
+        """Test if get_task_res return TaskRes containing node unavailable error."""
+        # Prepare
+        state: State = self.state_factory()
+        run_id = state.create_run()
+        node_id_0 = state.create_node(ping_interval=90)
+        node_id_1 = state.create_node(ping_interval=30)
+        # Create and store TaskIns
+        task_ins_0 = create_task_ins(
+            consumer_node_id=node_id_0, anonymous=False, run_id=run_id
+        )
+        task_ins_1 = create_task_ins(
+            consumer_node_id=node_id_1, anonymous=False, run_id=run_id
+        )
+        task_id_0 = state.store_task_ins(task_ins=task_ins_0)
+        task_id_1 = state.store_task_ins(task_ins=task_ins_1)
+        assert task_id_0 is not None and task_id_1 is not None
+
+        # Get TaskIns to mark them delivered
+        state.get_task_ins(node_id=node_id_0, limit=None)
+
+        # Create and store TaskRes
+        task_res_0 = create_task_res(
+            producer_node_id=100,
+            anonymous=False,
+            ancestry=[str(task_id_0)],
+            run_id=run_id,
+        )
+        state.store_task_res(task_res_0)
+
+        # Execute
+        current_time = time.time()
+        task_res_list: List[TaskRes] = []
+        with patch("time.time", side_effect=lambda: current_time + 50):
+            task_res_list = state.get_task_res({task_id_0, task_id_1}, limit=None)
+
+        # Assert
+        assert len(task_res_list) == 2
+        err_taskres = task_res_list[1]
+        assert err_taskres.task.HasField("error")
+        assert err_taskres.task.error.code == ErrorCode.NODE_UNAVAILABLE
+
 
 def create_task_ins(
     consumer_node_id: int,

src/py/flwr/server/superlink/state/utils.py

@@ -0,0 +1,56 @@
+# Copyright 2024 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Utility functions for State."""
+
+
+import time
+from logging import ERROR
+from uuid import uuid4
+
+from flwr.common import log
+from flwr.common.constant import ErrorCode
+from flwr.proto.error_pb2 import Error  # pylint: disable=E0611
+from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
+from flwr.proto.task_pb2 import Task, TaskIns, TaskRes  # pylint: disable=E0611
+
+NODE_UNAVAILABLE_ERROR_REASON = (
+    "Error: Node Unavailable - The destination node is currently unavailable. "
+    "It exceeds the time limit specified in its last ping."
+)
+
+
+def make_node_unavailable_taskres(ref_taskins: TaskIns) -> TaskRes:
+    """Generate a TaskRes with a node unavailable error from a TaskIns."""
+    current_time = time.time()
+    ttl = ref_taskins.task.ttl - (current_time - ref_taskins.task.created_at)
+    if ttl < 0:
+        log(ERROR, "Creating TaskRes for TaskIns that exceeds its TTL.")
+        ttl = 0
+    return TaskRes(
+        task_id=str(uuid4()),
+        group_id=ref_taskins.group_id,
+        run_id=ref_taskins.run_id,
+        task=Task(
+            producer=Node(node_id=ref_taskins.task.consumer.node_id, anonymous=False),
+            consumer=Node(node_id=ref_taskins.task.producer.node_id, anonymous=False),
+            created_at=current_time,
+            ttl=ttl,
+            ancestry=[ref_taskins.task_id],
+            task_type=ref_taskins.task.task_type,
+            error=Error(
+                code=ErrorCode.NODE_UNAVAILABLE, reason=NODE_UNAVAILABLE_ERROR_REASON
+            ),
+        ),
+    )