flwr-nightly 1.9.0.dev20240417__py3-none-any.whl → 1.9.0.dev20240507__py3-none-any.whl

flwr/client/rest_client/connection.py

@@ -21,7 +21,10 @@ import threading
 from contextlib import contextmanager
 from copy import copy
 from logging import ERROR, INFO, WARN
-from typing import Callable, Iterator, Optional, Tuple, Union
+from typing import Callable, Iterator, Optional, Tuple, Type, TypeVar, Union
+
+from cryptography.hazmat.primitives.asymmetric import ec
+from google.protobuf.message import Message as GrpcMessage
 
 from flwr.client.heartbeat import start_ping_loop
 from flwr.client.message_handler.message_handler import validate_out_message
@@ -42,6 +45,9 @@ from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
     CreateNodeRequest,
     CreateNodeResponse,
     DeleteNodeRequest,
+    DeleteNodeResponse,
+    GetRunRequest,
+    GetRunResponse,
     PingRequest,
     PingResponse,
     PullTaskInsRequest,
@@ -63,10 +69,13 @@ PATH_DELETE_NODE: str = "api/v0/fleet/delete-node"
 PATH_PULL_TASK_INS: str = "api/v0/fleet/pull-task-ins"
 PATH_PUSH_TASK_RES: str = "api/v0/fleet/push-task-res"
 PATH_PING: str = "api/v0/fleet/ping"
+PATH_GET_RUN: str = "/api/v0/fleet/get-run"
+
+T = TypeVar("T", bound=GrpcMessage)
 
 
 @contextmanager
-def http_request_response(  # pylint: disable=R0914, R0915
+def http_request_response(  # pylint: disable=,R0913, R0914, R0915
     server_address: str,
     insecure: bool,  # pylint: disable=unused-argument
     retry_invoker: RetryInvoker,
@@ -74,12 +83,16 @@ def http_request_response(  # pylint: disable=R0914, R0915
     root_certificates: Optional[
         Union[bytes, str]
     ] = None,  # pylint: disable=unused-argument
+    authentication_keys: Optional[  # pylint: disable=unused-argument
+        Tuple[ec.EllipticCurvePrivateKey, ec.EllipticCurvePublicKey]
+    ] = None,
 ) -> Iterator[
     Tuple[
         Callable[[], Optional[Message]],
         Callable[[Message], None],
         Optional[Callable[[], None]],
         Optional[Callable[[], None]],
+        Optional[Callable[[int], Tuple[str, str]]],
     ]
 ]:
     """Primitives for request/response-based interaction with a server.
@@ -141,55 +154,72 @@ def http_request_response(  # pylint: disable=R0914, R0915
     ping_stop_event = threading.Event()
 
     ###########################################################################
-    # ping/create_node/delete_node/receive/send functions
+    # ping/create_node/delete_node/receive/send/get_run functions
     ###########################################################################
 
-    def ping() -> None:
-        # Get Node
-        if node is None:
-            log(ERROR, "Node instance missing")
-            return
-
-        # Construct the ping request
-        req = PingRequest(node=node, ping_interval=PING_DEFAULT_INTERVAL)
-        req_bytes: bytes = req.SerializeToString()
+    def _request(
+        req: GrpcMessage, res_type: Type[T], api_path: str, retry: bool = True
+    ) -> Optional[T]:
+        # Serialize the request
+        req_bytes = req.SerializeToString()
 
         # Send the request
-        res = requests.post(
-            url=f"{base_url}/{PATH_PING}",
-            headers={
-                "Accept": "application/protobuf",
-                "Content-Type": "application/protobuf",
-            },
-            data=req_bytes,
-            verify=verify,
-            timeout=PING_CALL_TIMEOUT,
-        )
+        def post() -> requests.Response:
+            return requests.post(
+                f"{base_url}/{api_path}",
+                data=req_bytes,
+                headers={
+                    "Accept": "application/protobuf",
+                    "Content-Type": "application/protobuf",
+                },
+                verify=verify,
+                timeout=None,
+            )
+
+        if retry:
+            res: requests.Response = retry_invoker.invoke(post)
+        else:
+            res = post()
 
         # Check status code and headers
         if res.status_code != 200:
-            return
+            return None
         if "content-type" not in res.headers:
             log(
                 WARN,
                 "[Node] POST /%s: missing header `Content-Type`",
-                PATH_PULL_TASK_INS,
+                api_path,
             )
-            return
+            return None
         if res.headers["content-type"] != "application/protobuf":
             log(
                 WARN,
                 "[Node] POST /%s: header `Content-Type` has wrong value",
-                PATH_PULL_TASK_INS,
+                api_path,
             )
-            return
+            return None
 
         # Deserialize ProtoBuf from bytes
-        ping_res = PingResponse()
-        ping_res.ParseFromString(res.content)
+        grpc_res = res_type()
+        grpc_res.ParseFromString(res.content)
+        return grpc_res
+
+    def ping() -> None:
+        # Get Node
+        if node is None:
+            log(ERROR, "Node instance missing")
+            return
+
+        # Construct the ping request
+        req = PingRequest(node=node, ping_interval=PING_DEFAULT_INTERVAL)
+
+        # Send the request
+        res = _request(req, PingResponse, PATH_PING, retry=False)
+        if res is None:
+            return
 
         # Check if success
-        if not ping_res.success:
+        if not res.success:
             raise RuntimeError("Ping failed unexpectedly.")
 
         # Wait
@@ -201,46 +231,16 @@ def http_request_response(  # pylint: disable=R0914, R0915
 
     def create_node() -> None:
         """Set create_node."""
-        create_node_req_proto = CreateNodeRequest(ping_interval=PING_DEFAULT_INTERVAL)
-        create_node_req_bytes: bytes = create_node_req_proto.SerializeToString()
-
-        res = retry_invoker.invoke(
-            requests.post,
-            url=f"{base_url}/{PATH_CREATE_NODE}",
-            headers={
-                "Accept": "application/protobuf",
-                "Content-Type": "application/protobuf",
-            },
-            data=create_node_req_bytes,
-            verify=verify,
-            timeout=None,
-        )
+        req = CreateNodeRequest(ping_interval=PING_DEFAULT_INTERVAL)
 
-        # Check status code and headers
-        if res.status_code != 200:
-            return
-        if "content-type" not in res.headers:
-            log(
-                WARN,
-                "[Node] POST /%s: missing header `Content-Type`",
-                PATH_PULL_TASK_INS,
-            )
-            return
-        if res.headers["content-type"] != "application/protobuf":
-            log(
-                WARN,
-                "[Node] POST /%s: header `Content-Type` has wrong value",
-                PATH_PULL_TASK_INS,
-            )
+        # Send the request
+        res = _request(req, CreateNodeResponse, PATH_CREATE_NODE)
+        if res is None:
             return
 
-        # Deserialize ProtoBuf from bytes
-        create_node_response_proto = CreateNodeResponse()
-        create_node_response_proto.ParseFromString(res.content)
-
         # Remember the node and the ping-loop thread
         nonlocal node, ping_thread
-        node = create_node_response_proto.node
+        node = res.node
         ping_thread = start_ping_loop(ping, ping_stop_event)
 
     def delete_node() -> None:
@@ -256,36 +256,12 @@ def http_request_response(  # pylint: disable=R0914, R0915
             ping_thread.join()
 
         # Send DeleteNode request
-        delete_node_req_proto = DeleteNodeRequest(node=node)
-        delete_node_req_req_bytes: bytes = delete_node_req_proto.SerializeToString()
-        res = retry_invoker.invoke(
-            requests.post,
-            url=f"{base_url}/{PATH_DELETE_NODE}",
-            headers={
-                "Accept": "application/protobuf",
-                "Content-Type": "application/protobuf",
-            },
-            data=delete_node_req_req_bytes,
-            verify=verify,
-            timeout=None,
-        )
+        req = DeleteNodeRequest(node=node)
 
-        # Check status code and headers
-        if res.status_code != 200:
-            return
-        if "content-type" not in res.headers:
-            log(
-                WARN,
-                "[Node] POST /%s: missing header `Content-Type`",
-                PATH_PULL_TASK_INS,
-            )
+        # Send the request
+        res = _request(req, DeleteNodeResponse, PATH_CREATE_NODE)
+        if res is None:
             return
-        if res.headers["content-type"] != "application/protobuf":
-            log(
-                WARN,
-                "[Node] POST /%s: header `Content-Type` has wrong value",
-                PATH_PULL_TASK_INS,
-            )
 
         # Cleanup
         node = None
@@ -298,46 +274,15 @@ def http_request_response(  # pylint: disable=R0914, R0915
             return None
 
         # Request instructions (task) from server
-        pull_task_ins_req_proto = PullTaskInsRequest(node=node)
-        pull_task_ins_req_bytes: bytes = pull_task_ins_req_proto.SerializeToString()
+        req = PullTaskInsRequest(node=node)
 
-        # Request instructions (task) from server
-        res = retry_invoker.invoke(
-            requests.post,
-            url=f"{base_url}/{PATH_PULL_TASK_INS}",
-            headers={
-                "Accept": "application/protobuf",
-                "Content-Type": "application/protobuf",
-            },
-            data=pull_task_ins_req_bytes,
-            verify=verify,
-            timeout=None,
-        )
-
-        # Check status code and headers
-        if res.status_code != 200:
-            return None
-        if "content-type" not in res.headers:
-            log(
-                WARN,
-                "[Node] POST /%s: missing header `Content-Type`",
-                PATH_PULL_TASK_INS,
-            )
-            return None
-        if res.headers["content-type"] != "application/protobuf":
-            log(
-                WARN,
-                "[Node] POST /%s: header `Content-Type` has wrong value",
-                PATH_PULL_TASK_INS,
-            )
+        # Send the request
+        res = _request(req, PullTaskInsResponse, PATH_PULL_TASK_INS)
+        if res is None:
             return None
 
-        # Deserialize ProtoBuf from bytes
-        pull_task_ins_response_proto = PullTaskInsResponse()
-        pull_task_ins_response_proto.ParseFromString(res.content)
-
         # Get the current TaskIns
-        task_ins: Optional[TaskIns] = get_task_ins(pull_task_ins_response_proto)
+        task_ins: Optional[TaskIns] = get_task_ins(res)
 
         # Discard the current TaskIns if not valid
         if task_ins is not None and not (
@@ -372,61 +317,39 @@ def http_request_response(  # pylint: disable=R0914, R0915
         if not validate_out_message(message, metadata):
             log(ERROR, "Invalid out message")
             return
+        metadata = None
 
         # Construct TaskRes
         task_res = message_to_taskres(message)
 
         # Serialize ProtoBuf to bytes
-        push_task_res_request_proto = PushTaskResRequest(task_res_list=[task_res])
-        push_task_res_request_bytes: bytes = (
-            push_task_res_request_proto.SerializeToString()
-        )
-
-        # Send ClientMessage to server
-        res = retry_invoker.invoke(
-            requests.post,
-            url=f"{base_url}/{PATH_PUSH_TASK_RES}",
-            headers={
-                "Accept": "application/protobuf",
-                "Content-Type": "application/protobuf",
-            },
-            data=push_task_res_request_bytes,
-            verify=verify,
-            timeout=None,
-        )
-
-        metadata = None
+        req = PushTaskResRequest(task_res_list=[task_res])
 
-        # Check status code and headers
-        if res.status_code != 200:
-            return
-        if "content-type" not in res.headers:
-            log(
-                WARN,
-                "[Node] POST /%s: missing header `Content-Type`",
-                PATH_PUSH_TASK_RES,
-            )
-            return
-        if res.headers["content-type"] != "application/protobuf":
-            log(
-                WARN,
-                "[Node] POST /%s: header `Content-Type` has wrong value",
-                PATH_PUSH_TASK_RES,
-            )
+        # Send the request
+        res = _request(req, PushTaskResResponse, PATH_PUSH_TASK_RES)
+        if res is None:
             return
 
-        # Deserialize ProtoBuf from bytes
-        push_task_res_response_proto = PushTaskResResponse()
-        push_task_res_response_proto.ParseFromString(res.content)
         log(
             INFO,
             "[Node] POST /%s: success, created result %s",
             PATH_PUSH_TASK_RES,
-            push_task_res_response_proto.results,  # pylint: disable=no-member
+            res.results,  # pylint: disable=no-member
         )
 
+    def get_run(run_id: int) -> Tuple[str, str]:
+        # Construct the request
+        req = GetRunRequest(run_id=run_id)
+
+        # Send the request
+        res = _request(req, GetRunResponse, PATH_GET_RUN)
+        if res is None:
+            return "", ""
+
+        return res.run.fab_id, res.run.fab_version
+
     try:
         # Yield methods
-        yield (receive, send, create_node, delete_node)
+        yield (receive, send, create_node, delete_node, get_run)
     except Exception as exc:  # pylint: disable=broad-except
         log(ERROR, exc)

flwr/client/supernode/__init__.py

@@ -0,0 +1,24 @@
+# Copyright 2024 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Flower SuperNode."""
+
+
+from .app import run_client_app as run_client_app
+from .app import run_supernode as run_supernode
+
+__all__ = [
+    "run_client_app",
+    "run_supernode",
+]

flwr/client/supernode/app.py

@@ -0,0 +1,281 @@
+# Copyright 2024 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Flower SuperNode."""
+
+import argparse
+import sys
+from logging import DEBUG, INFO, WARN
+from pathlib import Path
+from typing import Callable, Optional, Tuple
+
+from cryptography.hazmat.primitives.asymmetric import ec
+from cryptography.hazmat.primitives.serialization import (
+    load_ssh_private_key,
+    load_ssh_public_key,
+)
+
+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.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
+
+
+def run_supernode() -> None:
+    """Run Flower SuperNode."""
+    log(INFO, "Starting Flower SuperNode")
+
+    event(EventType.RUN_SUPERNODE_ENTER)
+
+    _ = _parse_args_run_supernode().parse_args()
+
+    log(
+        DEBUG,
+        "Flower SuperNode starting...",
+    )
+
+    # Graceful shutdown
+    register_exit_handlers(
+        event_type=EventType.RUN_SUPERNODE_LEAVE,
+    )
+
+
+def run_client_app() -> None:
+    """Run Flower client app."""
+    log(INFO, "Long-running Flower client starting")
+
+    event(EventType.RUN_CLIENT_APP_ENTER)
+
+    args = _parse_args_run_client_app().parse_args()
+
+    root_certificates = _get_certificates(args)
+    log(
+        DEBUG,
+        "Flower will load ClientApp `%s`",
+        getattr(args, "client-app"),
+    )
+    load_fn = _get_load_client_app_fn(args)
+    authentication_keys = _try_setup_client_authentication(args)
+
+    _start_client_internal(
+        server_address=args.server,
+        load_client_app_fn=load_fn,
+        transport="rest" if args.rest else "grpc-rere",
+        root_certificates=root_certificates,
+        insecure=args.insecure,
+        authentication_keys=authentication_keys,
+        max_retries=args.max_retries,
+        max_wait_time=args.max_wait_time,
+    )
+    register_exit_handlers(event_type=EventType.RUN_CLIENT_APP_LEAVE)
+
+
+def _get_certificates(args: argparse.Namespace) -> Optional[bytes]:
+    """Load certificates if specified in args."""
+    # Obtain certificates
+    if args.insecure:
+        if args.root_certificates is not None:
+            sys.exit(
+                "Conflicting options: The '--insecure' flag disables HTTPS, "
+                "but '--root-certificates' was also specified. Please remove "
+                "the '--root-certificates' option when running in insecure mode, "
+                "or omit '--insecure' to use HTTPS."
+            )
+        log(
+            WARN,
+            "Option `--insecure` was set. "
+            "Starting insecure HTTP client connected to %s.",
+            args.server,
+        )
+        root_certificates = None
+    else:
+        # Load the certificates if provided, or load the system certificates
+        cert_path = args.root_certificates
+        if cert_path is None:
+            root_certificates = None
+        else:
+            root_certificates = Path(cert_path).read_bytes()
+        log(
+            DEBUG,
+            "Starting secure HTTPS client connected to %s "
+            "with the following certificates: %s.",
+            args.server,
+            cert_path,
+        )
+    return root_certificates
+
+
+def _get_load_client_app_fn(
+    args: argparse.Namespace,
+) -> Callable[[], ClientApp]:
+    """Get the load_client_app_fn function."""
+    client_app_dir = args.dir
+    if client_app_dir is not None:
+        sys.path.insert(0, client_app_dir)
+
+    app_ref: str = getattr(args, "client-app")
+    valid, error_msg = validate(app_ref)
+    if not valid and error_msg:
+        raise LoadClientAppError(error_msg) from None
+
+    def _load() -> ClientApp:
+        client_app = load_app(app_ref, LoadClientAppError)
+
+        if not isinstance(client_app, ClientApp):
+            raise LoadClientAppError(
+                f"Attribute {app_ref} is not of type {ClientApp}",
+            ) from None
+
+        return client_app
+
+    return _load
+
+
+def _parse_args_run_supernode() -> argparse.ArgumentParser:
+    """Parse flower-supernode command line arguments."""
+    parser = argparse.ArgumentParser(
+        description="Start a Flower SuperNode",
+    )
+
+    parser.add_argument(
+        "client-app",
+        nargs="?",
+        default="",
+        help="For example: `client:app` or `project.package.module:wrapper.app`. "
+        "This is optional and serves as the default ClientApp to be loaded when "
+        "the ServerApp does not specify `fab_id` and `fab_version`. "
+        "If not provided, defaults to an empty string.",
+    )
+    _parse_args_common(parser)
+    parser.add_argument(
+        "--flwr-dir",
+        default=None,
+        help="""The path containing installed Flower Apps.
+    By default, this value isequal to:
+
+        - `$FLWR_HOME/` if `$FLWR_HOME` is defined
+        - `$XDG_DATA_HOME/.flwr/` if `$XDG_DATA_HOME` is defined
+        - `$HOME/.flwr/` in all other cases
+    """,
+    )
+
+    return parser
+
+
+def _parse_args_run_client_app() -> argparse.ArgumentParser:
+    """Parse flower-client-app command line arguments."""
+    parser = argparse.ArgumentParser(
+        description="Start a Flower client app",
+    )
+
+    parser.add_argument(
+        "client-app",
+        help="For example: `client:app` or `project.package.module:wrapper.app`",
+    )
+    _parse_args_common(parser=parser)
+
+    return parser
+
+
+def _parse_args_common(parser: argparse.ArgumentParser) -> None:
+    parser.add_argument(
+        "--insecure",
+        action="store_true",
+        help="Run the client without HTTPS. By default, the client runs with "
+        "HTTPS enabled. Use this flag only if you understand the risks.",
+    )
+    parser.add_argument(
+        "--rest",
+        action="store_true",
+        help="Use REST as a transport layer for the client.",
+    )
+    parser.add_argument(
+        "--root-certificates",
+        metavar="ROOT_CERT",
+        type=str,
+        help="Specifies the path to the PEM-encoded root certificate file for "
+        "establishing secure HTTPS connections.",
+    )
+    parser.add_argument(
+        "--server",
+        default="0.0.0.0:9092",
+        help="Server address",
+    )
+    parser.add_argument(
+        "--max-retries",
+        type=int,
+        default=None,
+        help="The maximum number of times the client will try to connect to the"
+        "server before giving up in case of a connection error. By default,"
+        "it is set to None, meaning there is no limit to the number of tries.",
+    )
+    parser.add_argument(
+        "--max-wait-time",
+        type=float,
+        default=None,
+        help="The maximum duration before the client stops trying to"
+        "connect to the server in case of connection error. By default, it"
+        "is set to None, meaning there is no limit to the total time.",
+    )
+    parser.add_argument(
+        "--dir",
+        default="",
+        help="Add specified directory to the PYTHONPATH and load Flower "
+        "app from there."
+        " Default: current working directory.",
+    )
+    parser.add_argument(
+        "--authentication-keys",
+        nargs=2,
+        metavar=("CLIENT_PRIVATE_KEY", "CLIENT_PUBLIC_KEY"),
+        type=str,
+        help="Provide two file paths: (1) the client's private "
+        "key file, and (2) the client's public key file.",
+    )
+
+
+def _try_setup_client_authentication(
+    args: argparse.Namespace,
+) -> Optional[Tuple[ec.EllipticCurvePrivateKey, ec.EllipticCurvePublicKey]]:
+    if not args.authentication_keys:
+        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())
+
+    try:
+        client_private_key, client_public_key = ssh_types_to_elliptic_curve(
+            ssh_private_key, ssh_public_key
+        )
+    except TypeError:
+        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'."
+        )
+
+    return (
+        client_private_key,
+        client_public_key,
+    )