Add IAM cred caching for OIDC flow #3712
Conversation
|
The latest updates on your projects. Learn more about Vercel for Git ↗︎
|
There was a problem hiding this comment.
#3858 blocks this pr
since the refactoring to make the call async will require a refactor here
|
@Manouchehri please can we add some testing for this? couple users use the sts flow and i want to make sure there's no regressions |
Definitely, that's a good idea, I'll do. We sorta ish already test it, but it's more by luck and not intentional. Two questions:
|
212c953 to
3410367
Compare
|
This shaves a lot of time off the current setup, between like 0.25 to 1 second per Bedrock call. :) The STS fix is also needed. Ready for merging! |
|
Bump on this? It's becoming a bit of a burden to maintain out-of-tree patches for weeks. |
|
hey @Manouchehri missed this - sorry for the delay on my end. Saw you made a couple changes - could you walk me through it? |
|
i'm also happy to hop on a call and help get all the oidc pr's done https://calendly.com/d/4mp-gd3-k5k/litellm-1-1-onboarding-chat |
Gimme a moment to finish testing this and #3861. =) |
|
You free now? |
|
|
||
| iam_cache = DualCache() |
There was a problem hiding this comment.
switch to in memory cache
There was a problem hiding this comment.
InMemoryCache -
Line 59 in fb96f07
| aws_session_token=sts_response["Credentials"]["SessionToken"], | ||
| region_name=aws_region_name, | ||
| ) | ||
| iam_cache.set_cache(key=iam_creds_cache_key, value=json.dumps(iam_creds_dict), ttl=3600 - 60) |
There was a problem hiding this comment.
can we have ttl be an enum at the top of the file, so it's easier to know
| and aws_session_name is not None | ||
| ): | ||
| oidc_token = get_secret(aws_web_identity_token) | ||
| if aws_web_identity_token is not None and aws_role_name is not None and aws_session_name is not None: |
There was a problem hiding this comment.
can you add a comment explaining when you'd enter this flow
|
Self note: I can confirm that this is indeed working, I now only see one 
Previously (before this PR was merged), I had multiple 
|
Title
This should improve the perf, as this allows avoiding making an extra call to STS on every invocation of Bedrock.
Should be merged after #3688. Doing this as a separate PR in case you don't like my caching method; I don't want this PR to slow down #3688 from being merged.
Type
🆕 New Feature
Changes
Just uses DualCache to speed things up!
[REQUIRED] Testing - Attach a screenshot of any new tests passing locall
If UI changes, send a screenshot/GIF of working UI fixes
I ran the same request multiple times, and could see the cache being used in LiteLLM's logs. (It's kinda difficult to post a screenshot without including a lot of PII.)