Fix memory leak issue in SuperLink

adap · Feb 13, 2025 · 1d0cd13 · 1d0cd13
1 parent 50ac168
commit 1d0cd13
Showing 1 changed file with 7 additions and 7 deletions.
diff --git a/src/py/flwr/server/superlink/driver/serverappio_servicer.py b/src/py/flwr/server/superlink/driver/serverappio_servicer.py
@@ -208,6 +208,13 @@ def PullMessages(
         # Read from state
         task_res_list: list[TaskRes] = state.get_task_res(task_ids=message_ids)
 
+        # Delete the TaskIns/TaskRes pairs if TaskRes is found
+        task_ins_ids_to_delete = {
+            UUID(task_res.task.ancestry[0]) for task_res in task_res_list
+        }
+
+        state.delete_tasks(task_ins_ids=task_ins_ids_to_delete)
+
         # Convert to Messages
         messages_list = []
         while task_res_list:
@@ -220,13 +227,6 @@ def PullMessages(
             message = message_from_taskres(taskres=task_res)
             messages_list.append(message_to_proto(message))
 
-        # Delete the TaskIns/TaskRes pairs if TaskRes is found
-        task_ins_ids_to_delete = {
-            UUID(task_res.task.ancestry[0]) for task_res in task_res_list
-        }
-
-        state.delete_tasks(task_ins_ids=task_ins_ids_to_delete)
-
         return PullResMessagesResponse(messages_list=messages_list)
 
     def GetRun(