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

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

@@ -37,13 +37,15 @@ 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_factory import FfsFactory
 from flwr.server.superlink.fleet.message_handler import message_handler
-from flwr.server.superlink.state import StateFactory
+from flwr.server.superlink.linkstate import LinkStateFactory
 
 
 class FleetServicer(fleet_pb2_grpc.FleetServicer):
     """Fleet API servicer."""
 
-    def __init__(self, state_factory: StateFactory, ffs_factory: FfsFactory) -> None:
+    def __init__(
+        self, state_factory: LinkStateFactory, ffs_factory: FfsFactory
+    ) -> None:
         self.state_factory = state_factory
         self.ffs_factory = ffs_factory
 

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.state import State
+from flwr.server.superlink.linkstate import LinkState
 
 _PUBLIC_KEY_HEADER = "public-key"
 _AUTH_TOKEN_HEADER = "auth-token"
@@ -84,7 +84,7 @@ def _get_value_from_tuples(
 class AuthenticateServerInterceptor(grpc.ServerInterceptor):  # type: ignore
     """Server interceptor for node authentication."""
 
-    def __init__(self, state: State):
+    def __init__(self, state: LinkState):
         self.state = state
 
         self.node_public_keys = state.get_node_public_keys()

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

@@ -43,12 +43,12 @@ 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.state import State
+from flwr.server.superlink.linkstate import LinkState
 
 
 def create_node(
     request: CreateNodeRequest,  # pylint: disable=unused-argument
-    state: State,
+    state: LinkState,
 ) -> CreateNodeResponse:
     """."""
     # Create node
@@ -56,7 +56,7 @@ def create_node(
     return CreateNodeResponse(node=Node(node_id=node_id, anonymous=False))
 
 
-def delete_node(request: DeleteNodeRequest, state: State) -> DeleteNodeResponse:
+def delete_node(request: DeleteNodeRequest, state: LinkState) -> DeleteNodeResponse:
     """."""
     # Validate node_id
     if request.node.anonymous or request.node.node_id == 0:
@@ -69,14 +69,14 @@ def delete_node(request: DeleteNodeRequest, state: State) -> DeleteNodeResponse:
 
 def ping(
     request: PingRequest,  # pylint: disable=unused-argument
-    state: State,  # pylint: disable=unused-argument
+    state: LinkState,  # pylint: disable=unused-argument
 ) -> PingResponse:
     """."""
     res = state.acknowledge_ping(request.node.node_id, request.ping_interval)
     return PingResponse(success=res)
 
 
-def pull_task_ins(request: PullTaskInsRequest, state: State) -> PullTaskInsResponse:
+def pull_task_ins(request: PullTaskInsRequest, state: LinkState) -> PullTaskInsResponse:
     """Pull TaskIns handler."""
     # Get node_id if client node is not anonymous
     node = request.node  # pylint: disable=no-member
@@ -92,7 +92,7 @@ def pull_task_ins(request: PullTaskInsRequest, state: State) -> PullTaskInsRespo
     return response
 
 
-def push_task_res(request: PushTaskResRequest, state: State) -> PushTaskResResponse:
+def push_task_res(request: PushTaskResRequest, state: LinkState) -> PushTaskResResponse:
     """Push TaskRes handler."""
     # pylint: disable=no-member
     task_res: TaskRes = request.task_res_list[0]
@@ -113,7 +113,7 @@ def push_task_res(request: PushTaskResRequest, state: State) -> PushTaskResRespo
 
 
 def get_run(
-    request: GetRunRequest, state: State  # pylint: disable=W0613
+    request: GetRunRequest, state: LinkState  # pylint: disable=W0613
 ) -> GetRunResponse:
     """Get run information."""
     run = state.get_run(request.run_id)

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

@@ -40,7 +40,7 @@ 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.fleet.message_handler import message_handler
-from flwr.server.superlink.state import State
+from flwr.server.superlink.linkstate import LinkState
 
 try:
     from starlette.applications import Starlette
@@ -90,7 +90,7 @@ def rest_request_response(
 async def create_node(request: CreateNodeRequest) -> CreateNodeResponse:
     """Create Node."""
     # Get state from app
-    state: State = app.state.STATE_FACTORY.state()
+    state: LinkState = app.state.STATE_FACTORY.state()
 
     # Handle message
     return message_handler.create_node(request=request, state=state)
@@ -100,7 +100,7 @@ async def create_node(request: CreateNodeRequest) -> CreateNodeResponse:
 async def delete_node(request: DeleteNodeRequest) -> DeleteNodeResponse:
     """Delete Node Id."""
     # Get state from app
-    state: State = app.state.STATE_FACTORY.state()
+    state: LinkState = app.state.STATE_FACTORY.state()
 
     # Handle message
     return message_handler.delete_node(request=request, state=state)
@@ -110,7 +110,7 @@ async def delete_node(request: DeleteNodeRequest) -> DeleteNodeResponse:
 async def pull_task_ins(request: PullTaskInsRequest) -> PullTaskInsResponse:
     """Pull TaskIns."""
     # Get state from app
-    state: State = app.state.STATE_FACTORY.state()
+    state: LinkState = app.state.STATE_FACTORY.state()
 
     # Handle message
     return message_handler.pull_task_ins(request=request, state=state)
@@ -121,7 +121,7 @@ async def pull_task_ins(request: PullTaskInsRequest) -> PullTaskInsResponse:
 async def push_task_res(request: PushTaskResRequest) -> PushTaskResResponse:
     """Push TaskRes."""
     # Get state from app
-    state: State = app.state.STATE_FACTORY.state()
+    state: LinkState = app.state.STATE_FACTORY.state()
 
     # Handle message
     return message_handler.push_task_res(request=request, state=state)
@@ -131,7 +131,7 @@ async def push_task_res(request: PushTaskResRequest) -> PushTaskResResponse:
 async def ping(request: PingRequest) -> PingResponse:
     """Ping."""
     # Get state from app
-    state: State = app.state.STATE_FACTORY.state()
+    state: LinkState = app.state.STATE_FACTORY.state()
 
     # Handle message
     return message_handler.ping(request=request, state=state)
@@ -141,7 +141,7 @@ async def ping(request: PingRequest) -> PingResponse:
 async def get_run(request: GetRunRequest) -> GetRunResponse:
     """GetRun."""
     # Get state from app
-    state: State = app.state.STATE_FACTORY.state()
+    state: LinkState = app.state.STATE_FACTORY.state()
 
     # Handle message
     return message_handler.get_run(request=request, state=state)

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

@@ -28,7 +28,7 @@ from typing import Callable, Optional
 
 from flwr.client.client_app import ClientApp, ClientAppException, LoadClientAppError
 from flwr.client.clientapp.utils import get_load_client_app_fn
-from flwr.client.node_state import NodeState
+from flwr.client.run_info_store import DeprecatedRunInfoStore
 from flwr.common.constant import (
     NUM_PARTITIONS_KEY,
     PARTITION_ID_KEY,
@@ -40,7 +40,7 @@ from flwr.common.message import Error
 from flwr.common.serde import message_from_taskins, message_to_taskres
 from flwr.common.typing import Run
 from flwr.proto.task_pb2 import TaskIns, TaskRes  # pylint: disable=E0611
-from flwr.server.superlink.state import State, StateFactory
+from flwr.server.superlink.linkstate import LinkState, LinkStateFactory
 
 from .backend import Backend, error_messages_backends, supported_backends
 
@@ -48,7 +48,7 @@ NodeToPartitionMapping = dict[int, int]
 
 
 def _register_nodes(
-    num_nodes: int, state_factory: StateFactory
+    num_nodes: int, state_factory: LinkStateFactory
 ) -> NodeToPartitionMapping:
     """Register nodes with the StateFactory and create node-id:partition-id mapping."""
     nodes_mapping: NodeToPartitionMapping = {}
@@ -60,16 +60,16 @@ def _register_nodes(
     return nodes_mapping
 
 
-def _register_node_states(
+def _register_node_info_stores(
     nodes_mapping: NodeToPartitionMapping,
     run: Run,
     app_dir: Optional[str] = None,
-) -> dict[int, NodeState]:
-    """Create NodeState objects and pre-register the context for the run."""
-    node_states: dict[int, NodeState] = {}
+) -> dict[int, DeprecatedRunInfoStore]:
+    """Create DeprecatedRunInfoStore objects and register the context for the run."""
+    node_info_store: dict[int, DeprecatedRunInfoStore] = {}
     num_partitions = len(set(nodes_mapping.values()))
     for node_id, partition_id in nodes_mapping.items():
-        node_states[node_id] = NodeState(
+        node_info_store[node_id] = DeprecatedRunInfoStore(
             node_id=node_id,
             node_config={
                 PARTITION_ID_KEY: partition_id,
@@ -78,18 +78,18 @@ def _register_node_states(
         )
 
         # Pre-register Context objects
-        node_states[node_id].register_context(
+        node_info_store[node_id].register_context(
             run_id=run.run_id, run=run, app_dir=app_dir
         )
 
-    return node_states
+    return node_info_store
 
 
 # pylint: disable=too-many-arguments,too-many-locals
 def worker(
     taskins_queue: "Queue[TaskIns]",
     taskres_queue: "Queue[TaskRes]",
-    node_states: dict[int, NodeState],
+    node_info_store: dict[int, DeprecatedRunInfoStore],
     backend: Backend,
     f_stop: threading.Event,
 ) -> None:
@@ -103,7 +103,7 @@ def worker(
             node_id = task_ins.task.consumer.node_id
 
             # Retrieve context
-            context = node_states[node_id].retrieve_context(run_id=task_ins.run_id)
+            context = node_info_store[node_id].retrieve_context(run_id=task_ins.run_id)
 
             # Convert TaskIns to Message
             message = message_from_taskins(task_ins)
@@ -112,7 +112,7 @@ def worker(
             out_mssg, updated_context = backend.process_message(message, context)
 
             # Update Context
-            node_states[node_id].update_context(
+            node_info_store[node_id].update_context(
                 task_ins.run_id, context=updated_context
             )
         except Empty:
@@ -145,7 +145,7 @@ def worker(
 
 
 def add_taskins_to_queue(
-    state: State,
+    state: LinkState,
     queue: "Queue[TaskIns]",
     nodes_mapping: NodeToPartitionMapping,
     f_stop: threading.Event,
@@ -160,7 +160,7 @@ def add_taskins_to_queue(
 
 
 def put_taskres_into_state(
-    state: State, queue: "Queue[TaskRes]", f_stop: threading.Event
+    state: LinkState, queue: "Queue[TaskRes]", f_stop: threading.Event
 ) -> None:
     """Put TaskRes into State from a queue."""
     while not f_stop.is_set():
@@ -177,8 +177,8 @@ def run_api(
     app_fn: Callable[[], ClientApp],
     backend_fn: Callable[[], Backend],
     nodes_mapping: NodeToPartitionMapping,
-    state_factory: StateFactory,
-    node_states: dict[int, NodeState],
+    state_factory: LinkStateFactory,
+    node_info_stores: dict[int, DeprecatedRunInfoStore],
     f_stop: threading.Event,
 ) -> None:
     """Run the VCE."""
@@ -223,7 +223,7 @@ def run_api(
                     worker,
                     taskins_queue,
                     taskres_queue,
-                    node_states,
+                    node_info_stores,
                     backend,
                     f_stop,
                 )
@@ -264,7 +264,7 @@ def start_vce(
     client_app: Optional[ClientApp] = None,
     client_app_attr: Optional[str] = None,
     num_supernodes: Optional[int] = None,
-    state_factory: Optional[StateFactory] = None,
+    state_factory: Optional[LinkStateFactory] = None,
     existing_nodes_mapping: Optional[NodeToPartitionMapping] = None,
 ) -> None:
     """Start Fleet API with the Simulation Engine."""
@@ -303,7 +303,7 @@ def start_vce(
     if not state_factory:
         log(INFO, "A StateFactory was not supplied to the SimulationEngine.")
         # Create an empty in-memory state factory
-        state_factory = StateFactory(":flwr-in-memory-state:")
+        state_factory = LinkStateFactory(":flwr-in-memory-state:")
         log(INFO, "Created new %s.", state_factory.__class__.__name__)
 
     if num_supernodes:
@@ -312,8 +312,8 @@ def start_vce(
             num_nodes=num_supernodes, state_factory=state_factory
         )
 
-    # Construct mapping of NodeStates
-    node_states = _register_node_states(
+    # Construct mapping of DeprecatedRunInfoStore
+    node_info_stores = _register_node_info_stores(
         nodes_mapping=nodes_mapping, run=run, app_dir=app_dir if is_app else None
     )
 
@@ -376,7 +376,7 @@ def start_vce(
             backend_fn,
             nodes_mapping,
             state_factory,
-            node_states,
+            node_info_stores,
             f_stop,
         )
     except LoadClientAppError as loadapp_ex:

flwr/server/superlink/linkstate/__init__.py

@@ -0,0 +1,28 @@
+# Copyright 2024 Flower Labs GmbH. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Flower LinkState."""
+
+
+from .in_memory_linkstate import InMemoryLinkState as InMemoryLinkState
+from .linkstate import LinkState as LinkState
+from .linkstate_factory import LinkStateFactory as LinkStateFactory
+from .sqlite_linkstate import SqliteLinkState as SqliteLinkState
+
+__all__ = [
+    "InMemoryLinkState",
+    "LinkState",
+    "LinkStateFactory",
+    "SqliteLinkState",
+]

flwr/server/superlink/{state/in_memory_state.py → linkstate/in_memory_linkstate.py}

@@ -12,31 +12,54 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""In-memory State implementation."""
+"""In-memory LinkState implementation."""
 
 
 import threading
 import time
+from bisect import bisect_right
+from dataclasses import dataclass, field
 from logging import ERROR, WARNING
 from typing import Optional
 from uuid import UUID, uuid4
 
-from flwr.common import log, now
+from flwr.common import Context, log, now
 from flwr.common.constant import (
     MESSAGE_TTL_TOLERANCE,
     NODE_ID_NUM_BYTES,
     RUN_ID_NUM_BYTES,
+    Status,
 )
-from flwr.common.typing import Run, UserConfig
+from flwr.common.record import ConfigsRecord
+from flwr.common.typing import Run, RunStatus, UserConfig
 from flwr.proto.task_pb2 import TaskIns, TaskRes  # pylint: disable=E0611
-from flwr.server.superlink.state.state import State
+from flwr.server.superlink.linkstate.linkstate import LinkState
 from flwr.server.utils import validate_task_ins_or_res
 
-from .utils import generate_rand_int_from_bytes, make_node_unavailable_taskres
+from .utils import (
+    generate_rand_int_from_bytes,
+    has_valid_sub_status,
+    is_valid_transition,
+    make_node_unavailable_taskres,
+)
+
+
+@dataclass
+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)
 
-class InMemoryState(State):  # pylint: disable=R0902,R0904
-    """In-memory State implementation."""
+
+class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
+    """In-memory LinkState implementation."""
 
     def __init__(self) -> None:
 
@@ -44,8 +67,10 @@ class InMemoryState(State):  # pylint: disable=R0902,R0904
         self.node_ids: dict[int, tuple[float, float]] = {}
         self.public_key_to_node_id: dict[bytes, int] = {}
 
-        # Map run_id to (fab_id, fab_version)
-        self.run_ids: dict[int, Run] = {}
+        # Map run_id to RunRecord
+        self.run_ids: dict[int, RunRecord] = {}
+        self.contexts: dict[int, Context] = {}
+        self.federation_options: dict[int, ConfigsRecord] = {}
         self.task_ins_store: dict[UUID, TaskIns] = {}
         self.task_res_store: dict[UUID, TaskRes] = {}
 
@@ -64,8 +89,25 @@ class InMemoryState(State):  # pylint: disable=R0902,R0904
             return None
         # Validate run_id
         if task_ins.run_id not in self.run_ids:
-            log(ERROR, "`run_id` is invalid")
+            log(ERROR, "Invalid run ID for TaskIns: %s", task_ins.run_id)
             return None
+        # Validate source node ID
+        if task_ins.task.producer.node_id != 0:
+            log(
+                ERROR,
+                "Invalid source node ID for TaskIns: %s",
+                task_ins.task.producer.node_id,
+            )
+            return None
+        # Validate destination node ID
+        if not task_ins.task.consumer.anonymous:
+            if task_ins.task.consumer.node_id not in self.node_ids:
+                log(
+                    ERROR,
+                    "Invalid destination node ID for TaskIns: %s",
+                    task_ins.task.consumer.node_id,
+                )
+                return None
 
         # Create task_id
         task_id = uuid4()
@@ -277,7 +319,7 @@ class InMemoryState(State):  # pylint: disable=R0902,R0904
     def create_node(
         self, ping_interval: float, public_key: Optional[bytes] = None
     ) -> int:
-        """Create, store in state, and return `node_id`."""
+        """Create, store in the link state, and return `node_id`."""
         # Sample a random int64 as node_id
         node_id = generate_rand_int_from_bytes(NODE_ID_NUM_BYTES)
 
@@ -338,12 +380,14 @@ class InMemoryState(State):  # pylint: disable=R0902,R0904
         """Retrieve stored `node_id` filtered by `node_public_keys`."""
         return self.public_key_to_node_id.get(node_public_key)
 
+    # pylint: disable=too-many-arguments,too-many-positional-arguments
     def create_run(
         self,
         fab_id: Optional[str],
         fab_version: Optional[str],
         fab_hash: Optional[str],
         override_config: UserConfig,
+        federation_options: ConfigsRecord,
     ) -> int:
         """Create a new run for the specified `fab_hash`."""
         # Sample a random int64 as run_id
@@ -351,13 +395,25 @@ class InMemoryState(State):  # pylint: disable=R0902,R0904
             run_id = generate_rand_int_from_bytes(RUN_ID_NUM_BYTES)
 
             if run_id not in self.run_ids:
-                self.run_ids[run_id] = Run(
-                    run_id=run_id,
-                    fab_id=fab_id if fab_id else "",
-                    fab_version=fab_version if fab_version else "",
-                    fab_hash=fab_hash if fab_hash else "",
-                    override_config=override_config,
+                run_record = RunRecord(
+                    run=Run(
+                        run_id=run_id,
+                        fab_id=fab_id if fab_id else "",
+                        fab_version=fab_version if fab_version else "",
+                        fab_hash=fab_hash if fab_hash else "",
+                        override_config=override_config,
+                    ),
+                    status=RunStatus(
+                        status=Status.PENDING,
+                        sub_status="",
+                        details="",
+                    ),
+                    pending_at=now().isoformat(),
                 )
+                self.run_ids[run_id] = run_record
+
+                # Record federation options. Leave empty if not passed
+                self.federation_options[run_id] = federation_options
                 return run_id
         log(ERROR, "Unexpected run creation failure.")
         return 0
@@ -365,7 +421,7 @@ class InMemoryState(State):  # pylint: disable=R0902,R0904
     def store_server_private_public_key(
         self, private_key: bytes, public_key: bytes
     ) -> None:
-        """Store `server_private_key` and `server_public_key` in state."""
+        """Store `server_private_key` and `server_public_key` in the link state."""
         with self.lock:
             if self.server_private_key is None and self.server_public_key is None:
                 self.server_private_key = private_key
@@ -382,12 +438,12 @@ class InMemoryState(State):  # pylint: disable=R0902,R0904
         return self.server_public_key
 
     def store_node_public_keys(self, public_keys: set[bytes]) -> None:
-        """Store a set of `node_public_keys` in state."""
+        """Store a set of `node_public_keys` in the link state."""
         with self.lock:
             self.node_public_keys = public_keys
 
     def store_node_public_key(self, public_key: bytes) -> None:
-        """Store a `node_public_key` in state."""
+        """Store a `node_public_key` in the link state."""
         with self.lock:
             self.node_public_keys.add(public_key)
 
@@ -401,7 +457,77 @@ class InMemoryState(State):  # pylint: disable=R0902,R0904
             if run_id not in self.run_ids:
                 log(ERROR, "`run_id` is invalid")
                 return None
-            return self.run_ids[run_id]
+            return self.run_ids[run_id].run
+
+    def get_run_status(self, run_ids: set[int]) -> dict[int, RunStatus]:
+        """Retrieve the statuses for the specified runs."""
+        with self.lock:
+            return {
+                run_id: self.run_ids[run_id].status
+                for run_id in set(run_ids)
+                if run_id in self.run_ids
+            }
+
+    def update_run_status(self, run_id: int, new_status: RunStatus) -> bool:
+        """Update the status of the run with the specified `run_id`."""
+        with self.lock:
+            # Check if the run_id exists
+            if run_id not in self.run_ids:
+                log(ERROR, "`run_id` is invalid")
+                return False
+
+            # Check if the status transition is valid
+            current_status = self.run_ids[run_id].status
+            if not is_valid_transition(current_status, new_status):
+                log(
+                    ERROR,
+                    'Invalid status transition: from "%s" to "%s"',
+                    current_status.status,
+                    new_status.status,
+                )
+                return False
+
+            # Check if the sub-status is valid
+            if not has_valid_sub_status(current_status):
+                log(
+                    ERROR,
+                    'Invalid sub-status "%s" for status "%s"',
+                    current_status.sub_status,
+                    current_status.status,
+                )
+                return False
+
+            # Update the status
+            run_record = self.run_ids[run_id]
+            if new_status.status == Status.STARTING:
+                run_record.starting_at = now().isoformat()
+            elif new_status.status == Status.RUNNING:
+                run_record.running_at = now().isoformat()
+            elif new_status.status == Status.FINISHED:
+                run_record.finished_at = now().isoformat()
+            run_record.status = new_status
+            return True
+
+    def get_pending_run_id(self) -> Optional[int]:
+        """Get the `run_id` of a run with `Status.PENDING` status, if any."""
+        pending_run_id = None
+
+        # 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:
+                pending_run_id = run_id
+                break
+
+        return pending_run_id
+
+    def get_federation_options(self, run_id: int) -> Optional[ConfigsRecord]:
+        """Retrieve the federation options for the specified `run_id`."""
+        with self.lock:
+            if run_id not in self.run_ids:
+                log(ERROR, "`run_id` is invalid")
+                return None
+            return self.federation_options[run_id]
 
     def acknowledge_ping(self, node_id: int, ping_interval: float) -> bool:
         """Acknowledge a ping received from a node, serving as a heartbeat."""
@@ -410,3 +536,36 @@ class InMemoryState(State):  # pylint: disable=R0902,R0904
                 self.node_ids[node_id] = (time.time() + ping_interval, ping_interval)
                 return True
         return False
+
+    def get_serverapp_context(self, run_id: int) -> Optional[Context]:
+        """Get the context for the specified `run_id`."""
+        return self.contexts.get(run_id)
+
+    def set_serverapp_context(self, run_id: int, context: Context) -> None:
+        """Set the context for the specified `run_id`."""
+        if run_id not in self.run_ids:
+            raise ValueError(f"Run {run_id} not found")
+        self.contexts[run_id] = context
+
+    def add_serverapp_log(self, run_id: int, log_message: str) -> None:
+        """Add a log entry to the serverapp logs for the specified `run_id`."""
+        if run_id not in self.run_ids:
+            raise ValueError(f"Run {run_id} not found")
+        run = self.run_ids[run_id]
+        with run.log_lock:
+            run.logs.append((now().timestamp(), log_message))
+
+    def get_serverapp_log(
+        self, run_id: int, after_timestamp: Optional[float]
+    ) -> tuple[str, float]:
+        """Get the serverapp logs for the specified `run_id`."""
+        if run_id not in self.run_ids:
+            raise ValueError(f"Run {run_id} not found")
+        run = self.run_ids[run_id]
+        if after_timestamp is None:
+            after_timestamp = 0.0
+        with run.log_lock:
+            # Find the index where the timestamp would be inserted
+            index = bisect_right(run.logs, (after_timestamp, ""))
+            latest_timestamp = run.logs[-1][0] if index < len(run.logs) else 0.0
+            return "".join(log for _, log in run.logs[index:]), latest_timestamp