[Identity] Add custom IMDS retry policy (#42330)

This custom policy adds logic to increase the retry backoff factor in
the event that a response is a 410. It can take up to 70 seconds for
IMDS to become available.

Signed-off-by: Paul Van Eck <[email protected]>

Author: pvaneck

sdk/identity/azure-identity/CHANGELOG.md

@@ -12,6 +12,8 @@
 
 ### Other Changes
 
+- `ManagedIdentityCredential` now retries IMDS 410 status responses for at least 70 seconds total duration as required by [Azure IMDS documentation](https://learn.microsoft.com/azure/virtual-machines/instance-metadata-service?tabs=windows#errors-and-debugging).  ([#42330](https://github.com/Azure/azure-sdk-for-python/pull/42330))
+
 ## 1.24.0b1 (2025-07-17)
 
 ### Features Added

sdk/identity/azure-identity/azure/identity/_credentials/imds.py

@@ -6,9 +6,11 @@
 import json
 from typing import Any, Optional, Dict
 
+from azure.core.pipeline import PipelineResponse
 from azure.core.exceptions import ClientAuthenticationError, HttpResponseError
 from azure.core.pipeline.transport import HttpRequest
 from azure.core.credentials import AccessTokenInfo
+from azure.core.pipeline.policies import RetryPolicy
 
 from .. import CredentialUnavailableError
 from .._constants import EnvironmentVariables
@@ -31,6 +33,28 @@
 }
 
 
+class ImdsRetryPolicy(RetryPolicy):
+    """Custom retry policy for IMDS credential with extended retry duration for 410 responses.
+
+    This policy ensures that specifically for 410 status codes, the total exponential backoff duration
+    is at least 70 seconds to handle temporary IMDS endpoint unavailability.
+    For other status codes, it uses the standard retry behavior.
+    """
+
+    def __init__(self, **kwargs: Any) -> None:
+        # Increased backoff factor to ensure at least 70 seconds retry duration for 410 responses.
+        # Five retries, with each retry sleeping for [0.0s, 5.0s, 10.0s, 20.0s, 40.0s] between attempts (75s total)
+        self.backoff_factor_for_410 = 2.5
+        super().__init__(**kwargs)
+
+    def is_retry(self, settings: Dict[str, Any], response: PipelineResponse[Any, Any]) -> bool:
+        if response.http_response.status_code == 410:
+            settings["backoff"] = self.backoff_factor_for_410
+        else:
+            settings["backoff"] = self.backoff_factor
+        return super().is_retry(settings, response)
+
+
 def _get_request(scope: str, identity_config: Dict) -> HttpRequest:
     url = (
         os.environ.get(EnvironmentVariables.AZURE_POD_IDENTITY_AUTHORITY_HOST, IMDS_AUTHORITY).strip("/")
@@ -58,7 +82,7 @@ def _check_forbidden_response(ex: HttpResponseError) -> None:
 
 class ImdsCredential(MsalManagedIdentityClient):
     def __init__(self, **kwargs: Any) -> None:
-        super(ImdsCredential, self).__init__(**kwargs)
+        super().__init__(retry_policy_class=ImdsRetryPolicy, **dict(PIPELINE_SETTINGS, **kwargs))
         self._config = kwargs
 
         if EnvironmentVariables.AZURE_POD_IDENTITY_AUTHORITY_HOST in os.environ:

sdk/identity/azure-identity/azure/identity/_internal/pipeline.py

@@ -63,7 +63,8 @@ def _get_policies(config, _per_retry_policies=None, **kwargs):
 def build_pipeline(transport=None, policies=None, **kwargs):
     if not policies:
         config = _get_config(**kwargs)
-        config.retry_policy = RetryPolicy(**kwargs)
+        retry_policy_class = kwargs.pop("retry_policy_class", None)
+        config.retry_policy = retry_policy_class(**kwargs) if retry_policy_class else RetryPolicy(**kwargs)
         policies = _get_policies(config, **kwargs)
     if not transport:
         from azure.core.pipeline.transport import (  # pylint: disable=non-abstract-transport-import, no-name-in-module
@@ -82,7 +83,8 @@ def build_async_pipeline(transport=None, policies=None, **kwargs):
         from azure.core.pipeline.policies import AsyncRetryPolicy
 
         config = _get_config(**kwargs)
-        config.retry_policy = AsyncRetryPolicy(**kwargs)
+        retry_policy_class = kwargs.pop("retry_policy_class", None)
+        config.retry_policy = retry_policy_class(**kwargs) if retry_policy_class else AsyncRetryPolicy(**kwargs)
         policies = _get_policies(config, **kwargs)
     if not transport:
         from azure.core.pipeline.transport import (  # pylint: disable=non-abstract-transport-import, no-name-in-module

sdk/identity/azure-identity/azure/identity/aio/_credentials/imds.py

@@ -3,10 +3,13 @@
 # Licensed under the MIT License.
 # ------------------------------------
 import os
-from typing import Optional, Any
+from typing import Optional, Any, Dict
 
 from azure.core.exceptions import ClientAuthenticationError, HttpResponseError
 from azure.core.credentials import AccessTokenInfo
+from azure.core.pipeline.policies import AsyncRetryPolicy
+from azure.core.pipeline import PipelineResponse
+
 from ... import CredentialUnavailableError
 from ..._constants import EnvironmentVariables
 from .._internal import AsyncContextManager
@@ -16,10 +19,33 @@
 from ..._credentials.imds import _get_request, _check_forbidden_response, PIPELINE_SETTINGS
 
 
+class AsyncImdsRetryPolicy(AsyncRetryPolicy):
+    """Async custom retry policy for IMDS credential with extended retry duration for 410 responses.
+
+    This policy ensures that specifically for 410 status codes, the total exponential backoff duration
+    is at least 70 seconds to handle temporary IMDS endpoint unavailability.
+    For other status codes, it uses the standard retry behavior.
+    """
+
+    def __init__(self, **kwargs: Any) -> None:
+        # Increased backoff factor to ensure at least 70 seconds retry duration for 410 responses.
+        # Five retries, with each retry sleeping for [0.0s, 5.0s, 10.0s, 20.0s, 40.0s] between attempts (75s total)
+        self.backoff_factor_for_410 = 2.5
+        super().__init__(**kwargs)
+
+    def is_retry(self, settings: Dict[str, Any], response: PipelineResponse[Any, Any]) -> bool:
+        if response.http_response.status_code == 410:
+            settings["backoff"] = self.backoff_factor_for_410
+        else:
+            settings["backoff"] = self.backoff_factor
+        return super().is_retry(settings, response)
+
+
 class ImdsCredential(AsyncContextManager, GetTokenMixin):
     def __init__(self, **kwargs: Any) -> None:
         super().__init__()
 
+        kwargs["retry_policy_class"] = AsyncImdsRetryPolicy
         self._client = AsyncManagedIdentityClient(_get_request, **dict(PIPELINE_SETTINGS, **kwargs))
         if EnvironmentVariables.AZURE_POD_IDENTITY_AUTHORITY_HOST in os.environ:
             self._endpoint_available: Optional[bool] = True

sdk/identity/azure-identity/tests/test_imds_credential.py

@@ -4,10 +4,20 @@
 # ------------------------------------
 from itertools import product
 import time
+from unittest.mock import Mock
 
 from azure.identity import CredentialUnavailableError
-from azure.identity._credentials.imds import IMDS_TOKEN_PATH, ImdsCredential, IMDS_AUTHORITY
+from azure.identity._credentials.imds import (
+    IMDS_TOKEN_PATH,
+    ImdsCredential,
+    ImdsRetryPolicy,
+    IMDS_AUTHORITY,
+    PIPELINE_SETTINGS,
+)
 from azure.identity._internal.utils import within_credential_chain
+from azure.core.pipeline import PipelineResponse
+from azure.core.pipeline.policies import RetryPolicy
+from azure.core.pipeline.transport import HttpRequest, HttpResponse
 import pytest
 
 from helpers import mock, mock_response, Request, validating_transport, GET_TOKEN_METHODS
@@ -135,3 +145,53 @@ def test_managed_identity_aci_probe(self, get_token_method):
         token = getattr(credential, get_token_method)(scope)
         assert token.token == expected_token
         within_credential_chain.set(False)
+
+    def test_imds_credential_uses_custom_retry_policy(self):
+        credential = ImdsCredential()
+        policies = credential._client._pipeline._impl_policies
+        assert any(isinstance(policy, ImdsRetryPolicy) for policy in policies)
+        # Only one retry policy should be present
+        assert sum(isinstance(policy, RetryPolicy) for policy in policies) == 1
+
+    def test_imds_retry_policy(self):
+        retry_policy = ImdsRetryPolicy(**PIPELINE_SETTINGS)
+
+        # Create a shared mock HttpRequest
+        request = Mock(spec=HttpRequest, body=None, files=None)
+        request.method = "GET"
+
+        # Helper to create HttpResponse and PipelineResponse mocks
+        def make_pipeline_response(status_code):
+            response = Mock(spec=HttpResponse, status_code=status_code, http_request=request)
+            response.headers = {}
+            pipeline_response = Mock(spec=PipelineResponse, http_request=request, http_response=response)
+            return pipeline_response
+
+        pipeline_response_410 = make_pipeline_response(410)
+        pipeline_response_404 = make_pipeline_response(404)
+
+        # Simulate 5 retries for 410 response
+        settings_410 = retry_policy.configure_retries({})
+        total_time_410 = 0
+        for _ in range(5):
+            if retry_policy.is_retry(settings_410, pipeline_response_410):
+                retry_policy.increment(settings_410, response=pipeline_response_410, error=None)
+                backoff_time = retry_policy.get_backoff_time(settings_410)
+                total_time_410 += backoff_time
+
+        assert (
+            total_time_410 >= 70
+        ), f"Total retry time for 410 responses should be at least 70 seconds, got {total_time_410:.2f} seconds"
+
+        # Simulate 5 retries for 404 response
+        settings_404 = retry_policy.configure_retries({})
+        total_time_404 = 0
+        for _ in range(5):
+            if retry_policy.is_retry(settings_404, pipeline_response_404):
+                retry_policy.increment(settings_404, response=pipeline_response_404, error=None)
+                backoff_time = retry_policy.get_backoff_time(settings_404)
+                total_time_404 += backoff_time
+
+        assert (
+            total_time_404 < 30
+        ), f"Total retry time for 404 responses should use standard backoff, got {total_time_404:.2f} seconds"

sdk/identity/azure-identity/tests/test_imds_credential_async.py

@@ -6,13 +6,18 @@
 import json
 import time
 from unittest import mock
+from unittest.mock import Mock
 
 from azure.core.exceptions import ClientAuthenticationError
 from azure.identity import CredentialUnavailableError
 from azure.identity._constants import EnvironmentVariables
 from azure.identity._credentials.imds import IMDS_AUTHORITY, IMDS_TOKEN_PATH
 from azure.identity._internal.user_agent import USER_AGENT
-from azure.identity.aio._credentials.imds import ImdsCredential, PIPELINE_SETTINGS
+from azure.identity.aio._credentials.imds import ImdsCredential, AsyncImdsRetryPolicy
+from azure.identity._credentials.imds import PIPELINE_SETTINGS
+from azure.core.pipeline import PipelineResponse
+from azure.core.pipeline.policies import AsyncRetryPolicy
+from azure.core.pipeline.transport import HttpRequest, HttpResponse
 from azure.identity._internal.utils import within_credential_chain
 import pytest
 
@@ -346,3 +351,53 @@ async def test_managed_identity_aci_probe(self, get_token_method):
         token = await getattr(credential, get_token_method)(scope)
         assert token.token == expected_token
         within_credential_chain.set(False)
+
+    async def test_imds_credential_uses_custom_retry_policy(self):
+        credential = ImdsCredential()
+        policies = credential._client._pipeline._impl_policies  # type: ignore
+        assert any(isinstance(policy, AsyncImdsRetryPolicy) for policy in policies)
+        # Only one retry policy should be present
+        assert sum(isinstance(policy, AsyncRetryPolicy) for policy in policies) == 1
+
+    def test_imds_retry_policy(self):
+        retry_policy = AsyncImdsRetryPolicy(**PIPELINE_SETTINGS)
+
+        # Create a shared mock HttpRequest
+        request = Mock(spec=HttpRequest, body=None, files=None)
+        request.method = "GET"
+
+        # Helper to create HttpResponse and PipelineResponse mocks
+        def make_pipeline_response(status_code):
+            response = Mock(spec=HttpResponse, status_code=status_code, http_request=request)
+            response.headers = {}
+            pipeline_response = Mock(spec=PipelineResponse, http_request=request, http_response=response)
+            return pipeline_response
+
+        pipeline_response_410 = make_pipeline_response(410)
+        pipeline_response_404 = make_pipeline_response(404)
+
+        # Simulate 5 retries for 410 response
+        settings_410 = retry_policy.configure_retries({})
+        total_time_410 = 0
+        for _ in range(5):
+            if retry_policy.is_retry(settings_410, pipeline_response_410):
+                retry_policy.increment(settings_410, response=pipeline_response_410, error=None)
+                backoff_time = retry_policy.get_backoff_time(settings_410)
+                total_time_410 += backoff_time
+
+        assert (
+            total_time_410 >= 70
+        ), f"Total retry time for 410 responses should be at least 70 seconds, got {total_time_410:.2f} seconds"
+
+        # Simulate 5 retries for 404 response
+        settings_404 = retry_policy.configure_retries({})
+        total_time_404 = 0
+        for _ in range(5):
+            if retry_policy.is_retry(settings_404, pipeline_response_404):
+                retry_policy.increment(settings_404, response=pipeline_response_404, error=None)
+                backoff_time = retry_policy.get_backoff_time(settings_404)
+                total_time_404 += backoff_time
+
+        assert (
+            total_time_404 < 30
+        ), f"Total retry time for 404 responses should use standard backoff, got {total_time_404:.2f} seconds"