ACCOUNT_ID=123456789012                         # your AWS account ID
REGION=us-east-1                                # your region
DOMAINID=d-12345678abcd                         # your SageMaker Studio domain ID.
ROLE_ARN='arn:aws:iam::<account-id>:role/service-role/AmazonSageMaker-ExecutionRole-123456789'                                             # your SageMaker execution role to create image
IMAGE_NAME=r-4.1.3-rstudio                      # the image name to be shown in SageMaker Studio
REPO=rstudio-custom                             # ECR repository name
TAG=$IMAGE_NAME                                 # image tag 
DISPLAY_NAME=RSession-r-4.1.3-rstudio           # Display name in RStudio
IMAGE_CONFIG_NAME=r-4-1-3-rstudio               # app image configuration name


# login to own ECR and public DLC's ECR
aws --region ${REGION} ecr get-login-password | docker login --username AWS --password-stdin ${ACCOUNT_ID}.dkr.ecr.${REGION}.amazonaws.com

# create repo
aws --region ${REGION} ecr create-repository --repository-name ${REPO}

# build docker image and push to ECR
docker build . -t ${REPO}:${TAG} -t ${ACCOUNT_ID}.dkr.ecr.${REGION}.amazonaws.com/${REPO}:${TAG}

docker push ${ACCOUNT_ID}.dkr.ecr.${REGION}.amazonaws.com/${REPO}:${TAG}

# Will throw an error gracefully if $IMAGE_NAME exists
aws --region ${REGION} sagemaker create-image \
    --image-name ${IMAGE_NAME} \
    --display-name ${DISPLAY_NAME} \
    --role-arn ${ROLE_ARN}

aws --region ${REGION} sagemaker create-image-version \
    --image-name ${IMAGE_NAME} \
    --base-image "${ACCOUNT_ID}.dkr.ecr.${REGION}.amazonaws.com/${REPO}:${TAG}"
    
# Verify the image-version is created successfully. Do NOT proceed if image-version is in CREATE_FAILED state or in any other state apart from CREATED.
aws --region ${REGION} sagemaker describe-image-version \
    --image-name ${IMAGE_NAME}

## Create AppImageConfig for this image
aws --region ${REGION} sagemaker create-app-image-config \
    --app-image-config-name ${IMAGE_CONFIG_NAME}

## Update the Domain, providing the Image and AppImageConfig
aws --region ${REGION} sagemaker update-domain \
    --domain-id ${DOMAINID} \
    --cli-input-json file://default-user-settings.json