tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-05 20:23:18 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(bedrock): skip AssumeRole when ECS/EC2 already running as target IAM role

Browse Source 
When aws_role_name is configured but the environment (ECS task role, EC2 instance
profile) is already running as that role, AssumeRole is unnecessary and can fail
with AccessDenied. This adds same-role detection for ECS/EC2 (extending existing
IRSA support) and a fallback to ambient credentials when AssumeRole returns
AccessDenied.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
mubashir1osmani
2026-02-11 15:45:47 +05:30
committed by Sameer Kankute
co-authored by Claude
parent 2b9b5302ef
commit 70a1bf92e2
3 changed files with 270 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/llms/bedrock/base_aws_llm.py
+80 -18
View File
@@ -211,25 +211,13 @@ class BaseAWSLLM:
aws_external_id=aws_external_id,
)
elif aws_role_name is not None:
# Check if we're in IRSA and trying to assume the same role we already have
current_role_arn = os.getenv("AWS_ROLE_ARN")
web_identity_token_file = os.getenv("AWS_WEB_IDENTITY_TOKEN_FILE")
# In IRSA environments, we should skip role assumption if we're already running as the target role
# This is true when:
# 1. We have AWS_ROLE_ARN set (current role)
# 2. We have AWS_WEB_IDENTITY_TOKEN_FILE set (IRSA environment)
# 3. The current role matches the requested role
if (
current_role_arn
and web_identity_token_file
and current_role_arn == aws_role_name
):
# Check if we're already running as the target role and can skip assumption
# This handles IRSA (EKS), ECS task roles, and EC2 instance profiles
if self._is_already_running_as_role(aws_role_name):
verbose_logger.debug(
"Using IRSA same-role optimization: calling _auth_with_env_vars"
"Already running as target role %s, using ambient credentials",
aws_role_name,
)
# We're already running as this role via IRSA, no need to assume it again
# Use the default boto3 credentials (which will use the IRSA credentials)
credentials, _cache_ttl = self._auth_with_env_vars()
else:
verbose_logger.debug(
@@ -553,6 +541,65 @@ class BaseAWSLLM:
aws_region_name = "us-west-2"
return aws_region_name
def _is_already_running_as_role(self, aws_role_name: str) -> bool:
"""
Check if the current environment is already running as the target IAM role.
This handles multiple AWS environments:
- IRSA (EKS): AWS_ROLE_ARN + AWS_WEB_IDENTITY_TOKEN_FILE are set
- ECS task roles: Uses sts:GetCallerIdentity to check current role ARN
- EC2 instance profiles: Uses sts:GetCallerIdentity to check current role ARN
Returns True if the current identity matches the target role, meaning
we can skip sts:AssumeRole and use ambient credentials directly.
"""
# Fast path: IRSA environment check (no API call needed)
current_role_arn = os.getenv("AWS_ROLE_ARN")
web_identity_token_file = os.getenv("AWS_WEB_IDENTITY_TOKEN_FILE")
if current_role_arn and web_identity_token_file:
return current_role_arn == aws_role_name
# For ECS/EC2: call sts:GetCallerIdentity to check if already running as the role
try:
import boto3
with tracer.trace("boto3.client(sts).get_caller_identity"):
sts_client = boto3.client("sts")
identity = sts_client.get_caller_identity()
caller_arn = identity.get("Arn", "")
# The caller ARN for an ECS task role looks like:
# arn:aws:sts::123456789012:assumed-role/MyRole/session-name
# The target role ARN looks like:
# arn:aws:iam::123456789012:role/MyRole
# We need to compare the role name portion
if ":assumed-role/" in caller_arn:
# Extract role name from assumed-role ARN
# Format: arn:aws:sts::ACCOUNT:assumed-role/ROLE_NAME/SESSION
caller_role_name = caller_arn.split(":assumed-role/")[1].split("/")[0]
# Extract role name from target role ARN
# Format: arn:aws:iam::ACCOUNT:role/ROLE_NAME or
# arn:aws:iam::ACCOUNT:role/path/ROLE_NAME
if ":role/" in aws_role_name:
target_role_name = aws_role_name.split(":role/")[-1].split("/")[-1]
else:
target_role_name = aws_role_name
if caller_role_name == target_role_name:
verbose_logger.debug(
"Current identity already matches target role: %s",
aws_role_name,
)
return True
except Exception as e:
verbose_logger.debug(
"Could not determine current role identity: %s", str(e)
)
return False
@tracer.wrap()
def _auth_with_web_identity_token(
self,
@@ -867,7 +914,22 @@ class BaseAWSLLM:
if aws_external_id is not None:
assume_role_params["ExternalId"] = aws_external_id
sts_response = sts_client.assume_role(**assume_role_params)
try:
sts_response = sts_client.assume_role(**assume_role_params)
except Exception as e:
error_str = str(e)
# If AssumeRole fails because the caller already IS the role
# (e.g., ECS task role, root account, or same-role scenario),
# fall back to using ambient credentials directly
if "AccessDenied" in error_str:
verbose_logger.warning(
"AssumeRole failed for %s (%s). "
"Falling back to ambient credentials (boto3 default chain).",
aws_role_name,
error_str,
)
return self._auth_with_env_vars()
raise
# Extract the credentials from the response and convert to Session Credentials
sts_credentials = sts_response["Credentials"]
litellm/utils.py
+6
View File
@@ -6140,6 +6140,12 @@ def validate_environment( # noqa: PLR0915
if (
"AWS_ACCESS_KEY_ID" in os.environ
and "AWS_SECRET_ACCESS_KEY" in os.environ
) or (
# IAM role, profile, or web identity auth don't require access keys
"AWS_ROLE_ARN" in os.environ
or "AWS_PROFILE" in os.environ
or "AWS_WEB_IDENTITY_TOKEN_FILE" in os.environ
or "AWS_CONTAINER_CREDENTIALS_RELATIVE_URI" in os.environ # ECS task role
):
keys_in_environment = True
else:
tests/test_litellm/llms/bedrock/test_base_aws_llm.py
+184 -12
View File
@@ -853,29 +853,66 @@ def test_role_assumption_ttl_calculation():
assert 3500 <= ttl <= 3600 # Allow some variance for test execution time
def test_role_assumption_error_handling():
def test_role_assumption_access_denied_falls_back_to_env_vars():
"""
Test that role assumption errors are properly propagated.
Test that when AssumeRole fails with AccessDenied, we fall back to ambient credentials.
This handles ECS task roles, root accounts, and same-role scenarios where
AssumeRole is unnecessary because the caller already has the role's permissions.
"""
base_aws_llm = BaseAWSLLM()
# Mock the boto3 STS client to raise an exception
# Mock the boto3 STS client to raise AccessDenied
mock_sts_client = MagicMock()
mock_sts_client.assume_role.side_effect = Exception("AccessDenied: User is not authorized to perform sts:AssumeRole")
mock_sts_client.assume_role.side_effect = Exception(
"An error occurred (AccessDenied) when calling the AssumeRole operation: "
"Roles may not be assumed by root accounts."
)
# Mock _auth_with_env_vars to return fallback credentials
mock_creds = MagicMock()
mock_creds.access_key = "fallback-access-key"
mock_creds.secret_key = "fallback-secret-key"
with patch("boto3.client", return_value=mock_sts_client):
with patch.object(
base_aws_llm, "_auth_with_env_vars", return_value=(mock_creds, None)
) as mock_env_auth:
credentials, ttl = base_aws_llm._auth_with_aws_role(
aws_access_key_id=None,
aws_secret_access_key=None,
aws_session_token=None,
aws_role_name="arn:aws:iam::1111111111111:role/UnauthorizedRole",
aws_session_name="error-test-session",
)
# Should have fallen back to env vars
mock_env_auth.assert_called_once()
assert credentials.access_key == "fallback-access-key"
def test_role_assumption_non_access_denied_error_propagated():
"""
Test that non-AccessDenied errors from AssumeRole are still propagated.
"""
base_aws_llm = BaseAWSLLM()
# Mock the boto3 STS client to raise a non-AccessDenied exception
mock_sts_client = MagicMock()
mock_sts_client.assume_role.side_effect = Exception(
"An error occurred (MalformedPolicyDocument) when calling the AssumeRole operation"
)
with patch("boto3.client", return_value=mock_sts_client):
# Should raise the exception
with pytest.raises(Exception) as exc_info:
base_aws_llm._auth_with_aws_role(
aws_access_key_id=None,
aws_secret_access_key=None,
aws_session_token=None,
aws_role_name="arn:aws:iam::1111111111111:role/UnauthorizedRole",
aws_session_name="error-test-session"
aws_role_name="arn:aws:iam::1111111111111:role/BadPolicyRole",
aws_session_name="error-test-session",
)
assert "AccessDenied" in str(exc_info.value)
assert "MalformedPolicyDocument" in str(exc_info.value)
def test_multiple_role_assumptions_in_sequence():
@@ -1195,3 +1232,138 @@ def test_converse_handler_external_id_extraction():
assert hasattr(mock_get_credentials, 'called_kwargs')
assert "aws_external_id" in mock_get_credentials.called_kwargs
assert mock_get_credentials.called_kwargs["aws_external_id"] == "TestExternalID123"
def test_is_already_running_as_role_irsa_same_role():
"""Test IRSA fast path: when AWS_ROLE_ARN matches target role."""
base_aws_llm = BaseAWSLLM()
with patch.dict(os.environ, {
"AWS_ROLE_ARN": "arn:aws:iam::123456789012:role/MyRole",
"AWS_WEB_IDENTITY_TOKEN_FILE": "/var/run/secrets/token",
}):
assert base_aws_llm._is_already_running_as_role(
"arn:aws:iam::123456789012:role/MyRole"
) is True
def test_is_already_running_as_role_irsa_different_role():
"""Test IRSA fast path: when AWS_ROLE_ARN does NOT match target role."""
base_aws_llm = BaseAWSLLM()
with patch.dict(os.environ, {
"AWS_ROLE_ARN": "arn:aws:iam::123456789012:role/MyRole",
"AWS_WEB_IDENTITY_TOKEN_FILE": "/var/run/secrets/token",
}):
assert base_aws_llm._is_already_running_as_role(
"arn:aws:iam::999999999999:role/OtherRole"
) is False
def test_is_already_running_as_role_ecs_task_role():
"""Test ECS/EC2 path: GetCallerIdentity shows assumed-role matching target."""
base_aws_llm = BaseAWSLLM()
mock_sts_client = MagicMock()
mock_sts_client.get_caller_identity.return_value = {
"Arn": "arn:aws:sts::123456789012:assumed-role/MyEcsTaskRole/ecs-task-id"
}
with patch.dict(os.environ, {}, clear=False):
# Ensure no IRSA env vars
env = {k: v for k, v in os.environ.items() if k not in ("AWS_ROLE_ARN", "AWS_WEB_IDENTITY_TOKEN_FILE")}
with patch.dict(os.environ, env, clear=True):
with patch("boto3.client", return_value=mock_sts_client):
assert base_aws_llm._is_already_running_as_role(
"arn:aws:iam::123456789012:role/MyEcsTaskRole"
) is True
def test_is_already_running_as_role_ecs_different_role():
"""Test ECS/EC2 path: GetCallerIdentity shows a different role."""
base_aws_llm = BaseAWSLLM()
mock_sts_client = MagicMock()
mock_sts_client.get_caller_identity.return_value = {
"Arn": "arn:aws:sts::123456789012:assumed-role/MyEcsTaskRole/ecs-task-id"
}
with patch.dict(os.environ, {}, clear=False):
env = {k: v for k, v in os.environ.items() if k not in ("AWS_ROLE_ARN", "AWS_WEB_IDENTITY_TOKEN_FILE")}
with patch.dict(os.environ, env, clear=True):
with patch("boto3.client", return_value=mock_sts_client):
assert base_aws_llm._is_already_running_as_role(
"arn:aws:iam::999999999999:role/DifferentRole"
) is False
def test_is_already_running_as_role_ecs_role_with_path():
"""Test ECS path with role that has a path prefix (e.g., /service-role/MyRole)."""
base_aws_llm = BaseAWSLLM()
mock_sts_client = MagicMock()
mock_sts_client.get_caller_identity.return_value = {
"Arn": "arn:aws:sts::123456789012:assumed-role/MyEcsTaskRole/ecs-task-id"
}
with patch.dict(os.environ, {}, clear=False):
env = {k: v for k, v in os.environ.items() if k not in ("AWS_ROLE_ARN", "AWS_WEB_IDENTITY_TOKEN_FILE")}
with patch.dict(os.environ, env, clear=True):
with patch("boto3.client", return_value=mock_sts_client):
# Role ARN with path
assert base_aws_llm._is_already_running_as_role(
"arn:aws:iam::123456789012:role/service-role/MyEcsTaskRole"
) is True
def test_is_already_running_as_role_get_caller_identity_fails():
"""Test that when GetCallerIdentity fails, we return False (don't crash)."""
base_aws_llm = BaseAWSLLM()
mock_sts_client = MagicMock()
mock_sts_client.get_caller_identity.side_effect = Exception("No credentials found")
with patch.dict(os.environ, {}, clear=False):
env = {k: v for k, v in os.environ.items() if k not in ("AWS_ROLE_ARN", "AWS_WEB_IDENTITY_TOKEN_FILE")}
with patch.dict(os.environ, env, clear=True):
with patch("boto3.client", return_value=mock_sts_client):
assert base_aws_llm._is_already_running_as_role(
"arn:aws:iam::123456789012:role/SomeRole"
) is False
def test_get_credentials_ecs_same_role_skips_assume_role():
"""
End-to-end test: when running on ECS with the same role as aws_role_name,
get_credentials should use ambient credentials and NOT call AssumeRole.
"""
base_aws_llm = BaseAWSLLM()
mock_creds = MagicMock()
mock_creds.access_key = "ecs-access-key"
mock_creds.secret_key = "ecs-secret-key"
mock_creds.token = "ecs-session-token"
with patch.object(
base_aws_llm,
"_is_already_running_as_role",
return_value=True,
):
with patch.object(
base_aws_llm,
"_auth_with_env_vars",
return_value=(mock_creds, None),
) as mock_env_auth:
with patch.object(
base_aws_llm,
"_auth_with_aws_role",
) as mock_role_auth:
credentials = base_aws_llm.get_credentials(
aws_role_name="arn:aws:iam::123456789012:role/MyEcsTaskRole",
aws_region_name="us-east-1",
)
# Should use env vars, NOT role assumption
mock_env_auth.assert_called_once()
mock_role_auth.assert_not_called()
assert credentials.access_key == "ecs-access-key"