flwr-nightly 1.19.0.dev20250609__py3-none-any.whl → 1.19.0.dev20250611__py3-none-any.whl

flwr/client/grpc_rere_client/connection.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 """Contextmanager for a gRPC request-response channel to the Flower server."""
 
+
 from collections.abc import Iterator, Sequence
 from contextlib import contextmanager
 from copy import copy
@@ -32,9 +33,11 @@ from flwr.common.grpc import create_channel, on_channel_state_change
 from flwr.common.heartbeat import HeartbeatSender
 from flwr.common.inflatable import get_all_nested_objects
 from flwr.common.inflatable_grpc_utils import (
-    inflate_object_from_contents,
     make_pull_object_fn_grpc,
     make_push_object_fn_grpc,
+)
+from flwr.common.inflatable_utils import (
+    inflate_object_from_contents,
     pull_objects,
     push_objects,
 )

flwr/client/rest_client/connection.py

@@ -14,12 +14,11 @@
 # ==============================================================================
 """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 ERROR, INFO, WARN
-from typing import Callable, Optional, TypeVar, Union
+from logging import DEBUG, ERROR, INFO, WARN
+from typing import Callable, Optional, TypeVar, Union, cast
 
 from cryptography.hazmat.primitives.asymmetric import ec
 from google.protobuf.message import Message as GrpcMessage
@@ -31,10 +30,20 @@ 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
+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.logger import log
-from flwr.common.message import Message
+from flwr.common.message import Message, remove_content_from_message
 from flwr.common.retry_invoker import RetryInvoker
-from flwr.common.serde import message_from_proto, message_to_proto, run_from_proto
+from flwr.common.serde import 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,6 +60,13 @@ from flwr.proto.heartbeat_pb2 import (  # pylint: disable=E0611
     SendNodeHeartbeatRequest,
     SendNodeHeartbeatResponse,
 )
+from flwr.proto.message_pb2 import ObjectIDs  # pylint: disable=E0611
+from flwr.proto.message_pb2 import (  # pylint: disable=E0611
+    PullObjectRequest,
+    PullObjectResponse,
+    PushObjectRequest,
+    PushObjectResponse,
+)
 from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
 from flwr.proto.run_pb2 import GetRunRequest, GetRunResponse  # pylint: disable=E0611
 
@@ -64,6 +80,8 @@ PATH_CREATE_NODE: str = "api/v0/fleet/create-node"
 PATH_DELETE_NODE: str = "api/v0/fleet/delete-node"
 PATH_PULL_MESSAGES: str = "/api/v0/fleet/pull-messages"
 PATH_PUSH_MESSAGES: str = "/api/v0/fleet/push-messages"
+PATH_PULL_OBJECT: str = "/api/v0/fleet/pull-object"
+PATH_PUSH_OBJECT: str = "/api/v0/fleet/push-object"
 PATH_SEND_NODE_HEARTBEAT: str = "api/v0/fleet/send-node-heartbeat"
 PATH_GET_RUN: str = "/api/v0/fleet/get-run"
 PATH_GET_FAB: str = "/api/v0/fleet/get-fab"
@@ -296,14 +314,48 @@ def http_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
         ):
             message_proto = None
 
-        # Return the Message if available
-        nonlocal metadata
-        message = None
-        if message_proto is not None:
-            message = message_from_proto(message_proto)
-            metadata = copy(message.metadata)
+        # Construct the Message
+        in_message: Optional[Message] = None
+
+        if message_proto:
             log(INFO, "[Node] POST /%s: success", PATH_PULL_MESSAGES)
-        return message
+            msg_id = message_proto.metadata.message_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=message_proto.metadata.run_id,
+                    ),
+                )
+            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
+
+        return in_message
 
     def send(message: Message) -> None:
         """Send Message result back to server."""
@@ -318,29 +370,69 @@ def http_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
             log(ERROR, "No current message")
             return
 
+        # 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
-        metadata = None
 
-        # Serialize ProtoBuf to bytes
-        message_proto = message_to_proto(message=message)
-
-        # Serialize ProtoBuf to bytes
-        req = PushMessagesRequest(node=node, messages_list=[message_proto])
+        # Get all nested objects
+        all_objects = get_all_nested_objects(message)
+        all_object_ids = list(all_objects.keys())
+        msg_id = all_object_ids[-1]  # Last object is the message itself
+        descendant_ids = all_object_ids[:-1]  # All but the last object are descendants
+
+        # Serialize Message
+        message_proto = message_to_proto(message=remove_content_from_message(message))
+        req = PushMessagesRequest(
+            node=node,
+            messages_list=[message_proto],
+            msg_to_descendant_mapping={msg_id: ObjectIDs(object_ids=descendant_ids)},
+        )
 
         # Send the request
         res = _request(req, PushMessagesResponse, PATH_PUSH_MESSAGES)
-        if res is None:
-            return
+        if res:
+            log(
+                INFO,
+                "[Node] POST /%s: success, created result %s",
+                PATH_PUSH_MESSAGES,
+                res.results,  # pylint: disable=no-member
+            )
 
-        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))
+
+        # Cleanup
+        metadata = None
 
     def get_run(run_id: int) -> Run:
         # Construct the request

flwr/common/auth_plugin/auth_plugin.py

@@ -20,7 +20,7 @@ from collections.abc import Sequence
 from pathlib import Path
 from typing import Optional, Union
 
-from flwr.common.typing import UserInfo
+from flwr.common.typing import AccountInfo
 from flwr.proto.exec_pb2_grpc import ExecStub
 
 from ..typing import UserAuthCredentials, UserAuthLoginDetails
@@ -53,7 +53,7 @@ class ExecAuthPlugin(ABC):
     @abstractmethod
     def validate_tokens_in_metadata(
         self, metadata: Sequence[tuple[str, Union[str, bytes]]]
-    ) -> tuple[bool, Optional[UserInfo]]:
+    ) -> tuple[bool, Optional[AccountInfo]]:
         """Validate authentication tokens in the provided metadata."""
 
     @abstractmethod
@@ -63,7 +63,9 @@ class ExecAuthPlugin(ABC):
     @abstractmethod
     def refresh_tokens(
         self, metadata: Sequence[tuple[str, Union[str, bytes]]]
-    ) -> tuple[Optional[Sequence[tuple[str, Union[str, bytes]]]], Optional[UserInfo]]:
+    ) -> tuple[
+        Optional[Sequence[tuple[str, Union[str, bytes]]]], Optional[AccountInfo]
+    ]:
         """Refresh authentication tokens in the provided metadata."""
 
 
@@ -84,7 +86,7 @@ class ExecAuthzPlugin(ABC):  # pylint: disable=too-few-public-methods
         """Abstract constructor."""
 
     @abstractmethod
-    def verify_user_authorization(self, user_info: UserInfo) -> bool:
+    def verify_user_authorization(self, account_info: AccountInfo) -> bool:
         """Verify user authorization request."""
 
 

flwr/common/event_log_plugin/event_log_plugin.py

@@ -21,7 +21,7 @@ from typing import Optional, Union
 import grpc
 from google.protobuf.message import Message as GrpcMessage
 
-from flwr.common.typing import LogEntry, UserInfo
+from flwr.common.typing import AccountInfo, LogEntry
 
 
 class EventLogWriterPlugin(ABC):
@@ -36,7 +36,7 @@ class EventLogWriterPlugin(ABC):
         self,
         request: GrpcMessage,
         context: grpc.ServicerContext,
-        user_info: Optional[UserInfo],
+        account_info: Optional[AccountInfo],
         method_name: str,
     ) -> LogEntry:
         """Compose pre-event log entry from the provided request and context."""
@@ -46,7 +46,7 @@ class EventLogWriterPlugin(ABC):
         self,
         request: GrpcMessage,
         context: grpc.ServicerContext,
-        user_info: Optional[UserInfo],
+        account_info: Optional[AccountInfo],
         method_name: str,
         response: Optional[Union[GrpcMessage, BaseException]],
     ) -> LogEntry:

flwr/common/inflatable.py

@@ -18,6 +18,9 @@
 from __future__ import annotations
 
 import hashlib
+import threading
+from collections.abc import Iterator
+from contextlib import contextmanager
 from typing import TypeVar, cast
 
 from .constant import HEAD_BODY_DIVIDER, HEAD_VALUE_DIVIDER
@@ -33,6 +36,33 @@ class UnexpectedObjectContentError(Exception):
         )
 
 
+_ctx = threading.local()
+
+
+def _is_recompute_enabled() -> bool:
+    """Check if recomputing object IDs is enabled."""
+    return getattr(_ctx, "recompute_object_id_enabled", True)
+
+
+def _get_computed_object_ids() -> set[str]:
+    """Get the set of computed object IDs."""
+    return getattr(_ctx, "computed_object_ids", set())
+
+
+@contextmanager
+def no_object_id_recompute() -> Iterator[None]:
+    """Context manager to disable recomputing object IDs."""
+    old_value = _is_recompute_enabled()
+    old_set = _get_computed_object_ids()
+    _ctx.recompute_object_id_enabled = False
+    _ctx.computed_object_ids = set()
+    try:
+        yield
+    finally:
+        _ctx.recompute_object_id_enabled = old_value
+        _ctx.computed_object_ids = old_set
+
+
 class InflatableObject:
     """Base class for inflatable objects."""
 
@@ -65,8 +95,23 @@ class InflatableObject:
     @property
     def object_id(self) -> str:
         """Get object_id."""
+        # If recomputing object ID is disabled and the object ID is already computed,
+        # return the cached object ID.
+        if (
+            not _is_recompute_enabled()
+            and (obj_id := self.__dict__.get("_object_id"))
+            in _get_computed_object_ids()
+        ):
+            return cast(str, obj_id)
+
         if self.is_dirty or "_object_id" not in self.__dict__:
-            self.__dict__["_object_id"] = get_object_id(self.deflate())
+            obj_id = get_object_id(self.deflate())
+            self.__dict__["_object_id"] = obj_id
+
+            # If recomputing object ID is disabled, add the object ID to the set of
+            # computed object IDs to avoid recomputing it within the context.
+            if not _is_recompute_enabled():
+                _get_computed_object_ids().add(obj_id)
         return cast(str, self.__dict__["_object_id"])
 
     @property

flwr/common/inflatable_grpc_utils.py

@@ -12,14 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""InflatableObject utils."""
+"""InflatableObject gRPC utils."""
 
 
-import concurrent.futures
-import random
-import threading
-import time
-from typing import Callable, Optional
+from typing import Callable
 
 from flwr.proto.message_pb2 import (  # pylint: disable=E0611
     PullObjectRequest,
@@ -29,42 +25,7 @@ from flwr.proto.message_pb2 import (  # pylint: disable=E0611
 )
 from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
 
-from .constant import (
-    MAX_CONCURRENT_PULLS,
-    MAX_CONCURRENT_PUSHES,
-    PULL_BACKOFF_CAP,
-    PULL_INITIAL_BACKOFF,
-    PULL_MAX_TIME,
-    PULL_MAX_TRIES_PER_OBJECT,
-)
-from .inflatable import InflatableObject, get_object_head_values_from_object_content
-from .message import Message
-from .record import Array, ArrayRecord, ConfigRecord, MetricRecord, RecordDict
-
-# Helper registry that maps names of classes to their type
-inflatable_class_registry: dict[str, type[InflatableObject]] = {
-    Array.__qualname__: Array,
-    ArrayRecord.__qualname__: ArrayRecord,
-    ConfigRecord.__qualname__: ConfigRecord,
-    Message.__qualname__: Message,
-    MetricRecord.__qualname__: MetricRecord,
-    RecordDict.__qualname__: RecordDict,
-}
-
-
-class ObjectUnavailableError(Exception):
-    """Exception raised when an object has been pre-registered but is not yet
-    available."""
-
-    def __init__(self, object_id: str):
-        super().__init__(f"Object with ID '{object_id}' is not yet available.")
-
-
-class ObjectIdNotPreregisteredError(Exception):
-    """Exception raised when an object ID is not pre-registered."""
-
-    def __init__(self, object_id: str):
-        super().__init__(f"Object with ID '{object_id}' could not be found.")
+from .inflatable_utils import ObjectIdNotPreregisteredError, ObjectUnavailableError
 
 
 def make_pull_object_fn_grpc(
@@ -136,227 +97,3 @@ def make_push_object_fn_grpc(
             raise ObjectIdNotPreregisteredError(object_id)
 
     return push_object_fn
-
-
-def push_objects(
-    objects: dict[str, InflatableObject],
-    push_object_fn: Callable[[str, bytes], None],
-    *,
-    object_ids_to_push: Optional[set[str]] = None,
-    keep_objects: bool = False,
-    max_concurrent_pushes: int = MAX_CONCURRENT_PUSHES,
-) -> None:
-    """Push multiple objects to the servicer.
-
-    Parameters
-    ----------
-    objects : dict[str, InflatableObject]
-        A dictionary of objects to push, where keys are object IDs and values are
-        `InflatableObject` instances.
-    push_object_fn : Callable[[str, bytes], None]
-        A function that takes an object ID and its content as bytes, and pushes
-        it to the servicer. This function should raise `ObjectIdNotPreregisteredError`
-        if the object ID is not pre-registered.
-    object_ids_to_push : Optional[set[str]] (default: None)
-        A set of object IDs to push. If not provided, all objects will be pushed.
-    keep_objects : bool (default: False)
-        If `True`, the original objects will be kept in the `objects` dictionary
-        after pushing. If `False`, they will be removed from the dictionary to avoid
-        high memory usage.
-    max_concurrent_pushes : int (default: MAX_CONCURRENT_PUSHES)
-        The maximum number of concurrent pushes to perform.
-    """
-    if object_ids_to_push is not None:
-        # Filter objects to push only those with IDs in the set
-        objects = {k: v for k, v in objects.items() if k in object_ids_to_push}
-
-    lock = threading.Lock()
-
-    def push(obj_id: str) -> None:
-        """Push a single object."""
-        object_content = objects[obj_id].deflate()
-        if not keep_objects:
-            with lock:
-                del objects[obj_id]
-        push_object_fn(obj_id, object_content)
-
-    with concurrent.futures.ThreadPoolExecutor(
-        max_workers=max_concurrent_pushes
-    ) as executor:
-        list(executor.map(push, list(objects.keys())))
-
-
-def pull_objects(  # pylint: disable=too-many-arguments
-    object_ids: list[str],
-    pull_object_fn: Callable[[str], bytes],
-    *,
-    max_concurrent_pulls: int = MAX_CONCURRENT_PULLS,
-    max_time: Optional[float] = PULL_MAX_TIME,
-    max_tries_per_object: Optional[int] = PULL_MAX_TRIES_PER_OBJECT,
-    initial_backoff: float = PULL_INITIAL_BACKOFF,
-    backoff_cap: float = PULL_BACKOFF_CAP,
-) -> dict[str, bytes]:
-    """Pull multiple objects from the servicer.
-
-    Parameters
-    ----------
-    object_ids : list[str]
-        A list of object IDs to pull.
-    pull_object_fn : Callable[[str], bytes]
-        A function that takes an object ID and returns the object content as bytes.
-        The function should raise `ObjectUnavailableError` if the object is not yet
-        available, or `ObjectIdNotPreregisteredError` if the object ID is not
-        pre-registered.
-    max_concurrent_pulls : int (default: MAX_CONCURRENT_PULLS)
-        The maximum number of concurrent pulls to perform.
-    max_time : Optional[float] (default: PULL_MAX_TIME)
-        The maximum time to wait for all pulls to complete. If `None`, waits
-        indefinitely.
-    max_tries_per_object : Optional[int] (default: PULL_MAX_TRIES_PER_OBJECT)
-        The maximum number of attempts to pull each object. If `None`, pulls
-        indefinitely until the object is available.
-    initial_backoff : float (default: PULL_INITIAL_BACKOFF)
-        The initial backoff time in seconds for retrying pulls after an
-        `ObjectUnavailableError`.
-    backoff_cap : float (default: PULL_BACKOFF_CAP)
-        The maximum backoff time in seconds. Backoff times will not exceed this value.
-
-    Returns
-    -------
-    dict[str, bytes]
-        A dictionary where keys are object IDs and values are the pulled
-        object contents.
-    """
-    if max_tries_per_object is None:
-        max_tries_per_object = int(1e9)
-    if max_time is None:
-        max_time = float("inf")
-
-    results: dict[str, bytes] = {}
-    results_lock = threading.Lock()
-    err_to_raise: Optional[Exception] = None
-    early_stop = threading.Event()
-    start = time.monotonic()
-
-    def pull_with_retries(object_id: str) -> None:
-        """Attempt to pull a single object with retry and backoff."""
-        nonlocal err_to_raise
-        tries = 0
-        delay = initial_backoff
-
-        while not early_stop.is_set():
-            try:
-                object_content = pull_object_fn(object_id)
-                with results_lock:
-                    results[object_id] = object_content
-                return
-
-            except ObjectUnavailableError as err:
-                tries += 1
-                if (
-                    tries >= max_tries_per_object
-                    or time.monotonic() - start >= max_time
-                ):
-                    # Stop all work if one object exhausts retries
-                    early_stop.set()
-                    with results_lock:
-                        if err_to_raise is None:
-                            err_to_raise = err
-                    return
-
-                # Apply exponential backoff with ±20% jitter
-                sleep_time = delay * (1 + random.uniform(-0.2, 0.2))
-                early_stop.wait(sleep_time)
-                delay = min(delay * 2, backoff_cap)
-
-            except ObjectIdNotPreregisteredError as err:
-                # Permanent failure: object ID is invalid
-                early_stop.set()
-                with results_lock:
-                    if err_to_raise is None:
-                        err_to_raise = err
-                return
-
-    # Submit all pull tasks concurrently
-    with concurrent.futures.ThreadPoolExecutor(
-        max_workers=max_concurrent_pulls
-    ) as executor:
-        futures = {
-            executor.submit(pull_with_retries, obj_id): obj_id for obj_id in object_ids
-        }
-
-        # Wait for completion
-        concurrent.futures.wait(futures)
-
-    if err_to_raise is not None:
-        raise err_to_raise
-
-    return results
-
-
-def inflate_object_from_contents(
-    object_id: str,
-    object_contents: dict[str, bytes],
-    *,
-    keep_object_contents: bool = False,
-    objects: Optional[dict[str, InflatableObject]] = None,
-) -> InflatableObject:
-    """Inflate an object from object contents.
-
-    Parameters
-    ----------
-    object_id : str
-        The ID of the object to inflate.
-    object_contents : dict[str, bytes]
-        A dictionary mapping object IDs to their contents as bytes.
-        All descendant objects must be present in this dictionary.
-    keep_object_contents : bool (default: False)
-        If `True`, the object content will be kept in the `object_contents`
-        dictionary after inflation. If `False`, the object content will be
-        removed from the dictionary to save memory.
-    objects : Optional[dict[str, InflatableObject]] (default: None)
-        No need to provide this parameter. A dictionary to store already
-        inflated objects, mapping object IDs to their corresponding
-        `InflatableObject` instances.
-
-    Returns
-    -------
-    InflatableObject
-        The inflated object.
-    """
-    if objects is None:
-        # Initialize objects dictionary
-        objects = {}
-
-    if object_id in objects:
-        # If the object is already in the objects dictionary, return it
-        return objects[object_id]
-
-    # Extract object class and object_ids of children
-    object_content = object_contents[object_id]
-    obj_type, children_obj_ids, _ = get_object_head_values_from_object_content(
-        object_content=object_contents[object_id]
-    )
-
-    # Remove the object content from the dictionary to save memory
-    if not keep_object_contents:
-        del object_contents[object_id]
-
-    # Resolve object class
-    cls_type = inflatable_class_registry[obj_type]
-
-    # Inflate all children objects
-    children: dict[str, InflatableObject] = {}
-    for child_obj_id in children_obj_ids:
-        children[child_obj_id] = inflate_object_from_contents(
-            child_obj_id,
-            object_contents,
-            keep_object_contents=keep_object_contents,
-            objects=objects,
-        )
-
-    # Inflate object passing its children
-    obj = cls_type.inflate(object_content, children=children)
-    del object_content  # Free memory after inflation
-    objects[object_id] = obj
-    return obj