feat(util.py/azure.py): Add OIDC support when running LiteLLM on Azure + Azure Upstream caching

[Manouchehri]

Title

This adds OIDC support for AKS (Managed Kubernetes Service).

Relevant issues

Resolves #1852.

Type

🆕 New Feature

Changes

TBD

[REQUIRED] Testing - Attach a screenshot of any new tests passing locally

TBD.

@RyoYang or @olad32, could you please try this config below and let me know what happens? (Feel free to use any Azure OpenAI model.)

model_list:
  - model_name: gpt-4-0125-preview
    litellm_params:
      model: azure/gpt-4-0125-preview
      api_version: "2024-05-01-preview"
      azure_ad_token: "oidc/azure/api://AzureADTokenExchange"
      api_base: "https://example.openai.azure.com"
    model_info:
      base_model: azure/gpt-4-0125-preview

[vercel]

The latest updates on your projects. Learn more about Vercel for Git ↗︎

Name	Status	Preview	Comments	Updated (UTC)
litellm	✅ Ready (Inspect)	Visit Preview	💬 Add feedback	Jun 11, 2024 3:57pm

[RyoYang]

@Manouchehri Thanks for the quick updates, the functionality looks good in my environments.

[Manouchehri]

You successfully ran it?

[RyoYang]

Yea, it ran well in my AKS.

[Manouchehri]

Oh wow, that’s awesome I guessed it right on the first try. Thanks for confirming!

[RyoYang]

Oh wow, that’s awesome I guessed it right on the first try. Thanks for confirming!

Great work! How long can you checkin the change into main branch?

[Manouchehri]

@RyoYang Could you help me get some more info so I can write a unit test as well?

Option 1: Could you share AZURE_AUTHORITY_HOST and the contents of AZURE_FEDERATED_TOKEN_FILE with me?

For AZURE_FEDERATED_TOKEN_FILE, make sure to remove the part after the last dot so you don't give me access to your real infrastructure.

e.g. if you have:

eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQ.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c

Please only shared with me:

eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQ.SIGNATURE_REMOVED

Option 2: Run with --detailed_debug. You should see the JWT/OpenID token nearby a line that says oidc_token.

https://github.com/BerriAI/litellm/pull/3861/files#diff-fedcce684ff5b9a48140f0db2fb56e423c036e45a550c9f22c6ad3c20372862dR159-R164

Do the same thing (remove the last bit after the last .) as mentioned above.

[Manouchehri]

@krrishdholakia This is ready for merging. The CircleCI + Bedrock tests currently cover most of what this does too; that said, I don't have an easy way to test this myself (as it would require setting LiteLLM up on Azure Kubernetes Service).

[RyoYang]

@RyoYang Could you help me get some more info so I can write a unit test as well?

Option 1: Could you share AZURE_AUTHORITY_HOST and the contents of AZURE_FEDERATED_TOKEN_FILE with me?

For AZURE_FEDERATED_TOKEN_FILE, make sure to remove the part after the last dot so you don't give me access to your real infrastructure.

e.g. if you have:

eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQ.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c

Please only shared with me:

eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQ.SIGNATURE_REMOVED

Option 2: Run with --detailed_debug. You should see the JWT/OpenID token nearby a line that says oidc_token.

https://github.com/BerriAI/litellm/pull/3861/files#diff-fedcce684ff5b9a48140f0db2fb56e423c036e45a550c9f22c6ad3c20372862dR159-R164

Do the same thing (remove the last bit after the last .) as mentioned above.

Hey @Manouchehri Due to the security consideration, I cannot share those string to public repo. Do we have any other way for testing?

[RyoYang]

Also, can we merge this PR at first, as we will make this repo as submodule with this change, without this PR, we can not move forward. @Manouchehri @krrishdholakia

[Manouchehri]

Bump?