#!/bin/bash
# ============================================================
# Common SLURM/DeepSpeed utilities for multinode training
# Usage: source this file in your training script after
#        defining job-specific variables.
# ============================================================

# --------------------------------------------------
# detect_slurm_account: Auto-detect user's SLURM account
# Returns account name via stdout
# --------------------------------------------------
detect_slurm_account() {
    local account
    account=$(sacctmgr -nP show assoc where user=$(whoami) format=account | head -1)
    if [ -z "$account" ]; then
        echo "ERROR: Could not detect SLURM account for user $(whoami)" >&2
        return 1
    fi
    echo "$account"
}

# --------------------------------------------------
# setup_slurm_nodes: Parse SLURM node info
# Sets: head_node, head_node_ip, slurm_name, idx, total
# --------------------------------------------------
setup_slurm_nodes() {
    idx=$SLURM_ARRAY_TASK_ID
    total=$SLURM_ARRAY_TASK_COUNT

    nodes=( $( scontrol show hostnames $SLURM_JOB_NODELIST ) )
    nodes_array=($nodes)
    head_node=${nodes_array[0]}
    head_node_ip=$(srun --nodes=1 --ntasks=1 -w "$head_node" hostname --ip-address)
    slurm_name=${idx}-of-${total}
}

# --------------------------------------------------
# setup_environment: Export common environment variables
# --------------------------------------------------
setup_environment() {
    export LOGLEVEL=INFO
    export CUDA_HOME=${CUDA_HOME:-/cm/shared/apps/cuda12.4/toolkit/12.4.1}
    export MDM_VL_KEEP_VISION=1
    export TOKENIZERS_PARALLELISM=false
    export DATASET_BACKEND="webdataset_multi_modal"
}

# --------------------------------------------------
# print_training_info: Display training configuration
# Requires: model_name_or_path, tokenizer_name, dataset_path,
#           output_dir, job_name, and any extra_info lines
# --------------------------------------------------
print_training_info() {
    local description="${1:-Training}"
    echo "=========================================="
    echo "${description}"
    echo "=========================================="
    echo "Nodes allocated: $SLURM_JOB_NODELIST"
    echo "Number of nodes: $SLURM_NNODES"
    echo "head_node: $head_node"
    echo "head_node_ip: $head_node_ip"
    echo "slurm_name: $slurm_name"
    echo ""
    echo "Configuration:"
    echo "  Model: ${model_name_or_path}"
    echo "  Tokenizer: ${tokenizer_name}"
    echo "  Dataset: ${dataset_path}"
    echo "  Output: ${output_dir}"
    # Print any extra config lines passed as arguments
    shift
    for line in "$@"; do
        echo "  ${line}"
    done
    echo "=========================================="
}

# --------------------------------------------------
# detect_checkpoint: Auto-detect latest checkpoint in output_dir
# Sets: latest_checkpoint, resume_arg
# --------------------------------------------------
detect_checkpoint() {
    latest_checkpoint=""
    resume_arg=""
    if [ -d "${output_dir}" ]; then
        latest_checkpoint=$(find "${output_dir}" -name "checkpoint-*" -type d | sort -V | tail -1)
        if [ -n "${latest_checkpoint}" ]; then
            echo "Found checkpoint: ${latest_checkpoint}"
            resume_arg="--resume_from_checkpoint ${latest_checkpoint}"
        else
            echo "No checkpoint found in ${output_dir}"
        fi
    else
        echo "Output dir ${output_dir} does not exist, training from scratch"
    fi
}

# --------------------------------------------------
# setup_logging: Create dirs, start GPU monitor, clean hostfiles
# Requires: job_name, output_dir
# Sets: log_dir
# --------------------------------------------------
setup_logging() {
    log_dir=log/${job_name}
    mkdir -p ${output_dir} ${log_dir}

    nvidia-smi --query-gpu=timestamp,memory.used,memory.total,utilization.gpu,pstate \
               --format=csv,noheader,nounits -l 60 &

    rm -rf ${log_dir}/deepspeed-hostfile
    rm -rf ${log_dir}/temp-hostfile
}

# --------------------------------------------------
# run_multinode_deepspeed: Build and execute the srun+deepspeed command
# Arguments: all extra training arguments as a single string
# Requires: head_node, log_dir, deepspeed_args, model_name_or_path,
#           tokenizer_name, trust_remote_code, dataset_path,
#           output_dir, resume_arg, conversation_template, slurm_name
# --------------------------------------------------
run_multinode_deepspeed() {
    local extra_train_args="$1"

    cmd='
export PDSH_RCMD_TYPE=ssh;
export CUDA_HOME=/cm/shared/apps/cuda12.4/toolkit/12.4.1;
export MDM_VL_KEEP_VISION=1;
export TOKENIZERS_PARALLELISM=false;
export PYTHONPATH="'"${PWD}"'/src:'"${PWD}"'/third_party/vila:'"${PWD}"'/examples:${PYTHONPATH:-}";
export DATASET_BACKEND="webdataset_multi_modal";
echo "SLURMD_NODENAME:= " $SLURMD_NODENAME 1>&2;
echo "head_node:= " '"${head_node}"' 1>&2;
if [ "$SLURMD_NODENAME" = '"${head_node}"' ]; then
  echo $SLURM_JOB_NODELIST > '"${log_dir}"'/temp-hostfile;
  read master_node slave_node <<< $(./expand_hostfile.sh '"${log_dir}"');
fi
# Wait for all nodes
sleep 60;
./setup_ssh.sh '"${log_dir}"';
echo "MASTER_NODE:= " $master_node 1>&2;
# Only master node runs DeepSpeed
if [ "$SLURMD_NODENAME" = '"${head_node}"' ]; then
  echo "Running on master node. Executing DeepSpeed..." 1>&2;
  deepspeed --master_addr=${master_node} --hostfile '"${log_dir}"'/deepspeed-hostfile '"${deepspeed_args}"' \
    examples/finetune_qwenvl.py \
    --model_name_or_path '"${model_name_or_path}"' \
    --tokenizer_name '"${tokenizer_name}"' \
    --trust_remote_code '"${trust_remote_code}"' \
    --dataset_path '"${dataset_path}"' \
    --output_dir '"${output_dir}"' \
    '"${resume_arg}"' \
    --conversation_template '"${conversation_template}"' \
    '"${extra_train_args}"'
else
    echo "Not running on master node. Skipping DeepSpeed execution." 1>&2;
fi
'

    echo $cmd
    mkdir -p ${log_dir}/slurm_log/train
    srun -o ${log_dir}/slurm_log/train/slurm_%A_${slurm_name}.log \
         -e ${log_dir}/slurm_log/train/error_%A_${slurm_name}.log \
        bash -c "${cmd}"
}
