#!/bin/bash
#SBATCH --job-name={{ job_name }}
#SBATCH --nodes={{ total_nodes }}
#SBATCH --ntasks={{ total_nodes }}
#SBATCH --ntasks-per-node=1
#SBATCH --account={{ account }}
#SBATCH --time={{ time_limit }}
#SBATCH --output=logs/%j/log.out
#SBATCH --error=logs/%j/log.err
#SBATCH --partition={{ partition }}

# Constants
PREFILL_NODES={{ prefill_nodes }}
DECODE_NODES={{ decode_nodes }}
TOTAL_NODES=$((PREFILL_NODES + DECODE_NODES))
GPUS_PER_NODE={{ gpus_per_node }}
LOG_DIR="${SLURM_SUBMIT_DIR}/logs/${SLURM_JOB_ID}/"
SCRIPT_DIR="${SLURM_SUBMIT_DIR}/scripts"
OUTPUT_DIR="${SLURM_SUBMIT_DIR}/outputs"
MODEL_DIR="{{ model_dir }}"
CONFIG_DIR="{{ config_dir }}"
CONTAINER_IMAGE="{{ container_image }}"
NETWORK_INTERFACE="{{ network_interface }}"
GPU_TYPE="{{ gpu_type | default('h100') }}"
USE_SGLANG_COMMANDS="{{ use_sglang_commands | default(false) }}"

{% raw %}

mkdir -p "${OUTPUT_DIR}" "${LOG_DIR}"

nodes=($(scontrol show hostnames $SLURM_NODELIST))
if [ ${#nodes[@]} -ne $TOTAL_NODES ]; then
    echo "Error: Expected $TOTAL_NODES nodes but got ${#nodes[@]} nodes"
    exit 1
fi

# Print node information
for i in "${!nodes[@]}"; do
    echo "Node $i: ${nodes[$i]}"
done

PREFILL_HOST_IP=$(srun --nodes=1 --ntasks=1 --nodelist=${nodes[0]} ip route get $(getent ahosts ${nodes[0]} | grep STREAM | head -1 | awk '{print $1}') | awk '{for(i=1;i<=NF;i++) if($i=="src") print $(i+1)}')
if [ -z "$PREFILL_HOST_IP" ]; then
    echo "Error: Could not retrieve IP address for prefill host ${nodes[0]} on interface $NETWORK_INTERFACE"
    exit 1
fi
echo "Prefill host IP address: $PREFILL_HOST_IP"

DECODE_HOST_IP=$(srun --nodes=1 --ntasks=1 --nodelist=${nodes[$PREFILL_NODES]} ip route get $(getent ahosts ${nodes[$PREFILL_NODES]} | grep STREAM | head -1 | awk '{print $1}') | awk '{for(i=1;i<=NF;i++) if($i=="src") print $(i+1)}')
if [ -z "$DECODE_HOST_IP" ]; then
    echo "Error: Could not retrieve IP address for decode host ${nodes[$PREFILL_NODES]} on interface $NETWORK_INTERFACE"
    exit 1
fi
echo "Decode host IP address: $DECODE_HOST_IP"

# Prepare enroot arguments to pass to srun commands
ENROOT_ARGS="\
    --container-image=${CONTAINER_IMAGE} \
    --no-container-entrypoint \
    --no-container-mount-home \
    --container-mounts=${MODEL_DIR}:/model/,${CONFIG_DIR}:/configs/,${SCRIPT_DIR}:/scripts/,${OUTPUT_DIR}:/outputs/,${LOG_DIR}:/logs/ \
"

# Build common worker arguments
WORKER_ARGS="--gpu_type ${GPU_TYPE} --gpus_per_node ${GPUS_PER_NODE}"
if [ "$USE_SGLANG_COMMANDS" = "True" ]; then
    WORKER_ARGS="${WORKER_ARGS} --use-sglang-commands"
fi

# Launch prefill tasks on the first PREFILL_NODES nodes
for i in $(seq 0 $((PREFILL_NODES - 1))); do
    node=${nodes[$i]}
    rank=$i
    echo "Launching prefill task on node ${i} (rank ${rank}): $node"

    cmd="srun $ENROOT_ARGS --nodes=1 --ntasks=1 --nodelist=$node --output=${LOG_DIR}/${node}_prefill.out --error=${LOG_DIR}/${node}_prefill.err python /scripts/worker_setup.py --prefill_host_ip ${PREFILL_HOST_IP} --decode_host_ip ${DECODE_HOST_IP} --rank ${rank} --total_nodes ${PREFILL_NODES} --worker_type prefill --gpu_utilization_log /logs/${node}_prefill_gpu_utilization.log ${WORKER_ARGS}"
    echo "$cmd"
    $cmd &
done

# Launch decode tasks on the next DECODE_NODES nodes
for i in $(seq $PREFILL_NODES $((PREFILL_NODES + DECODE_NODES - 1))); do
    node=${nodes[$i]}
    rank=$((i - PREFILL_NODES))
    echo "Launching decode task on node ${i} (rank ${rank}): $node"

    cmd="srun $ENROOT_ARGS --nodes=1 --ntasks=1 --nodelist=$node --output=${LOG_DIR}/${node}_decode.out --error=${LOG_DIR}/${node}_decode.err python /scripts/worker_setup.py --decode_host_ip ${DECODE_HOST_IP} --prefill_host_ip ${PREFILL_HOST_IP} --rank ${rank} --total_nodes ${DECODE_NODES} --worker_type decode --gpu_utilization_log /logs/${node}_decode_gpu_utilization.log ${WORKER_ARGS}"
    echo "$cmd"
    $cmd &
done

echo ""
echo "To connect to the host prefill node:"
echo "srun $ENROOT_ARGS --jobid $SLURM_JOB_ID -w ${nodes[0]} --overlap --pty bash"

echo ""
echo "Make sure to cancel the job at the end:"
echo "scancel $SLURM_JOB_ID"

# Wait for all tasks to complete
wait
echo "Script finished at $(date)"

{% endraw %}
