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

flwr/server/compat/driver_client_proxy.py

@@ -25,7 +25,7 @@ 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
+from ..driver.driver import GrpcDriverHelper
 
 SLEEP_TIME = 1
 
@@ -33,7 +33,9 @@ SLEEP_TIME = 1
 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: GrpcDriverHelper, anonymous: bool, run_id: int
+    ):
         super().__init__(str(node_id))
         self.node_id = node_id
         self.driver = driver

flwr/server/driver/__init__.py

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

flwr/server/driver/abc_driver.py

@@ -0,0 +1,140 @@
+# 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Driver (abstract base class)."""
+
+
+from abc import ABC, abstractmethod
+from typing import Iterable, List, Optional
+
+from flwr.common import Message, RecordSet
+
+
+class Driver(ABC):
+    """Abstract base Driver class for the Driver API."""
+
+    @abstractmethod
+    def create_message(  # pylint: disable=too-many-arguments
+        self,
+        content: RecordSet,
+        message_type: str,
+        dst_node_id: int,
+        group_id: str,
+        ttl: Optional[float] = None,
+    ) -> Message:
+        """Create a new message with specified parameters.
+
+        This method constructs a new `Message` with given content and metadata.
+        The `run_id` and `src_node_id` will be set automatically.
+
+        Parameters
+        ----------
+        content : RecordSet
+            The content for the new message. This holds records that are to be sent
+            to the destination node.
+        message_type : str
+            The type of the message, defining the action to be executed on
+            the receiving end.
+        dst_node_id : int
+            The ID of the destination node to which the message is being sent.
+        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 : 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 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.
+        """
+
+    @abstractmethod
+    def get_node_ids(self) -> List[int]:
+        """Get node IDs."""
+
+    @abstractmethod
+    def push_messages(self, messages: Iterable[Message]) -> Iterable[str]:
+        """Push messages to specified node IDs.
+
+        This method takes an iterable of messages and sends each message
+        to the node specified in `dst_node_id`.
+
+        Parameters
+        ----------
+        messages : Iterable[Message]
+            An iterable of messages to be sent.
+
+        Returns
+        -------
+        message_ids : Iterable[str]
+            An iterable of IDs for the messages that were sent, which can be used
+            to pull replies.
+        """
+
+    @abstractmethod
+    def pull_messages(self, message_ids: Iterable[str]) -> Iterable[Message]:
+        """Pull messages based on message IDs.
+
+        This method is used to collect messages from the SuperLink
+        that correspond to a set of given message IDs.
+
+        Parameters
+        ----------
+        message_ids : Iterable[str]
+            An iterable of message IDs for which reply messages are to be retrieved.
+
+        Returns
+        -------
+        messages : Iterable[Message]
+            An iterable of messages received.
+        """
+
+    @abstractmethod
+    def send_and_receive(
+        self,
+        messages: Iterable[Message],
+        *,
+        timeout: Optional[float] = None,
+    ) -> Iterable[Message]:
+        """Push messages to specified node IDs and pull the reply messages.
+
+        This method sends a list of messages to their destination node IDs and then
+        waits for the replies. It continues to pull replies until either all
+        replies are received or the specified timeout duration is exceeded.
+
+        Parameters
+        ----------
+        messages : Iterable[Message]
+            An iterable of messages to be sent.
+        timeout : Optional[float] (default: None)
+            The timeout duration in seconds. If specified, the method will wait for
+            replies for this duration. If `None`, there is no time limit and the method
+            will wait until replies for all messages are received.
+
+        Returns
+        -------
+        replies : Iterable[Message]
+            An iterable of reply messages received from the SuperLink.
+
+        Notes
+        -----
+        This method uses `push_messages` to send the messages and `pull_messages`
+        to collect the replies. If `timeout` is set, the method may not return
+        replies for all sent messages. A message remains valid until its TTL,
+        which is not affected by `timeout`.
+        """

flwr/server/driver/driver.py

@@ -16,20 +16,121 @@
 
 import time
 import warnings
+from logging import DEBUG, ERROR, WARNING
 from typing import Iterable, List, Optional, Tuple
 
-from flwr.common import DEFAULT_TTL, Message, Metadata, RecordSet
+import grpc
+
+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 .grpc_driver import DEFAULT_SERVER_ADDRESS_DRIVER, GrpcDriver
+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:
@@ -57,22 +158,22 @@ class Driver:
     ) -> None:
         self.addr = driver_service_address
         self.root_certificates = root_certificates
-        self.grpc_driver: Optional[GrpcDriver] = None
+        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_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:
+    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 = GrpcDriver(
+            self.grpc_driver_helper = GrpcDriverHelper(
                 driver_service_address=self.addr,
                 root_certificates=self.root_certificates,
             )
-            self.grpc_driver.connect()
-            res = self.grpc_driver.create_run(CreateRunRequest())
+            self.grpc_driver_helper.connect()
+            res = self.grpc_driver_helper.create_run(CreateRunRequest())
             self.run_id = res.run_id
-        return self.grpc_driver, self.run_id
+        return self.grpc_driver_helper, self.run_id
 
     def _check_message(self, message: Message) -> None:
         # Check if the message is valid
@@ -122,7 +223,7 @@ class Driver:
         message : Message
             A new `Message` instance with the specified content and metadata.
         """
-        _, run_id = self._get_grpc_driver_and_run_id()
+        _, 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 "
@@ -146,9 +247,9 @@ class Driver:
 
     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))
+        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]
 
     def push_messages(self, messages: Iterable[Message]) -> Iterable[str]:
@@ -168,7 +269,7 @@ 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()
+        grpc_driver_helper, _ = self._get_grpc_driver_helper_and_run_id()
         # Construct TaskIns
         task_ins_list: List[TaskIns] = []
         for msg in messages:
@@ -178,8 +279,10 @@ class Driver:
             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))
+        # Call GrpcDriverHelper method
+        res = grpc_driver_helper.push_task_ins(
+            PushTaskInsRequest(task_ins_list=task_ins_list)
+        )
         return list(res.task_ids)
 
     def pull_messages(self, message_ids: Iterable[str]) -> Iterable[Message]:
@@ -198,7 +301,7 @@ class Driver:
         messages : Iterable[Message]
             An iterable of messages received.
         """
-        grpc_driver, _ = self._get_grpc_driver_and_run_id()
+        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)
@@ -260,8 +363,8 @@ class Driver:
 
     def close(self) -> None:
         """Disconnect from the SuperLink if connected."""
-        # Check if GrpcDriver is initialized
-        if self.grpc_driver is None:
+        # Check if GrpcDriverHelper is initialized
+        if self.grpc_driver_helper is None:
             return
         # Disconnect
-        self.grpc_driver.disconnect()
+        self.grpc_driver_helper.disconnect()

flwr/server/superlink/driver/driver_servicer.py

@@ -64,7 +64,7 @@ class DriverServicer(driver_pb2_grpc.DriverServicer):
         """Create run ID."""
         log(INFO, "DriverServicer.CreateRun")
         state: State = self.state_factory.state()
-        run_id = state.create_run()
+        run_id = state.create_run("None/None", "None")
         return CreateRunResponse(run_id=run_id)
 
     def PushTaskIns(

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

@@ -33,6 +33,7 @@ from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
     PushTaskResRequest,
     PushTaskResResponse,
     Reconnect,
+    Run,
 )
 from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
 from flwr.proto.task_pb2 import TaskIns, TaskRes  # pylint: disable=E0611
@@ -109,4 +110,6 @@ def get_run(
     request: GetRunRequest, state: State  # pylint: disable=W0613
 ) -> GetRunResponse:
     """Get run information."""
-    return GetRunResponse()
+    run_id, fab_id, fab_version = state.get_run(request.run_id)
+    run = Run(run_id=run_id, fab_id=fab_id, fab_version=fab_version)
+    return GetRunResponse(run=run)

flwr/server/superlink/state/in_memory_state.py

@@ -36,7 +36,8 @@ class InMemoryState(State):
     def __init__(self) -> None:
         # Map node_id to (online_until, ping_interval)
         self.node_ids: Dict[int, Tuple[float, float]] = {}
-        self.run_ids: Set[int] = set()
+        # 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.lock = threading.Lock()
@@ -238,18 +239,26 @@ class InMemoryState(State):
                 if online_until > current_time
             }
 
-    def create_run(self) -> int:
-        """Create one run."""
+    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
         with self.lock:
             run_id: int = int.from_bytes(os.urandom(8), "little", signed=True)
 
             if run_id not in self.run_ids:
-                self.run_ids.add(run_id)
+                self.run_ids[run_id] = (fab_id, fab_version)
                 return run_id
         log(ERROR, "Unexpected run creation failure.")
         return 0
 
+    def get_run(self, run_id: int) -> Tuple[int, str, str]:
+        """Retrieve information about the run with the specified `run_id`."""
+        with self.lock:
+            if run_id not in self.run_ids:
+                log(ERROR, "`run_id` is invalid")
+                return 0, "", ""
+            return run_id, *self.run_ids[run_id]
+
     def acknowledge_ping(self, node_id: int, ping_interval: float) -> bool:
         """Acknowledge a ping received from a node, serving as a heartbeat."""
         with self.lock:

flwr/server/superlink/state/sqlite_state.py

@@ -46,7 +46,9 @@ CREATE INDEX IF NOT EXISTS idx_online_until ON node (online_until);
 
 SQL_CREATE_TABLE_RUN = """
 CREATE TABLE IF NOT EXISTS run(
-    run_id INTEGER UNIQUE
+    run_id          INTEGER UNIQUE,
+    fab_id          TEXT,
+    fab_version     TEXT
 );
 """
 
@@ -558,8 +560,8 @@ class SqliteState(State):
         result: Set[int] = {row["node_id"] for row in rows}
         return result
 
-    def create_run(self) -> int:
-        """Create one run and store it in state."""
+    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
         run_id: int = int.from_bytes(os.urandom(8), "little", signed=True)
 
@@ -567,12 +569,22 @@ class SqliteState(State):
         query = "SELECT COUNT(*) FROM run WHERE run_id = ?;"
         # If run_id does not exist
         if self.query(query, (run_id,))[0]["COUNT(*)"] == 0:
-            query = "INSERT INTO run VALUES(:run_id);"
-            self.query(query, {"run_id": run_id})
+            query = "INSERT INTO run (run_id, fab_id, fab_version) VALUES (?, ?, ?);"
+            self.query(query, (run_id, fab_id, fab_version))
             return run_id
         log(ERROR, "Unexpected run creation failure.")
         return 0
 
+    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 = ?;"
+        try:
+            row = self.query(query, (run_id,))[0]
+            return run_id, row["fab_id"], row["fab_version"]
+        except sqlite3.IntegrityError:
+            log(ERROR, "`run_id` does not exist.")
+            return 0, "", ""
+
     def acknowledge_ping(self, node_id: int, ping_interval: float) -> bool:
         """Acknowledge a ping received from a node, serving as a heartbeat."""
         # Update `online_until` and `ping_interval` for the given `node_id`

flwr/server/superlink/state/state.py

@@ -16,7 +16,7 @@
 
 
 import abc
-from typing import List, Optional, Set
+from typing import List, Optional, Set, Tuple
 from uuid import UUID
 
 from flwr.proto.task_pb2 import TaskIns, TaskRes  # pylint: disable=E0611
@@ -150,8 +150,26 @@ class State(abc.ABC):
         """
 
     @abc.abstractmethod
-    def create_run(self) -> int:
-        """Create one run."""
+    def create_run(self, fab_id: str, fab_version: str) -> int:
+        """Create a new run for the specified `fab_id` and `fab_version`."""
+
+    @abc.abstractmethod
+    def get_run(self, run_id: int) -> Tuple[int, str, str]:
+        """Retrieve information about the run with the specified `run_id`.
+
+        Parameters
+        ----------
+        run_id : int
+            The identifier of the run.
+
+        Returns
+        -------
+        Tuple[int, str, str]
+            A tuple containing three elements:
+            - `run_id`: The identifier of the run, same as the specified `run_id`.
+            - `fab_id`: The identifier of the FAB used in the specified run.
+            - `fab_version`: The version of the FAB used in the specified run.
+        """
 
     @abc.abstractmethod
     def acknowledge_ping(self, node_id: int, ping_interval: float) -> bool:

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

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