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

flwr/supernode/runtime/run_clientapp.py

@@ -16,9 +16,6 @@
 
 
 import gc
-import os
-import threading
-import time
 from logging import DEBUG, ERROR, INFO
 from typing import Optional
 
@@ -32,34 +29,46 @@ from flwr.common import Context, Message
 from flwr.common.config import get_flwr_dir
 from flwr.common.constant import ErrorCode
 from flwr.common.grpc import create_channel, on_channel_state_change
+from flwr.common.inflatable import (
+    get_all_nested_objects,
+    get_object_tree,
+    no_object_id_recompute,
+)
+from flwr.common.inflatable_protobuf_utils import (
+    make_confirm_message_received_fn_protobuf,
+    make_pull_object_fn_protobuf,
+    make_push_object_fn_protobuf,
+)
+from flwr.common.inflatable_utils import pull_and_inflate_object_from_tree, push_objects
 from flwr.common.logger import log
+from flwr.common.message import remove_content_from_message
 from flwr.common.retry_invoker import _make_simple_grpc_retry_invoker, _wrap_stub
 from flwr.common.serde import (
     context_from_proto,
     context_to_proto,
     fab_from_proto,
-    message_from_proto,
     message_to_proto,
     run_from_proto,
 )
 from flwr.common.typing import Fab, Run
-
-# pylint: disable=E0611
-from flwr.proto.clientappio_pb2 import (
-    GetTokenRequest,
-    GetTokenResponse,
-    PullClientAppInputsRequest,
-    PullClientAppInputsResponse,
-    PushClientAppOutputsRequest,
-    PushClientAppOutputsResponse,
+from flwr.proto.appio_pb2 import (  # pylint: disable=E0611
+    PullAppInputsRequest,
+    PullAppInputsResponse,
+    PullAppMessagesRequest,
+    PullAppMessagesResponse,
+    PushAppMessagesRequest,
+    PushAppOutputsRequest,
+    PushAppOutputsResponse,
 )
 from flwr.proto.clientappio_pb2_grpc import ClientAppIoStub
+from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
+from flwr.supercore.app_utils import start_parent_process_monitor
+from flwr.supercore.utils import mask_string
 
 
 def run_clientapp(  # pylint: disable=R0913, R0914, R0917
     clientappio_api_address: str,
-    run_once: bool,
-    token: Optional[int] = None,
+    token: str,
     flwr_dir: Optional[str] = None,
     certificates: Optional[bytes] = None,
     parent_pid: Optional[int] = None,
@@ -82,124 +91,95 @@ def run_clientapp(  # pylint: disable=R0913, R0914, R0917
         stub = ClientAppIoStub(channel)
         _wrap_stub(stub, _make_simple_grpc_retry_invoker())
 
-        while True:
-            # If token is not set, loop until token is received from SuperNode
-            while token is None:
-                token = get_token(stub)
-                time.sleep(1)
-
-            # Pull Message, Context, Run and (optional) FAB from SuperNode
-            message, context, run, fab = pull_clientappinputs(stub=stub, token=token)
-
-            # Install FAB, if provided
-            if fab:
-                log(DEBUG, "[flwr-clientapp] Start FAB installation.")
-                install_from_fab(fab.content, flwr_dir=flwr_dir_, skip_prompt=True)
-
-            load_client_app_fn = get_load_client_app_fn(
-                default_app_ref="",
-                app_path=None,
-                multi_app=True,
-                flwr_dir=str(flwr_dir_),
-            )
+        # Pull Message, Context, Run and (optional) FAB from SuperNode
+        message, context, run, fab = pull_clientappinputs(stub=stub, token=token)
 
-            try:
-                # Load ClientApp
-                log(DEBUG, "[flwr-clientapp] Start `ClientApp` Loading.")
-                client_app: ClientApp = load_client_app_fn(
-                    run.fab_id, run.fab_version, fab.hash_str if fab else ""
-                )
+        # Install FAB, if provided
+        if fab:
+            log(DEBUG, "[flwr-clientapp] Start FAB installation.")
+            install_from_fab(fab.content, flwr_dir=flwr_dir_, skip_prompt=True)
 
-                # Execute ClientApp
-                reply_message = client_app(message=message, context=context)
+        load_client_app_fn = get_load_client_app_fn(
+            default_app_ref="",
+            app_path=None,
+            multi_app=True,
+            flwr_dir=str(flwr_dir_),
+        )
 
-            except Exception as ex:  # pylint: disable=broad-exception-caught
-                # Don't update/change NodeState
+        try:
+            # Load ClientApp
+            log(DEBUG, "[flwr-clientapp] Start `ClientApp` Loading.")
+            client_app: ClientApp = load_client_app_fn(
+                run.fab_id, run.fab_version, fab.hash_str if fab else ""
+            )
 
-                e_code = ErrorCode.CLIENT_APP_RAISED_EXCEPTION
-                # Ex fmt: "<class 'ZeroDivisionError'>:<'division by zero'>"
-                reason = str(type(ex)) + ":<'" + str(ex) + "'>"
-                exc_entity = "ClientApp"
-                if isinstance(ex, LoadClientAppError):
-                    reason = (
-                        "An exception was raised when attempting to load `ClientApp`"
-                    )
-                    e_code = ErrorCode.LOAD_CLIENT_APP_EXCEPTION
+            # Execute ClientApp
+            reply_message = client_app(message=message, context=context)
 
-                log(ERROR, "%s raised an exception", exc_entity, exc_info=ex)
+        except Exception as ex:  # pylint: disable=broad-exception-caught
+            # Don't update/change NodeState
 
-                # Create error message
-                reply_message = Message(
-                    Error(code=e_code, reason=reason), reply_to=message
-                )
+            e_code = ErrorCode.CLIENT_APP_RAISED_EXCEPTION
+            # Ex fmt: "<class 'ZeroDivisionError'>:<'division by zero'>"
+            reason = str(type(ex)) + ":<'" + str(ex) + "'>"
+            exc_entity = "ClientApp"
+            if isinstance(ex, LoadClientAppError):
+                reason = "An exception was raised when attempting to load `ClientApp`"
+                e_code = ErrorCode.LOAD_CLIENT_APP_EXCEPTION
 
-            # Push Message and Context to SuperNode
-            _ = push_clientappoutputs(
-                stub=stub, token=token, message=reply_message, context=context
-            )
+            log(ERROR, "%s raised an exception", exc_entity, exc_info=ex)
 
-            del client_app, message, context, run, fab, reply_message
-            gc.collect()
+            # Create error message
+            reply_message = Message(Error(code=e_code, reason=reason), reply_to=message)
 
-            # Reset token to `None` to prevent flwr-clientapp from trying to pull the
-            # same inputs again
-            token = None
+        # Push Message and Context to SuperNode
+        _ = push_clientappoutputs(
+            stub=stub, token=token, message=reply_message, context=context
+        )
 
-            # Stop the loop if `flwr-clientapp` is expected to process only a single
-            # message
-            if run_once:
-                break
+        del client_app, message, context, run, fab, reply_message
+        gc.collect()
 
-    except KeyboardInterrupt:
-        log(INFO, "Closing connection")
     except grpc.RpcError as e:
         log(ERROR, "GRPC error occurred: %s", str(e))
     finally:
         channel.close()
 
 
-def start_parent_process_monitor(
-    parent_pid: int,
-) -> None:
-    """Monitor the parent process and exit if it terminates."""
-
-    def monitor() -> None:
-        while True:
-            time.sleep(0.2)
-            if os.getppid() != parent_pid:
-                os.kill(os.getpid(), 9)
-
-    threading.Thread(target=monitor, daemon=True).start()
-
-
-def get_token(stub: grpc.Channel) -> Optional[int]:
-    """Get a token from SuperNode."""
-    log(DEBUG, "[flwr-clientapp] Request token")
-    try:
-        res: GetTokenResponse = stub.GetToken(GetTokenRequest())
-        log(DEBUG, "[GetToken] Received token: %s", res.token)
-        return res.token
-    except grpc.RpcError as e:
-        if e.code() == grpc.StatusCode.FAILED_PRECONDITION:  # pylint: disable=no-member
-            log(DEBUG, "[GetToken] No token available yet")
-        else:
-            log(ERROR, "[GetToken] gRPC error occurred: %s", str(e))
-        return None
-
-
 def pull_clientappinputs(
-    stub: grpc.Channel, token: int
+    stub: ClientAppIoStub, token: str
 ) -> tuple[Message, Context, Run, Optional[Fab]]:
     """Pull ClientAppInputs from SuperNode."""
-    log(INFO, "[flwr-clientapp] Pull `ClientAppInputs` for token %s", token)
+    masked_token = mask_string(token)
+    log(INFO, "[flwr-clientapp] Pull `ClientAppInputs` for token %s", masked_token)
     try:
-        res: PullClientAppInputsResponse = stub.PullClientAppInputs(
-            PullClientAppInputsRequest(token=token)
+        # Pull Context, Run and (optional) FAB
+        res: PullAppInputsResponse = stub.PullClientAppInputs(
+            PullAppInputsRequest(token=token)
         )
-        message = message_from_proto(res.message)
         context = context_from_proto(res.context)
         run = run_from_proto(res.run)
         fab = fab_from_proto(res.fab) if res.fab else None
+
+        # Pull and inflate the message
+        pull_msg_res: PullAppMessagesResponse = stub.PullMessage(
+            PullAppMessagesRequest(token=token)
+        )
+        run_id = context.run_id
+        node = Node(node_id=context.node_id)
+        object_tree = pull_msg_res.message_object_trees[0]
+        message = pull_and_inflate_object_from_tree(
+            object_tree,
+            make_pull_object_fn_protobuf(stub.PullObject, node, run_id),
+            make_confirm_message_received_fn_protobuf(
+                stub.ConfirmMessageReceived, node, run_id
+            ),
+            return_type=Message,
+        )
+
+        # Set the message ID
+        # The deflated message doesn't contain the message_id (its own object_id)
+        message.metadata.__dict__["_message_id"] = object_tree.object_id
         return message, context, run, fab
     except grpc.RpcError as e:
         log(ERROR, "[PullClientAppInputs] gRPC error occurred: %s", str(e))
@@ -207,18 +187,52 @@ def pull_clientappinputs(
 
 
 def push_clientappoutputs(
-    stub: grpc.Channel, token: int, message: Message, context: Context
-) -> PushClientAppOutputsResponse:
+    stub: ClientAppIoStub, token: str, message: Message, context: Context
+) -> PushAppOutputsResponse:
     """Push ClientAppOutputs to SuperNode."""
-    log(INFO, "[flwr-clientapp] Push `ClientAppOutputs` for token %s", token)
-    proto_message = message_to_proto(message)
+    masked_token = mask_string(token)
+    log(INFO, "[flwr-clientapp] Push `ClientAppOutputs` for token %s", masked_token)
+    # Set message ID
+    message.metadata.__dict__["_message_id"] = message.object_id
+    proto_message = message_to_proto(remove_content_from_message(message))
     proto_context = context_to_proto(context)
 
     try:
-        res: PushClientAppOutputsResponse = stub.PushClientAppOutputs(
-            PushClientAppOutputsRequest(
-                token=token, message=proto_message, context=proto_context
+
+        with no_object_id_recompute():
+            # Get object tree and all objects to push
+            object_tree = get_object_tree(message)
+
+            # Push Message
+            # This is temporary. The message should not contain its content
+            push_msg_res = stub.PushMessage(
+                PushAppMessagesRequest(
+                    token=token,
+                    messages_list=[proto_message],
+                    message_object_trees=[object_tree],
+                )
+            )
+            del proto_message
+
+            # Retrieve the object IDs to push
+            object_ids_to_push = set(push_msg_res.objects_to_push)
+
+            # Push all objects
+            all_objects = get_all_nested_objects(message)
+            del message
+            push_objects(
+                all_objects,
+                make_push_object_fn_protobuf(
+                    stub.PushObject,
+                    Node(node_id=context.node_id),
+                    run_id=context.run_id,
+                ),
+                object_ids_to_push=object_ids_to_push,
             )
+
+        # Push Context
+        res: PushAppOutputsResponse = stub.PushClientAppOutputs(
+            PushAppOutputsRequest(token=token, context=proto_context)
         )
         return res
     except grpc.RpcError as e:

flwr/supernode/servicer/clientappio/__init__.py

@@ -15,10 +15,8 @@
 """ClientAppIo API Servicer."""
 
 
-from .clientappio_servicer import ClientAppInputs, ClientAppIoServicer, ClientAppOutputs
+from .clientappio_servicer import ClientAppIoServicer
 
 __all__ = [
-    "ClientAppInputs",
     "ClientAppIoServicer",
-    "ClientAppOutputs",
 ]