Address PR comments

NVIDIA · Jan 6, 2025 · 42fe31e · 42fe31e
1 parent 7d9f9ba
commit 42fe31e
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/nvflare/app_opt/confidential_computing/aci_authorizer.py b/nvflare/app_opt/confidential_computing/aci_authorizer.py
@@ -25,6 +25,10 @@
 
 
 class ACIAuthorizer(CCAuthorizer):
+    def __init__(self, retry_count=5, retry_sleep=2):
+        self.retry_count = retry_count
+        self.retry_sleep = retry_sleep
+
     def generate(self):
         count = 0
         token = ""
@@ -40,9 +44,9 @@ def generate(self):
                     token = r.json().get("token")
                 break
             except:
-                if count > 5:
+                if count > self.retry_count:
                     break
-                time.sleep(2)
+                time.sleep(self.retry_sleep)
         return token
 
     def verify(self, token):