fix: call DescribeDomain as fallback in get_execution_role (#2435)

* fix: call DescribeDomain as fallback in get_execution_role

* reformat

Co-authored-by: Wang Napieralski <[email protected]>
Co-authored-by: Ahsan Khan <[email protected]>

Author: 3 people

src/sagemaker/session.py

@@ -3515,7 +3515,11 @@ def get_caller_identity_arn(self):
                 user_profile_desc = self.sagemaker_client.describe_user_profile(
                     DomainId=domain_id, UserProfileName=user_profile_name
                 )
-                return user_profile_desc["UserSettings"]["ExecutionRole"]
+                if user_profile_desc.get("UserSettings") is not None:
+                    return user_profile_desc["UserSettings"]["ExecutionRole"]
+
+                domain_desc = self.sagemaker_client.describe_domain(DomainId=domain_id)
+                return domain_desc["DefaultUserSettings"]["ExecutionRole"]
             except ClientError:
                 LOGGER.debug(
                     "Couldn't call 'describe_notebook_instance' to get the Role "

tests/unit/test_session.py

@@ -340,6 +340,33 @@ def test_get_caller_identity_arn_from_describe_user_profile(boto_session):
     )
 
 
+@patch(
+    "six.moves.builtins.open",
+    mock_open(
+        read_data='{"ResourceName": "SageMakerInstance", '
+        '"DomainId": "d-kbnw5yk6tg8j", '
+        '"UserProfileName": "default-1617915559064"}'
+    ),
+)
+@patch("os.path.exists", side_effect=mock_exists(NOTEBOOK_METADATA_FILE, True))
+def test_get_caller_identity_arn_from_describe_domain(boto_session):
+    sess = Session(boto_session)
+    expected_role = "arn:aws:iam::369233609183:role/service-role/SageMakerRole-20171129T072388"
+    sess.sagemaker_client.describe_user_profile.return_value = {}
+    sess.sagemaker_client.describe_domain.return_value = {
+        "DefaultUserSettings": {"ExecutionRole": expected_role}
+    }
+
+    actual = sess.get_caller_identity_arn()
+
+    assert actual == expected_role
+    sess.sagemaker_client.describe_user_profile.assert_called_once_with(
+        DomainId="d-kbnw5yk6tg8j",
+        UserProfileName="default-1617915559064",
+    )
+    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)