flwr-nightly 1.9.0.dev20240420__py3-none-any.whl → 1.9.0.dev20240507__py3-none-any.whl

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

@@ -15,7 +15,7 @@
 """Ray backend for the Fleet API using the Simulation Engine."""
 
 import pathlib
-from logging import ERROR, INFO
+from logging import DEBUG, ERROR, INFO
 from typing import Callable, Dict, List, Tuple, Union
 
 import ray
@@ -46,7 +46,7 @@ class RayBackend(Backend):
     ) -> None:
         """Prepare RayBackend by initialising Ray and creating the ActorPool."""
         log(INFO, "Initialising: %s", self.__class__.__name__)
-        log(INFO, "Backend config: %s", backend_config)
+        log(DEBUG, "Backend config: %s", backend_config)
 
         if not pathlib.Path(work_dir).exists():
             raise ValueError(f"Specified work_dir {work_dir} does not exist.")
@@ -109,7 +109,7 @@ class RayBackend(Backend):
         else:
             client_resources = {"num_cpus": 2, "num_gpus": 0.0}
             log(
-                INFO,
+                DEBUG,
                 "`%s` not specified in backend config. Applying default setting: %s",
                 self.client_resources_key,
                 client_resources,
@@ -129,7 +129,7 @@ class RayBackend(Backend):
     async def build(self) -> None:
         """Build pool of Ray actors that this backend will submit jobs to."""
         await self.pool.add_actors_to_pool(self.pool.actors_capacity)
-        log(INFO, "Constructed ActorPool with: %i actors", self.pool.num_actors)
+        log(DEBUG, "Constructed ActorPool with: %i actors", self.pool.num_actors)
 
     async def process_message(
         self,
@@ -173,4 +173,4 @@ class RayBackend(Backend):
         """Terminate all actors in actor pool."""
         await self.pool.terminate_all_actors()
         ray.shutdown()
-        log(INFO, "Terminated %s", self.__class__.__name__)
+        log(DEBUG, "Terminated %s", self.__class__.__name__)

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

@@ -293,7 +293,7 @@ def start_vce(
         node_states[node_id] = NodeState()
 
     # Load backend config
-    log(INFO, "Supported backends: %s", list(supported_backends.keys()))
+    log(DEBUG, "Supported backends: %s", list(supported_backends.keys()))
     backend_config = json.loads(backend_config_json_stream)
 
     try:

flwr/server/superlink/state/in_memory_state.py

@@ -30,16 +30,24 @@ from flwr.server.utils import validate_task_ins_or_res
 from .utils import make_node_unavailable_taskres
 
 
-class InMemoryState(State):
+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]:
@@ -202,23 +210,46 @@ class InMemoryState(State):
         """
         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
+
+            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) -> None:
+    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]:
@@ -239,6 +270,10 @@ class InMemoryState(State):
                 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
@@ -251,6 +286,39 @@ class InMemoryState(State):
         log(ERROR, "Unexpected run creation failure.")
         return 0
 
+    def store_server_private_public_key(
+        self, private_key: bytes, public_key: bytes
+    ) -> None:
+        """Store `server_private_key` and `server_public_key` in state."""
+        with self.lock:
+            if self.server_private_key is None and self.server_public_key is None:
+                self.server_private_key = private_key
+                self.server_public_key = public_key
+            else:
+                raise RuntimeError("Server private and public key already set")
+
+    def get_server_private_key(self) -> Optional[bytes]:
+        """Retrieve `server_private_key` in urlsafe bytes."""
+        return self.server_private_key
+
+    def get_server_public_key(self) -> Optional[bytes]:
+        """Retrieve `server_public_key` in urlsafe bytes."""
+        return self.server_public_key
+
+    def store_client_public_keys(self, public_keys: Set[bytes]) -> None:
+        """Store a set of `client_public_keys` in state."""
+        with self.lock:
+            self.client_public_keys = public_keys
+
+    def store_client_public_key(self, public_key: bytes) -> None:
+        """Store a `client_public_key` in state."""
+        with self.lock:
+            self.client_public_keys.add(public_key)
+
+    def get_client_public_keys(self) -> Set[bytes]:
+        """Retrieve all currently stored `client_public_keys` as a set."""
+        return self.client_public_keys
+
     def get_run(self, run_id: int) -> Tuple[int, str, str]:
         """Retrieve information about the run with the specified `run_id`."""
         with self.lock:

flwr/server/superlink/state/sqlite_state.py

@@ -20,7 +20,7 @@ import re
 import sqlite3
 import time
 from logging import DEBUG, ERROR
-from typing import Any, Dict, List, Optional, Set, Tuple, Union, cast
+from typing import Any, Dict, List, Optional, Sequence, Set, Tuple, Union, cast
 from uuid import UUID, uuid4
 
 from flwr.common import log, now
@@ -36,7 +36,21 @@ 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
+);
+"""
+
+SQL_CREATE_TABLE_CREDENTIAL = """
+CREATE TABLE IF NOT EXISTS credential(
+    private_key BLOB PRIMARY KEY,
+    public_key BLOB
+);
+"""
+
+SQL_CREATE_TABLE_PUBLIC_KEY = """
+CREATE TABLE IF NOT EXISTS public_key(
+    public_key BLOB UNIQUE
 );
 """
 
@@ -72,7 +86,6 @@ CREATE TABLE IF NOT EXISTS task_ins(
 );
 """
 
-
 SQL_CREATE_TABLE_TASK_RES = """
 CREATE TABLE IF NOT EXISTS task_res(
     task_id                 TEXT UNIQUE,
@@ -96,7 +109,7 @@ CREATE TABLE IF NOT EXISTS task_res(
 DictOrTuple = Union[Tuple[Any, ...], Dict[str, Any]]
 
 
-class SqliteState(State):
+class SqliteState(State):  # pylint: disable=R0904
     """SQLite-based state implementation."""
 
     def __init__(
@@ -134,6 +147,8 @@ class SqliteState(State):
         cur.execute(SQL_CREATE_TABLE_TASK_INS)
         cur.execute(SQL_CREATE_TABLE_TASK_RES)
         cur.execute(SQL_CREATE_TABLE_NODE)
+        cur.execute(SQL_CREATE_TABLE_CREDENTIAL)
+        cur.execute(SQL_CREATE_TABLE_PUBLIC_KEY)
         cur.execute(SQL_CREATE_INDEX_ONLINE_UNTIL)
         res = cur.execute("SELECT name FROM sqlite_schema;")
 
@@ -142,7 +157,7 @@ class SqliteState(State):
     def query(
         self,
         query: str,
-        data: Optional[Union[List[DictOrTuple], DictOrTuple]] = None,
+        data: Optional[Union[Sequence[DictOrTuple], DictOrTuple]] = None,
     ) -> List[Dict[str, Any]]:
         """Execute a SQL query."""
         if self.conn is None:
@@ -520,26 +535,54 @@ class SqliteState(State):
 
         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.
@@ -560,6 +603,15 @@ class SqliteState(State):
         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
@@ -575,6 +627,59 @@ class SqliteState(State):
         log(ERROR, "Unexpected run creation failure.")
         return 0
 
+    def store_server_private_public_key(
+        self, private_key: bytes, public_key: bytes
+    ) -> None:
+        """Store `server_private_key` and `server_public_key` in state."""
+        query = "SELECT COUNT(*) FROM credential"
+        count = self.query(query)[0]["COUNT(*)"]
+        if count < 1:
+            query = (
+                "INSERT OR REPLACE INTO credential (private_key, public_key) "
+                "VALUES (:private_key, :public_key)"
+            )
+            self.query(query, {"private_key": private_key, "public_key": public_key})
+        else:
+            raise RuntimeError("Server private and public key already set")
+
+    def get_server_private_key(self) -> Optional[bytes]:
+        """Retrieve `server_private_key` in urlsafe bytes."""
+        query = "SELECT private_key FROM credential"
+        rows = self.query(query)
+        try:
+            private_key: Optional[bytes] = rows[0]["private_key"]
+        except IndexError:
+            private_key = None
+        return private_key
+
+    def get_server_public_key(self) -> Optional[bytes]:
+        """Retrieve `server_public_key` in urlsafe bytes."""
+        query = "SELECT public_key FROM credential"
+        rows = self.query(query)
+        try:
+            public_key: Optional[bytes] = rows[0]["public_key"]
+        except IndexError:
+            public_key = None
+        return public_key
+
+    def store_client_public_keys(self, public_keys: Set[bytes]) -> None:
+        """Store a set of `client_public_keys` in state."""
+        query = "INSERT INTO public_key (public_key) VALUES (?)"
+        data = [(key,) for key in public_keys]
+        self.query(query, data)
+
+    def store_client_public_key(self, public_key: bytes) -> None:
+        """Store a `client_public_key` in state."""
+        query = "INSERT INTO public_key (public_key) VALUES (:public_key)"
+        self.query(query, {"public_key": public_key})
+
+    def get_client_public_keys(self) -> Set[bytes]:
+        """Retrieve all currently stored `client_public_keys` as a set."""
+        query = "SELECT public_key FROM public_key"
+        rows = self.query(query)
+        result: Set[bytes] = {row["public_key"] for row in rows}
+        return result
+
     def get_run(self, run_id: int) -> Tuple[int, str, str]:
         """Retrieve information about the run with the specified `run_id`."""
         query = "SELECT * FROM run WHERE run_id = ?;"

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`."""
@@ -171,6 +177,32 @@ class State(abc.ABC):
             - `fab_version`: The version of the FAB used in the specified run.
         """
 
+    @abc.abstractmethod
+    def store_server_private_public_key(
+        self, private_key: bytes, public_key: bytes
+    ) -> None:
+        """Store `server_private_key` and `server_public_key` in state."""
+
+    @abc.abstractmethod
+    def get_server_private_key(self) -> Optional[bytes]:
+        """Retrieve `server_private_key` in urlsafe bytes."""
+
+    @abc.abstractmethod
+    def get_server_public_key(self) -> Optional[bytes]:
+        """Retrieve `server_public_key` in urlsafe bytes."""
+
+    @abc.abstractmethod
+    def store_client_public_keys(self, public_keys: Set[bytes]) -> None:
+        """Store a set of `client_public_keys` in state."""
+
+    @abc.abstractmethod
+    def store_client_public_key(self, public_key: bytes) -> None:
+        """Store a `client_public_key` in state."""
+
+    @abc.abstractmethod
+    def get_client_public_keys(self) -> Set[bytes]:
+        """Retrieve all currently stored `client_public_keys` as a set."""
+
     @abc.abstractmethod
     def acknowledge_ping(self, node_id: int, ping_interval: float) -> bool:
         """Acknowledge a ping received from a node, serving as a heartbeat.

flwr/simulation/__init__.py

@@ -17,7 +17,7 @@
 
 import importlib
 
-from flwr.simulation.run_simulation import run_simulation, run_simulation_from_cli
+from flwr.simulation.run_simulation import run_simulation
 
 is_ray_installed = importlib.util.find_spec("ray") is not None
 
@@ -36,4 +36,4 @@ To install the necessary dependencies, install `flwr` with the `simulation` extr
         raise ImportError(RAY_IMPORT_ERROR)
 
 
-__all__ = ["start_simulation", "run_simulation_from_cli", "run_simulation"]
+__all__ = ["start_simulation", "run_simulation"]

flwr/simulation/app.py

@@ -15,6 +15,8 @@
 """Flower simulation app."""
 
 
+import asyncio
+import logging
 import sys
 import threading
 import traceback
@@ -27,7 +29,7 @@ from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy
 
 from flwr.client import ClientFn
 from flwr.common import EventType, event
-from flwr.common.logger import log
+from flwr.common.logger import log, set_logger_propagation
 from flwr.server.client_manager import ClientManager
 from flwr.server.history import History
 from flwr.server.server import Server, init_defaults, run_fl
@@ -156,6 +158,7 @@ def start_simulation(
         is an advanced feature. For all details, please refer to the Ray documentation:
         https://docs.ray.io/en/latest/ray-core/scheduling/index.html
 
+
     Returns
     -------
     hist : flwr.server.history.History
@@ -167,6 +170,18 @@ def start_simulation(
         {"num_clients": len(clients_ids) if clients_ids is not None else num_clients},
     )
 
+    # Set logger propagation
+    loop: Optional[asyncio.AbstractEventLoop] = None
+    try:
+        loop = asyncio.get_running_loop()
+    except RuntimeError:
+        loop = None
+    finally:
+        if loop and loop.is_running():
+            # Set logger propagation to False to prevent duplicated log output in Colab.
+            logger = logging.getLogger("flwr")
+            _ = set_logger_propagation(logger, False)
+
     # Initialize server and server config
     initialized_server, initialized_config = init_defaults(
         server=server,

flwr/simulation/run_simulation.py

@@ -28,8 +28,9 @@ import grpc
 
 from flwr.client import ClientApp
 from flwr.common import EventType, event, log
+from flwr.common.logger import set_logger_propagation, update_console_handler
 from flwr.common.typing import ConfigsRecordValues
-from flwr.server.driver.driver import Driver
+from flwr.server.driver import Driver, GrpcDriver
 from flwr.server.run_serverapp import run
 from flwr.server.server_app import ServerApp
 from flwr.server.superlink.driver.driver_grpc import run_driver_api_grpc
@@ -154,7 +155,7 @@ def run_serverapp_th(
             # Upon completion, trigger stop event if one was passed
             if stop_event is not None:
                 stop_event.set()
-                log(WARNING, "Triggered stop event for Simulation Engine.")
+                log(DEBUG, "Triggered stop event for Simulation Engine.")
 
     serverapp_th = threading.Thread(
         target=server_th_with_start_checks,
@@ -204,7 +205,7 @@ def _main_loop(
     serverapp_th = None
     try:
         # Initialize Driver
-        driver = Driver(
+        driver = GrpcDriver(
             driver_service_address=driver_api_address,
             root_certificates=None,
         )
@@ -248,7 +249,7 @@ def _main_loop(
         if serverapp_th:
             serverapp_th.join()
 
-    log(INFO, "Stopping Simulation Engine now.")
+    log(DEBUG, "Stopping Simulation Engine now.")
 
 
 # pylint: disable=too-many-arguments,too-many-locals
@@ -317,9 +318,9 @@ def _run_simulation(
         enabled, DEBUG-level logs will be displayed.
     """
     # Set logging level
-    if not verbose_logging:
-        logger = logging.getLogger("flwr")
-        logger.setLevel(INFO)
+    logger = logging.getLogger("flwr")
+    if verbose_logging:
+        update_console_handler(level=DEBUG, timestamps=True, colored=True)
 
     if backend_config is None:
         backend_config = {}
@@ -364,6 +365,8 @@ def _run_simulation(
 
     finally:
         if run_in_thread:
+            # Set logger propagation to False to prevent duplicated log output in Colab.
+            logger = set_logger_propagation(logger, False)
             log(DEBUG, "Starting Simulation Engine on a new thread.")
             simulation_engine_th = threading.Thread(target=_main_loop, args=args)
             simulation_engine_th.start()

{flwr_nightly-1.9.0.dev20240420.dist-info → flwr_nightly-1.9.0.dev20240507.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flwr-nightly
-Version: 1.9.0.dev20240420
+Version: 1.9.0.dev20240507
 Summary: Flower: A Friendly Federated Learning Framework
 Home-page: https://flower.ai
 License: Apache-2.0
@@ -36,6 +36,7 @@ Requires-Dist: cryptography (>=42.0.4,<43.0.0)
 Requires-Dist: grpcio (>=1.60.0,<2.0.0)
 Requires-Dist: iterators (>=0.0.2,<0.0.3)
 Requires-Dist: numpy (>=1.21.0,<2.0.0)
+Requires-Dist: pathspec (>=0.12.1,<0.13.0)
 Requires-Dist: protobuf (>=4.25.2,<5.0.0)
 Requires-Dist: pycryptodome (>=3.18.0,<4.0.0)
 Requires-Dist: ray (==2.6.3) ; (python_version >= "3.8" and python_version < "3.12") and (extra == "simulation")
@@ -193,7 +194,7 @@ Other [examples](https://github.com/adap/flower/tree/main/examples):
 - [PyTorch: From Centralized to Federated](https://github.com/adap/flower/tree/main/examples/pytorch-from-centralized-to-federated)
 - [Vertical FL](https://github.com/adap/flower/tree/main/examples/vertical-fl)
 - [Federated Finetuning of OpenAI's Whisper](https://github.com/adap/flower/tree/main/examples/whisper-federated-finetuning)
-- [Federated Finetuning of Large Language Model](https://github.com/adap/flower/tree/main/examples/fedllm-finetune)
+- [Federated Finetuning of Large Language Model](https://github.com/adap/flower/tree/main/examples/llm-flowertune)
 - [Federated Finetuning of a Vision Transformer](https://github.com/adap/flower/tree/main/examples/vit-finetune)
 - [Advanced Flower with TensorFlow/Keras](https://github.com/adap/flower/tree/main/examples/advanced-tensorflow)
 - [Advanced Flower with PyTorch](https://github.com/adap/flower/tree/main/examples/advanced-pytorch)