refactor(framework) Wrap ClientAppIoStub with RetryInvoker (#4637)

adap · Dec 16, 2024 · fae9df9 · fae9df9
1 parent 86825e5
commit fae9df9
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/py/flwr/client/clientapp/app.py b/src/py/flwr/client/clientapp/app.py
@@ -32,6 +32,7 @@
 from flwr.common.grpc import create_channel
 from flwr.common.logger import log
 from flwr.common.message import Error
+from flwr.common.retry_invoker import _make_simple_grpc_retry_invoker, _wrap_stub
 from flwr.common.serde import (
     context_from_proto,
     context_to_proto,
@@ -106,9 +107,9 @@ def run_clientapp(  # pylint: disable=R0914
 
     # Resolve directory where FABs are installed
     flwr_dir_ = get_flwr_dir(flwr_dir)
-
     try:
         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