feat(framework) Generate error TaskRes for invalid cases in LinkState.get_task_res

src/py/flwr/common/constant.py

@@ -147,6 +147,8 @@ class ErrorCode:
     UNKNOWN = 0
     LOAD_CLIENT_APP_EXCEPTION = 1
     CLIENT_APP_RAISED_EXCEPTION = 2
+    MESSAGE_UNAVAILABLE = 3
+    REPLY_MESSAGE_UNAVAILABLE = 4
 
     def __new__(cls) -> ErrorCode:
         """Prevent instantiation."""

src/py/flwr/server/driver/inmemory_driver_test.py

@@ -277,6 +277,6 @@ def test_task_store_consistency_after_push_pull_inmemory_state(self) -> None:
         reply_tos = get_replies(self.driver, msg_ids, node_id)
 
         # Assert
-        self.assertEqual(reply_tos, msg_ids)
+        self.assertEqual(set(reply_tos), set(msg_ids))
         self.assertEqual(len(state.task_res_store), 0)
         self.assertEqual(len(state.task_ins_store), 0)

src/py/flwr/server/superlink/linkstate/in_memory_linkstate.py

@@ -40,6 +40,8 @@
     generate_rand_int_from_bytes,
     has_valid_sub_status,
     is_valid_transition,
+    verify_found_taskres,
+    verify_taskins_ids,
 )
 
 
@@ -67,12 +69,13 @@ def __init__(self) -> None:
         self.federation_options: dict[int, ConfigsRecord] = {}
         self.task_ins_store: dict[UUID, TaskIns] = {}
         self.task_res_store: dict[UUID, TaskRes] = {}
+        self.task_ins_id_to_task_res_id: dict[UUID, UUID] = {}
 
         self.node_public_keys: set[bytes] = set()
         self.server_public_key: Optional[bytes] = None
         self.server_private_key: Optional[bytes] = None
 
-        self.lock = threading.Lock()
+        self.lock = threading.RLock()
 
     def store_task_ins(self, task_ins: TaskIns) -> Optional[UUID]:
         """Store one TaskIns."""
@@ -222,42 +225,50 @@ def store_task_res(self, task_res: TaskRes) -> Optional[UUID]:
         task_res.task_id = str(task_id)
         with self.lock:
             self.task_res_store[task_id] = task_res
+            self.task_ins_id_to_task_res_id[UUID(task_ins_id)] = task_id
 
         # Return the new task_id
         return task_id
 
     def get_task_res(self, task_ids: set[UUID]) -> list[TaskRes]:
-        """Get all TaskRes that have not been delivered yet."""
+        """Get TaskRes for the given TaskIns IDs."""
+        ret: dict[UUID, TaskRes] = {}
+
         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():
-                reply_to = UUID(task_res.task.ancestry[0])
-
-                # Check if corresponding TaskIns exists and is not expired
-                task_ins = self.task_ins_store.get(reply_to)
-                if task_ins is None:
-                    log(WARNING, "TaskIns with task_id %s does not exist.", reply_to)
-                    task_ids.remove(reply_to)
-                    continue
-
-                if task_ins.task.created_at + task_ins.task.ttl <= time.time():
-                    log(WARNING, "TaskIns with task_id %s is expired.", reply_to)
-                    task_ids.remove(reply_to)
-                    continue
-
-                if reply_to in task_ids and task_res.task.delivered_at == "":
-                    task_res_list.append(task_res)
-                    replied_task_ids.add(reply_to)
-
-            # Mark all of them as delivered
+            current = time.time()
+
+            # Verify TaskIns IDs
+            ret = verify_taskins_ids(
+                inquired_taskins_ids=task_ids,
+                found_taskins_dict=self.task_ins_store,
+                current_time=current,
+            )
+
+            # Find all TaskRes
+            task_res_found: list[TaskRes] = []
+            for task_id in task_ids:
+                # If TaskRes exists and is not delivered, add it to the list
+                if task_res_id := self.task_ins_id_to_task_res_id.get(task_id):
+                    task_res = self.task_res_store[task_res_id]
+                    if task_res.task.delivered_at == "":
+                        task_res_found.append(task_res)
+            tmp_ret_dict = verify_found_taskres(
+                inquired_taskins_ids=task_ids,
+                found_taskins_dict=self.task_ins_store,
+                found_taskres_list=task_res_found,
+                current_time=current,
+            )
+            ret.update(tmp_ret_dict)
+
+            # Mark existing TaskRes to be returned as delivered
             delivered_at = now().isoformat()
-            for task_res in task_res_list:
+            for task_res in task_res_found:
                 task_res.task.delivered_at = delivered_at
 
-            # Return TaskRes
-            return task_res_list
+            # Cleanup
+            self._force_delete_tasks_by_ids(set(ret.keys()))
+
+        return list(ret.values())
 
     def delete_tasks(self, task_ids: set[UUID]) -> None:
         """Delete all delivered TaskIns/TaskRes pairs."""
@@ -278,9 +289,25 @@ def delete_tasks(self, task_ids: set[UUID]) -> None:
 
             for task_id in task_ins_to_be_deleted:
                 del self.task_ins_store[task_id]
+                del self.task_ins_id_to_task_res_id[task_id]
             for task_id in task_res_to_be_deleted:
                 del self.task_res_store[task_id]
 
+    def _force_delete_tasks_by_ids(self, task_ids: set[UUID]) -> None:
+        """Delete tasks based on a set of TaskIns IDs."""
+        if not task_ids:
+            return
+
+        with self.lock:
+            for task_id in task_ids:
+                # Delete TaskIns
+                if task_id in self.task_ins_store:
+                    del self.task_ins_store[task_id]
+                # Delete TaskRes
+                if task_id in self.task_ins_id_to_task_res_id:
+                    task_res_id = self.task_ins_id_to_task_res_id.pop(task_id)
+                    del self.task_res_store[task_res_id]
+
     def num_task_ins(self) -> int:
         """Calculate the number of task_ins in store.
 

src/py/flwr/server/superlink/linkstate/linkstate.py

@@ -101,13 +101,27 @@ def store_task_res(self, task_res: TaskRes) -> Optional[UUID]:
 
     @abc.abstractmethod
     def get_task_res(self, task_ids: set[UUID]) -> list[TaskRes]:
-        """Get TaskRes for task_ids.
+        """Get TaskRes for the given TaskIns IDs.
 
-        Usually, the ServerAppIo API calls this method to get results for instructions
-        it has previously scheduled.
+        This method is typically called by the ServerAppIo API to obtain
+        results (TaskRes) for previously scheduled instructions (TaskIns).
+        For each task_id provided, this method returns one of the following responses:
 
-        Retrieves all TaskRes for the given `task_ids` and returns and empty list of
-        none could be found.
+        - An error TaskRes if the corresponding TaskIns does not exist or has expired.
+        - An error TaskRes if the corresponding TaskRes exists but has expired.
+        - The valid TaskRes if the TaskIns has a corresponding valid TaskRes.
+        - Nothing if the TaskIns is still valid and waiting for a TaskRes.
+
+        Parameters
+        ----------
+        task_ids : set[UUID]
+            A set of TaskIns IDs for which to retrieve results (TaskRes).
+
+        Returns
+        -------
+        list[TaskRes]
+            A list of TaskRes corresponding to the given task IDs. If no
+            TaskRes could be found for any of the task IDs, an empty list is returned.
         """
 
     @abc.abstractmethod

src/py/flwr/server/superlink/linkstate/linkstate_test.py

@@ -350,13 +350,6 @@ def test_store_and_delete_tasks(self) -> None:
         # - State has three TaskIns, all of them delivered
         # - State has two TaskRes, one of the delivered, the other not
 
-        assert state.num_task_ins() == 3
-        assert state.num_task_res() == 2
-
-        # Execute
-        state.delete_tasks(task_ids={task_id_0, task_id_1, task_id_2})
-
-        # Assert
         assert state.num_task_ins() == 2
         assert state.num_task_res() == 1
 
@@ -932,8 +925,8 @@ def test_get_task_ins_not_return_expired(self) -> None:
             task_ins_list = state.get_task_ins(node_id=1, limit=None)
             assert len(task_ins_list) == 0
 
-    def test_get_task_res_not_return_expired(self) -> None:
-        """Test get_task_res not to return TaskRes if its TaskIns is expired."""
+    def test_get_task_res_expired_task_ins(self) -> None:
+        """Test get_task_res to return error TaskRes if its TaskIns has expired."""
         # Prepare
         state = self.state_factory()
         node_id = state.create_node(1e3)
@@ -961,7 +954,9 @@ def test_get_task_res_not_return_expired(self) -> None:
             task_res_list = state.get_task_res(task_ids={task_id})
 
             # Assert
-            assert len(task_res_list) == 0
+            assert len(task_res_list) == 1
+            assert task_res_list[0].task.HasField("error")
+            assert state.num_task_ins() == state.num_task_res() == 0
 
     def test_get_task_res_returns_empty_for_missing_taskins(self) -> None:
         """Test that get_task_res returns an empty result when the corresponding TaskIns
@@ -983,7 +978,9 @@ def test_get_task_res_returns_empty_for_missing_taskins(self) -> None:
         task_res_list = state.get_task_res(task_ids={UUID(task_ins_id)})
 
         # Assert
-        assert len(task_res_list) == 0
+        assert len(task_res_list) == 1
+        assert task_res_list[0].task.HasField("error")
+        assert state.num_task_ins() == state.num_task_res() == 0
 
     def test_get_task_res_return_if_not_expired(self) -> None:
         """Test get_task_res to return TaskRes if its TaskIns exists and is not