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

flwr/supernode/servicer/clientappio/clientappio_servicer.py

@@ -15,16 +15,15 @@
 """ClientAppIo API servicer."""
 
 
-from dataclasses import dataclass
 from logging import DEBUG, ERROR
-from typing import Optional, cast
+from typing import cast
 
 import grpc
 
-from flwr.common import Context, Message, typing
+from flwr.common import Context
+from flwr.common.inflatable import UnexpectedObjectContentError
 from flwr.common.logger import log
 from flwr.common.serde import (
-    clientappstatus_to_proto,
     context_from_proto,
     context_to_proto,
     fab_to_proto,
@@ -36,209 +35,269 @@ from flwr.common.typing import Fab, Run
 
 # pylint: disable=E0611
 from flwr.proto import clientappio_pb2_grpc
+from flwr.proto.appio_pb2 import (  # pylint: disable=E0401
+    PullAppInputsRequest,
+    PullAppInputsResponse,
+    PullAppMessagesRequest,
+    PullAppMessagesResponse,
+    PushAppMessagesRequest,
+    PushAppMessagesResponse,
+    PushAppOutputsRequest,
+    PushAppOutputsResponse,
+)
 from flwr.proto.clientappio_pb2 import (  # pylint: disable=E0401
-    GetTokenRequest,
-    GetTokenResponse,
-    PullClientAppInputsRequest,
-    PullClientAppInputsResponse,
-    PushClientAppOutputsRequest,
-    PushClientAppOutputsResponse,
+    GetRunIdsWithPendingMessagesRequest,
+    GetRunIdsWithPendingMessagesResponse,
+    RequestTokenRequest,
+    RequestTokenResponse,
+)
+from flwr.proto.message_pb2 import (
+    ConfirmMessageReceivedRequest,
+    ConfirmMessageReceivedResponse,
+    PullObjectRequest,
+    PullObjectResponse,
+    PushObjectRequest,
+    PushObjectResponse,
 )
+from flwr.proto.run_pb2 import GetRunRequest, GetRunResponse  # pylint: disable=E0611
 
+# pylint: disable=E0601
+from flwr.supercore.ffs import FfsFactory
+from flwr.supercore.object_store import NoObjectInStoreError, ObjectStoreFactory
+from flwr.supercore.object_store.utils import store_mapping_and_register_objects
+from flwr.supernode.nodestate import NodeStateFactory
 
-@dataclass
-class ClientAppInputs:
-    """Specify the inputs to the ClientApp."""
 
-    message: Message
-    context: Context
-    run: Run
-    fab: Optional[Fab]
-    token: int
+# pylint: disable=C0103,W0613,W0201
+class ClientAppIoServicer(clientappio_pb2_grpc.ClientAppIoServicer):
+    """ClientAppIo API servicer."""
 
+    def __init__(
+        self,
+        state_factory: NodeStateFactory,
+        ffs_factory: FfsFactory,
+        objectstore_factory: ObjectStoreFactory,
+    ) -> None:
+        self.state_factory = state_factory
+        self.ffs_factory = ffs_factory
+        self.objectstore_factory = objectstore_factory
 
-@dataclass
-class ClientAppOutputs:
-    """Specify the outputs from the ClientApp."""
+    def GetRunIdsWithPendingMessages(
+        self,
+        request: GetRunIdsWithPendingMessagesRequest,
+        context: grpc.ServicerContext,
+    ) -> GetRunIdsWithPendingMessagesResponse:
+        """Get run IDs with pending messages."""
+        log(DEBUG, "ClientAppIo.GetRunIdsWithPendingMessages")
 
-    message: Message
-    context: Context
+        # Initialize state connection
+        state = self.state_factory.state()
 
+        # Get run IDs with pending messages
+        run_ids = state.get_run_ids_with_pending_messages()
 
-# pylint: disable=C0103,W0613,W0201
-class ClientAppIoServicer(clientappio_pb2_grpc.ClientAppIoServicer):
-    """ClientAppIo API servicer."""
+        # Return run IDs
+        return GetRunIdsWithPendingMessagesResponse(run_ids=run_ids)
 
-    def __init__(self) -> None:
-        self.clientapp_input: Optional[ClientAppInputs] = None
-        self.clientapp_output: Optional[ClientAppOutputs] = None
-        self.token_returned: bool = False
-        self.inputs_returned: bool = False
+    def RequestToken(
+        self, request: RequestTokenRequest, context: grpc.ServicerContext
+    ) -> RequestTokenResponse:
+        """Request token."""
+        log(DEBUG, "ClientAppIo.RequestToken")
 
-    def GetToken(
-        self, request: GetTokenRequest, context: grpc.ServicerContext
-    ) -> GetTokenResponse:
-        """Get token."""
-        log(DEBUG, "ClientAppIo.GetToken")
+        # Initialize state connection
+        state = self.state_factory.state()
 
-        # Fail if no ClientAppInputs are available
-        if self.clientapp_input is None:
-            context.abort(
-                grpc.StatusCode.FAILED_PRECONDITION,
-                "No inputs available.",
-            )
-        clientapp_input = cast(ClientAppInputs, self.clientapp_input)
+        # Attempt to create a token for the provided run ID
+        try:
+            token = state.create_token(request.run_id)
+        except ValueError:
+            # Return an empty token if A token already exists for this run ID,
+            # indicating the run is in progress
+            return RequestTokenResponse(token="")
 
-        # Fail if token was already returned in a previous call
-        if self.token_returned:
-            context.abort(
-                grpc.StatusCode.FAILED_PRECONDITION,
-                "Token already returned. A token can be returned only once.",
-            )
+        # Return the token
+        return RequestTokenResponse(token=token)
+
+    def GetRun(
+        self, request: GetRunRequest, context: grpc.ServicerContext
+    ) -> GetRunResponse:
+        """Get run information."""
+        log(DEBUG, "ClientAppIo.GetRun")
 
-        # If
-        # - ClientAppInputs is set, and
-        # - token hasn't been returned before,
-        # return token
-        self.token_returned = True
-        return GetTokenResponse(token=clientapp_input.token)
+        # Initialize state connection
+        state = self.state_factory.state()
+
+        # Retrieve run information
+        run = state.get_run(request.run_id)
+
+        if run is None:
+            return GetRunResponse()
+
+        return GetRunResponse(run=run_to_proto(run))
 
     def PullClientAppInputs(
-        self, request: PullClientAppInputsRequest, context: grpc.ServicerContext
-    ) -> PullClientAppInputsResponse:
+        self, request: PullAppInputsRequest, context: grpc.ServicerContext
+    ) -> PullAppInputsResponse:
         """Pull Message, Context, and Run."""
         log(DEBUG, "ClientAppIo.PullClientAppInputs")
 
-        # Fail if no ClientAppInputs are available
-        if self.clientapp_input is None:
-            context.abort(
-                grpc.StatusCode.FAILED_PRECONDITION,
-                "No inputs available.",
-            )
-        clientapp_input = cast(ClientAppInputs, self.clientapp_input)
+        # Initialize state and ffs connection
+        state = self.state_factory.state()
+        ffs = self.ffs_factory.ffs()
 
-        # Fail if token wasn't returned in a previous call
-        if not self.token_returned:
+        # Validate the token
+        run_id = state.get_run_id_by_token(request.token)
+        if run_id is None or not state.verify_token(run_id, request.token):
             context.abort(
-                grpc.StatusCode.FAILED_PRECONDITION,
-                "Token hasn't been returned."
-                "Token must be returned before can be returned only once.",
+                grpc.StatusCode.PERMISSION_DENIED,
+                "Invalid token.",
             )
+            raise RuntimeError("This line should never be reached.")
 
-        # Fail if token isn't matching
-        if request.token != clientapp_input.token:
-            context.abort(
-                grpc.StatusCode.INVALID_ARGUMENT,
-                "Mismatch between ClientApp and SuperNode token",
-            )
+        # Retrieve context, run and fab for this run
+        context = cast(Context, state.get_context(run_id))
+        run = cast(Run, state.get_run(run_id))
+        fab = Fab(run.fab_hash, ffs.get(run.fab_hash)[0])  # type: ignore
 
-        # Success
-        self.inputs_returned = True
-        return PullClientAppInputsResponse(
-            message=message_to_proto(clientapp_input.message),
-            context=context_to_proto(clientapp_input.context),
-            run=run_to_proto(clientapp_input.run),
-            fab=fab_to_proto(clientapp_input.fab) if clientapp_input.fab else None,
+        return PullAppInputsResponse(
+            context=context_to_proto(context),
+            run=run_to_proto(run),
+            fab=fab_to_proto(fab),
         )
 
     def PushClientAppOutputs(
-        self, request: PushClientAppOutputsRequest, context: grpc.ServicerContext
-    ) -> PushClientAppOutputsResponse:
+        self, request: PushAppOutputsRequest, context: grpc.ServicerContext
+    ) -> PushAppOutputsResponse:
         """Push Message and Context."""
         log(DEBUG, "ClientAppIo.PushClientAppOutputs")
 
-        # Fail if no ClientAppInputs are available
-        if not self.clientapp_input:
-            context.abort(
-                grpc.StatusCode.FAILED_PRECONDITION,
-                "No inputs available.",
-            )
-        clientapp_input = cast(ClientAppInputs, self.clientapp_input)
+        # Initialize state connection
+        state = self.state_factory.state()
 
-        # Fail if token wasn't returned in a previous call
-        if not self.token_returned:
+        # Validate the token
+        run_id = state.get_run_id_by_token(request.token)
+        if run_id is None or not state.verify_token(run_id, request.token):
             context.abort(
-                grpc.StatusCode.FAILED_PRECONDITION,
-                "Token hasn't been returned."
-                "Token must be returned before can be returned only once.",
+                grpc.StatusCode.PERMISSION_DENIED,
+                "Invalid token.",
             )
+            raise RuntimeError("This line should never be reached.")
+
+        # Save the context to the state
+        state.store_context(context_from_proto(request.context))
 
-        # Fail if inputs weren't delivered in a previous call
-        if not self.inputs_returned:
+        # Remove the token to make the run eligible for processing
+        # A run associated with a token cannot be handled until its token is cleared
+        state.delete_token(run_id)
+
+        return PushAppOutputsResponse()
+
+    def PullMessage(
+        self, request: PullAppMessagesRequest, context: grpc.ServicerContext
+    ) -> PullAppMessagesResponse:
+        """Pull one Message."""
+        # Initialize state and store connection
+        state = self.state_factory.state()
+        store = self.objectstore_factory.store()
+
+        # Validate the token
+        run_id = state.get_run_id_by_token(request.token)
+        if run_id is None or not state.verify_token(run_id, request.token):
             context.abort(
-                grpc.StatusCode.FAILED_PRECONDITION,
-                "Inputs haven't been delivered."
-                "Inputs must be delivered before can be returned only once.",
+                grpc.StatusCode.PERMISSION_DENIED,
+                "Invalid token.",
             )
+            raise RuntimeError("This line should never be reached.")
+
+        # Retrieve message for this run
+        message = state.get_messages(run_ids=[run_id], is_reply=False)[0]
+
+        # Retrieve the object tree for the message
+        object_tree = store.get_object_tree(message.metadata.message_id)
 
-        # Fail if token isn't matching
-        if request.token != clientapp_input.token:
+        return PullAppMessagesResponse(
+            messages_list=[message_to_proto(message)],
+            message_object_trees=[object_tree],
+        )
+
+    def PushMessage(
+        self, request: PushAppMessagesRequest, context: grpc.ServicerContext
+    ) -> PushAppMessagesResponse:
+        """Push one Message."""
+        # Initialize state and store connection
+        state = self.state_factory.state()
+        store = self.objectstore_factory.store()
+
+        # Validate the token
+        run_id = state.get_run_id_by_token(request.token)
+        if run_id is None or not state.verify_token(run_id, request.token):
             context.abort(
-                grpc.StatusCode.INVALID_ARGUMENT,
-                "Mismatch between ClientApp and SuperNode token",
+                grpc.StatusCode.PERMISSION_DENIED,
+                "Invalid token.",
             )
+            raise RuntimeError("This line should never be reached.")
+
+        # Save the message to the state
+        state.store_message(message_from_proto(request.messages_list[0]))
+
+        # Store Message object to descendants mapping and preregister objects
+        objects_to_push = store_mapping_and_register_objects(store, request=request)
+
+        return PushAppMessagesResponse(objects_to_push=objects_to_push)
 
-        # Preconditions met
+    def PushObject(
+        self, request: PushObjectRequest, context: grpc.ServicerContext
+    ) -> PushObjectResponse:
+        """Push an object to the ObjectStore."""
+        log(DEBUG, "ServerAppIoServicer.PushObject")
+
+        # Init state and store
+        store = self.objectstore_factory.store()
+
+        # Insert in store
+        stored = False
         try:
-            # Update Message and Context
-            self.clientapp_output = ClientAppOutputs(
-                message=message_from_proto(request.message),
-                context=context_from_proto(request.context),
+            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
+            context.abort(grpc.StatusCode.FAILED_PRECONDITION, str(e))
+
+        return PushObjectResponse(stored=stored)
+
+    def PullObject(
+        self, request: PullObjectRequest, context: grpc.ServicerContext
+    ) -> PullObjectResponse:
+        """Pull an object from the ObjectStore."""
+        log(DEBUG, "ServerAppIoServicer.PullObject")
+
+        # Init state and 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,
             )
+        return PullObjectResponse(object_found=False, object_available=False)
 
-            # Set status
-            code = typing.ClientAppOutputCode.SUCCESS
-            status = typing.ClientAppOutputStatus(code=code, message="Success")
-        except Exception as e:  # pylint: disable=broad-exception-caught
-            log(ERROR, "ClientApp failed to push message to SuperNode, %s", e)
-            code = typing.ClientAppOutputCode.UNKNOWN_ERROR
-            status = typing.ClientAppOutputStatus(code=code, message="Unkonwn error")
+    def ConfirmMessageReceived(
+        self, request: ConfirmMessageReceivedRequest, context: grpc.ServicerContext
+    ) -> ConfirmMessageReceivedResponse:
+        """Confirm message received."""
+        log(DEBUG, "ServerAppIoServicer.ConfirmMessageReceived")
 
-        # Return status to ClientApp process
-        proto_status = clientappstatus_to_proto(status=status)
-        return PushClientAppOutputsResponse(status=proto_status)
+        # Init state and store
+        store = self.objectstore_factory.store()
 
-    def set_inputs(
-        self, clientapp_input: ClientAppInputs, token_returned: bool
-    ) -> None:
-        """Set ClientApp inputs.
-
-        Parameters
-        ----------
-        clientapp_input : ClientAppInputs
-            The inputs to the ClientApp.
-        token_returned : bool
-            A boolean indicating if the token has been returned.
-            Set to `True` when passing the token to `flwr-clientap`
-            and `False` otherwise.
-        """
-        if (
-            self.clientapp_input is not None
-            or self.clientapp_output is not None
-            or self.token_returned
-        ):
-            raise ValueError(
-                "ClientAppInputs and ClientAppOutputs must not be set before "
-                "calling `set_inputs`."
-            )
-        log(DEBUG, "ClientAppInputs set (token: %s)", clientapp_input.token)
-        self.clientapp_input = clientapp_input
-        self.token_returned = token_returned
-
-    def has_outputs(self) -> bool:
-        """Check if ClientAppOutputs are available."""
-        return self.clientapp_output is not None
-
-    def get_outputs(self) -> ClientAppOutputs:
-        """Get ClientApp outputs."""
-        if self.clientapp_output is None:
-            raise ValueError("ClientAppOutputs not set before calling `get_outputs`.")
-
-        # Set outputs to a local variable and clear state
-        output: ClientAppOutputs = self.clientapp_output
-        self.clientapp_input = None
-        self.clientapp_output = None
-        self.token_returned = False
-        self.inputs_returned = False
-
-        return output
+        # Delete the message object
+        store.delete(request.message_object_id)
+
+        return ConfirmMessageReceivedResponse()