Skip to content

Commit 5806f8b

Browse files
jgoyani1root
authored andcommitted
fix: update get_execution_role_arn from metadata file if present (aws#4388)
1 parent 55df3fc commit 5806f8b

File tree

2 files changed

+25
-1
lines changed

2 files changed

+25
-1
lines changed

src/sagemaker/session.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5412,14 +5412,19 @@ def get_caller_identity_arn(self):
54125412
domain_id = metadata.get("DomainId")
54135413
user_profile_name = metadata.get("UserProfileName")
54145414
space_name = metadata.get("SpaceName")
5415+
execution_role_arn = metadata.get("ExecutionRoleArn")
54155416
try:
54165417
if domain_id is None:
54175418
instance_desc = self.sagemaker_client.describe_notebook_instance(
54185419
NotebookInstanceName=instance_name
54195420
)
54205421
return instance_desc["RoleArn"]
54215422

5422-
# In Space app, find execution role from DefaultSpaceSettings on domain level
5423+
# find execution role from the metadata file if present
5424+
if execution_role_arn is not None:
5425+
return execution_role_arn
5426+
5427+
# In Shared Space app, find execution role from DefaultSpaceSettings on domain level
54235428
if space_name is not None:
54245429
domain_desc = self.sagemaker_client.describe_domain(DomainId=domain_id)
54255430
return domain_desc["DefaultSpaceSettings"]["ExecutionRole"]

tests/unit/test_session.py

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -698,6 +698,25 @@ def test_fallback_to_domain_if_role_unavailable_in_user_settings(boto_session):
698698
sess.sagemaker_client.describe_domain.assert_called_once_with(DomainId="d-kbnw5yk6tg8j")
699699

700700

701+
@patch(
702+
"six.moves.builtins.open",
703+
mock_open(
704+
read_data='{"ResourceName": "SageMakerInstance", '
705+
'"DomainId": "d-kbnw5yk6tg8j", '
706+
'"ExecutionRoleArn": "arn:aws:iam::369233609183:role/service-role/SageMakerRole-20171129T072388", '
707+
'"SpaceName": "space_name"}'
708+
),
709+
)
710+
@patch("os.path.exists", side_effect=mock_exists(NOTEBOOK_METADATA_FILE, True))
711+
def test_get_caller_identity_arn_from_metadata_file_for_space(boto_session):
712+
sess = Session(boto_session)
713+
expected_role = "arn:aws:iam::369233609183:role/service-role/SageMakerRole-20171129T072388"
714+
715+
actual = sess.get_caller_identity_arn()
716+
717+
assert actual == expected_role
718+
719+
701720
@patch(
702721
"six.moves.builtins.open",
703722
mock_open(

0 commit comments

Comments
 (0)