#!/usr/bin/env bash
set -euo pipefail

MASTER_ADDR="10.5.1.44"
MASTER_PORT="6002"
NNODES="4"
GPUS_PER_NODE="8"
IFNAME="eth0"

NODE_RANK="0"
MODEL_DIR="/home/bit/data/bd/DeepSeek-V3-bf16"
SERVED_MODEL_NAME="DeepSeek-V3-bf16"
IMAGE="vllm-metax-test:latest"
CONTAINER_NAME="vllm-metax-test"

TENSOR_PARALLEL_SIZE="8"
PIPELINE_PARALLEL_SIZE="4"
DTYPE="bfloat16"
GPU_MEMORY_UTILIZATION="0.85"
MAX_MODEL_LEN="4096"
DISTRIBUTED_TIMEOUT_SECONDS="1800"

TOKENIZER_MODE="auto"

HOST_BIND="0.0.0.0"
PORT_BIND="8000"

print_usage() {
  echo "Usage: $0 [--node-rank R] [--nnodes N] [--master-addr IP --master-port PORT] [--ifname IF] [--model-dir DIR] [--host 0.0.0.0 --port 8000]"
  echo "Examples:"
  echo "  On worker: $0 --node-rank 1"
  echo "  On master: $0 --node-rank 0 --host 0.0.0.0 --port 8000"
}

while [[ $# -gt 0 ]]; do
  case "$1" in
    --node-rank)
      NODE_RANK="${2:-}"; shift 2;;
    --master-addr)
      MASTER_ADDR="${2:-}"; shift 2;;
    --master-port)
      MASTER_PORT="${2:-}"; shift 2;;
    --nnodes)
      NNODES="${2:-}"; shift 2;;
    --gpus-per-node)
      GPUS_PER_NODE="${2:-}"; shift 2;;
    --ifname)
      IFNAME="${2:-}"; shift 2;;
    --model-dir)
      MODEL_DIR="${2:-}"; shift 2;;
    --served-model-name)
      SERVED_MODEL_NAME="${2:-}"; shift 2;;
    --image)
      IMAGE="${2:-}"; shift 2;;
    --container-name)
      CONTAINER_NAME="${2:-}"; shift 2;;
    --tp)
      TENSOR_PARALLEL_SIZE="${2:-}"; shift 2;;
    --pp)
      PIPELINE_PARALLEL_SIZE="${2:-}"; shift 2;;
    --dtype)
      DTYPE="${2:-}"; shift 2;;
    --gpu-mem)
      GPU_MEMORY_UTILIZATION="${2:-}"; shift 2;;
    --max-model-len)
      MAX_MODEL_LEN="${2:-}"; shift 2;;
    --timeout)
      DISTRIBUTED_TIMEOUT_SECONDS="${2:-}"; shift 2;;
    --tokenizer-mode)
      TOKENIZER_MODE="${2:-}"; shift 2;;
    -h|--help)
      print_usage; exit 0;;
    *)
      echo "Unknown argument: $1"
      print_usage
      exit 2;;
  esac
done

if [[ "${NODE_RANK}" == "0" ]]; then
  :
else
  HOST_BIND=""
  PORT_BIND=""
fi

export MASTER_ADDR MASTER_PORT NNODES NODE_RANK
export GLOO_SOCKET_IFNAME="$IFNAME"
export GLOO_PORT_RANGE="50000:65000"
export MCCL_DEBUG=INFO
export MCCL_DEBUG_SUBSYS=INIT,NET
export MCCL_ASYNC_ERROR_HANDLING=1
export MCCL_SOCKET_IFNAME="eth0"
export MCCL_IB_TIMEOUT=25
export MCCL_IB_HCA="tcu_0:1,tcu_1:1"
export MCCL_IB_GID_INDEX=1
export MCCL_NET_PLUGIN=none
export MACA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
export MCCL_IB_DISABLE=0
export MCCL_IB_GDR_LEVEL=1
export MCCL_NET_GDR_LEVEL=1
export MCCL_DMABUF_ENABLE=0
export MCCL_IB_PCI_RELAXED_ORDERING=1
export FLASHINFER_DISABLE_VERSION_CHECK="1"
export VLLM_ALLREDUCE_USE_FLASHINFER="0"

headless_args=()
if [[ "${NODE_RANK}" != "0" ]]; then
  headless_args+=("--headless")
fi

docker rm -f "${CONTAINER_NAME}" 2>/dev/null || true

extra_args=()
if [[ -n "${HOST_BIND}" ]]; then
  extra_args+=("--host" "${HOST_BIND}")
fi
if [[ -n "${PORT_BIND}" ]]; then
  extra_args+=("--port" "${PORT_BIND}")
fi

# 注意：多行命令中间绝对不能加注释，会破坏bash换行解析！
docker run -d \
  --name "${CONTAINER_NAME}" \
  --restart unless-stopped \
  --device=/dev/dri/ \
  --device=/dev/mxcd/ \
  --ipc=host \
  --privileged \
  --network host \
  --ulimit memlock=-1 \
  --ulimit stack=67108864 \
  --cap-add=IPC_LOCK \
  --device-cgroup-rule='c 231:* rmw' \
  --device-cgroup-rule='c 10:122 rmw' \
  -e MASTER_ADDR -e MASTER_PORT -e NNODES -e NODE_RANK \
  -e GLOO_SOCKET_IFNAME -e GLOO_PORT_RANGE \
  -e MCCL_SOCKET_IFNAME -e MCCL_DEBUG -e MCCL_DEBUG_SUBSYS -e MCCL_ASYNC_ERROR_HANDLING \
  -e MCCL_IB_HCA -e MCCL_IB_GID_INDEX -e MCCL_NET_PLUGIN -e MCCL_IB_DISABLE \
  -e MCCL_IB_GDR_LEVEL -e MCCL_NET_GDR_LEVEL \
  -e MCCL_DMABUF_ENABLE -e MCCL_IB_PCI_RELAXED_ORDERING \
  -e MACA_VISIBLE_DEVICES -e MCCL_IB_TIMEOUT \
  -e FLASHINFER_DISABLE_VERSION_CHECK -e VLLM_ALLREDUCE_USE_FLASHINFER \
  -v /dev/infiniband:/dev/infiniband \
  -v /sys/class/infiniband:/sys/class/infiniband:ro \
  -v /sys/class/net:/sys/class/net:ro \
  -v "${MODEL_DIR}:/models/DeepSeek-V3:ro" \
  "${IMAGE}" \
/opt/conda/bin/vllm serve /models/DeepSeek-V3 \
--enforce-eager \
--trust-remote-code \
--tokenizer-mode "${TOKENIZER_MODE}" \
--served-model-name "${SERVED_MODEL_NAME}" \
--distributed-executor-backend mp \
--nnodes "${NNODES}" --node-rank "${NODE_RANK}" \
--master-addr "${MASTER_ADDR}" --master-port "${MASTER_PORT}" \
--tensor-parallel-size "${TENSOR_PARALLEL_SIZE}" --pipeline-parallel-size "${PIPELINE_PARALLEL_SIZE}" \
--dtype "${DTYPE}" \
--gpu-memory-utilization "${GPU_MEMORY_UTILIZATION}" \
--max-model-len "${MAX_MODEL_LEN}" \
--disable-custom-all-reduce \
--distributed-timeout-seconds "${DISTRIBUTED_TIMEOUT_SECONDS}" \
"${headless_args[@]}" \
"${extra_args[@]}"

echo "Started container: ${CONTAINER_NAME}"
echo "Follow logs: docker logs -f ${CONTAINER_NAME}"