Enable token caching in azure-identity

[gdippolito]

Is your feature request related to a problem? Please describe. Running DefaultAzureCredentials authentication on different threads do not cache the token and would make separate authentication requests. This is especially annoying when using ManagedIdentity when running with many threads. The managed identity endpoint will rate limit the request and therefore cause authentication errors on some of the threads.

Describe the solution you'd like I would like the token retrieved to be cached and reused from different threads. This will significantly reduce the number of requests on the authentication endpoint.

Describe alternatives you've considered I have created a python singleton to avoid re-initialising the same object. However, this has been a bit painful to manage.

[kristapratico]

Hi @gdippolito thanks for your feedback, we'll look into this as soon as possible.

[xiangyan99]

Thanks for reaching out.

Could you give more details about your scenario?

And if possible, can you share the code how you used the ManagedIdentityCredential?

[gdippolito]

@xiangyan99 I have create a simple repository to illustrate how we have workaround the rate limit on the managed identity endpoint here

Our use case is mostly during AI training. When we start a training job we will spin up multiple dataloader processes for each GPU (in our case 8 GPUs) and several threads within those processes to fetch data from several storage accounts (at least 8).

Please let me know if the small repo I have created (sorry it is a bit rough) helps.

[xiangyan99]

@gdippolito Thank you for sharing the code.

Our token caching mechanism is per-instance caching.

If every time you create a new credential instance, the token caching will not work as you expected.

We can use single credential or single storage client to make the caching work.

Please let me know if it works.

Thank you.

[gdippolito]

Hi @xiangyan99 thanks for the response.

Would you mind clarifying what single credential or single storage client means? In my case we spawn many thread and processes I'm not sure how I could make a single client?

Thanks.

[xiangyan99]

e.g.

Shared credential instance:

credential = DefaultAzureCredential()
client1 = BlogStorageClient(credential=credential)
client2 = BlogStorageClient(credential=credential)

The token will be cached.

But if we use

client1 = BlogStorageClient(credential=DefaultAzureCredential())
client2 = BlogStorageClient(credential=DefaultAzureCredential())

The token will not be cached.

Thanks.

[gdippolito]

Hi @xiangyan99 thanks for the response. I think this will only work if you define multiple clients sequentially. However, the token won't be cached when using multiple Threads and processes. I think a strategy for caching token already exists for Interactive authentication, can this not be added to managed identity as well?

[xiangyan99]

As you are aware, our managed identity and interactive identity are implemented in different ways.

Interactive identity uses msal underneath while managed identity does not.

There is a tracking request to add managed identity support in msal. https://github.com/AzureAD/microsoft-authentication-library-for-python/issues/58

When it is implemented, we can move azure.identity.managedidentity to use msal too therefore you will see same behavior.

[ghost]

Hi, we're sending this friendly reminder because we haven't heard back from you in a while. We need more information about this issue to help address it. Please be sure to give us your input within the next 7 days. If we don't hear back from you within 14 days of this comment the issue will be automatically closed. Thank you!