flwr-nightly 1.8.0.dev20240315__py3-none-any.whl → 1.15.0.dev20250114__py3-none-any.whl

flwr/client/numpy_client.py

@@ -16,12 +16,11 @@
 
 
 from abc import ABC
-from typing import Callable, Dict, Tuple
+from typing import Callable
 
 from flwr.client.client import Client
 from flwr.common import (
     Config,
-    Context,
     NDArrays,
     Scalar,
     ndarrays_to_parameters,
@@ -70,9 +69,7 @@ Example
 class NumPyClient(ABC):
     """Abstract base class for Flower clients using NumPy."""
 
-    context: Context
-
-    def get_properties(self, config: Config) -> Dict[str, Scalar]:
+    def get_properties(self, config: Config) -> dict[str, Scalar]:
         """Return a client's set of properties.
 
         Parameters
@@ -92,7 +89,7 @@ class NumPyClient(ABC):
         _ = (self, config)
         return {}
 
-    def get_parameters(self, config: Dict[str, Scalar]) -> NDArrays:
+    def get_parameters(self, config: dict[str, Scalar]) -> NDArrays:
         """Return the current local model parameters.
 
         Parameters
@@ -111,8 +108,8 @@ class NumPyClient(ABC):
         return []
 
     def fit(
-        self, parameters: NDArrays, config: Dict[str, Scalar]
-    ) -> Tuple[NDArrays, int, Dict[str, Scalar]]:
+        self, parameters: NDArrays, config: dict[str, Scalar]
+    ) -> tuple[NDArrays, int, dict[str, Scalar]]:
         """Train the provided parameters using the locally held dataset.
 
         Parameters
@@ -140,8 +137,8 @@ class NumPyClient(ABC):
         return [], 0, {}
 
     def evaluate(
-        self, parameters: NDArrays, config: Dict[str, Scalar]
-    ) -> Tuple[float, int, Dict[str, Scalar]]:
+        self, parameters: NDArrays, config: dict[str, Scalar]
+    ) -> tuple[float, int, dict[str, Scalar]]:
         """Evaluate the provided parameters using the locally held dataset.
 
         Parameters
@@ -174,14 +171,6 @@ class NumPyClient(ABC):
         _ = (self, parameters, config)
         return 0.0, 0, {}
 
-    def get_context(self) -> Context:
-        """Get the run context from this client."""
-        return self.context
-
-    def set_context(self, context: Context) -> None:
-        """Apply a run context to this client."""
-        self.context = context
-
     def to_client(self) -> Client:
         """Convert to object to Client type and return it."""
         return _wrap_numpy_client(client=self)
@@ -278,21 +267,9 @@ def _evaluate(self: Client, ins: EvaluateIns) -> EvaluateRes:
     )
 
 
-def _get_context(self: Client) -> Context:
-    """Return context of underlying NumPyClient."""
-    return self.numpy_client.get_context()  # type: ignore
-
-
-def _set_context(self: Client, context: Context) -> None:
-    """Apply context to underlying NumPyClient."""
-    self.numpy_client.set_context(context)  # type: ignore
-
-
 def _wrap_numpy_client(client: NumPyClient) -> Client:
-    member_dict: Dict[str, Callable] = {  # type: ignore
+    member_dict: dict[str, Callable] = {  # type: ignore
         "__init__": _constructor,
-        "get_context": _get_context,
-        "set_context": _set_context,
     }
 
     # Add wrapper type methods (if overridden)

flwr/client/rest_client/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2020 Flower Labs GmbH. All Rights Reserved.
+# Copyright 2023 Flower Labs GmbH. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

flwr/client/rest_client/connection.py

@@ -1,4 +1,4 @@
-# Copyright 2020 Flower Labs GmbH. All Rights Reserved.
+# Copyright 2023 Flower Labs GmbH. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -15,31 +15,48 @@
 """Contextmanager for a REST request-response channel to the Flower server."""
 
 
+import random
 import sys
+import threading
+from collections.abc import Iterator
 from contextlib import contextmanager
 from copy import copy
 from logging import ERROR, INFO, WARN
-from typing import Callable, Dict, Iterator, Optional, Tuple, Union, cast
+from typing import Callable, Optional, 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
-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 MISSING_EXTRA_REST
+from flwr.common.constant import (
+    MISSING_EXTRA_REST,
+    PING_BASE_MULTIPLIER,
+    PING_CALL_TIMEOUT,
+    PING_DEFAULT_INTERVAL,
+    PING_RANDOM_RANGE,
+)
 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
+from flwr.common.serde import 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,
-    PullTaskInsRequest,
-    PullTaskInsResponse,
-    PushTaskResRequest,
-    PushTaskResResponse,
+    DeleteNodeResponse,
+    PingRequest,
+    PingResponse,
+    PullMessagesRequest,
+    PullMessagesResponse,
+    PushMessagesRequest,
+    PushMessagesResponse,
 )
 from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
-from flwr.proto.task_pb2 import TaskIns  # pylint: disable=E0611
+from flwr.proto.run_pb2 import GetRunRequest, GetRunResponse  # pylint: disable=E0611
 
 try:
     import requests
@@ -47,19 +64,21 @@ except ModuleNotFoundError:
     sys.exit(MISSING_EXTRA_REST)
 
 
-KEY_NODE = "node"
-KEY_METADATA = "in_message_metadata"
-
-
 PATH_CREATE_NODE: str = "api/v0/fleet/create-node"
 PATH_DELETE_NODE: str = "api/v0/fleet/delete-node"
 PATH_PULL_TASK_INS: str = "api/v0/fleet/pull-task-ins"
+PATH_PULL_MESSAGES: str = "/api/v0/fleet/pull-messages"
 PATH_PUSH_TASK_RES: str = "api/v0/fleet/push-task-res"
+PATH_PUSH_MESSAGES: str = "/api/v0/fleet/push-messages"
+PATH_PING: str = "api/v0/fleet/ping"
+PATH_GET_RUN: str = "/api/v0/fleet/get-run"
+PATH_GET_FAB: str = "/api/v0/fleet/get-fab"
+
+T = TypeVar("T", bound=GrpcMessage)
 
 
 @contextmanager
-# pylint: disable-next=too-many-statements
-def http_request_response(
+def http_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
     server_address: str,
     insecure: bool,  # pylint: disable=unused-argument
     retry_invoker: RetryInvoker,
@@ -67,12 +86,17 @@ def http_request_response(
     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[
+    tuple[
         Callable[[], Optional[Message]],
         Callable[[Message], None],
+        Optional[Callable[[], Optional[int]]],
         Optional[Callable[[], None]],
-        Optional[Callable[[], None]],
+        Optional[Callable[[int], Run]],
+        Optional[Callable[[str, int], Fab]],
     ]
 ]:
     """Primitives for request/response-based interaction with a server.
@@ -98,10 +122,16 @@ def http_request_response(
         Path of the root certificate. If provided, a secure
         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.
 
     Returns
     -------
-    receive, send : Callable, Callable
+    receive : Callable
+    send : Callable
+    create_node : Optional[Callable]
+    delete_node : Optional[Callable]
+    get_run : Optional[Callable]
     """
     log(
         WARN,
@@ -126,233 +156,226 @@ def http_request_response(
             "For the REST API, the root certificates "
             "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.")
 
-    # Necessary state to validate messages to be sent
-    state: Dict[str, Optional[Metadata]] = {KEY_METADATA: None}
-
-    # Enable create_node and delete_node to store node
-    node_store: Dict[str, Optional[Node]] = {KEY_NODE: None}
+    # Shared variables for inner functions
+    metadata: Optional[Metadata] = None
+    node: Optional[Node] = None
+    ping_thread: Optional[threading.Thread] = None
+    ping_stop_event = threading.Event()
 
     ###########################################################################
-    # receive/send functions
+    # ping/create_node/delete_node/receive/send/get_run functions
     ###########################################################################
 
-    def create_node() -> None:
-        """Set create_node."""
-        create_node_req_proto = CreateNodeRequest()
-        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,
-        )
+    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
+        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
-        create_node_response_proto = CreateNodeResponse()
-        create_node_response_proto.ParseFromString(res.content)
-        # pylint: disable-next=no-member
-        node_store[KEY_NODE] = create_node_response_proto.node
+        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 res.success:
+            raise RuntimeError("Ping failed unexpectedly.")
+
+        # Wait
+        rd = random.uniform(*PING_RANDOM_RANGE)
+        next_interval: float = PING_DEFAULT_INTERVAL - PING_CALL_TIMEOUT
+        next_interval *= PING_BASE_MULTIPLIER + rd
+        if not ping_stop_event.is_set():
+            ping_stop_event.wait(next_interval)
+
+    def create_node() -> Optional[int]:
+        """Set create_node."""
+        req = CreateNodeRequest(ping_interval=PING_DEFAULT_INTERVAL)
+
+        # Send the request
+        res = _request(req, CreateNodeResponse, PATH_CREATE_NODE)
+        if res is None:
+            return None
+
+        # Remember the node and the ping-loop thread
+        nonlocal node, ping_thread
+        node = res.node
+        ping_thread = start_ping_loop(ping, ping_stop_event)
+        return node.node_id
 
     def delete_node() -> None:
         """Set delete_node."""
-        if node_store[KEY_NODE] is None:
+        nonlocal node
+        if node is None:
             log(ERROR, "Node instance missing")
             return
-        node: Node = cast(Node, node_store[KEY_NODE])
-        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,
-        )
 
-        # 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,
-            )
+        # Stop the ping-loop thread
+        ping_stop_event.set()
+        if ping_thread is not None:
+            ping_thread.join()
+
+        # Send DeleteNode request
+        req = DeleteNodeRequest(node=node)
+
+        # Send the request
+        res = _request(req, DeleteNodeResponse, PATH_DELETE_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
 
     def receive() -> Optional[Message]:
         """Receive next task from server."""
         # Get Node
-        if node_store[KEY_NODE] is None:
+        if node is None:
             log(ERROR, "Node instance missing")
             return None
-        node: Node = cast(Node, node_store[KEY_NODE])
-
-        # 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()
-
-        # 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,
-            )
-            return None
+        # Request instructions (message) from server
+        req = PullMessagesRequest(node=node)
 
-        # Deserialize ProtoBuf from bytes
-        pull_task_ins_response_proto = PullTaskInsResponse()
-        pull_task_ins_response_proto.ParseFromString(res.content)
+        # Send the request
+        res = _request(req, PullMessagesResponse, PATH_PULL_MESSAGES)
+        if res is None:
+            return None
 
-        # Get the current TaskIns
-        task_ins: Optional[TaskIns] = get_task_ins(pull_task_ins_response_proto)
+        # Get the current Messages
+        message_proto = None if len(res.messages_list) == 0 else res.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
 
         # Return the Message if available
+        nonlocal metadata
         message = None
-        state[KEY_METADATA] = None
-        if task_ins is not None:
-            message = message_from_taskins(task_ins)
-            state[KEY_METADATA] = copy(message.metadata)
+        if message_proto is not None:
+            message = message_from_proto(message_proto)
+            metadata = copy(message.metadata)
             log(INFO, "[Node] POST /%s: success", PATH_PULL_TASK_INS)
         return message
 
     def send(message: Message) -> None:
         """Send task result back to server."""
         # Get Node
-        if node_store[KEY_NODE] is None:
+        if node is None:
             log(ERROR, "Node instance missing")
             return
 
         # Get incoming message
-        in_metadata = state[KEY_METADATA]
-        if in_metadata is None:
+        nonlocal metadata
+        if metadata is None:
             log(ERROR, "No current message")
             return
 
         # Validate out message
-        if not validate_out_message(message, in_metadata):
+        if not validate_out_message(message, metadata):
             log(ERROR, "Invalid out message")
             return
-
-        # Construct TaskRes
-        task_res = message_to_taskres(message)
+        metadata = None
 
         # 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,
-        )
+        message_proto = message_to_proto(message=message)
 
-        state[KEY_METADATA] = None
+        # Serialize ProtoBuf to bytes
+        req = PushMessagesRequest(node=node, messages_list=[message_proto])
 
-        # 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, PushMessagesResponse, PATH_PUSH_MESSAGES)
+        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) -> Run:
+        # Construct the request
+        req = GetRunRequest(node=node, run_id=run_id)
+
+        # Send the request
+        res = _request(req, GetRunResponse, PATH_GET_RUN)
+        if res is None:
+            return Run.create_empty(run_id)
+
+        return run_from_proto(res.run)
+
+    def get_fab(fab_hash: str, run_id: int) -> Fab:
+        # Construct the request
+        req = GetFabRequest(node=node, hash_str=fab_hash, run_id=run_id)
+
+        # Send the request
+        res = _request(req, GetFabResponse, PATH_GET_FAB)
+        if res is None:
+            return Fab("", b"")
+
+        return Fab(
+            res.fab.hash_str,
+            res.fab.content,
         )
 
     try:
         # Yield methods
-        yield (receive, send, create_node, delete_node)
+        yield (receive, send, create_node, delete_node, get_run, get_fab)
     except Exception as exc:  # pylint: disable=broad-except
         log(ERROR, exc)