# TODO: Merge this and run_sagemaker.sh to avoid code duplication 

# NOTE: these unset are important; if not done AWS_PROFILE set in launcher will be ignored
unset AWS_ACCESS_KEY_ID
unset AWS_SECRET_ACCESS_KEY
unset AWS_SESSION_TOKEN
unset AWS_PROFILE

ARGS_FILE=$1
INSTANCE_COUNT=$2
NAME=$3

INSTANCE_TYPE=${4:-g6e}
# INSTANCE_TYPE=g5
# INSTANCE_TYPE=g6e-small

USER_NAME=dianchen
ENTRY_POINT=externals/vidar/scripts/datasets/webdataset/create_multinode_balanced.py

BUILD_TYPE=full
# BUILD_TYPE=update
# BUILD_TYPE=none

PROFILE=west
REGION=us-west-2
ARN=arn:aws:iam::124224456861:role/service-role/SageMaker-SageMakerAllAccess
S3_REMOTE_SYNC=s3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/s3_remote_sync/
PRIORITY=10

AWS_DEFAULT_REGION=${REGION}                            \
    python custom/sagemaker/launch_sagemaker_cache.py   \
    --base-job-name=${USER_NAME}-cosmos-latent-cache    \
    --entry_point=${ENTRY_POINT}                        \
    --user=${USER_NAME}                                 \
    --args_file=${ARGS_FILE}                            \
    --instance-type=${INSTANCE_TYPE}                    \
    --instance-count=${INSTANCE_COUNT}                  \
    --profile=${PROFILE}                                \
    --build-type=${BUILD_TYPE}                          \
    --region=${REGION}                                  \
    --arn=${ARN}                                        \
    --s3-remote-sync=${S3_REMOTE_SYNC}                  \
    --priority=${PRIORITY}                              \
    --name=${NAME}