flwr 1.18.0__py3-none-any.whl → 1.20.0__py3-none-any.whl

flwr/supernode/nodestate/in_memory_nodestate.py

@@ -0,0 +1,199 @@
+# Copyright 2025 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.
+# ==============================================================================
+"""In-memory NodeState implementation."""
+
+
+import secrets
+from collections.abc import Sequence
+from dataclasses import dataclass
+from threading import Lock
+from typing import Optional
+
+from flwr.common import Context, Message
+from flwr.common.constant import FLWR_APP_TOKEN_LENGTH
+from flwr.common.typing import Run
+
+from .nodestate import NodeState
+
+
+@dataclass
+class MessageEntry:
+    """Data class to represent a message entry."""
+
+    message: Message
+    is_retrieved: bool = False
+
+
+class InMemoryNodeState(NodeState):  # pylint: disable=too-many-instance-attributes
+    """In-memory NodeState implementation."""
+
+    def __init__(self) -> None:
+        # Store node_id
+        self.node_id: Optional[int] = None
+        # Store Object ID to MessageEntry mapping
+        self.msg_store: dict[str, MessageEntry] = {}
+        self.lock_msg_store = Lock()
+        # Store run ID to Run mapping
+        self.run_store: dict[int, Run] = {}
+        self.lock_run_store = Lock()
+        # Store run ID to Context mapping
+        self.ctx_store: dict[int, Context] = {}
+        self.lock_ctx_store = Lock()
+        # 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 = Lock()
+
+    def set_node_id(self, node_id: Optional[int]) -> None:
+        """Set the node ID."""
+        self.node_id = node_id
+
+    def get_node_id(self) -> int:
+        """Get the node ID."""
+        if self.node_id is None:
+            raise ValueError("Node ID not set")
+        return self.node_id
+
+    def store_message(self, message: Message) -> Optional[str]:
+        """Store a message."""
+        with self.lock_msg_store:
+            msg_id = message.metadata.message_id
+            if msg_id == "" or msg_id in self.msg_store:
+                return None
+            self.msg_store[msg_id] = MessageEntry(message=message)
+            return msg_id
+
+    def get_messages(
+        self,
+        *,
+        run_ids: Optional[Sequence[int]] = None,
+        is_reply: Optional[bool] = None,
+        limit: Optional[int] = None,
+    ) -> Sequence[Message]:
+        """Retrieve messages based on the specified filters."""
+        selected_messages: list[Message] = []
+
+        with self.lock_msg_store:
+            # Iterate through all messages in the store
+            for object_id in list(self.msg_store.keys()):
+                entry = self.msg_store[object_id]
+                message = entry.message
+
+                # Skip messages that have already been retrieved
+                if entry.is_retrieved:
+                    continue
+
+                # Skip messages whose run_id doesn't match the filter
+                if run_ids is not None:
+                    if message.metadata.run_id not in run_ids:
+                        continue
+
+                # If is_reply filter is set, filter for reply/non-reply messages
+                if is_reply is not None:
+                    is_reply_message = message.metadata.reply_to_message_id != ""
+                    # XOR logic to filter mismatched types (reply vs non-reply)
+                    if is_reply ^ is_reply_message:
+                        continue
+
+                # Add the message to the result set
+                selected_messages.append(message)
+
+                # Mark the message as retrieved
+                entry.is_retrieved = True
+
+                # Stop if the number of collected messages reaches the limit
+                if limit is not None and len(selected_messages) >= limit:
+                    break
+
+        return selected_messages
+
+    def delete_messages(
+        self,
+        *,
+        message_ids: Optional[Sequence[str]] = None,
+    ) -> None:
+        """Delete messages based on the specified filters."""
+        with self.lock_msg_store:
+            if message_ids is None:
+                # If no message IDs are provided, clear the entire store
+                self.msg_store.clear()
+                return
+
+            # Remove specified messages from the store
+            for msg_id in message_ids:
+                self.msg_store.pop(msg_id, None)
+
+    def store_run(self, run: Run) -> None:
+        """Store a run."""
+        with self.lock_run_store:
+            self.run_store[run.run_id] = run
+
+    def get_run(self, run_id: int) -> Optional[Run]:
+        """Retrieve a run by its ID."""
+        with self.lock_run_store:
+            return self.run_store.get(run_id)
+
+    def store_context(self, context: Context) -> None:
+        """Store a context."""
+        with self.lock_ctx_store:
+            self.ctx_store[context.run_id] = context
+
+    def get_context(self, run_id: int) -> Optional[Context]:
+        """Retrieve a context by its run ID."""
+        with self.lock_ctx_store:
+            return self.ctx_store.get(run_id)
+
+    def get_run_ids_with_pending_messages(self) -> Sequence[int]:
+        """Retrieve run IDs that have at least one pending message."""
+        # Collect run IDs from messages
+        with self.lock_msg_store:
+            ret = {
+                entry.message.metadata.run_id
+                for entry in self.msg_store.values()
+                if entry.message.metadata.reply_to_message_id == ""
+                and not entry.is_retrieved
+            }
+
+        # Remove run IDs that have tokens stored (indicating they are in progress)
+        with self.lock_token_store:
+            ret -= set(self.token_store.keys())
+            return list(ret)
+
+    def create_token(self, run_id: int) -> 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:
+                raise ValueError("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)

flwr/supernode/nodestate/nodestate.py

@@ -0,0 +1,227 @@
+# Copyright 2025 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.
+# ==============================================================================
+"""Abstract base class NodeState."""
+
+
+from abc import ABC, abstractmethod
+from collections.abc import Sequence
+from typing import Optional
+
+from flwr.common import Context, Message
+from flwr.common.typing import Run
+
+
+class NodeState(ABC):
+    """Abstract base class for node state."""
+
+    @abstractmethod
+    def set_node_id(self, node_id: int) -> None:
+        """Set the node ID."""
+
+    @abstractmethod
+    def get_node_id(self) -> int:
+        """Get the node ID."""
+
+    @abstractmethod
+    def store_message(self, message: Message) -> Optional[str]:
+        """Store a message.
+
+        Parameters
+        ----------
+        message : Message
+            The message to store.
+
+        Returns
+        -------
+        Optional[str]
+            The object ID of the stored message, or None if storage failed.
+        """
+
+    @abstractmethod
+    def get_messages(
+        self,
+        *,
+        run_ids: Optional[Sequence[int]] = None,
+        is_reply: Optional[bool] = None,
+        limit: Optional[int] = None,
+    ) -> Sequence[Message]:
+        """Retrieve messages 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
+        ----------
+        run_ids : Optional[Sequence[int]] (default: None)
+            Sequence of run IDs to filter by. If a sequence is provided,
+            it is treated as an OR condition.
+        is_reply : Optional[bool] (default: None)
+            If True, filter for reply messages; if False, filter for non-reply
+            (instruction) messages.
+        limit : Optional[int] (default: None)
+            Maximum number of messages to return. If None, no limit is applied.
+
+        Returns
+        -------
+        Sequence[Message]
+            A sequence of messages matching the specified filters.
+
+        Notes
+        -----
+        **IMPORTANT:** Retrieved messages will **NOT** be returned again by subsequent
+        calls to this method, even if the filters match them.
+        """
+
+    @abstractmethod
+    def delete_messages(
+        self,
+        *,
+        message_ids: Optional[Sequence[str]] = None,
+    ) -> None:
+        """Delete messages 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
+        ----------
+        message_ids : Optional[Sequence[str]] (default: None)
+            Sequence of message (object) IDs to filter by. If a sequence is provided,
+            it is treated as an OR condition.
+
+        Notes
+        -----
+        **IMPORTANT:** **All messages** will be deleted if no filters are provided.
+        """
+
+    @abstractmethod
+    def store_run(self, run: Run) -> None:
+        """Store a run.
+
+        Parameters
+        ----------
+        run : Run
+            The `Run` instance to store.
+        """
+
+    @abstractmethod
+    def get_run(self, run_id: int) -> Optional[Run]:
+        """Retrieve a run by its ID.
+
+        Parameters
+        ----------
+        run_id : int
+            The ID of the run to retrieve.
+
+        Returns
+        -------
+        Optional[Run]
+            The `Run` instance if found, otherwise None.
+        """
+
+    @abstractmethod
+    def store_context(self, context: Context) -> None:
+        """Store a context.
+
+        Parameters
+        ----------
+        context : Context
+            The context to store.
+        """
+
+    @abstractmethod
+    def get_context(self, run_id: int) -> Optional[Context]:
+        """Retrieve a context by its run ID.
+
+        Parameters
+        ----------
+        run_id : int
+            The ID of the run with which the context is associated.
+
+        Returns
+        -------
+        Optional[Context]
+            The `Context` instance if found, otherwise None.
+        """
+
+    @abstractmethod
+    def get_run_ids_with_pending_messages(self) -> Sequence[int]:
+        """Retrieve run IDs that have at least one pending message.
+
+        Run IDs that are currently in progress (i.e., those associated with tokens)
+        will not be returned, even if they have pending messages.
+
+        Returns
+        -------
+        Sequence[int]
+            Sequence of run IDs with pending messages.
+        """
+
+    @abstractmethod
+    def create_token(self, run_id: int) -> str:
+        """Create a token for the given run ID.
+
+        Parameters
+        ----------
+        run_id : int
+            The ID of the run for which to create a token.
+
+        Returns
+        -------
+        str
+            A unique token associated with the run ID.
+        """
+
+    @abstractmethod
+    def verify_token(self, run_id: int, token: str) -> bool:
+        """Verify a token for the given run ID.
+
+        Parameters
+        ----------
+        run_id : int
+            The ID of the run for which to verify the token.
+        token : str
+            The token to verify.
+
+        Returns
+        -------
+        bool
+            True if the token is valid for the run ID, False otherwise.
+        """
+
+    @abstractmethod
+    def delete_token(self, run_id: int) -> None:
+        """Delete the token for the given run ID.
+
+        Parameters
+        ----------
+        run_id : int
+            The ID of the run for which to delete the token.
+        """
+
+    @abstractmethod
+    def get_run_id_by_token(self, token: str) -> Optional[int]:
+        """Get the run ID associated with a given token.
+
+        Parameters
+        ----------
+        token : str
+            The token to look up.
+
+        Returns
+        -------
+        Optional[int]
+            The run ID if the token is valid, otherwise None.
+        """

flwr/supernode/runtime/__init__.py

@@ -0,0 +1,15 @@
+# Copyright 2025 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 SuperNode components."""