flwr-nightly 1.8.0.dev20240227__py3-none-any.whl → 1.8.0.dev20240229__py3-none-any.whl

flwr/proto/task_pb2.pyi

@@ -3,6 +3,7 @@
 isort:skip_file
 """
 import builtins
+import flwr.proto.error_pb2
 import flwr.proto.node_pb2
 import flwr.proto.recordset_pb2
 import google.protobuf.descriptor
@@ -23,6 +24,7 @@ class Task(google.protobuf.message.Message):
     ANCESTRY_FIELD_NUMBER: builtins.int
     TASK_TYPE_FIELD_NUMBER: builtins.int
     RECORDSET_FIELD_NUMBER: builtins.int
+    ERROR_FIELD_NUMBER: builtins.int
     @property
     def producer(self) -> flwr.proto.node_pb2.Node: ...
     @property
@@ -35,6 +37,8 @@ class Task(google.protobuf.message.Message):
     task_type: typing.Text
     @property
     def recordset(self) -> flwr.proto.recordset_pb2.RecordSet: ...
+    @property
+    def error(self) -> flwr.proto.error_pb2.Error: ...
     def __init__(self,
         *,
         producer: typing.Optional[flwr.proto.node_pb2.Node] = ...,
@@ -45,9 +49,10 @@ class Task(google.protobuf.message.Message):
         ancestry: typing.Optional[typing.Iterable[typing.Text]] = ...,
         task_type: typing.Text = ...,
         recordset: typing.Optional[flwr.proto.recordset_pb2.RecordSet] = ...,
+        error: typing.Optional[flwr.proto.error_pb2.Error] = ...,
         ) -> None: ...
-    def HasField(self, field_name: typing_extensions.Literal["consumer",b"consumer","producer",b"producer","recordset",b"recordset"]) -> builtins.bool: ...
-    def ClearField(self, field_name: typing_extensions.Literal["ancestry",b"ancestry","consumer",b"consumer","created_at",b"created_at","delivered_at",b"delivered_at","producer",b"producer","recordset",b"recordset","task_type",b"task_type","ttl",b"ttl"]) -> None: ...
+    def HasField(self, field_name: typing_extensions.Literal["consumer",b"consumer","error",b"error","producer",b"producer","recordset",b"recordset"]) -> builtins.bool: ...
+    def ClearField(self, field_name: typing_extensions.Literal["ancestry",b"ancestry","consumer",b"consumer","created_at",b"created_at","delivered_at",b"delivered_at","error",b"error","producer",b"producer","recordset",b"recordset","task_type",b"task_type","ttl",b"ttl"]) -> None: ...
 global___Task = Task
 
 class TaskIns(google.protobuf.message.Message):

flwr/server/__init__.py

@@ -16,12 +16,14 @@
 
 
 from . import strategy
+from . import workflow as workflow
 from .app import run_driver_api as run_driver_api
 from .app import run_fleet_api as run_fleet_api
 from .app import run_superlink as run_superlink
 from .app import start_server as start_server
 from .client_manager import ClientManager as ClientManager
 from .client_manager import SimpleClientManager as SimpleClientManager
+from .compat import LegacyContext as LegacyContext
 from .compat import start_driver as start_driver
 from .driver import Driver as Driver
 from .history import History as History
@@ -34,6 +36,7 @@ __all__ = [
     "ClientManager",
     "Driver",
     "History",
+    "LegacyContext",
     "run_driver_api",
     "run_fleet_api",
     "run_server_app",
@@ -45,4 +48,5 @@ __all__ = [
     "start_driver",
     "start_server",
     "strategy",
+    "workflow",
 ]

flwr/server/app.py

@@ -14,8 +14,8 @@
 # ==============================================================================
 """Flower server app."""
 
-
 import argparse
+import asyncio
 import importlib.util
 import sys
 import threading
@@ -36,9 +36,6 @@ from flwr.common.constant import (
 )
 from flwr.common.exit_handlers import register_exit_handlers
 from flwr.common.logger import log
-from flwr.proto.driver_pb2_grpc import (  # pylint: disable=E0611
-    add_DriverServicer_to_server,
-)
 from flwr.proto.fleet_pb2_grpc import (  # pylint: disable=E0611
     add_FleetServicer_to_server,
 )
@@ -48,7 +45,7 @@ from .history import History
 from .server import Server, init_defaults, run_fl
 from .server_config import ServerConfig
 from .strategy import Strategy
-from .superlink.driver.driver_servicer import DriverServicer
+from .superlink.driver.driver_grpc import run_driver_api_grpc
 from .superlink.fleet.grpc_bidi.grpc_server import (
     generic_create_grpc_server,
     start_grpc_server,
@@ -204,7 +201,7 @@ def run_driver_api() -> None:
     state_factory = StateFactory(args.database)
 
     # Start server
-    grpc_server: grpc.Server = _run_driver_api_grpc(
+    grpc_server: grpc.Server = run_driver_api_grpc(
         address=address,
         state_factory=state_factory,
         certificates=certificates,
@@ -313,7 +310,7 @@ def run_superlink() -> None:
     state_factory = StateFactory(args.database)
 
     # Start Driver API
-    driver_server: grpc.Server = _run_driver_api_grpc(
+    driver_server: grpc.Server = run_driver_api_grpc(
         address=address,
         state_factory=state_factory,
         certificates=certificates,
@@ -362,6 +359,7 @@ def run_superlink() -> None:
         )
         grpc_servers.append(fleet_server)
     elif args.fleet_api_type == TRANSPORT_TYPE_VCE:
+        f_stop = asyncio.Event()  # Does nothing
         _run_fleet_api_vce(
             num_supernodes=args.num_supernodes,
             client_app_module_name=args.client_app,
@@ -369,6 +367,7 @@ def run_superlink() -> None:
             backend_config_json_stream=args.backend_config,
             working_dir=args.dir,
             state_factory=state_factory,
+            f_stop=f_stop,
         )
     else:
         raise ValueError(f"Unknown fleet_api_type: {args.fleet_api_type}")
@@ -412,30 +411,6 @@ def _try_obtain_certificates(
     return certificates
 
 
-def _run_driver_api_grpc(
-    address: str,
-    state_factory: StateFactory,
-    certificates: Optional[Tuple[bytes, bytes, bytes]],
-) -> grpc.Server:
-    """Run Driver API (gRPC, request-response)."""
-    # Create Driver API gRPC server
-    driver_servicer: grpc.Server = DriverServicer(
-        state_factory=state_factory,
-    )
-    driver_add_servicer_to_server_fn = add_DriverServicer_to_server
-    driver_grpc_server = generic_create_grpc_server(
-        servicer_and_add_fn=(driver_servicer, driver_add_servicer_to_server_fn),
-        server_address=address,
-        max_message_length=GRPC_MAX_MESSAGE_LENGTH,
-        certificates=certificates,
-    )
-
-    log(INFO, "Flower ECE: Starting Driver API (gRPC-rere) on %s", address)
-    driver_grpc_server.start()
-
-    return driver_grpc_server
-
-
 def _run_fleet_api_grpc_rere(
     address: str,
     state_factory: StateFactory,
@@ -468,6 +443,7 @@ def _run_fleet_api_vce(
     backend_config_json_stream: str,
     working_dir: str,
     state_factory: StateFactory,
+    f_stop: asyncio.Event,
 ) -> None:
     log(INFO, "Flower VCE: Starting Fleet API (VirtualClientEngine)")
 
@@ -478,6 +454,7 @@ def _run_fleet_api_vce(
         backend_config_json_stream=backend_config_json_stream,
         state_factory=state_factory,
         working_dir=working_dir,
+        f_stop=f_stop,
     )
 
 

flwr/server/client_proxy.py

@@ -47,6 +47,7 @@ class ClientProxy(ABC):
         self,
         ins: GetPropertiesIns,
         timeout: Optional[float],
+        group_id: Optional[int],
     ) -> GetPropertiesRes:
         """Return the client's properties."""
 
@@ -55,6 +56,7 @@ class ClientProxy(ABC):
         self,
         ins: GetParametersIns,
         timeout: Optional[float],
+        group_id: Optional[int],
     ) -> GetParametersRes:
         """Return the current local model parameters."""
 
@@ -63,6 +65,7 @@ class ClientProxy(ABC):
         self,
         ins: FitIns,
         timeout: Optional[float],
+        group_id: Optional[int],
     ) -> FitRes:
         """Refine the provided parameters using the locally held dataset."""
 
@@ -71,6 +74,7 @@ class ClientProxy(ABC):
         self,
         ins: EvaluateIns,
         timeout: Optional[float],
+        group_id: Optional[int],
     ) -> EvaluateRes:
         """Evaluate the provided parameters using the locally held dataset."""
 
@@ -79,5 +83,6 @@ class ClientProxy(ABC):
         self,
         ins: ReconnectIns,
         timeout: Optional[float],
+        group_id: Optional[int],
     ) -> DisconnectRes:
         """Disconnect and (optionally) reconnect later."""

flwr/server/compat/__init__.py

@@ -16,7 +16,9 @@
 
 
 from .app import start_driver as start_driver
+from .legacy_context import LegacyContext as LegacyContext
 
 __all__ = [
+    "LegacyContext",
     "start_driver",
 ]

flwr/server/compat/app.py

@@ -16,16 +16,13 @@
 
 
 import sys
-import threading
-import time
 from logging import INFO
 from pathlib import Path
-from typing import Dict, Optional, Union
+from typing import Optional, Union
 
 from flwr.common import EventType, event
 from flwr.common.address import parse_address
 from flwr.common.logger import log, warn_deprecated_feature
-from flwr.proto import driver_pb2  # pylint: disable=E0611
 from flwr.server.client_manager import ClientManager
 from flwr.server.history import History
 from flwr.server.server import Server, init_defaults, run_fl
@@ -33,8 +30,7 @@ from flwr.server.server_config import ServerConfig
 from flwr.server.strategy import Strategy
 
 from ..driver import Driver
-from ..driver.grpc_driver import GrpcDriver
-from .driver_client_proxy import DriverClientProxy
+from .app_utils import start_update_client_manager_thread
 
 DEFAULT_SERVER_ADDRESS_DRIVER = "[::]:9091"
 
@@ -104,11 +100,7 @@ def start_driver(  # pylint: disable=too-many-arguments, too-many-locals
     """
     event(EventType.START_DRIVER_ENTER)
 
-    if driver:
-        # pylint: disable=protected-access
-        grpc_driver, _ = driver._get_grpc_driver_and_run_id()
-        # pylint: enable=protected-access
-    else:
+    if driver is None:
         # Not passing a `Driver` object is deprecated
         warn_deprecated_feature("start_driver")
 
@@ -122,12 +114,9 @@ def start_driver(  # pylint: disable=too-many-arguments, too-many-locals
         # Create the Driver
         if isinstance(root_certificates, str):
             root_certificates = Path(root_certificates).read_bytes()
-        grpc_driver = GrpcDriver(
+        driver = Driver(
             driver_service_address=address, root_certificates=root_certificates
         )
-        grpc_driver.connect()
-
-    lock = threading.Lock()
 
     # Initialize the Driver API server and config
     initialized_server, initialized_config = init_defaults(
@@ -142,18 +131,10 @@ def start_driver(  # pylint: disable=too-many-arguments, too-many-locals
         initialized_config,
     )
 
-    f_stop = threading.Event()
     # Start the thread updating nodes
-    thread = threading.Thread(
-        target=update_client_manager,
-        args=(
-            grpc_driver,
-            initialized_server.client_manager(),
-            lock,
-            f_stop,
-        ),
+    thread, f_stop = start_update_client_manager_thread(
+        driver, initialized_server.client_manager()
     )
-    thread.start()
 
     # Start training
     hist = run_fl(
@@ -164,72 +145,10 @@ def start_driver(  # pylint: disable=too-many-arguments, too-many-locals
     f_stop.set()
 
     # Stop the Driver API server and the thread
-    with lock:
-        if driver:
-            del driver
-        else:
-            grpc_driver.disconnect()
+    del driver
 
     thread.join()
 
     event(EventType.START_SERVER_LEAVE)
 
     return hist
-
-
-def update_client_manager(
-    driver: GrpcDriver,
-    client_manager: ClientManager,
-    lock: threading.Lock,
-    f_stop: threading.Event,
-) -> None:
-    """Update the nodes list in the client manager.
-
-    This function periodically communicates with the associated driver to get all
-    node_ids. Each node_id is then converted into a `DriverClientProxy` instance
-    and stored in the `registered_nodes` dictionary with node_id as key.
-
-    New nodes will be added to the ClientManager via `client_manager.register()`,
-    and dead nodes will be removed from the ClientManager via
-    `client_manager.unregister()`.
-    """
-    # Request for run_id
-    run_id = driver.create_run(
-        driver_pb2.CreateRunRequest()  # pylint: disable=E1101
-    ).run_id
-
-    # Loop until the driver is disconnected
-    registered_nodes: Dict[int, DriverClientProxy] = {}
-    while not f_stop.is_set():
-        with lock:
-            # End the while loop if the driver is disconnected
-            if driver.stub is None:
-                break
-            get_nodes_res = driver.get_nodes(
-                req=driver_pb2.GetNodesRequest(run_id=run_id)  # pylint: disable=E1101
-            )
-        all_node_ids = {node.node_id for node in get_nodes_res.nodes}
-        dead_nodes = set(registered_nodes).difference(all_node_ids)
-        new_nodes = all_node_ids.difference(registered_nodes)
-
-        # Unregister dead nodes
-        for node_id in dead_nodes:
-            client_proxy = registered_nodes[node_id]
-            client_manager.unregister(client_proxy)
-            del registered_nodes[node_id]
-
-        # Register new nodes
-        for node_id in new_nodes:
-            client_proxy = DriverClientProxy(
-                node_id=node_id,
-                driver=driver,
-                anonymous=False,
-                run_id=run_id,
-            )
-            if client_manager.register(client_proxy):
-                registered_nodes[node_id] = client_proxy
-            else:
-                raise RuntimeError("Could not register node.")
-
-        # Sleep for 3 seconds
-        time.sleep(3)

flwr/server/compat/app_utils.py

@@ -0,0 +1,102 @@
+# 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.
+# ==============================================================================
+"""Utility functions for the `start_driver`."""
+
+
+import threading
+import time
+from typing import Dict, Tuple
+
+from ..client_manager import ClientManager
+from ..compat.driver_client_proxy import DriverClientProxy
+from ..driver import Driver
+
+
+def start_update_client_manager_thread(
+    driver: Driver,
+    client_manager: ClientManager,
+) -> 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
+    get all node_ids. Each node_id is then converted into a `DriverClientProxy`
+    instance and stored in the `registered_nodes` dictionary with node_id as key.
+
+    New nodes will be added to the ClientManager via `client_manager.register()`,
+    and dead nodes will be removed from the ClientManager via
+    `client_manager.unregister()`.
+
+    Parameters
+    ----------
+    driver : Driver
+        The Driver object to use.
+    client_manager : ClientManager
+        The ClientManager object to be updated.
+
+    Returns
+    -------
+    threading.Thread
+        A thread that updates the ClientManager and handles the stop event.
+    threading.Event
+        An event that, when set, signals the thread to stop.
+    """
+    f_stop = threading.Event()
+    thread = threading.Thread(
+        target=_update_client_manager,
+        args=(
+            driver,
+            client_manager,
+            f_stop,
+        ),
+    )
+    thread.start()
+
+    return thread, f_stop
+
+
+def _update_client_manager(
+    driver: Driver,
+    client_manager: ClientManager,
+    f_stop: threading.Event,
+) -> None:
+    """Update the nodes list in the client manager."""
+    # Loop until the driver is disconnected
+    registered_nodes: Dict[int, DriverClientProxy] = {}
+    while not f_stop.is_set():
+        all_node_ids = set(driver.get_node_ids())
+        dead_nodes = set(registered_nodes).difference(all_node_ids)
+        new_nodes = all_node_ids.difference(registered_nodes)
+
+        # Unregister dead nodes
+        for node_id in dead_nodes:
+            client_proxy = registered_nodes[node_id]
+            client_manager.unregister(client_proxy)
+            del registered_nodes[node_id]
+
+        # Register new nodes
+        for node_id in new_nodes:
+            client_proxy = DriverClientProxy(
+                node_id=node_id,
+                driver=driver.grpc_driver,  # type: ignore
+                anonymous=False,
+                run_id=driver.run_id,  # type: ignore
+            )
+            if client_manager.register(client_proxy):
+                registered_nodes[node_id] = client_proxy
+            else:
+                raise RuntimeError("Could not register node.")
+
+        # Sleep for 3 seconds
+        time.sleep(3)

flwr/server/compat/driver_client_proxy.py

@@ -47,57 +47,68 @@ class DriverClientProxy(ClientProxy):
         self.anonymous = anonymous
 
     def get_properties(
-        self, ins: common.GetPropertiesIns, timeout: Optional[float]
+        self,
+        ins: common.GetPropertiesIns,
+        timeout: Optional[float],
+        group_id: Optional[int],
     ) -> common.GetPropertiesRes:
         """Return client's properties."""
         # Ins to RecordSet
         out_recordset = compat.getpropertiesins_to_recordset(ins)
         # Fetch response
         in_recordset = self._send_receive_recordset(
-            out_recordset, MESSAGE_TYPE_GET_PROPERTIES, timeout
+            out_recordset, MESSAGE_TYPE_GET_PROPERTIES, timeout, group_id
         )
         # RecordSet to Res
         return compat.recordset_to_getpropertiesres(in_recordset)
 
     def get_parameters(
-        self, ins: common.GetParametersIns, timeout: Optional[float]
+        self,
+        ins: common.GetParametersIns,
+        timeout: Optional[float],
+        group_id: Optional[int],
     ) -> common.GetParametersRes:
         """Return the current local model parameters."""
         # Ins to RecordSet
         out_recordset = compat.getparametersins_to_recordset(ins)
         # Fetch response
         in_recordset = self._send_receive_recordset(
-            out_recordset, MESSAGE_TYPE_GET_PARAMETERS, timeout
+            out_recordset, MESSAGE_TYPE_GET_PARAMETERS, timeout, group_id
         )
         # RecordSet to Res
         return compat.recordset_to_getparametersres(in_recordset, False)
 
-    def fit(self, ins: common.FitIns, timeout: Optional[float]) -> common.FitRes:
+    def fit(
+        self, ins: common.FitIns, timeout: Optional[float], group_id: Optional[int]
+    ) -> common.FitRes:
         """Train model parameters on the locally held dataset."""
         # Ins to RecordSet
         out_recordset = compat.fitins_to_recordset(ins, keep_input=True)
         # Fetch response
         in_recordset = self._send_receive_recordset(
-            out_recordset, MESSAGE_TYPE_FIT, timeout
+            out_recordset, MESSAGE_TYPE_FIT, timeout, group_id
         )
         # RecordSet to Res
         return compat.recordset_to_fitres(in_recordset, keep_input=False)
 
     def evaluate(
-        self, ins: common.EvaluateIns, timeout: Optional[float]
+        self, ins: common.EvaluateIns, timeout: Optional[float], group_id: Optional[int]
     ) -> common.EvaluateRes:
         """Evaluate model parameters on the locally held dataset."""
         # Ins to RecordSet
         out_recordset = compat.evaluateins_to_recordset(ins, keep_input=True)
         # Fetch response
         in_recordset = self._send_receive_recordset(
-            out_recordset, MESSAGE_TYPE_EVALUATE, timeout
+            out_recordset, MESSAGE_TYPE_EVALUATE, timeout, group_id
         )
         # RecordSet to Res
         return compat.recordset_to_evaluateres(in_recordset)
 
     def reconnect(
-        self, ins: common.ReconnectIns, timeout: Optional[float]
+        self,
+        ins: common.ReconnectIns,
+        timeout: Optional[float],
+        group_id: Optional[int],
     ) -> common.DisconnectRes:
         """Disconnect and (optionally) reconnect later."""
         return common.DisconnectRes(reason="")  # Nothing to do here (yet)
@@ -107,10 +118,11 @@ class DriverClientProxy(ClientProxy):
         recordset: RecordSet,
         task_type: str,
         timeout: Optional[float],
+        group_id: Optional[int],
     ) -> RecordSet:
         task_ins = task_pb2.TaskIns(  # pylint: disable=E1101
             task_id="",
-            group_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

flwr/server/compat/legacy_context.py

@@ -0,0 +1,55 @@
+# 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.
+# ==============================================================================
+"""Legacy Context."""
+
+
+from dataclasses import dataclass
+from typing import Optional
+
+from flwr.common import Context, RecordSet
+
+from ..client_manager import ClientManager, SimpleClientManager
+from ..history import History
+from ..server_config import ServerConfig
+from ..strategy import FedAvg, Strategy
+
+
+@dataclass
+class LegacyContext(Context):
+    """Legacy Context."""
+
+    config: ServerConfig
+    strategy: Strategy
+    client_manager: ClientManager
+    history: History
+
+    def __init__(
+        self,
+        state: RecordSet,
+        config: Optional[ServerConfig] = None,
+        strategy: Optional[Strategy] = None,
+        client_manager: Optional[ClientManager] = None,
+    ) -> None:
+        if config is None:
+            config = ServerConfig()
+        if strategy is None:
+            strategy = FedAvg()
+        if client_manager is None:
+            client_manager = SimpleClientManager()
+        self.config = config
+        self.strategy = strategy
+        self.client_manager = client_manager
+        self.history = History()
+        super().__init__(state)

flwr/server/run_serverapp.py

@@ -107,7 +107,7 @@ def run_server_app() -> None:
     run(server_app_attr, driver, server_app_dir)
 
     # Clean up
-    del driver
+    driver.__del__()  # pylint: disable=unnecessary-dunder-call
 
     event(EventType.RUN_SERVER_APP_LEAVE)