Skip to content

Commit b45d79c

Browse files
authored
Fix Estimator role expansion (aws#68)
Instead of manually constructing the role ARN, use the IAM boto client to do it. This properly expands service-roles and regular roles.
1 parent e1d79d5 commit b45d79c

File tree

1 file changed

+3
-4
lines changed

1 file changed

+3
-4
lines changed

src/sagemaker/session.py

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -522,8 +522,8 @@ def endpoint_from_production_variants(self, name, production_variants, wait=True
522522
def expand_role(self, role):
523523
"""Expand an IAM role name into an ARN.
524524
525-
If the role is already in the form of an ARN, then the role is simply returned. Otherwise, the role
526-
is formatted as an ARN, using the current account as the IAM role's AWS account.
525+
If the role is already in the form of an ARN, then the role is simply returned. Otherwise we retrieve the full
526+
ARN and return it.
527527
528528
Args:
529529
role (str): An AWS IAM role (either name or full ARN).
@@ -534,8 +534,7 @@ def expand_role(self, role):
534534
if '/' in role:
535535
return role
536536
else:
537-
account = self.boto_session.client('sts').get_caller_identity()['Account']
538-
return 'arn:aws:iam::{}:role/{}'.format(account, role)
537+
return boto3.resource("iam").Role(role).arn
539538

540539
def get_caller_identity_arn(self):
541540
"""Returns the ARN user or role whose credentials are used to call the API.

0 commit comments

Comments
 (0)