flwr-nightly 1.9.0.dev20240502__tar.gz → 1.9.0.dev20240506__tar.gz

{flwr_nightly-1.9.0.dev20240502 → flwr_nightly-1.9.0.dev20240506}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flwr-nightly
-Version: 1.9.0.dev20240502
+Version: 1.9.0.dev20240506
 Summary: Flower: A Friendly Federated Learning Framework
 Home-page: https://flower.ai
 License: Apache-2.0

{flwr_nightly-1.9.0.dev20240502 → flwr_nightly-1.9.0.dev20240506}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "flwr-nightly"
-version = "1.9.0.dev20240502"
+version = "1.9.0.dev20240506"
 description = "Flower: A Friendly Federated Learning Framework"
 license = "Apache-2.0"
 authors = ["The Flower Authors <hello@flower.ai>"]

{flwr_nightly-1.9.0.dev20240502 → flwr_nightly-1.9.0.dev20240506}/src/py/flwr/common/message.py

@@ -18,7 +18,6 @@ from __future__ import annotations
 
 import time
 import warnings
-from dataclasses import dataclass
 from typing import Optional, cast
 
 from .record import RecordSet
@@ -26,7 +25,6 @@ from .record import RecordSet
 DEFAULT_TTL = 3600
 
 
-@dataclass
 class Metadata:  # pylint: disable=too-many-instance-attributes
     """A dataclass holding metadata associated with the current message.
 
@@ -161,8 +159,18 @@ class Metadata:  # pylint: disable=too-many-instance-attributes
         """Set partition_id."""
         self.__dict__["_partition_id"] = value
 
+    def __repr__(self) -> str:
+        """Return a string representation of this instance."""
+        view = ", ".join([f"{k.lstrip('_')}={v!r}" for k, v in self.__dict__.items()])
+        return f"{self.__class__.__qualname__}({view})"
+
+    def __eq__(self, other: object) -> bool:
+        """Compare two instances of the class."""
+        if not isinstance(other, self.__class__):
+            raise NotImplementedError
+        return self.__dict__ == other.__dict__
+
 
-@dataclass
 class Error:
     """A dataclass that stores information about an error that occurred.
 
@@ -191,8 +199,18 @@ class Error:
         """Reason reported about the error."""
         return cast(Optional[str], self.__dict__["_reason"])
 
+    def __repr__(self) -> str:
+        """Return a string representation of this instance."""
+        view = ", ".join([f"{k.lstrip('_')}={v!r}" for k, v in self.__dict__.items()])
+        return f"{self.__class__.__qualname__}({view})"
+
+    def __eq__(self, other: object) -> bool:
+        """Compare two instances of the class."""
+        if not isinstance(other, self.__class__):
+            raise NotImplementedError
+        return self.__dict__ == other.__dict__
+
 
-@dataclass
 class Message:
     """State of your application from the viewpoint of the entity using it.
 
@@ -357,6 +375,17 @@ class Message:
 
         return message
 
+    def __repr__(self) -> str:
+        """Return a string representation of this instance."""
+        view = ", ".join(
+            [
+                f"{k.lstrip('_')}={v!r}"
+                for k, v in self.__dict__.items()
+                if v is not None
+            ]
+        )
+        return f"{self.__class__.__qualname__}({view})"
+
 
 def _create_reply_metadata(msg: Message, ttl: float) -> Metadata:
     """Construct metadata for a reply message."""

{flwr_nightly-1.9.0.dev20240502 → flwr_nightly-1.9.0.dev20240506}/src/py/flwr/common/record/parametersrecord.py

@@ -82,7 +82,6 @@ def _check_value(value: Array) -> None:
         )
 
 
-@dataclass
 class ParametersRecord(TypedDict[str, Array]):
     """Parameters record.
 

{flwr_nightly-1.9.0.dev20240502 → flwr_nightly-1.9.0.dev20240506}/src/py/flwr/common/record/recordset.py

@@ -83,7 +83,6 @@ class RecordSetData:
             )
 
 
-@dataclass
 class RecordSet:
     """RecordSet stores groups of parameters, metrics and configs."""
 
@@ -117,3 +116,15 @@ class RecordSet:
         """Dictionary holding ConfigsRecord instances."""
         data = cast(RecordSetData, self.__dict__["_data"])
         return data.configs_records
+
+    def __repr__(self) -> str:
+        """Return a string representation of this instance."""
+        flds = ("parameters_records", "metrics_records", "configs_records")
+        view = ", ".join([f"{fld}={getattr(self, fld)!r}" for fld in flds])
+        return f"{self.__class__.__qualname__}({view})"
+
+    def __eq__(self, other: object) -> bool:
+        """Compare two instances of the class."""
+        if not isinstance(other, self.__class__):
+            raise NotImplementedError
+        return self.__dict__ == other.__dict__

{flwr_nightly-1.9.0.dev20240502 → flwr_nightly-1.9.0.dev20240506}/src/py/flwr/server/app.py

@@ -43,6 +43,7 @@ from flwr.common.constant import (
 from flwr.common.exit_handlers import register_exit_handlers
 from flwr.common.logger import log
 from flwr.common.secure_aggregation.crypto.symmetric_encryption import (
+    private_key_to_bytes,
     public_key_to_bytes,
     ssh_types_to_elliptic_curve,
 )
@@ -374,13 +375,18 @@ def run_superlink() -> None:
                 server_private_key,
                 server_public_key,
             ) = maybe_keys
-            interceptors = [
-                AuthenticateServerInterceptor(
-                    client_public_keys,
-                    server_private_key,
-                    server_public_key,
-                )
-            ]
+            state = state_factory.state()
+            state.store_client_public_keys(client_public_keys)
+            state.store_server_private_public_key(
+                private_key_to_bytes(server_private_key),
+                public_key_to_bytes(server_public_key),
+            )
+            log(
+                INFO,
+                "Client authentication enabled with %d known public keys",
+                len(client_public_keys),
+            )
+            interceptors = [AuthenticateServerInterceptor(state)]
 
         fleet_server = _run_fleet_api_grpc_rere(
             address=address,

{flwr_nightly-1.9.0.dev20240502 → flwr_nightly-1.9.0.dev20240506}/src/py/flwr/server/superlink/fleet/grpc_rere/server_interceptor.py

@@ -16,17 +16,17 @@
 
 
 import base64
-from logging import INFO
-from typing import Any, Callable, Sequence, Set, Tuple, Union
+from logging import WARNING
+from typing import Any, Callable, Optional, Sequence, Tuple, Union
 
 import grpc
 from cryptography.hazmat.primitives.asymmetric import ec
 
 from flwr.common.logger import log
 from flwr.common.secure_aggregation.crypto.symmetric_encryption import (
+    bytes_to_private_key,
     bytes_to_public_key,
     generate_shared_key,
-    public_key_to_bytes,
     verify_hmac,
 )
 from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
@@ -43,6 +43,8 @@ from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
     PushTaskResRequest,
     PushTaskResResponse,
 )
+from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
+from flwr.server.superlink.state import State
 
 _PUBLIC_KEY_HEADER = "public-key"
 _AUTH_TOKEN_HEADER = "auth-token"
@@ -79,22 +81,21 @@ def _get_value_from_tuples(
 class AuthenticateServerInterceptor(grpc.ServerInterceptor):  # type: ignore
     """Server interceptor for client authentication."""
 
-    def __init__(
-        self,
-        client_public_keys: Set[bytes],
-        private_key: ec.EllipticCurvePrivateKey,
-        public_key: ec.EllipticCurvePublicKey,
-    ):
-        self.server_private_key = private_key
-        self.client_public_keys = client_public_keys
-        self.encoded_server_public_key = base64.urlsafe_b64encode(
-            public_key_to_bytes(public_key)
-        )
-        log(
-            INFO,
-            "Client authentication enabled with %d known public keys",
-            len(client_public_keys),
-        )
+    def __init__(self, state: State):
+        self.state = state
+
+        self.client_public_keys = state.get_client_public_keys()
+        if len(self.client_public_keys) == 0:
+            log(WARNING, "Authentication enabled, but no known public keys configured")
+
+        private_key = self.state.get_server_private_key()
+        public_key = self.state.get_server_public_key()
+
+        if private_key is None or public_key is None:
+            raise ValueError("Error loading authentication keys")
+
+        self.server_private_key = bytes_to_private_key(private_key)
+        self.encoded_server_public_key = base64.urlsafe_b64encode(public_key)
 
     def intercept_service(
         self,
@@ -124,45 +125,29 @@ class AuthenticateServerInterceptor(grpc.ServerInterceptor):  # type: ignore
                     _PUBLIC_KEY_HEADER, context.invocation_metadata()
                 )
             )
-            is_public_key_known = client_public_key_bytes in self.client_public_keys
-            if not is_public_key_known:
+            if client_public_key_bytes not in self.client_public_keys:
                 context.abort(grpc.StatusCode.UNAUTHENTICATED, "Access denied")
 
             if isinstance(request, CreateNodeRequest):
-                context.send_initial_metadata(
-                    (
-                        (
-                            _PUBLIC_KEY_HEADER,
-                            self.encoded_server_public_key,
-                        ),
-                    )
+                return self._create_authenticated_node(
+                    client_public_key_bytes, request, context
                 )
-            elif isinstance(
-                request,
-                (
-                    DeleteNodeRequest,
-                    PullTaskInsRequest,
-                    PushTaskResRequest,
-                    GetRunRequest,
-                    PingRequest,
-                ),
-            ):
-                hmac_value = base64.urlsafe_b64decode(
-                    _get_value_from_tuples(
-                        _AUTH_TOKEN_HEADER, context.invocation_metadata()
-                    )
-                )
-                client_public_key = bytes_to_public_key(client_public_key_bytes)
-                shared_secret = generate_shared_key(
-                    self.server_private_key,
-                    client_public_key,
-                )
-                verify = verify_hmac(
-                    shared_secret, request.SerializeToString(True), hmac_value
+
+            # Verify hmac value
+            hmac_value = base64.urlsafe_b64decode(
+                _get_value_from_tuples(
+                    _AUTH_TOKEN_HEADER, context.invocation_metadata()
                 )
-                if not verify:
-                    context.abort(grpc.StatusCode.UNAUTHENTICATED, "Access denied")
-            else:
+            )
+            public_key = bytes_to_public_key(client_public_key_bytes)
+
+            if not self._verify_hmac(public_key, request, hmac_value):
+                context.abort(grpc.StatusCode.UNAUTHENTICATED, "Access denied")
+
+            # Verify node_id
+            node_id = self.state.get_node_id(client_public_key_bytes)
+
+            if not self._verify_node_id(node_id, request):
                 context.abort(grpc.StatusCode.UNAUTHENTICATED, "Access denied")
 
             return method_handler.unary_unary(request, context)  # type: ignore
@@ -172,3 +157,59 @@ class AuthenticateServerInterceptor(grpc.ServerInterceptor):  # type: ignore
             request_deserializer=method_handler.request_deserializer,
             response_serializer=method_handler.response_serializer,
         )
+
+    def _verify_node_id(
+        self,
+        node_id: Optional[int],
+        request: Union[
+            DeleteNodeRequest,
+            PullTaskInsRequest,
+            PushTaskResRequest,
+            GetRunRequest,
+            PingRequest,
+        ],
+    ) -> bool:
+        if node_id is None:
+            return False
+        if isinstance(request, PushTaskResRequest):
+            if len(request.task_res_list) == 0:
+                return False
+            return request.task_res_list[0].task.producer.node_id == node_id
+        if isinstance(request, GetRunRequest):
+            return node_id in self.state.get_nodes(request.run_id)
+        return request.node.node_id == node_id
+
+    def _verify_hmac(
+        self, public_key: ec.EllipticCurvePublicKey, request: Request, hmac_value: bytes
+    ) -> bool:
+        shared_secret = generate_shared_key(self.server_private_key, public_key)
+        return verify_hmac(shared_secret, request.SerializeToString(True), hmac_value)
+
+    def _create_authenticated_node(
+        self,
+        public_key_bytes: bytes,
+        request: CreateNodeRequest,
+        context: grpc.ServicerContext,
+    ) -> CreateNodeResponse:
+        context.send_initial_metadata(
+            (
+                (
+                    _PUBLIC_KEY_HEADER,
+                    self.encoded_server_public_key,
+                ),
+            )
+        )
+
+        node_id = self.state.get_node_id(public_key_bytes)
+
+        # Handle `CreateNode` here instead of calling the default method handler
+        # Return previously assigned `node_id` for the provided `public_key`
+        if node_id is not None:
+            self.state.acknowledge_ping(node_id, request.ping_interval)
+            return CreateNodeResponse(node=Node(node_id=node_id, anonymous=False))
+
+        # No `node_id` exists for the provided `public_key`
+        # Handle `CreateNode` here instead of calling the default method handler
+        # Note: the innermost `CreateNode` method will never be called
+        node_id = self.state.create_node(request.ping_interval, public_key_bytes)
+        return CreateNodeResponse(node=Node(node_id=node_id, anonymous=False))

{flwr_nightly-1.9.0.dev20240502 → flwr_nightly-1.9.0.dev20240506}/src/py/flwr/server/superlink/state/in_memory_state.py

@@ -30,19 +30,24 @@ from flwr.server.utils import validate_task_ins_or_res
 from .utils import make_node_unavailable_taskres
 
 
-class InMemoryState(State):  # pylint: disable=R0902
+class InMemoryState(State):  # pylint: disable=R0902,R0904
     """In-memory State implementation."""
 
     def __init__(self) -> None:
+
         # Map node_id to (online_until, ping_interval)
         self.node_ids: Dict[int, Tuple[float, float]] = {}
+        self.public_key_to_node_id: Dict[bytes, int] = {}
+
         # Map run_id to (fab_id, fab_version)
         self.run_ids: Dict[int, Tuple[str, str]] = {}
         self.task_ins_store: Dict[UUID, TaskIns] = {}
         self.task_res_store: Dict[UUID, TaskRes] = {}
+
         self.client_public_keys: Set[bytes] = set()
         self.server_public_key: Optional[bytes] = None
         self.server_private_key: Optional[bytes] = None
+
         self.lock = threading.Lock()
 
     def store_task_ins(self, task_ins: TaskIns) -> Optional[UUID]:
@@ -205,23 +210,46 @@ class InMemoryState(State):  # pylint: disable=R0902
         """
         return len(self.task_res_store)
 
-    def create_node(self, ping_interval: float) -> int:
+    def create_node(
+        self, ping_interval: float, public_key: Optional[bytes] = None
+    ) -> int:
         """Create, store in state, and return `node_id`."""
         # Sample a random int64 as node_id
         node_id: int = int.from_bytes(os.urandom(8), "little", signed=True)
 
         with self.lock:
-            if node_id not in self.node_ids:
-                self.node_ids[node_id] = (time.time() + ping_interval, ping_interval)
-                return node_id
-        log(ERROR, "Unexpected node registration failure.")
-        return 0
+            if node_id in self.node_ids:
+                log(ERROR, "Unexpected node registration failure.")
+                return 0
 
-    def delete_node(self, node_id: int) -> None:
+            if public_key is not None:
+                if (
+                    public_key in self.public_key_to_node_id
+                    or node_id in self.public_key_to_node_id.values()
+                ):
+                    log(ERROR, "Unexpected node registration failure.")
+                    return 0
+
+                self.public_key_to_node_id[public_key] = node_id
+
+            self.node_ids[node_id] = (time.time() + ping_interval, ping_interval)
+            return node_id
+
+    def delete_node(self, node_id: int, public_key: Optional[bytes] = None) -> None:
         """Delete a client node."""
         with self.lock:
             if node_id not in self.node_ids:
                 raise ValueError(f"Node {node_id} not found")
+
+            if public_key is not None:
+                if (
+                    public_key not in self.public_key_to_node_id
+                    or node_id not in self.public_key_to_node_id.values()
+                ):
+                    raise ValueError("Public key or node_id not found")
+
+                del self.public_key_to_node_id[public_key]
+
             del self.node_ids[node_id]
 
     def get_nodes(self, run_id: int) -> Set[int]:
@@ -242,6 +270,10 @@ class InMemoryState(State):  # pylint: disable=R0902
                 if online_until > current_time
             }
 
+    def get_node_id(self, client_public_key: bytes) -> Optional[int]:
+        """Retrieve stored `node_id` filtered by `client_public_keys`."""
+        return self.public_key_to_node_id.get(client_public_key)
+
     def create_run(self, fab_id: str, fab_version: str) -> int:
         """Create a new run for the specified `fab_id` and `fab_version`."""
         # Sample a random int64 as run_id

{flwr_nightly-1.9.0.dev20240502 → flwr_nightly-1.9.0.dev20240506}/src/py/flwr/server/superlink/state/sqlite_state.py

@@ -36,7 +36,8 @@ SQL_CREATE_TABLE_NODE = """
 CREATE TABLE IF NOT EXISTS node(
     node_id         INTEGER UNIQUE,
     online_until    REAL,
-    ping_interval   REAL
+    ping_interval   REAL,
+    public_key      BLOB
 );
 """
 
@@ -534,26 +535,54 @@ class SqliteState(State):  # pylint: disable=R0904
 
         return None
 
-    def create_node(self, ping_interval: float) -> int:
+    def create_node(
+        self, ping_interval: float, public_key: Optional[bytes] = None
+    ) -> int:
         """Create, store in state, and return `node_id`."""
         # Sample a random int64 as node_id
         node_id: int = int.from_bytes(os.urandom(8), "little", signed=True)
 
+        query = "SELECT node_id FROM node WHERE public_key = :public_key;"
+        row = self.query(query, {"public_key": public_key})
+
+        if len(row) > 0:
+            log(ERROR, "Unexpected node registration failure.")
+            return 0
+
         query = (
-            "INSERT INTO node (node_id, online_until, ping_interval) VALUES (?, ?, ?)"
+            "INSERT INTO node "
+            "(node_id, online_until, ping_interval, public_key) "
+            "VALUES (?, ?, ?, ?)"
         )
 
         try:
-            self.query(query, (node_id, time.time() + ping_interval, ping_interval))
+            self.query(
+                query, (node_id, time.time() + ping_interval, ping_interval, public_key)
+            )
         except sqlite3.IntegrityError:
             log(ERROR, "Unexpected node registration failure.")
             return 0
         return node_id
 
-    def delete_node(self, node_id: int) -> None:
+    def delete_node(self, node_id: int, public_key: Optional[bytes] = None) -> None:
         """Delete a client node."""
-        query = "DELETE FROM node WHERE node_id = :node_id;"
-        self.query(query, {"node_id": node_id})
+        query = "DELETE FROM node WHERE node_id = ?"
+        params = (node_id,)
+
+        if public_key is not None:
+            query += " AND public_key = ?"
+            params += (public_key,)  # type: ignore
+
+        if self.conn is None:
+            raise AttributeError("State is not initialized.")
+
+        try:
+            with self.conn:
+                rows = self.conn.execute(query, params)
+                if rows.rowcount < 1:
+                    raise ValueError("Public key or node_id not found")
+        except KeyError as exc:
+            log(ERROR, {"query": query, "data": params, "exception": exc})
 
     def get_nodes(self, run_id: int) -> Set[int]:
         """Retrieve all currently stored node IDs as a set.
@@ -574,6 +603,15 @@ class SqliteState(State):  # pylint: disable=R0904
         result: Set[int] = {row["node_id"] for row in rows}
         return result
 
+    def get_node_id(self, client_public_key: bytes) -> Optional[int]:
+        """Retrieve stored `node_id` filtered by `client_public_keys`."""
+        query = "SELECT node_id FROM node WHERE public_key = :public_key;"
+        row = self.query(query, {"public_key": client_public_key})
+        if len(row) > 0:
+            node_id: int = row[0]["node_id"]
+            return node_id
+        return None
+
     def create_run(self, fab_id: str, fab_version: str) -> int:
         """Create a new run for the specified `fab_id` and `fab_version`."""
         # Sample a random int64 as run_id

{flwr_nightly-1.9.0.dev20240502 → flwr_nightly-1.9.0.dev20240506}/src/py/flwr/server/superlink/state/state.py

@@ -22,7 +22,7 @@ from uuid import UUID
 from flwr.proto.task_pb2 import TaskIns, TaskRes  # pylint: disable=E0611
 
 
-class State(abc.ABC):
+class State(abc.ABC):  # pylint: disable=R0904
     """Abstract State."""
 
     @abc.abstractmethod
@@ -132,11 +132,13 @@ class State(abc.ABC):
         """Delete all delivered TaskIns/TaskRes pairs."""
 
     @abc.abstractmethod
-    def create_node(self, ping_interval: float) -> int:
+    def create_node(
+        self, ping_interval: float, public_key: Optional[bytes] = None
+    ) -> int:
         """Create, store in state, and return `node_id`."""
 
     @abc.abstractmethod
-    def delete_node(self, node_id: int) -> None:
+    def delete_node(self, node_id: int, public_key: Optional[bytes] = None) -> None:
         """Remove `node_id` from state."""
 
     @abc.abstractmethod
@@ -149,6 +151,10 @@ class State(abc.ABC):
         an empty `Set` MUST be returned.
         """
 
+    @abc.abstractmethod
+    def get_node_id(self, client_public_key: bytes) -> Optional[int]:
+        """Retrieve stored `node_id` filtered by `client_public_keys`."""
+
     @abc.abstractmethod
     def create_run(self, fab_id: str, fab_version: str) -> int:
         """Create a new run for the specified `fab_id` and `fab_version`."""