#!/bin/bash
#SBATCH --ntasks-per-node=1
#SBATCH --no-requeue
#SBATCH --partition=production-cluster

# ----------------- Set up environment -----------------
set -x -e

source /fsx/m4/start-m4-user
conda activate base
conda activate $CONDA_ENV_NAME

pushd $WORKING_DIR

python -c 'import torch; cuda=torch.version.cuda; assert cuda.startswith("11"), f"cuda-11.x is needed for bf16, got {cuda}"'


# export WANDB_MODE=offline
export WANDB_DIR=/fsx/m4/experiments


GIT_PYTHON_GIT_EXECUTABLE=`which git`
export GIT_PYTHON_GIT_EXECUTABLE

export AWS_MAX_ATTEMPTS=20
export PYTHONPATH=$WORKING_DIR:$PYTHONPATH
# ------------------------------------------------------

# ----------------- Define paths -----------------
BASE_S3_PATH="s3://m4-datasets"
BASE_DATA_PATH_PMD="$BASE_S3_PATH/general_pmd_tar_20_04_2023/image"
BASE_DATA_PATH_CM4="$BASE_S3_PATH/webdocs/web_document_dataset_filtered_imgurldedup_nsfwfiltered_urldedup_linededup_finalcleaning_setimgurlsdedup_12_05_2023_tar"
BASE_DATA_PATH_LAOIN_1="$BASE_S3_PATH/laion/laion_2b_en_filtered_dedup_tar_20_04_2023"
BASE_DATA_PATH_LAOIN_2="$BASE_S3_PATH/LAION_data/laion_dataset_filtered_dedup_8_05_2023_tar"
BASE_DATA_PATH_WIKI="$BASE_S3_PATH/enwiki/enwiki_v2_tar_20_04_2023"
SAVE_DIR="/fsx/m4/experiments/local_experiment_dir/$RUN_NAME"
ACCELERATE_CONFIG_FILE="$SAVE_DIR/${SLURM_JOB_ID}_accelerate_config.yaml.autogenerated"
DEEPSPEED_CONFIG_FILE="$SAVE_DIR/${SLURM_JOB_ID}_ds_config.json.autogenerated"
CONFIG_FILE="$TRAINING_CONFIGS_DIR/config.yaml"
# -------------------------------------------------

# ----------------- Create accelerate config -----------------
# Auto-generate the accelerate config
NUM_GPUS=$((NUM_GPUS_PER_NODE*SLURM_NNODES))
MASTER_ADDR=`scontrol show hostnames $SLURM_JOB_NODELIST | head -n 1`
# From https://i.hsfzxjy.site/2021-03-10-obtain-a-random-unused-tcp-port-with-bash/
function unused_port() {
    N=${1:-1}
    comm -23 \
        <(seq "1025" "65535" | sort) \
        <(ss -Htan |
            awk '{print $4}' |
            cut -d':' -f2 |
            sort -u) |
        shuf |
        head -n "$N"
}
MASTER_PORT=$(unused_port)


cat << EOT > $ACCELERATE_CONFIG_FILE
# WARNING: do not edit this file as this is an slurm-auto-generated file
compute_environment: LOCAL_MACHINE
deepspeed_config:
  deepspeed_multinode_launcher: standard
  deepspeed_config_file: $DEEPSPEED_CONFIG_FILE
  zero3_init_flag: true
distributed_type: DEEPSPEED
fsdp_config: {}
machine_rank: 0
main_process_ip: $MASTER_ADDR
main_process_port: $MASTER_PORT
main_training_function: main
num_machines: $SLURM_NNODES
num_processes: $NUM_GPUS
use_cpu: false
EOT
# -------------------------------------------------

# ----------------- Create deepspeed config -----------------
# Auto-generate the DS config
cat << EOT > $DEEPSPEED_CONFIG_FILE
{
    "communication_data_type": "fp32",
    "bf16": {
        "enabled": true
    },
    "fp16": {
        "enabled": false,
        "auto_cast": true,
        "loss_scale": 0.0,
        "initial_scale_power": 32,
        "loss_scale_window": 1000,
        "hysteresis": 2,
        "min_loss_scale": 1
    },
    "zero_optimization": {
        "stage": 3,
        "allgather_partitions": true,
        "allgather_bucket_size": 5e8,
        "overlap_comm": false,
        "reduce_scatter": true,
        "reduce_bucket_size": "auto",
        "contiguous_gradients": true,
        "stage3_gather_16bit_weights_on_model_save": false,
        "offload_optimizer": {
            "device": "none"
        },
        "offload_param": {
            "device": "none"
        }
    },
    "gradient_clipping": "auto",
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "steps_per_print": 2000000
}
EOT
# -------------------------------------------------

# ----------------- Create commands to access data -----------------
# Get PMD shards
pmd_train_subsets=(
"coco"
"conceptual_12m"
"conceptual_captions"
"localized_narratives__ADE20k"
"localized_narratives__coco"
"localized_narratives__flickr30k"
"localized_narratives__openimages"
"red_caps"
# # "sbu_captions"
"visual_genome"
"wit"
"yfcc100m"
)
pmd_validation_subsets=(
"coco"
)


all_pmd_train_shards=()
for subset in ${pmd_train_subsets[@]}
do
    new_train_shards=$(aws s3 ls --recursive $BASE_DATA_PATH_PMD/$subset/train/ | grep "\.tar"| cut -c32- | uniq)
    all_pmd_train_shards+=$(printf -- "pipe:bash ${WORKING_DIR}/experiments/pretraining/vloom/common/webdataset_get_file.sh $BASE_S3_PATH/%s \n" $new_train_shards)
done
all_pmd_val_shards=()
for subset in ${pmd_validation_subsets[@]}
do
    new_val_shards=$(aws s3 ls --recursive $BASE_DATA_PATH_PMD/$subset/validation/ | grep "\.tar"| cut -c32- | uniq)
    all_pmd_val_shards+=$(printf -- "pipe:bash ${WORKING_DIR}/experiments/pretraining/vloom/common/webdataset_get_file.sh $BASE_S3_PATH/%s \n" $new_val_shards)
done

# Get CM4 shards
all_cm4_train_shards=$(aws s3 ls --recursive $BASE_DATA_PATH_CM4/train/ | grep "\.tar"| cut -c32- | uniq)
all_cm4_train_shards=$(printf -- "pipe:bash ${WORKING_DIR}/experiments/pretraining/vloom/common/webdataset_get_file.sh $BASE_S3_PATH/%s \n" $all_cm4_train_shards)
# all_cm4_validation_shards=$(aws s3 ls --recursive $BASE_DATA_PATH_CM4/valid/ | grep "\.tar"| cut -c32- | uniq)
# all_cm4_validation_shards=$(printf -- "pipe:bash ${WORKING_DIR}/experiments/pretraining/vloom/common/webdataset_get_file.sh $BASE_S3_PATH/%s \n" $all_cm4_validation_shards)

# Get Wiki shards
all_wiki_train_shards=$(aws s3 ls --recursive $BASE_DATA_PATH_WIKI/train/ | grep "\.tar"| cut -c32- | uniq)
all_wiki_train_shards=$(printf -- "pipe:bash ${WORKING_DIR}/experiments/pretraining/vloom/common/webdataset_get_file.sh $BASE_S3_PATH/%s \n" $all_wiki_train_shards)
# all_wiki_validation_shards=$(aws s3 ls --recursive $BASE_DATA_PATH_WIKI/valid/ | grep "\.tar"| cut -c32- | uniq)
# all_wiki_validation_shards=$(printf -- "pipe:bash ${WORKING_DIR}/experiments/pretraining/vloom/common/webdataset_get_file.sh $BASE_S3_PATH/%s \n" $all_wiki_validation_shards)

# Get Laion shards
all_laion_train_shards_1=$(aws s3 ls --recursive $BASE_DATA_PATH_LAOIN_1/train/ | grep "\.tar"| cut -c32- | uniq)
all_laion_train_shards_1=$(printf -- "pipe:bash ${WORKING_DIR}/experiments/pretraining/vloom/common/webdataset_get_file.sh $BASE_S3_PATH/%s \n" $all_laion_train_shards_1)
all_laion_train_shards_2=$(aws s3 ls --recursive $BASE_DATA_PATH_LAOIN_2/train/ | grep "\.tar"| cut -c32- | uniq)
all_laion_train_shards_2=$(printf -- "pipe:bash ${WORKING_DIR}/experiments/pretraining/vloom/common/webdataset_get_file.sh $BASE_S3_PATH/%s \n" $all_laion_train_shards_2)

# If we have too many training shards, we get an `Argument list too long` error
# when launching the training. So we save the paths to the training shards in
# a txt file, and open it later in `m4/training/dataset.py`
ALL_CM4_TRAIN_SHARDS_TXT_FILE="$SAVE_DIR/all_cm4_train_shards.txt"
printf "%s\n" "${all_cm4_train_shards[@]}" > $ALL_CM4_TRAIN_SHARDS_TXT_FILE
# ALL_CM4_VAL_SHARDS_TXT_FILE="$SAVE_DIR/all_cm4_val_shards.txt"
# printf "%s\n" "${all_cm4_validation_shards[@]}" > $ALL_CM4_VAL_SHARDS_TXT_FILE
# ALL_WIKI_TRAIN_SHARDS_TXT_FILE="$SAVE_DIR/all_wiki_train_shards.txt"
printf "%s\n" "${all_wiki_train_shards[@]}" >> $ALL_CM4_TRAIN_SHARDS_TXT_FILE
printf "%s\n" "${all_wiki_train_shards[@]}" >> $ALL_CM4_TRAIN_SHARDS_TXT_FILE
printf "%s\n" "${all_wiki_train_shards[@]}" >> $ALL_CM4_TRAIN_SHARDS_TXT_FILE
# ALL_WIKI_VAL_SHARDS_TXT_FILE="$SAVE_DIR/all_wiki_val_shards.txt"
# printf "%s\n" "${all_wiki_validation_shards[@]}" > $ALL_WIKI_VAL_SHARDS_TXT_FILE
ALL_PMD_TRAIN_SHARDS_TXT_FILE="$SAVE_DIR/all_pmd_train_shards.txt"
printf "%s\n" "${all_pmd_train_shards[@]}" > $ALL_PMD_TRAIN_SHARDS_TXT_FILE
printf "%s\n" "${all_pmd_train_shards[@]}" >> $ALL_PMD_TRAIN_SHARDS_TXT_FILE
printf "%s\n" "${all_pmd_train_shards[@]}" >> $ALL_PMD_TRAIN_SHARDS_TXT_FILE
# ALL_PMD_VAL_SHARDS_TXT_FILE="$SAVE_DIR/all_pmd_val_shards.txt"
# printf "%s\n" "${all_pmd_val_shards[@]}" > $ALL_PMD_VAL_SHARDS_TXT_FILE
# ALL_LAION_TRAIN_SHARDS_TXT_FILE="$SAVE_DIR/all_laion_train_shards.txt"
printf "%s\n" "${all_laion_train_shards_1[@]}" >> $ALL_PMD_TRAIN_SHARDS_TXT_FILE
printf "%s\n" "${all_laion_train_shards_2[@]}" >> $ALL_PMD_TRAIN_SHARDS_TXT_FILE
# -------------------------------------------------

# ----------------- Create commands to launch training -----------------
pip freeze > $SAVE_DIR/${SLURM_JOB_ID}_requirements.txt


# Note: it is important to escape `$SLURM_PROCID` since we want the srun on each node to evaluate this variable
export LAUNCHER="accelerate launch \
    --rdzv_conf "rdzv_backend=c10d,rdzv_endpoint=$MASTER_ADDR:$MASTER_PORT" \
    --config_file $ACCELERATE_CONFIG_FILE \
    --main_process_ip $MASTER_ADDR \
    --main_process_port $MASTER_PORT \
    --machine_rank \$SLURM_PROCID \
    --role \$(hostname -s): --tee 3 \
    "
export PROGRAM="m4/training/main.py \
        --config $CONFIG_FILE \
        --pmd.training_datasets_paths $ALL_PMD_TRAIN_SHARDS_TXT_FILE \
        --cm4.training_datasets_paths $ALL_CM4_TRAIN_SHARDS_TXT_FILE \
        --job_id $SLURM_JOB_ID \
        --jz_job_time_sec $JOB_TIME_SEC \
        --save_dir $SAVE_DIR \
    "
export M4_DATA_CMD="rm -rf /scratch/m4data && mkdir -p /scratch/m4data && "

export CMD="$M4_DATA_CMD $LAUNCHER $PROGRAM"


# makes everything very slow
#export CUDA_LAUNCH_BLOCKING=1
        # --laion.training_datasets_paths $ALL_LAION_TRAIN_SHARDS_TXT_FILE \
        # --wiki.training_datasets_paths $ALL_WIKI_TRAIN_SHARDS_TXT_FILE \
# force crashing on nccl issues like hanging broadcast
#export NCCL_ASYNC_ERROR_HANDLING=1

echo $CMD

# srun error handling:
# --wait=60: wait 60 sec after the first task terminates before terminating all remaining tasks
# --kill-on-bad-exit=1: terminate a step if any task exits with a non-zero exit code
SRUN_ARGS=" \
    --wait=60 \
    --kill-on-bad-exit=1 \
    "
# ---------------------------------------------------------
# ----------------- Launch training -----------------

srun $SRUN_ARGS --jobid $SLURM_JOBID bash -c "$CMD" 2>&1 | tee -a $SAVE_DIR/logs/main_log.txt