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

flwr/server/serverapp/app.py

@@ -23,14 +23,18 @@ from typing import Optional
 
 from flwr.cli.config_utils import get_fab_metadata
 from flwr.cli.install import install_from_fab
-from flwr.common.args import add_args_flwr_app_common, try_obtain_certificates
+from flwr.common.args import add_args_flwr_app_common, try_obtain_root_certificates
 from flwr.common.config import (
     get_flwr_dir,
     get_fused_config_from_dir,
     get_project_config,
     get_project_dir,
 )
-from flwr.common.constant import Status, SubStatus
+from flwr.common.constant import (
+    SERVERAPPIO_API_DEFAULT_CLIENT_ADDRESS,
+    Status,
+    SubStatus,
+)
 from flwr.common.logger import (
     log,
     mirror_output_to_queue,
@@ -62,33 +66,18 @@ def flwr_serverapp() -> None:
     log_queue: Queue[Optional[str]] = Queue()
     mirror_output_to_queue(log_queue)
 
-    parser = argparse.ArgumentParser(
-        description="Run a Flower ServerApp",
-    )
-    parser.add_argument(
-        "--superlink",
-        type=str,
-        help="Address of SuperLink's ServerAppIo API",
-    )
-    parser.add_argument(
-        "--run-once",
-        action="store_true",
-        help="When set, this process will start a single ServerApp for a pending Run. "
-        "If there is no pending Run, the process will exit.",
-    )
-    add_args_flwr_app_common(parser=parser)
-    args = parser.parse_args()
+    args = _parse_args_run_flwr_serverapp().parse_args()
 
     log(INFO, "Starting Flower ServerApp")
-    certificates = try_obtain_certificates(args)
+    certificates = try_obtain_root_certificates(args, args.serverappio_api_address)
 
     log(
         DEBUG,
         "Starting isolated `ServerApp` connected to SuperLink's ServerAppIo API at %s",
-        args.superlink,
+        args.serverappio_api_address,
     )
     run_serverapp(
-        superlink=args.superlink,
+        serverappio_api_address=args.serverappio_api_address,
         log_queue=log_queue,
         run_once=args.run_once,
         flwr_dir=args.flwr_dir,
@@ -100,7 +89,7 @@ def flwr_serverapp() -> None:
 
 
 def run_serverapp(  # pylint: disable=R0914, disable=W0212
-    superlink: str,
+    serverappio_api_address: str,
     log_queue: Queue[Optional[str]],
     run_once: bool,
     flwr_dir: Optional[str] = None,
@@ -108,7 +97,7 @@ def run_serverapp(  # pylint: disable=R0914, disable=W0212
 ) -> None:
     """Run Flower ServerApp process."""
     driver = GrpcDriver(
-        serverappio_service_address=superlink,
+        serverappio_service_address=serverappio_api_address,
         root_certificates=certificates,
     )
 
@@ -212,3 +201,25 @@ def run_serverapp(  # pylint: disable=R0914, disable=W0212
         # Stop the loop if `flwr-serverapp` is expected to process a single run
         if run_once:
             break
+
+
+def _parse_args_run_flwr_serverapp() -> argparse.ArgumentParser:
+    """Parse flwr-serverapp command line arguments."""
+    parser = argparse.ArgumentParser(
+        description="Run a Flower ServerApp",
+    )
+    parser.add_argument(
+        "--serverappio-api-address",
+        default=SERVERAPPIO_API_DEFAULT_CLIENT_ADDRESS,
+        type=str,
+        help="Address of SuperLink's ServerAppIo API (IPv4, IPv6, or a domain name)."
+        f"By default, it is set to {SERVERAPPIO_API_DEFAULT_CLIENT_ADDRESS}.",
+    )
+    parser.add_argument(
+        "--run-once",
+        action="store_true",
+        help="When set, this process will start a single ServerApp for a pending Run. "
+        "If there is no pending Run, the process will exit.",
+    )
+    add_args_flwr_app_common(parser=parser)
+    return parser

flwr/server/superlink/fleet/rest_rere/rest_api.py

@@ -19,7 +19,7 @@ from __future__ import annotations
 
 import sys
 from collections.abc import Awaitable
-from typing import Callable, TypeVar
+from typing import Callable, TypeVar, cast
 
 from google.protobuf.message import Message as GrpcMessage
 
@@ -39,8 +39,9 @@ from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
 )
 from flwr.proto.run_pb2 import GetRunRequest, GetRunResponse  # pylint: disable=E0611
 from flwr.server.superlink.ffs.ffs import Ffs
+from flwr.server.superlink.ffs.ffs_factory import FfsFactory
 from flwr.server.superlink.fleet.message_handler import message_handler
-from flwr.server.superlink.linkstate import LinkState
+from flwr.server.superlink.linkstate import LinkState, LinkStateFactory
 
 try:
     from starlette.applications import Starlette
@@ -90,7 +91,7 @@ def rest_request_response(
 async def create_node(request: CreateNodeRequest) -> CreateNodeResponse:
     """Create Node."""
     # Get state from app
-    state: LinkState = app.state.STATE_FACTORY.state()
+    state: LinkState = cast(LinkStateFactory, app.state.STATE_FACTORY).state()
 
     # Handle message
     return message_handler.create_node(request=request, state=state)
@@ -100,7 +101,7 @@ async def create_node(request: CreateNodeRequest) -> CreateNodeResponse:
 async def delete_node(request: DeleteNodeRequest) -> DeleteNodeResponse:
     """Delete Node Id."""
     # Get state from app
-    state: LinkState = app.state.STATE_FACTORY.state()
+    state: LinkState = cast(LinkStateFactory, app.state.STATE_FACTORY).state()
 
     # Handle message
     return message_handler.delete_node(request=request, state=state)
@@ -110,7 +111,7 @@ async def delete_node(request: DeleteNodeRequest) -> DeleteNodeResponse:
 async def pull_task_ins(request: PullTaskInsRequest) -> PullTaskInsResponse:
     """Pull TaskIns."""
     # Get state from app
-    state: LinkState = app.state.STATE_FACTORY.state()
+    state: LinkState = cast(LinkStateFactory, app.state.STATE_FACTORY).state()
 
     # Handle message
     return message_handler.pull_task_ins(request=request, state=state)
@@ -121,7 +122,7 @@ async def pull_task_ins(request: PullTaskInsRequest) -> PullTaskInsResponse:
 async def push_task_res(request: PushTaskResRequest) -> PushTaskResResponse:
     """Push TaskRes."""
     # Get state from app
-    state: LinkState = app.state.STATE_FACTORY.state()
+    state: LinkState = cast(LinkStateFactory, app.state.STATE_FACTORY).state()
 
     # Handle message
     return message_handler.push_task_res(request=request, state=state)
@@ -131,7 +132,7 @@ async def push_task_res(request: PushTaskResRequest) -> PushTaskResResponse:
 async def ping(request: PingRequest) -> PingResponse:
     """Ping."""
     # Get state from app
-    state: LinkState = app.state.STATE_FACTORY.state()
+    state: LinkState = cast(LinkStateFactory, app.state.STATE_FACTORY).state()
 
     # Handle message
     return message_handler.ping(request=request, state=state)
@@ -141,7 +142,7 @@ async def ping(request: PingRequest) -> PingResponse:
 async def get_run(request: GetRunRequest) -> GetRunResponse:
     """GetRun."""
     # Get state from app
-    state: LinkState = app.state.STATE_FACTORY.state()
+    state: LinkState = cast(LinkStateFactory, app.state.STATE_FACTORY).state()
 
     # Handle message
     return message_handler.get_run(request=request, state=state)
@@ -151,7 +152,7 @@ async def get_run(request: GetRunRequest) -> GetRunResponse:
 async def get_fab(request: GetFabRequest) -> GetFabResponse:
     """GetRun."""
     # Get ffs from app
-    ffs: Ffs = app.state.FFS_FACTORY.state()
+    ffs: Ffs = cast(FfsFactory, app.state.FFS_FACTORY).ffs()
 
     # Handle message
     return message_handler.get_fab(request=request, ffs=ffs)

flwr/server/superlink/linkstate/in_memory_linkstate.py

@@ -40,6 +40,8 @@ from .utils import (
     generate_rand_int_from_bytes,
     has_valid_sub_status,
     is_valid_transition,
+    verify_found_taskres,
+    verify_taskins_ids,
 )
 
 
@@ -48,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)
 
@@ -72,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."""
@@ -227,42 +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)
-
-            # 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."""
@@ -283,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.
 
@@ -386,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
 
@@ -452,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
             }
@@ -466,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,
@@ -489,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]:
@@ -504,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

flwr/server/superlink/linkstate/sqlite_linkstate.py

@@ -57,6 +57,8 @@ from .utils import (
     generate_rand_int_from_bytes,
     has_valid_sub_status,
     is_valid_transition,
+    verify_found_taskres,
+    verify_taskins_ids,
 )
 
 SQL_CREATE_TABLE_NODE = """
@@ -510,136 +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 = ""
-        """
-
-        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});
+            WHERE ancestry IN ({",".join(["?"] * len(task_ids))})
+            AND delivered_at = "";
         """
-        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)]
+        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)
 
-        # 2. Query: Select offline nodes
-        placeholders = ",".join([f":id_{i}" for i in range(len(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 node_id
-            FROM node
-            WHERE node_id IN ({placeholders})
-            AND online_until < :time;
+            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(node_ids)}
-        data["time"] = time.time()
-        offline_node_ids = [int(row["node_id"]) for row in self.query(query, data)]
+        data: list[Any] = [delivered_at] + task_res_ids
+        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)
+        # 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.
@@ -699,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:
@@ -922,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
@@ -1255,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