flwr 1.23.0__py3-none-any.whl → 1.25.0__py3-none-any.whl

flwr/server/server.py

@@ -19,7 +19,6 @@ import concurrent.futures
 import io
 import timeit
 from logging import INFO, WARN
-from typing import Optional, Union
 
 from flwr.common import (
     Code,
@@ -43,15 +42,15 @@ from .server_config import ServerConfig
 
 FitResultsAndFailures = tuple[
     list[tuple[ClientProxy, FitRes]],
-    list[Union[tuple[ClientProxy, FitRes], BaseException]],
+    list[tuple[ClientProxy, FitRes] | BaseException],
 ]
 EvaluateResultsAndFailures = tuple[
     list[tuple[ClientProxy, EvaluateRes]],
-    list[Union[tuple[ClientProxy, EvaluateRes], BaseException]],
+    list[tuple[ClientProxy, EvaluateRes] | BaseException],
 ]
 ReconnectResultsAndFailures = tuple[
     list[tuple[ClientProxy, DisconnectRes]],
-    list[Union[tuple[ClientProxy, DisconnectRes], BaseException]],
+    list[tuple[ClientProxy, DisconnectRes] | BaseException],
 ]
 
 
@@ -62,16 +61,16 @@ class Server:
         self,
         *,
         client_manager: ClientManager,
-        strategy: Optional[Strategy] = None,
+        strategy: Strategy | None = None,
     ) -> None:
         self._client_manager: ClientManager = client_manager
         self.parameters: Parameters = Parameters(
             tensors=[], tensor_type="numpy.ndarray"
         )
         self.strategy: Strategy = strategy if strategy is not None else FedAvg()
-        self.max_workers: Optional[int] = None
+        self.max_workers: int | None = None
 
-    def set_max_workers(self, max_workers: Optional[int]) -> None:
+    def set_max_workers(self, max_workers: int | None) -> None:
         """Set the max_workers used by ThreadPoolExecutor."""
         self.max_workers = max_workers
 
@@ -84,7 +83,7 @@ class Server:
         return self._client_manager
 
     # pylint: disable=too-many-locals
-    def fit(self, num_rounds: int, timeout: Optional[float]) -> tuple[History, float]:
+    def fit(self, num_rounds: int, timeout: float | None) -> tuple[History, float]:
         """Run federated averaging for a number of rounds."""
         history = History()
 
@@ -161,10 +160,8 @@ class Server:
     def evaluate_round(
         self,
         server_round: int,
-        timeout: Optional[float],
-    ) -> Optional[
-        tuple[Optional[float], dict[str, Scalar], EvaluateResultsAndFailures]
-    ]:
+        timeout: float | None,
+    ) -> tuple[float | None, dict[str, Scalar], EvaluateResultsAndFailures] | None:
         """Validate current global model on a number of clients."""
         # Get clients and their respective instructions from strategy
         client_instructions = self.strategy.configure_evaluate(
@@ -198,7 +195,7 @@ class Server:
 
         # Aggregate the evaluation results
         aggregated_result: tuple[
-            Optional[float],
+            float | None,
             dict[str, Scalar],
         ] = self.strategy.aggregate_evaluate(server_round, results, failures)
 
@@ -208,10 +205,8 @@ class Server:
     def fit_round(
         self,
         server_round: int,
-        timeout: Optional[float],
-    ) -> Optional[
-        tuple[Optional[Parameters], dict[str, Scalar], FitResultsAndFailures]
-    ]:
+        timeout: float | None,
+    ) -> tuple[Parameters | None, dict[str, Scalar], FitResultsAndFailures] | None:
         """Perform a single round of federated averaging."""
         # Get clients and their respective instructions from strategy
         client_instructions = self.strategy.configure_fit(
@@ -246,14 +241,14 @@ class Server:
 
         # Aggregate training results
         aggregated_result: tuple[
-            Optional[Parameters],
+            Parameters | None,
             dict[str, Scalar],
         ] = self.strategy.aggregate_fit(server_round, results, failures)
 
         parameters_aggregated, metrics_aggregated = aggregated_result
         return parameters_aggregated, metrics_aggregated, (results, failures)
 
-    def disconnect_all_clients(self, timeout: Optional[float]) -> None:
+    def disconnect_all_clients(self, timeout: float | None) -> None:
         """Send shutdown signal to all clients."""
         all_clients = self._client_manager.all()
         clients = [all_clients[k] for k in all_clients.keys()]
@@ -266,11 +261,11 @@ class Server:
         )
 
     def _get_initial_parameters(
-        self, server_round: int, timeout: Optional[float]
+        self, server_round: int, timeout: float | None
     ) -> Parameters:
         """Get initial parameters from one of the available clients."""
         # Server-side parameter initialization
-        parameters: Optional[Parameters] = self.strategy.initialize_parameters(
+        parameters: Parameters | None = self.strategy.initialize_parameters(
             client_manager=self._client_manager
         )
         if parameters is not None:
@@ -297,8 +292,8 @@ class Server:
 
 def reconnect_clients(
     client_instructions: list[tuple[ClientProxy, ReconnectIns]],
-    max_workers: Optional[int],
-    timeout: Optional[float],
+    max_workers: int | None,
+    timeout: float | None,
 ) -> ReconnectResultsAndFailures:
     """Instruct clients to disconnect and never reconnect."""
     with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
@@ -313,7 +308,7 @@ def reconnect_clients(
 
     # Gather results
     results: list[tuple[ClientProxy, DisconnectRes]] = []
-    failures: list[Union[tuple[ClientProxy, DisconnectRes], BaseException]] = []
+    failures: list[tuple[ClientProxy, DisconnectRes] | BaseException] = []
     for future in finished_fs:
         failure = future.exception()
         if failure is not None:
@@ -327,7 +322,7 @@ def reconnect_clients(
 def reconnect_client(
     client: ClientProxy,
     reconnect: ReconnectIns,
-    timeout: Optional[float],
+    timeout: float | None,
 ) -> tuple[ClientProxy, DisconnectRes]:
     """Instruct client to disconnect and (optionally) reconnect later."""
     disconnect = client.reconnect(
@@ -340,8 +335,8 @@ def reconnect_client(
 
 def fit_clients(
     client_instructions: list[tuple[ClientProxy, FitIns]],
-    max_workers: Optional[int],
-    timeout: Optional[float],
+    max_workers: int | None,
+    timeout: float | None,
     group_id: int,
 ) -> FitResultsAndFailures:
     """Refine parameters concurrently on all selected clients."""
@@ -357,7 +352,7 @@ def fit_clients(
 
     # Gather results
     results: list[tuple[ClientProxy, FitRes]] = []
-    failures: list[Union[tuple[ClientProxy, FitRes], BaseException]] = []
+    failures: list[tuple[ClientProxy, FitRes] | BaseException] = []
     for future in finished_fs:
         _handle_finished_future_after_fit(
             future=future, results=results, failures=failures
@@ -366,7 +361,7 @@ def fit_clients(
 
 
 def fit_client(
-    client: ClientProxy, ins: FitIns, timeout: Optional[float], group_id: int
+    client: ClientProxy, ins: FitIns, timeout: float | None, group_id: int
 ) -> tuple[ClientProxy, FitRes]:
     """Refine parameters on a single client."""
     fit_res = client.fit(ins, timeout=timeout, group_id=group_id)
@@ -376,7 +371,7 @@ def fit_client(
 def _handle_finished_future_after_fit(
     future: concurrent.futures.Future,  # type: ignore
     results: list[tuple[ClientProxy, FitRes]],
-    failures: list[Union[tuple[ClientProxy, FitRes], BaseException]],
+    failures: list[tuple[ClientProxy, FitRes] | BaseException],
 ) -> None:
     """Convert finished future into either a result or a failure."""
     # Check if there was an exception
@@ -400,8 +395,8 @@ def _handle_finished_future_after_fit(
 
 def evaluate_clients(
     client_instructions: list[tuple[ClientProxy, EvaluateIns]],
-    max_workers: Optional[int],
-    timeout: Optional[float],
+    max_workers: int | None,
+    timeout: float | None,
     group_id: int,
 ) -> EvaluateResultsAndFailures:
     """Evaluate parameters concurrently on all selected clients."""
@@ -417,7 +412,7 @@ def evaluate_clients(
 
     # Gather results
     results: list[tuple[ClientProxy, EvaluateRes]] = []
-    failures: list[Union[tuple[ClientProxy, EvaluateRes], BaseException]] = []
+    failures: list[tuple[ClientProxy, EvaluateRes] | BaseException] = []
     for future in finished_fs:
         _handle_finished_future_after_evaluate(
             future=future, results=results, failures=failures
@@ -428,7 +423,7 @@ def evaluate_clients(
 def evaluate_client(
     client: ClientProxy,
     ins: EvaluateIns,
-    timeout: Optional[float],
+    timeout: float | None,
     group_id: int,
 ) -> tuple[ClientProxy, EvaluateRes]:
     """Evaluate parameters on a single client."""
@@ -439,7 +434,7 @@ def evaluate_client(
 def _handle_finished_future_after_evaluate(
     future: concurrent.futures.Future,  # type: ignore
     results: list[tuple[ClientProxy, EvaluateRes]],
-    failures: list[Union[tuple[ClientProxy, EvaluateRes], BaseException]],
+    failures: list[tuple[ClientProxy, EvaluateRes] | BaseException],
 ) -> None:
     """Convert finished future into either a result or a failure."""
     # Check if there was an exception
@@ -462,10 +457,10 @@ def _handle_finished_future_after_evaluate(
 
 
 def init_defaults(
-    server: Optional[Server],
-    config: Optional[ServerConfig],
-    strategy: Optional[Strategy],
-    client_manager: Optional[ClientManager],
+    server: Server | None,
+    config: ServerConfig | None,
+    strategy: Strategy | None,
+    client_manager: ClientManager | None,
 ) -> tuple[Server, ServerConfig]:
     """Create server instance if none was given."""
     if server is None:

flwr/server/server_app.py

@@ -16,9 +16,8 @@
 
 
 import inspect
-from collections.abc import Iterator
+from collections.abc import Callable, Iterator
 from contextlib import contextmanager
-from typing import Callable, Optional
 
 from flwr.common import Context
 from flwr.common.logger import warn_deprecated_feature_with_example
@@ -118,11 +117,11 @@ class ServerApp:  # pylint: disable=too-many-instance-attributes
     # pylint: disable=too-many-arguments,too-many-positional-arguments
     def __init__(
         self,
-        server: Optional[Server] = None,
-        config: Optional[ServerConfig] = None,
-        strategy: Optional[Strategy] = None,
-        client_manager: Optional[ClientManager] = None,
-        server_fn: Optional[ServerFn] = None,
+        server: Server | None = None,
+        config: ServerConfig | None = None,
+        strategy: Strategy | None = None,
+        client_manager: ClientManager | None = None,
+        server_fn: ServerFn | None = None,
     ) -> None:
         if any([server, config, strategy, client_manager]):
             warn_deprecated_feature_with_example(
@@ -148,7 +147,7 @@ class ServerApp:  # pylint: disable=too-many-instance-attributes
         self._strategy = strategy
         self._client_manager = client_manager
         self._server_fn = server_fn
-        self._main: Optional[ServerAppCallable] = None
+        self._main: ServerAppCallable | None = None
         self._lifespan = _empty_lifespan
 
     def __call__(self, grid: Grid, context: Context) -> None:

flwr/server/server_config.py

@@ -16,7 +16,6 @@
 
 
 from dataclasses import dataclass
-from typing import Optional
 
 
 @dataclass
@@ -28,7 +27,7 @@ class ServerConfig:
     """
 
     num_rounds: int = 1
-    round_timeout: Optional[float] = None
+    round_timeout: float | None = None
 
     def __repr__(self) -> str:
         """Return the string representation of the ServerConfig."""

flwr/server/serverapp/app.py

@@ -19,7 +19,8 @@ import argparse
 from logging import DEBUG, ERROR, INFO
 from pathlib import Path
 from queue import Queue
-from typing import Optional
+
+import grpc
 
 from flwr.app.exception import AppExitException
 from flwr.cli.config_utils import get_fab_metadata
@@ -38,8 +39,7 @@ from flwr.common.constant import (
     Status,
     SubStatus,
 )
-from flwr.common.exit import ExitCode, add_exit_handler, flwr_exit
-from flwr.common.heartbeat import HeartbeatSender, get_grpc_app_heartbeat_fn
+from flwr.common.exit import ExitCode, flwr_exit, register_signal_handlers
 from flwr.common.logger import (
     log,
     mirror_output_to_queue,
@@ -66,6 +66,7 @@ from flwr.proto.serverappio_pb2_grpc import ServerAppIoStub
 from flwr.server.grid.grpc_grid import GrpcGrid
 from flwr.server.run_serverapp import run as run_
 from flwr.supercore.app_utils import start_parent_process_monitor
+from flwr.supercore.heartbeat import HeartbeatSender, make_app_heartbeat_fn_grpc
 from flwr.supercore.superexec.plugin import ServerAppExecPlugin
 from flwr.supercore.superexec.run_superexec import run_with_deprecation_warning
 
@@ -73,7 +74,7 @@ from flwr.supercore.superexec.run_superexec import run_with_deprecation_warning
 def flwr_serverapp() -> None:
     """Run process-isolated Flower ServerApp."""
     # Capture stdout/stderr
-    log_queue: Queue[Optional[str]] = Queue()
+    log_queue: Queue[str | None] = Queue()
     mirror_output_to_queue(log_queue)
 
     args = _parse_args_run_flwr_serverapp().parse_args()
@@ -120,21 +121,22 @@ def flwr_serverapp() -> None:
 
 def run_serverapp(  # pylint: disable=R0913, R0914, R0915, R0917, W0212
     serverappio_api_address: str,
-    log_queue: Queue[Optional[str]],
+    log_queue: Queue[str | None],
     token: str,
-    flwr_dir: Optional[str] = None,
-    certificates: Optional[bytes] = None,
-    parent_pid: Optional[int] = None,
+    flwr_dir: str | None = None,
+    certificates: bytes | None = None,
+    parent_pid: int | None = None,
 ) -> None:
     """Run Flower ServerApp process."""
     # Monitor the main process in case of SIGKILL
     if parent_pid is not None:
         start_parent_process_monitor(parent_pid)
 
-    # Resolve directory where FABs are installed
+    # Initialize variables for exit handler
     flwr_dir_ = get_flwr_dir(flwr_dir)
     log_uploader = None
     hash_run_id = None
+    run = None
     run_status = None
     heartbeat_sender = None
     grid = None
@@ -143,7 +145,7 @@ def run_serverapp(  # pylint: disable=R0913, R0914, R0915, R0917, W0212
 
     def on_exit() -> None:
         # Stop heartbeat sender
-        if heartbeat_sender:
+        if heartbeat_sender and heartbeat_sender.is_running:
             heartbeat_sender.stop()
 
         # Stop log uploader for this run and upload final logs
@@ -151,7 +153,7 @@ def run_serverapp(  # pylint: disable=R0913, R0914, R0915, R0917, W0212
             stop_log_uploader(log_queue, log_uploader)
 
         # Update run status
-        if run_status and grid:
+        if run and run_status and grid:
             run_status_proto = run_status_to_proto(run_status)
             grid._stub.UpdateRunStatus(
                 UpdateRunStatusRequest(run_id=run.run_id, run_status=run_status_proto)
@@ -161,7 +163,12 @@ def run_serverapp(  # pylint: disable=R0913, R0914, R0915, R0917, W0212
         if grid:
             grid.close()
 
-    add_exit_handler(on_exit)
+    # Register signal handlers for graceful shutdown
+    register_signal_handlers(
+        event_type=EventType.FLWR_SERVERAPP_RUN_LEAVE,
+        exit_message="Run stopped by user.",
+        exit_handlers=[on_exit],
+    )
 
     try:
         # Initialize the GrpcGrid
@@ -171,9 +178,14 @@ def run_serverapp(  # pylint: disable=R0913, R0914, R0915, R0917, W0212
         )
 
         # Pull ServerAppInputs from LinkState
-        req = PullAppInputsRequest(token=token)
-        log(DEBUG, "[flwr-serverapp] Pull ServerAppInputs")
-        res: PullAppInputsResponse = grid._stub.PullAppInputs(req)
+        try:
+            log(DEBUG, "[flwr-serverapp] Pull ServerAppInputs")
+            req = PullAppInputsRequest(token=token)
+            res: PullAppInputsResponse = grid._stub.PullAppInputs(req)
+        except grpc.RpcError as ex:
+            if ex.code() == grpc.StatusCode.FAILED_PRECONDITION:
+                raise RuntimeError("Failed to start the run.") from ex
+            raise
         context = context_from_proto(res.context)
         run = run_from_proto(res.run)
         fab = fab_from_proto(res.fab)
@@ -214,25 +226,15 @@ def run_serverapp(  # pylint: disable=R0913, R0914, R0915, R0917, W0212
             app_path,
         )
 
-        # Change status to Running
-        run_status_proto = run_status_to_proto(RunStatus(Status.RUNNING, "", ""))
-        grid._stub.UpdateRunStatus(
-            UpdateRunStatusRequest(run_id=run.run_id, run_status=run_status_proto)
-        )
-
         event(
             EventType.FLWR_SERVERAPP_RUN_ENTER,
             event_details={"run-id-hash": hash_run_id},
         )
 
         # Set up heartbeat sender
-        heartbeat_fn = get_grpc_app_heartbeat_fn(
-            grid._stub,
-            run.run_id,
-            failure_message="Heartbeat failed unexpectedly. The SuperLink could "
-            "not find the provided run ID, or the run status is invalid.",
+        heartbeat_sender = HeartbeatSender(
+            make_app_heartbeat_fn_grpc(grid._stub, token)
         )
-        heartbeat_sender = HeartbeatSender(heartbeat_fn)
         heartbeat_sender.start()
 
         # Load and run the ServerApp with the Grid
@@ -256,11 +258,15 @@ def run_serverapp(  # pylint: disable=R0913, R0914, R0915, R0917, W0212
     # Raised when the run is already stopped by the user
     except RunNotRunningException:
         log(INFO, "")
-        log(INFO, "Run ID %s stopped.", run.run_id)
+        log(INFO, "Run ID %s stopped.", run.run_id)  # type: ignore[union-attr]
         log(INFO, "")
         run_status = None
         # No need to update the exit code since this is expected behavior
 
+    except RuntimeError:
+        log(ERROR, "Failed to start run.")
+        exit_code = ExitCode.SERVERAPP_RUN_START_REJECTED
+
     except Exception as ex:  # pylint: disable=broad-exception-caught
         exc_entity = "ServerApp"
         log(ERROR, "%s raised an exception", exc_entity, exc_info=ex)

flwr/server/serverapp_components.py

@@ -16,7 +16,6 @@
 
 
 from dataclasses import dataclass
-from typing import Optional
 
 from .client_manager import ClientManager
 from .server import Server
@@ -46,7 +45,7 @@ class ServerAppComponents:  # pylint: disable=too-many-instance-attributes
         will be used.
     """
 
-    server: Optional[Server] = None
-    config: Optional[ServerConfig] = None
-    strategy: Optional[Strategy] = None
-    client_manager: Optional[ClientManager] = None
+    server: Server | None = None
+    config: ServerConfig | None = None
+    strategy: Strategy | None = None
+    client_manager: ClientManager | None = None

flwr/server/strategy/aggregate.py

@@ -15,8 +15,9 @@
 """Aggregation functions for strategy implementations."""
 # mypy: disallow_untyped_calls=False
 
+from collections.abc import Callable
 from functools import partial, reduce
-from typing import Any, Callable, Union
+from typing import Any
 
 import numpy as np
 
@@ -37,7 +38,7 @@ def aggregate(results: list[tuple[NDArrays, int]]) -> NDArrays:
     # Compute average weights of each layer
     weights_prime: NDArrays = [
         reduce(np.add, layer_updates) / num_examples_total
-        for layer_updates in zip(*weighted_weights)
+        for layer_updates in zip(*weighted_weights, strict=True)
     ]
     return weights_prime
 
@@ -53,7 +54,7 @@ def aggregate_inplace(results: list[tuple[ClientProxy, FitRes]]) -> NDArrays:
     )
 
     def _try_inplace(
-        x: NDArray, y: Union[NDArray, np.float64], np_binary_op: np.ufunc
+        x: NDArray, y: NDArray | np.float64, np_binary_op: np.ufunc
     ) -> NDArray:
         return (  # type: ignore[no-any-return]
             np_binary_op(x, y, out=x)
@@ -75,7 +76,7 @@ def aggregate_inplace(results: list[tuple[ClientProxy, FitRes]]) -> NDArrays:
         )
         params = [
             reduce(partial(_try_inplace, np_binary_op=np.add), layer_updates)
-            for layer_updates in zip(params, res)
+            for layer_updates in zip(params, res, strict=True)
         ]
 
     return params
@@ -88,7 +89,7 @@ def aggregate_median(results: list[tuple[NDArrays, int]]) -> NDArrays:
 
     # Compute median weight of each layer
     median_w: NDArrays = [
-        np.median(np.asarray(layer), axis=0) for layer in zip(*weights)
+        np.median(np.asarray(layer), axis=0) for layer in zip(*weights, strict=True)
     ]
     return median_w
 
@@ -235,7 +236,7 @@ def aggregate_qffl(
         for j in range(1, len(deltas)):
             tmp += scaled_deltas[j][i]
         updates.append(tmp)
-    new_parameters = [(u - v) * 1.0 for u, v in zip(parameters, updates)]
+    new_parameters = [(u - v) * 1.0 for u, v in zip(parameters, updates, strict=True)]
     return new_parameters
 
 
@@ -287,7 +288,7 @@ def aggregate_trimmed_avg(
 
     trimmed_w: NDArrays = [
         _trim_mean(np.asarray(layer), proportiontocut=proportiontocut)
-        for layer in zip(*weights)
+        for layer in zip(*weights, strict=True)
     ]
 
     return trimmed_w
@@ -299,7 +300,7 @@ def _check_weights_equality(weights1: NDArrays, weights2: NDArrays) -> bool:
         return False
     return all(
         np.array_equal(layer_weights1, layer_weights2)
-        for layer_weights1, layer_weights2 in zip(weights1, weights2)
+        for layer_weights1, layer_weights2 in zip(weights1, weights2, strict=True)
     )
 
 

flwr/server/strategy/bulyan.py

@@ -18,8 +18,9 @@ Paper: arxiv.org/abs/1802.07927
 """
 
 
+from collections.abc import Callable
 from logging import WARNING
-from typing import Any, Callable, Optional, Union
+from typing import Any
 
 from flwr.common import (
     FitRes,
@@ -84,18 +85,19 @@ class Bulyan(FedAvg):
         min_evaluate_clients: int = 2,
         min_available_clients: int = 2,
         num_malicious_clients: int = 0,
-        evaluate_fn: Optional[
+        evaluate_fn: (
             Callable[
                 [int, NDArrays, dict[str, Scalar]],
-                Optional[tuple[float, dict[str, Scalar]]],
+                tuple[float, dict[str, Scalar]] | None,
             ]
-        ] = None,
-        on_fit_config_fn: Optional[Callable[[int], dict[str, Scalar]]] = None,
-        on_evaluate_config_fn: Optional[Callable[[int], dict[str, Scalar]]] = None,
+            | None
+        ) = None,
+        on_fit_config_fn: Callable[[int], dict[str, Scalar]] | None = None,
+        on_evaluate_config_fn: Callable[[int], dict[str, Scalar]] | None = None,
         accept_failures: bool = True,
-        initial_parameters: Optional[Parameters] = None,
-        fit_metrics_aggregation_fn: Optional[MetricsAggregationFn] = None,
-        evaluate_metrics_aggregation_fn: Optional[MetricsAggregationFn] = None,
+        initial_parameters: Parameters | None = None,
+        fit_metrics_aggregation_fn: MetricsAggregationFn | None = None,
+        evaluate_metrics_aggregation_fn: MetricsAggregationFn | None = None,
         first_aggregation_rule: Callable = aggregate_krum,  # type: ignore
         **aggregation_rule_kwargs: Any,
     ) -> None:
@@ -126,8 +128,8 @@ class Bulyan(FedAvg):
         self,
         server_round: int,
         results: list[tuple[ClientProxy, FitRes]],
-        failures: list[Union[tuple[ClientProxy, FitRes], BaseException]],
-    ) -> tuple[Optional[Parameters], dict[str, Scalar]]:
+        failures: list[tuple[ClientProxy, FitRes] | BaseException],
+    ) -> tuple[Parameters | None, dict[str, Scalar]]:
         """Aggregate fit results using Bulyan."""
         if not results:
             return None, {}