diff --git a/litellm/constants.py b/litellm/constants.py index 23e92d26a59..57b8abd9553 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -400,10 +400,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 diff --git a/tests/test_litellm/test_constants.py b/tests/test_litellm/test_constants.py index 12e473f68a4..0ae9f150269 100644 --- a/tests/test_litellm/test_constants.py +++ b/tests/test_litellm/test_constants.py @@ -68,3 +68,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)