flwr-nightly 1.8.0.dev20240315__py3-none-any.whl → 1.15.0.dev20250114__py3-none-any.whl

flwr/server/client_manager.py

@@ -19,7 +19,7 @@ import random
 import threading
 from abc import ABC, abstractmethod
 from logging import INFO
-from typing import Dict, List, Optional
+from typing import Optional
 
 from flwr.common.logger import log
 
@@ -47,6 +47,7 @@ class ClientManager(ABC):
         Parameters
         ----------
         client : flwr.server.client_proxy.ClientProxy
+            The ClientProxy of the Client to register.
 
         Returns
         -------
@@ -64,10 +65,11 @@ class ClientManager(ABC):
         Parameters
         ----------
         client : flwr.server.client_proxy.ClientProxy
+            The ClientProxy of the Client to unregister.
         """
 
     @abstractmethod
-    def all(self) -> Dict[str, ClientProxy]:
+    def all(self) -> dict[str, ClientProxy]:
         """Return all available clients."""
 
     @abstractmethod
@@ -80,7 +82,7 @@ class ClientManager(ABC):
         num_clients: int,
         min_num_clients: Optional[int] = None,
         criterion: Optional[Criterion] = None,
-    ) -> List[ClientProxy]:
+    ) -> list[ClientProxy]:
         """Sample a number of Flower ClientProxy instances."""
 
 
@@ -88,7 +90,7 @@ class SimpleClientManager(ClientManager):
     """Provides a pool of available clients."""
 
     def __init__(self) -> None:
-        self.clients: Dict[str, ClientProxy] = {}
+        self.clients: dict[str, ClientProxy] = {}
         self._cv = threading.Condition()
 
     def __len__(self) -> int:
@@ -170,7 +172,7 @@ class SimpleClientManager(ClientManager):
             with self._cv:
                 self._cv.notify_all()
 
-    def all(self) -> Dict[str, ClientProxy]:
+    def all(self) -> dict[str, ClientProxy]:
         """Return all available clients."""
         return self.clients
 
@@ -179,7 +181,7 @@ class SimpleClientManager(ClientManager):
         num_clients: int,
         min_num_clients: Optional[int] = None,
         criterion: Optional[Criterion] = None,
-    ) -> List[ClientProxy]:
+    ) -> list[ClientProxy]:
         """Sample a number of Flower ClientProxy instances."""
         # Block until at least num_clients are connected.
         if min_num_clients is None:

flwr/server/compat/app.py

@@ -1,4 +1,4 @@
-# Copyright 2022 Flower Labs GmbH. All Rights Reserved.
+# Copyright 2024 Flower Labs GmbH. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -15,14 +15,10 @@
 """Flower driver app."""
 
 
-import sys
 from logging import INFO
-from pathlib import Path
-from typing import Optional, Union
+from typing import Optional
 
-from flwr.common import EventType, event
-from flwr.common.address import parse_address
-from flwr.common.logger import log, warn_deprecated_feature
+from flwr.common.logger import log
 from flwr.server.client_manager import ClientManager
 from flwr.server.history import History
 from flwr.server.server import Server, init_defaults, run_fl
@@ -32,33 +28,21 @@ from flwr.server.strategy import Strategy
 from ..driver import Driver
 from .app_utils import start_update_client_manager_thread
 
-DEFAULT_SERVER_ADDRESS_DRIVER = "[::]:9091"
-
-ERROR_MESSAGE_DRIVER_NOT_CONNECTED = """
-[Driver] Error: Not connected.
-
-Call `connect()` on the `Driver` instance before calling any of the other `Driver`
-methods.
-"""
-
 
 def start_driver(  # pylint: disable=too-many-arguments, too-many-locals
     *,
-    server_address: str = DEFAULT_SERVER_ADDRESS_DRIVER,
+    driver: Driver,
     server: Optional[Server] = None,
     config: Optional[ServerConfig] = None,
     strategy: Optional[Strategy] = None,
     client_manager: Optional[ClientManager] = None,
-    root_certificates: Optional[Union[bytes, str]] = None,
-    driver: Optional[Driver] = None,
 ) -> History:
     """Start a Flower Driver API server.
 
     Parameters
     ----------
-    server_address : Optional[str]
-        The IPv4 or IPv6 address of the Driver API server.
-        Defaults to `"[::]:8080"`.
+    driver : Driver
+        The Driver object to use.
     server : Optional[flwr.server.Server] (default: None)
         A server implementation, either `flwr.server.Server` or a subclass
         thereof. If no instance is provided, then `start_driver` will create
@@ -74,50 +58,12 @@ def start_driver(  # pylint: disable=too-many-arguments, too-many-locals
         An implementation of the class `flwr.server.ClientManager`. If no
         implementation is provided, then `start_driver` will use
         `flwr.server.SimpleClientManager`.
-    root_certificates : Optional[Union[bytes, str]] (default: None)
-        The PEM-encoded root certificates as a byte string or a path string.
-        If provided, a secure connection using the certificates will be
-        established to an SSL-enabled Flower server.
-    driver : Optional[Driver] (default: None)
-        The Driver object to use.
 
     Returns
     -------
     hist : flwr.server.history.History
         Object containing training and evaluation metrics.
-
-    Examples
-    --------
-    Starting a driver that connects to an insecure server:
-
-    >>> start_driver()
-
-    Starting a driver that connects to an SSL-enabled server:
-
-    >>> start_driver(
-    >>>     root_certificates=Path("/crts/root.pem").read_bytes()
-    >>> )
     """
-    event(EventType.START_DRIVER_ENTER)
-
-    if driver is None:
-        # Not passing a `Driver` object is deprecated
-        warn_deprecated_feature("start_driver")
-
-        # Parse IP address
-        parsed_address = parse_address(server_address)
-        if not parsed_address:
-            sys.exit(f"Server IP address ({server_address}) cannot be parsed.")
-        host, port, is_v6 = parsed_address
-        address = f"[{host}]:{port}" if is_v6 else f"{host}:{port}"
-
-        # Create the Driver
-        if isinstance(root_certificates, str):
-            root_certificates = Path(root_certificates).read_bytes()
-        driver = Driver(
-            driver_service_address=address, root_certificates=root_certificates
-        )
-
     # Initialize the Driver API server and config
     initialized_server, initialized_config = init_defaults(
         server=server,
@@ -147,6 +93,4 @@ def start_driver(  # pylint: disable=too-many-arguments, too-many-locals
     f_stop.set()
     thread.join()
 
-    event(EventType.START_SERVER_LEAVE)
-
     return hist

flwr/server/compat/app_utils.py

@@ -16,8 +16,8 @@
 
 
 import threading
-import time
-from typing import Dict, Tuple
+
+from flwr.common.typing import RunNotRunningException
 
 from ..client_manager import ClientManager
 from ..compat.driver_client_proxy import DriverClientProxy
@@ -27,7 +27,7 @@ from ..driver import Driver
 def start_update_client_manager_thread(
     driver: Driver,
     client_manager: ClientManager,
-) -> Tuple[threading.Thread, threading.Event]:
+) -> tuple[threading.Thread, threading.Event]:
     """Periodically update the nodes list in the client manager in a thread.
 
     This function starts a thread that periodically uses the associated driver to
@@ -60,6 +60,7 @@ def start_update_client_manager_thread(
             client_manager,
             f_stop,
         ),
+        daemon=True,
     )
     thread.start()
 
@@ -73,9 +74,13 @@ def _update_client_manager(
 ) -> None:
     """Update the nodes list in the client manager."""
     # Loop until the driver is disconnected
-    registered_nodes: Dict[int, DriverClientProxy] = {}
+    registered_nodes: dict[int, DriverClientProxy] = {}
     while not f_stop.is_set():
-        all_node_ids = set(driver.get_node_ids())
+        try:
+            all_node_ids = set(driver.get_node_ids())
+        except RunNotRunningException:
+            f_stop.set()
+            break
         dead_nodes = set(registered_nodes).difference(all_node_ids)
         new_nodes = all_node_ids.difference(registered_nodes)
 
@@ -89,9 +94,9 @@ def _update_client_manager(
         for node_id in new_nodes:
             client_proxy = DriverClientProxy(
                 node_id=node_id,
-                driver=driver.grpc_driver,  # type: ignore
+                driver=driver,
                 anonymous=False,
-                run_id=driver.run_id,  # type: ignore
+                run_id=driver.run.run_id,
             )
             if client_manager.register(client_proxy):
                 registered_nodes[node_id] = client_proxy
@@ -99,4 +104,5 @@ def _update_client_manager(
                 raise RuntimeError("Could not register node.")
 
         # Sleep for 3 seconds
-        time.sleep(3)
+        if not f_stop.is_set():
+            f_stop.wait(3)

flwr/server/compat/driver_client_proxy.py

@@ -1,4 +1,4 @@
-# Copyright 2020 Flower Labs GmbH. All Rights Reserved.
+# Copyright 2024 Flower Labs GmbH. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -15,25 +15,20 @@
 """Flower ClientProxy implementation for Driver API."""
 
 
-import time
-from typing import List, Optional
+from typing import Optional
 
 from flwr import common
-from flwr.common import MessageType, MessageTypeLegacy, RecordSet
+from flwr.common import Message, MessageType, MessageTypeLegacy, RecordSet
 from flwr.common import recordset_compat as compat
-from flwr.common import serde
-from flwr.proto import driver_pb2, node_pb2, task_pb2  # pylint: disable=E0611
 from flwr.server.client_proxy import ClientProxy
 
-from ..driver.grpc_driver import GrpcDriver
-
-SLEEP_TIME = 1
+from ..driver.driver import Driver
 
 
 class DriverClientProxy(ClientProxy):
     """Flower client proxy which delegates work using the Driver API."""
 
-    def __init__(self, node_id: int, driver: GrpcDriver, anonymous: bool, run_id: int):
+    def __init__(self, node_id: int, driver: Driver, anonymous: bool, run_id: int):
         super().__init__(str(node_id))
         self.node_id = node_id
         self.driver = driver
@@ -114,56 +109,28 @@ class DriverClientProxy(ClientProxy):
         timeout: Optional[float],
         group_id: Optional[int],
     ) -> RecordSet:
-        task_ins = task_pb2.TaskIns(  # pylint: disable=E1101
-            task_id="",
-            group_id=str(group_id) if group_id is not None else "",
-            run_id=self.run_id,
-            task=task_pb2.Task(  # pylint: disable=E1101
-                producer=node_pb2.Node(  # pylint: disable=E1101
-                    node_id=0,
-                    anonymous=True,
-                ),
-                consumer=node_pb2.Node(  # pylint: disable=E1101
-                    node_id=self.node_id,
-                    anonymous=self.anonymous,
-                ),
-                task_type=task_type,
-                recordset=serde.recordset_to_proto(recordset),
-            ),
-        )
-        push_task_ins_req = driver_pb2.PushTaskInsRequest(  # pylint: disable=E1101
-            task_ins_list=[task_ins]
-        )
-
-        # Send TaskIns to Driver API
-        push_task_ins_res = self.driver.push_task_ins(req=push_task_ins_req)
-
-        if len(push_task_ins_res.task_ids) != 1:
-            raise ValueError("Unexpected number of task_ids")
 
-        task_id = push_task_ins_res.task_ids[0]
-        if task_id == "":
-            raise ValueError(f"Failed to schedule task for node {self.node_id}")
-
-        if timeout:
-            start_time = time.time()
+        # Create message
+        message = self.driver.create_message(
+            content=recordset,
+            message_type=task_type,
+            dst_node_id=self.node_id,
+            group_id=str(group_id) if group_id else "",
+            ttl=timeout,
+        )
 
-        while True:
-            pull_task_res_req = driver_pb2.PullTaskResRequest(  # pylint: disable=E1101
-                node=node_pb2.Node(node_id=0, anonymous=True),  # pylint: disable=E1101
-                task_ids=[task_id],
-            )
+        # Send message and wait for reply
+        messages = list(self.driver.send_and_receive(messages=[message]))
 
-            # Ask Driver API for TaskRes
-            pull_task_res_res = self.driver.pull_task_res(req=pull_task_res_req)
+        # A single reply is expected
+        if len(messages) != 1:
+            raise ValueError(f"Expected one Message but got: {len(messages)}")
 
-            task_res_list: List[task_pb2.TaskRes] = list(  # pylint: disable=E1101
-                pull_task_res_res.task_res_list
+        # Only messages without errors can be handled beyond these point
+        msg: Message = messages[0]
+        if msg.has_error():
+            raise ValueError(
+                f"Message contains an Error (reason: {msg.error.reason}). "
+                "It originated during client-side execution of a message."
             )
-            if len(task_res_list) == 1:
-                task_res = task_res_list[0]
-                return serde.recordset_from_proto(task_res.task.recordset)
-
-            if timeout is not None and time.time() > start_time + timeout:
-                raise RuntimeError("Timeout reached")
-            time.sleep(SLEEP_TIME)
+        return msg.content

flwr/server/compat/legacy_context.py

@@ -18,7 +18,7 @@
 from dataclasses import dataclass
 from typing import Optional
 
-from flwr.common import Context, RecordSet
+from flwr.common import Context
 
 from ..client_manager import ClientManager, SimpleClientManager
 from ..history import History
@@ -35,9 +35,9 @@ class LegacyContext(Context):
     client_manager: ClientManager
     history: History
 
-    def __init__(
+    def __init__(  # pylint: disable=too-many-arguments
         self,
-        state: RecordSet,
+        context: Context,
         config: Optional[ServerConfig] = None,
         strategy: Optional[Strategy] = None,
         client_manager: Optional[ClientManager] = None,
@@ -52,4 +52,5 @@ class LegacyContext(Context):
         self.strategy = strategy
         self.client_manager = client_manager
         self.history = History()
-        super().__init__(state)
+
+        super().__init__(**vars(context))

flwr/server/driver/__init__.py

@@ -17,8 +17,10 @@
 
 from .driver import Driver
 from .grpc_driver import GrpcDriver
+from .inmemory_driver import InMemoryDriver
 
 __all__ = [
     "Driver",
     "GrpcDriver",
+    "InMemoryDriver",
 ]

flwr/server/driver/driver.py

@@ -1,4 +1,4 @@
-# Copyright 2022 Flower Labs GmbH. All Rights Reserved.
+# Copyright 2024 Flower Labs GmbH. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -12,85 +12,47 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Flower driver service client."""
+"""Driver (abstract base class)."""
 
 
-import time
-from typing import Iterable, List, Optional, Tuple
+from abc import ABC, abstractmethod
+from collections.abc import Iterable
+from typing import Optional
 
-from flwr.common import Message, Metadata, RecordSet
-from flwr.common.serde import message_from_taskres, message_to_taskins
-from flwr.proto.driver_pb2 import (  # pylint: disable=E0611
-    CreateRunRequest,
-    GetNodesRequest,
-    PullTaskResRequest,
-    PushTaskInsRequest,
-)
-from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
-from flwr.proto.task_pb2 import TaskIns  # pylint: disable=E0611
+from flwr.common import Message, RecordSet
+from flwr.common.typing import Run
 
-from .grpc_driver import DEFAULT_SERVER_ADDRESS_DRIVER, GrpcDriver
 
+class Driver(ABC):
+    """Abstract base Driver class for the ServerAppIo API."""
 
-class Driver:
-    """`Driver` class provides an interface to the Driver API.
+    @abstractmethod
+    def set_run(self, run_id: int) -> None:
+        """Request a run to the SuperLink with a given `run_id`.
 
-    Parameters
-    ----------
-    driver_service_address : Optional[str]
-        The IPv4 or IPv6 address of the Driver API server.
-        Defaults to `"[::]:9091"`.
-    certificates : bytes (default: None)
-        Tuple containing root certificate, server certificate, and private key
-        to start a secure SSL-enabled server. The tuple is expected to have
-        three bytes elements in the following order:
+        If a Run with the specified `run_id` exists, a local Run
+        object will be created. It enables further functionality
+        in the driver, such as sending `Messages`.
 
-            * CA certificate.
-            * server certificate.
-            * server private key.
-    """
-
-    def __init__(
-        self,
-        driver_service_address: str = DEFAULT_SERVER_ADDRESS_DRIVER,
-        root_certificates: Optional[bytes] = None,
-    ) -> None:
-        self.addr = driver_service_address
-        self.root_certificates = root_certificates
-        self.grpc_driver: Optional[GrpcDriver] = None
-        self.run_id: Optional[int] = None
-        self.node = Node(node_id=0, anonymous=True)
-
-    def _get_grpc_driver_and_run_id(self) -> Tuple[GrpcDriver, int]:
-        # Check if the GrpcDriver is initialized
-        if self.grpc_driver is None or self.run_id is None:
-            # Connect and create run
-            self.grpc_driver = GrpcDriver(
-                driver_service_address=self.addr,
-                root_certificates=self.root_certificates,
-            )
-            self.grpc_driver.connect()
-            res = self.grpc_driver.create_run(CreateRunRequest())
-            self.run_id = res.run_id
-        return self.grpc_driver, self.run_id
+        Parameters
+        ----------
+        run_id : int
+            The `run_id` of the Run this Driver object operates in.
+        """
 
-    def _check_message(self, message: Message) -> None:
-        # Check if the message is valid
-        if not (
-            message.metadata.run_id == self.run_id
-            and message.metadata.src_node_id == self.node.node_id
-            and message.metadata.message_id == ""
-            and message.metadata.reply_to_message == ""
-        ):
-            raise ValueError(f"Invalid message: {message}")
+    @property
+    @abstractmethod
+    def run(self) -> Run:
+        """Run information."""
 
-    def create_message(  # pylint: disable=too-many-arguments
+    @abstractmethod
+    def create_message(  # pylint: disable=too-many-arguments,R0917
         self,
         content: RecordSet,
         message_type: str,
         dst_node_id: int,
         group_id: str,
-        ttl: str,
+        ttl: Optional[float] = None,
     ) -> Message:
         """Create a new message with specified parameters.
 
@@ -110,36 +72,23 @@ class Driver:
         group_id : str
             The ID of the group to which this message is associated. In some settings,
             this is used as the FL round.
-        ttl : str
+        ttl : Optional[float] (default: None)
             Time-to-live for the round trip of this message, i.e., the time from sending
-            this message to receiving a reply. It specifies the duration for which the
-            message and its potential reply are considered valid.
+            this message to receiving a reply. It specifies in seconds the duration for
+            which the message and its potential reply are considered valid. If unset,
+            the default TTL (i.e., `common.DEFAULT_TTL`) will be used.
 
         Returns
         -------
         message : Message
             A new `Message` instance with the specified content and metadata.
         """
-        _, run_id = self._get_grpc_driver_and_run_id()
-        metadata = Metadata(
-            run_id=run_id,
-            message_id="",  # Will be set by the server
-            src_node_id=self.node.node_id,
-            dst_node_id=dst_node_id,
-            reply_to_message="",
-            group_id=group_id,
-            ttl=ttl,
-            message_type=message_type,
-        )
-        return Message(metadata=metadata, content=content)
 
-    def get_node_ids(self) -> List[int]:
+    @abstractmethod
+    def get_node_ids(self) -> list[int]:
         """Get node IDs."""
-        grpc_driver, run_id = self._get_grpc_driver_and_run_id()
-        # Call GrpcDriver method
-        res = grpc_driver.get_nodes(GetNodesRequest(run_id=run_id))
-        return [node.node_id for node in res.nodes]
 
+    @abstractmethod
     def push_messages(self, messages: Iterable[Message]) -> Iterable[str]:
         """Push messages to specified node IDs.
 
@@ -157,20 +106,8 @@ class Driver:
             An iterable of IDs for the messages that were sent, which can be used
             to pull replies.
         """
-        grpc_driver, _ = self._get_grpc_driver_and_run_id()
-        # Construct TaskIns
-        task_ins_list: List[TaskIns] = []
-        for msg in messages:
-            # Check message
-            self._check_message(msg)
-            # Convert Message to TaskIns
-            taskins = message_to_taskins(msg)
-            # Add to list
-            task_ins_list.append(taskins)
-        # Call GrpcDriver method
-        res = grpc_driver.push_task_ins(PushTaskInsRequest(task_ins_list=task_ins_list))
-        return list(res.task_ids)
 
+    @abstractmethod
     def pull_messages(self, message_ids: Iterable[str]) -> Iterable[Message]:
         """Pull messages based on message IDs.
 
@@ -187,15 +124,8 @@ class Driver:
         messages : Iterable[Message]
             An iterable of messages received.
         """
-        grpc_driver, _ = self._get_grpc_driver_and_run_id()
-        # Pull TaskRes
-        res = grpc_driver.pull_task_res(
-            PullTaskResRequest(node=self.node, task_ids=message_ids)
-        )
-        # Convert TaskRes to Message
-        msgs = [message_from_taskres(taskres) for taskres in res.task_res_list]
-        return msgs
 
+    @abstractmethod
     def send_and_receive(
         self,
         messages: Iterable[Message],
@@ -229,28 +159,3 @@ class Driver:
         replies for all sent messages. A message remains valid until its TTL,
         which is not affected by `timeout`.
         """
-        # Push messages
-        msg_ids = set(self.push_messages(messages))
-
-        # Pull messages
-        end_time = time.time() + (timeout if timeout is not None else 0.0)
-        ret: List[Message] = []
-        while timeout is None or time.time() < end_time:
-            res_msgs = self.pull_messages(msg_ids)
-            ret.extend(res_msgs)
-            msg_ids.difference_update(
-                {msg.metadata.reply_to_message for msg in res_msgs}
-            )
-            if len(msg_ids) == 0:
-                break
-            # Sleep
-            time.sleep(3)
-        return ret
-
-    def close(self) -> None:
-        """Disconnect from the SuperLink if connected."""
-        # Check if GrpcDriver is initialized
-        if self.grpc_driver is None:
-            return
-        # Disconnect
-        self.grpc_driver.disconnect()