if [ $# -ne 1 ]; then
    echo "Usage: $0 <ray_jobid>"
    echo "Example: $0 160852"
    exit 1
fi

RAY_JOBID=$1
HEAD_NODE=$(scontrol show hostnames $(squeue -j ${RAY_JOBID} -h -o %R) | head -n 1)
HEAD_NODE_IP=$(dig +short ${HEAD_NODE})

module reset
module load brics/nccl
module list

source .venv/bin/activate

MODEL_NAME="openai/gpt-oss-120b"

echo STARTING VLLM BENCH AGAINST SERVER AT $HEAD_NODE_IP

srun \
    --overlap \
    --jobid=${RAY_JOBID} \
    --nodelist=${HEAD_NODE} \
    --nodes=1 \
    --gpus=1 \
    --ntasks-per-node=1 \
    vllm bench serve \
    --model $MODEL_NAME \
    --host $HEAD_NODE_IP \
    --port 8000 \
    --trust-remote-code \
    --dataset-name random \
    --random-input-len 1024 \
    --random-output-len 1024 \
    --ignore-eos \
    --max-concurrency 512 \
    --num-prompts 2560