def METHOD_NAME(
    sagemaker_session, role_arn=None, alias=KEY_ALIAS, sagemaker_role="SageMakerRole"
):
    kms_client = sagemaker_session.boto_session.client("kms")
    kms_key_arn = _get_kms_key_arn(kms_client, alias)
    region = sagemaker_session.boto_region_name
    sts_client = sagemaker_session.boto_session.client(
        "sts", region_name=region, endpoint_url=utils.sts_regional_endpoint(region)
    )
    account_id = sts_client.get_caller_identity()["Account"]
    if kms_key_arn is None:
        return _create_kms_key(kms_client, account_id, region, role_arn, sagemaker_role, alias)
    if role_arn:
        _add_role_to_policy(kms_client, account_id, role_arn, region, alias, sagemaker_role)
    return kms_key_arn