flwr-nightly 1.13.0.dev20241106__py3-none-any.whl → 1.13.0.dev20241117__py3-none-any.whl

flwr/server/superlink/linkstate/in_memory_linkstate.py

@@ -40,7 +40,8 @@ from .utils import (
     generate_rand_int_from_bytes,
     has_valid_sub_status,
     is_valid_transition,
-    make_node_unavailable_taskres,
+    verify_found_taskres,
+    verify_taskins_ids,
 )
 
 
@@ -49,11 +50,6 @@ class RunRecord:  # pylint: disable=R0902
     """The record of a specific run, including its status and timestamps."""
 
     run: Run
-    status: RunStatus
-    pending_at: str = ""
-    starting_at: str = ""
-    running_at: str = ""
-    finished_at: str = ""
     logs: list[tuple[float, str]] = field(default_factory=list)
     log_lock: threading.Lock = field(default_factory=threading.Lock)
 
@@ -73,12 +69,13 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         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."""
@@ -228,57 +225,50 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         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)
-
-            # Check if the node is offline
-            for task_id in task_ids - replied_task_ids:
-                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
+            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."""
@@ -299,9 +289,25 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
 
             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.
 
@@ -402,13 +408,16 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
                         fab_version=fab_version if fab_version else "",
                         fab_hash=fab_hash if fab_hash else "",
                         override_config=override_config,
+                        pending_at=now().isoformat(),
+                        starting_at="",
+                        running_at="",
+                        finished_at="",
+                        status=RunStatus(
+                            status=Status.PENDING,
+                            sub_status="",
+                            details="",
+                        ),
                     ),
-                    status=RunStatus(
-                        status=Status.PENDING,
-                        sub_status="",
-                        details="",
-                    ),
-                    pending_at=now().isoformat(),
                 )
                 self.run_ids[run_id] = run_record
 
@@ -451,6 +460,11 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         """Retrieve all currently stored `node_public_keys` as a set."""
         return self.node_public_keys
 
+    def get_run_ids(self) -> set[int]:
+        """Retrieve all run IDs."""
+        with self.lock:
+            return set(self.run_ids.keys())
+
     def get_run(self, run_id: int) -> Optional[Run]:
         """Retrieve information about the run with the specified `run_id`."""
         with self.lock:
@@ -463,7 +477,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         """Retrieve the statuses for the specified runs."""
         with self.lock:
             return {
-                run_id: self.run_ids[run_id].status
+                run_id: self.run_ids[run_id].run.status
                 for run_id in set(run_ids)
                 if run_id in self.run_ids
             }
@@ -477,7 +491,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
                 return False
 
             # Check if the status transition is valid
-            current_status = self.run_ids[run_id].status
+            current_status = self.run_ids[run_id].run.status
             if not is_valid_transition(current_status, new_status):
                 log(
                     ERROR,
@@ -500,12 +514,12 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
             # Update the status
             run_record = self.run_ids[run_id]
             if new_status.status == Status.STARTING:
-                run_record.starting_at = now().isoformat()
+                run_record.run.starting_at = now().isoformat()
             elif new_status.status == Status.RUNNING:
-                run_record.running_at = now().isoformat()
+                run_record.run.running_at = now().isoformat()
             elif new_status.status == Status.FINISHED:
-                run_record.finished_at = now().isoformat()
-            run_record.status = new_status
+                run_record.run.finished_at = now().isoformat()
+            run_record.run.status = new_status
             return True
 
     def get_pending_run_id(self) -> Optional[int]:
@@ -515,7 +529,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         # Loop through all registered runs
         for run_id, run_rec in self.run_ids.items():
             # Break once a pending run is found
-            if run_rec.status.status == Status.PENDING:
+            if run_rec.run.status.status == Status.PENDING:
                 pending_run_id = run_id
                 break
 

flwr/server/superlink/linkstate/linkstate.py

@@ -101,13 +101,27 @@ class LinkState(abc.ABC):  # pylint: disable=R0904
 
     @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
@@ -163,6 +177,10 @@ class LinkState(abc.ABC):  # pylint: disable=R0904
     ) -> int:
         """Create a new run for the specified `fab_hash`."""
 
+    @abc.abstractmethod
+    def get_run_ids(self) -> set[int]:
+        """Retrieve all run IDs."""
+
     @abc.abstractmethod
     def get_run(self, run_id: int) -> Optional[Run]:
         """Retrieve information about the run with the specified `run_id`.
@@ -175,10 +193,7 @@ class LinkState(abc.ABC):  # pylint: disable=R0904
         Returns
         -------
         Optional[Run]
-            A dataclass instance containing three elements if `run_id` is valid:
-            - `run_id`: The identifier of the run, same as the specified `run_id`.
-            - `fab_id`: The identifier of the FAB used in the specified run.
-            - `fab_version`: The version of the FAB used in the specified run.
+            The `Run` instance if found; otherwise, `None`.
         """
 
     @abc.abstractmethod

flwr/server/superlink/linkstate/sqlite_linkstate.py

@@ -57,7 +57,8 @@ from .utils import (
     generate_rand_int_from_bytes,
     has_valid_sub_status,
     is_valid_transition,
-    make_node_unavailable_taskres,
+    verify_found_taskres,
+    verify_taskins_ids,
 )
 
 SQL_CREATE_TABLE_NODE = """
@@ -511,150 +512,67 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
 
     # pylint: disable-next=R0912,R0915,R0914
     def get_task_res(self, task_ids: set[UUID]) -> list[TaskRes]:
-        """Get TaskRes for task_ids.
+        """Get TaskRes for the given TaskIns IDs."""
+        ret: dict[UUID, TaskRes] = {}
 
-        Usually, the ServerAppIo API calls this method to get results for instructions
-        it has previously scheduled.
-
-        Retrieves all TaskRes for the given `task_ids` and returns and empty list if
-        none could be found.
-
-        Constraints
-        -----------
-        If `limit` is not `None`, return, at most, `limit` number of TaskRes. The limit
-        will only take effect if enough task_ids are in the set AND are currently
-        available. If `limit` is set, it has to be greater than zero.
-        """
-        # Check if corresponding TaskIns exists and is not expired
-        task_ids_placeholders = ",".join([f":id_{i}" for i in range(len(task_ids))])
+        # Verify TaskIns IDs
+        current = time.time()
         query = f"""
             SELECT *
             FROM task_ins
-            WHERE task_id IN ({task_ids_placeholders})
-            AND (created_at + ttl) > CAST(strftime('%s', 'now') AS REAL)
+            WHERE task_id IN ({",".join(["?"] * len(task_ids))});
         """
-        query += ";"
-
-        task_ins_data = {}
-        for index, task_id in enumerate(task_ids):
-            task_ins_data[f"id_{index}"] = str(task_id)
-
-        task_ins_rows = self.query(query, task_ins_data)
-
-        if not task_ins_rows:
-            return []
-
-        for row in task_ins_rows:
-            # Convert values from sint64 to uint64
+        rows = self.query(query, tuple(str(task_id) for task_id in task_ids))
+        found_task_ins_dict: dict[UUID, TaskIns] = {}
+        for row in rows:
             convert_sint64_values_in_dict_to_uint64(
                 row, ["run_id", "producer_node_id", "consumer_node_id"]
             )
-            task_ins = dict_to_task_ins(row)
-            if task_ins.task.created_at + task_ins.task.ttl <= time.time():
-                log(WARNING, "TaskIns with task_id %s is expired.", task_ins.task_id)
-                task_ids.remove(UUID(task_ins.task_id))
+            found_task_ins_dict[UUID(row["task_id"])] = dict_to_task_ins(row)
 
-        # Retrieve all anonymous Tasks
-        if len(task_ids) == 0:
-            return []
+        ret = verify_taskins_ids(
+            inquired_taskins_ids=task_ids,
+            found_taskins_dict=found_task_ins_dict,
+            current_time=current,
+        )
 
-        placeholders = ",".join([f":id_{i}" for i in range(len(task_ids))])
+        # Find all TaskRes
         query = f"""
             SELECT *
             FROM task_res
-            WHERE ancestry IN ({placeholders})
-            AND delivered_at = ""
+            WHERE ancestry IN ({",".join(["?"] * len(task_ids))})
+            AND delivered_at = "";
         """
+        rows = self.query(query, tuple(str(task_id) for task_id in task_ids))
+        for row in rows:
+            convert_sint64_values_in_dict_to_uint64(
+                row, ["run_id", "producer_node_id", "consumer_node_id"]
+            )
+        tmp_ret_dict = verify_found_taskres(
+            inquired_taskins_ids=task_ids,
+            found_taskins_dict=found_task_ins_dict,
+            found_taskres_list=[dict_to_task_res(row) for row in rows],
+            current_time=current,
+        )
+        ret.update(tmp_ret_dict)
 
-        data: dict[str, Union[str, float, int]] = {}
-
-        query += ";"
-
-        for index, task_id in enumerate(task_ids):
-            data[f"id_{index}"] = str(task_id)
-
-        rows = self.query(query, data)
-
-        if rows:
-            # Prepare query
-            found_task_ids = [row["task_id"] for row in rows]
-            placeholders = ",".join([f":id_{i}" for i in range(len(found_task_ids))])
-            query = f"""
-                UPDATE task_res
-                SET delivered_at = :delivered_at
-                WHERE task_id IN ({placeholders})
-                RETURNING *;
-            """
-
-            # Prepare data for query
-            delivered_at = now().isoformat()
-            data = {"delivered_at": delivered_at}
-            for index, task_id in enumerate(found_task_ids):
-                data[f"id_{index}"] = str(task_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", "producer_node_id", "consumer_node_id"]
-                )
-
-        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))])
+        # Mark existing TaskRes to be returned as delivered
+        delivered_at = now().isoformat()
+        for task_res in ret.values():
+            task_res.task.delivered_at = delivered_at
+        task_res_ids = [task_res.task_id for task_res in ret.values()]
         query = f"""
-            SELECT *
-            FROM task_ins
-            WHERE consumer_node_id IN ({placeholders});
+            UPDATE task_res
+            SET delivered_at = ?
+            WHERE task_id IN ({",".join(["?"] * len(task_res_ids))});
         """
-        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:
-            for row in rows:
-                # Convert values from sint64 to uint64
-                convert_sint64_values_in_dict_to_uint64(
-                    row, ["run_id", "producer_node_id", "consumer_node_id"]
-                )
+        data: list[Any] = [delivered_at] + task_res_ids
+        self.query(query, data)
 
-            task_ins = dict_to_task_ins(row)
-            err_taskres = make_node_unavailable_taskres(
-                ref_taskins=task_ins,
-            )
-            result.append(err_taskres)
+        # Cleanup
+        self._force_delete_tasks_by_ids(set(ret.keys()))
 
-        return result
+        return list(ret.values())
 
     def num_task_ins(self) -> int:
         """Calculate the number of task_ins in store.
@@ -714,6 +632,32 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
 
         return None
 
+    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
+        if self.conn is None:
+            raise AttributeError("LinkState not initialized")
+
+        placeholders = ",".join([f":id_{index}" for index in range(len(task_ids))])
+        data = {f"id_{index}": str(task_id) for index, task_id in enumerate(task_ids)}
+
+        # Delete task_ins
+        query_1 = f"""
+            DELETE FROM task_ins
+            WHERE task_id IN ({placeholders});
+        """
+
+        # Delete task_res
+        query_2 = f"""
+            DELETE FROM task_res
+            WHERE ancestry IN ({placeholders});
+        """
+
+        with self.conn:
+            self.conn.execute(query_1, data)
+            self.conn.execute(query_2, data)
+
     def create_node(
         self, ping_interval: float, public_key: Optional[bytes] = None
     ) -> int:
@@ -917,6 +861,12 @@ 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)
+        return {convert_sint64_to_uint64(row["run_id"]) for row in rows}
+
     def get_run(self, run_id: int) -> Optional[Run]:
         """Retrieve information about the run with the specified `run_id`."""
         # Convert the uint64 value to sint64 for SQLite
@@ -931,6 +881,15 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
                 fab_version=row["fab_version"],
                 fab_hash=row["fab_hash"],
                 override_config=json.loads(row["override_config"]),
+                pending_at=row["pending_at"],
+                starting_at=row["starting_at"],
+                running_at=row["running_at"],
+                finished_at=row["finished_at"],
+                status=RunStatus(
+                    status=determine_run_status(row),
+                    sub_status=row["sub_status"],
+                    details=row["details"],
+                ),
             )
         log(ERROR, "`run_id` does not exist.")
         return None
@@ -1264,10 +1223,10 @@ def dict_to_task_res(task_dict: dict[str, Any]) -> TaskRes:
 def determine_run_status(row: dict[str, Any]) -> str:
     """Determine the status of the run based on timestamp fields."""
     if row["pending_at"]:
+        if row["finished_at"]:
+            return Status.FINISHED
         if row["starting_at"]:
             if row["running_at"]:
-                if row["finished_at"]:
-                    return Status.FINISHED
                 return Status.RUNNING
             return Status.STARTING
         return Status.PENDING