flwr-nightly 1.15.0.dev20250104__py3-none-any.whl → 1.15.0.dev20250123__py3-none-any.whl

flwr/client/app.py

@@ -15,13 +15,14 @@
 """Flower client app."""
 
 
-import signal
-import subprocess
+import multiprocessing
+import os
 import sys
+import threading
 import time
 from contextlib import AbstractContextManager
-from dataclasses import dataclass
 from logging import ERROR, INFO, WARN
+from os import urandom
 from pathlib import Path
 from typing import Callable, Optional, Union, cast
 
@@ -33,6 +34,7 @@ from flwr.cli.config_utils import get_fab_metadata
 from flwr.cli.install import install_from_fab
 from flwr.client.client import Client
 from flwr.client.client_app import ClientApp, LoadClientAppError
+from flwr.client.clientapp.app import flwr_clientapp
 from flwr.client.nodestate.nodestate_factory import NodeStateFactory
 from flwr.client.typing import ClientFnExt
 from flwr.common import GRPC_MAX_MESSAGE_LENGTH, Context, EventType, Message, event
@@ -43,7 +45,6 @@ from flwr.common.constant import (
     ISOLATION_MODE_PROCESS,
     ISOLATION_MODE_SUBPROCESS,
     MAX_RETRY_DELAY,
-    MISSING_EXTRA_REST,
     RUN_ID_NUM_BYTES,
     SERVER_OCTET,
     TRANSPORT_TYPE_GRPC_ADAPTER,
@@ -53,13 +54,13 @@ from flwr.common.constant import (
     TRANSPORT_TYPES,
     ErrorCode,
 )
+from flwr.common.exit import ExitCode, flwr_exit
+from flwr.common.grpc import generic_create_grpc_server
 from flwr.common.logger import log, warn_deprecated_feature
 from flwr.common.message import Error
 from flwr.common.retry_invoker import RetryInvoker, RetryState, exponential
 from flwr.common.typing import Fab, Run, RunNotRunningException, UserConfig
 from flwr.proto.clientappio_pb2_grpc import add_ClientAppIoServicer_to_server
-from flwr.server.superlink.fleet.grpc_bidi.grpc_server import generic_create_grpc_server
-from flwr.server.superlink.linkstate.utils import generate_rand_int_from_bytes
 
 from .clientapp.clientappio_servicer import ClientAppInputs, ClientAppIoServicer
 from .grpc_adapter_client.connection import grpc_adapter
@@ -345,10 +346,7 @@ 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,
@@ -358,7 +356,6 @@ def start_client_internal(
             )
 
     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:
@@ -391,10 +388,11 @@ def start_client_internal(
     run_info_store: Optional[DeprecatedRunInfoStore] = None
     state_factory = NodeStateFactory()
     state = state_factory.state()
+    mp_spawn_context = multiprocessing.get_context("spawn")
 
     runs: dict[int, Run] = {}
 
-    while not app_state_tracker.interrupt:
+    while True:
         sleep_duration: int = 0
         with connection(
             address,
@@ -433,9 +431,8 @@ def start_client_internal(
                         node_config=node_config,
                     )
 
-            app_state_tracker.register_signal_handler()
             # pylint: disable=too-many-nested-blocks
-            while not app_state_tracker.interrupt:
+            while True:
                 try:
                     # Receive
                     message = receive()
@@ -513,7 +510,7 @@ def start_client_internal(
                             #    Docker container.
 
                             # Generate SuperNode token
-                            token: int = generate_rand_int_from_bytes(RUN_ID_NUM_BYTES)
+                            token = int.from_bytes(urandom(RUN_ID_NUM_BYTES), "little")
 
                             # Mode 1: SuperNode starts ClientApp as subprocess
                             start_subprocess = isolation == ISOLATION_MODE_SUBPROCESS
@@ -549,12 +546,13 @@ def start_client_internal(
                                 ]
                                 command.append("--insecure")
 
-                                subprocess.run(
-                                    command,
-                                    stdout=None,
-                                    stderr=None,
-                                    check=True,
+                                proc = mp_spawn_context.Process(
+                                    target=_run_flwr_clientapp,
+                                    args=(command, os.getpid()),
+                                    daemon=True,
                                 )
+                                proc.start()
+                                proc.join()
                             else:
                                 # Wait for output to become available
                                 while not clientappio_servicer.has_outputs():
@@ -592,10 +590,7 @@ def start_client_internal(
                             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
-                            )
+                        log(ERROR, "%s raised an exception", exc_entity, exc_info=ex)
 
                         # Create error message
                         reply_message = message.create_error_reply(
@@ -621,19 +616,14 @@ def start_client_internal(
                         run_id,
                     )
                     log(INFO, "")
-
-                except StopIteration:
-                    sleep_duration = 0
-                    break
             # pylint: enable=too-many-nested-blocks
 
             # Unregister node
-            if delete_node is not None and app_state_tracker.is_connected:
+            if delete_node is not None:
                 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
@@ -773,7 +763,10 @@ def _init_connection(transport: Optional[str], server_address: str) -> tuple[
     # Parse IP address
     parsed_address = parse_address(server_address)
     if not parsed_address:
-        sys.exit(f"Server address ({server_address}) cannot be parsed.")
+        flwr_exit(
+            ExitCode.COMMON_ADDRESS_INVALID,
+            f"SuperLink address ({server_address}) cannot be parsed.",
+        )
     host, port, is_v6 = parsed_address
     address = f"[{host}]:{port}" if is_v6 else f"{host}:{port}"
 
@@ -788,12 +781,9 @@ def _init_connection(transport: Optional[str], server_address: str) -> tuple[
 
             from .rest_client.connection import http_request_response
         except ModuleNotFoundError:
-            sys.exit(MISSING_EXTRA_REST)
+            flwr_exit(ExitCode.COMMON_MISSING_EXTRA_REST)
         if server_address[:4] != "http":
-            sys.exit(
-                "When using the REST API, please provide `https://` or "
-                "`http://` before the server address (e.g. `http://127.0.0.1:8080`)"
-            )
+            flwr_exit(ExitCode.SUPERNODE_REST_ADDRESS_INVALID)
         connection, error_type = http_request_response, RequestsConnectionError
     elif transport == TRANSPORT_TYPE_GRPC_RERE:
         connection, error_type = grpc_request_response, RpcError
@@ -809,21 +799,19 @@ 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 _run_flwr_clientapp(args: list[str], main_pid: int) -> None:
+    # Monitor the main process in case of SIGKILL
+    def main_process_monitor() -> None:
+        while True:
+            time.sleep(1)
+            if os.getppid() != main_pid:
+                os.kill(os.getpid(), 9)
 
-        def signal_handler(sig, frame):  # type: ignore
-            # pylint: disable=unused-argument
-            self.interrupt = True
-            raise StopIteration from None
+    threading.Thread(target=main_process_monitor, daemon=True).start()
 
-        signal.signal(signal.SIGINT, signal_handler)
-        signal.signal(signal.SIGTERM, signal_handler)
+    # Run the command
+    sys.argv = args
+    flwr_clientapp()
 
 
 def run_clientappio_api_grpc(

flwr/client/clientapp/app.py

@@ -16,7 +16,6 @@
 
 
 import argparse
-import sys
 import time
 from logging import DEBUG, ERROR, INFO
 from typing import Optional
@@ -29,6 +28,7 @@ from flwr.common import Context, Message
 from flwr.common.args import add_args_flwr_app_common
 from flwr.common.config import get_flwr_dir
 from flwr.common.constant import CLIENTAPPIO_API_DEFAULT_CLIENT_ADDRESS, ErrorCode
+from flwr.common.exit import ExitCode, flwr_exit
 from flwr.common.grpc import create_channel
 from flwr.common.logger import log
 from flwr.common.message import Error
@@ -61,12 +61,10 @@ def flwr_clientapp() -> None:
     """Run process-isolated Flower ClientApp."""
     args = _parse_args_run_flwr_clientapp().parse_args()
     if not args.insecure:
-        log(
-            ERROR,
-            "flwr-clientapp does not support TLS yet. "
-            "Please use the '--insecure' flag.",
+        flwr_exit(
+            ExitCode.COMMON_TLS_NOT_SUPPORTED,
+            "flwr-clientapp does not support TLS yet.",
         )
-        sys.exit(1)
 
     log(INFO, "Starting Flower ClientApp")
     log(

flwr/client/clientapp/utils.py

@@ -66,7 +66,7 @@ def get_load_client_app_fn(
         # `fab_hash` is not required since the app is loaded from `runtime_app_dir`.
         elif app_path is not None:
             config = get_project_config(runtime_app_dir)
-            this_fab_version, this_fab_id = get_metadata_from_config(config)
+            this_fab_id, this_fab_version = get_metadata_from_config(config)
 
             if this_fab_version != fab_version or this_fab_id != fab_id:
                 raise LoadClientAppError(

flwr/client/grpc_client/connection.py

@@ -47,12 +47,6 @@ from flwr.proto.transport_pb2 import (  # pylint: disable=E0611
 )
 from flwr.proto.transport_pb2_grpc import FlowerServiceStub  # pylint: disable=E0611
 
-# The following flags can be uncommented for debugging. Other possible values:
-# https://github.com/grpc/grpc/blob/master/doc/environment_variables.md
-# import os
-# os.environ["GRPC_VERBOSITY"] = "debug"
-# os.environ["GRPC_TRACE"] = "tcp,http"
-
 
 def on_channel_state_change(channel_connectivity: str) -> None:
     """Log channel connectivity."""

flwr/client/grpc_rere_client/client_interceptor.py

@@ -15,67 +15,18 @@
 """Flower client interceptor."""
 
 
-import base64
-import collections
-from collections.abc import Sequence
-from logging import WARNING
-from typing import Any, Callable, Optional, Union
+from typing import Any, Callable
 
 import grpc
 from cryptography.hazmat.primitives.asymmetric import ec
+from google.protobuf.message import Message as GrpcMessage
 
-from flwr.common.logger import log
+from flwr.common import now
+from flwr.common.constant import PUBLIC_KEY_HEADER, SIGNATURE_HEADER, TIMESTAMP_HEADER
 from flwr.common.secure_aggregation.crypto.symmetric_encryption import (
-    bytes_to_public_key,
-    compute_hmac,
-    generate_shared_key,
     public_key_to_bytes,
+    sign_message,
 )
-from flwr.proto.fab_pb2 import GetFabRequest  # pylint: disable=E0611
-from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
-    CreateNodeRequest,
-    DeleteNodeRequest,
-    PingRequest,
-    PullTaskInsRequest,
-    PushTaskResRequest,
-)
-from flwr.proto.run_pb2 import GetRunRequest  # pylint: disable=E0611
-
-_PUBLIC_KEY_HEADER = "public-key"
-_AUTH_TOKEN_HEADER = "auth-token"
-
-Request = Union[
-    CreateNodeRequest,
-    DeleteNodeRequest,
-    PullTaskInsRequest,
-    PushTaskResRequest,
-    GetRunRequest,
-    PingRequest,
-    GetFabRequest,
-]
-
-
-def _get_value_from_tuples(
-    key_string: str, tuples: Sequence[tuple[str, Union[str, bytes]]]
-) -> bytes:
-    value = next((value for key, value in tuples if key == key_string), "")
-    if isinstance(value, str):
-        return value.encode()
-
-    return value
-
-
-class _ClientCallDetails(
-    collections.namedtuple(
-        "_ClientCallDetails", ("method", "timeout", "metadata", "credentials")
-    ),
-    grpc.ClientCallDetails,  # type: ignore
-):
-    """Details for each client call.
-
-    The class will be passed on as the first argument in continuation function.
-    In our case, `AuthenticateClientInterceptor` adds new metadata to the construct.
-    """
 
 
 class AuthenticateClientInterceptor(grpc.UnaryUnaryClientInterceptor):  # type: ignore
@@ -87,84 +38,33 @@ class AuthenticateClientInterceptor(grpc.UnaryUnaryClientInterceptor):  # type:
         public_key: ec.EllipticCurvePublicKey,
     ):
         self.private_key = private_key
-        self.public_key = public_key
-        self.shared_secret: Optional[bytes] = None
-        self.server_public_key: Optional[ec.EllipticCurvePublicKey] = None
-        self.encoded_public_key = base64.urlsafe_b64encode(
-            public_key_to_bytes(self.public_key)
-        )
+        self.public_key_bytes = public_key_to_bytes(public_key)
 
     def intercept_unary_unary(
         self,
         continuation: Callable[[Any, Any], Any],
         client_call_details: grpc.ClientCallDetails,
-        request: Request,
+        request: GrpcMessage,
     ) -> grpc.Call:
         """Flower client interceptor.
 
         Intercept unary call from client and add necessary authentication header in the
         RPC metadata.
         """
-        metadata = []
-        postprocess = False
-        if client_call_details.metadata is not None:
-            metadata = list(client_call_details.metadata)
-
-        # Always add the public key header
-        metadata.append(
-            (
-                _PUBLIC_KEY_HEADER,
-                self.encoded_public_key,
-            )
-        )
-
-        if isinstance(request, CreateNodeRequest):
-            postprocess = True
-        elif isinstance(
-            request,
-            (
-                DeleteNodeRequest,
-                PullTaskInsRequest,
-                PushTaskResRequest,
-                GetRunRequest,
-                PingRequest,
-                GetFabRequest,
-            ),
-        ):
-            if self.shared_secret is None:
-                raise RuntimeError("Failure to compute hmac")
-
-            message_bytes = request.SerializeToString(deterministic=True)
-            metadata.append(
-                (
-                    _AUTH_TOKEN_HEADER,
-                    base64.urlsafe_b64encode(
-                        compute_hmac(self.shared_secret, message_bytes)
-                    ),
-                )
-            )
+        metadata = list(client_call_details.metadata or [])
 
-        client_call_details = _ClientCallDetails(
-            client_call_details.method,
-            client_call_details.timeout,
-            metadata,
-            client_call_details.credentials,
-        )
+        # Add the public key
+        metadata.append((PUBLIC_KEY_HEADER, self.public_key_bytes))
 
-        response = continuation(client_call_details, request)
-        if postprocess:
-            server_public_key_bytes = base64.urlsafe_b64decode(
-                _get_value_from_tuples(_PUBLIC_KEY_HEADER, response.initial_metadata())
-            )
+        # Add timestamp
+        timestamp = now().isoformat()
+        metadata.append((TIMESTAMP_HEADER, timestamp))
 
-            if server_public_key_bytes != b"":
-                self.server_public_key = bytes_to_public_key(server_public_key_bytes)
-            else:
-                log(WARNING, "Can't get server public key, SuperLink may be offline")
+        # Sign and add the signature
+        signature = sign_message(self.private_key, timestamp.encode("ascii"))
+        metadata.append((SIGNATURE_HEADER, signature))
 
-            if self.server_public_key is not None:
-                self.shared_secret = generate_shared_key(
-                    self.private_key, self.server_public_key
-                )
+        # Overwrite the metadata
+        details = client_call_details._replace(metadata=metadata)
 
-        return response
+        return continuation(details, request)

flwr/client/grpc_rere_client/connection.py

@@ -29,7 +29,6 @@ from cryptography.hazmat.primitives.asymmetric import ec
 
 from flwr.client.heartbeat import start_ping_loop
 from flwr.client.message_handler.message_handler import validate_out_message
-from flwr.client.message_handler.task_handler import get_task_ins, validate_task_ins
 from flwr.common import GRPC_MAX_MESSAGE_LENGTH
 from flwr.common.constant import (
     PING_BASE_MULTIPLIER,
@@ -41,7 +40,7 @@ from flwr.common.grpc import create_channel
 from flwr.common.logger import log
 from flwr.common.message import Message, Metadata
 from flwr.common.retry_invoker import RetryInvoker
-from flwr.common.serde import message_from_taskins, message_to_taskres, run_from_proto
+from flwr.common.serde import message_from_proto, message_to_proto, run_from_proto
 from flwr.common.typing import Fab, Run, RunNotRunningException
 from flwr.proto.fab_pb2 import GetFabRequest, GetFabResponse  # pylint: disable=E0611
 from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
@@ -49,13 +48,13 @@ from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
     DeleteNodeRequest,
     PingRequest,
     PingResponse,
-    PullTaskInsRequest,
-    PushTaskResRequest,
+    PullMessagesRequest,
+    PullMessagesResponse,
+    PushMessagesRequest,
 )
 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
 from .grpc_adapter import GrpcAdapter
@@ -227,28 +226,31 @@ def grpc_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
         node = None
 
     def receive() -> Optional[Message]:
-        """Receive next task from server."""
+        """Receive next message from server."""
         # Get Node
         if node is None:
             log(ERROR, "Node instance missing")
             return None
 
-        # Request instructions (task) from server
-        request = PullTaskInsRequest(node=node)
-        response = retry_invoker.invoke(stub.PullTaskIns, request=request)
+        # Request instructions (message) from server
+        request = PullMessagesRequest(node=node)
+        response: PullMessagesResponse = retry_invoker.invoke(
+            stub.PullMessages, request=request
+        )
 
-        # Get the current TaskIns
-        task_ins: Optional[TaskIns] = get_task_ins(response)
+        # Get the current Messages
+        message_proto = (
+            None if len(response.messages_list) == 0 else response.messages_list[0]
+        )
 
-        # Discard the current TaskIns if not valid
-        if task_ins is not None and not (
-            task_ins.task.consumer.node_id == node.node_id
-            and validate_task_ins(task_ins)
+        # Discard the current message if not valid
+        if message_proto is not None and not (
+            message_proto.metadata.dst_node_id == node.node_id
         ):
-            task_ins = None
+            message_proto = None
 
         # Construct the Message
-        in_message = message_from_taskins(task_ins) if task_ins else None
+        in_message = message_from_proto(message_proto) if message_proto else None
 
         # Remember `metadata` of the in message
         nonlocal metadata
@@ -258,7 +260,7 @@ def grpc_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
         return in_message
 
     def send(message: Message) -> None:
-        """Send task result back to server."""
+        """Send message reply to server."""
         # Get Node
         if node is None:
             log(ERROR, "Node instance missing")
@@ -275,12 +277,10 @@ def grpc_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
             log(ERROR, "Invalid out message")
             return
 
-        # Construct TaskRes
-        task_res = message_to_taskres(message)
-
-        # Serialize ProtoBuf to bytes
-        request = PushTaskResRequest(node=node, task_res_list=[task_res])
-        _ = retry_invoker.invoke(stub.PushTaskRes, request)
+        # Serialize Message
+        message_proto = message_to_proto(message=message)
+        request = PushMessagesRequest(node=node, messages_list=[message_proto])
+        _ = retry_invoker.invoke(stub.PushMessages, request)
 
         # Cleanup
         metadata = None
@@ -311,3 +311,13 @@ def grpc_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
         yield (receive, send, create_node, delete_node, get_run, get_fab)
     except Exception as exc:  # pylint: disable=broad-except
         log(ERROR, exc)
+    # Cleanup
+    finally:
+        try:
+            if node is not None:
+                # Disable retrying
+                retry_invoker.max_tries = 1
+                delete_node()
+        except grpc.RpcError:
+            pass
+        channel.close()

flwr/client/grpc_rere_client/grpc_adapter.py

@@ -40,8 +40,12 @@ from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
     DeleteNodeResponse,
     PingRequest,
     PingResponse,
+    PullMessagesRequest,
+    PullMessagesResponse,
     PullTaskInsRequest,
     PullTaskInsResponse,
+    PushMessagesRequest,
+    PushMessagesResponse,
     PushTaskResRequest,
     PushTaskResResponse,
 )
@@ -132,12 +136,24 @@ class GrpcAdapter:
         """."""
         return self._send_and_receive(request, PullTaskInsResponse, **kwargs)
 
+    def PullMessages(  # pylint: disable=C0103
+        self, request: PullMessagesRequest, **kwargs: Any
+    ) -> PullMessagesResponse:
+        """."""
+        return self._send_and_receive(request, PullMessagesResponse, **kwargs)
+
     def PushTaskRes(  # pylint: disable=C0103
         self, request: PushTaskResRequest, **kwargs: Any
     ) -> PushTaskResResponse:
         """."""
         return self._send_and_receive(request, PushTaskResResponse, **kwargs)
 
+    def PushMessages(  # pylint: disable=C0103
+        self, request: PushMessagesRequest, **kwargs: Any
+    ) -> PushMessagesResponse:
+        """."""
+        return self._send_and_receive(request, PushMessagesResponse, **kwargs)
+
     def GetRun(  # pylint: disable=C0103
         self, request: GetRunRequest, **kwargs: Any
     ) -> GetRunResponse: