Introduce flower-supernode (#3240)

Co-authored-by: Heng Pan <[email protected]>

pyproject.toml

@@ -56,6 +56,7 @@ flwr = "flwr.cli.app:app"
 flower-driver-api = "flwr.server:run_driver_api"
 flower-fleet-api = "flwr.server:run_fleet_api"
 flower-superlink = "flwr.server:run_superlink"
+flower-supernode = "flwr.client:run_supernode"
 flower-client-app = "flwr.client:run_client_app"
 flower-server-app = "flwr.server:run_server_app"
 flower-simulation = "flwr.simulation:run_simulation_from_cli"

src/py/flwr/client/__init__.py

@@ -21,6 +21,7 @@
 from .client import Client as Client
 from .client_app import ClientApp as ClientApp
 from .numpy_client import NumPyClient as NumPyClient
+from .supernode import run_supernode as run_supernode
 from .typing import ClientFn as ClientFn
 
 __all__ = [
@@ -29,6 +30,7 @@
     "ClientFn",
     "NumPyClient",
     "run_client_app",
+    "run_supernode",
     "start_client",
     "start_numpy_client",
 ]

src/py/flwr/client/app.py

@@ -47,14 +47,15 @@
 from .message_handler.message_handler import handle_control_message
 from .node_state import NodeState
 from .numpy_client import NumPyClient
+from .supernode.app import parse_args_run_client_app
 
 
 def run_client_app() -> None:
     """Run Flower client app."""
-    event(EventType.RUN_CLIENT_APP_ENTER)
-
     log(INFO, "Long-running Flower client starting")
 
+    event(EventType.RUN_CLIENT_APP_ENTER)
+
     args = _parse_args_run_client_app().parse_args()
 
     # Obtain certificates
@@ -131,56 +132,7 @@ def _parse_args_run_client_app() -> argparse.ArgumentParser:
         description="Start a Flower client app",
     )
 
-    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.",
-    )
+    parse_args_run_client_app(parser=parser)
 
     return parser
 

src/py/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",
+]

src/py/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.",
+    )

src/py/flwr/common/telemetry.py

@@ -160,6 +160,10 @@ def _generate_next_value_(name: str, start: int, count: int, last_values: List[A
     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.

src/py/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)