dataset=$1
ARGS_FILE=anydata/sagemaker/webdatasetize_config/${dataset}.yaml

INSTANCE_COUNT=1
NAME=$2

INSTANCE_TYPE=${3:-g6e}
QUEUE_NAME=${4:-none}

USER_NAME=${5:?username required, e.g. fzc}
SPLIT_JSON=${6:-}

# User namespace for SageMaker job names, ECR image names, and S3 output paths.
echo "Using SageMaker user namespace: ${USER_NAME}"
ENTRY_POINT=anydata/webdataset/web_sagemaker.py

BUILD_TYPE=${BUILD_TYPE:-full}

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

EXTRA_ARGS=""
if [[ -n "${SPLIT_JSON}" ]]; then
    echo "Using split json: ${SPLIT_JSON}"
    EXTRA_ARGS="--split_json ${SPLIT_JSON}"
else
    echo "No split json provided, processing full dataset."
fi

AWS_DEFAULT_REGION=${REGION}                            \
    python3 anydata/sagemaker/launch_webdatasetize_sm.py            \
    --base-job-name=${USER_NAME}-webdatasetize    \
    --entry_point=${ENTRY_POINT}                        \
    --user=${USER_NAME}                                   \
    --args_file=${ARGS_FILE}                            \
    --instance-count=${INSTANCE_COUNT}              \
    --profile=${PROFILE}                            \
    --region=${REGION}             \
    --arn=${ARN}                                    \
    --s3-remote-sync=${S3_REMOTE_SYNC}              \
    --priority=${PRIORITY}                          \
    --name=${NAME}                                  \
    --version=${VERSION}                            \
    --instance-type=${INSTANCE_TYPE}                \
    --build-type=${BUILD_TYPE}                      \
    --queue=${QUEUE_NAME}       		            \
    ${EXTRA_ARGS}