flwr-nightly 1.13.0.dev20241021__py3-none-any.whl → 1.13.0.dev20241111__py3-none-any.whl

flwr/server/driver/grpc_driver.py

@@ -23,7 +23,7 @@ from typing import Optional, cast
 import grpc
 
 from flwr.common import DEFAULT_TTL, Message, Metadata, RecordSet
-from flwr.common.constant import DRIVER_API_DEFAULT_ADDRESS
+from flwr.common.constant import SERVERAPPIO_API_DEFAULT_ADDRESS
 from flwr.common.grpc import create_channel
 from flwr.common.logger import log
 from flwr.common.serde import (
@@ -32,7 +32,9 @@ from flwr.common.serde import (
     user_config_from_proto,
 )
 from flwr.common.typing import Run
-from flwr.proto.driver_pb2 import (  # pylint: disable=E0611
+from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
+from flwr.proto.run_pb2 import GetRunRequest, GetRunResponse  # pylint: disable=E0611
+from flwr.proto.serverappio_pb2 import (  # pylint: disable=E0611
     GetNodesRequest,
     GetNodesResponse,
     PullTaskResRequest,
@@ -40,9 +42,7 @@ from flwr.proto.driver_pb2 import (  # pylint: disable=E0611
     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.run_pb2 import GetRunRequest, GetRunResponse  # pylint: disable=E0611
+from flwr.proto.serverappio_pb2_grpc import ServerAppIoStub  # pylint: disable=E0611
 from flwr.proto.task_pb2 import TaskIns  # pylint: disable=E0611
 
 from .driver import Driver
@@ -56,14 +56,12 @@ Call `connect()` on the `GrpcDriverStub` instance before calling any of the othe
 
 
 class GrpcDriver(Driver):
-    """`GrpcDriver` provides an interface to the Driver API.
+    """`GrpcDriver` provides an interface to the ServerAppIo API.
 
     Parameters
     ----------
-    run_id : int
-        The identifier of the run.
-    driver_service_address : str (default: "[::]:9091")
-        The address (URL, IPv6, IPv4) of the SuperLink Driver API service.
+    serverappio_service_address : str (default: "[::]:9091")
+        The address (URL, IPv6, IPv4) of the SuperLink ServerAppIo API service.
     root_certificates : Optional[bytes] (default: None)
         The PEM-encoded root certificates as a byte string.
         If provided, a secure connection using the certificates will be
@@ -72,25 +70,23 @@ class GrpcDriver(Driver):
 
     def __init__(  # pylint: disable=too-many-arguments
         self,
-        run_id: int,
-        driver_service_address: str = DRIVER_API_DEFAULT_ADDRESS,
+        serverappio_service_address: str = SERVERAPPIO_API_DEFAULT_ADDRESS,
         root_certificates: Optional[bytes] = None,
     ) -> None:
-        self._run_id = run_id
-        self._addr = driver_service_address
+        self._addr = serverappio_service_address
         self._cert = root_certificates
         self._run: Optional[Run] = None
-        self._grpc_stub: Optional[DriverStub] = None
+        self._grpc_stub: Optional[ServerAppIoStub] = None
         self._channel: Optional[grpc.Channel] = None
         self.node = Node(node_id=0, anonymous=True)
 
     @property
     def _is_connected(self) -> bool:
-        """Check if connected to the Driver API server."""
+        """Check if connected to the ServerAppIo API server."""
         return self._channel is not None
 
     def _connect(self) -> None:
-        """Connect to the Driver API.
+        """Connect to the ServerAppIo API.
 
         This will not call GetRun.
         """
@@ -102,11 +98,11 @@ class GrpcDriver(Driver):
             insecure=(self._cert is None),
             root_certificates=self._cert,
         )
-        self._grpc_stub = DriverStub(self._channel)
+        self._grpc_stub = ServerAppIoStub(self._channel)
         log(DEBUG, "[Driver] Connected to %s", self._addr)
 
     def _disconnect(self) -> None:
-        """Disconnect from the Driver API."""
+        """Disconnect from the ServerAppIo API."""
         if not self._is_connected:
             log(DEBUG, "Already disconnected")
             return
@@ -116,15 +112,13 @@ class GrpcDriver(Driver):
         channel.close()
         log(DEBUG, "[Driver] Disconnected")
 
-    def _init_run(self) -> None:
-        # Check if is initialized
-        if self._run is not None:
-            return
+    def set_run(self, run_id: int) -> None:
+        """Set the run."""
         # Get the run info
-        req = GetRunRequest(run_id=self._run_id)
+        req = GetRunRequest(run_id=run_id)
         res: GetRunResponse = self._stub.GetRun(req)
         if not res.HasField("run"):
-            raise RuntimeError(f"Cannot find the run with ID: {self._run_id}")
+            raise RuntimeError(f"Cannot find the run with ID: {run_id}")
         self._run = Run(
             run_id=res.run.run_id,
             fab_id=res.run.fab_id,
@@ -136,21 +130,20 @@ class GrpcDriver(Driver):
     @property
     def run(self) -> Run:
         """Run information."""
-        self._init_run()
         return Run(**vars(self._run))
 
     @property
-    def _stub(self) -> DriverStub:
-        """Driver stub."""
+    def _stub(self) -> ServerAppIoStub:
+        """ServerAppIo stub."""
         if not self._is_connected:
             self._connect()
-        return cast(DriverStub, self._grpc_stub)
+        return cast(ServerAppIoStub, self._grpc_stub)
 
     def _check_message(self, message: Message) -> None:
         # Check if the message is valid
         if not (
             # Assume self._run being initialized
-            message.metadata.run_id == self._run_id
+            message.metadata.run_id == cast(Run, self._run).run_id
             and message.metadata.src_node_id == self.node.node_id
             and message.metadata.message_id == ""
             and message.metadata.reply_to_message == ""
@@ -171,7 +164,6 @@ class GrpcDriver(Driver):
         This method constructs a new `Message` with given content and metadata.
         The `run_id` and `src_node_id` will be set automatically.
         """
-        self._init_run()
         if ttl:
             warnings.warn(
                 "A custom TTL was set, but note that the SuperLink does not enforce "
@@ -182,7 +174,7 @@ class GrpcDriver(Driver):
 
         ttl_ = DEFAULT_TTL if ttl is None else ttl
         metadata = Metadata(
-            run_id=self._run_id,
+            run_id=cast(Run, self._run).run_id,
             message_id="",  # Will be set by the server
             src_node_id=self.node.node_id,
             dst_node_id=dst_node_id,
@@ -195,10 +187,9 @@ class GrpcDriver(Driver):
 
     def get_node_ids(self) -> list[int]:
         """Get node IDs."""
-        self._init_run()
         # Call GrpcDriverStub method
         res: GetNodesResponse = self._stub.GetNodes(
-            GetNodesRequest(run_id=self._run_id)
+            GetNodesRequest(run_id=cast(Run, self._run).run_id)
         )
         return [node.node_id for node in res.nodes]
 
@@ -208,7 +199,6 @@ class GrpcDriver(Driver):
         This method takes an iterable of messages and sends each message
         to the node specified in `dst_node_id`.
         """
-        self._init_run()
         # Construct TaskIns
         task_ins_list: list[TaskIns] = []
         for msg in messages:
@@ -230,7 +220,6 @@ class GrpcDriver(Driver):
         This method is used to collect messages from the SuperLink that correspond to a
         set of given message IDs.
         """
-        self._init_run()
         # Pull TaskRes
         res: PullTaskResResponse = self._stub.PullTaskRes(
             PullTaskResRequest(node=self.node, task_ids=message_ids)

flwr/server/driver/inmemory_driver.py

@@ -25,18 +25,16 @@ from flwr.common import DEFAULT_TTL, Message, Metadata, RecordSet
 from flwr.common.serde import message_from_taskres, message_to_taskins
 from flwr.common.typing import Run
 from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
-from flwr.server.superlink.state import StateFactory
+from flwr.server.superlink.linkstate import LinkStateFactory
 
 from .driver import Driver
 
 
 class InMemoryDriver(Driver):
-    """`InMemoryDriver` class provides an interface to the Driver API.
+    """`InMemoryDriver` class provides an interface to the ServerAppIo API.
 
     Parameters
     ----------
-    run_id : int
-        The identifier of the run.
     state_factory : StateFactory
         A StateFactory embedding a state that this driver can interface with.
     pull_interval : float (default=0.1)
@@ -45,18 +43,15 @@ class InMemoryDriver(Driver):
 
     def __init__(
         self,
-        run_id: int,
-        state_factory: StateFactory,
+        state_factory: LinkStateFactory,
         pull_interval: float = 0.1,
     ) -> None:
-        self._run_id = run_id
         self._run: Optional[Run] = None
         self.state = state_factory.state()
         self.pull_interval = pull_interval
         self.node = Node(node_id=0, anonymous=True)
 
     def _check_message(self, message: Message) -> None:
-        self._init_run()
         # Check if the message is valid
         if not (
             message.metadata.run_id == cast(Run, self._run).run_id
@@ -67,19 +62,16 @@ class InMemoryDriver(Driver):
         ):
             raise ValueError(f"Invalid message: {message}")
 
-    def _init_run(self) -> None:
+    def set_run(self, run_id: int) -> None:
         """Initialize the run."""
-        if self._run is not None:
-            return
-        run = self.state.get_run(self._run_id)
+        run = self.state.get_run(run_id)
         if run is None:
-            raise RuntimeError(f"Cannot find the run with ID: {self._run_id}")
+            raise RuntimeError(f"Cannot find the run with ID: {run_id}")
         self._run = run
 
     @property
     def run(self) -> Run:
         """Run ID."""
-        self._init_run()
         return Run(**vars(cast(Run, self._run)))
 
     def create_message(  # pylint: disable=too-many-arguments,R0917
@@ -95,7 +87,6 @@ class InMemoryDriver(Driver):
         This method constructs a new `Message` with given content and metadata.
         The `run_id` and `src_node_id` will be set automatically.
         """
-        self._init_run()
         if ttl:
             warnings.warn(
                 "A custom TTL was set, but note that the SuperLink does not enforce "
@@ -119,7 +110,6 @@ class InMemoryDriver(Driver):
 
     def get_node_ids(self) -> list[int]:
         """Get node IDs."""
-        self._init_run()
         return list(self.state.get_nodes(cast(Run, self._run).run_id))
 
     def push_messages(self, messages: Iterable[Message]) -> Iterable[str]:

flwr/server/run_serverapp.py

@@ -31,10 +31,9 @@ from flwr.common.config import (
     get_project_config,
     get_project_dir,
 )
-from flwr.common.constant import DRIVER_API_DEFAULT_ADDRESS
+from flwr.common.constant import SERVERAPPIO_API_DEFAULT_ADDRESS
 from flwr.common.logger import log, update_console_handler, warn_deprecated_feature
 from flwr.common.object_ref import load_app
-from flwr.common.typing import UserConfig
 from flwr.proto.fab_pb2 import GetFabRequest, GetFabResponse  # pylint: disable=E0611
 from flwr.proto.run_pb2 import (  # pylint: disable=E0611
     CreateRunRequest,
@@ -48,11 +47,11 @@ from .server_app import LoadServerAppError, ServerApp
 
 def run(
     driver: Driver,
+    context: Context,
     server_app_dir: str,
-    server_app_run_config: UserConfig,
     server_app_attr: Optional[str] = None,
     loaded_server_app: Optional[ServerApp] = None,
-) -> None:
+) -> Context:
     """Run ServerApp with a given Driver."""
     if not (server_app_attr is None) ^ (loaded_server_app is None):
         raise ValueError(
@@ -78,15 +77,11 @@ def run(
 
     server_app = _load()
 
-    # Initialize Context
-    context = Context(
-        node_id=0, node_config={}, state=RecordSet(), run_config=server_app_run_config
-    )
-
     # Call ServerApp
     server_app(driver=driver, context=context)
 
     log(DEBUG, "ServerApp finished running.")
+    return context
 
 
 # pylint: disable-next=too-many-branches,too-many-statements,too-many-locals
@@ -111,18 +106,18 @@ def run_server_app() -> None:
             "app by executing `flwr new` and following the prompt."
         )
 
-    if args.server != DRIVER_API_DEFAULT_ADDRESS:
+    if args.server != SERVERAPPIO_API_DEFAULT_ADDRESS:
         warn = "Passing flag --server is deprecated. Use --superlink instead."
         warn_deprecated_feature(warn)
 
-        if args.superlink != DRIVER_API_DEFAULT_ADDRESS:
+        if args.superlink != SERVERAPPIO_API_DEFAULT_ADDRESS:
             # if `--superlink` also passed, then
             # warn user that this argument overrides what was passed with `--server`
             log(
                 WARN,
                 "Both `--server` and `--superlink` were passed. "
-                "`--server` will be ignored. Connecting to the Superlink Driver API "
-                "at %s.",
+                "`--server` will be ignored. Connecting to the "
+                "SuperLink ServerAppIo API at %s.",
                 args.superlink,
             )
         else:
@@ -175,11 +170,11 @@ def run_server_app() -> None:
     if app_path is None:
         # User provided `--run-id`, but not `app_dir`
         driver = GrpcDriver(
-            run_id=args.run_id,
-            driver_service_address=args.superlink,
+            serverappio_service_address=args.superlink,
             root_certificates=root_certificates,
         )
         flwr_dir = get_flwr_dir(args.flwr_dir)
+        driver.set_run(args.run_id)
         run_ = driver.run
         if not run_.fab_hash:
             raise ValueError("FAB hash not provided.")
@@ -193,12 +188,12 @@ def run_server_app() -> None:
 
         app_path = str(get_project_dir(fab_id, fab_version, run_.fab_hash, flwr_dir))
         config = get_project_config(app_path)
+        run_id = run_.run_id
     else:
         # User provided `app_dir`, but not `--run-id`
         # Create run if run_id is not provided
         driver = GrpcDriver(
-            run_id=0,  # Will be overwritten
-            driver_service_address=args.superlink,
+            serverappio_service_address=args.superlink,
             root_certificates=root_certificates,
         )
         # Load config from the project directory
@@ -208,8 +203,9 @@ def run_server_app() -> None:
         # Create run
         req = CreateRunRequest(fab_id=fab_id, fab_version=fab_version)
         res: CreateRunResponse = driver._stub.CreateRun(req)  # pylint: disable=W0212
-        # Overwrite driver._run_id
-        driver._run_id = res.run_id  # pylint: disable=W0212
+        # Fetch full `Run` using `run_id`
+        driver.set_run(res.run_id)  # pylint: disable=W0212
+        run_id = res.run_id
 
     # Obtain server app reference and the run config
     server_app_attr = config["tool"]["flwr"]["app"]["components"]["serverapp"]
@@ -225,11 +221,20 @@ def run_server_app() -> None:
         root_certificates,
     )
 
+    # Initialize Context
+    context = Context(
+        run_id=run_id,
+        node_id=0,
+        node_config={},
+        state=RecordSet(),
+        run_config=server_app_run_config,
+    )
+
     # Run the ServerApp with the Driver
     run(
         driver=driver,
+        context=context,
         server_app_dir=app_path,
-        server_app_run_config=server_app_run_config,
         server_app_attr=server_app_attr,
     )
 
@@ -272,13 +277,14 @@ def _parse_args_run_server_app() -> argparse.ArgumentParser:
     )
     parser.add_argument(
         "--server",
-        default=DRIVER_API_DEFAULT_ADDRESS,
+        default=SERVERAPPIO_API_DEFAULT_ADDRESS,
         help="Server address",
     )
     parser.add_argument(
         "--superlink",
-        default=DRIVER_API_DEFAULT_ADDRESS,
-        help="SuperLink Driver API (gRPC-rere) address (IPv4, IPv6, or a domain name)",
+        default=SERVERAPPIO_API_DEFAULT_ADDRESS,
+        help="SuperLink ServerAppIo API (gRPC-rere) address "
+        "(IPv4, IPv6, or a domain name)",
     )
     parser.add_argument(
         "--run-id",

flwr/server/{superlink/state → serverapp}/__init__.py

@@ -12,17 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Flower server state."""
+"""Flower AppIO service."""
 
 
-from .in_memory_state import InMemoryState as InMemoryState
-from .sqlite_state import SqliteState as SqliteState
-from .state import State as State
-from .state_factory import StateFactory as StateFactory
+from .app import flwr_serverapp as flwr_serverapp
 
 __all__ = [
-    "InMemoryState",
-    "SqliteState",
-    "State",
-    "StateFactory",
+    "flwr_serverapp",
 ]

flwr/server/serverapp/app.py

@@ -0,0 +1,214 @@
+# 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.
+# ==============================================================================
+"""Flower ServerApp process."""
+
+import argparse
+from logging import DEBUG, ERROR, INFO
+from pathlib import Path
+from queue import Queue
+from time import sleep
+from typing import Optional
+
+from flwr.cli.config_utils import get_fab_metadata
+from flwr.cli.install import install_from_fab
+from flwr.common.args import add_args_flwr_app_common, try_obtain_certificates
+from flwr.common.config import (
+    get_flwr_dir,
+    get_fused_config_from_dir,
+    get_project_config,
+    get_project_dir,
+)
+from flwr.common.constant import Status, SubStatus
+from flwr.common.logger import (
+    log,
+    mirror_output_to_queue,
+    restore_output,
+    start_log_uploader,
+    stop_log_uploader,
+)
+from flwr.common.serde import (
+    context_from_proto,
+    context_to_proto,
+    fab_from_proto,
+    run_from_proto,
+    run_status_to_proto,
+)
+from flwr.common.typing import RunStatus
+from flwr.proto.run_pb2 import UpdateRunStatusRequest  # pylint: disable=E0611
+from flwr.proto.serverappio_pb2 import (  # pylint: disable=E0611
+    PullServerAppInputsRequest,
+    PullServerAppInputsResponse,
+    PushServerAppOutputsRequest,
+)
+from flwr.server.driver.grpc_driver import GrpcDriver
+from flwr.server.run_serverapp import run as run_
+
+
+def flwr_serverapp() -> None:
+    """Run process-isolated Flower ServerApp."""
+    # Capture stdout/stderr
+    log_queue: Queue[Optional[str]] = Queue()
+    mirror_output_to_queue(log_queue)
+
+    parser = argparse.ArgumentParser(
+        description="Run a Flower ServerApp",
+    )
+    parser.add_argument(
+        "--superlink",
+        type=str,
+        help="Address of SuperLink's ServerAppIo API",
+    )
+    parser.add_argument(
+        "--run-once",
+        action="store_true",
+        help="When set, this process will start a single ServerApp for a pending Run. "
+        "If there is no pending Run, the process will exit.",
+    )
+    add_args_flwr_app_common(parser=parser)
+    args = parser.parse_args()
+
+    log(INFO, "Starting Flower ServerApp")
+    certificates = try_obtain_certificates(args)
+
+    log(
+        DEBUG,
+        "Starting isolated `ServerApp` connected to SuperLink's ServerAppIo API at %s",
+        args.superlink,
+    )
+    run_serverapp(
+        superlink=args.superlink,
+        log_queue=log_queue,
+        run_once=args.run_once,
+        flwr_dir=args.flwr_dir,
+        certificates=certificates,
+    )
+
+    # Restore stdout/stderr
+    restore_output()
+
+
+def run_serverapp(  # pylint: disable=R0914, disable=W0212
+    superlink: str,
+    log_queue: Queue[Optional[str]],
+    run_once: bool,
+    flwr_dir: Optional[str] = None,
+    certificates: Optional[bytes] = None,
+) -> None:
+    """Run Flower ServerApp process."""
+    driver = GrpcDriver(
+        serverappio_service_address=superlink,
+        root_certificates=certificates,
+    )
+
+    # Resolve directory where FABs are installed
+    flwr_dir_ = get_flwr_dir(flwr_dir)
+    log_uploader = None
+
+    while True:
+
+        try:
+            # Pull ServerAppInputs from LinkState
+            req = PullServerAppInputsRequest()
+            res: PullServerAppInputsResponse = driver._stub.PullServerAppInputs(req)
+            if not res.HasField("run"):
+                sleep(3)
+                run_status = None
+                continue
+
+            context = context_from_proto(res.context)
+            run = run_from_proto(res.run)
+            fab = fab_from_proto(res.fab)
+
+            driver.set_run(run.run_id)
+
+            # Start log uploader for this run
+            log_uploader = start_log_uploader(
+                log_queue=log_queue,
+                node_id=0,
+                run_id=run.run_id,
+                stub=driver._stub,
+            )
+
+            log(DEBUG, "ServerApp process starts FAB installation.")
+            install_from_fab(fab.content, flwr_dir=flwr_dir_, skip_prompt=True)
+
+            fab_id, fab_version = get_fab_metadata(fab.content)
+
+            app_path = str(
+                get_project_dir(fab_id, fab_version, fab.hash_str, flwr_dir_)
+            )
+            config = get_project_config(app_path)
+
+            # Obtain server app reference and the run config
+            server_app_attr = config["tool"]["flwr"]["app"]["components"]["serverapp"]
+            server_app_run_config = get_fused_config_from_dir(
+                Path(app_path), run.override_config
+            )
+
+            # Update run_config in context
+            context.run_config = server_app_run_config
+
+            log(
+                DEBUG,
+                "Flower will load ServerApp `%s` in %s",
+                server_app_attr,
+                app_path,
+            )
+
+            # Change status to Running
+            run_status_proto = run_status_to_proto(RunStatus(Status.RUNNING, "", ""))
+            driver._stub.UpdateRunStatus(
+                UpdateRunStatusRequest(run_id=run.run_id, run_status=run_status_proto)
+            )
+
+            # Load and run the ServerApp with the Driver
+            updated_context = run_(
+                driver=driver,
+                server_app_dir=app_path,
+                server_app_attr=server_app_attr,
+                context=context,
+            )
+
+            # Send resulting context
+            context_proto = context_to_proto(updated_context)
+            out_req = PushServerAppOutputsRequest(
+                run_id=run.run_id, context=context_proto
+            )
+            _ = driver._stub.PushServerAppOutputs(out_req)
+
+            run_status = RunStatus(Status.FINISHED, SubStatus.COMPLETED, "")
+
+        except Exception as ex:  # pylint: disable=broad-exception-caught
+            exc_entity = "ServerApp"
+            log(ERROR, "%s raised an exception", exc_entity, exc_info=ex)
+            run_status = RunStatus(Status.FINISHED, SubStatus.FAILED, str(ex))
+
+        finally:
+            if run_status:
+                run_status_proto = run_status_to_proto(run_status)
+                driver._stub.UpdateRunStatus(
+                    UpdateRunStatusRequest(
+                        run_id=run.run_id, run_status=run_status_proto
+                    )
+                )
+
+            # Stop log uploader for this run
+            if log_uploader:
+                stop_log_uploader(log_queue, log_uploader)
+                log_uploader = None
+
+        # Stop the loop if `flwr-serverapp` is expected to process a single run
+        if run_once:
+            break

flwr/server/strategy/aggregate.py

@@ -48,12 +48,12 @@ def aggregate_inplace(results: list[tuple[ClientProxy, FitRes]]) -> NDArrays:
     num_examples_total = sum(fit_res.num_examples for (_, fit_res) in results)
 
     # Compute scaling factors for each result
-    scaling_factors = [
-        fit_res.num_examples / num_examples_total for _, fit_res in results
-    ]
+    scaling_factors = np.asarray(
+        [fit_res.num_examples / num_examples_total for _, fit_res in results]
+    )
 
     def _try_inplace(
-        x: NDArray, y: Union[NDArray, float], np_binary_op: np.ufunc
+        x: NDArray, y: Union[NDArray, np.float64], np_binary_op: np.ufunc
     ) -> NDArray:
         return (  # type: ignore[no-any-return]
             np_binary_op(x, y, out=x)

flwr/server/strategy/fedadam.py

@@ -170,8 +170,18 @@ class FedAdam(FedOpt):
             for x, y in zip(self.v_t, delta_t)
         ]
 
+        # Compute the bias-corrected learning rate, `eta_norm` for improving convergence
+        # in the early rounds of FL training. This `eta_norm` is `\alpha_t` in Kingma &
+        # Ba, 2014 (http://arxiv.org/abs/1412.6980) "Adam: A Method for Stochastic
+        # Optimization" in the formula line right before Section 2.1.
+        eta_norm = (
+            self.eta
+            * np.sqrt(1 - np.power(self.beta_2, server_round + 1.0))
+            / (1 - np.power(self.beta_1, server_round + 1.0))
+        )
+
         new_weights = [
-            x + self.eta * y / (np.sqrt(z) + self.tau)
+            x + eta_norm * y / (np.sqrt(z) + self.tau)
             for x, y, z in zip(self.current_weights, self.m_t, self.v_t)
         ]
 

flwr/server/superlink/driver/__init__.py

@@ -12,4 +12,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Flower driver service."""
+"""Flower ServerAppIo service."""