Skip to content

Commit 0ad2fae

Browse files
author
Jun Lyu
committed
fix: update get_execution_role to directly return the ExecutionRoleArn if it presents in the resource metadata file
1 parent 2e2b27c commit 0ad2fae

File tree

2 files changed

+23
-4
lines changed

2 files changed

+23
-4
lines changed

src/sagemaker/session.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -6285,16 +6285,16 @@ def get_caller_identity_arn(self):
62856285
user_profile_name = metadata.get("UserProfileName")
62866286
execution_role_arn = metadata.get("ExecutionRoleArn")
62876287
try:
6288+
# find execution role from the metadata file if present
6289+
if execution_role_arn is not None:
6290+
return execution_role_arn
6291+
62886292
if domain_id is None:
62896293
instance_desc = self.sagemaker_client.describe_notebook_instance(
62906294
NotebookInstanceName=instance_name
62916295
)
62926296
return instance_desc["RoleArn"]
62936297

6294-
# find execution role from the metadata file if present
6295-
if execution_role_arn is not None:
6296-
return execution_role_arn
6297-
62986298
user_profile_desc = self.sagemaker_client.describe_user_profile(
62996299
DomainId=domain_id, UserProfileName=user_profile_name
63006300
)

tests/unit/test_session.py

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -728,6 +728,25 @@ def test_get_caller_identity_arn_from_metadata_file_for_space(boto_session):
728728
assert actual == expected_role
729729

730730

731+
@patch(
732+
"six.moves.builtins.open",
733+
mock_open(
734+
read_data='{"ResourceName": "SageMakerInstance", '
735+
'"ExecutionRoleArn": "arn:aws:iam::369233609183:role/service-role/SageMakerRole-20171129T072388"}'
736+
),
737+
)
738+
@patch("os.path.exists", side_effect=mock_exists(NOTEBOOK_METADATA_FILE, True))
739+
def test_get_caller_identity_arn_from_metadata_file_with_no_domain_id(boto_session):
740+
sess = Session(boto_session)
741+
expected_role = "arn:aws:iam::369233609183:role/service-role/SageMakerRole-20171129T072388"
742+
743+
actual = sess.get_caller_identity_arn()
744+
745+
assert actual == expected_role
746+
# Should not call describe_notebook_instance since ExecutionRoleArn is available
747+
sess.sagemaker_client.describe_notebook_instance.assert_not_called()
748+
749+
731750
@patch(
732751
"six.moves.builtins.open",
733752
mock_open(read_data='{"ResourceName": "SageMakerInstance"}'),

0 commit comments

Comments
 (0)