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

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

@@ -14,12 +14,12 @@
 # ==============================================================================
 """Fleet API message handlers."""
 
-
+from logging import ERROR
 from typing import Optional
-from uuid import UUID
 
-from flwr.common import Message
+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,
@@ -33,23 +33,35 @@ from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
     CreateNodeResponse,
     DeleteNodeRequest,
     DeleteNodeResponse,
-    PingRequest,
-    PingResponse,
     PullMessagesRequest,
     PullMessagesResponse,
     PushMessagesRequest,
     PushMessagesResponse,
     Reconnect,
 )
+from flwr.proto.heartbeat_pb2 import (  # pylint: disable=E0611
+    SendNodeHeartbeatRequest,
+    SendNodeHeartbeatResponse,
+)
+from flwr.proto.message_pb2 import (  # pylint: disable=E0611
+    ConfirmMessageReceivedRequest,
+    ConfirmMessageReceivedResponse,
+    PullObjectRequest,
+    PullObjectResponse,
+    PushObjectRequest,
+    PushObjectResponse,
+)
 from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
 from flwr.proto.run_pb2 import (  # pylint: disable=E0611
     GetRunRequest,
     GetRunResponse,
     Run,
 )
-from flwr.server.superlink.ffs.ffs import Ffs
 from flwr.server.superlink.linkstate import LinkState
 from flwr.server.superlink.utils import check_abort
+from flwr.supercore.ffs import Ffs
+from flwr.supercore.object_store import NoObjectInStoreError, ObjectStore
+from flwr.supercore.object_store.utils import store_mapping_and_register_objects
 
 
 def create_node(
@@ -58,7 +70,7 @@ def create_node(
 ) -> CreateNodeResponse:
     """."""
     # Create node
-    node_id = state.create_node(ping_interval=request.ping_interval)
+    node_id = state.create_node(heartbeat_interval=request.heartbeat_interval)
     return CreateNodeResponse(node=Node(node_id=node_id))
 
 
@@ -73,17 +85,21 @@ def delete_node(request: DeleteNodeRequest, state: LinkState) -> DeleteNodeRespo
     return DeleteNodeResponse()
 
 
-def ping(
-    request: PingRequest,  # pylint: disable=unused-argument
+def send_node_heartbeat(
+    request: SendNodeHeartbeatRequest,  # pylint: disable=unused-argument
     state: LinkState,  # pylint: disable=unused-argument
-) -> PingResponse:
+) -> SendNodeHeartbeatResponse:
     """."""
-    res = state.acknowledge_ping(request.node.node_id, request.ping_interval)
-    return PingResponse(success=res)
+    res = state.acknowledge_node_heartbeat(
+        request.node.node_id, request.heartbeat_interval
+    )
+    return SendNodeHeartbeatResponse(success=res)
 
 
 def pull_messages(
-    request: PullMessagesRequest, state: LinkState
+    request: PullMessagesRequest,
+    state: LinkState,
+    store: ObjectStore,
 ) -> PullMessagesResponse:
     """Pull Messages handler."""
     # Get node_id if client node is not anonymous
@@ -95,14 +111,28 @@ def pull_messages(
 
     # Convert to Messages
     msg_proto = []
+    trees = []
     for msg in message_list:
-        msg_proto.append(message_to_proto(msg))
+        try:
+            # Retrieve Message object tree from ObjectStore
+            msg_object_id = msg.metadata.message_id
+            obj_tree = store.get_object_tree(msg_object_id)
+
+            # Add Message and its object tree to the response
+            msg_proto.append(message_to_proto(msg))
+            trees.append(obj_tree)
+        except NoObjectInStoreError as e:
+            log(ERROR, e.message)
+            # Delete message ins from state
+            state.delete_messages(message_ins_ids={msg_object_id})
 
-    return PullMessagesResponse(messages_list=msg_proto)
+    return PullMessagesResponse(messages_list=msg_proto, message_object_trees=trees)
 
 
 def push_messages(
-    request: PushMessagesRequest, state: LinkState
+    request: PushMessagesRequest,
+    state: LinkState,
+    store: ObjectStore,
 ) -> PushMessagesResponse:
     """Push Messages handler."""
     # Convert Message from proto
@@ -113,22 +143,29 @@ def push_messages(
         msg.metadata.run_id,
         [Status.PENDING, Status.STARTING, Status.FINISHED],
         state,
+        store,
     )
     if abort_msg:
         raise InvalidRunStatusException(abort_msg)
 
     # Store Message in State
-    message_id: Optional[UUID] = state.store_message_res(message=msg)
+    message_id: Optional[str] = state.store_message_res(message=msg)
+
+    # Store Message object to descendants mapping and preregister objects
+    objects_to_push = store_mapping_and_register_objects(store, request=request)
 
     # Build response
     response = PushMessagesResponse(
         reconnect=Reconnect(reconnect=5),
         results={str(message_id): 0},
+        objects_to_push=objects_to_push,
     )
     return response
 
 
-def get_run(request: GetRunRequest, state: LinkState) -> GetRunResponse:
+def get_run(
+    request: GetRunRequest, state: LinkState, store: ObjectStore
+) -> GetRunResponse:
     """Get run information."""
     run = state.get_run(request.run_id)
 
@@ -140,6 +177,7 @@ def get_run(request: GetRunRequest, state: LinkState) -> GetRunResponse:
         request.run_id,
         [Status.PENDING, Status.STARTING, Status.FINISHED],
         state,
+        store,
     )
     if abort_msg:
         raise InvalidRunStatusException(abort_msg)
@@ -156,7 +194,7 @@ def get_run(request: GetRunRequest, state: LinkState) -> GetRunResponse:
 
 
 def get_fab(
-    request: GetFabRequest, ffs: Ffs, state: LinkState  # pylint: disable=W0613
+    request: GetFabRequest, ffs: Ffs, state: LinkState, store: ObjectStore
 ) -> GetFabResponse:
     """Get FAB."""
     # Abort if the run is not running
@@ -164,6 +202,7 @@ def get_fab(
         request.run_id,
         [Status.PENDING, Status.STARTING, Status.FINISHED],
         state,
+        store,
     )
     if abort_msg:
         raise InvalidRunStatusException(abort_msg)
@@ -173,3 +212,75 @@ 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,
+        store,
+    )
+    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,
+        store,
+    )
+    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)
+
+
+def confirm_message_received(
+    request: ConfirmMessageReceivedRequest,
+    state: LinkState,
+    store: ObjectStore,
+) -> ConfirmMessageReceivedResponse:
+    """Confirm message received handler."""
+    abort_msg = check_abort(
+        request.run_id,
+        [Status.PENDING, Status.STARTING, Status.FINISHED],
+        state,
+        store,
+    )
+    if abort_msg:
+        raise InvalidRunStatusException(abort_msg)
+
+    # Delete the message object
+    store.delete(request.message_object_id)
+
+    return ConfirmMessageReceivedResponse()

flwr/server/superlink/fleet/rest_rere/rest_api.py

@@ -29,18 +29,28 @@ from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
     CreateNodeResponse,
     DeleteNodeRequest,
     DeleteNodeResponse,
-    PingRequest,
-    PingResponse,
     PullMessagesRequest,
     PullMessagesResponse,
     PushMessagesRequest,
     PushMessagesResponse,
 )
+from flwr.proto.heartbeat_pb2 import (  # pylint: disable=E0611
+    SendNodeHeartbeatRequest,
+    SendNodeHeartbeatResponse,
+)
+from flwr.proto.message_pb2 import (  # pylint: disable=E0611
+    ConfirmMessageReceivedRequest,
+    ConfirmMessageReceivedResponse,
+    PullObjectRequest,
+    PullObjectResponse,
+    PushObjectRequest,
+    PushObjectResponse,
+)
 from flwr.proto.run_pb2 import GetRunRequest, GetRunResponse  # pylint: disable=E0611
-from flwr.server.superlink.ffs.ffs import Ffs
-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 LinkState, LinkStateFactory
+from flwr.supercore.ffs import Ffs, FfsFactory
+from flwr.supercore.object_store import ObjectStore, ObjectStoreFactory
 
 try:
     from starlette.applications import Starlette
@@ -111,9 +121,10 @@ async def pull_message(request: PullMessagesRequest) -> PullMessagesResponse:
     """Pull PullMessages."""
     # Get state from app
     state: LinkState = cast(LinkStateFactory, app.state.STATE_FACTORY).state()
+    store: ObjectStore = cast(ObjectStoreFactory, app.state.OBJECTSTORE_FACTORY).store()
 
     # Handle message
-    return message_handler.pull_messages(request=request, state=state)
+    return message_handler.pull_messages(request=request, state=state, store=store)
 
 
 @rest_request_response(PushMessagesRequest)
@@ -121,19 +132,44 @@ async def push_message(request: PushMessagesRequest) -> PushMessagesResponse:
     """Pull PushMessages."""
     # Get state from app
     state: LinkState = cast(LinkStateFactory, app.state.STATE_FACTORY).state()
+    store: ObjectStore = cast(ObjectStoreFactory, app.state.OBJECTSTORE_FACTORY).store()
+
+    # Handle message
+    return message_handler.push_messages(request=request, state=state, store=store)
+
+
+@rest_request_response(PullObjectRequest)
+async def pull_object(request: PullObjectRequest) -> PullObjectResponse:
+    """Pull PullObject."""
+    # Get state from app
+    state: LinkState = cast(LinkStateFactory, app.state.STATE_FACTORY).state()
+    store: ObjectStore = cast(ObjectStoreFactory, app.state.OBJECTSTORE_FACTORY).store()
+
+    # Handle message
+    return message_handler.pull_object(request=request, state=state, store=store)
+
+
+@rest_request_response(PushObjectRequest)
+async def push_object(request: PushObjectRequest) -> PushObjectResponse:
+    """Pull PushObject."""
+    # Get state from app
+    state: LinkState = cast(LinkStateFactory, app.state.STATE_FACTORY).state()
+    store: ObjectStore = cast(ObjectStoreFactory, app.state.OBJECTSTORE_FACTORY).store()
 
     # Handle message
-    return message_handler.push_messages(request=request, state=state)
+    return message_handler.push_object(request=request, state=state, store=store)
 
 
-@rest_request_response(PingRequest)
-async def ping(request: PingRequest) -> PingResponse:
-    """Ping."""
+@rest_request_response(SendNodeHeartbeatRequest)
+async def send_node_heartbeat(
+    request: SendNodeHeartbeatRequest,
+) -> SendNodeHeartbeatResponse:
+    """Send node heartbeat."""
     # Get state from app
     state: LinkState = cast(LinkStateFactory, app.state.STATE_FACTORY).state()
 
     # Handle message
-    return message_handler.ping(request=request, state=state)
+    return message_handler.send_node_heartbeat(request=request, state=state)
 
 
 @rest_request_response(GetRunRequest)
@@ -141,9 +177,10 @@ async def get_run(request: GetRunRequest) -> GetRunResponse:
     """GetRun."""
     # Get state from app
     state: LinkState = cast(LinkStateFactory, app.state.STATE_FACTORY).state()
+    store: ObjectStore = cast(ObjectStoreFactory, app.state.OBJECTSTORE_FACTORY).store()
 
     # Handle message
-    return message_handler.get_run(request=request, state=state)
+    return message_handler.get_run(request=request, state=state, store=store)
 
 
 @rest_request_response(GetFabRequest)
@@ -154,9 +191,25 @@ async def get_fab(request: GetFabRequest) -> GetFabResponse:
 
     # Get state from app
     state: LinkState = cast(LinkStateFactory, app.state.STATE_FACTORY).state()
+    store: ObjectStore = cast(ObjectStoreFactory, app.state.OBJECTSTORE_FACTORY).store()
+
+    # Handle message
+    return message_handler.get_fab(request=request, ffs=ffs, state=state, store=store)
+
+
+@rest_request_response(ConfirmMessageReceivedRequest)
+async def confirm_message_received(
+    request: ConfirmMessageReceivedRequest,
+) -> ConfirmMessageReceivedResponse:
+    """Confirm message received."""
+    # Get state from app
+    state: LinkState = cast(LinkStateFactory, app.state.STATE_FACTORY).state()
+    store: ObjectStore = cast(ObjectStoreFactory, app.state.OBJECTSTORE_FACTORY).store()
 
     # Handle message
-    return message_handler.get_fab(request=request, ffs=ffs, state=state)
+    return message_handler.confirm_message_received(
+        request=request, state=state, store=store
+    )
 
 
 routes = [
@@ -164,9 +217,16 @@ routes = [
     Route("/api/v0/fleet/delete-node", delete_node, methods=["POST"]),
     Route("/api/v0/fleet/pull-messages", pull_message, methods=["POST"]),
     Route("/api/v0/fleet/push-messages", push_message, methods=["POST"]),
-    Route("/api/v0/fleet/ping", ping, methods=["POST"]),
+    Route("/api/v0/fleet/pull-object", pull_object, methods=["POST"]),
+    Route("/api/v0/fleet/push-object", push_object, methods=["POST"]),
+    Route("/api/v0/fleet/send-node-heartbeat", send_node_heartbeat, methods=["POST"]),
     Route("/api/v0/fleet/get-run", get_run, methods=["POST"]),
     Route("/api/v0/fleet/get-fab", get_fab, methods=["POST"]),
+    Route(
+        "/api/v0/fleet/confirm-message-received",
+        confirm_message_received,
+        methods=["POST"],
+    ),
 ]
 
 app: Starlette = Starlette(

flwr/server/superlink/fleet/vce/vce_api.py

@@ -25,19 +25,20 @@ from pathlib import Path
 from queue import Empty, Queue
 from time import sleep
 from typing import Callable, Optional
+from uuid import uuid4
 
+from flwr.app.error import Error
 from flwr.client.client_app import ClientApp, ClientAppException, LoadClientAppError
 from flwr.client.clientapp.utils import get_load_client_app_fn
 from flwr.client.run_info_store import DeprecatedRunInfoStore
 from flwr.common import Message
 from flwr.common.constant import (
+    HEARTBEAT_MAX_INTERVAL,
     NUM_PARTITIONS_KEY,
     PARTITION_ID_KEY,
-    PING_MAX_INTERVAL,
     ErrorCode,
 )
 from flwr.common.logger import log
-from flwr.common.message import Error
 from flwr.common.typing import Run
 from flwr.server.superlink.linkstate import LinkState, LinkStateFactory
 
@@ -53,7 +54,7 @@ def _register_nodes(
     nodes_mapping: NodeToPartitionMapping = {}
     state = state_factory.state()
     for i in range(num_nodes):
-        node_id = state.create_node(ping_interval=PING_MAX_INTERVAL)
+        node_id = state.create_node(heartbeat_interval=HEARTBEAT_MAX_INTERVAL)
         nodes_mapping[node_id] = i
     log(DEBUG, "Registered %i nodes", len(nodes_mapping))
     return nodes_mapping
@@ -134,6 +135,8 @@ def worker(
 
         finally:
             if out_mssg:
+                # Assign a message_id
+                out_mssg.metadata.__dict__["_message_id"] = str(uuid4())
                 # Store reply Messages in state
                 messageres_queue.put(out_mssg)