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

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

@@ -16,6 +16,7 @@
 
 
 import json
+import secrets
 import threading
 import time
 import traceback
@@ -27,12 +28,13 @@ from typing import Callable, Optional
 from uuid import uuid4
 
 from flwr.app.error import Error
-from flwr.client.client_app import ClientApp, ClientAppException, LoadClientAppError
-from flwr.client.clientapp.utils import get_load_client_app_fn
 from flwr.client.run_info_store import DeprecatedRunInfoStore
+from flwr.clientapp.client_app import ClientApp, ClientAppException, LoadClientAppError
+from flwr.clientapp.utils import get_load_client_app_fn
 from flwr.common import Message
 from flwr.common.constant import (
-    HEARTBEAT_MAX_INTERVAL,
+    HEARTBEAT_INTERVAL_INF,
+    NOOP_FLWR_AID,
     NUM_PARTITIONS_KEY,
     PARTITION_ID_KEY,
     ErrorCode,
@@ -40,6 +42,7 @@ from flwr.common.constant import (
 from flwr.common.logger import log
 from flwr.common.typing import Run
 from flwr.server.superlink.linkstate import LinkState, LinkStateFactory
+from flwr.supercore.constant import FLWR_IN_MEMORY_DB_NAME
 
 from .backend import Backend, error_messages_backends, supported_backends
 
@@ -53,7 +56,17 @@ def _register_nodes(
     nodes_mapping: NodeToPartitionMapping = {}
     state = state_factory.state()
     for i in range(num_nodes):
-        node_id = state.create_node(heartbeat_interval=HEARTBEAT_MAX_INTERVAL)
+        node_id = state.create_node(
+            # No node authentication in simulation;
+            # use NOOP_FLWR_AID as owner_aid and
+            # use random bytes as public key
+            NOOP_FLWR_AID,
+            secrets.token_bytes(32),
+            heartbeat_interval=HEARTBEAT_INTERVAL_INF,
+        )
+        state.acknowledge_node_heartbeat(
+            node_id=node_id, heartbeat_interval=HEARTBEAT_INTERVAL_INF
+        )
         nodes_mapping[node_id] = i
     log(DEBUG, "Registered %i nodes", len(nodes_mapping))
     return nodes_mapping
@@ -300,7 +313,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 = LinkStateFactory(":flwr-in-memory-state:")
+        state_factory = LinkStateFactory(FLWR_IN_MEMORY_DB_NAME)
         log(INFO, "Created new %s.", state_factory.__class__.__name__)
 
     if num_supernodes:

flwr/server/superlink/linkstate/in_memory_linkstate.py

@@ -17,17 +17,18 @@
 
 import secrets
 import threading
-import time
 from bisect import bisect_right
 from collections import defaultdict
+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_MAX_INTERVAL,
+    HEARTBEAT_INTERVAL_INF,
     HEARTBEAT_PATIENCE,
     MESSAGE_TTL_TOLERANCE,
     NODE_ID_NUM_BYTES,
@@ -39,8 +40,10 @@ from flwr.common.constant import (
 )
 from flwr.common.record import ConfigRecord
 from flwr.common.typing import Run, RunStatus, UserConfig
+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 .utils import (
     check_node_availability_for_in_message,
@@ -69,10 +72,10 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
 
     def __init__(self) -> None:
 
-        # Map node_id to (online_until, heartbeat_interval)
-        self.node_ids: dict[int, tuple[float, float]] = {}
-        self.public_key_to_node_id: dict[bytes, int] = {}
-        self.node_id_to_public_key: dict[int, bytes] = {}
+        # Map node_id to NodeInfo
+        self.nodes: dict[int, NodeInfo] = {}
+        self.node_public_key_to_node_id: dict[bytes, int] = {}
+        self.owner_to_node_ids: dict[str, set[int]] = {}  # Quick lookup
 
         # Map run_id to RunRecord
         self.run_ids: dict[int, RunRecord] = {}
@@ -114,7 +117,11 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
             )
             return None
         # Validate destination node ID
-        if message.metadata.dst_node_id not in self.node_ids:
+        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,
+        ]:
             log(
                 ERROR,
                 "Invalid destination node ID for Message: %s",
@@ -136,7 +143,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
 
         # Find Message for node_id that were not delivered yet
         message_ins_list: list[Message] = []
-        current_time = time.time()
+        current_time = now().timestamp()
         with self.lock:
             for _, msg_ins in self.message_ins_store.items():
                 if (
@@ -190,7 +197,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
                 return None
 
             ins_metadata = msg_ins.metadata
-            if ins_metadata.created_at + ins_metadata.ttl <= time.time():
+            if ins_metadata.created_at + ins_metadata.ttl <= now().timestamp():
                 log(
                     ERROR,
                     "Failed to store Message: the message it is replying to "
@@ -238,7 +245,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         ret: dict[str, Message] = {}
 
         with self.lock:
-            current = time.time()
+            current = now().timestamp()
 
             # Verify Message IDs
             ret = verify_message_ids(
@@ -256,9 +263,10 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
                 inquired_in_message_ids=message_ids,
                 found_in_message_dict=self.message_ins_store,
                 node_id_to_online_until={
-                    node_id: self.node_ids[node_id][0]
+                    node_id: self.nodes[node_id].online_until
                     for node_id in dst_node_ids
-                    if node_id in self.node_ids
+                    if node_id in self.nodes
+                    and self.nodes[node_id].status != NodeStatus.UNREGISTERED
                 },
                 current_time=current,
             )
@@ -330,7 +338,9 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         """
         return len(self.message_res_store)
 
-    def create_node(self, heartbeat_interval: float) -> int:
+    def create_node(
+        self, owner_aid: 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
         node_id = generate_rand_int_from_bytes(
@@ -338,28 +348,88 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         )
 
         with self.lock:
-            if node_id in self.node_ids:
+            if node_id in self.nodes:
                 log(ERROR, "Unexpected node registration failure.")
                 return 0
+            if public_key in self.node_public_key_to_node_id:
+                raise ValueError("Public key already in use")
 
-            # Mark the node online until time.time() + heartbeat_interval
-            self.node_ids[node_id] = (
-                time.time() + heartbeat_interval,
-                heartbeat_interval,
+            # The node is not activated upon creation
+            self.nodes[node_id] = NodeInfo(
+                node_id=node_id,
+                owner_aid=owner_aid,
+                status=NodeStatus.REGISTERED,
+                registered_at=now().isoformat(),
+                last_activated_at=None,
+                last_deactivated_at=None,
+                unregistered_at=None,
+                online_until=None,
+                heartbeat_interval=heartbeat_interval,
+                public_key=public_key,
             )
+            self.node_public_key_to_node_id[public_key] = node_id
+            self.owner_to_node_ids.setdefault(owner_aid, set()).add(node_id)
             return node_id
 
-    def delete_node(self, node_id: int) -> None:
+    def delete_node(self, owner_aid: str, node_id: int) -> None:
         """Delete a node."""
         with self.lock:
-            if node_id not in self.node_ids:
-                raise ValueError(f"Node {node_id} not found")
+            if (
+                not (node := self.nodes.get(node_id))
+                or node.status == NodeStatus.UNREGISTERED
+                or owner_aid != self.nodes[node_id].owner_aid
+            ):
+                raise ValueError(
+                    f"Node ID {node_id} already unregistered, not found or "
+                    "the request was unauthorized."
+                )
+
+            node.status = NodeStatus.UNREGISTERED
+            current = now()
+            node.unregistered_at = current.isoformat()
+            # Set online_until to current timestamp on deletion, if it is in the future
+            node.online_until = min(node.online_until, current.timestamp())
+
+    def activate_node(self, node_id: int, heartbeat_interval: float) -> bool:
+        """Activate the node with the specified `node_id`."""
+        with self.lock:
+            self._check_and_tag_offline_nodes(node_ids=[node_id])
 
-            # Remove node ID <> public key mappings
-            if pk := self.node_id_to_public_key.pop(node_id, None):
-                del self.public_key_to_node_id[pk]
+            # Check if the node exists
+            if not (node := self.nodes.get(node_id)):
+                return False
+
+            # Only activate if the node is currently registered or offline
+            current_dt = now()
+            if node.status in (NodeStatus.REGISTERED, NodeStatus.OFFLINE):
+                node.status = NodeStatus.ONLINE
+                node.last_activated_at = current_dt.isoformat()
+                node.online_until = (
+                    current_dt.timestamp() + HEARTBEAT_PATIENCE * heartbeat_interval
+                )
+                node.heartbeat_interval = heartbeat_interval
+                return True
+            return False
+
+    def deactivate_node(self, node_id: int) -> bool:
+        """Deactivate the node with the specified `node_id`."""
+        with self.lock:
+            self._check_and_tag_offline_nodes(node_ids=[node_id])
+
+            # Check if the node exists
+            if not (node := self.nodes.get(node_id)):
+                return False
 
-            del self.node_ids[node_id]
+            # Only deactivate if the node is currently online
+            current_dt = now()
+            if node.status == NodeStatus.ONLINE:
+                node.status = NodeStatus.OFFLINE
+                node.last_deactivated_at = current_dt.isoformat()
+
+                # Set online_until to current timestamp
+                node.online_until = current_dt.timestamp()
+                return True
+            return False
 
     def get_nodes(self, run_id: int) -> set[int]:
         """Return all available nodes.
@@ -372,36 +442,70 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         with self.lock:
             if run_id not in self.run_ids:
                 return set()
-            current_time = time.time()
             return {
-                node_id
-                for node_id, (online_until, _) in self.node_ids.items()
-                if online_until > current_time
+                node.node_id
+                for node in self.get_node_info(statuses=[NodeStatus.ONLINE])
             }
 
-    def set_node_public_key(self, node_id: int, public_key: bytes) -> None:
-        """Set `public_key` for the specified `node_id`."""
+    def get_node_info(
+        self,
+        *,
+        node_ids: Optional[Sequence[int]] = None,
+        owner_aids: Optional[Sequence[str]] = None,
+        statuses: Optional[Sequence[str]] = None,
+    ) -> Sequence[NodeInfo]:
+        """Retrieve information about nodes based on the specified filters."""
         with self.lock:
-            if node_id not in self.node_ids:
-                raise ValueError(f"Node {node_id} not found")
-
-            if public_key in self.public_key_to_node_id:
-                raise ValueError("Public key already in use")
-
-            self.public_key_to_node_id[public_key] = node_id
-            self.node_id_to_public_key[node_id] = public_key
-
-    def get_node_public_key(self, node_id: int) -> Optional[bytes]:
+            self._check_and_tag_offline_nodes()
+            result = []
+            for node_id in self.nodes.keys() if node_ids is None else node_ids:
+                if (node := self.nodes.get(node_id)) is None:
+                    continue
+                if owner_aids is not None and node.owner_aid not in owner_aids:
+                    continue
+                if statuses is not None and node.status not in statuses:
+                    continue
+                result.append(node)
+            return result
+
+    def _check_and_tag_offline_nodes(
+        self, node_ids: Optional[list[int]] = None
+    ) -> None:
+        with self.lock:
+            # Set all nodes of "online" status to "offline" if they've offline
+            current_ts = now().timestamp()
+            for node_id in node_ids or self.nodes.keys():
+                if (node := self.nodes.get(node_id)) is None:
+                    continue
+                if node.status == NodeStatus.ONLINE:
+                    if node.online_until <= current_ts:
+                        node.status = NodeStatus.OFFLINE
+                        node.last_deactivated_at = datetime.fromtimestamp(
+                            node.online_until, tz=timezone.utc
+                        ).isoformat()
+
+    def get_node_public_key(self, node_id: int) -> bytes:
         """Get `public_key` for the specified `node_id`."""
         with self.lock:
-            if node_id not in self.node_ids:
-                raise ValueError(f"Node {node_id} not found")
+            if (
+                node := self.nodes.get(node_id)
+            ) is None or node.status == NodeStatus.UNREGISTERED:
+                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]:
+        """Get `node_id` for the specified `public_key` if it exists and is not
+        deleted."""
+        with self.lock:
+            node_id = self.node_public_key_to_node_id.get(public_key)
 
-            return self.node_id_to_public_key.get(node_id)
+            if node_id is None:
+                return None
 
-    def get_node_id(self, node_public_key: bytes) -> Optional[int]:
-        """Retrieve stored `node_id` filtered by `node_public_keys`."""
-        return self.public_key_to_node_id.get(node_public_key)
+            node_info = self.nodes[node_id]
+            if node_info.status == NodeStatus.UNREGISTERED:
+                return None
+            return node_id
 
     # pylint: disable=too-many-arguments,too-many-positional-arguments
     def create_run(
@@ -449,26 +553,6 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         log(ERROR, "Unexpected run creation failure.")
         return 0
 
-    def clear_supernode_auth_keys(self) -> None:
-        """Clear stored `node_public_keys` in the link state if any."""
-        with self.lock:
-            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:
-            self.node_public_keys.update(public_keys)
-
-    def store_node_public_key(self, public_key: bytes) -> None:
-        """Store a `node_public_key` in the link state."""
-        with self.lock:
-            self.node_public_keys.add(public_key)
-
-    def get_node_public_keys(self) -> set[bytes]:
-        """Retrieve all currently stored `node_public_keys` as a set."""
-        with self.lock:
-            return self.node_public_keys.copy()
-
     def get_run_ids(self, flwr_aid: Optional[str]) -> set[int]:
         """Retrieve all run IDs if `flwr_aid` is not specified.
 
@@ -561,7 +645,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
             current = now()
             run_record = self.run_ids[run_id]
             if new_status.status in (Status.STARTING, Status.RUNNING):
-                run_record.heartbeat_interval = HEARTBEAT_MAX_INTERVAL
+                run_record.heartbeat_interval = HEARTBEAT_INTERVAL_INF
                 run_record.active_until = (
                     current.timestamp() + run_record.heartbeat_interval
                 )
@@ -608,13 +692,23 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         the node is marked as offline.
         """
         with self.lock:
-            if node_id in self.node_ids:
-                self.node_ids[node_id] = (
-                    time.time() + HEARTBEAT_PATIENCE * heartbeat_interval,
-                    heartbeat_interval,
+            if (
+                node := self.nodes.get(node_id)
+            ) and node.status != NodeStatus.UNREGISTERED:
+                current_dt = now()
+
+                # Set timestamp if the status changes
+                if node.status != NodeStatus.ONLINE:  # offline or registered
+                    node.status = NodeStatus.ONLINE
+                    node.last_activated_at = current_dt.isoformat()
+
+                # Refresh `online_until` and `heartbeat_interval`
+                node.online_until = (
+                    current_dt.timestamp() + HEARTBEAT_PATIENCE * heartbeat_interval
                 )
+                node.heartbeat_interval = heartbeat_interval
                 return True
-        return False
+            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.

flwr/server/superlink/linkstate/linkstate.py

@@ -16,11 +16,13 @@
 
 
 import abc
+from collections.abc import Sequence
 from typing import Optional
 
 from flwr.common import Context, Message
 from flwr.common.record import ConfigRecord
 from flwr.common.typing import Run, RunStatus, UserConfig
+from flwr.proto.node_pb2 import NodeInfo  # pylint: disable=E0611
 from flwr.supercore.corestate import CoreState
 
 
@@ -128,13 +130,54 @@ class LinkState(CoreState):  # pylint: disable=R0904
         """Get all instruction Message IDs for the given run_id."""
 
     @abc.abstractmethod
-    def create_node(self, heartbeat_interval: float) -> int:
+    def create_node(
+        self, owner_aid: str, public_key: bytes, heartbeat_interval: float
+    ) -> int:
         """Create, store in the link state, and return `node_id`."""
 
     @abc.abstractmethod
-    def delete_node(self, node_id: int) -> None:
+    def delete_node(self, owner_aid: str, node_id: int) -> None:
         """Remove `node_id` from the link state."""
 
+    @abc.abstractmethod
+    def activate_node(self, node_id: int, heartbeat_interval: float) -> bool:
+        """Activate the node with the specified `node_id`.
+
+        Transitions the node status to "online". The transition will fail
+        if the current status is not "registered" or "offline".
+
+        Parameters
+        ----------
+        node_id : int
+            The identifier of the node to activate.
+        heartbeat_interval : float
+            The interval (in seconds) from the current timestamp within which
+            the next heartbeat from this node is expected to be received.
+
+        Returns
+        -------
+        bool
+            True if the status transition was successful, False otherwise.
+        """
+
+    @abc.abstractmethod
+    def deactivate_node(self, node_id: int) -> bool:
+        """Deactivate the node with the specified `node_id`.
+
+        Transitions the node status to "offline". The transition will fail
+        if the current status is not "online".
+
+        Parameters
+        ----------
+        node_id : int
+            The identifier of the node to deactivate.
+
+        Returns
+        -------
+        bool
+            True if the status transition was successful, False otherwise.
+        """
+
     @abc.abstractmethod
     def get_nodes(self, run_id: int) -> set[int]:
         """Retrieve all currently stored node IDs as a set.
@@ -146,16 +189,72 @@ class LinkState(CoreState):  # pylint: disable=R0904
         """
 
     @abc.abstractmethod
-    def set_node_public_key(self, node_id: int, public_key: bytes) -> None:
-        """Set `public_key` for the specified `node_id`."""
+    def get_node_id_by_public_key(self, public_key: bytes) -> Optional[int]:
+        """Get `node_id` for the specified `public_key` if it exists and is not deleted.
+
+        Parameters
+        ----------
+        public_key : bytes
+            The public key of the node whose information is to be retrieved.
+
+        Returns
+        -------
+        Optional[int]
+            The `node_id` associated with the specified `public_key` if it exists
+            and is not deleted; otherwise, `None`.
+        """
 
     @abc.abstractmethod
-    def get_node_public_key(self, node_id: int) -> Optional[bytes]:
-        """Get `public_key` for the specified `node_id`."""
+    def get_node_info(
+        self,
+        *,
+        node_ids: Optional[Sequence[int]] = None,
+        owner_aids: Optional[Sequence[str]] = None,
+        statuses: Optional[Sequence[str]] = None,
+    ) -> Sequence[NodeInfo]:
+        """Retrieve information about nodes based on the specified filters.
+
+        If a filter is set to None, it is ignored.
+        If multiple filters are provided, they are combined using AND logic.
+
+        Parameters
+        ----------
+        node_ids : Optional[Sequence[int]] (default: None)
+            Sequence of node IDs to filter by. If a sequence is provided,
+            it is treated as an OR condition.
+        owner_aids : Optional[Sequence[str]] (default: None)
+            Sequence of owner account IDs to filter by. If a sequence is provided,
+            it is treated as an OR condition.
+        statuses : Optional[Sequence[str]] (default: None)
+            Sequence of node status values (e.g., "created", "activated")
+            to filter by. If a sequence is provided, it is treated as an OR condition.
+
+        Returns
+        -------
+        Sequence[NodeInfo]
+            A sequence of NodeInfo objects representing the nodes matching
+            the specified filters.
+        """
 
     @abc.abstractmethod
-    def get_node_id(self, node_public_key: bytes) -> Optional[int]:
-        """Retrieve stored `node_id` filtered by `node_public_keys`."""
+    def get_node_public_key(self, node_id: int) -> bytes:
+        """Get `public_key` for the specified `node_id`.
+
+        Parameters
+        ----------
+        node_id : int
+            The identifier of the node whose public key is to be retrieved.
+
+        Returns
+        -------
+        bytes
+            The public key associated with the specified `node_id`.
+
+        Raises
+        ------
+        ValueError
+            If the specified `node_id` does not exist in the link state.
+        """
 
     @abc.abstractmethod
     def create_run(  # pylint: disable=too-many-arguments,too-many-positional-arguments
@@ -254,22 +353,6 @@ class LinkState(CoreState):  # pylint: disable=R0904
             The federation options for the run if it exists; None otherwise.
         """
 
-    @abc.abstractmethod
-    def clear_supernode_auth_keys(self) -> None:
-        """Clear stored `node_public_keys` in the link state if any."""
-
-    @abc.abstractmethod
-    def store_node_public_keys(self, public_keys: set[bytes]) -> None:
-        """Store a set of `node_public_keys` in the link state."""
-
-    @abc.abstractmethod
-    def store_node_public_key(self, public_key: bytes) -> None:
-        """Store a `node_public_key` in the link state."""
-
-    @abc.abstractmethod
-    def get_node_public_keys(self) -> set[bytes]:
-        """Retrieve all currently stored `node_public_keys` as a set."""
-
     @abc.abstractmethod
     def acknowledge_node_heartbeat(
         self, node_id: int, heartbeat_interval: float

flwr/server/superlink/linkstate/linkstate_factory.py

@@ -19,6 +19,7 @@ from logging import DEBUG
 from typing import Optional
 
 from flwr.common.logger import log
+from flwr.supercore.constant import FLWR_IN_MEMORY_DB_NAME
 
 from .in_memory_linkstate import InMemoryLinkState
 from .linkstate import LinkState
@@ -44,7 +45,7 @@ class LinkStateFactory:
     def state(self) -> LinkState:
         """Return a State instance and create it, if necessary."""
         # InMemoryState
-        if self.database == ":flwr-in-memory-state:":
+        if self.database == FLWR_IN_MEMORY_DB_NAME:
             if self.state_instance is None:
                 self.state_instance = InMemoryLinkState()
             log(DEBUG, "Using InMemoryState")