[2.4] Integrated ReliableMessage with XGBoost (#2399)

* Integrated ReliableMessage with XGBoost

* Addressed comments in PR

* Removed RequestSender and moved Sender to nvflare.apis.util

* Removed _check_reply

* Fixed a test error caused by not having a mock engine

* Removed the extra self.engine initialization

* Removed duplicate assignment of self.engine

Author: nvidianz

nvflare/apis/utils/reliable_message.py

@@ -304,7 +304,7 @@ def _send_request(
         # keep sending the request until a positive ack or result is received
         num_tries = 0
         while True:
-            if abort_signal.triggered:
+            if abort_signal and abort_signal.triggered:
                 return make_reply(ReturnCode.TASK_ABORTED)
 
             ack = engine.send_aux_request(
@@ -338,7 +338,7 @@ def _send_request(
                 return _error_reply(ReturnCode.COMMUNICATION_ERROR, f"Max send retries ({cls._max_retries}) reached")
             start = time.time()
             while time.time() - start < cls._query_interval:
-                if abort_signal.triggered:
+                if abort_signal and abort_signal.triggered:
                     return make_reply(ReturnCode.TASK_ABORTED)
                 time.sleep(0.1)
 
@@ -355,6 +355,7 @@ def _query_result(
     ) -> Shareable:
 
         # Querying phase - try to get result
+        engine = fl_ctx.get_engine()
         query = Shareable()
         query.set_header(HEADER_TX, receiver.tx_id)
         query.set_header(HEADER_OP, OP_QUERY)
@@ -367,7 +368,7 @@ def _query_result(
                 # check other condition and/or send query to ask for result.
                 return receiver.result
 
-            if abort_signal.triggered:
+            if abort_signal and abort_signal.triggered:
                 return make_reply(ReturnCode.TASK_ABORTED)
 
             # send a query. The ack of the query could be the result itself, or a status report.

nvflare/apis/utils/reliable_sender.py

@@ -0,0 +1,54 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from nvflare.apis.fl_context import FLContext
+from nvflare.apis.shareable import Shareable
+from nvflare.apis.signal import Signal
+from nvflare.apis.utils.reliable_message import ReliableMessage
+from nvflare.apis.utils.sender import Sender
+from nvflare.fuel.f3.cellnet.fqcn import FQCN
+
+
+class ReliableSender(Sender):
+    def __init__(self, max_request_workers=20, query_interval=5, max_retries=5, max_tx_time=300.0):
+        """Constructor
+
+        Args:
+            max_request_workers: Number of concurrent request worker threads
+            query_interval: Retry/query interval
+            max_retries: Number of retries
+            max_tx_time: Max transmitting time
+        """
+
+        super().__init__()
+        self.max_request_workers = max_request_workers
+        self.query_interval = query_interval
+        self.max_retries = max_retries
+        self.max_tx_time = max_tx_time
+        self.enabled = False
+
+    def send_request(
+        self, target: str, topic: str, req: Shareable, timeout: float, fl_ctx: FLContext, abort_signal: Signal
+    ) -> Shareable:
+
+        if not self.enabled:
+            ReliableMessage.enable(
+                fl_ctx,
+                max_request_workers=self.max_request_workers,
+                query_interval=self.query_interval,
+                max_retries=self.max_retries,
+                max_tx_time=self.max_tx_time,
+            )
+            self.enabled = True
+
+        return ReliableMessage.send_request(FQCN.ROOT_SERVER, topic, req, timeout, abort_signal, fl_ctx)

nvflare/apis/utils/sender.py

@@ -0,0 +1,86 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from abc import ABC, abstractmethod
+from typing import Optional
+
+from nvflare.apis.fl_component import FLComponent
+from nvflare.apis.fl_context import FLContext
+from nvflare.apis.shareable import Shareable
+from nvflare.apis.signal import Signal
+from nvflare.fuel.f3.cellnet.fqcn import FQCN
+
+
+class Sender(FLComponent, ABC):
+    """An abstract class to send request"""
+
+    @abstractmethod
+    def send_request(
+        self, target: str, topic: str, req: Shareable, timeout: float, fl_ctx: FLContext, abort_signal: Signal
+    ) -> Optional[Shareable]:
+        """Send a request to target. This is an abstract method. Derived class must implement this method
+
+         Args:
+             target: The destination
+             topic: Topic for the request
+             req: the request Shareable
+             timeout: Timeout of the request in seconds
+             fl_ctx: FLContext for the transaction
+             abort_signal: used for checking whether the job is aborted.
+
+        Returns:
+            The reply in Shareable
+
+        """
+        pass
+
+    def send_to_server(
+        self, topic: str, req: Shareable, timeout: float, fl_ctx: FLContext, abort_signal: Signal
+    ) -> Optional[Shareable]:
+        """Send an XGB request to the server.
+
+        Args:
+            topic: The topic of the request
+            req: the request Shareable
+            timeout: The timeout value for the request
+            fl_ctx: The FLContext for the request
+            abort_signal: used for checking whether the job is aborted.
+
+        Returns: reply from the server
+        """
+
+        return self.send_request(FQCN.ROOT_SERVER, topic, req, timeout, fl_ctx, abort_signal)
+
+
+class SimpleSender(Sender):
+    def __init__(self):
+        super().__init__()
+
+    def send_request(
+        self, target: str, topic: str, req: Shareable, timeout: float, fl_ctx: FLContext, abort_signal: Signal
+    ) -> Optional[Shareable]:
+
+        engine = fl_ctx.get_engine()
+        reply = engine.send_aux_request(
+            targets=[target],
+            topic=topic,
+            request=req,
+            timeout=timeout,
+            fl_ctx=fl_ctx,
+        )
+
+        # send_aux_request returns multiple replies in a dict
+        if reply:
+            return reply.get(target)
+        else:
+            return None

nvflare/app_opt/xgboost/histogram_based_v2/adaptor.py

@@ -18,12 +18,13 @@
 from typing import Tuple
 
 from nvflare.apis.fl_component import FLComponent
+from nvflare.apis.fl_constant import ReturnCode
 from nvflare.apis.fl_context import FLContext
 from nvflare.apis.shareable import Shareable
 from nvflare.apis.signal import Signal
+from nvflare.apis.utils.sender import Sender
 from nvflare.app_opt.xgboost.histogram_based_v2.defs import Constant
 from nvflare.app_opt.xgboost.histogram_based_v2.runner import XGBRunner
-from nvflare.app_opt.xgboost.histogram_based_v2.sender import Sender
 from nvflare.fuel.utils.validation_utils import check_non_negative_int, check_object_type, check_positive_int
 
 
@@ -279,7 +280,7 @@ class XGBClientAdaptor(XGBAdaptor, ABC):
     XGBClientAdaptor specifies commonly required methods for client adaptor implementations.
     """
 
-    def __init__(self):
+    def __init__(self, req_timeout: float):
         """Constructor of XGBClientAdaptor"""
         XGBAdaptor.__init__(self)
         self.engine = None
@@ -288,6 +289,7 @@ def __init__(self):
         self.rank = None
         self.num_rounds = None
         self.world_size = None
+        self.req_timeout = req_timeout
 
     def set_sender(self, sender: Sender):
         """Set the sender to be used to send XGB operation requests to the server.
@@ -314,6 +316,8 @@ def configure(self, config: dict, fl_ctx: FLContext):
         Returns:
             None
         """
+        self.engine = fl_ctx.get_engine()
+
         ws = config.get(Constant.CONF_KEY_WORLD_SIZE)
         if not ws:
             raise RuntimeError("world_size is not configured")
@@ -345,8 +349,22 @@ def _send_request(self, op: str, req: Shareable) -> bytes:
         Returns:
             operation result
         """
-        reply = self.sender.send_to_server(op, req, self.abort_signal)
+        req.set_header(Constant.MSG_KEY_XGB_OP, op)
+
+        with self.engine.new_context() as fl_ctx:
+            reply = self.sender.send_to_server(
+                Constant.TOPIC_XGB_REQUEST, req, self.req_timeout, fl_ctx, self.abort_signal
+            )
+
         if isinstance(reply, Shareable):
+            rc = reply.get_return_code()
+            if rc != ReturnCode.OK:
+                raise RuntimeError(f"received error return code: {rc}")
+
+            reply_op = reply.get_header(Constant.MSG_KEY_XGB_OP)
+            if reply_op != op:
+                raise RuntimeError(f"received op {reply_op} != expected op {op}")
+
             rcv_buf = reply.get(Constant.PARAM_KEY_RCV_BUF)
             if not isinstance(rcv_buf, bytes):
                 raise RuntimeError(f"invalid rcv_buf for {op=}: expect bytes but got {type(rcv_buf)}")

nvflare/app_opt/xgboost/histogram_based_v2/adaptor_controller.py

@@ -22,6 +22,7 @@
 from nvflare.apis.impl.controller import ClientTask, Controller, Task
 from nvflare.apis.shareable import ReturnCode, Shareable, make_reply
 from nvflare.apis.signal import Signal
+from nvflare.apis.utils.reliable_message import ReliableMessage
 from nvflare.app_opt.xgboost.histogram_based_v2.adaptor import XGBServerAdaptor
 from nvflare.app_opt.xgboost.histogram_based_v2.defs import Constant
 from nvflare.fuel.utils.validation_utils import check_number_range, check_object_type, check_positive_number, check_str
@@ -218,6 +219,16 @@ def start_controller(self, fl_ctx: FLContext):
             message_handle_func=self._process_client_done,
         )
 
+        ReliableMessage.enable(fl_ctx)
+        ReliableMessage.register_request_handler(
+            topic=Constant.TOPIC_XGB_REQUEST,
+            handler_f=self._process_xgb_request,
+        )
+        ReliableMessage.register_request_handler(
+            topic=Constant.TOPIC_CLIENT_DONE,
+            handler_f=self._process_client_done,
+        )
+
     def _trigger_stop(self, fl_ctx: FLContext, error=None):
         # first trigger the abort_signal to tell all components (mainly the controller's control_flow and adaptor)
         # that check this signal to abort.

nvflare/app_opt/xgboost/histogram_based_v2/adaptor_executor.py

@@ -11,24 +11,25 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
 from nvflare.apis.event_type import EventType
 from nvflare.apis.executor import Executor
 from nvflare.apis.fl_constant import ReturnCode
 from nvflare.apis.fl_context import FLContext
 from nvflare.apis.shareable import Shareable, make_reply
 from nvflare.apis.signal import Signal
+from nvflare.apis.utils.sender import Sender, SimpleSender
 from nvflare.app_opt.xgboost.histogram_based_v2.adaptor import XGBClientAdaptor
 from nvflare.app_opt.xgboost.histogram_based_v2.defs import Constant
-from nvflare.app_opt.xgboost.histogram_based_v2.sender import Sender
 from nvflare.fuel.f3.cellnet.fqcn import FQCN
+from nvflare.fuel.utils.validation_utils import check_str
 from nvflare.security.logging import secure_format_exception
 
 
 class XGBExecutor(Executor):
     def __init__(
         self,
         adaptor_component_id: str,
+        sender_id: str = None,
         configure_task_name=Constant.CONFIG_TASK_NAME,
         start_task_name=Constant.START_TASK_NAME,
         req_timeout=100.0,
@@ -37,11 +38,17 @@ def __init__(
 
         Args:
             adaptor_component_id: the component ID of client target adaptor
+            sender_id: The sender component id
             configure_task_name: name of the config task
             start_task_name: name of the start task
         """
         Executor.__init__(self)
         self.adaptor_component_id = adaptor_component_id
+
+        if sender_id:
+            check_str("sender_id", sender_id)
+        self.sender_id = sender_id
+
         self.req_timeout = req_timeout
         self.configure_task_name = configure_task_name
         self.start_task_name = start_task_name
@@ -78,9 +85,12 @@ def handle_event(self, event_type: str, fl_ctx: FLContext):
                 )
                 return
 
+            sender = self._get_sender(fl_ctx)
+            if not sender:
+                return
+
             adaptor.set_abort_signal(self.abort_signal)
-            engine = fl_ctx.get_engine()
-            adaptor.set_sender(Sender(engine, self.req_timeout))
+            adaptor.set_sender(sender)
             adaptor.initialize(fl_ctx)
             self.adaptor = adaptor
         elif event_type == EventType.END_RUN:
@@ -168,3 +178,31 @@ def _notify_client_done(self, rc, fl_ctx: FLContext):
             fl_ctx=fl_ctx,
             optional=True,
         )
+
+    def _get_sender(self, fl_ctx: FLContext) -> Sender:
+        """Get request sender to be used by this executor.
+
+        Args:
+            fl_ctx: the FL context
+
+        Returns:
+            A sender object
+        """
+
+        if self.sender_id:
+            engine = fl_ctx.get_engine()
+            sender = engine.get_component(self.sender_id)
+            if not sender:
+                self.system_panic(f"cannot get component for {self.sender_id}", fl_ctx)
+            else:
+                if not isinstance(sender, Sender):
+                    self.system_panic(
+                        f"invalid component '{self.sender_id}': expect {Sender.__name__} but got {type(sender)}",
+                        fl_ctx,
+                    )
+                    sender = None
+
+        else:
+            sender = SimpleSender()
+
+        return sender

nvflare/app_opt/xgboost/histogram_based_v2/adaptors/grpc_client_adaptor.py

@@ -83,15 +83,17 @@ def __init__(
         self,
         int_server_grpc_options=None,
         in_process=False,
+        req_timeout=100,
     ):
         """Constructor method to initialize the object.
 
         Args:
             int_server_grpc_options: An optional list of key-value pairs (`channel_arguments`
                 in gRPC Core runtime) to configure the gRPC channel of internal `GrpcServer`.
             in_process (bool): Specifies whether to start the `XGBRunner` in the same process or not.
+            req_timeout: Request timeout
         """
-        XGBClientAdaptor.__init__(self)
+        XGBClientAdaptor.__init__(self, req_timeout)
         self.int_server_grpc_options = int_server_grpc_options
         self.in_process = in_process
         self.internal_xgb_server = None