flwr 1.24.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

@@ -127,7 +127,7 @@ def send_node_heartbeat(
     return SendNodeHeartbeatResponse(success=res)
 
 
-def pull_messages(
+def pull_messages(  # pylint: disable=too-many-locals
     request: PullMessagesRequest,
     state: LinkState,
     store: ObjectStore,
@@ -143,6 +143,8 @@ def pull_messages(
     # Convert to Messages
     msg_proto = []
     trees = []
+    run_id_to_record: int | None = None
+
     for msg in message_list:
         try:
             # Retrieve Message object tree from ObjectStore
@@ -152,12 +154,30 @@ def pull_messages(
             # Add Message and its object tree to the response
             msg_proto.append(message_to_proto(msg))
             trees.append(obj_tree)
+
+            # Track run_id for traffic recording
+            run_id_to_record = msg.metadata.run_id
+
         except NoObjectInStoreError as e:
             log(ERROR, e.message)
             # Delete message ins from state
             state.delete_messages(message_ins_ids={msg_object_id})
 
-    return PullMessagesResponse(messages_list=msg_proto, message_object_trees=trees)
+    response = PullMessagesResponse(messages_list=msg_proto, message_object_trees=trees)
+
+    # Record incoming traffic size
+    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 = response.ByteSize()
+        state.store_traffic(
+            run_id_to_record, bytes_sent=bytes_sent, bytes_recv=bytes_recv
+        )
+
+    return response
 
 
 def push_messages(
@@ -170,6 +190,9 @@ def push_messages(
     msg = message_from_proto(message_proto=request.messages_list[0])
     run_id = msg.metadata.run_id
 
+    # Record incoming traffic size
+    bytes_recv = request.ByteSize()
+
     # Abort if the run is not running
     abort_msg = check_abort(
         run_id,
@@ -193,6 +216,16 @@ def push_messages(
         results={str(message_id): 0},
         objects_to_push=objects_to_push,
     )
+
+    # Record outgoing traffic size
+    bytes_sent = response.ByteSize()
+
+    # Record traffic only if message was successfully processed
+    # Only one message is processed per request
+    state.store_traffic(run_id, bytes_sent=bytes_sent, bytes_recv=bytes_recv)
+    if request.clientapp_runtime_list:
+        state.add_clientapp_runtime(run_id, request.clientapp_runtime_list[0])
+
     return response
 
 
@@ -257,6 +290,10 @@ def push_object(
     try:
         store.put(request.object_id, request.object_content)
         stored = True
+        # Record bytes traffic pushed from SuperNode
+        state.store_traffic(
+            request.run_id, bytes_sent=0, bytes_recv=len(request.object_content)
+        )
     except (NoObjectInStoreError, ValueError) as e:
         log(ERROR, str(e))
     except UnexpectedObjectContentError as e:
@@ -283,6 +320,9 @@ def pull_object(
     content = store.get(request.object_id)
     if content is not None:
         object_available = content != b""
+        # Record bytes traffic pulled by SuperNode
+        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."""
@@ -576,6 +568,9 @@ class InMemoryLinkState(LinkState, InMemoryCoreState):  # pylint: disable=R0902,
                         ),
                         flwr_aid=flwr_aid if flwr_aid else "",
                         federation=federation,
+                        bytes_sent=0,
+                        bytes_recv=0,
+                        clientapp_runtime=0.0,
                     ),
                 )
                 self.run_ids[run_id] = run_record
@@ -771,3 +766,34 @@ class InMemoryLinkState(LinkState, InMemoryCoreState):  # pylint: disable=R0902,
             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 store_traffic(self, run_id: int, *, bytes_sent: int, bytes_recv: int) -> None:
+        """Store traffic data for the specified `run_id`."""
+        # Validate non-negative values
+        if bytes_sent < 0 or bytes_recv < 0:
+            raise ValueError(
+                f"Negative traffic values for run {run_id}: "
+                f"bytes_sent={bytes_sent}, bytes_recv={bytes_recv}"
+            )
+
+        if bytes_sent == 0 and bytes_recv == 0:
+            raise ValueError(
+                f"Both bytes_sent and bytes_recv cannot be zero for run {run_id}"
+            )
+
+        with self.lock:
+            if run_id not in self.run_ids:
+                raise ValueError(f"Run {run_id} not found")
+            run_record = self.run_ids[run_id]
+
+        with run_record.lock:
+            run = run_record.run
+            run.bytes_sent += bytes_sent
+            run.bytes_recv += bytes_recv
+
+    def add_clientapp_runtime(self, run_id: int, runtime: float) -> None:
+        """Add ClientApp runtime to the cumulative total for the specified `run_id`."""
+        with self.lock:
+            if run_id not in self.run_ids:
+                raise ValueError(f"Run {run_id} not found")
+            self.run_ids[run_id].run.clientapp_runtime += runtime

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,
@@ -480,3 +461,35 @@ class LinkState(CoreState):  # pylint: disable=R0904
             - The timestamp of the latest log entry in the returned logs.
               Returns `0` if no logs are returned.
         """
+
+    @abc.abstractmethod
+    def store_traffic(self, run_id: int, *, bytes_sent: int, bytes_recv: int) -> None:
+        """Store traffic data for the specified `run_id`.
+
+        Parameters
+        ----------
+        run_id : int
+            The identifier of the run for which to store traffic data.
+        bytes_sent : int
+            The number of bytes pulled by SuperNodes from the SuperLink to add to the
+            run's total.
+        bytes_recv : int
+            The number of bytes received by SuperLink from SuperNodes to add to the
+            run's total.
+        """
+
+    @abc.abstractmethod
+    def add_clientapp_runtime(self, run_id: int, runtime: float) -> None:
+        """Add ClientApp runtime to the cumulative total for the specified `run_id`.
+
+        This method accumulates the runtime by adding the provided value to the
+        existing total runtime for the run. Multiple ClientApps can contribute
+        to the same run's total runtime.
+
+        Parameters
+        ----------
+        run_id : int
+            The identifier of the run for which to store each ClientApp's runtime.
+        runtime : float
+            The runtime in seconds to add to the `run_id`'s cumulative total.
+        """

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