flwr-nightly 1.15.0.dev20250104__py3-none-any.whl → 1.15.0.dev20250123__py3-none-any.whl

flwr/client/rest_client/connection.py

@@ -16,7 +16,6 @@
 
 
 import random
-import sys
 import threading
 from collections.abc import Iterator
 from contextlib import contextmanager
@@ -26,22 +25,22 @@ from typing import Callable, Optional, TypeVar, Union
 
 from cryptography.hazmat.primitives.asymmetric import ec
 from google.protobuf.message import Message as GrpcMessage
+from requests.exceptions import ConnectionError as RequestsConnectionError
 
 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,
     PING_BASE_MULTIPLIER,
     PING_CALL_TIMEOUT,
     PING_DEFAULT_INTERVAL,
     PING_RANDOM_RANGE,
 )
+from flwr.common.exit import ExitCode, flwr_exit
 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, run_from_proto
+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
@@ -51,25 +50,26 @@ from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
     DeleteNodeResponse,
     PingRequest,
     PingResponse,
-    PullTaskInsRequest,
-    PullTaskInsResponse,
-    PushTaskResRequest,
-    PushTaskResResponse,
+    PullMessagesRequest,
+    PullMessagesResponse,
+    PushMessagesRequest,
+    PushMessagesResponse,
 )
 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:
     import requests
 except ModuleNotFoundError:
-    sys.exit(MISSING_EXTRA_REST)
+    flwr_exit(ExitCode.COMMON_MISSING_EXTRA_REST)
 
 
 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"
@@ -286,29 +286,28 @@ def http_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
             log(ERROR, "Node instance missing")
             return None
 
-        # Request instructions (task) from server
-        req = PullTaskInsRequest(node=node)
+        # Request instructions (message) from server
+        req = PullMessagesRequest(node=node)
 
         # Send the request
-        res = _request(req, PullTaskInsResponse, PATH_PULL_TASK_INS)
+        res = _request(req, PullMessagesResponse, PATH_PULL_MESSAGES)
         if res is None:
             return None
 
-        # Get the current TaskIns
-        task_ins: Optional[TaskIns] = get_task_ins(res)
+        # 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
-        if task_ins is not None:
-            message = message_from_taskins(task_ins)
+        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
@@ -332,14 +331,14 @@ def http_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
             return
         metadata = None
 
-        # Construct TaskRes
-        task_res = message_to_taskres(message)
+        # Serialize ProtoBuf to bytes
+        message_proto = message_to_proto(message=message)
 
         # Serialize ProtoBuf to bytes
-        req = PushTaskResRequest(node=node, task_res_list=[task_res])
+        req = PushMessagesRequest(node=node, messages_list=[message_proto])
 
         # Send the request
-        res = _request(req, PushTaskResResponse, PATH_PUSH_TASK_RES)
+        res = _request(req, PushMessagesResponse, PATH_PUSH_MESSAGES)
         if res is None:
             return
 
@@ -380,3 +379,12 @@ def http_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
         yield (receive, send, create_node, delete_node, get_run, get_fab)
     except Exception as exc:  # pylint: disable=broad-except
         log(ERROR, exc)
+    # Cleanup
+    finally:
+        try:
+            if node is not None:
+                # Disable retrying
+                retry_invoker.max_tries = 1
+                delete_node()
+        except RequestsConnectionError:
+            pass

flwr/client/supernode/app.py

@@ -40,6 +40,7 @@ from flwr.common.constant import (
     TRANSPORT_TYPE_GRPC_RERE,
     TRANSPORT_TYPE_REST,
 )
+from flwr.common.exit import ExitCode, flwr_exit
 from flwr.common.exit_handlers import register_exit_handlers
 from flwr.common.logger import log, warn_deprecated_feature
 
@@ -86,6 +87,12 @@ def run_supernode() -> None:
 
     log(DEBUG, "Isolation mode: %s", args.isolation)
 
+    # Register handlers for graceful shutdown
+    register_exit_handlers(
+        event_type=EventType.RUN_SUPERNODE_LEAVE,
+        exit_message="SuperNode terminated gracefully.",
+    )
+
     start_client_internal(
         server_address=args.superlink,
         load_client_app_fn=load_fn,
@@ -103,11 +110,6 @@ def run_supernode() -> None:
         clientappio_api_address=args.clientappio_api_address,
     )
 
-    # Graceful shutdown
-    register_exit_handlers(
-        event_type=EventType.RUN_SUPERNODE_LEAVE,
-    )
-
 
 def run_client_app() -> None:
     """Run Flower client app."""
@@ -280,11 +282,7 @@ def _try_setup_client_authentication(
         return None
 
     if not args.auth_supernode_private_key or not args.auth_supernode_public_key:
-        sys.exit(
-            "Authentication requires file paths to both "
-            "'--auth-supernode-private-key' and '--auth-supernode-public-key'"
-            "to be provided (providing only one of them is not sufficient)."
-        )
+        flwr_exit(ExitCode.SUPERNODE_NODE_AUTH_KEYS_REQUIRED)
 
     try:
         ssh_private_key = load_ssh_private_key(
@@ -294,11 +292,9 @@ def _try_setup_client_authentication(
         if not isinstance(ssh_private_key, ec.EllipticCurvePrivateKey):
             raise ValueError()
     except (ValueError, UnsupportedAlgorithm):
-        sys.exit(
-            "Error: Unable to parse the private key file in "
-            "'--auth-supernode-private-key'. Authentication requires elliptic "
-            "curve private and public key pair. Please ensure that the file "
-            "path points to a valid private key file and try again."
+        flwr_exit(
+            ExitCode.SUPERNODE_NODE_AUTH_KEYS_INVALID,
+            "Unable to parse the private key file.",
         )
 
     try:
@@ -308,11 +304,9 @@ def _try_setup_client_authentication(
         if not isinstance(ssh_public_key, ec.EllipticCurvePublicKey):
             raise ValueError()
     except (ValueError, UnsupportedAlgorithm):
-        sys.exit(
-            "Error: Unable to parse the public key file in "
-            "'--auth-supernode-public-key'. Authentication requires elliptic "
-            "curve private and public key pair. Please ensure that the file "
-            "path points to a valid public key file and try again."
+        flwr_exit(
+            ExitCode.SUPERNODE_NODE_AUTH_KEYS_INVALID,
+            "Unable to parse the public key file.",
         )
 
     return (

flwr/common/auth_plugin/auth_plugin.py

@@ -18,26 +18,32 @@
 from abc import ABC, abstractmethod
 from collections.abc import Sequence
 from pathlib import Path
-from typing import Any, Optional, Union
+from typing import Optional, Union
 
 from flwr.proto.exec_pb2_grpc import ExecStub
 
+from ..typing import UserAuthCredentials, UserAuthLoginDetails
+
 
 class ExecAuthPlugin(ABC):
     """Abstract Flower Auth Plugin class for ExecServicer.
 
     Parameters
     ----------
-    config : dict[str, Any]
-        The authentication configuration loaded from a YAML file.
+    user_auth_config_path : Path
+        Path to the YAML file containing the authentication configuration.
     """
 
     @abstractmethod
-    def __init__(self, config: dict[str, Any]):
+    def __init__(
+        self,
+        user_auth_config_path: Path,
+        verify_tls_cert: bool,
+    ):
         """Abstract constructor."""
 
     @abstractmethod
-    def get_login_details(self) -> dict[str, str]:
+    def get_login_details(self) -> Optional[UserAuthLoginDetails]:
         """Get the login details."""
 
     @abstractmethod
@@ -47,7 +53,7 @@ class ExecAuthPlugin(ABC):
         """Validate authentication tokens in the provided metadata."""
 
     @abstractmethod
-    def get_auth_tokens(self, auth_details: dict[str, str]) -> dict[str, str]:
+    def get_auth_tokens(self, device_code: str) -> Optional[UserAuthCredentials]:
         """Get authentication tokens."""
 
     @abstractmethod
@@ -62,50 +68,55 @@ class CliAuthPlugin(ABC):
 
     Parameters
     ----------
-    user_auth_config_path : Path
-        The path to the user's authentication configuration file.
+    credentials_path : Path
+        Path to the user's authentication credentials file.
     """
 
     @staticmethod
     @abstractmethod
     def login(
-        login_details: dict[str, str],
+        login_details: UserAuthLoginDetails,
         exec_stub: ExecStub,
-    ) -> dict[str, Any]:
-        """Authenticate the user with the SuperLink.
+    ) -> UserAuthCredentials:
+        """Authenticate the user and retrieve authentication credentials.
 
         Parameters
         ----------
-        login_details : dict[str, str]
-            A dictionary containing the user's login details.
+        login_details : UserAuthLoginDetails
+            An object containing the user's login details.
         exec_stub : ExecStub
-            An instance of `ExecStub` used for communication with the SuperLink.
+            A stub for executing RPC calls to the server.
 
         Returns
         -------
-        user_auth_config : dict[str, Any]
-            A dictionary containing the user's authentication configuration
-            in JSON format.
+        UserAuthCredentials
+            The authentication credentials obtained after login.
         """
 
     @abstractmethod
-    def __init__(self, user_auth_config_path: Path):
+    def __init__(self, credentials_path: Path):
         """Abstract constructor."""
 
     @abstractmethod
-    def store_tokens(self, user_auth_config: dict[str, Any]) -> None:
-        """Store authentication tokens from the provided user_auth_config.
+    def store_tokens(self, credentials: UserAuthCredentials) -> None:
+        """Store authentication tokens to the `credentials_path`.
 
-        The configuration, including tokens, will be saved as a JSON file
-        at `user_auth_config_path`.
+        The credentials, including tokens, will be saved as a JSON file
+        at `credentials_path`.
         """
 
     @abstractmethod
     def load_tokens(self) -> None:
-        """Load authentication tokens from the user_auth_config_path."""
+        """Load authentication tokens from the `credentials_path`."""
 
     @abstractmethod
     def write_tokens_to_metadata(
         self, metadata: Sequence[tuple[str, Union[str, bytes]]]
     ) -> Sequence[tuple[str, Union[str, bytes]]]:
         """Write authentication tokens to the provided metadata."""
+
+    @abstractmethod
+    def read_tokens_from_metadata(
+        self, metadata: Sequence[tuple[str, Union[str, bytes]]]
+    ) -> Optional[UserAuthCredentials]:
+        """Read authentication tokens from the provided metadata."""

flwr/common/config.py

@@ -17,13 +17,13 @@
 
 import os
 import re
+import zipfile
+from io import BytesIO
 from pathlib import Path
-from typing import Any, Optional, Union, cast, get_args
+from typing import IO, Any, Optional, TypeVar, Union, cast, get_args
 
 import tomli
 
-from flwr.cli.config_utils import get_fab_config, validate_fields
-from flwr.common import ConfigsRecord
 from flwr.common.constant import (
     APP_DIR,
     FAB_CONFIG_FILE,
@@ -33,6 +33,10 @@ from flwr.common.constant import (
 )
 from flwr.common.typing import Run, UserConfig, UserConfigValue
 
+from . import ConfigsRecord, object_ref
+
+T_dict = TypeVar("T_dict", bound=dict[str, Any])  # pylint: disable=invalid-name
+
 
 def get_flwr_dir(provided_path: Optional[str] = None) -> Path:
     """Return the Flower home directory based on env variables."""
@@ -80,7 +84,7 @@ def get_project_config(project_dir: Union[str, Path]) -> dict[str, Any]:
         config = tomli.loads(toml_file.read())
 
     # Validate pyproject.toml fields
-    is_valid, errors, _ = validate_fields(config)
+    is_valid, errors, _ = validate_fields_in_config(config)
     if not is_valid:
         error_msg = "\n".join([f"  - {error}" for error in errors])
         raise ValueError(
@@ -91,19 +95,28 @@ def get_project_config(project_dir: Union[str, Path]) -> dict[str, Any]:
 
 
 def fuse_dicts(
-    main_dict: UserConfig,
-    override_dict: UserConfig,
-) -> UserConfig:
+    main_dict: T_dict,
+    override_dict: T_dict,
+    check_keys: bool = True,
+) -> T_dict:
     """Merge a config with the overrides.
 
-    Remove the nesting by adding the nested keys as prefixes separated by dots, and fuse
-    it with the override dict.
+    If `check_keys` is set to True, an error will be raised if the override
+    dictionary contains keys that are not present in the main dictionary.
+    Otherwise, only the keys present in the main dictionary will be updated.
     """
-    fused_dict = main_dict.copy()
+    if not isinstance(main_dict, dict) or not isinstance(override_dict, dict):
+        raise ValueError("Both dictionaries must be of type dict")
+
+    fused_dict = cast(T_dict, main_dict.copy())
 
     for key, value in override_dict.items():
         if key in main_dict:
+            if isinstance(value, dict):
+                fused_dict[key] = fuse_dicts(main_dict[key], value)
             fused_dict[key] = value
+        elif check_keys:
+            raise ValueError(f"Key '{key}' is not present in the main dictionary")
 
     return fused_dict
 
@@ -192,8 +205,8 @@ def unflatten_dict(flat_dict: dict[str, Any]) -> dict[str, Any]:
 
 
 def parse_config_args(
-    config: Optional[list[str]],
-) -> UserConfig:
+    config: Optional[list[str]], flatten: bool = True
+) -> dict[str, Any]:
     """Parse separator separated list of key-value pairs separated by '='."""
     overrides: UserConfig = {}
 
@@ -221,16 +234,16 @@ def parse_config_args(
             matches = pattern.findall(config_line)
             toml_str = "\n".join(f"{k} = {v}" for k, v in matches)
             overrides.update(tomli.loads(toml_str))
-            flat_overrides = flatten_dict(overrides)
+            flat_overrides = flatten_dict(overrides) if flatten else overrides
 
     return flat_overrides
 
 
 def get_metadata_from_config(config: dict[str, Any]) -> tuple[str, str]:
-    """Extract `fab_version` and `fab_id` from a project config."""
+    """Extract `fab_id` and `fab_version` from a project config."""
     return (
-        config["project"]["version"],
         f"{config['tool']['flwr']['app']['publisher']}/{config['project']['name']}",
+        config["project"]["version"],
     )
 
 
@@ -241,3 +254,127 @@ def user_config_to_configsrecord(config: UserConfig) -> ConfigsRecord:
         c_record[k] = v
 
     return c_record
+
+
+def get_fab_config(fab_file: Union[Path, bytes]) -> dict[str, Any]:
+    """Extract the config from a FAB file or path.
+
+    Parameters
+    ----------
+    fab_file : Union[Path, bytes]
+        The Flower App Bundle file to validate and extract the metadata from.
+        It can either be a path to the file or the file itself as bytes.
+
+    Returns
+    -------
+    Dict[str, Any]
+        The `config` of the given Flower App Bundle.
+    """
+    fab_file_archive: Union[Path, IO[bytes]]
+    if isinstance(fab_file, bytes):
+        fab_file_archive = BytesIO(fab_file)
+    elif isinstance(fab_file, Path):
+        fab_file_archive = fab_file
+    else:
+        raise ValueError("fab_file must be either a Path or bytes")
+
+    with zipfile.ZipFile(fab_file_archive, "r") as zipf:
+        with zipf.open("pyproject.toml") as file:
+            toml_content = file.read().decode("utf-8")
+        try:
+            conf = tomli.loads(toml_content)
+        except tomli.TOMLDecodeError:
+            raise ValueError("Invalid TOML content in pyproject.toml") from None
+
+        is_valid, errors, _ = validate_config(conf, check_module=False)
+        if not is_valid:
+            raise ValueError(errors)
+
+        return conf
+
+
+def _validate_run_config(config_dict: dict[str, Any], errors: list[str]) -> None:
+    for key, value in config_dict.items():
+        if isinstance(value, dict):
+            _validate_run_config(config_dict[key], errors)
+        elif not isinstance(value, get_args(UserConfigValue)):
+            raise ValueError(
+                f"The value for key {key} needs to be of type `int`, `float`, "
+                "`bool, `str`, or  a `dict` of those.",
+            )
+
+
+# pylint: disable=too-many-branches
+def validate_fields_in_config(
+    config: dict[str, Any]
+) -> tuple[bool, list[str], list[str]]:
+    """Validate pyproject.toml fields."""
+    errors = []
+    warnings = []
+
+    if "project" not in config:
+        errors.append("Missing [project] section")
+    else:
+        if "name" not in config["project"]:
+            errors.append('Property "name" missing in [project]')
+        if "version" not in config["project"]:
+            errors.append('Property "version" missing in [project]')
+        if "description" not in config["project"]:
+            warnings.append('Recommended property "description" missing in [project]')
+        if "license" not in config["project"]:
+            warnings.append('Recommended property "license" missing in [project]')
+        if "authors" not in config["project"]:
+            warnings.append('Recommended property "authors" missing in [project]')
+
+    if (
+        "tool" not in config
+        or "flwr" not in config["tool"]
+        or "app" not in config["tool"]["flwr"]
+    ):
+        errors.append("Missing [tool.flwr.app] section")
+    else:
+        if "publisher" not in config["tool"]["flwr"]["app"]:
+            errors.append('Property "publisher" missing in [tool.flwr.app]')
+        if "config" in config["tool"]["flwr"]["app"]:
+            _validate_run_config(config["tool"]["flwr"]["app"]["config"], errors)
+        if "components" not in config["tool"]["flwr"]["app"]:
+            errors.append("Missing [tool.flwr.app.components] section")
+        else:
+            if "serverapp" not in config["tool"]["flwr"]["app"]["components"]:
+                errors.append(
+                    'Property "serverapp" missing in [tool.flwr.app.components]'
+                )
+            if "clientapp" not in config["tool"]["flwr"]["app"]["components"]:
+                errors.append(
+                    'Property "clientapp" missing in [tool.flwr.app.components]'
+                )
+
+    return len(errors) == 0, errors, warnings
+
+
+def validate_config(
+    config: dict[str, Any],
+    check_module: bool = True,
+    project_dir: Optional[Union[str, Path]] = None,
+) -> tuple[bool, list[str], list[str]]:
+    """Validate pyproject.toml."""
+    is_valid, errors, warnings = validate_fields_in_config(config)
+
+    if not is_valid:
+        return False, errors, warnings
+
+    # Validate serverapp
+    serverapp_ref = config["tool"]["flwr"]["app"]["components"]["serverapp"]
+    is_valid, reason = object_ref.validate(serverapp_ref, check_module, project_dir)
+
+    if not is_valid and isinstance(reason, str):
+        return False, [reason], []
+
+    # Validate clientapp
+    clientapp_ref = config["tool"]["flwr"]["app"]["components"]["clientapp"]
+    is_valid, reason = object_ref.validate(clientapp_ref, check_module, project_dir)
+
+    if not is_valid and isinstance(reason, str):
+        return False, [reason], []
+
+    return True, [], []

flwr/common/constant.py

@@ -17,14 +17,6 @@
 
 from __future__ import annotations
 
-MISSING_EXTRA_REST = """
-Extra dependencies required for using the REST-based Fleet API are missing.
-
-To use the REST API, install `flwr` with the `rest` extra:
-
-    `pip install flwr[rest]`.
-"""
-
 TRANSPORT_TYPE_GRPC_BIDI = "grpc-bidi"
 TRANSPORT_TYPE_GRPC_RERE = "grpc-rere"
 TRANSPORT_TYPE_GRPC_ADAPTER = "grpc-adapter"
@@ -83,6 +75,9 @@ FAB_HASH_TRUNCATION = 8
 FLWR_DIR = ".flwr"  # The default Flower directory: ~/.flwr/
 FLWR_HOME = "FLWR_HOME"  # If set, override the default Flower directory
 
+# Constant for SuperLink
+SUPERLINK_NODE_ID = 1
+
 # Constants entries in Node config for Simulation
 PARTITION_ID_KEY = "partition-id"
 NUM_PARTITIONS_KEY = "num-partitions"
@@ -114,6 +109,14 @@ MAX_RETRY_DELAY = 20  # Maximum delay duration between two consecutive retries.
 # Constants for user authentication
 CREDENTIALS_DIR = ".credentials"
 AUTH_TYPE = "auth_type"
+ACCESS_TOKEN_KEY = "access_token"
+REFRESH_TOKEN_KEY = "refresh_token"
+
+# Constants for node authentication
+PUBLIC_KEY_HEADER = "public-key-bin"  # Must end with "-bin" for binary data
+SIGNATURE_HEADER = "signature-bin"  # Must end with "-bin" for binary data
+TIMESTAMP_HEADER = "timestamp"
+TIMESTAMP_TOLERANCE = 10  # Tolerance for timestamp verification
 
 
 class MessageType:

flwr/common/exit/__init__.py

@@ -0,0 +1,24 @@
+# Copyright 2025 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 exit functionality."""
+
+
+from .exit import flwr_exit
+from .exit_code import ExitCode
+
+__all__ = [
+    "ExitCode",
+    "flwr_exit",
+]