flwr-nightly 1.9.0.dev20240416__py3-none-any.whl → 1.9.0.dev20240420__py3-none-any.whl

flwr/client/rest_client/connection.py

@@ -21,7 +21,9 @@ 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 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 +44,9 @@ from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
     CreateNodeRequest,
     CreateNodeResponse,
     DeleteNodeRequest,
+    DeleteNodeResponse,
+    GetRunRequest,
+    GetRunResponse,
     PingRequest,
     PingResponse,
     PullTaskInsRequest,
@@ -63,6 +68,9 @@ 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
@@ -80,6 +88,7 @@ def http_request_response(  # pylint: disable=R0914, R0915
         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 +150,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 +227,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 +252,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 +270,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 +313,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,22 @@
+# 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_supernode as run_supernode
+
+__all__ = [
+    "run_supernode",
+]

flwr/client/supernode/app.py

@@ -0,0 +1,107 @@
+# 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
+from logging import DEBUG, INFO
+
+from flwr.common import EventType, event
+from flwr.common.exit_handlers import register_exit_handlers
+from flwr.common.logger import log
+
+
+def run_supernode() -> None:
+    """Run Flower SuperNode."""
+    log(INFO, "Starting Flower SuperNode")
+
+    event(EventType.RUN_SUPERNODE_ENTER)
+
+    args = _parse_args_run_supernode().parse_args()
+
+    log(
+        DEBUG,
+        "Flower will load ClientApp `%s`",
+        getattr(args, "client-app"),
+    )
+
+    # Graceful shutdown
+    register_exit_handlers(
+        event_type=EventType.RUN_SUPERNODE_LEAVE,
+    )
+
+
+def _parse_args_run_supernode() -> argparse.ArgumentParser:
+    """Parse flower-supernode command line arguments."""
+    parser = argparse.ArgumentParser(
+        description="Start a Flower SuperNode",
+    )
+
+    parse_args_run_client_app(parser=parser)
+
+    return parser
+
+
+def parse_args_run_client_app(parser: argparse.ArgumentParser) -> None:
+    """Parse command line arguments."""
+    parser.add_argument(
+        "client-app",
+        help="For example: `client:app` or `project.package.module:wrapper.app`",
+    )
+    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.",
+    )

flwr/common/record/recordset.py

@@ -16,23 +16,20 @@
 
 
 from dataclasses import dataclass
-from typing import Callable, Dict, Optional, Type, TypeVar
+from typing import Dict, Optional, cast
 
 from .configsrecord import ConfigsRecord
 from .metricsrecord import MetricsRecord
 from .parametersrecord import ParametersRecord
 from .typeddict import TypedDict
 
-T = TypeVar("T")
 
+class RecordSetData:
+    """Inner data container for the RecordSet class."""
 
-@dataclass
-class RecordSet:
-    """RecordSet stores groups of parameters, metrics and configs."""
-
-    _parameters_records: TypedDict[str, ParametersRecord]
-    _metrics_records: TypedDict[str, MetricsRecord]
-    _configs_records: TypedDict[str, ConfigsRecord]
+    parameters_records: TypedDict[str, ParametersRecord]
+    metrics_records: TypedDict[str, MetricsRecord]
+    configs_records: TypedDict[str, ConfigsRecord]
 
     def __init__(
         self,
@@ -40,40 +37,82 @@ class RecordSet:
         metrics_records: Optional[Dict[str, MetricsRecord]] = None,
         configs_records: Optional[Dict[str, ConfigsRecord]] = None,
     ) -> None:
-        def _get_check_fn(__t: Type[T]) -> Callable[[T], None]:
-            def _check_fn(__v: T) -> None:
-                if not isinstance(__v, __t):
-                    raise TypeError(f"Expected `{__t}`, but `{type(__v)}` was passed.")
-
-            return _check_fn
-
-        self._parameters_records = TypedDict[str, ParametersRecord](
-            _get_check_fn(str), _get_check_fn(ParametersRecord)
+        self.parameters_records = TypedDict[str, ParametersRecord](
+            self._check_fn_str, self._check_fn_params
         )
-        self._metrics_records = TypedDict[str, MetricsRecord](
-            _get_check_fn(str), _get_check_fn(MetricsRecord)
+        self.metrics_records = TypedDict[str, MetricsRecord](
+            self._check_fn_str, self._check_fn_metrics
         )
-        self._configs_records = TypedDict[str, ConfigsRecord](
-            _get_check_fn(str), _get_check_fn(ConfigsRecord)
+        self.configs_records = TypedDict[str, ConfigsRecord](
+            self._check_fn_str, self._check_fn_configs
         )
         if parameters_records is not None:
-            self._parameters_records.update(parameters_records)
+            self.parameters_records.update(parameters_records)
         if metrics_records is not None:
-            self._metrics_records.update(metrics_records)
+            self.metrics_records.update(metrics_records)
         if configs_records is not None:
-            self._configs_records.update(configs_records)
+            self.configs_records.update(configs_records)
+
+    def _check_fn_str(self, key: str) -> None:
+        if not isinstance(key, str):
+            raise TypeError(
+                f"Expected `{str.__name__}`, but "
+                f"received `{type(key).__name__}` for the key."
+            )
+
+    def _check_fn_params(self, record: ParametersRecord) -> None:
+        if not isinstance(record, ParametersRecord):
+            raise TypeError(
+                f"Expected `{ParametersRecord.__name__}`, but "
+                f"received `{type(record).__name__}` for the value."
+            )
+
+    def _check_fn_metrics(self, record: MetricsRecord) -> None:
+        if not isinstance(record, MetricsRecord):
+            raise TypeError(
+                f"Expected `{MetricsRecord.__name__}`, but "
+                f"received `{type(record).__name__}` for the value."
+            )
+
+    def _check_fn_configs(self, record: ConfigsRecord) -> None:
+        if not isinstance(record, ConfigsRecord):
+            raise TypeError(
+                f"Expected `{ConfigsRecord.__name__}`, but "
+                f"received `{type(record).__name__}` for the value."
+            )
+
+
+@dataclass
+class RecordSet:
+    """RecordSet stores groups of parameters, metrics and configs."""
+
+    def __init__(
+        self,
+        parameters_records: Optional[Dict[str, ParametersRecord]] = None,
+        metrics_records: Optional[Dict[str, MetricsRecord]] = None,
+        configs_records: Optional[Dict[str, ConfigsRecord]] = None,
+    ) -> None:
+        data = RecordSetData(
+            parameters_records=parameters_records,
+            metrics_records=metrics_records,
+            configs_records=configs_records,
+        )
+        setattr(self, "_data", data)  # noqa
 
     @property
     def parameters_records(self) -> TypedDict[str, ParametersRecord]:
         """Dictionary holding ParametersRecord instances."""
-        return self._parameters_records
+        data = cast(RecordSetData, getattr(self, "_data"))  # noqa
+        return data.parameters_records
 
     @property
     def metrics_records(self) -> TypedDict[str, MetricsRecord]:
         """Dictionary holding MetricsRecord instances."""
-        return self._metrics_records
+        data = cast(RecordSetData, getattr(self, "_data"))  # noqa
+        return data.metrics_records
 
     @property
     def configs_records(self) -> TypedDict[str, ConfigsRecord]:
         """Dictionary holding ConfigsRecord instances."""
-        return self._configs_records
+        data = cast(RecordSetData, getattr(self, "_data"))  # noqa
+        return data.configs_records

flwr/common/telemetry.py

@@ -160,6 +160,10 @@ class EventType(str, Enum):
     RUN_SERVER_APP_ENTER = auto()
     RUN_SERVER_APP_LEAVE = auto()
 
+    # SuperNode
+    RUN_SUPERNODE_ENTER = auto()
+    RUN_SUPERNODE_LEAVE = auto()
+
 
 # Use the ThreadPoolExecutor with max_workers=1 to have a queue
 # and also ensure that telemetry calls are not blocking.

flwr/server/app.py

@@ -291,9 +291,11 @@ def run_fleet_api() -> None:
 
 # pylint: disable=too-many-branches, too-many-locals, too-many-statements
 def run_superlink() -> None:
-    """Run Flower server (Driver API and Fleet API)."""
-    log(INFO, "Starting Flower server")
+    """Run Flower SuperLink (Driver API and Fleet API)."""
+    log(INFO, "Starting Flower SuperLink")
+
     event(EventType.RUN_SUPERLINK_ENTER)
+
     args = _parse_args_run_superlink().parse_args()
 
     # Parse IP address
@@ -568,9 +570,7 @@ def _parse_args_run_fleet_api() -> argparse.ArgumentParser:
 def _parse_args_run_superlink() -> argparse.ArgumentParser:
     """Parse command line arguments for both Driver API and Fleet API."""
     parser = argparse.ArgumentParser(
-        description="This will start a Flower server "
-        "(meaning, a Driver API and a Fleet API), "
-        "that clients will be able to connect to.",
+        description="Start a Flower SuperLink",
     )
 
     _add_args_common(parser=parser)

flwr/server/compat/app_utils.py

@@ -89,7 +89,7 @@ def _update_client_manager(
         for node_id in new_nodes:
             client_proxy = DriverClientProxy(
                 node_id=node_id,
-                driver=driver.grpc_driver,  # type: ignore
+                driver=driver.grpc_driver_helper,  # type: ignore
                 anonymous=False,
                 run_id=driver.run_id,  # type: ignore
             )