# AD2 WFM SageMaker launcher (account 385697366450, ad2-gaia-wfm).
# Queued via AWS Batch onto the cv-wfm-p5en reserved pool (this account requires the Batch queue --
# needs the queue-capable sagemaker; see README step 1).
# Queue defaults to cv-wfm (-> fss-cv-wfm-p5en-48xlarge, renamed from cv-ml 2026-05-26); override with SM_QUEUE=<alias>.
# Priority defaults to 220 (gaia); override with SM_PRIORITY=<n>.
# Usage: bash custom/sagemaker/run_sm_ad2.sh <INSTANCE_COUNT> <EXPERIMENT> <NAME> [BUILD_TYPE] [VERSION] [OVERRIDES...]
# Requires env vars: SM_USER (set in your exps/*.sh script)

unset AWS_ACCESS_KEY_ID
unset AWS_SECRET_ACCESS_KEY
unset AWS_SESSION_TOKEN
unset AWS_PROFILE

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

BUILD_TYPE=${4:-full}
VERSION=${5:-271-2stage}

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

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

AWS_DEFAULT_REGION=us-west-2                        \
    python3 custom/sagemaker/launch_sm.py           \
    --account=ad2-gaia-wfm                          \
    --base-job-name=${SM_USER}-any4d                \
    --entry_point=${ENTRY_POINT}                    \
    --config=${CONFIG}                              \
    --experiment=${EXPERIMENT}                      \
    --instance-count=${INSTANCE_COUNT}              \
    --priority=${SM_PRIORITY:-220}                  \
    --queue=${SM_QUEUE:-cv-wfm}                     \
    --name=${NAME}                                  \
    --version=${VERSION}                            \
    --build-type=${BUILD_TYPE}                      \
    "${OVERRIDES[@]}"