flwr 1.23.0__py3-none-any.whl → 1.24.0__py3-none-any.whl

flwr/server/superlink/linkstate/in_memory_linkstate.py

@@ -15,7 +15,6 @@
 """In-memory LinkState implementation."""
 
 
-import secrets
 import threading
 from bisect import bisect_right
 from collections import defaultdict
@@ -23,12 +22,9 @@ from collections.abc import Sequence
 from dataclasses import dataclass, field
 from datetime import datetime, timezone
 from logging import ERROR, WARNING
-from typing import Optional
 
 from flwr.common import Context, Message, log, now
 from flwr.common.constant import (
-    FLWR_APP_TOKEN_LENGTH,
-    HEARTBEAT_INTERVAL_INF,
     HEARTBEAT_PATIENCE,
     MESSAGE_TTL_TOLERANCE,
     NODE_ID_NUM_BYTES,
@@ -44,6 +40,9 @@ from flwr.proto.node_pb2 import NodeInfo  # pylint: disable=E0611
 from flwr.server.superlink.linkstate.linkstate import LinkState
 from flwr.server.utils import validate_message
 from flwr.supercore.constant import NodeStatus
+from flwr.supercore.corestate.in_memory_corestate import InMemoryCoreState
+from flwr.supercore.object_store.object_store import ObjectStore
+from flwr.superlink.federation import FederationManager
 
 from .utils import (
     check_node_availability_for_in_message,
@@ -60,17 +59,18 @@ class RunRecord:  # pylint: disable=R0902
     """The record of a specific run, including its status and timestamps."""
 
     run: Run
-    active_until: float = 0.0
-    heartbeat_interval: float = 0.0
     logs: list[tuple[float, str]] = field(default_factory=list)
     log_lock: threading.Lock = field(default_factory=threading.Lock)
     lock: threading.RLock = field(default_factory=threading.RLock)
 
 
-class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
+class InMemoryLinkState(LinkState, InMemoryCoreState):  # pylint: disable=R0902,R0904
     """In-memory LinkState implementation."""
 
-    def __init__(self) -> None:
+    def __init__(
+        self, federation_manager: FederationManager, object_store: ObjectStore
+    ) -> None:
+        super().__init__(object_store)
 
         # Map node_id to NodeInfo
         self.nodes: dict[int, NodeInfo] = {}
@@ -85,19 +85,21 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         self.message_res_store: dict[str, Message] = {}
         self.message_ins_id_to_message_res_id: dict[str, str] = {}
 
-        # Store run ID to token mapping and token to run ID mapping
-        self.token_store: dict[int, str] = {}
-        self.token_to_run_id: dict[str, int] = {}
-        self.lock_token_store = threading.Lock()
-
         # Map flwr_aid to run_ids for O(1) reverse index lookup
         self.flwr_aid_to_run_ids: dict[str, set[int]] = defaultdict(set)
 
         self.node_public_keys: set[bytes] = set()
 
         self.lock = threading.RLock()
+        federation_manager.linkstate = self
+        self._federation_manager = federation_manager
+
+    @property
+    def federation_manager(self) -> FederationManager:
+        """Get the FederationManager instance."""
+        return self._federation_manager
 
-    def store_message_ins(self, message: Message) -> Optional[str]:
+    def store_message_ins(self, message: Message) -> str | None:
         """Store one Message."""
         # Validate message
         errors = validate_message(message, is_reply_message=False)
@@ -108,6 +110,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         if message.metadata.run_id not in self.run_ids:
             log(ERROR, "Invalid run ID for Message: %s", message.metadata.run_id)
             return None
+        federation = self.run_ids[message.metadata.run_id].run.federation
         # Validate source node ID
         if message.metadata.src_node_id != SUPERLINK_NODE_ID:
             log(
@@ -118,10 +121,14 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
             return None
         # Validate destination node ID
         dst_node = self.nodes.get(message.metadata.dst_node_id)
-        if dst_node is None or dst_node.status not in [
-            NodeStatus.ONLINE,
-            NodeStatus.OFFLINE,
-        ]:
+        if (
+            # Node must exist
+            dst_node is None
+            # Node must be online or offline
+            or dst_node.status not in (NodeStatus.ONLINE, NodeStatus.OFFLINE)
+            # Node must belong to the same federation
+            or not self.federation_manager.has_node(dst_node.node_id, federation)
+        ):
             log(
                 ERROR,
                 "Invalid destination node ID for Message: %s",
@@ -136,21 +143,50 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         # Return the new message_id
         return message_id
 
-    def get_message_ins(self, node_id: int, limit: Optional[int]) -> list[Message]:
+    def _check_stored_messages(self, message_ids: set[str]) -> None:
+        """Check and delete the message if it's invalid."""
+        with self.lock:
+            invalid_msg_ids: set[str] = set()
+            current = now().timestamp()
+            for msg_id in message_ids:
+                if not (message := self.message_ins_store.get(msg_id)):
+                    continue
+
+                # Check if the message has expired
+                available_until = message.metadata.created_at + message.metadata.ttl
+                if available_until <= current:
+                    invalid_msg_ids.add(msg_id)
+                    continue
+
+                # Check if the destination node and the source node are still in the
+                # same federation
+                src_node_id = message.metadata.src_node_id
+                dst_node_id = message.metadata.dst_node_id
+                filtered = self.federation_manager.filter_nodes(
+                    {src_node_id, dst_node_id},
+                    self.run_ids[message.metadata.run_id].run.federation,
+                )
+                if len(filtered) != 2:  # Not both nodes are in the federation
+                    invalid_msg_ids.add(msg_id)
+
+            # Delete all invalid messages
+            self.delete_messages(invalid_msg_ids)
+
+    def get_message_ins(self, node_id: int, limit: int | None) -> list[Message]:
         """Get all Messages that have not been delivered yet."""
         if limit is not None and limit < 1:
             raise AssertionError("`limit` must be >= 1")
 
         # Find Message for node_id that were not delivered yet
         message_ins_list: list[Message] = []
-        current_time = now().timestamp()
         with self.lock:
-            for _, msg_ins in self.message_ins_store.items():
+            for msg_id in list(self.message_ins_store.keys()):
+                self._check_stored_messages({msg_id})
+
                 if (
-                    msg_ins.metadata.dst_node_id == node_id
+                    (msg_ins := self.message_ins_store.get(msg_id))
+                    and msg_ins.metadata.dst_node_id == node_id
                     and msg_ins.metadata.delivered_at == ""
-                    and msg_ins.metadata.created_at + msg_ins.metadata.ttl
-                    > current_time
                 ):
                     message_ins_list.append(msg_ins)
                 if limit and len(message_ins_list) == limit:
@@ -165,7 +201,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         return message_ins_list
 
     # pylint: disable=R0911
-    def store_message_res(self, message: Message) -> Optional[str]:
+    def store_message_res(self, message: Message) -> str | None:
         """Store one Message."""
         # Validate message
         errors = validate_message(message, is_reply_message=True)
@@ -177,6 +213,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         with self.lock:
             # Check if the Message it is replying to exists and is valid
             msg_ins_id = res_metadata.reply_to_message_id
+            self._check_stored_messages({msg_ins_id})
             msg_ins = self.message_ins_store.get(msg_ins_id)
 
             # Ensure that dst_node_id of original Message matches the src_node_id of
@@ -196,22 +233,13 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
                 )
                 return None
 
-            ins_metadata = msg_ins.metadata
-            if ins_metadata.created_at + ins_metadata.ttl <= now().timestamp():
-                log(
-                    ERROR,
-                    "Failed to store Message: the message it is replying to "
-                    "(with ID %s) has expired",
-                    msg_ins_id,
-                )
-                return None
-
             # Fail if the Message TTL exceeds the
             # expiration time of the Message it replies to.
             # Condition: ins_metadata.created_at + ins_metadata.ttl ≥
             #            res_metadata.created_at + res_metadata.ttl
             # A small tolerance is introduced to account
             # for floating-point precision issues.
+            ins_metadata = msg_ins.metadata
             max_allowed_ttl = (
                 ins_metadata.created_at + ins_metadata.ttl - res_metadata.created_at
             )
@@ -245,6 +273,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         ret: dict[str, Message] = {}
 
         with self.lock:
+            self._check_stored_messages(message_ids)
             current = now().timestamp()
 
             # Verify Message IDs
@@ -339,7 +368,11 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         return len(self.message_res_store)
 
     def create_node(
-        self, owner_aid: str, public_key: bytes, heartbeat_interval: float
+        self,
+        owner_aid: str,
+        owner_name: str,
+        public_key: bytes,
+        heartbeat_interval: float,
     ) -> int:
         """Create, store in the link state, and return `node_id`."""
         # Sample a random int64 as node_id
@@ -358,6 +391,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
             self.nodes[node_id] = NodeInfo(
                 node_id=node_id,
                 owner_aid=owner_aid,
+                owner_name=owner_name,
                 status=NodeStatus.REGISTERED,
                 registered_at=now().isoformat(),
                 last_activated_at=None,
@@ -442,17 +476,19 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         with self.lock:
             if run_id not in self.run_ids:
                 return set()
-            return {
+            federation = self.run_ids[run_id].run.federation
+            node_ids = {
                 node.node_id
                 for node in self.get_node_info(statuses=[NodeStatus.ONLINE])
             }
+            return self.federation_manager.filter_nodes(node_ids, federation)
 
     def get_node_info(
         self,
         *,
-        node_ids: Optional[Sequence[int]] = None,
-        owner_aids: Optional[Sequence[str]] = None,
-        statuses: Optional[Sequence[str]] = None,
+        node_ids: Sequence[int] | None = None,
+        owner_aids: Sequence[str] | None = None,
+        statuses: Sequence[str] | None = None,
     ) -> Sequence[NodeInfo]:
         """Retrieve information about nodes based on the specified filters."""
         with self.lock:
@@ -468,9 +504,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
                 result.append(node)
             return result
 
-    def _check_and_tag_offline_nodes(
-        self, node_ids: Optional[list[int]] = None
-    ) -> None:
+    def _check_and_tag_offline_nodes(self, node_ids: list[int] | None = None) -> None:
         with self.lock:
             # Set all nodes of "online" status to "offline" if they've offline
             current_ts = now().timestamp()
@@ -493,7 +527,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
                 raise ValueError(f"Node ID {node_id} not found")
             return node.public_key
 
-    def get_node_id_by_public_key(self, public_key: bytes) -> Optional[int]:
+    def get_node_id_by_public_key(self, public_key: bytes) -> int | None:
         """Get `node_id` for the specified `public_key` if it exists and is not
         deleted."""
         with self.lock:
@@ -510,14 +544,15 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
     # 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],
+        fab_id: str | None,
+        fab_version: str | None,
+        fab_hash: str | None,
         override_config: UserConfig,
+        federation: str,
         federation_options: ConfigRecord,
-        flwr_aid: Optional[str],
+        flwr_aid: str | None,
     ) -> int:
-        """Create a new run for the specified `fab_hash`."""
+        """Create a new run."""
         # Sample a random int64 as run_id
         with self.lock:
             run_id = generate_rand_int_from_bytes(RUN_ID_NUM_BYTES)
@@ -540,6 +575,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
                             details="",
                         ),
                         flwr_aid=flwr_aid if flwr_aid else "",
+                        federation=federation,
                     ),
                 )
                 self.run_ids[run_id] = run_record
@@ -553,7 +589,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         log(ERROR, "Unexpected run creation failure.")
         return 0
 
-    def get_run_ids(self, flwr_aid: Optional[str]) -> set[int]:
+    def get_run_ids(self, flwr_aid: str | None) -> set[int]:
         """Retrieve all run IDs if `flwr_aid` is not specified.
 
         Otherwise, retrieve all run IDs for the specified `flwr_aid`.
@@ -564,30 +600,10 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
                 return set(self.flwr_aid_to_run_ids.get(flwr_aid, ()))
             return set(self.run_ids.keys())
 
-    def _check_and_tag_inactive_run(self, run_ids: set[int]) -> None:
-        """Check if any runs are no longer active.
-
-        Marks runs with status 'starting' or 'running' as failed
-        if they have not sent a heartbeat before `active_until`.
-        """
-        current = now()
-        for record in (self.run_ids.get(run_id) for run_id in run_ids):
-            if record is None:
-                continue
-            with record.lock:
-                if record.run.status.status in (Status.STARTING, Status.RUNNING):
-                    if record.active_until < current.timestamp():
-                        record.run.status = RunStatus(
-                            status=Status.FINISHED,
-                            sub_status=SubStatus.FAILED,
-                            details=RUN_FAILURE_DETAILS_NO_HEARTBEAT,
-                        )
-                        record.run.finished_at = now().isoformat()
-
-    def get_run(self, run_id: int) -> Optional[Run]:
+    def get_run(self, run_id: int) -> Run | None:
         """Retrieve information about the run with the specified `run_id`."""
-        # Check if runs are still active
-        self._check_and_tag_inactive_run(run_ids={run_id})
+        # Clean up expired tokens; this will flag inactive runs as needed
+        self._cleanup_expired_tokens()
 
         with self.lock:
             if run_id not in self.run_ids:
@@ -597,8 +613,8 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
 
     def get_run_status(self, run_ids: set[int]) -> dict[int, RunStatus]:
         """Retrieve the statuses for the specified runs."""
-        # Check if runs are still active
-        self._check_and_tag_inactive_run(run_ids=run_ids)
+        # Clean up expired tokens; this will flag inactive runs as needed
+        self._cleanup_expired_tokens()
 
         with self.lock:
             return {
@@ -609,8 +625,8 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
 
     def update_run_status(self, run_id: int, new_status: RunStatus) -> bool:
         """Update the status of the run with the specified `run_id`."""
-        # Check if runs are still active
-        self._check_and_tag_inactive_run(run_ids={run_id})
+        # Clean up expired tokens; this will flag inactive runs as needed
+        self._cleanup_expired_tokens()
 
         with self.lock:
             # Check if the run_id exists
@@ -640,17 +656,9 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
                 )
                 return False
 
-            # Initialize heartbeat_interval and active_until
-            # when switching to starting or running
+            # Update the run status
             current = now()
             run_record = self.run_ids[run_id]
-            if new_status.status in (Status.STARTING, Status.RUNNING):
-                run_record.heartbeat_interval = HEARTBEAT_INTERVAL_INF
-                run_record.active_until = (
-                    current.timestamp() + run_record.heartbeat_interval
-                )
-
-            # Update the run status
             if new_status.status == Status.STARTING:
                 run_record.run.starting_at = current.isoformat()
             elif new_status.status == Status.RUNNING:
@@ -660,7 +668,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
             run_record.run.status = new_status
             return True
 
-    def get_pending_run_id(self) -> Optional[int]:
+    def get_pending_run_id(self) -> int | None:
         """Get the `run_id` of a run with `Status.PENDING` status, if any."""
         pending_run_id = None
 
@@ -673,7 +681,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
 
         return pending_run_id
 
-    def get_federation_options(self, run_id: int) -> Optional[ConfigRecord]:
+    def get_federation_options(self, run_id: int) -> ConfigRecord | None:
         """Retrieve the federation options for the specified `run_id`."""
         with self.lock:
             if run_id not in self.run_ids:
@@ -710,44 +718,28 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
                 return True
             return False
 
-    def acknowledge_app_heartbeat(self, run_id: int, heartbeat_interval: float) -> bool:
-        """Acknowledge a heartbeat received from a ServerApp for a given run.
+    def _on_tokens_expired(self, expired_records: list[tuple[int, float]]) -> None:
+        """Transition runs with expired tokens to failed status.
 
-        A run with status `"running"` is considered alive as long as it sends heartbeats
-        within the tolerated interval: HEARTBEAT_PATIENCE × heartbeat_interval.
-        HEARTBEAT_PATIENCE = N allows for N-1 missed heartbeat before the run is
-        marked as `"completed:failed"`.
+        Parameters
+        ----------
+        expired_records : list[tuple[int, float]]
+            List of tuples containing (run_id, active_until timestamp)
+            for expired tokens.
         """
-        with self.lock:
-            # Search for the run
-            record = self.run_ids.get(run_id)
-
-            # Check if the run_id exists
-            if record is None:
-                log(ERROR, "`run_id` is invalid")
-                return False
-
-        with record.lock:
-            # Check if runs are still active
-            self._check_and_tag_inactive_run(run_ids={run_id})
-
-            # Check if the run is of status "running"/"starting"
-            current_status = record.run.status
-            if current_status.status not in (Status.RUNNING, Status.STARTING):
-                log(
-                    ERROR,
-                    'Cannot acknowledge heartbeat for run with status "%s"',
-                    current_status.status,
+        for run_id, active_until in expired_records:
+            if not (run_record := self.run_ids.get(run_id)):
+                continue
+            with run_record.lock:
+                run_record.run.status = RunStatus(
+                    status=Status.FINISHED,
+                    sub_status=SubStatus.FAILED,
+                    details=RUN_FAILURE_DETAILS_NO_HEARTBEAT,
                 )
-                return False
+                active_until_dt = datetime.fromtimestamp(active_until, tz=timezone.utc)
+                run_record.run.finished_at = active_until_dt.isoformat()
 
-            # Update the `active_until` and `heartbeat_interval` for the given run
-            current = now().timestamp()
-            record.active_until = current + HEARTBEAT_PATIENCE * heartbeat_interval
-            record.heartbeat_interval = heartbeat_interval
-            return True
-
-    def get_serverapp_context(self, run_id: int) -> Optional[Context]:
+    def get_serverapp_context(self, run_id: int) -> Context | None:
         """Get the context for the specified `run_id`."""
         return self.contexts.get(run_id)
 
@@ -766,7 +758,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
             run.logs.append((now().timestamp(), log_message))
 
     def get_serverapp_log(
-        self, run_id: int, after_timestamp: Optional[float]
+        self, run_id: int, after_timestamp: float | None
     ) -> tuple[str, float]:
         """Get the serverapp logs for the specified `run_id`."""
         if run_id not in self.run_ids:
@@ -779,30 +771,3 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
             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
-
-    def create_token(self, run_id: int) -> Optional[str]:
-        """Create a token for the given run ID."""
-        token = secrets.token_hex(FLWR_APP_TOKEN_LENGTH)  # Generate a random token
-        with self.lock_token_store:
-            if run_id in self.token_store:
-                return None  # Token already created for this run ID
-            self.token_store[run_id] = token
-            self.token_to_run_id[token] = run_id
-        return token
-
-    def verify_token(self, run_id: int, token: str) -> bool:
-        """Verify a token for the given run ID."""
-        with self.lock_token_store:
-            return self.token_store.get(run_id) == token
-
-    def delete_token(self, run_id: int) -> None:
-        """Delete the token for the given run ID."""
-        with self.lock_token_store:
-            token = self.token_store.pop(run_id, None)
-            if token is not None:
-                self.token_to_run_id.pop(token, None)
-
-    def get_run_id_by_token(self, token: str) -> Optional[int]:
-        """Get the run ID associated with a given token."""
-        with self.lock_token_store:
-            return self.token_to_run_id.get(token)