scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-08 03:08:45 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 115

feat(bedrock): make the IAM credential cache bounds env configurable

Browse source 
BEDROCK_IAM_CACHE_MAX_ENTRIES and BEDROCK_IAM_CACHE_FETCH_LOCK_STRIPES were
compile-time literals, so an operator whose live credential-key cardinality
exceeds 1000 had no way to raise the bound short of patching the package.
The cache then evicts and re-assumes, restoring the sts:AssumeRole volume the
cache exists to remove.

Both now use the os.getenv pattern the rest of constants.py already uses,
with the current values as defaults.
This commit is contained in:
Aryan Pardeshi 2026-08-11 01:25:46 +05:30
parent 9bca9dfbb1
commit 861e15daba
2 changed files with 35 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/constants.py
View file

@ -374,10 +374,10 @@ MAX_STRING_LENGTH_PROMPT_IN_DB: Final = int(os.getenv("MAX_STRING_LENGTH_PROMPT_
BEDROCK_MAX_POLICY_SIZE: Final = int(os.getenv("BEDROCK_MAX_POLICY_SIZE", 75))
# One entry per distinct AWS credential-argument set. Per-user cost attribution passes the attributed
# identity as aws_session_name, so this bounds how many attributed identities keep a cached STS session.
BEDROCK_IAM_CACHE_MAX_ENTRIES: Final = 1000
BEDROCK_IAM_CACHE_MAX_ENTRIES: Final = int(os.getenv("BEDROCK_IAM_CACHE_MAX_ENTRIES", 1000))
# Single-flight lock stripes over that cache. Only keys landing on the same stripe wait for each
# other, so a burst of distinct identities still resolves its credentials in parallel.
BEDROCK_IAM_CACHE_FETCH_LOCK_STRIPES: Final = 64
BEDROCK_IAM_CACHE_FETCH_LOCK_STRIPES: Final = int(os.getenv("BEDROCK_IAM_CACHE_FETCH_LOCK_STRIPES", 64))
# Retire a cached STS credential this many seconds before AWS expires it, so a request that reads it
# still has a usable credential for the whole call.
STS_CREDENTIAL_EXPIRY_SAFETY_MARGIN_SECONDS: Final = 60

33
tests/test_litellm/test_constants.py
View file

@ -71,3 +71,36 @@ def _build_constant_env_var_map() -> dict[str, str]:
env_var_map[constant_name] = env_var_name
return env_var_map
@pytest.mark.parametrize(
"constant_name, override",
[
("BEDROCK_IAM_CACHE_MAX_ENTRIES", 4096),
("BEDROCK_IAM_CACHE_FETCH_LOCK_STRIPES", 128),
],
)
def test_bedrock_iam_cache_bounds_are_env_overridable(constant_name: str, override: int) -> None:
try:
with mock.patch.dict(os.environ, {constant_name: str(override)}):
reloaded = importlib.reload(constants)
assert getattr(reloaded, constant_name) == override
finally:
importlib.reload(constants)
@pytest.mark.parametrize(
"constant_name, expected_default",
[
("BEDROCK_IAM_CACHE_MAX_ENTRIES", 1000),
("BEDROCK_IAM_CACHE_FETCH_LOCK_STRIPES", 64),
],
)
def test_bedrock_iam_cache_bounds_keep_their_defaults(constant_name: str, expected_default: int) -> None:
try:
with mock.patch.dict(os.environ, {}, clear=False):
os.environ.pop(constant_name, None)
reloaded = importlib.reload(constants)
assert getattr(reloaded, constant_name) == expected_default
finally:
importlib.reload(constants)