# Robotics account SM launcher.
# Usage: bash custom/sagemaker/run_sm.sh <INSTANCE_COUNT> <EXPERIMENT> <NAME> <QUEUE> [BUILD_TYPE] [VERSION] [PRIORITY] [OVERRIDES...]
# Requires env vars: SM_USER, TRI_PROJECT, TRI_OWNER_EMAIL (set in your exps/*.sh script)
# Optional env var: SM_ACCOUNT (default: robotics-new). Set to robotics-old to launch via the
# legacy 124224456861 path while IT finishes onboarding the new 385697366450 path.
# 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

SM_ACCOUNT=${SM_ACCOUNT:-robotics-new}

INSTANCE_COUNT=$1
EXPERIMENT=$2
NAME=$3

QUEUE_NAME=${4:-cv-wfm}
BUILD_TYPE=${5:-full}
VERSION=${6:-271}

ENTRY_POINT=scripts/train.py
CONFIG=cosmos_predict2/configs/base/config.py

REGION=us-west-2
PRIORITY=${7:-${SM_PRIORITY:-20}}  # positional arg 7, else SM_PRIORITY env, else 20

# Edit this to fastfile when needed, default is s3.
INPUT_SOURCE=s3

shift 7 2>/dev/null || shift $#
OVERRIDES=("$@")

AWS_DEFAULT_REGION=${REGION}                        \
    python3 custom/sagemaker/launch_sm.py           \
    --account=${SM_ACCOUNT}                         \
    --base-job-name=${SM_USER}-any4d                \
    --entry_point=${ENTRY_POINT}                    \
    --config=${CONFIG}                              \
    --experiment=${EXPERIMENT}                      \
    --instance-count=${INSTANCE_COUNT}              \
    --region=${REGION}                              \
    --input-source=${INPUT_SOURCE}                  \
    --priority=${PRIORITY}                          \
    --queue=${QUEUE_NAME}                           \
    --name=${NAME}                                  \
    --version=${VERSION}                            \
    --build-type=${BUILD_TYPE}                      \
    "${OVERRIDES[@]}"