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

flwr/client/rest_client/connection.py

@@ -36,18 +36,27 @@ from flwr.common.inflatable_protobuf_utils import (
 from flwr.common.logger import log
 from flwr.common.message import Message, remove_content_from_message
 from flwr.common.retry_invoker import RetryInvoker
-from flwr.common.serde import message_from_proto, message_to_proto, run_from_proto
+from flwr.common.serde import (
+    fab_from_proto,
+    message_from_proto,
+    message_to_proto,
+    run_from_proto,
+)
 from flwr.common.typing import Fab, Run
 from flwr.proto.fab_pb2 import GetFabRequest, GetFabResponse  # pylint: disable=E0611
 from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
-    CreateNodeRequest,
-    CreateNodeResponse,
-    DeleteNodeRequest,
-    DeleteNodeResponse,
+    ActivateNodeRequest,
+    ActivateNodeResponse,
+    DeactivateNodeRequest,
+    DeactivateNodeResponse,
     PullMessagesRequest,
     PullMessagesResponse,
     PushMessagesRequest,
     PushMessagesResponse,
+    RegisterNodeFleetRequest,
+    RegisterNodeFleetResponse,
+    UnregisterNodeFleetRequest,
+    UnregisterNodeFleetResponse,
 )
 from flwr.proto.heartbeat_pb2 import (  # pylint: disable=E0611
     SendNodeHeartbeatRequest,
@@ -64,6 +73,7 @@ from flwr.proto.message_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.supercore.primitives.asymmetric import generate_key_pairs, public_key_to_bytes
 
 try:
     import requests
@@ -71,8 +81,10 @@ except ModuleNotFoundError:
     flwr_exit(ExitCode.COMMON_MISSING_EXTRA_REST)
 
 
-PATH_CREATE_NODE: str = "api/v0/fleet/create-node"
-PATH_DELETE_NODE: str = "api/v0/fleet/delete-node"
+PATH_REGISTER_NODE: str = "/api/v0/fleet/register-node"
+PATH_ACTIVATE_NODE: str = "/api/v0/fleet/activate-node"
+PATH_DEACTIVATE_NODE: str = "/api/v0/fleet/deactivate-node"
+PATH_UNREGISTER_NODE: str = "/api/v0/fleet/unregister-node"
 PATH_PULL_MESSAGES: str = "/api/v0/fleet/pull-messages"
 PATH_PUSH_MESSAGES: str = "/api/v0/fleet/push-messages"
 PATH_PULL_OBJECT: str = "/api/v0/fleet/pull-object"
@@ -99,10 +111,9 @@ def http_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
     ] = None,
 ) -> Iterator[
     tuple[
+        int,
         Callable[[], Optional[tuple[Message, ObjectTree]]],
         Callable[[Message, ObjectTree], set[str]],
-        Callable[[], Optional[int]],
-        Callable[[], None],
         Callable[[int], Run],
         Callable[[str, int], Fab],
         Callable[[int, str], bytes],
@@ -134,15 +145,15 @@ def http_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
         connection using the certificates will be established to an SSL-enabled
         Flower server. Bytes won't work for the REST API.
     authentication_keys : Optional[Tuple[PrivateKey, PublicKey]] (default: None)
-        Client authentication is not supported for this transport type.
+        SuperNode authentication is not supported for this transport type.
 
     Returns
     -------
-    receive : Callable
-    send : Callable
-    create_node : Optional[Callable]
-    delete_node : Optional[Callable]
-    get_run : Optional[Callable]
+    node_id : int
+    receive : Callable[[], Optional[tuple[Message, ObjectTree]]]
+    send : Callable[[Message, ObjectTree], set[str]]
+    get_run : Callable[[int], Run]
+    get_fab : Callable[[str, int], Fab]
     pull_object : Callable[[str], bytes]
     push_object : Callable[[str, bytes], None]
     confirm_message_received : Callable[[str], None]
@@ -171,7 +182,14 @@ def http_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
             "must be provided as a string path to the client.",
         )
     if authentication_keys is not None:
-        log(ERROR, "Client authentication is not supported for this transport type.")
+        log(ERROR, "SuperNode authentication is not supported for this transport type.")
+
+    # REST does NOT support node authentication
+    self_registered = False
+    if authentication_keys is None:
+        self_registered = True
+        authentication_keys = generate_key_pairs()
+    node_pk = public_key_to_bytes(authentication_keys[1])
 
     # Shared variables for inner functions
     node: Optional[Node] = None
@@ -180,7 +198,7 @@ def http_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
     retry_invoker.should_giveup = None
 
     ###########################################################################
-    # heartbeat/create_node/delete_node/receive/send/get_run functions
+    # SuperNode functions
     ###########################################################################
 
     def _request(
@@ -290,23 +308,35 @@ def http_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
 
     heartbeat_sender = HeartbeatSender(send_node_heartbeat)
 
-    def create_node() -> Optional[int]:
-        """Set create_node."""
-        req = CreateNodeRequest(heartbeat_interval=HEARTBEAT_DEFAULT_INTERVAL)
+    def register_node() -> None:
+        """Register node with SuperLink."""
+        req = RegisterNodeFleetRequest(public_key=node_pk)
 
         # Send the request
-        res = _request(req, CreateNodeResponse, PATH_CREATE_NODE)
+        res = _request(req, RegisterNodeFleetResponse, PATH_REGISTER_NODE)
         if res is None:
-            return None
+            raise RuntimeError("Failed to register node")
+
+    def activate_node() -> int:
+        """Activate node and start heartbeat."""
+        req = ActivateNodeRequest(
+            public_key=node_pk,
+            heartbeat_interval=HEARTBEAT_DEFAULT_INTERVAL,
+        )
+
+        # Send the request
+        res = _request(req, ActivateNodeResponse, PATH_ACTIVATE_NODE)
+        if res is None:
+            raise RuntimeError("Failed to activate node")
 
         # Remember the node and start the heartbeat sender
         nonlocal node
-        node = res.node
+        node = Node(node_id=res.node_id)
         heartbeat_sender.start()
         return node.node_id
 
-    def delete_node() -> None:
-        """Set delete_node."""
+    def deactivate_node() -> None:
+        """Deactivate node and stop heartbeat."""
         nonlocal node
         if node is None:
             raise RuntimeError("Node instance missing")
@@ -314,13 +344,27 @@ def http_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
         # Stop the heartbeat sender
         heartbeat_sender.stop()
 
-        # Send DeleteNode request
-        req = DeleteNodeRequest(node=node)
+        # Send DeactivateNode request
+        req = DeactivateNodeRequest(node_id=node.node_id)
 
         # Send the request
-        res = _request(req, DeleteNodeResponse, PATH_DELETE_NODE)
+        res = _request(req, DeactivateNodeResponse, PATH_DEACTIVATE_NODE)
         if res is None:
-            return
+            raise RuntimeError("Failed to deactivate node")
+
+    def unregister_node() -> None:
+        """Unregister node from SuperLink."""
+        nonlocal node
+        if node is None:
+            raise RuntimeError("Node instance missing")
+
+        # Send UnregisterNode request
+        req = UnregisterNodeFleetRequest(node_id=node.node_id)
+
+        # Send the request
+        res = _request(req, UnregisterNodeFleetResponse, PATH_UNREGISTER_NODE)
+        if res is None:
+            raise RuntimeError("Failed to unregister node")
 
         # Cleanup
         node = None
@@ -392,12 +436,9 @@ def http_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
         # Send the request
         res = _request(req, GetFabResponse, PATH_GET_FAB)
         if res is None:
-            return Fab("", b"")
+            return Fab("", b"", {})
 
-        return Fab(
-            res.fab.hash_str,
-            res.fab.content,
-        )
+        return fab_from_proto(res.fab)
 
     def pull_object(run_id: int, object_id: str) -> bytes:
         """Pull the object from the SuperLink."""
@@ -439,12 +480,14 @@ def http_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
         fn(object_id)
 
     try:
+        if self_registered:
+            register_node()
+        node_id = activate_node()
         # Yield methods
         yield (
+            node_id,
             receive,
             send,
-            create_node,
-            delete_node,
             get_run,
             get_fab,
             pull_object,
@@ -459,6 +502,8 @@ def http_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
             if node is not None:
                 # Disable retrying
                 retry_invoker.max_tries = 1
-                delete_node()
+                deactivate_node()
+                if self_registered:
+                    unregister_node()
         except RequestsConnectionError:
             pass

flwr/clientapp/__init__.py

@@ -15,9 +15,8 @@
 """Public Flower ClientApp APIs."""
 
 
-from flwr.client.client_app import ClientApp
-
 from . import mod
+from .client_app import ClientApp as ClientApp
 
 __all__ = [
     "ClientApp",

flwr/{client/clientapp → clientapp}/utils.py

@@ -19,7 +19,7 @@ from logging import DEBUG
 from pathlib import Path
 from typing import Callable, Optional
 
-from flwr.client.client_app import ClientApp, LoadClientAppError
+from flwr.clientapp.client_app import ClientApp, LoadClientAppError
 from flwr.common.config import (
     get_flwr_dir,
     get_metadata_from_config,

flwr/common/constant.py

@@ -17,6 +17,8 @@
 
 from __future__ import annotations
 
+import os
+
 TRANSPORT_TYPE_GRPC_BIDI = "grpc-bidi"
 TRANSPORT_TYPE_GRPC_RERE = "grpc-rere"
 TRANSPORT_TYPE_GRPC_ADAPTER = "grpc-adapter"
@@ -60,7 +62,9 @@ HEARTBEAT_DEFAULT_INTERVAL = 30
 HEARTBEAT_CALL_TIMEOUT = 5
 HEARTBEAT_BASE_MULTIPLIER = 0.8
 HEARTBEAT_RANDOM_RANGE = (-0.1, 0.1)
-HEARTBEAT_MAX_INTERVAL = 1e300
+HEARTBEAT_MIN_INTERVAL = 10
+HEARTBEAT_MAX_INTERVAL = 1800  # 30 minutes
+HEARTBEAT_INTERVAL_INF = 1e300  # Large value, disabling heartbeats
 HEARTBEAT_PATIENCE = 2
 RUN_FAILURE_DETAILS_NO_HEARTBEAT = "No heartbeat received from the run."
 
@@ -70,13 +74,23 @@ NODE_ID_NUM_BYTES = 8
 
 # Constants for FAB
 APP_DIR = "apps"
-FAB_ALLOWED_EXTENSIONS = {".py", ".toml", ".md"}
 FAB_CONFIG_FILE = "pyproject.toml"
 FAB_DATE = (2024, 10, 1, 0, 0, 0)
 FAB_HASH_TRUNCATION = 8
 FAB_MAX_SIZE = 10 * 1024 * 1024  # 10 MB
 FLWR_DIR = ".flwr"  # The default Flower directory: ~/.flwr/
 FLWR_HOME = "FLWR_HOME"  # If set, override the default Flower directory
+# FAB file include patterns (gitignore-style patterns)
+FAB_INCLUDE_PATTERNS = (
+    "**/*.py",
+    "**/*.toml",
+    "**/*.md",
+)
+# FAB file exclude patterns (gitignore-style patterns)
+FAB_EXCLUDE_PATTERNS = (
+    "**/__pycache__/**",
+    FAB_CONFIG_FILE,  # Exclude the original pyproject.toml
+)
 
 # Constant for SuperLink
 SUPERLINK_NODE_ID = 1
@@ -109,14 +123,14 @@ LOG_UPLOAD_INTERVAL = 0.2  # Minimum interval between two log uploads
 # Retry configurations
 MAX_RETRY_DELAY = 20  # Maximum delay duration between two consecutive retries.
 
-# Constants for user authentication
+# Constants for account authentication
 CREDENTIALS_DIR = ".credentials"
-AUTH_TYPE_JSON_KEY = "auth-type"  # For key name in JSON file
-AUTH_TYPE_YAML_KEY = "auth_type"  # For key name in YAML file
+AUTHN_TYPE_JSON_KEY = "authn-type"  # For key name in JSON file
+AUTHN_TYPE_YAML_KEY = "authn_type"  # For key name in YAML file
 ACCESS_TOKEN_KEY = "flwr-oidc-access-token"
 REFRESH_TOKEN_KEY = "flwr-oidc-refresh-token"
 
-# Constants for user authorization
+# Constants for account authorization
 AUTHZ_TYPE_YAML_KEY = "authz_type"  # For key name in YAML file
 
 # Constants for node authentication
@@ -135,7 +149,9 @@ GC_THRESHOLD = 200_000_000  # 200 MB
 # Constants for Inflatable
 HEAD_BODY_DIVIDER = b"\x00"
 HEAD_VALUE_DIVIDER = " "
-MAX_ARRAY_CHUNK_SIZE = 20_971_520  # 20 MB
+FLWR_PRIVATE_MAX_ARRAY_CHUNK_SIZE = int(
+    os.getenv("FLWR_PRIVATE_MAX_ARRAY_CHUNK_SIZE", "5242880")
+)  # 5 MB
 
 # Constants for serialization
 INT64_MAX_VALUE = 9223372036854775807  # (1 << 63) - 1
@@ -144,8 +160,12 @@ INT64_MAX_VALUE = 9223372036854775807  # (1 << 63) - 1
 FLWR_APP_TOKEN_LENGTH = 128  # Length of the token used
 
 # Constants for object pushing and pulling
-MAX_CONCURRENT_PUSHES = 8  # Default maximum number of concurrent pushes
-MAX_CONCURRENT_PULLS = 8  # Default maximum number of concurrent pulls
+FLWR_PRIVATE_MAX_CONCURRENT_OBJ_PUSHES = int(
+    os.getenv("FLWR_PRIVATE_MAX_CONCURRENT_OBJ_PUSHES", "2")
+)  # Default maximum number of concurrent pushes
+FLWR_PRIVATE_MAX_CONCURRENT_OBJ_PULLS = int(
+    os.getenv("FLWR_PRIVATE_MAX_CONCURRENT_OBJ_PULLS", "2")
+)  # Default maximum number of concurrent pulls
 PULL_MAX_TIME = 7200  # Default maximum time to wait for pulling objects
 PULL_MAX_TRIES_PER_OBJECT = 500  # Default maximum number of tries to pull an object
 PULL_INITIAL_BACKOFF = 1  # Initial backoff time for pulling objects
@@ -154,9 +174,13 @@ PULL_BACKOFF_CAP = 10  # Maximum backoff time for pulling objects
 
 # ControlServicer constants
 RUN_ID_NOT_FOUND_MESSAGE = "Run ID not found"
-NO_USER_AUTH_MESSAGE = "ControlServicer initialized without user authentication"
+NO_ACCOUNT_AUTH_MESSAGE = "ControlServicer initialized without account authentication"
 NO_ARTIFACT_PROVIDER_MESSAGE = "ControlServicer initialized without artifact provider"
 PULL_UNFINISHED_RUN_MESSAGE = "Cannot pull artifacts for an unfinished run"
+SUPERNODE_NOT_CREATED_FROM_CLI_MESSAGE = "Invalid SuperNode credentials"
+PUBLIC_KEY_ALREADY_IN_USE_MESSAGE = "Public key already in use"
+PUBLIC_KEY_NOT_VALID = "The provided public key is not valid"
+NODE_NOT_FOUND_MESSAGE = "Node ID not found for account"
 
 
 class MessageType:
@@ -245,12 +269,23 @@ class CliOutputFormat:
         raise TypeError(f"{cls.__name__} cannot be instantiated.")
 
 
-class AuthType:
-    """User authentication types."""
+class AuthnType:
+    """Account authentication types."""
 
+    NOOP = "noop"
     OIDC = "oidc"
 
-    def __new__(cls) -> AuthType:
+    def __new__(cls) -> AuthnType:
+        """Prevent instantiation."""
+        raise TypeError(f"{cls.__name__} cannot be instantiated.")
+
+
+class AuthzType:
+    """Account authorization types."""
+
+    NOOP = "noop"
+
+    def __new__(cls) -> AuthzType:
         """Prevent instantiation."""
         raise TypeError(f"{cls.__name__} cannot be instantiated.")
 
@@ -281,3 +316,8 @@ class ExecPluginType:
         """Return all SuperExec plugin types."""
         # Filter all constants (uppercase) of the class
         return [v for k, v in vars(ExecPluginType).items() if k.isupper()]
+
+
+# Constants for No-op auth plugins
+NOOP_FLWR_AID = "<none>"
+NOOP_ACCOUNT_NAME = "sys_noauth"

flwr/common/exit/exit_code.py

@@ -41,12 +41,16 @@ class ExitCode:
 
     # SuperNode-specific exit codes (300-399)
     SUPERNODE_REST_ADDRESS_INVALID = 300
-    SUPERNODE_NODE_AUTH_KEYS_REQUIRED = 301
-    SUPERNODE_NODE_AUTH_KEYS_INVALID = 302
+    # SUPERNODE_NODE_AUTH_KEYS_REQUIRED = 301 --- DELETED ---
+    SUPERNODE_NODE_AUTH_KEY_INVALID = 302
+    SUPERNODE_STARTED_WITHOUT_TLS_BUT_NODE_AUTH_ENABLED = 303
 
     # SuperExec-specific exit codes (400-499)
     SUPEREXEC_INVALID_PLUGIN_CONFIG = 400
 
+    # FlowerCLI-specific exit codes (500-599)
+    FLWRCLI_NODE_AUTH_PUBLIC_KEY_INVALID = 500
+
     # Common exit codes (600-699)
     COMMON_ADDRESS_INVALID = 600
     COMMON_MISSING_EXTRA_REST = 601
@@ -102,20 +106,26 @@ EXIT_CODE_HELP = {
         "When using the REST API, please provide `https://` or "
         "`http://` before the server address (e.g. `http://127.0.0.1:8080`)"
     ),
-    ExitCode.SUPERNODE_NODE_AUTH_KEYS_REQUIRED: (
-        "Node authentication requires file paths to both "
-        "'--auth-supernode-private-key' and '--auth-supernode-public-key' "
-        "to be provided (providing only one of them is not sufficient)."
-    ),
-    ExitCode.SUPERNODE_NODE_AUTH_KEYS_INVALID: (
-        "Node authentication requires elliptic curve private and public key pair. "
-        "Please ensure that the file path points to a valid private/public key "
+    ExitCode.SUPERNODE_NODE_AUTH_KEY_INVALID: (
+        "Node authentication requires elliptic curve private key. "
+        "Please ensure that the file path points to a valid private key "
         "file and try again."
     ),
+    ExitCode.SUPERNODE_STARTED_WITHOUT_TLS_BUT_NODE_AUTH_ENABLED: (
+        "The private key for SuperNode authentication was provided, but TLS is not "
+        "enabled. Node authentication can only be used when TLS is enabled."
+    ),
     # SuperExec-specific exit codes (400-499)
     ExitCode.SUPEREXEC_INVALID_PLUGIN_CONFIG: (
         "The YAML configuration for the SuperExec plugin is invalid."
     ),
+    # FlowerCLI-specific exit codes (500-599)
+    ExitCode.FLWRCLI_NODE_AUTH_PUBLIC_KEY_INVALID: (
+        "Node authentication requires a valid elliptic curve public key in the "
+        "SSH format and following a NIST standard elliptic curve (e.g. SECP384R1). "
+        "Please ensure that the file path points to a valid public key "
+        "file and try again."
+    ),
     # Common exit codes (600-699)
     ExitCode.COMMON_ADDRESS_INVALID: (
         "Please provide a valid URL, IPv4 or IPv6 address."

flwr/common/inflatable_utils.py

@@ -25,10 +25,10 @@ from typing import Callable, Optional, TypeVar
 from flwr.proto.message_pb2 import ObjectTree  # pylint: disable=E0611
 
 from .constant import (
+    FLWR_PRIVATE_MAX_CONCURRENT_OBJ_PULLS,
+    FLWR_PRIVATE_MAX_CONCURRENT_OBJ_PUSHES,
     HEAD_BODY_DIVIDER,
     HEAD_VALUE_DIVIDER,
-    MAX_CONCURRENT_PULLS,
-    MAX_CONCURRENT_PUSHES,
     PULL_BACKOFF_CAP,
     PULL_INITIAL_BACKOFF,
     PULL_MAX_TIME,
@@ -118,7 +118,7 @@ def push_objects(
     *,
     object_ids_to_push: Optional[set[str]] = None,
     keep_objects: bool = False,
-    max_concurrent_pushes: int = MAX_CONCURRENT_PUSHES,
+    max_concurrent_pushes: int = FLWR_PRIVATE_MAX_CONCURRENT_OBJ_PUSHES,
 ) -> None:
     """Push multiple objects to the servicer.
 
@@ -137,7 +137,7 @@ def push_objects(
         If `True`, the original objects will be kept in the `objects` dictionary
         after pushing. If `False`, they will be removed from the dictionary to avoid
         high memory usage.
-    max_concurrent_pushes : int (default: MAX_CONCURRENT_PUSHES)
+    max_concurrent_pushes : int (default: FLWR_PRIVATE_MAX_CONCURRENT_OBJ_PUSHES)
         The maximum number of concurrent pushes to perform.
     """
     lock = threading.Lock()
@@ -168,7 +168,7 @@ def push_object_contents_from_iterable(
     object_contents: Iterable[tuple[str, bytes]],
     push_object_fn: Callable[[str, bytes], None],
     *,
-    max_concurrent_pushes: int = MAX_CONCURRENT_PUSHES,
+    max_concurrent_pushes: int = FLWR_PRIVATE_MAX_CONCURRENT_OBJ_PUSHES,
 ) -> None:
     """Push multiple object contents to the servicer.
 
@@ -181,7 +181,7 @@ def push_object_contents_from_iterable(
         A function that takes an object ID and its content as bytes, and pushes
         it to the servicer. This function should raise `ObjectIdNotPreregisteredError`
         if the object ID is not pre-registered.
-    max_concurrent_pushes : int (default: MAX_CONCURRENT_PUSHES)
+    max_concurrent_pushes : int (default: FLWR_PRIVATE_MAX_CONCURRENT_OBJ_PUSHES)
         The maximum number of concurrent pushes to perform.
     """
 
@@ -210,7 +210,7 @@ def pull_objects(  # pylint: disable=too-many-arguments,too-many-locals
     object_ids: list[str],
     pull_object_fn: Callable[[str], bytes],
     *,
-    max_concurrent_pulls: int = MAX_CONCURRENT_PULLS,
+    max_concurrent_pulls: int = FLWR_PRIVATE_MAX_CONCURRENT_OBJ_PULLS,
     max_time: Optional[float] = PULL_MAX_TIME,
     max_tries_per_object: Optional[int] = PULL_MAX_TRIES_PER_OBJECT,
     initial_backoff: float = PULL_INITIAL_BACKOFF,
@@ -227,7 +227,7 @@ def pull_objects(  # pylint: disable=too-many-arguments,too-many-locals
         The function should raise `ObjectUnavailableError` if the object is not yet
         available, or `ObjectIdNotPreregisteredError` if the object ID is not
         pre-registered.
-    max_concurrent_pulls : int (default: MAX_CONCURRENT_PULLS)
+    max_concurrent_pulls : int (default: FLWR_PRIVATE_MAX_CONCURRENT_OBJ_PULLS)
         The maximum number of concurrent pulls to perform.
     max_time : Optional[float] (default: PULL_MAX_TIME)
         The maximum time to wait for all pulls to complete. If `None`, waits
@@ -442,7 +442,7 @@ def pull_and_inflate_object_from_tree(  # pylint: disable=R0913
     confirm_object_received_fn: Callable[[str], None],
     *,
     return_type: type[T] = InflatableObject,  # type: ignore
-    max_concurrent_pulls: int = MAX_CONCURRENT_PULLS,
+    max_concurrent_pulls: int = FLWR_PRIVATE_MAX_CONCURRENT_OBJ_PULLS,
     max_time: Optional[float] = PULL_MAX_TIME,
     max_tries_per_object: Optional[int] = PULL_MAX_TRIES_PER_OBJECT,
     initial_backoff: float = PULL_INITIAL_BACKOFF,
@@ -460,7 +460,7 @@ def pull_and_inflate_object_from_tree(  # pylint: disable=R0913
         A function to confirm that the object has been received.
     return_type : type[T] (default: InflatableObject)
         The type of the object to return. Must be a subclass of `InflatableObject`.
-    max_concurrent_pulls : int (default: MAX_CONCURRENT_PULLS)
+    max_concurrent_pulls : int (default: FLWR_PRIVATE_MAX_CONCURRENT_OBJ_PULLS)
         The maximum number of concurrent pulls to perform.
     max_time : Optional[float] (default: PULL_MAX_TIME)
         The maximum time to wait for all pulls to complete. If `None`, waits

flwr/common/record/array.py

@@ -25,7 +25,7 @@ from typing import TYPE_CHECKING, Any, cast, overload
 
 import numpy as np
 
-from ..constant import MAX_ARRAY_CHUNK_SIZE, SType
+from ..constant import FLWR_PRIVATE_MAX_ARRAY_CHUNK_SIZE, SType
 from ..inflatable import (
     InflatableObject,
     add_header_to_object_body,
@@ -272,8 +272,8 @@ class Array(InflatableObject):
         chunks: list[tuple[str, InflatableObject]] = []
         # memoryview allows for zero-copy slicing
         data_view = memoryview(self.data)
-        for start in range(0, len(data_view), MAX_ARRAY_CHUNK_SIZE):
-            end = min(start + MAX_ARRAY_CHUNK_SIZE, len(data_view))
+        for start in range(0, len(data_view), FLWR_PRIVATE_MAX_ARRAY_CHUNK_SIZE):
+            end = min(start + FLWR_PRIVATE_MAX_ARRAY_CHUNK_SIZE, len(data_view))
             ac = ArrayChunk(data_view[start:end])
             chunks.append((ac.object_id, ac))
 

flwr/common/record/arrayrecord.py

@@ -147,11 +147,20 @@ class ArrayRecord(TypedDict[str, Array], InflatableObject):
         keep_input: bool = True,
     ) -> None: ...
 
+    # This is also required for PyTorch state dict because they are not strongly typed
+    @overload
+    def __init__(  # noqa: E704
+        self,
+        torch_state_dict: dict[str, Any],
+        *,
+        keep_input: bool = True,
+    ) -> None: ...
+
     def __init__(  # pylint: disable=too-many-arguments
         self,
         *args: Any,
         numpy_ndarrays: list[NDArray] | None = None,
-        torch_state_dict: OrderedDict[str, torch.Tensor] | None = None,
+        torch_state_dict: OrderedDict[str, torch.Tensor] | dict[str, Any] | None = None,
         array_dict: OrderedDict[str, Array] | None = None,
         keep_input: bool = True,
     ) -> None:

flwr/common/secure_aggregation/crypto/symmetric_encryption.py

@@ -16,57 +16,14 @@
 
 
 import base64
-from typing import cast
 
 from cryptography.exceptions import InvalidSignature
 from cryptography.fernet import Fernet
-from cryptography.hazmat.primitives import hashes, hmac, serialization
+from cryptography.hazmat.primitives import hashes, hmac
 from cryptography.hazmat.primitives.asymmetric import ec
 from cryptography.hazmat.primitives.kdf.hkdf import HKDF
 
 
-def generate_key_pairs() -> (
-    tuple[ec.EllipticCurvePrivateKey, ec.EllipticCurvePublicKey]
-):
-    """Generate private and public key pairs with Cryptography."""
-    private_key = ec.generate_private_key(ec.SECP384R1())
-    public_key = private_key.public_key()
-    return private_key, public_key
-
-
-def private_key_to_bytes(private_key: ec.EllipticCurvePrivateKey) -> bytes:
-    """Serialize private key to bytes."""
-    return private_key.private_bytes(
-        encoding=serialization.Encoding.PEM,
-        format=serialization.PrivateFormat.PKCS8,
-        encryption_algorithm=serialization.NoEncryption(),
-    )
-
-
-def bytes_to_private_key(private_key_bytes: bytes) -> ec.EllipticCurvePrivateKey:
-    """Deserialize private key from bytes."""
-    return cast(
-        ec.EllipticCurvePrivateKey,
-        serialization.load_pem_private_key(data=private_key_bytes, password=None),
-    )
-
-
-def public_key_to_bytes(public_key: ec.EllipticCurvePublicKey) -> bytes:
-    """Serialize public key to bytes."""
-    return public_key.public_bytes(
-        encoding=serialization.Encoding.PEM,
-        format=serialization.PublicFormat.SubjectPublicKeyInfo,
-    )
-
-
-def bytes_to_public_key(public_key_bytes: bytes) -> ec.EllipticCurvePublicKey:
-    """Deserialize public key from bytes."""
-    return cast(
-        ec.EllipticCurvePublicKey,
-        serialization.load_pem_public_key(data=public_key_bytes),
-    )
-
-
 def generate_shared_key(
     private_key: ec.EllipticCurvePrivateKey, public_key: ec.EllipticCurvePublicKey
 ) -> bytes:
@@ -117,48 +74,3 @@ def verify_hmac(key: bytes, message: bytes, hmac_value: bytes) -> bool:
         return True
     except InvalidSignature:
         return False
-
-
-def sign_message(private_key: ec.EllipticCurvePrivateKey, message: bytes) -> bytes:
-    """Sign a message using the provided EC private key.
-
-    Parameters
-    ----------
-    private_key : ec.EllipticCurvePrivateKey
-        The EC private key to sign the message with.
-    message : bytes
-        The message to be signed.
-
-    Returns
-    -------
-    bytes
-        The signature of the message.
-    """
-    signature = private_key.sign(message, ec.ECDSA(hashes.SHA256()))
-    return signature
-
-
-def verify_signature(
-    public_key: ec.EllipticCurvePublicKey, message: bytes, signature: bytes
-) -> bool:
-    """Verify a signature against a message using the provided EC public key.
-
-    Parameters
-    ----------
-    public_key : ec.EllipticCurvePublicKey
-        The EC public key to verify the signature.
-    message : bytes
-        The original message.
-    signature : bytes
-        The signature to verify.
-
-    Returns
-    -------
-    bool
-        True if the signature is valid, False otherwise.
-    """
-    try:
-        public_key.verify(signature, message, ec.ECDSA(hashes.SHA256()))
-        return True
-    except InvalidSignature:
-        return False