Skip to content

change: move service-role path parsing for AmazonSageMaker-ExecutionRole for get_execution_role() into except block of IAM get_role() call and add warning message #2191

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Mar 8, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 18 additions & 8 deletions src/sagemaker/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -3498,14 +3498,6 @@ def get_caller_identity_arn(self):
endpoint_url=sts_regional_endpoint(self.boto_region_name),
).get_caller_identity()["Arn"]

if "AmazonSageMaker-ExecutionRole" in assumed_role:
role = re.sub(
r"^(.+)sts::(\d+):assumed-role/(.+?)/.*$",
r"\1iam::\2:role/service-role/\3",
assumed_role,
)
return role

role = re.sub(r"^(.+)sts::(\d+):assumed-role/(.+?)/.*$", r"\1iam::\2:role/\3", assumed_role)

# Call IAM to get the role's path
Expand All @@ -3518,6 +3510,24 @@ def get_caller_identity_arn(self):
role_name,
)

# This conditional has been present since the inception of SageMaker
# Guessing this conditional's purpose was to handle lack of IAM permissions
# https://github.com/aws/sagemaker-python-sdk/issues/2089#issuecomment-791802713
if "AmazonSageMaker-ExecutionRole" in assumed_role:
LOGGER.warning(
"Assuming role was created in SageMaker AWS console, "
"as the name contains `AmazonSageMaker-ExecutionRole`. "
"Defaulting to Role ARN with service-role in path. "
"If this Role ARN is incorrect, please add "
"IAM read permissions to your role or supply the "
"Role Arn directly."
)
role = re.sub(
r"^(.+)sts::(\d+):assumed-role/(.+?)/.*$",
r"\1iam::\2:role/service-role/\3",
assumed_role,
)

return role

def logs_for_job( # noqa: C901 - suppress complexity warning for this method
Expand Down
20 changes: 19 additions & 1 deletion tests/unit/test_session.py
Original file line number Diff line number Diff line change
Expand Up @@ -398,12 +398,30 @@ def test_get_caller_identity_arn_from_a_role(sts_regional_endpoint, boto_session
@patch("os.path.exists", side_effect=mock_exists(NOTEBOOK_METADATA_FILE, False))
@patch("sagemaker.session.sts_regional_endpoint", return_value=STS_ENDPOINT)
def test_get_caller_identity_arn_from_an_execution_role(sts_regional_endpoint, boto_session):
sess = Session(boto_session)
sts_arn = "arn:aws:sts::369233609183:assumed-role/AmazonSageMaker-ExecutionRole-20171129T072388/SageMaker"
sess.boto_session.client("sts", endpoint_url=STS_ENDPOINT).get_caller_identity.return_value = {
"Arn": sts_arn
}
iam_arn = "arn:aws:iam::369233609183:role/AmazonSageMaker-ExecutionRole-20171129T072388"
sess.boto_session.client("iam").get_role.return_value = {"Role": {"Arn": iam_arn}}

actual = sess.get_caller_identity_arn()
assert actual == iam_arn


@patch("os.path.exists", side_effect=mock_exists(NOTEBOOK_METADATA_FILE, False))
@patch("sagemaker.session.sts_regional_endpoint", return_value=STS_ENDPOINT)
def test_get_caller_identity_arn_from_a_sagemaker_execution_role_with_iam_client_error(
sts_regional_endpoint, boto_session
):
sess = Session(boto_session)
arn = "arn:aws:sts::369233609183:assumed-role/AmazonSageMaker-ExecutionRole-20171129T072388/SageMaker"
sess.boto_session.client("sts", endpoint_url=STS_ENDPOINT).get_caller_identity.return_value = {
"Arn": arn
}
sess.boto_session.client("iam").get_role.return_value = {"Role": {"Arn": arn}}

sess.boto_session.client("iam").get_role.side_effect = ClientError({}, {})

actual = sess.get_caller_identity_arn()
assert (
Expand Down