def METHOD_NAME(
    kms_client, account_id, role_arn, region, alias=KEY_ALIAS, sagemaker_role="SageMakerRole"
):
    key_id = _get_kms_key_id(kms_client, alias)
    policy = kms_client.get_key_policy(KeyId=key_id, PolicyName=POLICY_NAME)
    policy = json.loads(policy["Policy"])
    principal = policy["Statement"][0]["Principal"]["AWS"]
    if role_arn not in principal or sagemaker_role not in principal:
        principal = PRINCIPAL_TEMPLATE.format(
            partition=utils._aws_partition(region),
            account_id=account_id,
            role_arn=role_arn,
            sagemaker_role=sagemaker_role,
        )
        kms_client.put_key_policy(
            KeyId=key_id,
            PolicyName=POLICY_NAME,
            Policy=KEY_POLICY.format(id=POLICY_NAME, principal=principal),
        )