flwr 1.19.0__py3-none-any.whl → 1.20.0__py3-none-any.whl

flwr/client/rest_client/connection.py

@@ -14,40 +14,29 @@
 # ==============================================================================
 """Contextmanager for a REST request-response channel to the Flower server."""
 
+
 from collections.abc import Iterator
 from contextlib import contextmanager
-from copy import copy
-from logging import DEBUG, ERROR, INFO, WARN
-from typing import Callable, Optional, TypeVar, Union, cast
+from logging import ERROR, WARN
+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.app.metadata import Metadata
-from flwr.client.message_handler.message_handler import validate_out_message
 from flwr.common import GRPC_MAX_MESSAGE_LENGTH
 from flwr.common.constant import HEARTBEAT_DEFAULT_INTERVAL
 from flwr.common.exit import ExitCode, flwr_exit
 from flwr.common.heartbeat import HeartbeatSender
-from flwr.common.inflatable import (
-    get_all_nested_objects,
-    get_object_tree,
-    no_object_id_recompute,
-)
-from flwr.common.inflatable_rest_utils import (
-    make_pull_object_fn_rest,
-    make_push_object_fn_rest,
-)
-from flwr.common.inflatable_utils import (
-    inflate_object_from_contents,
-    pull_objects,
-    push_objects,
+from flwr.common.inflatable_protobuf_utils import (
+    make_confirm_message_received_fn_protobuf,
+    make_pull_object_fn_protobuf,
+    make_push_object_fn_protobuf,
 )
 from flwr.common.logger import log
 from flwr.common.message import Message, remove_content_from_message
 from flwr.common.retry_invoker import RetryInvoker
-from flwr.common.serde import message_to_proto, 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
@@ -67,6 +56,7 @@ from flwr.proto.heartbeat_pb2 import (  # pylint: disable=E0611
 from flwr.proto.message_pb2 import (  # pylint: disable=E0611
     ConfirmMessageReceivedRequest,
     ConfirmMessageReceivedResponse,
+    ObjectTree,
     PullObjectRequest,
     PullObjectResponse,
     PushObjectRequest,
@@ -109,12 +99,15 @@ def http_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
     ] = None,
 ) -> Iterator[
     tuple[
-        Callable[[], Optional[Message]],
-        Callable[[Message], None],
+        Callable[[], Optional[tuple[Message, ObjectTree]]],
+        Callable[[Message, ObjectTree], set[str]],
         Callable[[], Optional[int]],
         Callable[[], None],
         Callable[[int], Run],
         Callable[[str, int], Fab],
+        Callable[[int, str], bytes],
+        Callable[[int, str, bytes], None],
+        Callable[[int, str], None],
     ]
 ]:
     """Primitives for request/response-based interaction with a server.
@@ -150,6 +143,9 @@ def http_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
     create_node : Optional[Callable]
     delete_node : Optional[Callable]
     get_run : Optional[Callable]
+    pull_object : Callable[[str], bytes]
+    push_object : Callable[[str, bytes], None]
+    confirm_message_received : Callable[[str], None]
     """
     log(
         WARN,
@@ -178,7 +174,6 @@ def http_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
         log(ERROR, "Client authentication is not supported for this transport type.")
 
     # Shared variables for inner functions
-    metadata: Optional[Metadata] = None
     node: Optional[Node] = None
 
     ###########################################################################
@@ -232,6 +227,38 @@ def http_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
         grpc_res.ParseFromString(res.content)
         return grpc_res
 
+    def _pull_object_protobuf(request: PullObjectRequest) -> PullObjectResponse:
+        res = _request(
+            req=request,
+            res_type=PullObjectResponse,
+            api_path=PATH_PULL_OBJECT,
+        )
+        if res is None:
+            raise ValueError(f"{PullObjectResponse.__name__} is None.")
+        return res
+
+    def _push_object_protobuf(request: PushObjectRequest) -> PushObjectResponse:
+        res = _request(
+            req=request,
+            res_type=PushObjectResponse,
+            api_path=PATH_PUSH_OBJECT,
+        )
+        if res is None:
+            raise ValueError(f"{PushObjectResponse.__name__} is None.")
+        return res
+
+    def _confirm_message_received_protobuf(
+        request: ConfirmMessageReceivedRequest,
+    ) -> ConfirmMessageReceivedResponse:
+        res = _request(
+            req=request,
+            res_type=ConfirmMessageReceivedResponse,
+            api_path=PATH_CONFIRM_MESSAGE_RECEIVED,
+        )
+        if res is None:
+            raise ValueError(f"{ConfirmMessageReceivedResponse.__name__} is None.")
+        return res
+
     def send_node_heartbeat() -> bool:
         # Get Node
         if node is None:
@@ -279,8 +306,7 @@ def http_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
         """Set delete_node."""
         nonlocal node
         if node is None:
-            log(ERROR, "Node instance missing")
-            return
+            raise RuntimeError("Node instance missing")
 
         # Stop the heartbeat sender
         heartbeat_sender.stop()
@@ -296,162 +322,54 @@ def http_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
         # Cleanup
         node = None
 
-    def receive() -> Optional[Message]:
-        """Receive next Message from server."""
+    def receive() -> Optional[tuple[Message, ObjectTree]]:
+        """Pull a message with its ObjectTree from SuperLink."""
         # Get Node
         if node is None:
-            log(ERROR, "Node instance missing")
-            return None
+            raise RuntimeError("Node instance missing")
 
-        # Request instructions (message) from server
+        # Try to pull a message with its object tree from SuperLink
         req = PullMessagesRequest(node=node)
-
-        # Send the request
         res = _request(req, PullMessagesResponse, PATH_PULL_MESSAGES)
         if res is None:
-            return None
+            raise ValueError("PushMessagesResponse is None.")
 
-        # Get the current Messages
-        message_proto = None if len(res.messages_list) == 0 else res.messages_list[0]
+        # If no messages are available, return None
+        if len(res.messages_list) == 0:
+            return None
 
-        # Discard the current message if not valid
-        if message_proto is not None and not (
-            message_proto.metadata.dst_node_id == node.node_id
-        ):
-            message_proto = None
+        # Get the current Message and its object tree
+        message_proto = res.messages_list[0]
+        object_tree = res.message_object_trees[0]
 
         # Construct the Message
-        in_message: Optional[Message] = None
-
-        if message_proto:
-            log(INFO, "[Node] POST /%s: success", PATH_PULL_MESSAGES)
-            msg_id = message_proto.metadata.message_id
-            run_id = message_proto.metadata.run_id
-
-            def fn(request: PullObjectRequest) -> PullObjectResponse:
-                res = _request(
-                    req=request, res_type=PullObjectResponse, api_path=PATH_PULL_OBJECT
-                )
-                if res is None:
-                    raise ValueError("PushObjectResponse is None.")
-                return res
-
-            try:
-                all_object_contents = pull_objects(
-                    list(res.objects_to_pull[msg_id].object_ids) + [msg_id],
-                    pull_object_fn=make_pull_object_fn_rest(
-                        pull_object_rest=fn,
-                        node=node,
-                        run_id=run_id,
-                    ),
-                )
-
-                # Confirm that the message has been received
-                _request(
-                    req=ConfirmMessageReceivedRequest(
-                        node=node, run_id=run_id, message_object_id=msg_id
-                    ),
-                    res_type=ConfirmMessageReceivedResponse,
-                    api_path=PATH_CONFIRM_MESSAGE_RECEIVED,
-                )
-            except ValueError as e:
-                log(
-                    ERROR,
-                    "Pulling objects failed. Potential irrecoverable error: %s",
-                    str(e),
-                )
-            in_message = cast(
-                Message, inflate_object_from_contents(msg_id, all_object_contents)
-            )
-            # The deflated message doesn't contain the message_id (its own object_id)
-            # Inject
-            in_message.metadata.__dict__["_message_id"] = msg_id
-
-        # Remember `metadata` of the in message
-        nonlocal metadata
-        metadata = copy(in_message.metadata) if in_message else None
+        in_message = message_from_proto(message_proto)
 
-        return in_message
+        # Return the Message and its object tree
+        return in_message, object_tree
 
-    def send(message: Message) -> None:
-        """Send Message result back to server."""
+    def send(message: Message, object_tree: ObjectTree) -> set[str]:
+        """Send the message with its ObjectTree to SuperLink."""
         # Get Node
         if node is None:
-            log(ERROR, "Node instance missing")
-            return
+            raise RuntimeError("Node instance missing")
 
-        # Get incoming message
-        nonlocal metadata
-        if metadata is None:
-            log(ERROR, "No current message")
-            return
+        # Remove the content from the message if it has
+        if message.has_content():
+            message = remove_content_from_message(message)
 
-        # Set message_id
-        message.metadata.__dict__["_message_id"] = message.object_id
-        # Validate out message
-        if not validate_out_message(message, metadata):
-            log(ERROR, "Invalid out message")
-            return
-
-        with no_object_id_recompute():
-            # Get all nested objects
-            all_objects = get_all_nested_objects(message)
-            object_tree = get_object_tree(message)
-
-            # Serialize Message
-            message_proto = message_to_proto(
-                message=remove_content_from_message(message)
-            )
-            req = PushMessagesRequest(
-                node=node,
-                messages_list=[message_proto],
-                message_object_trees=[object_tree],
-            )
-
-            # Send the request
-            res = _request(req, PushMessagesResponse, PATH_PUSH_MESSAGES)
-            if res:
-                log(
-                    INFO,
-                    "[Node] POST /%s: success, created result %s",
-                    PATH_PUSH_MESSAGES,
-                    res.results,  # pylint: disable=no-member
-                )
-
-            if res and res.objects_to_push:
-                objs_to_push = set(res.objects_to_push[message.object_id].object_ids)
-
-                def fn(request: PushObjectRequest) -> PushObjectResponse:
-                    res = _request(
-                        req=request,
-                        res_type=PushObjectResponse,
-                        api_path=PATH_PUSH_OBJECT,
-                    )
-                    if res is None:
-                        raise ValueError("PushObjectResponse is None.")
-                    return res
-
-                try:
-                    push_objects(
-                        all_objects,
-                        push_object_fn=make_push_object_fn_rest(
-                            push_object_rest=fn,
-                            node=node,
-                            run_id=message_proto.metadata.run_id,
-                        ),
-                        object_ids_to_push=objs_to_push,
-                    )
-                    log(DEBUG, "Pushed %s objects to servicer.", len(objs_to_push))
-                except ValueError as e:
-                    log(
-                        ERROR,
-                        "Pushing objects failed. Potential irrecoverable error: %s",
-                        str(e),
-                    )
-                    log(ERROR, str(e))
+        # Send the message with its ObjectTree to SuperLink
+        req = PushMessagesRequest(
+            node=node,
+            messages_list=[message_to_proto(message)],
+            message_object_trees=[object_tree],
+        )
+        res = _request(req, PushMessagesResponse, PATH_PUSH_MESSAGES)
+        if res is None:
+            raise ValueError("PushMessagesResponse is None.")
 
-        # Cleanup
-        metadata = None
+        # Get and return the object IDs to push
+        return set(res.objects_to_push)
 
     def get_run(run_id: int) -> Run:
         # Construct the request
@@ -478,9 +396,58 @@ def http_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
             res.fab.content,
         )
 
+    def pull_object(run_id: int, object_id: str) -> bytes:
+        """Pull the object from the SuperLink."""
+        # Check Node
+        if node is None:
+            raise RuntimeError("Node instance missing")
+
+        fn = make_pull_object_fn_protobuf(
+            pull_object_protobuf=_pull_object_protobuf,
+            node=node,
+            run_id=run_id,
+        )
+        return fn(object_id)
+
+    def push_object(run_id: int, object_id: str, contents: bytes) -> None:
+        """Push the object to the SuperLink."""
+        # Check Node
+        if node is None:
+            raise RuntimeError("Node instance missing")
+
+        fn = make_push_object_fn_protobuf(
+            push_object_protobuf=_push_object_protobuf,
+            node=node,
+            run_id=run_id,
+        )
+        fn(object_id, contents)
+
+    def confirm_message_received(run_id: int, object_id: str) -> None:
+        """Confirm that the message has been received."""
+        # Check Node
+        if node is None:
+            raise RuntimeError("Node instance missing")
+
+        fn = make_confirm_message_received_fn_protobuf(
+            confirm_message_received_protobuf=_confirm_message_received_protobuf,
+            node=node,
+            run_id=run_id,
+        )
+        fn(object_id)
+
     try:
         # Yield methods
-        yield (receive, send, create_node, delete_node, get_run, get_fab)
+        yield (
+            receive,
+            send,
+            create_node,
+            delete_node,
+            get_run,
+            get_fab,
+            pull_object,
+            push_object,
+            confirm_message_received,
+        )
     except Exception as exc:  # pylint: disable=broad-except
         log(ERROR, exc)
     # Cleanup

flwr/common/constant.py

@@ -74,6 +74,7 @@ FAB_ALLOWED_EXTENSIONS = {".py", ".toml", ".md"}
 FAB_CONFIG_FILE = "pyproject.toml"
 FAB_DATE = (2024, 10, 1, 0, 0, 0)
 FAB_HASH_TRUNCATION = 8
+FAB_MAX_SIZE = 10 * 1024 * 1024  # 10 MB
 FLWR_DIR = ".flwr"  # The default Flower directory: ~/.flwr/
 FLWR_HOME = "FLWR_HOME"  # If set, override the default Flower directory
 
@@ -122,7 +123,7 @@ AUTHZ_TYPE_YAML_KEY = "authz_type"  # For key name in YAML file
 PUBLIC_KEY_HEADER = "flwr-public-key-bin"  # Must end with "-bin" for binary data
 SIGNATURE_HEADER = "flwr-signature-bin"  # Must end with "-bin" for binary data
 TIMESTAMP_HEADER = "flwr-timestamp"
-TIMESTAMP_TOLERANCE = 10  # General tolerance for timestamp verification
+TIMESTAMP_TOLERANCE = 300  # General tolerance for timestamp verification
 SYSTEM_TIME_TOLERANCE = 5  # Allowance for system time drift
 
 # Constants for grpc retry
@@ -134,6 +135,7 @@ GC_THRESHOLD = 200_000_000  # 200 MB
 # Constants for Inflatable
 HEAD_BODY_DIVIDER = b"\x00"
 HEAD_VALUE_DIVIDER = " "
+MAX_ARRAY_CHUNK_SIZE = 20_971_520  # 20 MB
 
 # Constants for serialization
 INT64_MAX_VALUE = 9223372036854775807  # (1 << 63) - 1

flwr/common/exit/exit_code.py

@@ -29,6 +29,9 @@ class ExitCode:
 
     # SuperLink-specific exit codes (100-199)
     SUPERLINK_THREAD_CRASH = 100
+    SUPERLINK_LICENSE_INVALID = 101
+    SUPERLINK_LICENSE_MISSING = 102
+    SUPERLINK_LICENSE_URL_INVALID = 103
 
     # ServerApp-specific exit codes (200-299)
 
@@ -60,6 +63,18 @@ EXIT_CODE_HELP = {
     ExitCode.GRACEFUL_EXIT_SIGTERM: "",
     # SuperLink-specific exit codes (100-199)
     ExitCode.SUPERLINK_THREAD_CRASH: "An important background thread has crashed.",
+    ExitCode.SUPERLINK_LICENSE_INVALID: (
+        "The license is invalid or has expired. "
+        "Please contact `hello@flower.ai` for assistance."
+    ),
+    ExitCode.SUPERLINK_LICENSE_MISSING: (
+        "The license is missing. Please specify the license key by setting the "
+        "environment variable `FLWR_LICENSE_KEY`."
+    ),
+    ExitCode.SUPERLINK_LICENSE_URL_INVALID: (
+        "The license URL is invalid. Please ensure that the `FLWR_LICENSE_URL` "
+        "environment variable is set to a valid URL."
+    ),
     # ServerApp-specific exit codes (200-299)
     # SuperNode-specific exit codes (300-399)
     ExitCode.SUPERNODE_REST_ADDRESS_INVALID: (
@@ -72,7 +87,7 @@ EXIT_CODE_HELP = {
         "to be provided (providing only one of them is not sufficient)."
     ),
     ExitCode.SUPERNODE_NODE_AUTH_KEYS_INVALID: (
-        "Node uthentication requires elliptic curve private and public key pair. "
+        "Node authentication requires elliptic curve private and public key pair. "
         "Please ensure that the file path points to a valid private/public key "
         "file and try again."
     ),

flwr/common/grpc.py

@@ -23,6 +23,9 @@ from logging import DEBUG, ERROR
 from typing import Any, Callable, Optional
 
 import grpc
+from grpc_health.v1.health_pb2_grpc import add_HealthServicer_to_server
+
+from flwr.supercore.grpc_health import SimpleHealthServicer
 
 from .address import is_port_in_use
 from .logger import log
@@ -98,7 +101,7 @@ def valid_certificates(certificates: tuple[bytes, bytes, bytes]) -> bool:
     return is_valid
 
 
-def generic_create_grpc_server(  # pylint: disable=too-many-arguments,R0917
+def generic_create_grpc_server(  # pylint: disable=too-many-arguments, R0914, R0917
     servicer_and_add_fn: tuple[Any, AddServicerToServerFn],
     server_address: str,
     max_concurrent_workers: int = 1000,
@@ -106,6 +109,7 @@ def generic_create_grpc_server(  # pylint: disable=too-many-arguments,R0917
     keepalive_time_ms: int = 210000,
     certificates: Optional[tuple[bytes, bytes, bytes]] = None,
     interceptors: Optional[Sequence[grpc.ServerInterceptor]] = None,
+    health_servicer: Optional[Any] = None,
 ) -> grpc.Server:
     """Create a gRPC server with a single servicer.
 
@@ -153,6 +157,10 @@ def generic_create_grpc_server(  # pylint: disable=too-many-arguments,R0917
             * server private key.
     interceptors : Optional[Sequence[grpc.ServerInterceptor]] (default: None)
         A list of gRPC interceptors.
+    health_servicer : Optional[Any] (default: None)
+        An optional health servicer to add to the server. If provided, it should be an
+        instance of a class that inherits the `HealthServicer` class.
+        If None is provided, `SimpleHealthServicer` will be used by default.
 
     Returns
     -------
@@ -203,6 +211,9 @@ def generic_create_grpc_server(  # pylint: disable=too-many-arguments,R0917
     )
     add_servicer_to_server_fn(servicer, server)
 
+    # Enable health service
+    add_HealthServicer_to_server(health_servicer or SimpleHealthServicer(), server)
+
     if certificates is not None:
         if not valid_certificates(certificates):
             sys.exit(1)

flwr/common/{inflatable_grpc_utils.py → inflatable_protobuf_utils.py}

@@ -18,6 +18,8 @@
 from typing import Callable
 
 from flwr.proto.message_pb2 import (  # pylint: disable=E0611
+    ConfirmMessageReceivedRequest,
+    ConfirmMessageReceivedResponse,
     PullObjectRequest,
     PullObjectResponse,
     PushObjectRequest,
@@ -27,9 +29,13 @@ from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
 
 from .inflatable_utils import ObjectIdNotPreregisteredError, ObjectUnavailableError
 
+ConfirmMessageReceivedProtobuf = Callable[
+    [ConfirmMessageReceivedRequest], ConfirmMessageReceivedResponse
+]
 
-def make_pull_object_fn_grpc(
-    pull_object_grpc: Callable[[PullObjectRequest], PullObjectResponse],
+
+def make_pull_object_fn_protobuf(
+    pull_object_protobuf: Callable[[PullObjectRequest], PullObjectResponse],
     node: Node,
     run_id: int,
 ) -> Callable[[str], bytes]:
@@ -37,8 +43,9 @@ def make_pull_object_fn_grpc(
 
     Parameters
     ----------
-    pull_object_grpc : Callable[[PullObjectRequest], PullObjectResponse]
-        The gRPC function to pull objects, e.g., `FleetStub.PullObject`.
+    pull_object_protobuf : Callable[[PullObjectRequest], PullObjectResponse]
+        A callable that takes a `PullObjectRequest` and returns a `PullObjectResponse`.
+        This function is typically backed by a gRPC client stub.
     node : Node
         The node making the request.
     run_id : int
@@ -54,7 +61,7 @@ def make_pull_object_fn_grpc(
 
     def pull_object_fn(object_id: str) -> bytes:
         request = PullObjectRequest(node=node, run_id=run_id, object_id=object_id)
-        response: PullObjectResponse = pull_object_grpc(request)
+        response: PullObjectResponse = pull_object_protobuf(request)
         if not response.object_found:
             raise ObjectIdNotPreregisteredError(object_id)
         if not response.object_available:
@@ -64,8 +71,8 @@ def make_pull_object_fn_grpc(
     return pull_object_fn
 
 
-def make_push_object_fn_grpc(
-    push_object_grpc: Callable[[PushObjectRequest], PushObjectResponse],
+def make_push_object_fn_protobuf(
+    push_object_protobuf: Callable[[PushObjectRequest], PushObjectResponse],
     node: Node,
     run_id: int,
 ) -> Callable[[str, bytes], None]:
@@ -73,8 +80,9 @@ def make_push_object_fn_grpc(
 
     Parameters
     ----------
-    push_object_grpc : Callable[[PushObjectRequest], PushObjectResponse]
-        The gRPC function to push objects, e.g., `FleetStub.PushObject`.
+    push_object_protobuf : Callable[[PushObjectRequest], PushObjectResponse]
+        A callable that takes a `PushObjectRequest` and returns a `PushObjectResponse`.
+        This function is typically backed by a gRPC client stub.
     node : Node
         The node making the request.
     run_id : int
@@ -92,8 +100,42 @@ def make_push_object_fn_grpc(
         request = PushObjectRequest(
             node=node, run_id=run_id, object_id=object_id, object_content=object_content
         )
-        response: PushObjectResponse = push_object_grpc(request)
+        response: PushObjectResponse = push_object_protobuf(request)
         if not response.stored:
             raise ObjectIdNotPreregisteredError(object_id)
 
     return push_object_fn
+
+
+def make_confirm_message_received_fn_protobuf(
+    confirm_message_received_protobuf: ConfirmMessageReceivedProtobuf,
+    node: Node,
+    run_id: int,
+) -> Callable[[str], None]:
+    """Create a confirm message received function that uses protobuf.
+
+    Parameters
+    ----------
+    confirm_message_received_protobuf : ConfirmMessageReceivedProtobuf
+        A callable that takes a `ConfirmMessageReceivedRequest` and returns a
+        `ConfirmMessageReceivedResponse`, confirming message receipt.
+        This function is typically backed by a gRPC client stub.
+    node : Node
+        The node making the request.
+    run_id : int
+        The run ID for the current message.
+
+    Returns
+    -------
+    Callable[[str], None]
+        A wrapper function that takes an object ID and confirms that
+        the message has been received.
+    """
+
+    def confirm_message_received_fn(object_id: str) -> None:
+        request = ConfirmMessageReceivedRequest(
+            node=node, run_id=run_id, message_object_id=object_id
+        )
+        confirm_message_received_protobuf(request)
+
+    return confirm_message_received_fn