flwr 1.25.0__py3-none-any.whl → 1.26.0__py3-none-any.whl

flwr/server/grid/grpc_grid.py

@@ -28,7 +28,6 @@ from flwr.common.constant import (
     SERVERAPPIO_API_DEFAULT_CLIENT_ADDRESS,
     SUPERLINK_NODE_ID,
     ErrorCode,
-    MessageType,
 )
 from flwr.common.grpc import create_channel, on_channel_state_change
 from flwr.common.inflatable import (
@@ -50,7 +49,7 @@ from flwr.common.inflatable_utils import (
 )
 from flwr.common.logger import log, warn_deprecated_feature
 from flwr.common.message import make_message, remove_content_from_message
-from flwr.common.retry_invoker import _make_simple_grpc_retry_invoker, _wrap_stub
+from flwr.common.retry_invoker import make_simple_grpc_retry_invoker, wrap_stub
 from flwr.common.serde import message_to_proto, run_from_proto
 from flwr.common.typing import Run
 from flwr.proto.appio_pb2 import (  # pylint: disable=E0611
@@ -69,6 +68,7 @@ from flwr.proto.serverappio_pb2 import (  # pylint: disable=E0611
     GetNodesResponse,
 )
 from flwr.proto.serverappio_pb2_grpc import ServerAppIoStub  # pylint: disable=E0611
+from flwr.supercore.constant import SYSTEM_MESSAGE_TYPE
 
 from .grid import Grid
 
@@ -127,7 +127,7 @@ class GrpcGrid(Grid):
         self._grpc_stub: ServerAppIoStub | None = None
         self._channel: grpc.Channel | None = None
         self.node = Node(node_id=SUPERLINK_NODE_ID)
-        self._retry_invoker = _make_simple_grpc_retry_invoker()
+        self._retry_invoker = make_simple_grpc_retry_invoker()
         super().__init__()
 
     @property
@@ -150,7 +150,7 @@ class GrpcGrid(Grid):
         )
         self._channel.subscribe(on_channel_state_change)
         self._grpc_stub = ServerAppIoStub(self._channel)
-        _wrap_stub(self._grpc_stub, self._retry_invoker)
+        wrap_stub(self._grpc_stub, self._retry_invoker)
         log(DEBUG, "[flwr-serverapp] Connected to %s", self._addr)
 
     def _disconnect(self) -> None:
@@ -341,7 +341,7 @@ class GrpcGrid(Grid):
                                 message_id="",
                                 src_node_id=self.node.node_id,
                                 dst_node_id=self.node.node_id,
-                                message_type=MessageType.SYSTEM,
+                                message_type=SYSTEM_MESSAGE_TYPE,
                                 group_id="",
                                 ttl=0,
                                 reply_to_message_id=msg_proto.metadata.reply_to_message_id,

flwr/server/serverapp/app.py

@@ -144,6 +144,10 @@ def run_serverapp(  # pylint: disable=R0913, R0914, R0915, R0917, W0212
     exit_code = ExitCode.SUCCESS
 
     def on_exit() -> None:
+        # Set Grpc max retries to 1 to avoid blocking on exit
+        if grid:
+            grid._retry_invoker.max_tries = 1
+
         # Stop heartbeat sender
         if heartbeat_sender and heartbeat_sender.is_running:
             heartbeat_sender.stop()
@@ -154,10 +158,13 @@ def run_serverapp(  # pylint: disable=R0913, R0914, R0915, R0917, W0212
 
         # Update run status
         if run and run_status and grid:
-            run_status_proto = run_status_to_proto(run_status)
-            grid._stub.UpdateRunStatus(
-                UpdateRunStatusRequest(run_id=run.run_id, run_status=run_status_proto)
-            )
+            try:
+                req = UpdateRunStatusRequest(
+                    run_id=run.run_id, run_status=run_status_to_proto(run_status)
+                )
+                grid._stub.UpdateRunStatus(req)
+            except grpc.RpcError:
+                pass
 
         # Close the Grpc connection
         if grid:

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

@@ -30,7 +30,6 @@ from flwr.common.constant import (
     GRPC_ADAPTER_METADATA_MESSAGE_QUALNAME_KEY,
 )
 from flwr.common.logger import log
-from flwr.common.version import package_name, package_version
 from flwr.proto import grpcadapter_pb2_grpc  # pylint: disable=E0611
 from flwr.proto.fab_pb2 import GetFabRequest  # pylint: disable=E0611
 from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
@@ -49,6 +48,7 @@ from flwr.proto.message_pb2 import (  # pylint: disable=E0611
     PushObjectRequest,
 )
 from flwr.proto.run_pb2 import GetRunRequest  # pylint: disable=E0611
+from flwr.supercore.version import package_name, package_version
 
 from ..grpc_rere.fleet_servicer import FleetServicer
 

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

@@ -108,7 +108,7 @@ class NodeAuthServerInterceptor(grpc.ServerInterceptor):  # type: ignore
     def _wrap_method_handler(
         self,
         method_handler: grpc.RpcMethodHandler,
-        expected_public_key: bytes,
+        received_public_key: bytes,
     ) -> grpc.RpcMethodHandler:
         def _generic_method_handler(
             request: GrpcMessage,
@@ -117,21 +117,22 @@ class NodeAuthServerInterceptor(grpc.ServerInterceptor):  # type: ignore
             # Note: This function runs in a different thread
             # than the `intercept_service` function.
 
-            # Retrieve the public key
-            if isinstance(request, (RegisterNodeFleetRequest | ActivateNodeRequest)):
-                actual_public_key = request.public_key
-            else:
+            # Skip registration and activation requests
+            if not isinstance(request, (RegisterNodeFleetRequest, ActivateNodeRequest)):
+                # Retrieve the node ID from the request
                 if hasattr(request, "node"):
-                    node_id = request.node.node_id
+                    received_node_id = request.node.node_id
                 else:
-                    node_id = request.node_id  # type: ignore[attr-defined]
-                actual_public_key = self.state_factory.state().get_node_public_key(
-                    node_id
+                    received_node_id = request.node_id  # type: ignore[attr-defined]
+
+                # Get the actual node ID based on the received public key
+                node_id = self.state_factory.state().get_node_id_by_public_key(
+                    received_public_key
                 )
 
-            # Verify the public key
-            if actual_public_key != expected_public_key:
-                context.abort(grpc.StatusCode.UNAUTHENTICATED, "Invalid node ID")
+                # Verify that the received node ID matches the actual node ID
+                if received_node_id != node_id:
+                    context.abort(grpc.StatusCode.UNAUTHENTICATED, "Invalid node ID")
 
             response: GrpcMessage = method_handler.unary_unary(request, context)
             return response

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

@@ -166,13 +166,13 @@ def pull_messages(  # pylint: disable=too-many-locals
     response = PullMessagesResponse(messages_list=msg_proto, message_object_trees=trees)
 
     # Record incoming traffic size
-    bytes_recv = len(request.SerializeToString())
+    bytes_recv = request.ByteSize()
 
     # Record traffic only if message was successfully processed
     # All messages in this request are assumed to belong to the same run
     if run_id_to_record is not None:
         # Record outgoing traffic size
-        bytes_sent = len(response.SerializeToString())
+        bytes_sent = response.ByteSize()
         state.store_traffic(
             run_id_to_record, bytes_sent=bytes_sent, bytes_recv=bytes_recv
         )
@@ -191,7 +191,7 @@ def push_messages(
     run_id = msg.metadata.run_id
 
     # Record incoming traffic size
-    bytes_recv = len(request.SerializeToString())
+    bytes_recv = request.ByteSize()
 
     # Abort if the run is not running
     abort_msg = check_abort(
@@ -218,7 +218,7 @@ def push_messages(
     )
 
     # Record outgoing traffic size
-    bytes_sent = len(response.SerializeToString())
+    bytes_sent = response.ByteSize()
 
     # Record traffic only if message was successfully processed
     # Only one message is processed per request
@@ -321,7 +321,8 @@ def pull_object(
     if content is not None:
         object_available = content != b""
         # Record bytes traffic pulled by SuperNode
-        state.store_traffic(request.run_id, bytes_sent=len(content), bytes_recv=0)
+        if object_available:
+            state.store_traffic(request.run_id, bytes_sent=len(content), bytes_recv=0)
         return PullObjectResponse(
             object_found=True,
             object_available=object_available,

flwr/server/superlink/linkstate/__init__.py

@@ -18,11 +18,11 @@
 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
+from .sql_linkstate import SqlLinkState as SqlLinkState
 
 __all__ = [
     "InMemoryLinkState",
     "LinkState",
     "LinkStateFactory",
-    "SqliteLinkState",
+    "SqlLinkState",
 ]

flwr/server/superlink/linkstate/in_memory_linkstate.py

@@ -23,6 +23,7 @@ from dataclasses import dataclass, field
 from datetime import datetime, timezone
 from logging import ERROR, WARNING
 
+from flwr.app.user_config import UserConfig
 from flwr.common import Context, Message, log, now
 from flwr.common.constant import (
     HEARTBEAT_PATIENCE,
@@ -35,7 +36,7 @@ from flwr.common.constant import (
     SubStatus,
 )
 from flwr.common.record import ConfigRecord
-from flwr.common.typing import Run, RunStatus, UserConfig
+from flwr.common.typing import Run, RunStatus
 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
@@ -518,15 +519,6 @@ class InMemoryLinkState(LinkState, InMemoryCoreState):  # pylint: disable=R0902,
                             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 := 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) -> int | None:
         """Get `node_id` for the specified `public_key` if it exists and is not
         deleted."""

flwr/server/superlink/linkstate/linkstate.py

@@ -18,9 +18,10 @@
 import abc
 from collections.abc import Sequence
 
+from flwr.app.user_config import UserConfig
 from flwr.common import Context, Message
 from flwr.common.record import ConfigRecord
-from flwr.common.typing import Run, RunStatus, UserConfig
+from flwr.common.typing import Run, RunStatus
 from flwr.proto.node_pb2 import NodeInfo  # pylint: disable=E0611
 from flwr.supercore.corestate import CoreState
 from flwr.superlink.federation import FederationManager
@@ -245,26 +246,6 @@ class LinkState(CoreState):  # pylint: disable=R0904
             the specified filters.
         """
 
-    @abc.abstractmethod
-    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
         self,

flwr/server/superlink/linkstate/linkstate_factory.py

@@ -24,7 +24,7 @@ from flwr.superlink.federation import FederationManager
 
 from .in_memory_linkstate import InMemoryLinkState
 from .linkstate import LinkState
-from .sqlite_linkstate import SqliteLinkState
+from .sql_linkstate import SqlLinkState
 
 
 class LinkStateFactory:
@@ -56,20 +56,28 @@ class LinkStateFactory:
 
     def state(self) -> LinkState:
         """Return a State instance and create it, if necessary."""
+        # Return cached state if it exists
+        if self.state_instance is not None:
+            if self.database == FLWR_IN_MEMORY_DB_NAME:
+                log(DEBUG, "Using InMemoryState")
+            else:
+                log(DEBUG, "Using SqlLinkState")
+            return self.state_instance
+
         # Get the ObjectStore instance
         object_store = self.objectstore_factory.store()
 
         # InMemoryState
         if self.database == FLWR_IN_MEMORY_DB_NAME:
-            if self.state_instance is None:
-                self.state_instance = InMemoryLinkState(
-                    self.federation_manager, object_store
-                )
+            self.state_instance = InMemoryLinkState(
+                self.federation_manager, object_store
+            )
             log(DEBUG, "Using InMemoryState")
             return self.state_instance
 
-        # SqliteState
-        state = SqliteLinkState(self.database, self.federation_manager, object_store)
+        # SqlLinkState
+        state = SqlLinkState(self.database, self.federation_manager, object_store)
         state.initialize()
-        log(DEBUG, "Using SqliteState")
+        self.state_instance = state
+        log(DEBUG, "Using SqlLinkState")
         return state