Skip to content

Commit 077e5f3

Browse files
committed
Create role if needed in get_execution_role
1 parent cebfd71 commit 077e5f3

File tree

1 file changed

+40
-2
lines changed

1 file changed

+40
-2
lines changed

src/sagemaker/session.py

+40-2
Original file line numberDiff line numberDiff line change
@@ -6780,13 +6780,16 @@ def production_variant(
67806780
return production_variant_configuration
67816781

67826782

6783-
def get_execution_role(sagemaker_session=None):
6783+
def get_execution_role(sagemaker_session=None, use_default=False):
67846784
"""Return the role ARN whose credentials are used to call the API.
67856785
67866786
Throws an exception if role doesn't exist.
67876787
67886788
Args:
6789-
sagemaker_session(Session): Current sagemaker session
6789+
sagemaker_session(Session): Current sagemaker session.
6790+
use_default(bool): Use a default role if `get_caller_identity_arn does not
6791+
return a correct role. This default role will be created if needed.
6792+
Defaults to ``False``.
67906793
67916794
Returns:
67926795
(str): The role ARN
@@ -6797,6 +6800,41 @@ def get_execution_role(sagemaker_session=None):
67976800

67986801
if ":role/" in arn:
67996802
return arn
6803+
6804+
if use_default:
6805+
default_role_name = "AmazonSageMaker-DefaultRole"
6806+
6807+
LOGGER.warning("Using default role: %s", default_role_name)
6808+
6809+
boto3_session = sagemaker_session.boto_session
6810+
permissions_policy = json.dumps(
6811+
{
6812+
"Version": "2012-10-17",
6813+
"Statement": [
6814+
{
6815+
"Effect": "Allow",
6816+
"Principal": {"Service": ["sagemaker.amazonaws.com"]},
6817+
"Action": "sts:AssumeRole",
6818+
}
6819+
],
6820+
}
6821+
)
6822+
iam_client = boto3_session.client("iam")
6823+
try:
6824+
iam_client.get_role(RoleName=default_role_name)
6825+
except iam_client.exceptions.NoSuchEntityException:
6826+
iam_client.create_role(
6827+
RoleName=default_role_name, AssumeRolePolicyDocument=str(permissions_policy)
6828+
)
6829+
6830+
LOGGER.warning("Created new sagemaker execution role: %s", default_role_name)
6831+
6832+
iam_client.attach_role_policy(
6833+
PolicyArn="arn:aws:iam::aws:policy/AmazonSageMakerFullAccess",
6834+
RoleName=default_role_name,
6835+
)
6836+
return iam_client.get_role(RoleName=default_role_name)["Role"]["Arn"]
6837+
68006838
message = (
68016839
"The current AWS identity is not a role: {}, therefore it cannot be used as a "
68026840
"SageMaker execution role"

0 commit comments

Comments
 (0)