flwr-nightly 1.9.0.dev20240520__py3-none-any.whl → 1.10.0.dev20240612__py3-none-any.whl

flwr/client/app.py

@@ -14,8 +14,10 @@
 # ==============================================================================
 """Flower client app."""
 
+import signal
 import sys
 import time
+from dataclasses import dataclass
 from logging import DEBUG, ERROR, INFO, WARN
 from typing import Callable, ContextManager, Optional, Tuple, Type, Union
 
@@ -37,7 +39,7 @@ from flwr.common.constant import (
 )
 from flwr.common.logger import log, warn_deprecated_feature
 from flwr.common.message import Error
-from flwr.common.retry_invoker import RetryInvoker, exponential
+from flwr.common.retry_invoker import RetryInvoker, RetryState, exponential
 
 from .grpc_client.connection import grpc_connection
 from .grpc_rere_client.connection import grpc_request_response
@@ -263,6 +265,29 @@ def _start_client_internal(
         transport, server_address
     )
 
+    app_state_tracker = _AppStateTracker()
+
+    def _on_sucess(retry_state: RetryState) -> None:
+        app_state_tracker.is_connected = True
+        if retry_state.tries > 1:
+            log(
+                INFO,
+                "Connection successful after %.2f seconds and %s tries.",
+                retry_state.elapsed_time,
+                retry_state.tries,
+            )
+
+    def _on_backoff(retry_state: RetryState) -> None:
+        app_state_tracker.is_connected = False
+        if retry_state.tries == 1:
+            log(WARN, "Connection attempt failed, retrying...")
+        else:
+            log(
+                DEBUG,
+                "Connection attempt failed, retrying in %.2f seconds",
+                retry_state.actual_wait,
+            )
+
     retry_invoker = RetryInvoker(
         wait_gen_factory=exponential,
         recoverable_exceptions=connection_error_type,
@@ -278,30 +303,13 @@ def _start_client_internal(
             if retry_state.tries > 1
             else None
         ),
-        on_success=lambda retry_state: (
-            log(
-                INFO,
-                "Connection successful after %.2f seconds and %s tries.",
-                retry_state.elapsed_time,
-                retry_state.tries,
-            )
-            if retry_state.tries > 1
-            else None
-        ),
-        on_backoff=lambda retry_state: (
-            log(WARN, "Connection attempt failed, retrying...")
-            if retry_state.tries == 1
-            else log(
-                DEBUG,
-                "Connection attempt failed, retrying in %.2f seconds",
-                retry_state.actual_wait,
-            )
-        ),
+        on_success=_on_sucess,
+        on_backoff=_on_backoff,
     )
 
     node_state = NodeState()
 
-    while True:
+    while not app_state_tracker.interrupt:
         sleep_duration: int = 0
         with connection(
             address,
@@ -318,99 +326,112 @@ def _start_client_internal(
             if create_node is not None:
                 create_node()  # pylint: disable=not-callable
 
-            while True:
-                # Receive
-                message = receive()
-                if message is None:
-                    time.sleep(3)  # Wait for 3s before asking again
-                    continue
-
-                log(INFO, "")
-                if len(message.metadata.group_id) > 0:
+            app_state_tracker.register_signal_handler()
+            while not app_state_tracker.interrupt:
+                try:
+                    # Receive
+                    message = receive()
+                    if message is None:
+                        time.sleep(3)  # Wait for 3s before asking again
+                        continue
+
+                    log(INFO, "")
+                    if len(message.metadata.group_id) > 0:
+                        log(
+                            INFO,
+                            "[RUN %s, ROUND %s]",
+                            message.metadata.run_id,
+                            message.metadata.group_id,
+                        )
                     log(
                         INFO,
-                        "[RUN %s, ROUND %s]",
-                        message.metadata.run_id,
-                        message.metadata.group_id,
+                        "Received: %s message %s",
+                        message.metadata.message_type,
+                        message.metadata.message_id,
                     )
-                log(
-                    INFO,
-                    "Received: %s message %s",
-                    message.metadata.message_type,
-                    message.metadata.message_id,
-                )
-
-                # Handle control message
-                out_message, sleep_duration = handle_control_message(message)
-                if out_message:
-                    send(out_message)
-                    break
-
-                # Register context for this run
-                node_state.register_context(run_id=message.metadata.run_id)
-
-                # Retrieve context for this run
-                context = node_state.retrieve_context(run_id=message.metadata.run_id)
 
-                # Create an error reply message that will never be used to prevent
-                # the used-before-assignment linting error
-                reply_message = message.create_error_reply(
-                    error=Error(code=ErrorCode.UNKNOWN, reason="Unknown")
-                )
+                    # Handle control message
+                    out_message, sleep_duration = handle_control_message(message)
+                    if out_message:
+                        send(out_message)
+                        break
 
-                # Handle app loading and task message
-                try:
-                    # Load ClientApp instance
-                    client_app: ClientApp = load_client_app_fn()
-
-                    # Execute ClientApp
-                    reply_message = client_app(message=message, context=context)
-                except Exception as ex:  # pylint: disable=broad-exception-caught
-
-                    # Legacy grpc-bidi
-                    if transport in ["grpc-bidi", None]:
-                        log(ERROR, "Client raised an exception.", exc_info=ex)
-                        # Raise exception, crash process
-                        raise ex
-
-                    # Don't update/change NodeState
-
-                    e_code = ErrorCode.CLIENT_APP_RAISED_EXCEPTION
-                    # Reason example: "<class 'ZeroDivisionError'>:<'division by zero'>"
-                    reason = str(type(ex)) + ":<'" + str(ex) + "'>"
-                    exc_entity = "ClientApp"
-                    if isinstance(ex, LoadClientAppError):
-                        reason = (
-                            "An exception was raised when attempting to load "
-                            "`ClientApp`"
-                        )
-                        e_code = ErrorCode.LOAD_CLIENT_APP_EXCEPTION
-                        exc_entity = "SuperNode"
+                    # Register context for this run
+                    node_state.register_context(run_id=message.metadata.run_id)
 
-                    log(ERROR, "%s raised an exception", exc_entity, exc_info=ex)
+                    # Retrieve context for this run
+                    context = node_state.retrieve_context(
+                        run_id=message.metadata.run_id
+                    )
 
-                    # Create error message
+                    # Create an error reply message that will never be used to prevent
+                    # the used-before-assignment linting error
                     reply_message = message.create_error_reply(
-                        error=Error(code=e_code, reason=reason)
-                    )
-                else:
-                    # No exception, update node state
-                    node_state.update_context(
-                        run_id=message.metadata.run_id,
-                        context=context,
+                        error=Error(code=ErrorCode.UNKNOWN, reason="Unknown")
                     )
 
-                # Send
-                send(reply_message)
-                log(INFO, "Sent reply")
+                    # Handle app loading and task message
+                    try:
+                        # Load ClientApp instance
+                        client_app: ClientApp = load_client_app_fn()
+
+                        # Execute ClientApp
+                        reply_message = client_app(message=message, context=context)
+                    except Exception as ex:  # pylint: disable=broad-exception-caught
+
+                        # Legacy grpc-bidi
+                        if transport in ["grpc-bidi", None]:
+                            log(ERROR, "Client raised an exception.", exc_info=ex)
+                            # Raise exception, crash process
+                            raise ex
+
+                        # Don't update/change NodeState
+
+                        e_code = ErrorCode.CLIENT_APP_RAISED_EXCEPTION
+                        # Ex fmt: "<class 'ZeroDivisionError'>:<'division by zero'>"
+                        reason = str(type(ex)) + ":<'" + str(ex) + "'>"
+                        exc_entity = "ClientApp"
+                        if isinstance(ex, LoadClientAppError):
+                            reason = (
+                                "An exception was raised when attempting to load "
+                                "`ClientApp`"
+                            )
+                            e_code = ErrorCode.LOAD_CLIENT_APP_EXCEPTION
+                            exc_entity = "SuperNode"
+
+                        if not app_state_tracker.interrupt:
+                            log(
+                                ERROR, "%s raised an exception", exc_entity, exc_info=ex
+                            )
+
+                        # Create error message
+                        reply_message = message.create_error_reply(
+                            error=Error(code=e_code, reason=reason)
+                        )
+                    else:
+                        # No exception, update node state
+                        node_state.update_context(
+                            run_id=message.metadata.run_id,
+                            context=context,
+                        )
+
+                    # Send
+                    send(reply_message)
+                    log(INFO, "Sent reply")
+
+                except StopIteration:
+                    sleep_duration = 0
+                    break
 
             # Unregister node
-            if delete_node is not None:
+            if delete_node is not None and app_state_tracker.is_connected:
                 delete_node()  # pylint: disable=not-callable
 
         if sleep_duration == 0:
             log(INFO, "Disconnect and shut down")
+            del app_state_tracker
             break
+
         # Sleep and reconnect afterwards
         log(
             INFO,
@@ -579,3 +600,20 @@ def _init_connection(transport: Optional[str], server_address: str) -> Tuple[
         )
 
     return connection, address, error_type
+
+
+@dataclass
+class _AppStateTracker:
+    interrupt: bool = False
+    is_connected: bool = False
+
+    def register_signal_handler(self) -> None:
+        """Register handlers for exit signals."""
+
+        def signal_handler(sig, frame):  # type: ignore
+            # pylint: disable=unused-argument
+            self.interrupt = True
+            raise StopIteration from None
+
+        signal.signal(signal.SIGINT, signal_handler)
+        signal.signal(signal.SIGTERM, signal_handler)

flwr/client/client_app.py

@@ -221,7 +221,7 @@ def _registration_error(fn_name: str) -> ValueError:
         >>> def client_fn(cid) -> Client:
         >>>     return FlowerClient().to_client()
         >>>
-        >>> app = ClientApp()
+        >>> app = ClientApp(
         >>>     client_fn=client_fn,
         >>> )
 

flwr/client/grpc_rere_client/client_interceptor.py

@@ -31,11 +31,11 @@ from flwr.common.secure_aggregation.crypto.symmetric_encryption import (
 from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
     CreateNodeRequest,
     DeleteNodeRequest,
-    GetRunRequest,
     PingRequest,
     PullTaskInsRequest,
     PushTaskResRequest,
 )
+from flwr.proto.run_pb2 import GetRunRequest  # pylint: disable=E0611
 
 _PUBLIC_KEY_HEADER = "public-key"
 _AUTH_TOKEN_HEADER = "auth-token"

flwr/client/grpc_rere_client/connection.py

@@ -21,7 +21,7 @@ from contextlib import contextmanager
 from copy import copy
 from logging import DEBUG, ERROR
 from pathlib import Path
-from typing import Callable, Iterator, Optional, Sequence, Tuple, Union, cast
+from typing import Callable, Iterator, Optional, Sequence, Tuple, Type, Union, cast
 
 import grpc
 from cryptography.hazmat.primitives.asymmetric import ec
@@ -44,8 +44,6 @@ from flwr.common.serde import message_from_taskins, message_to_taskres
 from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
     CreateNodeRequest,
     DeleteNodeRequest,
-    GetRunRequest,
-    GetRunResponse,
     PingRequest,
     PingResponse,
     PullTaskInsRequest,
@@ -53,6 +51,7 @@ from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
 )
 from flwr.proto.fleet_pb2_grpc import FleetStub  # pylint: disable=E0611
 from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
+from flwr.proto.run_pb2 import GetRunRequest, GetRunResponse  # pylint: disable=E0611
 from flwr.proto.task_pb2 import TaskIns  # pylint: disable=E0611
 
 from .client_interceptor import AuthenticateClientInterceptor
@@ -73,6 +72,7 @@ def grpc_request_response(  # pylint: disable=R0913, R0914, R0915
     authentication_keys: Optional[
         Tuple[ec.EllipticCurvePrivateKey, ec.EllipticCurvePublicKey]
     ] = None,
+    adapter_cls: Optional[Type[FleetStub]] = None,
 ) -> Iterator[
     Tuple[
         Callable[[], Optional[Message]],
@@ -133,7 +133,9 @@ def grpc_request_response(  # pylint: disable=R0913, R0914, R0915
     channel.subscribe(on_channel_state_change)
 
     # Shared variables for inner functions
-    stub = FleetStub(channel)
+    if adapter_cls is None:
+        adapter_cls = FleetStub
+    stub = adapter_cls(channel)
     metadata: Optional[Metadata] = None
     node: Optional[Node] = None
     ping_thread: Optional[threading.Thread] = None
@@ -190,8 +192,6 @@ def grpc_request_response(  # pylint: disable=R0913, R0914, R0915
 
         # Stop the ping-loop thread
         ping_stop_event.set()
-        if ping_thread is not None:
-            ping_thread.join()
 
         # Call FleetAPI
         delete_node_request = DeleteNodeRequest(node=node)

flwr/client/mod/__init__.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Mods."""
+"""Flower Built-in Mods."""
 
 
 from .centraldp_mods import adaptiveclipping_mod, fixedclipping_mod

flwr/client/rest_client/connection.py

@@ -46,8 +46,6 @@ from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
     CreateNodeResponse,
     DeleteNodeRequest,
     DeleteNodeResponse,
-    GetRunRequest,
-    GetRunResponse,
     PingRequest,
     PingResponse,
     PullTaskInsRequest,
@@ -56,6 +54,7 @@ from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
     PushTaskResResponse,
 )
 from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
+from flwr.proto.run_pb2 import GetRunRequest, GetRunResponse  # pylint: disable=E0611
 from flwr.proto.task_pb2 import TaskIns  # pylint: disable=E0611
 
 try:

flwr/client/supernode/app.py

@@ -20,6 +20,7 @@ from logging import DEBUG, INFO, WARN
 from pathlib import Path
 from typing import Callable, Optional, Tuple
 
+from cryptography.exceptions import UnsupportedAlgorithm
 from cryptography.hazmat.primitives.asymmetric import ec
 from cryptography.hazmat.primitives.serialization import (
     load_ssh_private_key,
@@ -29,14 +30,13 @@ from cryptography.hazmat.primitives.serialization import (
 from flwr.client.client_app import ClientApp, LoadClientAppError
 from flwr.common import EventType, event
 from flwr.common.exit_handlers import register_exit_handlers
-from flwr.common.logger import log
+from flwr.common.logger import log, warn_deprecated_feature
 from flwr.common.object_ref import load_app, validate
-from flwr.common.secure_aggregation.crypto.symmetric_encryption import (
-    ssh_types_to_elliptic_curve,
-)
 
 from ..app import _start_client_internal
 
+ADDRESS_FLEET_API_GRPC_RERE = "0.0.0.0:9092"
+
 
 def run_supernode() -> None:
     """Run Flower SuperNode."""
@@ -65,6 +65,23 @@ def run_client_app() -> None:
 
     args = _parse_args_run_client_app().parse_args()
 
+    if args.server != ADDRESS_FLEET_API_GRPC_RERE:
+        warn = "Passing flag --server is deprecated. Use --superlink instead."
+        warn_deprecated_feature(warn)
+
+        if args.superlink != ADDRESS_FLEET_API_GRPC_RERE:
+            # if `--superlink` also passed, then
+            # warn user that this argument overrides what was passed with `--server`
+            log(
+                WARN,
+                "Both `--server` and `--superlink` were passed. "
+                "`--server` will be ignored. Connecting to the Superlink Fleet API "
+                "at %s.",
+                args.superlink,
+            )
+        else:
+            args.superlink = args.server
+
     root_certificates = _get_certificates(args)
     log(
         DEBUG,
@@ -75,7 +92,7 @@ def run_client_app() -> None:
     authentication_keys = _try_setup_client_authentication(args)
 
     _start_client_internal(
-        server_address=args.server,
+        server_address=args.superlink,
         load_client_app_fn=load_fn,
         transport="rest" if args.rest else "grpc-rere",
         root_certificates=root_certificates,
@@ -102,7 +119,7 @@ def _get_certificates(args: argparse.Namespace) -> Optional[bytes]:
             WARN,
             "Option `--insecure` was set. "
             "Starting insecure HTTP client connected to %s.",
-            args.server,
+            args.superlink,
         )
         root_certificates = None
     else:
@@ -116,7 +133,7 @@ def _get_certificates(args: argparse.Namespace) -> Optional[bytes]:
             DEBUG,
             "Starting secure HTTPS client connected to %s "
             "with the following certificates: %s.",
-            args.server,
+            args.superlink,
             cert_path,
         )
     return root_certificates
@@ -215,9 +232,14 @@ def _parse_args_common(parser: argparse.ArgumentParser) -> None:
     )
     parser.add_argument(
         "--server",
-        default="0.0.0.0:9092",
+        default=ADDRESS_FLEET_API_GRPC_RERE,
         help="Server address",
     )
+    parser.add_argument(
+        "--superlink",
+        default=ADDRESS_FLEET_API_GRPC_RERE,
+        help="SuperLink Fleet API (gRPC-rere) address (IPv4, IPv6, or a domain name)",
+    )
     parser.add_argument(
         "--max-retries",
         type=int,
@@ -242,40 +264,60 @@ def _parse_args_common(parser: argparse.ArgumentParser) -> None:
         " Default: current working directory.",
     )
     parser.add_argument(
-        "--authentication-keys",
-        nargs=2,
-        metavar=("CLIENT_PRIVATE_KEY", "CLIENT_PUBLIC_KEY"),
+        "--auth-supernode-private-key",
+        type=str,
+        help="The SuperNode's private key (as a path str) to enable authentication.",
+    )
+    parser.add_argument(
+        "--auth-supernode-public-key",
         type=str,
-        help="Provide two file paths: (1) the client's private "
-        "key file, and (2) the client's public key file.",
+        help="The SuperNode's public key (as a path str) to enable authentication.",
     )
 
 
 def _try_setup_client_authentication(
     args: argparse.Namespace,
 ) -> Optional[Tuple[ec.EllipticCurvePrivateKey, ec.EllipticCurvePublicKey]]:
-    if not args.authentication_keys:
+    if not args.auth_supernode_private_key and not args.auth_supernode_public_key:
         return None
 
-    ssh_private_key = load_ssh_private_key(
-        Path(args.authentication_keys[0]).read_bytes(),
-        None,
-    )
-    ssh_public_key = load_ssh_public_key(Path(args.authentication_keys[1]).read_bytes())
+    if not args.auth_supernode_private_key or not args.auth_supernode_public_key:
+        sys.exit(
+            "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)."
+        )
+
+    try:
+        ssh_private_key = load_ssh_private_key(
+            Path(args.auth_supernode_private_key).read_bytes(),
+            None,
+        )
+        if not isinstance(ssh_private_key, ec.EllipticCurvePrivateKey):
+            raise ValueError()
+    except (ValueError, UnsupportedAlgorithm):
+        sys.exit(
+            "Error: Unable to parse the private key file in "
+            "'--auth-supernode-private-key'. Authentication requires elliptic "
+            "curve private and public key pair. Please ensure that the file "
+            "path points to a valid private key file and try again."
+        )
 
     try:
-        client_private_key, client_public_key = ssh_types_to_elliptic_curve(
-            ssh_private_key, ssh_public_key
+        ssh_public_key = load_ssh_public_key(
+            Path(args.auth_supernode_public_key).read_bytes()
         )
-    except TypeError:
+        if not isinstance(ssh_public_key, ec.EllipticCurvePublicKey):
+            raise ValueError()
+    except (ValueError, UnsupportedAlgorithm):
         sys.exit(
-            "The file paths provided could not be read as a private and public "
-            "key pair. Client authentication requires an elliptic curve public and "
-            "private key pair. Please provide the file paths containing elliptic "
-            "curve private and public keys to '--authentication-keys'."
+            "Error: Unable to parse the public key file in "
+            "'--auth-supernode-public-key'. Authentication requires elliptic "
+            "curve private and public key pair. Please ensure that the file "
+            "path points to a valid public key file and try again."
         )
 
     return (
-        client_private_key,
-        client_public_key,
+        ssh_private_key,
+        ssh_public_key,
     )

flwr/common/object_ref.py

@@ -30,6 +30,7 @@ attribute.
 
 def validate(
     module_attribute_str: str,
+    check_module: bool = True,
 ) -> Tuple[bool, Optional[str]]:
     """Validate object reference.
 
@@ -56,15 +57,18 @@ def validate(
             f"Missing attribute in {module_attribute_str}{OBJECT_REF_HELP_STR}",
         )
 
-    # Load module
-    module = find_spec(module_str)
-    if module and module.origin:
-        if not _find_attribute_in_module(module.origin, attributes_str):
-            return (
-                False,
-                f"Unable to find attribute {attributes_str} in module {module_str}"
-                f"{OBJECT_REF_HELP_STR}",
-            )
+    if check_module:
+        # Load module
+        module = find_spec(module_str)
+        if module and module.origin:
+            if not _find_attribute_in_module(module.origin, attributes_str):
+                return (
+                    False,
+                    f"Unable to find attribute {attributes_str} in module {module_str}"
+                    f"{OBJECT_REF_HELP_STR}",
+                )
+            return (True, None)
+    else:
         return (True, None)
 
     return (

flwr/common/recordset_compat.py

@@ -35,6 +35,8 @@ from .typing import (
     Status,
 )
 
+EMPTY_TENSOR_KEY = "_empty"
+
 
 def parametersrecord_to_parameters(
     record: ParametersRecord, keep_input: bool
@@ -59,7 +61,8 @@ def parametersrecord_to_parameters(
     parameters = Parameters(tensors=[], tensor_type="")
 
     for key in list(record.keys()):
-        parameters.tensors.append(record[key].data)
+        if key != EMPTY_TENSOR_KEY:
+            parameters.tensors.append(record[key].data)
 
         if not parameters.tensor_type:
             # Setting from first array in record. Recall the warning in the docstrings
@@ -103,6 +106,10 @@ def parameters_to_parametersrecord(
             data=tensor, dtype="", stype=tensor_type, shape=[]
         )
 
+    if num_arrays == 0:
+        ordered_dict[EMPTY_TENSOR_KEY] = Array(
+            data=b"", dtype="", stype=tensor_type, shape=[]
+        )
     return ParametersRecord(ordered_dict, keep_input=keep_input)
 
 

flwr/common/secure_aggregation/crypto/symmetric_encryption.py

@@ -117,18 +117,3 @@ def verify_hmac(key: bytes, message: bytes, hmac_value: bytes) -> bool:
         return True
     except InvalidSignature:
         return False
-
-
-def ssh_types_to_elliptic_curve(
-    private_key: serialization.SSHPrivateKeyTypes,
-    public_key: serialization.SSHPublicKeyTypes,
-) -> Tuple[ec.EllipticCurvePrivateKey, ec.EllipticCurvePublicKey]:
-    """Cast SSH key types to elliptic curve."""
-    if isinstance(private_key, ec.EllipticCurvePrivateKey) and isinstance(
-        public_key, ec.EllipticCurvePublicKey
-    ):
-        return (private_key, public_key)
-
-    raise TypeError(
-        "The provided key is not an EllipticCurvePrivateKey or EllipticCurvePublicKey"
-    )