[2.3] fix client death handling (#2484)

* fix client death handling

* remove unused imports

* improve HB logic

Author: yanchengnv

nvflare/apis/impl/controller.py

@@ -77,6 +77,12 @@ def _get_client_task(target, task: Task):
     return None
 
 
+class _DeadClientStatus:
+    def __init__(self):
+        self.report_time = time.time()
+        self.dead_time = None
+
+
 class Controller(Responder, ControllerSpec, ABC):
     def __init__(self, task_check_period=0.2):
         """Manage life cycles of tasks and their destinations.
@@ -92,7 +98,7 @@ def __init__(self, task_check_period=0.2):
         self._task_lock = Lock()
         self._task_monitor = threading.Thread(target=self._monitor_tasks, args=())
         self._task_check_period = task_check_period
-        self._dead_client_reports = {}  # clients that reported the job is dead on it: name => report time
+        self._dead_clients = {}  # clients that reported the job is dead on it: name => _DeadClientStatus
         self._dead_clients_lock = Lock()  # need lock since dead_clients can be modified from different threads
         # make sure _check_tasks, process_task_request, process_submission does not interfere with each other
         self._controller_lock = Lock()
@@ -115,6 +121,28 @@ def initialize_run(self, fl_ctx: FLContext):
         self.start_controller(fl_ctx)
         self._task_monitor.start()
 
+    def client_is_dead(self, client_name: str):
+        """This method is called when a client is deemed dead.
+
+        Args:
+            client_name: name of the client
+
+        Returns: None
+
+        """
+        pass
+
+    def client_is_revived(self, client_name: str):
+        """This method is called when a client is revived.
+
+        Args:
+            client_name: name of the client
+
+        Returns: None
+
+        """
+        pass
+
     def _try_again(self) -> Tuple[str, str, Shareable]:
         # TODO: how to tell client no shareable available now?
         return "", "", None
@@ -174,9 +202,6 @@ def _do_process_task_request(self, client: Client, fl_ctx: FLContext) -> Tuple[s
         if not isinstance(client, Client):
             raise TypeError("client must be an instance of Client, but got {}".format(type(client)))
 
-        with self._dead_clients_lock:
-            self._dead_client_reports.pop(client.name, None)
-
         if not isinstance(fl_ctx, FLContext):
             raise TypeError("fl_ctx must be an instance of FLContext, but got {}".format(type(fl_ctx)))
 
@@ -331,9 +356,10 @@ def handle_dead_job(self, client_name: str, fl_ctx: FLContext):
         """
         # record the report and to be used by the task monitor
         with self._dead_clients_lock:
-            self.log_info(fl_ctx, f"received dead job report from client {client_name}")
-            if not self._dead_client_reports.get(client_name):
-                self._dead_client_reports[client_name] = time.time()
+            self.log_warning(fl_ctx, f"received dead job report for client {client_name}")
+            if not self._dead_clients.get(client_name):
+                self.log_warning(fl_ctx, f"client {client_name} is placed on dead client watch list")
+                self._dead_clients[client_name] = _DeadClientStatus()
 
     def process_task_check(self, task_id: str, fl_ctx: FLContext):
         with self._task_lock:
@@ -369,13 +395,6 @@ def _do_process_submission(
         if not isinstance(client, Client):
             raise TypeError("client must be an instance of Client, but got {}".format(type(client)))
 
-        # reset the dead job report!
-        # note that due to potential race conditions, a client may fail to include the job id in its
-        # heartbeat (since the job hasn't started at the time of heartbeat report), but then includes
-        # the job ID later.
-        with self._dead_clients_lock:
-            self._dead_client_reports.pop(client.name, None)
-
         if not isinstance(fl_ctx, FLContext):
             raise TypeError("fl_ctx must be an instance of FLContext, but got {}".format(type(fl_ctx)))
         if not isinstance(result, Shareable):
@@ -651,10 +670,8 @@ def cancel_task(
         """Cancel the specified task.
 
         Change the task completion_status, which will inform task monitor to clean up this task
-
-        .. note::
-
-            We only mark the task as completed and leave it to the task monitor to clean up. This is to avoid potential deadlock of task_lock.
+        We only mark the task as completed and leave it to the task monitor to clean up.
+        This is to avoid potential deadlock of task_lock.
 
         Args:
             task (Task): the task to be cancelled
@@ -979,27 +996,64 @@ def wait_for_task(self, task: Task, abort_signal: Signal):
     def _check_dead_clients(self):
         if self._engine:
             clients = self._engine.get_clients()
+            dead_clients = []
             with self._dead_clients_lock:
                 for client in clients:
-                    if self._client_still_alive(client.name):
-                        return False
+                    if not self._client_still_alive(client.name):
+                        dead_clients.append(client.name)
 
-                # All the clients are dead, abort the job run.
+            if dead_clients and len(clients) == len(dead_clients):
                 with self._engine.new_context() as fl_ctx:
                     self.system_panic("All clients are dead", fl_ctx)
-            return True
+                    return True
         return False
 
     def _client_still_alive(self, client_name):
         now = time.time()
-        report_time = self._dead_client_reports.get(client_name, None)
-        grace_period = ConfigService.get_float_var(name=_CONFIG_VAR_DEAD_CLIENT_GRACE_PERIOD, default=30.0)
+        status = self._dead_clients.get(client_name, None)
+        grace_period = ConfigService.get_float_var(name=_CONFIG_VAR_DEAD_CLIENT_GRACE_PERIOD, default=60.0)
 
-        if not report_time:
+        if not status:
             # this client is still alive
             return True
-        elif now - report_time < grace_period:
+
+        assert isinstance(status, _DeadClientStatus)
+        if status.dead_time:
+            return False
+
+        if now - status.report_time < grace_period:
             # this report is still fresh - consider the client to be still alive
             return True
 
+        # consider client dead
+        status.dead_time = now
+        self.logger.error(f"Client {client_name} is deemed dead!")
+        self.client_is_dead(client_name)
         return False
+
+    def get_client_death_time(self, client_name: str):
+        """Get the time that the client was deemed dead
+
+        Args:
+            client_name: name of the client
+
+        Returns: time at which the client was deemed dead; or None if the client is not dead
+
+        """
+        status = self._dead_clients.get(client_name)
+        if status:
+            assert isinstance(status, _DeadClientStatus)
+            return status.dead_time
+        return None
+
+    def process_job_heartbeat(self, fl_ctx: FLContext):
+        peer_ctx = fl_ctx.get_peer_context()
+        assert isinstance(peer_ctx, FLContext)
+        client_name = peer_ctx.get_identity_name()
+        with self._dead_clients_lock:
+            if client_name in self._dead_clients:
+                self.log_info(fl_ctx, f"Client {client_name} is removed from watch list")
+                status = self._dead_clients.pop(client_name)
+                if status.dead_time:
+                    self.log_info(fl_ctx, f"Client {client_name} is revived")
+                    self.client_is_revived(client_name)

nvflare/apis/responder.py

@@ -76,6 +76,18 @@ def process_task_check(self, task_id: str, fl_ctx: FLContext):
         """
         pass
 
+    @abstractmethod
+    def process_job_heartbeat(self, fl_ctx: FLContext):
+        """Called by the Engine to handle heartbeat received from clients.
+
+        Args:
+            fl_ctx: the FLContext
+
+        Returns: None
+
+        """
+        pass
+
     @abstractmethod
     def handle_dead_job(self, client_name: str, fl_ctx: FLContext):
         """Called by the Engine to handle the case that the job on the client is dead.

nvflare/private/fed/client/client_runner.py

@@ -413,11 +413,13 @@ def _try_run(self):
 
                 time.sleep(task_fetch_interval)
 
-    def send_job_heartbeat(self, interval=30.0):
-        wait_times = int(interval / 2)
+    def send_job_heartbeat(self, interval=10.0):
         request = Shareable()
+        last_hb_time = 0.0
+        short_sleep = 1.0
+        fl_ctx = self.engine.new_context()
         while not self.asked_to_stop:
-            with self.engine.new_context() as fl_ctx:
+            if time.time() - last_hb_time >= interval:
                 self.engine.send_aux_request(
                     targets=[FQCN.ROOT_SERVER],
                     topic=ReservedTopic.JOB_HEART_BEAT,
@@ -426,11 +428,8 @@ def send_job_heartbeat(self, interval=30.0):
                     fl_ctx=fl_ctx,
                     optional=True,
                 )
-
-                for i in range(wait_times):
-                    time.sleep(2)
-                    if self.asked_to_stop:
-                        break
+                last_hb_time = time.time()
+            time.sleep(short_sleep)
 
     def fetch_and_run_one_task(self, fl_ctx) -> (float, bool):
         """Fetches and runs a task.

nvflare/private/fed/server/server_runner.py

@@ -328,6 +328,9 @@ def _try_to_get_task(self, client, fl_ctx, timeout=None, retry_interval=0.005):
                     self.log_info(fl_ctx, "no current workflow - asked client to try again later")
                     return "", "", None
 
+                if self.current_wf.responder:
+                    self.current_wf.responder.process_job_heartbeat(fl_ctx=fl_ctx)
+
                 task_name, task_id, task_data = self.current_wf.responder.process_task_request(client, fl_ctx)
 
                 if task_name and task_name != SpecialTaskName.TRY_AGAIN:
@@ -441,6 +444,9 @@ def process_submission(self, client: Client, task_name: str, task_id: str, resul
                     self.log_info(fl_ctx, "no current workflow - dropped submission.")
                     return
 
+                if self.current_wf.responder:
+                    self.current_wf.responder.process_job_heartbeat(fl_ctx)
+
                 wf_id = result.get_cookie(ReservedHeaderKey.WORKFLOW, None)
                 if wf_id is not None and wf_id != self.current_wf.id:
                     self.log_info(
@@ -500,6 +506,9 @@ def process_submission(self, client: Client, task_name: str, task_id: str, resul
 
     def _handle_job_heartbeat(self, topic: str, request: Shareable, fl_ctx: FLContext) -> Shareable:
         self.log_debug(fl_ctx, "received client job_heartbeat")
+        with self.wf_lock:
+            if self.current_wf and self.current_wf.responder:
+                self.current_wf.responder.process_job_heartbeat(fl_ctx=fl_ctx)
         return make_reply(ReturnCode.OK)
 
     def _handle_task_check(self, topic: str, request: Shareable, fl_ctx: FLContext) -> Shareable:
@@ -515,6 +524,9 @@ def _handle_task_check(self, topic: str, request: Shareable, fl_ctx: FLContext)
                 self.log_info(fl_ctx, "no current workflow - dropped task_check.")
                 return make_reply(ReturnCode.TASK_UNKNOWN)
 
+            if self.current_wf.responder:
+                self.current_wf.responder.process_job_heartbeat(fl_ctx)
+
             # filter task result
             task = self.current_wf.responder.process_task_check(task_id=task_id, fl_ctx=fl_ctx)
             if task: