fix(auth): retry transient ARK API key fetch failures (#989)

Co-authored-by: shanchunhua <shanchunhua@bytedance.com>

Author: xgtcode

tests/auth/veauth/test_ark_veauth.py

@@ -27,6 +27,22 @@ def _raw(api_key):
     return {"Result": {"ApiKey": api_key}}
 
 
+def _raw_error(code, message, code_number=None):
+    error = {"Code": code, "Message": message}
+    if code_number is not None:
+        error["CodeN"] = code_number
+    return {
+        "ResponseMetadata": {
+            "RequestId": "request-id",
+            "Action": "GetRawApiKey",
+            "Version": "2024-01-01",
+            "Service": "ark",
+            "Region": "cn-beijing",
+            "Error": error,
+        }
+    }
+
+
 def _key(id_, name):
     return {"Id": id_, "Name": name}
 
@@ -127,6 +143,47 @@ def test_numeric_id_is_sent_as_control_plane_integer():
     }
 
 
+def test_raw_api_key_retries_internal_service_timeout():
+    with (
+        patch("veadk.auth.veauth.ark_veauth.ve_request") as request,
+        patch("veadk.auth.veauth.ark_veauth.time.sleep") as sleep,
+    ):
+        request.side_effect = [
+            _raw_error(
+                "InternalServiceTimeout",
+                "Internal Service is timeout. Pls Contact With Admin",
+                code_number=100016,
+            ),
+            _raw_error(
+                "InternalServiceTimeout",
+                "Internal Service is timeout. Pls Contact With Admin",
+                code_number=100016,
+            ),
+            _raw("sk-SELECTED"),
+        ]
+
+        assert get_ark_token(api_key_id="4028965") == "sk-SELECTED"
+
+    assert request.call_count == 3
+    assert sleep.call_count == 2
+
+
+def test_raw_api_key_does_not_retry_non_transient_error():
+    with (
+        patch("veadk.auth.veauth.ark_veauth.ve_request") as request,
+        patch("veadk.auth.veauth.ark_veauth.time.sleep") as sleep,
+    ):
+        request.return_value = _raw_error(
+            "AccessDenied", "No permission to call GetRawApiKey."
+        )
+
+        with pytest.raises(ValueError, match="AccessDenied"):
+            get_ark_token(api_key_id="4028965")
+
+    assert request.call_count == 1
+    sleep.assert_not_called()
+
+
 def test_environment_sts_token_is_used_for_signed_requests(monkeypatch):
     monkeypatch.setenv("VOLCENGINE_SESSION_TOKEN", "temporary-session-token")
     with patch("veadk.auth.veauth.ark_veauth.ve_request") as request:

veadk/auth/veauth/ark_veauth.py

@@ -13,6 +13,10 @@
 # limitations under the License.
 
 import os
+import random
+import time
+
+import requests
 
 from veadk.auth.veauth.utils import get_credential_from_vefaas_iam
 from veadk.utils.logger import get_logger
@@ -25,6 +29,40 @@
 # name or exhaust the list.
 _ARK_PROJECT_NAME = "default"
 _ARK_PAGE_SIZE = 100
+_RAW_API_KEY_MAX_ATTEMPTS = 3
+_RAW_API_KEY_RETRY_DELAYS_SECONDS = (0.2, 0.5)
+_RAW_API_KEY_RETRYABLE_ERROR_CODES = {
+    "InternalError",
+    "InternalServiceTimeout",
+    "RequestTimeout",
+    "ServiceUnavailable",
+    "Throttling",
+    "TooManyRequests",
+}
+_RAW_API_KEY_RETRYABLE_ERROR_CODE_NUMBERS = {100016}
+
+
+def _response_error(res: dict) -> dict:
+    metadata = res.get("ResponseMetadata", {})
+    if not isinstance(metadata, dict):
+        return {}
+    error = metadata.get("Error", {})
+    return error if isinstance(error, dict) else {}
+
+
+def _is_retryable_raw_api_key_response(res: dict) -> bool:
+    error = _response_error(res)
+    return (
+        error.get("Code") in _RAW_API_KEY_RETRYABLE_ERROR_CODES
+        or error.get("CodeN") in _RAW_API_KEY_RETRYABLE_ERROR_CODE_NUMBERS
+    )
+
+
+def _sleep_before_raw_api_key_retry(attempt: int) -> None:
+    delay = _RAW_API_KEY_RETRY_DELAYS_SECONDS[
+        min(attempt - 1, len(_RAW_API_KEY_RETRY_DELAYS_SECONDS) - 1)
+    ]
+    time.sleep(delay + random.uniform(0, delay / 5))
 
 
 def get_ark_token(
@@ -142,20 +180,46 @@ def _list_api_keys(page_number: int) -> dict:
         if isinstance(target_id, str) and target_id.isdigit()
         else target_id
     )
-    res = ve_request(
-        request_body={"Id": request_key_id, "ProjectName": _ARK_PROJECT_NAME},
-        header={"X-Security-Token": session_token},
-        action="GetRawApiKey",
-        ak=access_key,
-        sk=secret_key,
-        service="ark",
-        version="2024-01-01",
-        region=region,
-        host=host,
-    )
-    try:
-        api_key = res["Result"]["ApiKey"]
-        logger.info("Successfully fetched ARK API Key.")
-        return api_key
-    except KeyError as error:
-        raise ValueError("Failed to get ARK API key.") from error
+    request_body = {"Id": request_key_id, "ProjectName": _ARK_PROJECT_NAME}
+    for attempt in range(1, _RAW_API_KEY_MAX_ATTEMPTS + 1):
+        try:
+            res = ve_request(
+                request_body=request_body,
+                header={"X-Security-Token": session_token},
+                action="GetRawApiKey",
+                ak=access_key,
+                sk=secret_key,
+                service="ark",
+                version="2024-01-01",
+                region=region,
+                host=host,
+            )
+        except requests.exceptions.RequestException as error:
+            if attempt >= _RAW_API_KEY_MAX_ATTEMPTS:
+                raise ValueError("Failed to get ARK api key.") from error
+            logger.warning(
+                "GetRawApiKey request failed; retrying "
+                f"({attempt}/{_RAW_API_KEY_MAX_ATTEMPTS}): {error}"
+            )
+            _sleep_before_raw_api_key_retry(attempt)
+            continue
+
+        try:
+            api_key = res["Result"]["ApiKey"]
+            logger.info("Successfully fetched ARK API Key.")
+            return api_key
+        except KeyError as error:
+            if (
+                attempt < _RAW_API_KEY_MAX_ATTEMPTS
+                and _is_retryable_raw_api_key_response(res)
+            ):
+                logger.warning(
+                    "GetRawApiKey returned a retryable error; retrying "
+                    f"({attempt}/{_RAW_API_KEY_MAX_ATTEMPTS}): "
+                    f"{_response_error(res)}"
+                )
+                _sleep_before_raw_api_key_retry(attempt)
+                continue
+            raise ValueError(f"Failed to get ARK api key: {res}") from error
+
+    raise ValueError("Failed to get ARK api key.")