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

flwr/serverapp/strategy/qfedavg.py

@@ -18,10 +18,8 @@ Paper: openreview.net/pdf?id=ByexElSYDr
 """
 
 
-from collections import OrderedDict
-from collections.abc import Iterable
+from collections.abc import Callable, Iterable
 from logging import INFO
-from typing import Callable, Optional
 
 import numpy as np
 
@@ -105,12 +103,12 @@ class QFedAvg(FedAvg):
         weighted_by_key: str = "num-examples",
         arrayrecord_key: str = "arrays",
         configrecord_key: str = "config",
-        train_metrics_aggr_fn: Optional[
-            Callable[[list[RecordDict], str], MetricRecord]
-        ] = None,
-        evaluate_metrics_aggr_fn: Optional[
-            Callable[[list[RecordDict], str], MetricRecord]
-        ] = None,
+        train_metrics_aggr_fn: (
+            Callable[[list[RecordDict], str], MetricRecord] | None
+        ) = None,
+        evaluate_metrics_aggr_fn: (
+            Callable[[list[RecordDict], str], MetricRecord] | None
+        ) = None,
     ) -> None:
         super().__init__(
             fraction_train=fraction_train,
@@ -127,7 +125,7 @@ class QFedAvg(FedAvg):
         self.q = q
         self.client_learning_rate = client_learning_rate
         self.train_loss_key = train_loss_key
-        self.current_arrays: Optional[ArrayRecord] = None
+        self.current_arrays: ArrayRecord | None = None
 
     def summary(self) -> None:
         """Log summary configuration of the strategy."""
@@ -148,7 +146,7 @@ class QFedAvg(FedAvg):
         self,
         server_round: int,
         replies: Iterable[Message],
-    ) -> tuple[Optional[ArrayRecord], Optional[MetricRecord]]:
+    ) -> tuple[ArrayRecord | None, MetricRecord | None]:
         """Aggregate ArrayRecords and MetricRecords in the received Messages."""
         # Call FedAvg aggregate_train to perform validation and aggregation
         valid_replies, _ = self._check_and_log_replies(replies, is_train=True)
@@ -184,7 +182,7 @@ class QFedAvg(FedAvg):
             if sum_delta is None:
                 sum_delta = delta
             else:
-                sum_delta = [sd + d for sd, d in zip(sum_delta, delta)]
+                sum_delta = [sd + d for sd, d in zip(sum_delta, delta, strict=True)]
             sum_h += h
 
         # Compute new global weights and convert to Array type
@@ -192,7 +190,7 @@ class QFedAvg(FedAvg):
         assert sum_delta is not None  # Make mypy happy
         array_list = [
             Array(np.asarray(gw - (d / sum_h)))
-            for gw, d in zip(global_weights, sum_delta)
+            for gw, d in zip(global_weights, sum_delta, strict=True)
         ]
 
         # Aggregate MetricRecords
@@ -200,13 +198,16 @@ class QFedAvg(FedAvg):
             [msg.content for msg in valid_replies],
             self.weighted_by_key,
         )
-        return ArrayRecord(OrderedDict(zip(array_keys, array_list))), metrics
+        return (
+            ArrayRecord(dict(zip(array_keys, array_list, strict=True))),
+            metrics,
+        )
 
 
 def get_train_loss(msg: Message, loss_key: str) -> float:
     """Extract training loss from a Message."""
     metrics = list(msg.content.metric_records.values())[0]
-    if (loss := metrics.get(loss_key)) is None or not isinstance(loss, (int, float)):
+    if (loss := metrics.get(loss_key)) is None or not isinstance(loss, (int | float)):
         raise AggregationError(
             "Missing or invalid training loss. "
             f"The strategy expected a float value for the key '{loss_key}' "
@@ -236,7 +237,7 @@ def compute_delta_and_h(
 ) -> tuple[list[NDArray], float]:
     """Compute delta and h used in q-FedAvg aggregation."""
     # Compute gradient_k = L * (w - w_k)
-    for gw, lw in zip(global_weights, local_weights):
+    for gw, lw in zip(global_weights, local_weights, strict=True):
         np.subtract(gw, lw, out=lw)
         lw *= L
     grad = local_weights  # After in-place operations, local_weights is now grad

flwr/serverapp/strategy/strategy.py

@@ -18,9 +18,8 @@
 import io
 import time
 from abc import ABC, abstractmethod
-from collections.abc import Iterable
+from collections.abc import Callable, Iterable
 from logging import INFO
-from typing import Callable, Optional
 
 from flwr.common import ArrayRecord, ConfigRecord, Message, MetricRecord, log
 from flwr.server import Grid
@@ -61,7 +60,7 @@ class Strategy(ABC):
         self,
         server_round: int,
         replies: Iterable[Message],
-    ) -> tuple[Optional[ArrayRecord], Optional[MetricRecord]]:
+    ) -> tuple[ArrayRecord | None, MetricRecord | None]:
         """Aggregate training results from client nodes.
 
         Parameters
@@ -109,7 +108,7 @@ class Strategy(ABC):
         self,
         server_round: int,
         replies: Iterable[Message],
-    ) -> Optional[MetricRecord]:
+    ) -> MetricRecord | None:
         """Aggregate evaluation metrics from client nodes.
 
         Parameters
@@ -138,11 +137,9 @@ class Strategy(ABC):
         initial_arrays: ArrayRecord,
         num_rounds: int = 3,
         timeout: float = 3600,
-        train_config: Optional[ConfigRecord] = None,
-        evaluate_config: Optional[ConfigRecord] = None,
-        evaluate_fn: Optional[
-            Callable[[int, ArrayRecord], Optional[MetricRecord]]
-        ] = None,
+        train_config: ConfigRecord | None = None,
+        evaluate_config: ConfigRecord | None = None,
+        evaluate_fn: Callable[[int, ArrayRecord], MetricRecord | None] | None = None,
     ) -> Result:
         """Execute the federated learning strategy.
 

flwr/serverapp/strategy/strategy_utils.py

@@ -17,10 +17,9 @@
 
 import json
 import random
-from collections import OrderedDict
 from logging import INFO
 from time import sleep
-from typing import Optional, cast
+from typing import cast
 
 import numpy as np
 
@@ -49,8 +48,8 @@ def config_to_str(config: ConfigRecord) -> str:
 def log_strategy_start_info(
     num_rounds: int,
     arrays: ArrayRecord,
-    train_config: Optional[ConfigRecord],
-    evaluate_config: Optional[ConfigRecord],
+    train_config: ConfigRecord | None,
+    evaluate_config: ConfigRecord | None,
 ) -> None:
     """Log information about the strategy start."""
     log(INFO, "\t├── Number of rounds: %d", num_rounds)
@@ -92,7 +91,7 @@ def aggregate_arrayrecords(
     # Perform weighted aggregation
     aggregated_np_arrays: dict[str, NDArray] = {}
 
-    for record, weight in zip(records, weight_factors):
+    for record, weight in zip(records, weight_factors, strict=True):
         for record_item in record.array_records.values():
             # aggregate in-place
             for key, value in record_item.items():
@@ -102,7 +101,7 @@ def aggregate_arrayrecords(
                     aggregated_np_arrays[key] += value.numpy() * weight
 
     return ArrayRecord(
-        OrderedDict({k: Array(np.asarray(v)) for k, v in aggregated_np_arrays.items()})
+        {k: Array(np.asarray(v)) for k, v in aggregated_np_arrays.items()}
     )
 
 
@@ -125,7 +124,7 @@ def aggregate_metricrecords(
     weight_factors = [w / total_weight for w in weights]
 
     aggregated_metrics = MetricRecord()
-    for record, weight in zip(records, weight_factors):
+    for record, weight in zip(records, weight_factors, strict=True):
         for record_item in record.metric_records.values():
             # aggregate in-place
             for key, value in record_item.items():
@@ -142,7 +141,7 @@ def aggregate_metricrecords(
                         current_list = cast(list[float], aggregated_metrics[key])
                         aggregated_metrics[key] = [
                             curr + val * weight
-                            for curr, val in zip(current_list, value)
+                            for curr, val in zip(current_list, value, strict=True)
                         ]
                     else:
                         current_value = cast(float, aggregated_metrics[key])

flwr/simulation/app.py

@@ -18,7 +18,6 @@
 import argparse
 from logging import DEBUG, ERROR, INFO
 from queue import Queue
-from typing import Optional
 
 from flwr.cli.config_utils import get_fab_metadata
 from flwr.cli.install import install_from_fab
@@ -38,8 +37,7 @@ from flwr.common.constant import (
     Status,
     SubStatus,
 )
-from flwr.common.exit import ExitCode, 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,
@@ -71,6 +69,7 @@ from flwr.server.superlink.fleet.vce.backend.backend import BackendConfig
 from flwr.simulation.run_simulation import _run_simulation
 from flwr.simulation.simulationio_connection import SimulationIoConnection
 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 SimulationExecPlugin
 from flwr.supercore.superexec.run_superexec import run_with_deprecation_warning
 
@@ -78,7 +77,7 @@ from flwr.supercore.superexec.run_superexec import run_with_deprecation_warning
 def flwr_simulation() -> None:
     """Run process-isolated Flower Simulation."""
     # 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_simulation().parse_args()
@@ -125,11 +124,11 @@ def flwr_simulation() -> None:
 
 def run_simulation_process(  # pylint: disable=R0913, R0914, R0915, R0917, W0212
     simulationio_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 Simulation process."""
     # Start monitoring the parent process if a PID is provided
@@ -141,11 +140,35 @@ def run_simulation_process(  # pylint: disable=R0913, R0914, R0915, R0917, W0212
         root_certificates=certificates,
     )
 
-    # Resolve directory where FABs are installed
+    # Initialize variables for finally block
     flwr_dir = get_flwr_dir(flwr_dir_)
     log_uploader = None
     heartbeat_sender = None
+    run = None
     run_status = None
+    exit_code = ExitCode.SUCCESS
+
+    def on_exit() -> None:
+        # Stop heartbeat sender
+        if heartbeat_sender and heartbeat_sender.is_running:
+            heartbeat_sender.stop()
+
+        # Stop log uploader for this run and upload final logs
+        if log_uploader:
+            stop_log_uploader(log_queue, log_uploader)
+
+        # Update run status
+        if run and run_status:
+            run_status_proto = run_status_to_proto(run_status)
+            conn._stub.UpdateRunStatus(
+                UpdateRunStatusRequest(run_id=run.run_id, run_status=run_status_proto)
+            )
+
+    register_signal_handlers(
+        event_type=EventType.FLWR_SIMULATION_RUN_LEAVE,
+        exit_message="Run stopped by user.",
+        exit_handlers=[on_exit],
+    )
 
     try:
         # Pull SimulationInputs from LinkState
@@ -193,12 +216,6 @@ def run_simulation_process(  # pylint: disable=R0913, R0914, R0915, R0917, W0212
             app_path,
         )
 
-        # Change status to Running
-        run_status_proto = run_status_to_proto(RunStatus(Status.RUNNING, "", ""))
-        conn._stub.UpdateRunStatus(
-            UpdateRunStatusRequest(run_id=run.run_id, run_status=run_status_proto)
-        )
-
         # Pull Federation Options
         fed_opt_res: GetFederationOptionsResponse = conn._stub.GetFederationOptions(
             GetFederationOptionsRequest(run_id=run.run_id)
@@ -216,23 +233,20 @@ def run_simulation_process(  # pylint: disable=R0913, R0914, R0915, R0917, W0212
         verbose: bool = fed_opt.get("verbose", False)
         enable_tf_gpu_growth: bool = fed_opt.get("enable_tf_gpu_growth", False)
 
+        run_id_hash = get_sha256_hash(run.run_id)
         event(
             EventType.FLWR_SIMULATION_RUN_ENTER,
             event_details={
                 "backend": "ray",
                 "num-supernodes": num_supernodes,
-                "run-id-hash": get_sha256_hash(run.run_id),
+                "run-id-hash": run_id_hash,
             },
         )
 
         # Set up heartbeat sender
-        heartbeat_fn = get_grpc_app_heartbeat_fn(
-            conn._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(conn._stub, token)
         )
-        heartbeat_sender = HeartbeatSender(heartbeat_fn)
         heartbeat_sender.start()
 
         # Launch the simulation
@@ -264,27 +278,17 @@ def run_simulation_process(  # pylint: disable=R0913, R0914, R0915, R0917, W0212
         log(ERROR, "%s raised an exception", exc_entity, exc_info=ex)
         run_status = RunStatus(Status.FINISHED, SubStatus.FAILED, str(ex))
 
-    finally:
-        # Stop heartbeat sender
-        if heartbeat_sender:
-            heartbeat_sender.stop()
-
-        # Stop log uploader for this run and upload final logs
-        if log_uploader:
-            stop_log_uploader(log_queue, log_uploader)
+        # General exit code
+        exit_code = ExitCode.SIMULATION_EXCEPTION
 
-        # Update run status
-        if run_status:
-            run_status_proto = run_status_to_proto(run_status)
-            conn._stub.UpdateRunStatus(
-                UpdateRunStatusRequest(run_id=run.run_id, run_status=run_status_proto)
-            )
-
-        # Clean up the Context if it exists
-        try:
-            del updated_context
-        except NameError:
-            pass
+    flwr_exit(
+        code=exit_code,
+        event_type=EventType.FLWR_SIMULATION_RUN_LEAVE,
+        event_details={
+            "run-id-hash": run_id_hash,
+            "success": exit_code == ExitCode.SUCCESS,
+        },
+    )
 
 
 def _parse_args_run_flwr_simulation() -> argparse.ArgumentParser:

flwr/simulation/legacy_app.py

@@ -22,7 +22,7 @@ import threading
 import traceback
 import warnings
 from logging import ERROR, INFO
-from typing import Any, Optional, Union
+from typing import Any
 
 import ray
 from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy
@@ -101,17 +101,17 @@ def start_simulation(
     *,
     client_fn: ClientFnExt,
     num_clients: int,
-    clients_ids: Optional[list[str]] = None,  # UNSUPPORTED, WILL BE REMOVED
-    client_resources: Optional[dict[str, float]] = None,
-    server: Optional[Server] = None,
-    config: Optional[ServerConfig] = None,
-    strategy: Optional[Strategy] = None,
-    client_manager: Optional[ClientManager] = None,
-    ray_init_args: Optional[dict[str, Any]] = None,
-    keep_initialised: Optional[bool] = False,
+    clients_ids: list[str] | None = None,  # UNSUPPORTED, WILL BE REMOVED
+    client_resources: dict[str, float] | None = None,
+    server: Server | None = None,
+    config: ServerConfig | None = None,
+    strategy: Strategy | None = None,
+    client_manager: ClientManager | None = None,
+    ray_init_args: dict[str, Any] | None = None,
+    keep_initialised: bool | None = False,
     actor_type: type[VirtualClientEngineActor] = ClientAppActor,
-    actor_kwargs: Optional[dict[str, Any]] = None,
-    actor_scheduling: Union[str, NodeAffinitySchedulingStrategy] = "DEFAULT",
+    actor_kwargs: dict[str, Any] | None = None,
+    actor_scheduling: str | NodeAffinitySchedulingStrategy = "DEFAULT",
 ) -> History:
     """Start a Ray-based Flower simulation server.
 
@@ -219,7 +219,7 @@ def start_simulation(
         sys.exit()
 
     # Set logger propagation
-    loop: Optional[asyncio.AbstractEventLoop] = None
+    loop: asyncio.AbstractEventLoop | None = None
     try:
         loop = asyncio.get_running_loop()
     except RuntimeError:

flwr/simulation/ray_transport/ray_actor.py

@@ -17,8 +17,9 @@
 
 import threading
 from abc import ABC
+from collections.abc import Callable
 from logging import DEBUG, ERROR, WARNING
-from typing import Any, Callable, Optional, Union
+from typing import Any
 
 import ray
 from ray import ObjectRef
@@ -76,13 +77,13 @@ class ClientAppActor(VirtualClientEngineActor):
         A function to execute upon actor initialization.
     """
 
-    def __init__(self, on_actor_init_fn: Optional[Callable[[], None]] = None) -> None:
+    def __init__(self, on_actor_init_fn: Callable[[], None] | None = None) -> None:
         super().__init__()
         if on_actor_init_fn:
             on_actor_init_fn()
 
 
-def pool_size_from_resources(client_resources: dict[str, Union[int, float]]) -> int:
+def pool_size_from_resources(client_resources: dict[str, int | float]) -> int:
     """Calculate number of Actors that fit in the cluster.
 
     For this we consider the resources available on each node and those required per
@@ -166,8 +167,8 @@ class VirtualClientEngineActorPool(ActorPool):
     def __init__(
         self,
         create_actor_fn: Callable[[], type[VirtualClientEngineActor]],
-        client_resources: dict[str, Union[int, float]],
-        actor_list: Optional[list[type[VirtualClientEngineActor]]] = None,
+        client_resources: dict[str, int | float],
+        actor_list: list[type[VirtualClientEngineActor]] | None = None,
     ):
         self.client_resources = client_resources
         self.create_actor_fn = create_actor_fn
@@ -186,9 +187,7 @@ class VirtualClientEngineActorPool(ActorPool):
 
         # A dict that maps cid to another dict containing: a reference to the remote job
         # and its status (i.e. whether it is ready or not)
-        self._cid_to_future: dict[
-            str, dict[str, Union[bool, Optional[ObjectRef[Any]]]]
-        ] = {}
+        self._cid_to_future: dict[str, dict[str, bool | ObjectRef[Any] | None]] = {}
         self.actor_to_remove: set[str] = set()  # a set
         self.num_actors = len(actors)
 
@@ -353,7 +352,7 @@ class VirtualClientEngineActorPool(ActorPool):
 
         return True
 
-    def process_unordered_future(self, timeout: Optional[float] = None) -> None:
+    def process_unordered_future(self, timeout: float | None = None) -> None:
         """Similar to parent's get_next_unordered() but without final ray.get()."""
         if not self.has_next():  # type: ignore
             raise StopIteration("No more results to get")
@@ -384,7 +383,7 @@ class VirtualClientEngineActorPool(ActorPool):
                     actor.terminate.remote()
 
     def get_client_result(
-        self, cid: str, timeout: Optional[float]
+        self, cid: str, timeout: float | None
     ) -> tuple[Message, Context]:
         """Get result from VirtualClient with specific cid."""
         # Loop until all jobs submitted to the pool are completed. Break early
@@ -407,7 +406,7 @@ class BasicActorPool:
     def __init__(
         self,
         actor_type: type[VirtualClientEngineActor],
-        client_resources: dict[str, Union[int, float]],
+        client_resources: dict[str, int | float],
         actor_kwargs: dict[str, Any],
     ):
         self.client_resources = client_resources

flwr/simulation/ray_transport/ray_client_proxy.py

@@ -17,7 +17,6 @@
 
 import traceback
 from logging import ERROR
-from typing import Optional
 
 from flwr import common
 from flwr.client import ClientFnExt
@@ -74,7 +73,7 @@ class RayActorClientProxy(ClientProxy):
             },
         )
 
-    def _submit_job(self, message: Message, timeout: Optional[float]) -> Message:
+    def _submit_job(self, message: Message, timeout: float | None) -> Message:
         """Sumbit a message to the ActorPool."""
         run_id = message.metadata.run_id
 
@@ -114,8 +113,8 @@ class RayActorClientProxy(ClientProxy):
         self,
         recorddict: RecordDict,
         message_type: str,
-        timeout: Optional[float],
-        group_id: Optional[int],
+        timeout: float | None,
+        group_id: int | None,
     ) -> Message:
         """Wrap a RecordDict inside a Message."""
         return make_message(
@@ -136,8 +135,8 @@ class RayActorClientProxy(ClientProxy):
     def get_properties(
         self,
         ins: common.GetPropertiesIns,
-        timeout: Optional[float],
-        group_id: Optional[int],
+        timeout: float | None,
+        group_id: int | None,
     ) -> common.GetPropertiesRes:
         """Return client's properties."""
         recorddict = getpropertiesins_to_recorddict(ins)
@@ -155,8 +154,8 @@ class RayActorClientProxy(ClientProxy):
     def get_parameters(
         self,
         ins: common.GetParametersIns,
-        timeout: Optional[float],
-        group_id: Optional[int],
+        timeout: float | None,
+        group_id: int | None,
     ) -> common.GetParametersRes:
         """Return the current local model parameters."""
         recorddict = getparametersins_to_recorddict(ins)
@@ -172,7 +171,7 @@ class RayActorClientProxy(ClientProxy):
         return recorddict_to_getparametersres(message_out.content, keep_input=False)
 
     def fit(
-        self, ins: common.FitIns, timeout: Optional[float], group_id: Optional[int]
+        self, ins: common.FitIns, timeout: float | None, group_id: int | None
     ) -> common.FitRes:
         """Train model parameters on the locally held dataset."""
         recorddict = fitins_to_recorddict(
@@ -190,7 +189,7 @@ class RayActorClientProxy(ClientProxy):
         return recorddict_to_fitres(message_out.content, keep_input=False)
 
     def evaluate(
-        self, ins: common.EvaluateIns, timeout: Optional[float], group_id: Optional[int]
+        self, ins: common.EvaluateIns, timeout: float | None, group_id: int | None
     ) -> common.EvaluateRes:
         """Evaluate model parameters on the locally held dataset."""
         recorddict = evaluateins_to_recorddict(
@@ -210,8 +209,8 @@ class RayActorClientProxy(ClientProxy):
     def reconnect(
         self,
         ins: common.ReconnectIns,
-        timeout: Optional[float],
-        group_id: Optional[int],
+        timeout: float | None,
+        group_id: int | None,
     ) -> common.DisconnectRes:
         """Disconnect and (optionally) reconnect later."""
         return common.DisconnectRes(reason="")  # Nothing to do here (yet)