flwr 1.13.1__py3-none-any.whl → 1.14.0__py3-none-any.whl

flwr/server/serverapp/app.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 """Flower ServerApp process."""
 
+
 import argparse
 import sys
 from logging import DEBUG, ERROR, INFO
@@ -24,6 +25,7 @@ from typing import Optional
 
 from flwr.cli.config_utils import get_fab_metadata
 from flwr.cli.install import install_from_fab
+from flwr.cli.utils import get_sha256_hash
 from flwr.common.args import add_args_flwr_app_common
 from flwr.common.config import (
     get_flwr_dir,
@@ -50,7 +52,8 @@ from flwr.common.serde import (
     run_from_proto,
     run_status_to_proto,
 )
-from flwr.common.typing import RunStatus
+from flwr.common.telemetry import EventType, event
+from flwr.common.typing import RunNotRunningException, RunStatus
 from flwr.proto.run_pb2 import UpdateRunStatusRequest  # pylint: disable=E0611
 from flwr.proto.serverappio_pb2 import (  # pylint: disable=E0611
     PullServerAppInputsRequest,
@@ -96,7 +99,7 @@ def flwr_serverapp() -> None:
     restore_output()
 
 
-def run_serverapp(  # pylint: disable=R0914, disable=W0212
+def run_serverapp(  # pylint: disable=R0914, disable=W0212, disable=R0915
     serverappio_api_address: str,
     log_queue: Queue[Optional[str]],
     run_once: bool,
@@ -112,7 +115,8 @@ def run_serverapp(  # pylint: disable=R0914, disable=W0212
     # Resolve directory where FABs are installed
     flwr_dir_ = get_flwr_dir(flwr_dir)
     log_uploader = None
-
+    success = True
+    hash_run_id = None
     while True:
 
         try:
@@ -128,6 +132,8 @@ def run_serverapp(  # pylint: disable=R0914, disable=W0212
             run = run_from_proto(res.run)
             fab = fab_from_proto(res.fab)
 
+            hash_run_id = get_sha256_hash(run.run_id)
+
             driver.set_run(run.run_id)
 
             # Start log uploader for this run
@@ -170,6 +176,11 @@ def run_serverapp(  # pylint: disable=R0914, disable=W0212
                 UpdateRunStatusRequest(run_id=run.run_id, run_status=run_status_proto)
             )
 
+            event(
+                EventType.FLWR_SERVERAPP_RUN_ENTER,
+                event_details={"run-id-hash": hash_run_id},
+            )
+
             # Load and run the ServerApp with the Driver
             updated_context = run_(
                 driver=driver,
@@ -186,11 +197,18 @@ def run_serverapp(  # pylint: disable=R0914, disable=W0212
             _ = driver._stub.PushServerAppOutputs(out_req)
 
             run_status = RunStatus(Status.FINISHED, SubStatus.COMPLETED, "")
+        except RunNotRunningException:
+            log(INFO, "")
+            log(INFO, "Run ID %s stopped.", run.run_id)
+            log(INFO, "")
+            run_status = None
+            success = False
 
         except Exception as ex:  # pylint: disable=broad-exception-caught
             exc_entity = "ServerApp"
             log(ERROR, "%s raised an exception", exc_entity, exc_info=ex)
             run_status = RunStatus(Status.FINISHED, SubStatus.FAILED, str(ex))
+            success = False
 
         finally:
             # Stop log uploader for this run and upload final logs
@@ -206,6 +224,10 @@ def run_serverapp(  # pylint: disable=R0914, disable=W0212
                         run_id=run.run_id, run_status=run_status_proto
                     )
                 )
+            event(
+                EventType.FLWR_SERVERAPP_RUN_LEAVE,
+                event_details={"run-id-hash": hash_run_id, "success": success},
+            )
 
         # Stop the loop if `flwr-serverapp` is expected to process a single run
         if run_once:

flwr/server/strategy/dpfedavg_fixed.py

@@ -17,6 +17,7 @@
 Paper: arxiv.org/pdf/1710.06963.pdf
 """
 
+
 from typing import Optional, Union
 
 from flwr.common import EvaluateIns, EvaluateRes, FitIns, FitRes, Parameters, Scalar

flwr/server/superlink/driver/serverappio_grpc.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 """ServerAppIo gRPC API."""
 
+
 from logging import INFO
 from typing import Optional
 

flwr/server/superlink/driver/serverappio_servicer.py

@@ -32,6 +32,7 @@ from flwr.common.serde import (
     fab_from_proto,
     fab_to_proto,
     run_status_from_proto,
+    run_status_to_proto,
     run_to_proto,
     user_config_from_proto,
 )
@@ -48,6 +49,8 @@ from flwr.proto.run_pb2 import (  # pylint: disable=E0611
     CreateRunResponse,
     GetRunRequest,
     GetRunResponse,
+    GetRunStatusRequest,
+    GetRunStatusResponse,
     UpdateRunStatusRequest,
     UpdateRunStatusResponse,
 )
@@ -67,6 +70,7 @@ from flwr.proto.task_pb2 import TaskRes  # pylint: disable=E0611
 from flwr.server.superlink.ffs.ffs import Ffs
 from flwr.server.superlink.ffs.ffs_factory import FfsFactory
 from flwr.server.superlink.linkstate import LinkState, LinkStateFactory
+from flwr.server.superlink.utils import abort_if
 from flwr.server.utils.validator import validate_task_ins_or_res
 
 
@@ -85,7 +89,18 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
     ) -> GetNodesResponse:
         """Get available nodes."""
         log(DEBUG, "ServerAppIoServicer.GetNodes")
+
+        # Init state
         state: LinkState = self.state_factory.state()
+
+        # Abort if the run is not running
+        abort_if(
+            request.run_id,
+            [Status.PENDING, Status.STARTING, Status.FINISHED],
+            state,
+            context,
+        )
+
         all_ids: set[int] = state.get_nodes(request.run_id)
         nodes: list[Node] = [
             Node(node_id=node_id, anonymous=False) for node_id in all_ids
@@ -123,6 +138,17 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
         """Push a set of TaskIns."""
         log(DEBUG, "ServerAppIoServicer.PushTaskIns")
 
+        # Init state
+        state: LinkState = self.state_factory.state()
+
+        # Abort if the run is not running
+        abort_if(
+            request.run_id,
+            [Status.PENDING, Status.STARTING, Status.FINISHED],
+            state,
+            context,
+        )
+
         # Set pushed_at (timestamp in seconds)
         pushed_at = time.time()
         for task_ins in request.task_ins_list:
@@ -133,9 +159,9 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
         for task_ins in request.task_ins_list:
             validation_errors = validate_task_ins_or_res(task_ins)
             _raise_if(bool(validation_errors), ", ".join(validation_errors))
-
-        # Init state
-        state: LinkState = self.state_factory.state()
+            _raise_if(
+                request.run_id != task_ins.run_id, "`task_ins` has mismatched `run_id`"
+            )
 
         # Store each TaskIns
         task_ids: list[Optional[UUID]] = []
@@ -153,33 +179,35 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
         """Pull a set of TaskRes."""
         log(DEBUG, "ServerAppIoServicer.PullTaskRes")
 
-        # Convert each task_id str to UUID
-        task_ids: set[UUID] = {UUID(task_id) for task_id in request.task_ids}
-
         # Init state
         state: LinkState = self.state_factory.state()
 
-        # Register callback
-        def on_rpc_done() -> None:
-            log(
-                DEBUG,
-                "ServerAppIoServicer.PullTaskRes callback: delete TaskIns/TaskRes",
-            )
-
-            if context.is_active():
-                return
-            if context.code() != grpc.StatusCode.OK:
-                return
-
-            # Delete delivered TaskIns and TaskRes
-            state.delete_tasks(task_ids=task_ids)
+        # Abort if the run is not running
+        abort_if(
+            request.run_id,
+            [Status.PENDING, Status.STARTING, Status.FINISHED],
+            state,
+            context,
+        )
 
-        context.add_callback(on_rpc_done)
+        # Convert each task_id str to UUID
+        task_ids: set[UUID] = {UUID(task_id) for task_id in request.task_ids}
 
         # Read from state
         task_res_list: list[TaskRes] = state.get_task_res(task_ids=task_ids)
 
-        context.set_code(grpc.StatusCode.OK)
+        # Validate request
+        for task_res in task_res_list:
+            _raise_if(
+                request.run_id != task_res.run_id, "`task_res` has mismatched `run_id`"
+            )
+
+        # 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 PullTaskResResponse(task_res_list=task_res_list)
 
     def GetRun(
@@ -255,7 +283,18 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
     ) -> PushServerAppOutputsResponse:
         """Push ServerApp process outputs."""
         log(DEBUG, "ServerAppIoServicer.PushServerAppOutputs")
+
+        # Init state
         state = self.state_factory.state()
+
+        # Abort if the run is not running
+        abort_if(
+            request.run_id,
+            [Status.PENDING, Status.STARTING, Status.FINISHED],
+            state,
+            context,
+        )
+
         state.set_serverapp_context(request.run_id, context_from_proto(request.context))
         return PushServerAppOutputsResponse()
 
@@ -263,9 +302,14 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
         self, request: UpdateRunStatusRequest, context: grpc.ServicerContext
     ) -> UpdateRunStatusResponse:
         """Update the status of a run."""
-        log(DEBUG, "ControlServicer.UpdateRunStatus")
+        log(DEBUG, "ServerAppIoServicer.UpdateRunStatus")
+
+        # Init state
         state = self.state_factory.state()
 
+        # Abort if the run is finished
+        abort_if(request.run_id, [Status.FINISHED], state, context)
+
         # Update the run status
         state.update_run_status(
             run_id=request.run_id, new_status=run_status_from_proto(request.run_status)
@@ -284,6 +328,21 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
         state.add_serverapp_log(request.run_id, merged_logs)
         return PushLogsResponse()
 
+    def GetRunStatus(
+        self, request: GetRunStatusRequest, context: grpc.ServicerContext
+    ) -> GetRunStatusResponse:
+        """Get the status of a run."""
+        log(DEBUG, "ServerAppIoServicer.GetRunStatus")
+        state = self.state_factory.state()
+
+        # Get run status from LinkState
+        run_statuses = state.get_run_status(set(request.run_ids))
+        run_status_dict = {
+            run_id: run_status_to_proto(run_status)
+            for run_id, run_status in run_statuses.items()
+        }
+        return GetRunStatusResponse(run_status_dict=run_status_dict)
+
 
 def _raise_if(validation_error: bool, detail: str) -> None:
     if validation_error:

flwr/server/superlink/ffs/disk_ffs.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 """Disk based Flower File Storage."""
 
+
 import hashlib
 import json
 from pathlib import Path

flwr/server/superlink/fleet/grpc_adapter/grpc_adapter_servicer.py

@@ -158,4 +158,5 @@ class GrpcAdapterServicer(grpcadapter_pb2_grpc.GrpcAdapterServicer):
         return message_handler.get_fab(
             request=request,
             ffs=self.ffs_factory.ffs(),
+            state=self.state_factory.state(),
         )

flwr/server/superlink/fleet/grpc_bidi/flower_service_servicer.py

@@ -18,6 +18,7 @@ Relevant knowledge for reading this modules code:
 - https://github.com/grpc/grpc/blob/master/doc/statuscodes.md
 """
 
+
 import uuid
 from collections.abc import Iterator
 from typing import Callable

flwr/server/superlink/fleet/grpc_rere/fleet_servicer.py

@@ -20,6 +20,7 @@ from logging import DEBUG, INFO
 import grpc
 
 from flwr.common.logger import log
+from flwr.common.typing import InvalidRunStatusException
 from flwr.proto import fleet_pb2_grpc  # pylint: disable=E0611
 from flwr.proto.fab_pb2 import GetFabRequest, GetFabResponse  # pylint: disable=E0611
 from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
@@ -38,6 +39,7 @@ from flwr.proto.run_pb2 import GetRunRequest, GetRunResponse  # pylint: disable=
 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 LinkStateFactory
+from flwr.server.superlink.utils import abort_grpc_context
 
 
 class FleetServicer(fleet_pb2_grpc.FleetServicer):
@@ -105,27 +107,45 @@ class FleetServicer(fleet_pb2_grpc.FleetServicer):
             )
         else:
             log(INFO, "[Fleet.PushTaskRes] No task results to push")
-        return message_handler.push_task_res(
-            request=request,
-            state=self.state_factory.state(),
-        )
+
+        try:
+            res = message_handler.push_task_res(
+                request=request,
+                state=self.state_factory.state(),
+            )
+        except InvalidRunStatusException as e:
+            abort_grpc_context(e.message, context)
+
+        return res
 
     def GetRun(
         self, request: GetRunRequest, context: grpc.ServicerContext
     ) -> GetRunResponse:
         """Get run information."""
         log(INFO, "[Fleet.GetRun] Requesting `Run` for run_id=%s", request.run_id)
-        return message_handler.get_run(
-            request=request,
-            state=self.state_factory.state(),
-        )
+
+        try:
+            res = message_handler.get_run(
+                request=request,
+                state=self.state_factory.state(),
+            )
+        except InvalidRunStatusException as e:
+            abort_grpc_context(e.message, context)
+
+        return res
 
     def GetFab(
         self, request: GetFabRequest, context: grpc.ServicerContext
     ) -> GetFabResponse:
         """Get FAB."""
         log(INFO, "[Fleet.GetFab] Requesting FAB for fab_hash=%s", request.hash_str)
-        return message_handler.get_fab(
-            request=request,
-            ffs=self.ffs_factory.ffs(),
-        )
+        try:
+            res = message_handler.get_fab(
+                request=request,
+                ffs=self.ffs_factory.ffs(),
+                state=self.state_factory.state(),
+            )
+        except InvalidRunStatusException as e:
+            abort_grpc_context(e.message, context)
+
+        return res

flwr/server/superlink/fleet/grpc_rere/server_interceptor.py

@@ -45,7 +45,7 @@ from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
 )
 from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
 from flwr.proto.run_pb2 import GetRunRequest, GetRunResponse  # pylint: disable=E0611
-from flwr.server.superlink.linkstate import LinkState
+from flwr.server.superlink.linkstate import LinkStateFactory
 
 _PUBLIC_KEY_HEADER = "public-key"
 _AUTH_TOKEN_HEADER = "auth-token"
@@ -84,15 +84,16 @@ def _get_value_from_tuples(
 class AuthenticateServerInterceptor(grpc.ServerInterceptor):  # type: ignore
     """Server interceptor for node authentication."""
 
-    def __init__(self, state: LinkState):
-        self.state = state
+    def __init__(self, state_factory: LinkStateFactory):
+        self.state_factory = state_factory
+        state = self.state_factory.state()
 
         self.node_public_keys = state.get_node_public_keys()
         if len(self.node_public_keys) == 0:
             log(WARNING, "Authentication enabled, but no known public keys configured")
 
-        private_key = self.state.get_server_private_key()
-        public_key = self.state.get_server_public_key()
+        private_key = state.get_server_private_key()
+        public_key = state.get_server_public_key()
 
         if private_key is None or public_key is None:
             raise ValueError("Error loading authentication keys")
@@ -154,7 +155,7 @@ class AuthenticateServerInterceptor(grpc.ServerInterceptor):  # type: ignore
                 context.abort(grpc.StatusCode.UNAUTHENTICATED, "Access denied")
 
             # Verify node_id
-            node_id = self.state.get_node_id(node_public_key_bytes)
+            node_id = self.state_factory.state().get_node_id(node_public_key_bytes)
 
             if not self._verify_node_id(node_id, request):
                 context.abort(grpc.StatusCode.UNAUTHENTICATED, "Access denied")
@@ -186,7 +187,7 @@ class AuthenticateServerInterceptor(grpc.ServerInterceptor):  # type: ignore
                 return False
             return request.task_res_list[0].task.producer.node_id == node_id
         if isinstance(request, GetRunRequest):
-            return node_id in self.state.get_nodes(request.run_id)
+            return node_id in self.state_factory.state().get_nodes(request.run_id)
         return request.node.node_id == node_id
 
     def _verify_hmac(
@@ -210,17 +211,17 @@ class AuthenticateServerInterceptor(grpc.ServerInterceptor):  # type: ignore
                 ),
             )
         )
-
-        node_id = self.state.get_node_id(public_key_bytes)
+        state = self.state_factory.state()
+        node_id = state.get_node_id(public_key_bytes)
 
         # Handle `CreateNode` here instead of calling the default method handler
         # Return previously assigned `node_id` for the provided `public_key`
         if node_id is not None:
-            self.state.acknowledge_ping(node_id, request.ping_interval)
+            state.acknowledge_ping(node_id, request.ping_interval)
             return CreateNodeResponse(node=Node(node_id=node_id, anonymous=False))
 
         # No `node_id` exists for the provided `public_key`
         # Handle `CreateNode` here instead of calling the default method handler
         # Note: the innermost `CreateNode` method will never be called
-        node_id = self.state.create_node(request.ping_interval, public_key_bytes)
+        node_id = state.create_node(request.ping_interval, public_key_bytes)
         return CreateNodeResponse(node=Node(node_id=node_id, anonymous=False))

flwr/server/superlink/fleet/message_handler/message_handler.py

@@ -19,8 +19,9 @@ import time
 from typing import Optional
 from uuid import UUID
 
+from flwr.common.constant import Status
 from flwr.common.serde import fab_to_proto, user_config_to_proto
-from flwr.common.typing import Fab
+from flwr.common.typing import Fab, InvalidRunStatusException
 from flwr.proto.fab_pb2 import GetFabRequest, GetFabResponse  # pylint: disable=E0611
 from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
     CreateNodeRequest,
@@ -44,6 +45,7 @@ from flwr.proto.run_pb2 import (  # pylint: disable=E0611
 from flwr.proto.task_pb2 import TaskIns, TaskRes  # pylint: disable=E0611
 from flwr.server.superlink.ffs.ffs import Ffs
 from flwr.server.superlink.linkstate import LinkState
+from flwr.server.superlink.utils import check_abort
 
 
 def create_node(
@@ -98,6 +100,15 @@ def push_task_res(request: PushTaskResRequest, state: LinkState) -> PushTaskResR
     task_res: TaskRes = request.task_res_list[0]
     # pylint: enable=no-member
 
+    # Abort if the run is not running
+    abort_msg = check_abort(
+        task_res.run_id,
+        [Status.PENDING, Status.STARTING, Status.FINISHED],
+        state,
+    )
+    if abort_msg:
+        raise InvalidRunStatusException(abort_msg)
+
     # Set pushed_at (timestamp in seconds)
     task_res.task.pushed_at = time.time()
 
@@ -112,15 +123,22 @@ def push_task_res(request: PushTaskResRequest, state: LinkState) -> PushTaskResR
     return response
 
 
-def get_run(
-    request: GetRunRequest, state: LinkState  # pylint: disable=W0613
-) -> GetRunResponse:
+def get_run(request: GetRunRequest, state: LinkState) -> GetRunResponse:
     """Get run information."""
     run = state.get_run(request.run_id)
 
     if run is None:
         return GetRunResponse()
 
+    # Abort if the run is not running
+    abort_msg = check_abort(
+        request.run_id,
+        [Status.PENDING, Status.STARTING, Status.FINISHED],
+        state,
+    )
+    if abort_msg:
+        raise InvalidRunStatusException(abort_msg)
+
     return GetRunResponse(
         run=Run(
             run_id=run.run_id,
@@ -133,9 +151,18 @@ def get_run(
 
 
 def get_fab(
-    request: GetFabRequest, ffs: Ffs  # pylint: disable=W0613
+    request: GetFabRequest, ffs: Ffs, state: LinkState  # pylint: disable=W0613
 ) -> GetFabResponse:
     """Get FAB."""
+    # Abort if the run is not running
+    abort_msg = check_abort(
+        request.run_id,
+        [Status.PENDING, Status.STARTING, Status.FINISHED],
+        state,
+    )
+    if abort_msg:
+        raise InvalidRunStatusException(abort_msg)
+
     if result := ffs.get(request.hash_str):
         fab = Fab(request.hash_str, result[0])
         return GetFabResponse(fab=fab_to_proto(fab))

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

@@ -154,8 +154,11 @@ async def get_fab(request: GetFabRequest) -> GetFabResponse:
     # Get ffs from app
     ffs: Ffs = cast(FfsFactory, app.state.FFS_FACTORY).ffs()
 
+    # Get state from app
+    state: LinkState = cast(LinkStateFactory, app.state.STATE_FACTORY).state()
+
     # Handle message
-    return message_handler.get_fab(request=request, ffs=ffs)
+    return message_handler.get_fab(request=request, ffs=ffs, state=state)
 
 
 routes = [

flwr/server/superlink/fleet/vce/__init__.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 """Fleet Simulation Engine side."""
 
+
 from .vce_api import start_vce
 
 __all__ = [

flwr/server/superlink/fleet/vce/backend/__init__.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 """Simulation Engine Backends."""
 
+
 import importlib
 
 from .backend import Backend, BackendConfig

flwr/server/superlink/fleet/vce/backend/raybackend.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 """Ray backend for the Fleet API using the Simulation Engine."""
 
+
 import sys
 from logging import DEBUG, ERROR
 from typing import Callable, Optional, Union

flwr/server/superlink/linkstate/in_memory_linkstate.py

@@ -265,41 +265,15 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
             for task_res in task_res_found:
                 task_res.task.delivered_at = delivered_at
 
-            # 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."""
-        task_ins_to_be_deleted: set[UUID] = set()
-        task_res_to_be_deleted: set[UUID] = set()
-
-        with self.lock:
-            for task_ins_id in task_ids:
-                # Find the task_id of the matching task_res
-                for task_res_id, task_res in self.task_res_store.items():
-                    if UUID(task_res.task.ancestry[0]) != task_ins_id:
-                        continue
-                    if task_res.task.delivered_at == "":
-                        continue
-
-                    task_ins_to_be_deleted.add(task_ins_id)
-                    task_res_to_be_deleted.add(task_res_id)
-
-            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:
+    def delete_tasks(self, task_ins_ids: set[UUID]) -> None:
+        """Delete TaskIns/TaskRes pairs based on provided TaskIns IDs."""
+        if not task_ins_ids:
             return
 
         with self.lock:
-            for task_id in task_ids:
+            for task_id in task_ins_ids:
                 # Delete TaskIns
                 if task_id in self.task_ins_store:
                     del self.task_ins_store[task_id]
@@ -308,6 +282,16 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
                     task_res_id = self.task_ins_id_to_task_res_id.pop(task_id)
                     del self.task_res_store[task_res_id]
 
+    def get_task_ids_from_run_id(self, run_id: int) -> set[UUID]:
+        """Get all TaskIns IDs for the given run_id."""
+        task_id_list: set[UUID] = set()
+        with self.lock:
+            for task_id, task_ins in self.task_ins_store.items():
+                if task_ins.run_id == run_id:
+                    task_id_list.add(task_id)
+
+        return task_id_list
+
     def num_task_ins(self) -> int:
         """Calculate the number of task_ins in store.
 
@@ -446,6 +430,13 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         """Retrieve `server_public_key` in urlsafe bytes."""
         return self.server_public_key
 
+    def clear_supernode_auth_keys_and_credentials(self) -> None:
+        """Clear stored `node_public_keys` and credentials in the link state if any."""
+        with self.lock:
+            self.server_private_key = None
+            self.server_public_key = None
+            self.node_public_keys.clear()
+
     def store_node_public_keys(self, public_keys: set[bytes]) -> None:
         """Store a set of `node_public_keys` in the link state."""
         with self.lock: