flwr 1.17.0__py3-none-any.whl → 1.19.0__py3-none-any.whl

flwr/common/secure_aggregation/crypto/shamir.py

@@ -1,4 +1,4 @@
-# Copyright 2023 Flower Labs GmbH. All Rights Reserved.
+# Copyright 2025 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.
@@ -15,61 +15,83 @@
 """Shamir's secret sharing."""
 
 
-import pickle
+import os
 from concurrent.futures import ThreadPoolExecutor
-from typing import cast
 
 from Crypto.Protocol.SecretSharing import Shamir
 from Crypto.Util.Padding import pad, unpad
 
 
 def create_shares(secret: bytes, threshold: int, num: int) -> list[bytes]:
-    """Return list of shares (bytes)."""
+    """Return a list of shares (bytes).
+
+    Shares are created from the provided secret using Shamir's secret sharing.
+    """
+    # Shamir's secret sharing requires the secret to be a multiple of 16 bytes
+    # (AES block size). Pad the secret to the next multiple of 16 bytes.
     secret_padded = pad(secret, 16)
-    secret_padded_chunk = [
-        (threshold, num, secret_padded[i : i + 16])
-        for i in range(0, len(secret_padded), 16)
-    ]
-    share_list: list[list[tuple[int, bytes]]] = [[] for _ in range(num)]
+    chunks = [secret_padded[i : i + 16] for i in range(0, len(secret_padded), 16)]
+
+    # The share list should contain shares of the secret, and each share consists of:
+    # <4 bytes of index><share of chunk1><share of chunk2>...<share of chunkN>
+    share_list: list[bytearray] = [bytearray() for _ in range(num)]
 
-    with ThreadPoolExecutor(max_workers=10) as executor:
+    # Create shares for each chunk in parallel
+    max_workers = min(len(chunks), os.cpu_count() or 1)
+    with ThreadPoolExecutor(max_workers=max_workers) as executor:
         for chunk_shares in executor.map(
-            lambda arg: _shamir_split(*arg), secret_padded_chunk
+            lambda chunk: _shamir_split(threshold, num, chunk), chunks
         ):
             for idx, share in chunk_shares:
-                # Index in `chunk_shares` starts from 1
-                share_list[idx - 1].append((idx, share))
+                # Initialize the share with the index if it is empty
+                if not share_list[idx - 1]:
+                    share_list[idx - 1] += idx.to_bytes(4, "little", signed=False)
 
-    return [pickle.dumps(shares) for shares in share_list]
+                # Append the share to the bytes
+                share_list[idx - 1] += share
+
+    return [bytes(share) for share in share_list]
 
 
 def _shamir_split(threshold: int, num: int, chunk: bytes) -> list[tuple[int, bytes]]:
+    """Create shares for a chunk using Shamir's secret sharing.
+
+    Each share is a tuple (index, share_bytes), where share_bytes is 16 bytes long.
+    """
     return Shamir.split(threshold, num, chunk, ssss=False)
 
 
-# Reconstructing secret with PyCryptodome
 def combine_shares(share_list: list[bytes]) -> bytes:
-    """Reconstruct secret from shares."""
-    unpickled_share_list: list[list[tuple[int, bytes]]] = [
-        cast(list[tuple[int, bytes]], pickle.loads(share)) for share in share_list
-    ]
+    """Reconstruct the secret from a list of shares."""
+    # Compute the number of chunks
+    # Each share contains 4 bytes of index and 16 bytes of share for each chunk
+    chunk_num = (len(share_list[0]) - 4) >> 4
 
-    chunk_num = len(unpickled_share_list[0])
     secret_padded = bytearray(0)
-    chunk_shares_list: list[list[tuple[int, bytes]]] = []
-    for i in range(chunk_num):
-        chunk_shares: list[tuple[int, bytes]] = []
-        for share in unpickled_share_list:
-            chunk_shares.append(share[i])
-        chunk_shares_list.append(chunk_shares)
-
-    with ThreadPoolExecutor(max_workers=10) as executor:
+    chunk_shares_list: list[list[tuple[int, bytes]]] = [[] for _ in range(chunk_num)]
+
+    # Split shares into chunks
+    for share in share_list:
+        # The first 4 bytes are the index
+        index = int.from_bytes(share[:4], "little", signed=False)
+        for i in range(chunk_num):
+            start = (i << 4) + 4
+            chunk_shares_list[i].append((index, share[start : start + 16]))
+
+    # Combine shares for each chunk in parallel
+    max_workers = min(chunk_num, os.cpu_count() or 1)
+    with ThreadPoolExecutor(max_workers=max_workers) as executor:
         for chunk in executor.map(_shamir_combine, chunk_shares_list):
             secret_padded += chunk
 
-    secret = unpad(secret_padded, 16)
-    return bytes(secret)
+    try:
+        secret = unpad(bytes(secret_padded), 16)
+    except ValueError:
+        # If unpadding fails, it means the shares are not valid
+        raise ValueError("Failed to combine shares") from None
+    return secret
 
 
 def _shamir_combine(shares: list[tuple[int, bytes]]) -> bytes:
+    """Reconstruct a chunk from shares using Shamir's secret sharing."""
     return Shamir.combine(shares, ssss=False)

flwr/common/secure_aggregation/crypto/symmetric_encryption.py

@@ -1,4 +1,4 @@
-# Copyright 2023 Flower Labs GmbH. All Rights Reserved.
+# Copyright 2025 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.

flwr/common/secure_aggregation/ndarrays_arithmetic.py

@@ -1,4 +1,4 @@
-# Copyright 2023 Flower Labs GmbH. All Rights Reserved.
+# Copyright 2025 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.

flwr/common/secure_aggregation/quantization.py

@@ -1,4 +1,4 @@
-# Copyright 2023 Flower Labs GmbH. All Rights Reserved.
+# Copyright 2025 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.

flwr/common/secure_aggregation/secaggplus_constants.py

@@ -1,4 +1,4 @@
-# Copyright 2023 Flower Labs GmbH. All Rights Reserved.
+# Copyright 2025 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.

flwr/common/secure_aggregation/secaggplus_utils.py

@@ -1,4 +1,4 @@
-# Copyright 2023 Flower Labs GmbH. All Rights Reserved.
+# Copyright 2025 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.

flwr/common/serde.py

@@ -1,4 +1,4 @@
-# Copyright 2020 Flower Labs GmbH. All Rights Reserved.
+# Copyright 2025 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.
@@ -16,28 +16,20 @@
 
 
 from collections import OrderedDict
-from collections.abc import MutableMapping
-from typing import Any, TypeVar, cast
-
-from google.protobuf.message import Message as GrpcMessage
+from typing import Any, cast
 
 # pylint: disable=E0611
 from flwr.proto.clientappio_pb2 import ClientAppOutputCode, ClientAppOutputStatus
-from flwr.proto.error_pb2 import Error as ProtoError
 from flwr.proto.fab_pb2 import Fab as ProtoFab
 from flwr.proto.message_pb2 import Context as ProtoContext
 from flwr.proto.message_pb2 import Message as ProtoMessage
-from flwr.proto.message_pb2 import Metadata as ProtoMetadata
 from flwr.proto.recorddict_pb2 import Array as ProtoArray
 from flwr.proto.recorddict_pb2 import ArrayRecord as ProtoArrayRecord
-from flwr.proto.recorddict_pb2 import BoolList, BytesList
 from flwr.proto.recorddict_pb2 import ConfigRecord as ProtoConfigRecord
 from flwr.proto.recorddict_pb2 import ConfigRecordValue as ProtoConfigRecordValue
-from flwr.proto.recorddict_pb2 import DoubleList
 from flwr.proto.recorddict_pb2 import MetricRecord as ProtoMetricRecord
 from flwr.proto.recorddict_pb2 import MetricRecordValue as ProtoMetricRecordValue
 from flwr.proto.recorddict_pb2 import RecordDict as ProtoRecordDict
-from flwr.proto.recorddict_pb2 import SintList, StringList, UintList
 from flwr.proto.run_pb2 import Run as ProtoRun
 from flwr.proto.run_pb2 import RunStatus as ProtoRunStatus
 from flwr.proto.transport_pb2 import (
@@ -60,8 +52,16 @@ from . import (
     RecordDict,
     typing,
 )
-from .message import Error, Message, Metadata, make_message
-from .record.typeddict import TypedDict
+from .constant import INT64_MAX_VALUE
+from .message import Message, make_message
+from .serde_utils import (
+    error_from_proto,
+    error_to_proto,
+    metadata_from_proto,
+    metadata_to_proto,
+    record_value_dict_from_proto,
+    record_value_dict_to_proto,
+)
 
 #  === Parameters message ===
 
@@ -339,7 +339,6 @@ def metrics_from_proto(proto: Any) -> typing.Metrics:
 
 
 # === Scalar messages ===
-INT64_MAX_VALUE = 9223372036854775807  # (1 << 63) - 1
 
 
 def scalar_to_proto(scalar: typing.Scalar) -> Scalar:
@@ -377,107 +376,21 @@ def scalar_from_proto(scalar_msg: Scalar) -> typing.Scalar:
 # === Record messages ===
 
 
-_type_to_field: dict[type, str] = {
-    float: "double",
-    int: "sint64",
-    bool: "bool",
-    str: "string",
-    bytes: "bytes",
-}
-_list_type_to_class_and_field: dict[type, tuple[type[GrpcMessage], str]] = {
-    float: (DoubleList, "double_list"),
-    int: (SintList, "sint_list"),
-    bool: (BoolList, "bool_list"),
-    str: (StringList, "string_list"),
-    bytes: (BytesList, "bytes_list"),
-}
-T = TypeVar("T")
-
-
-def _is_uint64(value: Any) -> bool:
-    """Check if a value is uint64."""
-    return isinstance(value, int) and value > INT64_MAX_VALUE
-
-
-def _record_value_to_proto(
-    value: Any, allowed_types: list[type], proto_class: type[T]
-) -> T:
-    """Serialize `*RecordValue` to ProtoBuf.
-
-    Note: `bool` MUST be put in the front of allowd_types if it exists.
-    """
-    arg = {}
-    for t in allowed_types:
-        # Single element
-        # Note: `isinstance(False, int) == True`.
-        if isinstance(value, t):
-            fld = _type_to_field[t]
-            if t is int and _is_uint64(value):
-                fld = "uint64"
-            arg[fld] = value
-            return proto_class(**arg)
-        # List
-        if isinstance(value, list) and all(isinstance(item, t) for item in value):
-            list_class, fld = _list_type_to_class_and_field[t]
-            # Use UintList if any element is of type `uint64`.
-            if t is int and any(_is_uint64(v) for v in value):
-                list_class, fld = UintList, "uint_list"
-            arg[fld] = list_class(vals=value)
-            return proto_class(**arg)
-    # Invalid types
-    raise TypeError(
-        f"The type of the following value is not allowed "
-        f"in '{proto_class.__name__}':\n{value}"
-    )
-
-
-def _record_value_from_proto(value_proto: GrpcMessage) -> Any:
-    """Deserialize `*RecordValue` from ProtoBuf."""
-    value_field = cast(str, value_proto.WhichOneof("value"))
-    if value_field.endswith("list"):
-        value = list(getattr(value_proto, value_field).vals)
-    else:
-        value = getattr(value_proto, value_field)
-    return value
-
-
-def _record_value_dict_to_proto(
-    value_dict: TypedDict[str, Any],
-    allowed_types: list[type],
-    value_proto_class: type[T],
-) -> dict[str, T]:
-    """Serialize the record value dict to ProtoBuf.
-
-    Note: `bool` MUST be put in the front of allowd_types if it exists.
-    """
-    # Move bool to the front
-    if bool in allowed_types and allowed_types[0] != bool:
-        allowed_types.remove(bool)
-        allowed_types.insert(0, bool)
-
-    def proto(_v: Any) -> T:
-        return _record_value_to_proto(_v, allowed_types, value_proto_class)
-
-    return {k: proto(v) for k, v in value_dict.items()}
-
-
-def _record_value_dict_from_proto(
-    value_dict_proto: MutableMapping[str, Any]
-) -> dict[str, Any]:
-    """Deserialize the record value dict from ProtoBuf."""
-    return {k: _record_value_from_proto(v) for k, v in value_dict_proto.items()}
-
-
 def array_to_proto(array: Array) -> ProtoArray:
     """Serialize Array to ProtoBuf."""
-    return ProtoArray(**vars(array))
+    return ProtoArray(
+        dtype=array.dtype,
+        shape=array.shape,
+        stype=array.stype,
+        data=array.data,
+    )
 
 
 def array_from_proto(array_proto: ProtoArray) -> Array:
     """Deserialize Array from ProtoBuf."""
     return Array(
         dtype=array_proto.dtype,
-        shape=list(array_proto.shape),
+        shape=tuple(array_proto.shape),
         stype=array_proto.stype,
         data=array_proto.data,
     )
@@ -486,8 +399,10 @@ def array_from_proto(array_proto: ProtoArray) -> Array:
 def array_record_to_proto(record: ArrayRecord) -> ProtoArrayRecord:
     """Serialize ArrayRecord to ProtoBuf."""
     return ProtoArrayRecord(
-        data_keys=record.keys(),
-        data_values=map(array_to_proto, record.values()),
+        items=[
+            ProtoArrayRecord.Item(key=k, value=array_to_proto(v))
+            for k, v in record.items()
+        ]
     )
 
 
@@ -497,7 +412,7 @@ def array_record_from_proto(
     """Deserialize ArrayRecord from ProtoBuf."""
     return ArrayRecord(
         array_dict=OrderedDict(
-            zip(record_proto.data_keys, map(array_from_proto, record_proto.data_values))
+            {item.key: array_from_proto(item.value) for item in record_proto.items}
         ),
         keep_input=False,
     )
@@ -505,17 +420,19 @@ def array_record_from_proto(
 
 def metric_record_to_proto(record: MetricRecord) -> ProtoMetricRecord:
     """Serialize MetricRecord to ProtoBuf."""
+    protos = record_value_dict_to_proto(record, [float, int], ProtoMetricRecordValue)
     return ProtoMetricRecord(
-        data=_record_value_dict_to_proto(record, [float, int], ProtoMetricRecordValue)
+        items=[ProtoMetricRecord.Item(key=k, value=v) for k, v in protos.items()]
     )
 
 
 def metric_record_from_proto(record_proto: ProtoMetricRecord) -> MetricRecord:
     """Deserialize MetricRecord from ProtoBuf."""
+    protos = {item.key: item.value for item in record_proto.items}
     return MetricRecord(
         metric_dict=cast(
             dict[str, typing.MetricRecordValues],
-            _record_value_dict_from_proto(record_proto.data),
+            record_value_dict_from_proto(protos),
         ),
         keep_input=False,
     )
@@ -523,68 +440,60 @@ def metric_record_from_proto(record_proto: ProtoMetricRecord) -> MetricRecord:
 
 def config_record_to_proto(record: ConfigRecord) -> ProtoConfigRecord:
     """Serialize ConfigRecord to ProtoBuf."""
+    protos = record_value_dict_to_proto(
+        record,
+        [bool, int, float, str, bytes],
+        ProtoConfigRecordValue,
+    )
     return ProtoConfigRecord(
-        data=_record_value_dict_to_proto(
-            record,
-            [bool, int, float, str, bytes],
-            ProtoConfigRecordValue,
-        )
+        items=[ProtoConfigRecord.Item(key=k, value=v) for k, v in protos.items()]
     )
 
 
 def config_record_from_proto(record_proto: ProtoConfigRecord) -> ConfigRecord:
     """Deserialize ConfigRecord from ProtoBuf."""
+    protos = {item.key: item.value for item in record_proto.items}
     return ConfigRecord(
         config_dict=cast(
             dict[str, typing.ConfigRecordValues],
-            _record_value_dict_from_proto(record_proto.data),
+            record_value_dict_from_proto(protos),
         ),
         keep_input=False,
     )
 
 
-# === Error message ===
-
-
-def error_to_proto(error: Error) -> ProtoError:
-    """Serialize Error to ProtoBuf."""
-    reason = error.reason if error.reason else ""
-    return ProtoError(code=error.code, reason=reason)
-
-
-def error_from_proto(error_proto: ProtoError) -> Error:
-    """Deserialize Error from ProtoBuf."""
-    reason = error_proto.reason if len(error_proto.reason) > 0 else None
-    return Error(code=error_proto.code, reason=reason)
-
-
 # === RecordDict message ===
 
 
 def recorddict_to_proto(recorddict: RecordDict) -> ProtoRecordDict:
     """Serialize RecordDict to ProtoBuf."""
-    return ProtoRecordDict(
-        arrays={
-            k: array_record_to_proto(v) for k, v in recorddict.array_records.items()
-        },
-        metrics={
-            k: metric_record_to_proto(v) for k, v in recorddict.metric_records.items()
-        },
-        configs={
-            k: config_record_to_proto(v) for k, v in recorddict.config_records.items()
-        },
-    )
+    item_cls = ProtoRecordDict.Item
+    items: list[ProtoRecordDict.Item] = []
+    for k, v in recorddict.items():
+        if isinstance(v, ArrayRecord):
+            items += [item_cls(key=k, array_record=array_record_to_proto(v))]
+        elif isinstance(v, MetricRecord):
+            items += [item_cls(key=k, metric_record=metric_record_to_proto(v))]
+        elif isinstance(v, ConfigRecord):
+            items += [item_cls(key=k, config_record=config_record_to_proto(v))]
+        else:
+            raise ValueError(f"Unsupported record type: {type(v)}")
+    return ProtoRecordDict(items=items)
 
 
 def recorddict_from_proto(recorddict_proto: ProtoRecordDict) -> RecordDict:
     """Deserialize RecordDict from ProtoBuf."""
     ret = RecordDict()
-    for k, arr_record_proto in recorddict_proto.arrays.items():
-        ret[k] = array_record_from_proto(arr_record_proto)
-    for k, m_record_proto in recorddict_proto.metrics.items():
-        ret[k] = metric_record_from_proto(m_record_proto)
-    for k, c_record_proto in recorddict_proto.configs.items():
-        ret[k] = config_record_from_proto(c_record_proto)
+    for item in recorddict_proto.items:
+        field = item.WhichOneof("value")
+        if field == "array_record":
+            ret[item.key] = array_record_from_proto(item.array_record)
+        elif field == "metric_record":
+            ret[item.key] = metric_record_from_proto(item.metric_record)
+        elif field == "config_record":
+            ret[item.key] = config_record_from_proto(item.config_record)
+        else:
+            raise ValueError(f"Unsupported record type: {field}")
     return ret
 
 
@@ -646,41 +555,6 @@ def user_config_value_from_proto(scalar_msg: Scalar) -> typing.UserConfigValue:
     return cast(typing.UserConfigValue, scalar)
 
 
-# === Metadata messages ===
-
-
-def metadata_to_proto(metadata: Metadata) -> ProtoMetadata:
-    """Serialize `Metadata` to ProtoBuf."""
-    proto = ProtoMetadata(  # pylint: disable=E1101
-        run_id=metadata.run_id,
-        message_id=metadata.message_id,
-        src_node_id=metadata.src_node_id,
-        dst_node_id=metadata.dst_node_id,
-        reply_to_message_id=metadata.reply_to_message_id,
-        group_id=metadata.group_id,
-        ttl=metadata.ttl,
-        message_type=metadata.message_type,
-        created_at=metadata.created_at,
-    )
-    return proto
-
-
-def metadata_from_proto(metadata_proto: ProtoMetadata) -> Metadata:
-    """Deserialize `Metadata` from ProtoBuf."""
-    metadata = Metadata(
-        run_id=metadata_proto.run_id,
-        message_id=metadata_proto.message_id,
-        src_node_id=metadata_proto.src_node_id,
-        dst_node_id=metadata_proto.dst_node_id,
-        reply_to_message_id=metadata_proto.reply_to_message_id,
-        group_id=metadata_proto.group_id,
-        created_at=metadata_proto.created_at,
-        ttl=metadata_proto.ttl,
-        message_type=metadata_proto.message_type,
-    )
-    return metadata
-
-
 # === Message messages ===
 
 
@@ -756,6 +630,7 @@ def run_to_proto(run: typing.Run) -> ProtoRun:
         running_at=run.running_at,
         finished_at=run.finished_at,
         status=run_status_to_proto(run.status),
+        flwr_aid=run.flwr_aid,
     )
     return proto
 
@@ -773,6 +648,7 @@ def run_from_proto(run_proto: ProtoRun) -> typing.Run:
         running_at=run_proto.running_at,
         finished_at=run_proto.finished_at,
         status=run_status_from_proto(run_proto.status),
+        flwr_aid=run_proto.flwr_aid,
     )
     return run