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

flwr/server/superlink/fleet/grpc_bidi/flower_service_servicer.py

@@ -20,8 +20,7 @@ Relevant knowledge for reading this modules code:
 
 
 import uuid
-from collections.abc import Iterator
-from typing import Callable
+from collections.abc import Callable, Iterator
 
 import grpc
 from iterators import TimeoutIterator

flwr/server/superlink/fleet/grpc_bidi/grpc_bridge.py

@@ -19,7 +19,6 @@ from collections.abc import Iterator
 from dataclasses import dataclass
 from enum import Enum
 from threading import Condition
-from typing import Optional
 
 from flwr.proto.transport_pb2 import (  # pylint: disable=E0611
     ClientMessage,
@@ -32,7 +31,7 @@ class InsWrapper:
     """Instruction wrapper class for a single server message."""
 
     server_message: ServerMessage
-    timeout: Optional[float]
+    timeout: float | None
 
 
 @dataclass
@@ -70,8 +69,8 @@ class GrpcBridge:
         # pylint: disable=unsubscriptable-object
         self._cv = Condition()  # cv stands for condition variable
         self._status = Status.AWAITING_INS_WRAPPER
-        self._ins_wrapper: Optional[InsWrapper] = None
-        self._res_wrapper: Optional[ResWrapper] = None
+        self._ins_wrapper: InsWrapper | None = None
+        self._res_wrapper: ResWrapper | None = None
 
     def _is_closed(self) -> bool:
         """Return True if closed and False otherwise."""

flwr/server/superlink/fleet/grpc_bidi/grpc_client_proxy.py

@@ -15,8 +15,6 @@
 """Flower ClientProxy implementation using gRPC bidirectional streaming."""
 
 
-from typing import Optional
-
 from flwr import common
 from flwr.common import serde
 from flwr.proto.transport_pb2 import (  # pylint: disable=E0611
@@ -45,8 +43,8 @@ class GrpcClientProxy(ClientProxy):
     def get_properties(
         self,
         ins: common.GetPropertiesIns,
-        timeout: Optional[float],
-        group_id: Optional[int],
+        timeout: float | None,
+        group_id: int | None,
     ) -> common.GetPropertiesRes:
         """Request client's set of internal properties."""
         get_properties_msg = serde.get_properties_ins_to_proto(ins)
@@ -65,8 +63,8 @@ class GrpcClientProxy(ClientProxy):
     def get_parameters(
         self,
         ins: common.GetParametersIns,
-        timeout: Optional[float],
-        group_id: Optional[int],
+        timeout: float | None,
+        group_id: int | None,
     ) -> common.GetParametersRes:
         """Return the current local model parameters."""
         get_parameters_msg = serde.get_parameters_ins_to_proto(ins)
@@ -85,8 +83,8 @@ class GrpcClientProxy(ClientProxy):
     def fit(
         self,
         ins: common.FitIns,
-        timeout: Optional[float],
-        group_id: Optional[int],
+        timeout: float | None,
+        group_id: int | None,
     ) -> common.FitRes:
         """Refine the provided parameters using the locally held dataset."""
         fit_ins_msg = serde.fit_ins_to_proto(ins)
@@ -104,8 +102,8 @@ class GrpcClientProxy(ClientProxy):
     def evaluate(
         self,
         ins: common.EvaluateIns,
-        timeout: Optional[float],
-        group_id: Optional[int],
+        timeout: float | None,
+        group_id: int | None,
     ) -> common.EvaluateRes:
         """Evaluate the provided parameters using the locally held dataset."""
         evaluate_msg = serde.evaluate_ins_to_proto(ins)
@@ -122,8 +120,8 @@ class GrpcClientProxy(ClientProxy):
     def reconnect(
         self,
         ins: common.ReconnectIns,
-        timeout: Optional[float],
-        group_id: Optional[int],
+        timeout: float | None,
+        group_id: int | None,
     ) -> common.DisconnectRes:
         """Disconnect and (optionally) reconnect later."""
         reconnect_ins_msg = serde.reconnect_ins_to_proto(ins)

flwr/server/superlink/fleet/grpc_bidi/grpc_server.py

@@ -15,8 +15,6 @@
 """Implements utility function to create a gRPC server."""
 
 
-from typing import Optional
-
 import grpc
 
 from flwr.common import GRPC_MAX_MESSAGE_LENGTH
@@ -36,7 +34,7 @@ def start_grpc_server(  # pylint: disable=too-many-arguments,R0917
     max_concurrent_workers: int = 1000,
     max_message_length: int = GRPC_MAX_MESSAGE_LENGTH,
     keepalive_time_ms: int = 210000,
-    certificates: Optional[tuple[bytes, bytes, bytes]] = None,
+    certificates: tuple[bytes, bytes, bytes] | None = None,
 ) -> grpc.Server:
     """Create and start a gRPC server running FlowerServiceServicer.
 

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

@@ -249,8 +249,8 @@ class FleetServicer(fleet_pb2_grpc.FleetServicer):
                 state=self.state_factory.state(),
                 store=self.objectstore_factory.store(),
             )
-        except InvalidRunStatusException as e:
-            abort_grpc_context(e.message, context)
+        except (InvalidRunStatusException, ValueError) as e:
+            abort_grpc_context(str(e), context)
 
         return res
 
@@ -266,8 +266,8 @@ class FleetServicer(fleet_pb2_grpc.FleetServicer):
                 state=self.state_factory.state(),
                 store=self.objectstore_factory.store(),
             )
-        except InvalidRunStatusException as e:
-            abort_grpc_context(e.message, context)
+        except (InvalidRunStatusException, ValueError) as e:
+            abort_grpc_context(str(e), context)
 
         return res
 

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

@@ -16,7 +16,8 @@
 
 
 import datetime
-from typing import Any, Callable, cast
+from collections.abc import Callable
+from typing import Any, cast
 
 import grpc
 from google.protobuf.message import Message as GrpcMessage
@@ -117,7 +118,7 @@ class NodeAuthServerInterceptor(grpc.ServerInterceptor):  # type: ignore
             # than the `intercept_service` function.
 
             # Retrieve the public key
-            if isinstance(request, (RegisterNodeFleetRequest, ActivateNodeRequest)):
+            if isinstance(request, (RegisterNodeFleetRequest | ActivateNodeRequest)):
                 actual_public_key = request.public_key
             else:
                 if hasattr(request, "node"):

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

@@ -15,12 +15,12 @@
 """Fleet API message handlers."""
 
 from logging import ERROR
-from typing import Optional
 
 from flwr.common import Message, log
 from flwr.common.constant import (
     HEARTBEAT_MAX_INTERVAL,
     HEARTBEAT_MIN_INTERVAL,
+    NOOP_ACCOUNT_NAME,
     NOOP_FLWR_AID,
     Status,
 )
@@ -29,9 +29,9 @@ from flwr.common.serde import (
     fab_to_proto,
     message_from_proto,
     message_to_proto,
-    user_config_to_proto,
+    run_to_proto,
 )
-from flwr.common.typing import Fab, InvalidRunStatusException
+from flwr.common.typing import Fab, InvalidRunStatusException, Run
 from flwr.proto.fab_pb2 import GetFabRequest, GetFabResponse  # pylint: disable=E0611
 from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
     ActivateNodeRequest,
@@ -60,16 +60,11 @@ from flwr.proto.message_pb2 import (  # pylint: disable=E0611
     PushObjectRequest,
     PushObjectResponse,
 )
-from flwr.proto.run_pb2 import (  # pylint: disable=E0611
-    GetRunRequest,
-    GetRunResponse,
-    Run,
-)
+from flwr.proto.run_pb2 import GetRunRequest, GetRunResponse  # pylint: disable=E0611
 from flwr.server.superlink.linkstate import LinkState
 from flwr.server.superlink.utils import check_abort
 from flwr.supercore.ffs import Ffs
 from flwr.supercore.object_store import NoObjectInStoreError, ObjectStore
-from flwr.supercore.object_store.utils import store_mapping_and_register_objects
 
 
 class InvalidHeartbeatIntervalError(Exception):
@@ -81,7 +76,7 @@ def register_node(
     state: LinkState,
 ) -> RegisterNodeFleetResponse:
     """Register a node (Fleet API only)."""
-    node_id = state.create_node(NOOP_FLWR_AID, request.public_key, 0)
+    node_id = state.create_node(NOOP_FLWR_AID, NOOP_ACCOUNT_NAME, request.public_key, 0)
     return RegisterNodeFleetResponse(node_id=node_id)
 
 
@@ -132,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,
@@ -148,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
@@ -157,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 = len(request.SerializeToString())
+
+    # 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())
+        state.store_traffic(
+            run_id_to_record, bytes_sent=bytes_sent, bytes_recv=bytes_recv
+        )
+
+    return response
 
 
 def push_messages(
@@ -173,10 +188,14 @@ def push_messages(
     """Push Messages handler."""
     # Convert Message from proto
     msg = message_from_proto(message_proto=request.messages_list[0])
+    run_id = msg.metadata.run_id
+
+    # Record incoming traffic size
+    bytes_recv = len(request.SerializeToString())
 
     # Abort if the run is not running
     abort_msg = check_abort(
-        msg.metadata.run_id,
+        run_id,
         [Status.PENDING, Status.STARTING, Status.FINISHED],
         state,
         store,
@@ -184,11 +203,12 @@ def push_messages(
     if abort_msg:
         raise InvalidRunStatusException(abort_msg)
 
-    # Store Message in State
-    message_id: Optional[str] = state.store_message_res(message=msg)
-
     # Store Message object to descendants mapping and preregister objects
-    objects_to_push = store_mapping_and_register_objects(store, request=request)
+    objects_to_push: set[str] = set()
+    for object_tree in request.message_object_trees:
+        objects_to_push |= set(store.preregister(run_id, object_tree))
+    # Store Message in State
+    message_id: str | None = state.store_message_res(message=msg)
 
     # Build response
     response = PushMessagesResponse(
@@ -196,6 +216,16 @@ def push_messages(
         results={str(message_id): 0},
         objects_to_push=objects_to_push,
     )
+
+    # Record outgoing traffic size
+    bytes_sent = len(response.SerializeToString())
+
+    # 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
 
 
@@ -203,10 +233,8 @@ def get_run(
     request: GetRunRequest, state: LinkState, store: ObjectStore
 ) -> GetRunResponse:
     """Get run information."""
-    run = state.get_run(request.run_id)
-
-    if run is None:
-        return GetRunResponse()
+    # Validate that the requesting SuperNode is part of the federation
+    run = _validate_node_in_federation(state, request.node.node_id, request.run_id)
 
     # Abort if the run is not running
     abort_msg = check_abort(
@@ -218,21 +246,16 @@ def get_run(
     if abort_msg:
         raise InvalidRunStatusException(abort_msg)
 
-    return GetRunResponse(
-        run=Run(
-            run_id=run.run_id,
-            fab_id=run.fab_id,
-            fab_version=run.fab_version,
-            override_config=user_config_to_proto(run.override_config),
-            fab_hash=run.fab_hash,
-        )
-    )
+    return GetRunResponse(run=run_to_proto(run))
 
 
 def get_fab(
     request: GetFabRequest, ffs: Ffs, state: LinkState, store: ObjectStore
 ) -> GetFabResponse:
     """Get FAB."""
+    # Validate that the requesting SuperNode is part of the federation
+    _validate_node_in_federation(state, request.node.node_id, request.run_id)
+
     # Abort if the run is not running
     abort_msg = check_abort(
         request.run_id,
@@ -267,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:
@@ -293,6 +320,8 @@ def pull_object(
     content = store.get(request.object_id)
     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)
         return PullObjectResponse(
             object_found=True,
             object_available=object_available,
@@ -329,3 +358,19 @@ def _validate_heartbeat_interval(interval: float) -> None:
             f"Heartbeat interval {interval} is out of bounds "
             f"[{HEARTBEAT_MIN_INTERVAL}, {HEARTBEAT_MAX_INTERVAL}]."
         )
+
+
+def _validate_node_in_federation(
+    state: LinkState,
+    node_id: int,
+    run_id: int,
+) -> Run:
+    """Raise if the requesting SuperNode is not part of the federation the run belongs
+    to."""
+    run = state.get_run(run_id)
+    if not run:
+        raise ValueError(f"Run ID not found: {run_id}")
+
+    if not state.federation_manager.has_node(node_id, run.federation):
+        raise ValueError(f"SuperNode is not part of the federation '{run.federation}'.")
+    return run

flwr/server/superlink/fleet/rest_rere/rest_api.py

@@ -17,8 +17,8 @@
 
 from __future__ import annotations
 
-from collections.abc import Awaitable
-from typing import Callable, TypeVar, cast
+from collections.abc import Awaitable, Callable
+from typing import TypeVar, cast
 
 from google.protobuf.message import Message as GrpcMessage
 

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

@@ -16,7 +16,7 @@
 
 
 from abc import ABC, abstractmethod
-from typing import Callable
+from collections.abc import Callable
 
 from flwr.clientapp.client_app import ClientApp
 from flwr.common.context import Context

flwr/server/superlink/fleet/vce/backend/raybackend.py

@@ -16,8 +16,8 @@
 
 
 import sys
+from collections.abc import Callable
 from logging import DEBUG, ERROR
-from typing import Callable, Optional, Union
 
 import ray
 
@@ -32,8 +32,8 @@ from flwr.simulation.ray_transport.utils import enable_tf_gpu_growth
 
 from .backend import Backend, BackendConfig
 
-ClientResourcesDict = dict[str, Union[int, float]]
-ActorArgsDict = dict[str, Union[int, float, Callable[[], None]]]
+ClientResourcesDict = dict[str, int | float]
+ActorArgsDict = dict[str, int | float | Callable[[], None]]
 
 
 class RayBackend(Backend):
@@ -57,9 +57,9 @@ class RayBackend(Backend):
 
         # Valide actor resources
         self.actor_kwargs = self._validate_actor_arguments(config=backend_config)
-        self.pool: Optional[BasicActorPool] = None
+        self.pool: BasicActorPool | None = None
 
-        self.app_fn: Optional[Callable[[], ClientApp]] = None
+        self.app_fn: Callable[[], ClientApp] | None = None
 
     def _validate_client_resources(self, config: BackendConfig) -> ClientResourcesDict:
         client_resources_config = config.get(self.client_resources_key)

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

@@ -20,11 +20,11 @@ import secrets
 import threading
 import time
 import traceback
+from collections.abc import Callable
 from concurrent.futures import ThreadPoolExecutor
 from logging import DEBUG, ERROR, INFO, WARN
 from pathlib import Path
 from queue import Empty, Queue
-from typing import Callable, Optional
 from uuid import uuid4
 
 from flwr.app.error import Error
@@ -34,6 +34,7 @@ from flwr.clientapp.utils import get_load_client_app_fn
 from flwr.common import Message
 from flwr.common.constant import (
     HEARTBEAT_INTERVAL_INF,
+    NOOP_ACCOUNT_NAME,
     NOOP_FLWR_AID,
     NUM_PARTITIONS_KEY,
     PARTITION_ID_KEY,
@@ -43,6 +44,8 @@ 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 flwr.supercore.object_store import ObjectStoreFactory
+from flwr.superlink.federation import NoOpFederationManager
 
 from .backend import Backend, error_messages_backends, supported_backends
 
@@ -61,6 +64,7 @@ def _register_nodes(
             # use NOOP_FLWR_AID as owner_aid and
             # use random bytes as public key
             NOOP_FLWR_AID,
+            NOOP_ACCOUNT_NAME,
             secrets.token_bytes(32),
             heartbeat_interval=HEARTBEAT_INTERVAL_INF,
         )
@@ -75,7 +79,7 @@ def _register_nodes(
 def _register_node_info_stores(
     nodes_mapping: NodeToPartitionMapping,
     run: Run,
-    app_dir: Optional[str] = None,
+    app_dir: str | None = None,
 ) -> dict[int, DeprecatedRunInfoStore]:
     """Create DeprecatedRunInfoStore objects and register the context for the run."""
     node_info_store: dict[int, DeprecatedRunInfoStore] = {}
@@ -270,12 +274,12 @@ def start_vce(
     is_app: bool,
     f_stop: threading.Event,
     run: Run,
-    flwr_dir: Optional[str] = None,
-    client_app: Optional[ClientApp] = None,
-    client_app_attr: Optional[str] = None,
-    num_supernodes: Optional[int] = None,
-    state_factory: Optional[LinkStateFactory] = None,
-    existing_nodes_mapping: Optional[NodeToPartitionMapping] = None,
+    flwr_dir: str | None = None,
+    client_app: ClientApp | None = None,
+    client_app_attr: str | None = None,
+    num_supernodes: int | None = None,
+    state_factory: LinkStateFactory | None = None,
+    existing_nodes_mapping: NodeToPartitionMapping | None = None,
 ) -> None:
     """Start Fleet API with the Simulation Engine."""
     nodes_mapping = {}
@@ -313,7 +317,9 @@ 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_DB_NAME)
+        state_factory = LinkStateFactory(
+            FLWR_IN_MEMORY_DB_NAME, NoOpFederationManager(), ObjectStoreFactory()
+        )
         log(INFO, "Created new %s.", state_factory.__class__.__name__)
 
     if num_supernodes: