flwr-nightly 1.8.0.dev20240314__py3-none-any.whl → 1.11.0.dev20240813__py3-none-any.whl

flwr/client/mod/secure_aggregation/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2023 Flower Labs GmbH. All Rights Reserved.
+# 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.

flwr/client/mod/secure_aggregation/secaggplus_mod.py

@@ -187,7 +187,7 @@ def secaggplus_mod(
 
     # Return message
     out_content.configs_records[RECORD_KEY_CONFIGS] = ConfigsRecord(res, False)
-    return msg.create_reply(out_content, ttl="")
+    return msg.create_reply(out_content)
 
 
 def check_stage(current_stage: str, configs: ConfigsRecord) -> None:

flwr/client/mod/utils.py

@@ -1,4 +1,4 @@
-# Copyright 2023 Flower Labs GmbH. All Rights Reserved.
+# 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.

flwr/client/node_state.py

@@ -15,27 +15,72 @@
 """Node state."""
 
 
-from typing import Any, Dict
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Dict, Optional
 
 from flwr.common import Context, RecordSet
+from flwr.common.config import get_fused_config, get_fused_config_from_dir
+from flwr.common.typing import Run, UserConfig
+
+
+@dataclass()
+class RunInfo:
+    """Contains the Context and initial run_config of a Run."""
+
+    context: Context
+    initial_run_config: UserConfig
 
 
 class NodeState:
     """State of a node where client nodes execute runs."""
 
-    def __init__(self) -> None:
-        self._meta: Dict[str, Any] = {}  # holds metadata about the node
-        self.run_contexts: Dict[int, Context] = {}
+    def __init__(
+        self,
+        node_id: int,
+        node_config: UserConfig,
+    ) -> None:
+        self.node_id = node_id
+        self.node_config = node_config
+        self.run_infos: Dict[int, RunInfo] = {}
 
-    def register_context(self, run_id: int) -> None:
+    def register_context(
+        self,
+        run_id: int,
+        run: Optional[Run] = None,
+        flwr_path: Optional[Path] = None,
+        app_dir: Optional[str] = None,
+    ) -> None:
         """Register new run context for this node."""
-        if run_id not in self.run_contexts:
-            self.run_contexts[run_id] = Context(state=RecordSet())
+        if run_id not in self.run_infos:
+            initial_run_config = {}
+            if app_dir:
+                # Load from app directory
+                app_path = Path(app_dir)
+                if app_path.is_dir():
+                    override_config = run.override_config if run else {}
+                    initial_run_config = get_fused_config_from_dir(
+                        app_path, override_config
+                    )
+                else:
+                    raise ValueError("The specified `app_dir` must be a directory.")
+            else:
+                # Load from .fab
+                initial_run_config = get_fused_config(run, flwr_path) if run else {}
+            self.run_infos[run_id] = RunInfo(
+                initial_run_config=initial_run_config,
+                context=Context(
+                    node_id=self.node_id,
+                    node_config=self.node_config,
+                    state=RecordSet(),
+                    run_config=initial_run_config.copy(),
+                ),
+            )
 
     def retrieve_context(self, run_id: int) -> Context:
         """Get run context given a run_id."""
-        if run_id in self.run_contexts:
-            return self.run_contexts[run_id]
+        if run_id in self.run_infos:
+            return self.run_infos[run_id].context
 
         raise RuntimeError(
             f"Context for run_id={run_id} doesn't exist."
@@ -45,4 +90,9 @@ class NodeState:
 
     def update_context(self, run_id: int, context: Context) -> None:
         """Update run context."""
-        self.run_contexts[run_id] = context
+        if context.run_config != self.run_infos[run_id].initial_run_config:
+            raise ValueError(
+                "The `run_config` field of the `Context` object cannot be "
+                f"modified (run_id: {run_id})."
+            )
+        self.run_infos[run_id].context = context

flwr/client/node_state_tests.py

@@ -41,7 +41,7 @@ def test_multirun_in_node_state() -> None:
     expected_values = {0: "1", 1: "1" * 3, 2: "1" * 2, 3: "1", 5: "1"}
 
     # NodeState
-    node_state = NodeState()
+    node_state = NodeState(node_id=0, node_config={})
 
     for task in tasks:
         run_id = task.run_id
@@ -59,7 +59,8 @@ def test_multirun_in_node_state() -> None:
         node_state.update_context(run_id=run_id, context=updated_state)
 
     # Verify values
-    for run_id, context in node_state.run_contexts.items():
+    for run_id, run_info in node_state.run_infos.items():
         assert (
-            context.state.configs_records["counter"]["count"] == expected_values[run_id]
+            run_info.context.state.configs_records["counter"]["count"]
+            == expected_values[run_id]
         )

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,30 +15,51 @@
 """Contextmanager for a REST request-response channel to the Flower server."""
 
 
+import random
 import sys
+import threading
 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, 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
 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_taskins,
+    message_to_taskres,
+    user_config_from_proto,
+)
+from flwr.common.typing import Fab, Run
 from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
     CreateNodeRequest,
     CreateNodeResponse,
     DeleteNodeRequest,
+    DeleteNodeResponse,
+    PingRequest,
+    PingResponse,
     PullTaskInsRequest,
     PullTaskInsResponse,
     PushTaskResRequest,
     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:
@@ -47,19 +68,18 @@ 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_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
-# pylint: disable-next=too-many-statements
-def http_request_response(
+def http_request_response(  # pylint: disable=,R0913, R0914, R0915
     server_address: str,
     insecure: bool,  # pylint: disable=unused-argument
     retry_invoker: RetryInvoker,
@@ -67,12 +87,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[
         Callable[[], Optional[Message]],
         Callable[[Message], None],
+        Optional[Callable[[], Optional[int]]],
         Optional[Callable[[], None]],
-        Optional[Callable[[], None]],
+        Optional[Callable[[int], Run]],
+        Optional[Callable[[str], Fab]],
     ]
 ]:
     """Primitives for request/response-based interaction with a server.
@@ -98,10 +123,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,144 +157,146 @@ 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_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
 
     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,
-        )
+        req = PullTaskInsRequest(node=node)
 
-        # 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 (
@@ -273,86 +306,73 @@ def http_request_response(
             task_ins = 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)
+            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
+        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()
-        )
+        req = PushTaskResRequest(task_res_list=[task_res])
 
-        # 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,
-        )
-
-        state[KEY_METADATA] = 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_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) -> Run:
+        # Construct the request
+        req = GetRunRequest(run_id=run_id)
+
+        # Send the request
+        res = _request(req, GetRunResponse, PATH_GET_RUN)
+        if res is None:
+            return Run(run_id, "", "", {})
+
+        return Run(
+            run_id,
+            res.run.fab_id,
+            res.run.fab_version,
+            user_config_from_proto(res.run.override_config),
+        )
+
+    def get_fab(fab_hash: str) -> Fab:
+        # Call FleetAPI
+        raise NotImplementedError
+
     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)

flwr/client/supernode/__init__.py

@@ -0,0 +1,26 @@
+# 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 flwr_clientapp as flwr_clientapp
+from .app import run_client_app as run_client_app
+from .app import run_supernode as run_supernode
+
+__all__ = [
+    "flwr_clientapp",
+    "run_client_app",
+    "run_supernode",
+]