#!/bin/bash
#SBATCH --job-name={{ job_name }}
#SBATCH --account={{ cluster_account }}
#SBATCH --time={{ time_limit }}
#SBATCH --exclusive
#SBATCH --nodes={{ nodes }}
#SBATCH --partition={{ partition }}
#SBATCH --output={{ home }}/.clserve/logs/%j/log.out
#SBATCH --error={{ home }}/.clserve/logs/%j/log.err


DP_SIZE={{ dp_size }}
TP_SIZE={{ tp_size }}
EP_SIZE={{ ep_size }}
CUDA_GRAPH_MAX_BS={{ cuda_graph_max_bs }}
GRAMMAR_BACKEND={{ grammar_backend }}
MODEL_PATH={{ model_path }}
NODES={{ nodes }}
NODES_PER_WORKER={{ nodes_per_worker }}
WORKERS={{ workers }}
LOG_DIR="${HOME}/.clserve/logs/${SLURM_JOB_ID}"
ENVIRONMENT={{ environment }}
USE_ROUTER={{ use_router | lower }}
ROUTER_ENVIRONMENT={{ router_environment }}
ROUTER_POLICY={{ router_policy }}
REASONING_PARSER="{{ reasoning_parser }}"
TOOL_CALL_PARSER="{{ tool_call_parser }}"
NUM_GPUS_PER_WORKER={{ num_gpus_per_worker }}
PROCESSES_PER_NODE=$((4 / NUM_GPUS_PER_WORKER))

{% raw %}

mkdir -p "${LOG_DIR}"

# Write metadata for clserve status command
cat > "${LOG_DIR}/metadata.txt" << METADATA_EOF
MODEL_PATH=${MODEL_PATH}
WORKERS=${WORKERS}
NODES_PER_WORKER=${NODES_PER_WORKER}
TP_SIZE=${TP_SIZE}
DP_SIZE=${DP_SIZE}
EP_SIZE=${EP_SIZE}
USE_ROUTER=${USE_ROUTER}
NUM_GPUS_PER_WORKER=${NUM_GPUS_PER_WORKER}
METADATA_EOF

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

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

# Collect all worker head IPs and URLs
worker_head_ips=()
worker_urls=()
node_ips=()

# First collect all node IPs
for node_idx in $(seq 0 $((NODES - 1))); do
    node=${nodes[$node_idx]}
    node_ip=$(srun --nodes=1 --ntasks=1 -w ${node} hostname -i)
    if [ -z "$node_ip" ]; then
        echo "Error: Could not retrieve IP address for node ${node}"
        exit 1
    fi
    node_ips+=("$node_ip")
    echo "Node $node_idx (${node}) IP: $node_ip"
done

# Build worker URLs based on num GPUs per worker
for worker_id in $(seq 0 $((WORKERS - 1))); do
    start_node=$((worker_id * NODES_PER_WORKER))
    worker_host_ip=${node_ips[$start_node]}
    worker_head_ips+=("$worker_host_ip")

    # For each node in this worker
    for node_offset in $(seq 0 $((NODES_PER_WORKER - 1))); do
        node_idx=$((start_node + node_offset))
        node_ip=${node_ips[$node_idx]}

        # For each process on this node
        for proc_id in $(seq 0 $((PROCESSES_PER_NODE - 1))); do
            port=$((5000 + proc_id))
            worker_urls+=("http://${node_ip}:${port}")
        done
    done
done

echo "All worker head IPs: $(printf '"%s"' "${worker_head_ips[@]}" | paste -sd ',' - | sed 's/,/, /g')"
echo "All worker URLs: $(printf '"%s"' "${worker_urls[@]}" | paste -sd ',' - | sed 's/,/, /g')"

# Launch workers
for worker_id in $(seq 0 $((WORKERS - 1))); do
    echo "Launching worker $worker_id"

    # Calculate node range for this worker
    start_node=$((worker_id * NODES_PER_WORKER))
    end_node=$((start_node + NODES_PER_WORKER - 1))

    # Get worker nodes
    worker_nodes=()
    for node_idx in $(seq $start_node $end_node); do
        worker_nodes+=("${nodes[$node_idx]}")
    done

    echo "Worker $worker_id nodes: ${worker_nodes[*]}"
    worker_host_ip=${worker_head_ips[$worker_id]}

    # Launch tasks for this worker
    if [ $NUM_GPUS_PER_WORKER -eq 4 ]; then
        # Original behavior: one process per node using all GPUs
        for local_rank in $(seq 0 $((NODES_PER_WORKER - 1))); do
            global_node_idx=$((start_node + local_rank))
            node=${nodes[$global_node_idx]}

            srun --nodes=1 --ntasks=1 --nodelist=$node --container-writable --environment=$ENVIRONMENT --output=${LOG_DIR}/worker${worker_id}_node${local_rank}_${node}.out --error=${LOG_DIR}/worker${worker_id}_node${local_rank}_${node}.err \
                bash --norc --noprofile -c "\
set -ex

export no_proxy=\"0.0.0.0,\$no_proxy\"
export NO_PROXY=\"0.0.0.0,\$NO_PROXY\"

export NCCL_DEBUG=WARN
export CUDA_DEVICE_ORDER=PCI_BUS_ID
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
export SGL_ENABLE_JIT_DEEPGEMM=false

python -m sglang.launch_server --model-path ${MODEL_PATH} --host 0.0.0.0 --port 5000 --dist-init-addr ${worker_host_ip}:5757 --nnodes ${NODES_PER_WORKER} --node-rank ${local_rank} --tp-size ${TP_SIZE} --dp-size ${DP_SIZE} --ep-size ${EP_SIZE} --cuda-graph-max-bs ${CUDA_GRAPH_MAX_BS} --grammar-backend ${GRAMMAR_BACKEND} --decode-log-interval 1 --trust-remote-code ${REASONING_PARSER} ${TOOL_CALL_PARSER}" &

            echo "Launched worker $worker_id node $local_rank on $node"
        done
    else
        # Multiple processes per node, each with a subset of GPUs
        for local_rank in $(seq 0 $((NODES_PER_WORKER - 1))); do
            global_node_idx=$((start_node + local_rank))
            node=${nodes[$global_node_idx]}

            for proc_id in $(seq 0 $((PROCESSES_PER_NODE - 1))); do
                port=$((5000 + proc_id))

                # Calculate GPU binding based on num_gpus_per_worker
                if [ $NUM_GPUS_PER_WORKER -eq 1 ]; then
                    gpu_bind="map_gpu:${proc_id}"
                    cuda_visible_devices="${proc_id}"
                else
                    # NUM_GPUS_PER_WORKER = 2
                    first_gpu=$((proc_id * 2))
                    second_gpu=$((first_gpu + 1))
                    gpu_bind="map_gpu:${first_gpu},${second_gpu}"
                    cuda_visible_devices="${first_gpu},${second_gpu}"
                fi

                srun --nodes=1 --ntasks=1 --nodelist=$node --container-writable --environment=$ENVIRONMENT --kill-on-bad-exit=1 --gpus-per-task=${NUM_GPUS_PER_WORKER} --cpus-per-task=50 --gpu-bind=${gpu_bind} --overlap --output=${LOG_DIR}/worker${worker_id}_node${local_rank}_proc${proc_id}_${node}.out --error=${LOG_DIR}/worker${worker_id}_node${local_rank}_proc${proc_id}_${node}.err \
                    bash --norc --noprofile -c "\
set -ex

export no_proxy=\"0.0.0.0,\$no_proxy\"
export NO_PROXY=\"0.0.0.0,\$NO_PROXY\"

export NCCL_DEBUG=WARN
export CUDA_DEVICE_ORDER=PCI_BUS_ID
export CUDA_VISIBLE_DEVICES=${cuda_visible_devices}
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
export SGL_ENABLE_JIT_DEEPGEMM=false

python -m sglang.launch_server --model-path ${MODEL_PATH} --host 0.0.0.0 --port ${port} --tp-size ${NUM_GPUS_PER_WORKER} --dp-size ${DP_SIZE} --cuda-graph-max-bs ${CUDA_GRAPH_MAX_BS} --grammar-backend ${GRAMMAR_BACKEND} --decode-log-interval 1 --trust-remote-code ${REASONING_PARSER} ${TOOL_CALL_PARSER}" &

                echo "Launched process $proc_id on node $node with GPUs ${cuda_visible_devices} on port ${port}"
            done
        done
    fi
done

# Launch router if enabled and there are multiple worker URLs (multiple workers or multiple processes per node)
total_processes=${#worker_urls[@]}
router_url=""
if [ "$USE_ROUTER" = "true" ] && [ $total_processes -gt 1 ]; then
    router_host_node=${nodes[0]}
    router_host_ip=${node_ips[0]}
    router_url="http://${router_host_ip}:30000"

    # Build worker URLs string for router
    worker_urls_str=""
    for url in "${worker_urls[@]}"; do
        worker_urls_str="$worker_urls_str $url"
    done

    echo "Starting router on ${router_host_node} (${router_host_ip})"
    echo "Router worker URLs:${worker_urls_str}"

    srun --nodes=1 --ntasks=1 --nodelist=$router_host_node --container-writable --environment=$ROUTER_ENVIRONMENT --cpus-per-task=50 --overlap --output=${LOG_DIR}/router_${router_host_node}.out --error=${LOG_DIR}/router_${router_host_node}.err \
        bash --norc --noprofile -c "\
set -ex

export no_proxy=\"0.0.0.0,\$no_proxy\"
export NO_PROXY=\"0.0.0.0,\$NO_PROXY\"

python -m sglang_router.launch_router --host 0.0.0.0 --port 30000 --worker-urls${worker_urls_str} --model-path ${MODEL_PATH} --policy ${ROUTER_POLICY} ${REASONING_PARSER} ${TOOL_CALL_PARSER} --worker-startup-timeout-secs 1200" &

    echo ""
    echo "Router URL: ${router_url}"
fi

# Write endpoint info for clserve status command
if [ -n "$router_url" ]; then
    echo "ENDPOINT_URL=${router_url}" >> "${LOG_DIR}/metadata.txt"
else
    echo "ENDPOINT_URL=${worker_urls[0]}" >> "${LOG_DIR}/metadata.txt"
fi

echo ""
echo "============================================"
echo "CLSERVE DEPLOYMENT INFO"
echo "============================================"
echo "Job ID: $SLURM_JOB_ID"
echo "Model: ${MODEL_PATH}"
if [ -n "$router_url" ]; then
    echo "Endpoint URL: ${router_url}"
else
    echo "Endpoint URL: ${worker_urls[0]}"
fi
echo "============================================"

echo ""
echo "To connect to the host node:"
echo "srun --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"
echo "Or use: clserve stop $SLURM_JOB_ID"

wait
echo "Script finished at $(date)"

{% endraw %}
