change: create role if needed in get_execution_role

src/sagemaker/session.py

@@ -6902,13 +6902,16 @@ def production_variant(
     return production_variant_configuration
 
 
-def get_execution_role(sagemaker_session=None):
+def get_execution_role(sagemaker_session=None, use_default=False):
     """Return the role ARN whose credentials are used to call the API.
 
     Throws an exception if role doesn't exist.
 
     Args:
-        sagemaker_session(Session): Current sagemaker session
+        sagemaker_session (Session): Current sagemaker session.
+        use_default (bool): Use a default role if ``get_caller_identity_arn`` does not
+            return a correct role. This default role will be created if needed.
+            Defaults to ``False``.
 
     Returns:
         (str): The role ARN
@@ -6919,6 +6922,41 @@ def get_execution_role(sagemaker_session=None):
 
     if ":role/" in arn:
         return arn
+
+    if use_default:
+        default_role_name = "AmazonSageMaker-DefaultRole"
+
+        LOGGER.warning("Using default role: %s", default_role_name)
+
+        boto3_session = sagemaker_session.boto_session
+        permissions_policy = json.dumps(
+            {
+                "Version": "2012-10-17",
+                "Statement": [
+                    {
+                        "Effect": "Allow",
+                        "Principal": {"Service": ["sagemaker.amazonaws.com"]},
+                        "Action": "sts:AssumeRole",
+                    }
+                ],
+            }
+        )
+        iam_client = boto3_session.client("iam")
+        try:
+            iam_client.get_role(RoleName=default_role_name)
+        except iam_client.exceptions.NoSuchEntityException:
+            iam_client.create_role(
+                RoleName=default_role_name, AssumeRolePolicyDocument=str(permissions_policy)
+            )
+
+            LOGGER.warning("Created new sagemaker execution role: %s", default_role_name)
+
+        iam_client.attach_role_policy(
+            PolicyArn="arn:aws:iam::aws:policy/AmazonSageMakerFullAccess",
+            RoleName=default_role_name,
+        )
+        return iam_client.get_role(RoleName=default_role_name)["Role"]["Arn"]
+
     message = (
         "The current AWS identity is not a role: {}, therefore it cannot be used as a "
         "SageMaker execution role"

tests/unit/test_session.py

@@ -15,6 +15,7 @@
 import copy
 import datetime
 import io
+import json
 import logging
 import os
 
@@ -532,6 +533,67 @@ def test_get_execution_role_throws_exception_if_arn_is_not_role_with_role_in_nam
     assert "The current AWS identity is not a role" in str(error.value)
 
 
+def test_get_execution_role_get_default_role(caplog):
+    session = Mock()
+    session.get_caller_identity_arn.return_value = "arn:aws:iam::369233609183:user/marcos"
+
+    iam_client = Mock()
+    iam_client.get_role.return_value = {"Role": {"Arn": "foo-role"}}
+    boto_session = Mock()
+    boto_session.client.return_value = iam_client
+
+    session.boto_session = boto_session
+    actual = get_execution_role(session, use_default=True)
+
+    iam_client.get_role.assert_called_with(RoleName="AmazonSageMaker-DefaultRole")
+    iam_client.attach_role_policy.assert_called_with(
+        PolicyArn="arn:aws:iam::aws:policy/AmazonSageMakerFullAccess",
+        RoleName="AmazonSageMaker-DefaultRole",
+    )
+    assert "Using default role: AmazonSageMaker-DefaultRole" in caplog.text
+    assert actual == "foo-role"
+
+
+def test_get_execution_role_create_default_role(caplog):
+    session = Mock()
+    session.get_caller_identity_arn.return_value = "arn:aws:iam::369233609183:user/marcos"
+    permissions_policy = json.dumps(
+        {
+            "Version": "2012-10-17",
+            "Statement": [
+                {
+                    "Effect": "Allow",
+                    "Principal": {"Service": ["sagemaker.amazonaws.com"]},
+                    "Action": "sts:AssumeRole",
+                }
+            ],
+        }
+    )
+    iam_client = Mock()
+    iam_client.exceptions.NoSuchEntityException = Exception
+    iam_client.get_role = Mock(side_effect=[Exception(), {"Role": {"Arn": "foo-role"}}])
+
+    boto_session = Mock()
+    boto_session.client.return_value = iam_client
+
+    session.boto_session = boto_session
+
+    actual = get_execution_role(session, use_default=True)
+
+    iam_client.create_role.assert_called_with(
+        RoleName="AmazonSageMaker-DefaultRole", AssumeRolePolicyDocument=str(permissions_policy)
+    )
+
+    iam_client.attach_role_policy.assert_called_with(
+        PolicyArn="arn:aws:iam::aws:policy/AmazonSageMakerFullAccess",
+        RoleName="AmazonSageMaker-DefaultRole",
+    )
+
+    assert "Created new sagemaker execution role: AmazonSageMaker-DefaultRole" in caplog.text
+
+    assert actual == "foo-role"
+
+
 @patch(
     "six.moves.builtins.open",
     mock_open(read_data='{"ResourceName": "SageMakerInstance"}'),