fix: update get_execution_role_arn to use role from DefaultSpaceSettings

Author: ZhangHan

src/sagemaker/session.py

@@ -3882,12 +3882,19 @@ def get_caller_identity_arn(self):
                 instance_name = metadata["ResourceName"]
                 domain_id = metadata.get("DomainId")
                 user_profile_name = metadata.get("UserProfileName")
+                space_name = metadata.get("SpaceName")
             try:
                 if domain_id is None:
                     instance_desc = self.sagemaker_client.describe_notebook_instance(
                         NotebookInstanceName=instance_name
                     )
                     return instance_desc["RoleArn"]
+
+                # In Space app, find execution role from DefaultSpaceSettings on domain level
+                if space_name is not None:
+                    domain_desc = self.sagemaker_client.describe_domain(DomainId=domain_id)
+                    return domain_desc["DefaultSpaceSettings"]["ExecutionRole"]
+
                 user_profile_desc = self.sagemaker_client.describe_user_profile(
                     DomainId=domain_id, UserProfileName=user_profile_name
                 )

tests/unit/test_session.py

@@ -397,6 +397,28 @@ def test_fallback_to_domain_if_role_unavailable_in_user_settings(boto_session):
     sess.sagemaker_client.describe_domain.assert_called_once_with(DomainId="d-kbnw5yk6tg8j")
 
 
+@patch(
+    "six.moves.builtins.open",
+    mock_open(
+        read_data='{"ResourceName": "SageMakerInstance", '
+        '"DomainId": "d-kbnw5yk6tg8j", '
+        '"SpaceName": "space_name"}'
+    ),
+)
+@patch("os.path.exists", side_effect=mock_exists(NOTEBOOK_METADATA_FILE, True))
+def test_get_caller_identity_arn_from_describe_domain_for_space(boto_session):
+    sess = Session(boto_session)
+    expected_role = "arn:aws:iam::369233609183:role/service-role/SageMakerRole-20171129T072388"
+    sess.sagemaker_client.describe_domain.return_value = {
+        "DefaultSpaceSettings": {"ExecutionRole": expected_role}
+    }
+
+    actual = sess.get_caller_identity_arn()
+
+    assert actual == expected_role
+    sess.sagemaker_client.describe_domain.assert_called_once_with(DomainId="d-kbnw5yk6tg8j")
+
+
 @patch("six.moves.builtins.open", mock_open(read_data='{"ResourceName": "SageMakerInstance"}'))
 @patch("os.path.exists", side_effect=mock_exists(NOTEBOOK_METADATA_FILE, True))
 @patch("sagemaker.session.sts_regional_endpoint", return_value=STS_ENDPOINT)