From c0f2379f1e4ed6797d22411810a20bee916feebe Mon Sep 17 00:00:00 2001
From: Charles Beauville <charles@flower.ai>
Date: Wed, 29 May 2024 22:10:41 +0200
Subject: [PATCH] feat(framework:skip) Add delete node on shutdown (#3524)

---
 src/py/flwr/client/app.py                         | 7 +++++++
 src/py/flwr/client/grpc_rere_client/connection.py | 5 -----
 2 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/src/py/flwr/client/app.py b/src/py/flwr/client/app.py
index e3487d352240..c9f45d14539c 100644
--- a/src/py/flwr/client/app.py
+++ b/src/py/flwr/client/app.py
@@ -268,6 +268,7 @@ def _load_client_app() -> ClientApp:
     run_tracker = _RunTracker()
 
     def _on_sucess(retry_state: RetryState) -> None:
+        run_tracker.is_connected = True
         if retry_state.tries > 1:
             log(
                 INFO,
@@ -279,6 +280,7 @@ def _on_sucess(retry_state: RetryState) -> None:
                 run_tracker.create_node()
 
     def _on_backoff(retry_state: RetryState) -> None:
+        run_tracker.is_connected = False
         if retry_state.tries == 1:
             log(WARN, "Connection attempt failed, retrying...")
         else:
@@ -426,6 +428,10 @@ def _on_backoff(retry_state: RetryState) -> None:
                     sleep_duration = 0
                     break
 
+            # Unregister node
+            if delete_node is not None and run_tracker.is_connected:
+                delete_node()  # pylint: disable=not-callable
+
         if sleep_duration == 0:
             log(INFO, "Disconnect and shut down")
             del run_tracker
@@ -605,6 +611,7 @@ def _init_connection(transport: Optional[str], server_address: str) -> Tuple[
 class _RunTracker:
     create_node: Optional[Callable[[], None]] = None
     interrupt: bool = False
+    is_connected: bool = False
 
     def register_signal_handler(self) -> None:
         """Register handlers for exit signals."""
diff --git a/src/py/flwr/client/grpc_rere_client/connection.py b/src/py/flwr/client/grpc_rere_client/connection.py
index 3778fd4061f9..d1769958863d 100644
--- a/src/py/flwr/client/grpc_rere_client/connection.py
+++ b/src/py/flwr/client/grpc_rere_client/connection.py
@@ -188,11 +188,6 @@ def delete_node() -> None:
             log(ERROR, "Node instance missing")
             return
 
-        # Stop the ping-loop thread
-        ping_stop_event.set()
-        if ping_thread is not None:
-            ping_thread.join()
-
         # Call FleetAPI
         delete_node_request = DeleteNodeRequest(node=node)
         retry_invoker.invoke(stub.DeleteNode, request=delete_node_request)