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

flwr/client/grpc_rere_client/connection.py

@@ -32,9 +32,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/common/inflatable_grpc_utils.py

@@ -15,11 +15,7 @@
 """InflatableObject 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

flwr/common/inflatable_utils.py

@@ -14,15 +14,281 @@
 # ==============================================================================
 """InflatableObject utilities."""
 
+import concurrent.futures
+import random
+import threading
+import time
+from typing import Callable, Optional
 
-from .constant import HEAD_BODY_DIVIDER, HEAD_VALUE_DIVIDER
+from .constant import (
+    HEAD_BODY_DIVIDER,
+    HEAD_VALUE_DIVIDER,
+    MAX_CONCURRENT_PULLS,
+    MAX_CONCURRENT_PUSHES,
+    PULL_BACKOFF_CAP,
+    PULL_INITIAL_BACKOFF,
+    PULL_MAX_TIME,
+    PULL_MAX_TRIES_PER_OBJECT,
+)
 from .inflatable import (
+    InflatableObject,
     UnexpectedObjectContentError,
     _get_object_head,
+    get_object_head_values_from_object_content,
     get_object_id,
     is_valid_sha256_hash,
 )
-from .inflatable_grpc_utils import inflatable_class_registry
+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.")
+
+
+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
 
 
 def validate_object_content(content: bytes) -> None:

flwr/server/grid/grpc_grid.py

@@ -30,9 +30,11 @@ from flwr.common.constant import (
 from flwr.common.grpc import create_channel, on_channel_state_change
 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/server/superlink/fleet/grpc_rere/fleet_servicer.py

@@ -15,12 +15,11 @@
 """Fleet API gRPC request-response servicer."""
 
 
-from logging import DEBUG, ERROR, INFO
+from logging import DEBUG, INFO
 
 import grpc
 from google.protobuf.json_format import MessageToDict
 
-from flwr.common.constant import Status
 from flwr.common.inflatable import UnexpectedObjectContentError
 from flwr.common.logger import log
 from flwr.common.typing import InvalidRunStatusException
@@ -50,9 +49,8 @@ from flwr.proto.run_pb2 import GetRunRequest, GetRunResponse  # pylint: disable=
 from flwr.server.superlink.ffs.ffs_factory import FfsFactory
 from flwr.server.superlink.fleet.message_handler import message_handler
 from flwr.server.superlink.linkstate import LinkStateFactory
-from flwr.server.superlink.utils import abort_grpc_context, check_abort
+from flwr.server.superlink.utils import abort_grpc_context
 from flwr.supercore.object_store import ObjectStoreFactory
-from flwr.supercore.object_store.object_store import NoObjectInStoreError
 
 
 class FleetServicer(fleet_pb2_grpc.FleetServicer):
@@ -185,36 +183,20 @@ class FleetServicer(fleet_pb2_grpc.FleetServicer):
             request.object_id,
         )
 
-        state = self.state_factory.state()
-
-        # Abort if the run is not running
-        abort_msg = check_abort(
-            request.run_id,
-            [Status.PENDING, Status.STARTING, Status.FINISHED],
-            state,
-        )
-        if abort_msg:
-            abort_grpc_context(abort_msg, context)
-
-        if request.node.node_id not in state.get_nodes(run_id=request.run_id):
-            # Cancel insertion in ObjectStore
-            context.abort(grpc.StatusCode.FAILED_PRECONDITION, "Unexpected node ID.")
-
-        # Init store
-        store = self.objectstore_factory.store()
-
-        # Insert in store
-        stored = False
         try:
-            store.put(request.object_id, request.object_content)
-            stored = True
-        except (NoObjectInStoreError, ValueError) as e:
-            log(ERROR, str(e))
+            # Insert in Store
+            res = message_handler.push_object(
+                request=request,
+                state=self.state_factory.state(),
+                store=self.objectstore_factory.store(),
+            )
+        except InvalidRunStatusException as e:
+            abort_grpc_context(e.message, context)
         except UnexpectedObjectContentError as e:
             # Object content is not valid
             context.abort(grpc.StatusCode.FAILED_PRECONDITION, str(e))
 
-        return PushObjectResponse(stored=stored)
+        return res
 
     def PullObject(
         self, request: PullObjectRequest, context: grpc.ServicerContext
@@ -226,31 +208,14 @@ class FleetServicer(fleet_pb2_grpc.FleetServicer):
             request.object_id,
         )
 
-        state = self.state_factory.state()
-
-        # Abort if the run is not running
-        abort_msg = check_abort(
-            request.run_id,
-            [Status.PENDING, Status.STARTING, Status.FINISHED],
-            state,
-        )
-        if abort_msg:
-            abort_grpc_context(abort_msg, context)
-
-        if request.node.node_id not in state.get_nodes(run_id=request.run_id):
-            # Cancel insertion in ObjectStore
-            context.abort(grpc.StatusCode.FAILED_PRECONDITION, "Unexpected node ID.")
-
-        # Init store
-        store = self.objectstore_factory.store()
-
-        # Fetch from store
-        content = store.get(request.object_id)
-        if content is not None:
-            object_available = content != b""
-            return PullObjectResponse(
-                object_found=True,
-                object_available=object_available,
-                object_content=content,
+        try:
+            # Fetch from store
+            res = message_handler.pull_object(
+                request=request,
+                state=self.state_factory.state(),
+                store=self.objectstore_factory.store(),
             )
-        return PullObjectResponse(object_found=False, object_available=False)
+        except InvalidRunStatusException as e:
+            abort_grpc_context(e.message, context)
+
+        return res

flwr/server/superlink/fleet/message_handler/message_handler.py

@@ -19,6 +19,7 @@ from typing import Optional
 
 from flwr.common import Message, log
 from flwr.common.constant import Status
+from flwr.common.inflatable import UnexpectedObjectContentError
 from flwr.common.serde import (
     fab_to_proto,
     message_from_proto,
@@ -42,7 +43,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
+    ObjectIDs,
+    PullObjectRequest,
+    PullObjectResponse,
+    PushObjectRequest,
+    PushObjectResponse,
+)
 from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
 from flwr.proto.run_pb2 import (  # pylint: disable=E0611
     GetRunRequest,
@@ -203,3 +210,52 @@ def get_fab(
         return GetFabResponse(fab=fab_to_proto(fab))
 
     raise ValueError(f"Found no FAB with hash: {request.hash_str}")
+
+
+def push_object(
+    request: PushObjectRequest, state: LinkState, store: ObjectStore
+) -> PushObjectResponse:
+    """Push Object."""
+    abort_msg = check_abort(
+        request.run_id,
+        [Status.PENDING, Status.STARTING, Status.FINISHED],
+        state,
+    )
+    if abort_msg:
+        raise InvalidRunStatusException(abort_msg)
+
+    stored = False
+    try:
+        store.put(request.object_id, request.object_content)
+        stored = True
+    except (NoObjectInStoreError, ValueError) as e:
+        log(ERROR, str(e))
+    except UnexpectedObjectContentError as e:
+        # Object content is not valid
+        log(ERROR, str(e))
+        raise
+    return PushObjectResponse(stored=stored)
+
+
+def pull_object(
+    request: PullObjectRequest, state: LinkState, store: ObjectStore
+) -> PullObjectResponse:
+    """Pull Object."""
+    abort_msg = check_abort(
+        request.run_id,
+        [Status.PENDING, Status.STARTING, Status.FINISHED],
+        state,
+    )
+    if abort_msg:
+        raise InvalidRunStatusException(abort_msg)
+
+    # Fetch from store
+    content = store.get(request.object_id)
+    if content is not None:
+        object_available = content != b""
+        return PullObjectResponse(
+            object_found=True,
+            object_available=object_available,
+            object_content=content,
+        )
+    return PullObjectResponse(object_found=False, object_available=False)

flwr/supernode/start_client_internal.py

@@ -52,9 +52,9 @@ from flwr.common.logger import log
 from flwr.common.retry_invoker import RetryInvoker, RetryState, exponential
 from flwr.common.typing import Fab, Run, RunNotRunningException, UserConfig
 from flwr.proto.clientappio_pb2_grpc import add_ClientAppIoServicer_to_server
-from flwr.server.superlink.ffs import FfsFactory
-from flwr.supercore.object_store import ObjectStoreFactory
-from flwr.supernode.nodestate import NodeStateFactory
+from flwr.server.superlink.ffs import Ffs, FfsFactory
+from flwr.supercore.object_store import ObjectStore, ObjectStoreFactory
+from flwr.supernode.nodestate import NodeState, NodeStateFactory
 from flwr.supernode.servicer.clientappio import ClientAppInputs, ClientAppIoServicer
 
 DEFAULT_FFS_DIR = get_flwr_dir() / "supernode" / "ffs"
@@ -145,7 +145,7 @@ def start_client_internal(
     # Initialize NodeState, Ffs, and ObjectStore
     state = state_factory.state()
     ffs = ffs_factory.ffs()
-    _store = object_store_factory.store()
+    store = object_store_factory.store()
 
     with _init_connection(
         transport=transport,
@@ -166,63 +166,20 @@ def start_client_internal(
 
         # pylint: disable=too-many-nested-blocks
         while True:
-            # Pull message
-            if (message := receive()) is None:
-                time.sleep(3)
-                continue
-
-            # Log message reception
-            log(INFO, "")
-            if message.metadata.group_id:
-                log(
-                    INFO,
-                    "[RUN %s, ROUND %s]",
-                    message.metadata.run_id,
-                    message.metadata.group_id,
-                )
-            else:
-                log(INFO, "[RUN %s]", message.metadata.run_id)
-            log(
-                INFO,
-                "Received: %s message %s",
-                message.metadata.message_type,
-                message.metadata.message_id,
+            # The signature of the function will change after
+            # completing the transition to the `NodeState`-based SuperNode
+            run_id = _pull_and_store_message(
+                state=state,
+                ffs=ffs,
+                object_store=store,
+                node_config=node_config,
+                receive=receive,
+                get_run=get_run,
+                get_fab=get_fab,
             )
 
-            # Ensure the run and FAB are available
-            run_id = message.metadata.run_id
-            try:
-                # Check if the message is from an unknown run
-                if (run_info := state.get_run(run_id)) is None:
-                    # Pull run info from SuperLink
-                    run_info = get_run(run_id)
-                    state.store_run(run_info)
-
-                    # Pull and store the FAB
-                    fab = get_fab(run_info.fab_hash, run_id)
-                    ffs.put(fab.content, {})
-
-                    # Initialize the context
-                    run_cfg = get_fused_config_from_fab(fab.content, run_info)
-                    run_ctx = Context(
-                        run_id=run_id,
-                        node_id=state.get_node_id(),
-                        node_config=node_config,
-                        state=RecordDict(),
-                        run_config=run_cfg,
-                    )
-                    state.store_context(run_ctx)
-
-                # Store the message in the state
-                state.store_message(message)
-            except RunNotRunningException:
-                log(
-                    INFO,
-                    "Run ID %s is not in `RUNNING` status. Ignoring message %s.",
-                    run_id,
-                    message.metadata.message_id,
-                )
-                time.sleep(3)
+            if run_id is None:
+                time.sleep(3)  # Wait for 3s before asking again
                 continue
 
             try:
@@ -302,6 +259,91 @@ def start_client_internal(
                 log(INFO, "")
 
 
+def _pull_and_store_message(  # pylint: disable=too-many-positional-arguments
+    state: NodeState,
+    ffs: Ffs,
+    object_store: ObjectStore,  # pylint: disable=unused-argument
+    node_config: UserConfig,
+    receive: Callable[[], Optional[Message]],
+    get_run: Callable[[int], Run],
+    get_fab: Callable[[str, int], Fab],
+) -> Optional[int]:
+    """Pull a message from the SuperLink and store it in the state.
+
+    This function current returns None if no message is received,
+    or run_id if a message is received and processed successfully.
+    This behavior will change in the future to return None after
+    completing transition to the `NodeState`-based SuperNode.
+    """
+    message = None
+    try:
+        # Pull message
+        if (message := receive()) is None:
+            return None
+
+        # Log message reception
+        log(INFO, "")
+        if message.metadata.group_id:
+            log(
+                INFO,
+                "[RUN %s, ROUND %s]",
+                message.metadata.run_id,
+                message.metadata.group_id,
+            )
+        else:
+            log(INFO, "[RUN %s]", message.metadata.run_id)
+        log(
+            INFO,
+            "Received: %s message %s",
+            message.metadata.message_type,
+            message.metadata.message_id,
+        )
+
+        # Ensure the run and FAB are available
+        run_id = message.metadata.run_id
+
+        # Check if the message is from an unknown run
+        if (run_info := state.get_run(run_id)) is None:
+            # Pull run info from SuperLink
+            run_info = get_run(run_id)
+            state.store_run(run_info)
+
+            # Pull and store the FAB
+            fab = get_fab(run_info.fab_hash, run_id)
+            ffs.put(fab.content, {})
+
+            # Initialize the context
+            run_cfg = get_fused_config_from_fab(fab.content, run_info)
+            run_ctx = Context(
+                run_id=run_id,
+                node_id=state.get_node_id(),
+                node_config=node_config,
+                state=RecordDict(),
+                run_config=run_cfg,
+            )
+            state.store_context(run_ctx)
+
+        # Store the message in the state
+        state.store_message(message)
+    except RunNotRunningException:
+        if message is None:
+            log(
+                INFO,
+                "Run transitioned to a non-`RUNNING` status while receiving a message. "
+                "Ignoring the message.",
+            )
+        else:
+            log(
+                INFO,
+                "Run ID %s is not in `RUNNING` status. Ignoring message %s.",
+                run_id,
+                message.metadata.message_id,
+            )
+        return None
+
+    return run_id
+
+
 @contextmanager
 def _init_connection(  # pylint: disable=too-many-positional-arguments
     transport: str,

{flwr_nightly-1.19.0.dev20250609.dist-info → flwr_nightly-1.19.0.dev20250610.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: flwr-nightly
-Version: 1.19.0.dev20250609
+Version: 1.19.0.dev20250610
 Summary: Flower: A Friendly Federated AI Framework
 License: Apache-2.0
 Keywords: Artificial Intelligence,Federated AI,Federated Analytics,Federated Evaluation,Federated Learning,Flower,Machine Learning

{flwr_nightly-1.19.0.dev20250609.dist-info → flwr_nightly-1.19.0.dev20250610.dist-info}/RECORD

@@ -84,7 +84,7 @@ flwr/client/grpc_adapter_client/__init__.py,sha256=RQWP5mFPROLHKgombiRvPXVWSoVrQ
 flwr/client/grpc_adapter_client/connection.py,sha256=aj5tTYyE8z2hQLXPPydsJiz8gBDIWLUhfWvqYkAL1L4,3966
 flwr/client/grpc_rere_client/__init__.py,sha256=i7iS0Lt8B7q0E2L72e4F_YrKm6ClRKnd71PNA6PW2O0,752
 flwr/client/grpc_rere_client/client_interceptor.py,sha256=zFaVHw6AxeNO-7eCKKb-RxrPa7zbM5Z-2-1Efc4adQY,2451
-flwr/client/grpc_rere_client/connection.py,sha256=nk4xSjBKG_mvoZsCVkkGLPPc6juxBId7RNcckgA6VSE,13994
+flwr/client/grpc_rere_client/connection.py,sha256=zxnr1Ppp3PpffvqhRlTBwTWHnek15lIDEBLfgqCSW40,14039
 flwr/client/grpc_rere_client/grpc_adapter.py,sha256=JvMZ7vCFTaTEo6AzKYh3zDmeQAU7VSjdysbC6t3ufWg,6351
 flwr/client/message_handler/__init__.py,sha256=0lyljDVqre3WljiZbPcwCCf8GiIaSVI_yo_ylEyPwSE,719
 flwr/client/message_handler/message_handler.py,sha256=X9SXX6et97Lw9_DGD93HKsEBGNjXClcFgc_5aLK0oiU,6541
@@ -123,8 +123,8 @@ flwr/common/exit_handlers.py,sha256=IaqJ60fXZuu7McaRYnoYKtlbH9t4Yl9goNExKqtmQbs,
 flwr/common/grpc.py,sha256=manTaHaPiyYngUq1ErZvvV2B2GxlXUUUGRy3jc3TBIQ,9798
 flwr/common/heartbeat.py,sha256=SyEpNDnmJ0lni0cWO67rcoJVKasCLmkNHm3dKLeNrLU,5749
 flwr/common/inflatable.py,sha256=gdAICtXklkQRMrxoTYEbzJl7AeFqZtUm4JU6f2it9FM,7264
-flwr/common/inflatable_grpc_utils.py,sha256=M982VQ8x1wh4QnLXR-IKWdtGuOFACvIsU6lYbbs1TI8,13040
-flwr/common/inflatable_utils.py,sha256=hPwap-P9pxHcxo-EtxtrnKNiFA2HsNGJaKX2b002PnE,2892
+flwr/common/inflatable_grpc_utils.py,sha256=-cdu-VNBK7aNdYRtjtTz3pnXmK-_q3XIzpr14vhanto,3549
+flwr/common/inflatable_utils.py,sha256=-GTdgR1zLS9WtXrbOGJMpaoyVEL8KmoQ2yF4HeLxTI0,12406
 flwr/common/logger.py,sha256=JbRf6E2vQxXzpDBq1T8IDUJo_usu3gjWEBPQ6uKcmdg,13049
 flwr/common/message.py,sha256=xAL7iZN5-n-xPQpgoSFvxNrzs8fmiiPfoU0DjNQEhRw,19953
 flwr/common/object_ref.py,sha256=p3SfTeqo3Aj16SkB-vsnNn01zswOPdGNBitcbRnqmUk,9134
@@ -237,7 +237,7 @@ flwr/server/criterion.py,sha256=G4e-6B48Pc7d5rmGVUpIzNKb6UF88O3VmTRuUltgjzM,1061
 flwr/server/fleet_event_log_interceptor.py,sha256=AkL7Y5d3xm2vRhL3ahmEVVoOvAP7PA7dRgB-je4v-Ys,3774
 flwr/server/grid/__init__.py,sha256=aWZHezoR2UGMJISB_gPMCm2N_2GSbm97A3lAp7ruhRQ,888
 flwr/server/grid/grid.py,sha256=naGCYt5J6dnmUvrcGkdNyKPe3MBd-0awGm1ALmgahqY,6625
-flwr/server/grid/grpc_grid.py,sha256=rVzPh0BGQOI3Dzhd8upBXNLEA5bKy82eBZkU4a1qFJY,13257
+flwr/server/grid/grpc_grid.py,sha256=qhJPS4tCWATHx24dDZoxtEPq6rjFu0lNggX70lstLkw,13302
 flwr/server/grid/inmemory_grid.py,sha256=RjejYT-d-hHuTs1KSs_5wvOdAWKLus8w5_UAcnGt4iw,6168
 flwr/server/history.py,sha256=cCkFhBN4GoHsYYNk5GG1Y089eKJh2DH_ZJbYPwLaGyk,5026
 flwr/server/run_serverapp.py,sha256=v0p6jXj2dFxlRUdoEeF1mnaFd9XRQi6dZCflPY6d3qI,2063
@@ -285,10 +285,10 @@ flwr/server/superlink/fleet/grpc_bidi/grpc_bridge.py,sha256=KouR9PUcrPmMtoLooF4O
 flwr/server/superlink/fleet/grpc_bidi/grpc_client_proxy.py,sha256=iSf0mbBAlig7G6subQwBSVjcUCgSihONKdZ1RmQPTOk,4887
 flwr/server/superlink/fleet/grpc_bidi/grpc_server.py,sha256=OsS-6GgCIzMMZDVu5Y-OKjynHVUrpdc_5OrtuB-IbU0,5174
 flwr/server/superlink/fleet/grpc_rere/__init__.py,sha256=ahDJJ1e-lDxBpeBMgPk7YZt2wB38_QltcpOC0gLbpFs,758
-flwr/server/superlink/fleet/grpc_rere/fleet_servicer.py,sha256=iD45MSaGxPmkmPp_Y6XR-crUaQIcAcfyw0qJZwt7C0o,9106
+flwr/server/superlink/fleet/grpc_rere/fleet_servicer.py,sha256=wp6I4BKv82dBlIJA8D2gx9CRgahkwnsVnKTkU6M06Og,7802
 flwr/server/superlink/fleet/grpc_rere/server_interceptor.py,sha256=DrHubsaLgJCwCeeJPYogQTiP0xYqjxwnT9rh7OP7BoU,6984
 flwr/server/superlink/fleet/message_handler/__init__.py,sha256=fHsRV0KvJ8HtgSA4_YBsEzuhJLjO8p6xx4aCY2oE1p4,731
-flwr/server/superlink/fleet/message_handler/message_handler.py,sha256=P43PapLZJKbZ0Oo0kP_KcO5zSMvO53SakQgPMiR5d1M,6500
+flwr/server/superlink/fleet/message_handler/message_handler.py,sha256=hbaukJ7EpfBiaXHa_R50MZKTkesxosx5IwFtDyIAO-0,8076
 flwr/server/superlink/fleet/rest_rere/__init__.py,sha256=Lzc93nA7tDqoy-zRUaPG316oqFiZX1HUCL5ELaXY_xw,735
 flwr/server/superlink/fleet/rest_rere/rest_api.py,sha256=jIljUNMvZ8dDvSlkyn1c2y9sDAf_QoBr-q_o3BWxJ7o,7199
 flwr/server/superlink/fleet/vce/__init__.py,sha256=XOKbAWOzlCqEOQ3M2cBYkH7HKA7PxlbCJMunt-ty-DY,784
@@ -357,8 +357,8 @@ flwr/supernode/runtime/run_clientapp.py,sha256=cvWSby7u31u97QapWHxJM-Wer6F1k6mbb
 flwr/supernode/servicer/__init__.py,sha256=lucTzre5WPK7G1YLCfaqg3rbFWdNSb7ZTt-ca8gxdEo,717
 flwr/supernode/servicer/clientappio/__init__.py,sha256=vJyOjO2FXZ2URbnthmdsgs6948wbYfdq1L1V8Um-Lr8,895
 flwr/supernode/servicer/clientappio/clientappio_servicer.py,sha256=LmzkxtNQBn5vVrHc0Bhq2WqaK6-LM2v4kfLBN0PiNNM,8522
-flwr/supernode/start_client_internal.py,sha256=5CwTNV-XmIhwR1jv3G7aQAXGhf6OFWS6U-vmxY1iKGA,16984
-flwr_nightly-1.19.0.dev20250609.dist-info/METADATA,sha256=HloLs2HPXsSLn_UVyxQf1G9yLasbznXJstMQ1Cr1mo4,15910
-flwr_nightly-1.19.0.dev20250609.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-flwr_nightly-1.19.0.dev20250609.dist-info/entry_points.txt,sha256=jNpDXGBGgs21RqUxelF_jwGaxtqFwm-MQyfz-ZqSjrA,367
-flwr_nightly-1.19.0.dev20250609.dist-info/RECORD,,
+flwr/supernode/start_client_internal.py,sha256=AkJ1FsBK6EpK7cmIGcae5WZazPhU71gileiSQogTZ-k,18164
+flwr_nightly-1.19.0.dev20250610.dist-info/METADATA,sha256=98gf7-LsujpWAZw4bsDkaitvVM4M0NT44a2lIGFtjCE,15910
+flwr_nightly-1.19.0.dev20250610.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+flwr_nightly-1.19.0.dev20250610.dist-info/entry_points.txt,sha256=jNpDXGBGgs21RqUxelF_jwGaxtqFwm-MQyfz-ZqSjrA,367
+flwr_nightly-1.19.0.dev20250610.dist-info/RECORD,,