set -e

HOST_ADDR=0.0.0.0
CONTROLER_PORT=28786
WORKER_BASE_PORT=30080

echo PYTHON_EXECUTABLE=$(which python3)
PYTHON_EXECUTABLE=$(which python3)

MODEL_BASE=/home/shared_data
CUDA_DEVICE_BASE=0
POLICY_MODEL_NAME=huggingface/llama/Meta-Llama-3-70B-Instruct-GPTQ-Int4
# VALUE_MODEL_NAME=/huggingface/peiyi9979/math-shepherd-mistral-7b-prm
VALUE_MODEL_NAME=/huggingface/Math-psa/checkpoint-2127/
APPROXIMATION_MODEL_NAME=Llama-3-8B-Instruct-GPTQ-4-Bit
MODEL_PATH=$MODEL_BASE/$POLICY_MODEL_NAME
VALUE_MODEL_PATH=$MODEL_BASE/$VALUE_MODEL_NAME
APPROXIMATION_MODEL_PATH=$MODEL_BASE/$APPROXIMATION_MODEL_NAME

LOGDIR=logs_fastchat

tmux start-server
tmux new-session -s FastChat321 -n controller -d
tmux send-keys "export LOGDIR=${LOGDIR}" Enter
tmux send-keys "$PYTHON_EXECUTABLE -m fastchat.serve.controller --port ${CONTROLER_PORT} --host $HOST_ADDR" Enter

NUM_LM_WORKER=1
NUM_RM_WORKER=1
NUM_AX_WORKER=1

echo "Wait 10 seconds ..."
sleep 5

# echo "Starting workers"
# for i in $(seq 0 $((NUM_LM_WORKER-1)))
# do
#   WORKER_PORT=$((WORKER_BASE_PORT+i))
#   tmux new-window -n policy_worker_$i
#   tmux send-keys "export LOGDIR=${LOGDIR}" Enter
#   # tmux send-keys "CUDA_VISIBLE_DEVICES=$((i+CUDA_DEVICE_BASE)) $PYTHON_EXECUTABLE -m reason.llm_service.workers.vllm_worker --model-path $MODEL_PATH --controller-address http://$HOST_ADDR:$CONTROLER_PORT --host $HOST_ADDR --port $WORKER_PORT --worker-address http://$HOST_ADDR:$WORKER_PORT --dtype bfloat16 --swap-space 32" Enter
#   tmux send-keys "CUDA_VISIBLE_DEVICES=3 $PYTHON_EXECUTABLE -m reason.llm_service.workers.vllm_worker --model-path $MODEL_PATH --controller-address http://$HOST_ADDR:$CONTROLER_PORT --host $HOST_ADDR --port $WORKER_PORT --worker-address http://$HOST_ADDR:$WORKER_PORT --dtype bfloat16 --swap-space 32 --seed 50" Enter
# done


# start value service
for i in $(seq 0 $((NUM_RM_WORKER-1)))
do
  WORKER_PORT=$((i+WORKER_BASE_PORT+NUM_LM_WORKER))
  tmux new-window -n value_worker
  tmux send-keys "export LOGDIR=${LOGDIR}" Enter
  # tmux send-keys "CUDA_VISIBLE_DEVICES=$((i+NUM_LM_WORKER+CUDA_DEVICE_BASE)) $PYTHON_EXECUTABLE -m reason.llm_service.workers.reward_model_worker --model-path $VALUE_MODEL_PATH --controller-address http://$HOST_ADDR:$CONTROLER_PORT --host $HOST_ADDR --port $WORKER_PORT --worker-address http://$HOST_ADDR:$WORKER_PORT" Enter
  tmux send-keys "CUDA_VISIBLE_DEVICES=1 $PYTHON_EXECUTABLE -m reason.llm_service.workers.reward_model_worker --model-path $VALUE_MODEL_PATH --controller-address http://$HOST_ADDR:$CONTROLER_PORT --host $HOST_ADDR --port $WORKER_PORT --worker-address http://$HOST_ADDR:$WORKER_PORT --seed 50" Enter
done


# start approximation service
for i in $(seq 0 $((NUM_AX_WORKER-1)))
do
  WORKER_PORT=$((WORKER_BASE_PORT+i+NUM_LM_WORKER+NUM_RM_WORKER))
  tmux new-window -n approximation_worker_$i
  tmux send-keys "export LOGDIR=${LOGDIR}" Enter
  # tmux send-keys "CUDA_VISIBLE_DEVICES=$((i+CUDA_DEVICE_BASE)) $PYTHON_EXECUTABLE -m reason.llm_service.workers.vllm_worker --model-path $MODEL_PATH --controller-address http://$HOST_ADDR:$CONTROLER_PORT --host $HOST_ADDR --port $WORKER_PORT --worker-address http://$HOST_ADDR:$WORKER_PORT --dtype bfloat16 --swap-space 32" Enter
  tmux send-keys "CUDA_VISIBLE_DEVICES=1 $PYTHON_EXECUTABLE -m reason.llm_service.workers.vllm_worker --model-path $APPROXIMATION_MODEL_PATH --controller-address http://$HOST_ADDR:$CONTROLER_PORT --host $HOST_ADDR --port $WORKER_PORT --worker-address http://$HOST_ADDR:$WORKER_PORT --dtype bfloat16 --swap-space 32 --seed 50 --gpu_memory_utilization 0.7" Enter
done


# echo "Starting workers"
# for i in $(seq 0 $((NUM_LM_WORKER-1)))
# do
#   WORKER_PORT=$((WORKER_BASE_PORT+i))
#   tmux new-window -n policy_worker_$i
#   tmux send-keys "export LOGDIR=${LOGDIR}" Enter
#   # tmux send-keys "CUDA_VISIBLE_DEVICES=$((i+CUDA_DEVICE_BASE)) $PYTHON_EXECUTABLE -m reason.llm_service.workers.vllm_worker --model-path $MODEL_PATH --controller-address http://$HOST_ADDR:$CONTROLER_PORT --host $HOST_ADDR --port $WORKER_PORT --worker-address http://$HOST_ADDR:$WORKER_PORT --dtype bfloat16 --swap-space 32" Enter
#   tmux send-keys "CUDA_VISIBLE_DEVICES=3 $PYTHON_EXECUTABLE -m reason.llm_service.workers.vllm_worker --model-path $MODEL_PATH --speculative-model-path $APPROXIMATION_MODEL_PATH --controller-address http://$HOST_ADDR:$CONTROLER_PORT --host $HOST_ADDR --port $WORKER_PORT --worker-address http://$HOST_ADDR:$WORKER_PORT --dtype bfloat16 --swap-space 32 --seed 50" Enter
# done