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

flwr/server/compat/driver_client_proxy.py

@@ -16,16 +16,14 @@
 
 
 import time
-from typing import List, Optional
+from typing import Optional
 
 from flwr import common
-from flwr.common import DEFAULT_TTL, 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.driver import GrpcDriverHelper
+from ..driver.driver import Driver
 
 SLEEP_TIME = 1
 
@@ -33,9 +31,7 @@ SLEEP_TIME = 1
 class DriverClientProxy(ClientProxy):
     """Flower client proxy which delegates work using the Driver API."""
 
-    def __init__(
-        self, node_id: int, driver: GrpcDriverHelper, 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
@@ -116,80 +112,39 @@ 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),
-                ttl=DEFAULT_TTL,
-            ),
-        )
-
-        # This would normally be recorded upon common.Message creation
-        # but this compatibility stack doesn't create Messages,
-        # so we need to inject `created_at` manually (needed for
-        # taskins validation by server.utils.validator)
-        task_ins.task.created_at = time.time()
 
-        push_task_ins_req = driver_pb2.PushTaskInsRequest(  # pylint: disable=E1101
-            task_ins_list=[task_ins]
+        # 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,
         )
 
-        # 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")
+        # Push message
+        message_ids = list(self.driver.push_messages(messages=[message]))
+        if len(message_ids) != 1:
+            raise ValueError("Unexpected number of message_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}")
+        message_id = message_ids[0]
+        if message_id == "":
+            raise ValueError(f"Failed to send message to node {self.node_id}")
 
         if timeout:
             start_time = time.time()
 
         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],
-            )
-
-            # Ask Driver API for TaskRes
-            pull_task_res_res = self.driver.pull_task_res(req=pull_task_res_req)
-
-            task_res_list: List[task_pb2.TaskRes] = list(  # pylint: disable=E1101
-                pull_task_res_res.task_res_list
-            )
-            if len(task_res_list) == 1:
-                task_res = task_res_list[0]
-
-                # This will raise an Exception if task_res carries an `error`
-                validate_task_res(task_res=task_res)
-
-                return serde.recordset_from_proto(task_res.task.recordset)
+            messages = list(self.driver.pull_messages(message_ids))
+            if len(messages) == 1:
+                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."
+                    )
+                return msg.content
 
             if timeout is not None and time.time() > start_time + timeout:
                 raise RuntimeError("Timeout reached")
             time.sleep(SLEEP_TIME)
-
-
-def validate_task_res(
-    task_res: task_pb2.TaskRes,  # pylint: disable=E1101
-) -> None:
-    """Validate if a TaskRes is empty or not."""
-    if not task_res.HasField("task"):
-        raise ValueError("Invalid TaskRes, field `task` missing")
-    if task_res.task.HasField("error"):
-        raise ValueError("Exception during client-side task execution")
-    if not task_res.task.HasField("recordset"):
-        raise ValueError("Invalid TaskRes, both `recordset` and `error` are missing")

flwr/server/driver/__init__.py

@@ -16,7 +16,10 @@
 
 
 from .driver import Driver
+from .grpc_driver import GrpcDriver, GrpcDriverHelper
 
 __all__ = [
     "Driver",
+    "GrpcDriver",
+    "GrpcDriverHelper",
 ]

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,180 +12,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Flower driver service client."""
+"""Driver (abstract base class)."""
 
-import time
-import warnings
-from logging import DEBUG, ERROR, WARNING
-from typing import Iterable, List, Optional, Tuple
 
-import grpc
+from abc import ABC, abstractmethod
+from typing import Iterable, List, Optional
 
-from flwr.common import DEFAULT_TTL, EventType, Message, Metadata, RecordSet, event
-from flwr.common.grpc import create_channel
-from flwr.common.logger import log
-from flwr.common.serde import message_from_taskres, message_to_taskins
-from flwr.proto.driver_pb2 import (  # pylint: disable=E0611
-    CreateRunRequest,
-    CreateRunResponse,
-    GetNodesRequest,
-    GetNodesResponse,
-    PullTaskResRequest,
-    PullTaskResResponse,
-    PushTaskInsRequest,
-    PushTaskInsResponse,
-)
-from flwr.proto.driver_pb2_grpc import DriverStub  # pylint: disable=E0611
-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
 
-DEFAULT_SERVER_ADDRESS_DRIVER = "[::]:9091"
 
-ERROR_MESSAGE_DRIVER_NOT_CONNECTED = """
-[Driver] Error: Not connected.
-
-Call `connect()` on the `GrpcDriverHelper` instance before calling any of the other
-`GrpcDriverHelper` methods.
-"""
-
-
-class GrpcDriverHelper:
-    """`GrpcDriverHelper` provides access to the gRPC Driver API/service."""
-
-    def __init__(
-        self,
-        driver_service_address: str = DEFAULT_SERVER_ADDRESS_DRIVER,
-        root_certificates: Optional[bytes] = None,
-    ) -> None:
-        self.driver_service_address = driver_service_address
-        self.root_certificates = root_certificates
-        self.channel: Optional[grpc.Channel] = None
-        self.stub: Optional[DriverStub] = None
-
-    def connect(self) -> None:
-        """Connect to the Driver API."""
-        event(EventType.DRIVER_CONNECT)
-        if self.channel is not None or self.stub is not None:
-            log(WARNING, "Already connected")
-            return
-        self.channel = create_channel(
-            server_address=self.driver_service_address,
-            insecure=(self.root_certificates is None),
-            root_certificates=self.root_certificates,
-        )
-        self.stub = DriverStub(self.channel)
-        log(DEBUG, "[Driver] Connected to %s", self.driver_service_address)
-
-    def disconnect(self) -> None:
-        """Disconnect from the Driver API."""
-        event(EventType.DRIVER_DISCONNECT)
-        if self.channel is None or self.stub is None:
-            log(DEBUG, "Already disconnected")
-            return
-        channel = self.channel
-        self.channel = None
-        self.stub = None
-        channel.close()
-        log(DEBUG, "[Driver] Disconnected")
-
-    def create_run(self, req: CreateRunRequest) -> CreateRunResponse:
-        """Request for run ID."""
-        # Check if channel is open
-        if self.stub is None:
-            log(ERROR, ERROR_MESSAGE_DRIVER_NOT_CONNECTED)
-            raise ConnectionError("`GrpcDriverHelper` instance not connected")
-
-        # Call Driver API
-        res: CreateRunResponse = self.stub.CreateRun(request=req)
-        return res
-
-    def get_nodes(self, req: GetNodesRequest) -> GetNodesResponse:
-        """Get client IDs."""
-        # Check if channel is open
-        if self.stub is None:
-            log(ERROR, ERROR_MESSAGE_DRIVER_NOT_CONNECTED)
-            raise ConnectionError("`GrpcDriverHelper` instance not connected")
-
-        # Call gRPC Driver API
-        res: GetNodesResponse = self.stub.GetNodes(request=req)
-        return res
-
-    def push_task_ins(self, req: PushTaskInsRequest) -> PushTaskInsResponse:
-        """Schedule tasks."""
-        # Check if channel is open
-        if self.stub is None:
-            log(ERROR, ERROR_MESSAGE_DRIVER_NOT_CONNECTED)
-            raise ConnectionError("`GrpcDriverHelper` instance not connected")
-
-        # Call gRPC Driver API
-        res: PushTaskInsResponse = self.stub.PushTaskIns(request=req)
-        return res
-
-    def pull_task_res(self, req: PullTaskResRequest) -> PullTaskResResponse:
-        """Get task results."""
-        # Check if channel is open
-        if self.stub is None:
-            log(ERROR, ERROR_MESSAGE_DRIVER_NOT_CONNECTED)
-            raise ConnectionError("`GrpcDriverHelper` instance not connected")
-
-        # Call Driver API
-        res: PullTaskResResponse = self.stub.PullTaskRes(request=req)
-        return res
-
-
-class Driver:
-    """`Driver` class provides an interface to the Driver API.
-
-    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:
-
-            * 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_helper: Optional[GrpcDriverHelper] = None
-        self.run_id: Optional[int] = None
-        self.node = Node(node_id=0, anonymous=True)
-
-    def _get_grpc_driver_helper_and_run_id(self) -> Tuple[GrpcDriverHelper, int]:
-        # Check if the GrpcDriverHelper is initialized
-        if self.grpc_driver_helper is None or self.run_id is None:
-            # Connect and create run
-            self.grpc_driver_helper = GrpcDriverHelper(
-                driver_service_address=self.addr,
-                root_certificates=self.root_certificates,
-            )
-            self.grpc_driver_helper.connect()
-            res = self.grpc_driver_helper.create_run(CreateRunRequest())
-            self.run_id = res.run_id
-        return self.grpc_driver_helper, self.run_id
-
-    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 == ""
-            and message.metadata.ttl > 0
-        ):
-            raise ValueError(f"Invalid message: {message}")
+class Driver(ABC):
+    """Abstract base Driver class for the Driver API."""
 
+    @abstractmethod
     def create_message(  # pylint: disable=too-many-arguments
         self,
         content: RecordSet,
@@ -223,35 +62,12 @@ class Driver:
         message : Message
             A new `Message` instance with the specified content and metadata.
         """
-        _, run_id = self._get_grpc_driver_helper_and_run_id()
-        if ttl:
-            warnings.warn(
-                "A custom TTL was set, but note that the SuperLink does not enforce "
-                "the TTL yet. The SuperLink will start enforcing the TTL in a future "
-                "version of Flower.",
-                stacklevel=2,
-            )
-
-        ttl_ = DEFAULT_TTL if ttl is None else ttl
-        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)
 
+    @abstractmethod
     def get_node_ids(self) -> List[int]:
         """Get node IDs."""
-        grpc_driver_helper, run_id = self._get_grpc_driver_helper_and_run_id()
-        # Call GrpcDriverHelper method
-        res = grpc_driver_helper.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.
 
@@ -269,22 +85,8 @@ class Driver:
             An iterable of IDs for the messages that were sent, which can be used
             to pull replies.
         """
-        grpc_driver_helper, _ = self._get_grpc_driver_helper_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 GrpcDriverHelper method
-        res = grpc_driver_helper.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.
 
@@ -301,15 +103,8 @@ class Driver:
         messages : Iterable[Message]
             An iterable of messages received.
         """
-        grpc_driver, _ = self._get_grpc_driver_helper_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],
@@ -343,28 +138,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 GrpcDriverHelper is initialized
-        if self.grpc_driver_helper is None:
-            return
-        # Disconnect
-        self.grpc_driver_helper.disconnect()