Add retries to allow for rate limits on WIF

anthropics · Feb 1, 2024 · 7b64e60 · 7b64e60
1 parent e8d35f0
commit 7b64e60
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 1 deletion.
diff --git a/blobfile/_gcp.py b/blobfile/_gcp.py
@@ -33,6 +33,12 @@
     strip_slashes,
 )
 
+from tenacity import (
+    retry,
+    retry_if_exception_type,
+    stop_after_attempt,
+)
+
 MAX_EXPIRATION = 7 * 24 * 60 * 60
 OAUTH_TOKEN = "oauth_token"
 ANONYMOUS = "anonymous"
@@ -474,7 +480,11 @@ def build_req() -> Request:
     result = json.loads(resp.data)
     return (OAUTH_TOKEN, result["access_token"]), now + float(result["expires_in"])
 
-
+@retry(
+    retry=retry_if_exception_type(ValueError),
+    before=get_wif_credentials,
+    stop=stop_after_attempt(1),
+)
 def _get_access_token(conf: Config, key: Any) -> Tuple[Any, float]:
     if os.environ.get("BLOBFILE_FORCE_GOOGLE_ANONYMOUS_AUTH", "0") == "1":
         return (ANONYMOUS, ""), float("inf")

diff --git a/env.yaml b/env.yaml
@@ -19,3 +19,4 @@ dependencies:
     - boto3==1.15.18
     - lxml-stubs==0.4.0
     - xmltodict==0.13.0
+    - tenacity==8.2.2