models/Bidirectional-Parallel-Refresh/start_vllm_server.sh

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

SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
MODEL_PATH="${SCRIPT_DIR}"
SERVED_MODEL_NAME="dmtd-qwen3-4b"
DEVICE=""
PORT=""
MAX_MODEL_LEN=40960
GPU_MEMORY_UTILIZATION=0.90
REFRESH_BACKEND="${DMTD_REFRESH_BACKEND:-merged}"
EXTRA_ARGS=()

usage() {
  cat <<'EOF'
Usage: ./start_vllm_server.sh [options] [-- extra-vllm-args]

Options:
  --device DEVICE              GPU to expose. For "01234567" or "0,1,...",
                               single-GPU mode selects the first GPU.
  --port PORT                  Server port. Default: choose a free non-core port.
  --max-model-len N            Maximum context length. Default: 40960.
  --gpu-memory-utilization F   vLLM GPU memory fraction. Default: 0.90.
  --served-model-name NAME     OpenAI API model name. Default: dmtd-qwen3-4b.
  --refresh-backend BACKEND    Refresh execution: merged or two_pass.
                               Default: merged.
  -h, --help                   Show this help.

Run this script inside the ms-swift container. It starts a single-GPU,
OpenAI-compatible Bidirectional-Parallel-Refresh vLLM server in the foreground.
The modified vLLM automatically selects FlexAttention for this checkpoint.
EOF
}

while (($#)); do
  case "$1" in
    --device)
      DEVICE="${2:?--device requires a value}"
      shift 2
      ;;
    --port)
      PORT="${2:?--port requires a value}"
      shift 2
      ;;
    --max-model-len)
      MAX_MODEL_LEN="${2:?--max-model-len requires a value}"
      shift 2
      ;;
    --gpu-memory-utilization)
      GPU_MEMORY_UTILIZATION="${2:?--gpu-memory-utilization requires a value}"
      shift 2
      ;;
    --served-model-name)
      SERVED_MODEL_NAME="${2:?--served-model-name requires a value}"
      shift 2
      ;;
    --refresh-backend)
      REFRESH_BACKEND="${2:?--refresh-backend requires a value}"
      shift 2
      ;;
    -h|--help)
      usage
      exit 0
      ;;
    --)
      shift
      EXTRA_ARGS=("$@")
      break
      ;;
    *)
      echo "Unknown argument: $1" >&2
      usage >&2
      exit 2
      ;;
  esac
done

if [[ "${REFRESH_BACKEND}" != "merged" && "${REFRESH_BACKEND}" != "two_pass" ]]; then
  echo "Invalid refresh backend: ${REFRESH_BACKEND}" >&2
  echo "Expected merged or two_pass." >&2
  exit 2
fi

if [[ -x /workspace/vllm/.venv/bin/vllm ]]; then
  VLLM_BIN=/workspace/vllm/.venv/bin/vllm
  PYTHON_BIN=/workspace/vllm/.venv/bin/python
elif command -v vllm >/dev/null 2>&1; then
  VLLM_BIN="$(command -v vllm)"
  PYTHON_BIN="$(command -v python3)"
else
  echo "vLLM was not found. Run this script inside the ms-swift container." >&2
  echo "Host example:" >&2
  echo "  sudo docker exec -it ms-swift bash" >&2
  echo "  cd /workspace/parallel-eval/models/Bidirectional-Parallel-Refresh" >&2
  exit 1
fi

select_first_device() {
  local value="$1"
  value="${value#cuda:}"
  value="${value%%,*}"
  value="${value%% *}"
  if [[ "${value}" =~ ^[0-9]+$ && ${#value} -gt 1 ]]; then
    value="${value:0:1}"
  fi
  if [[ ! "${value}" =~ ^[0-9]+$ ]]; then
    echo "Invalid GPU device: $1" >&2
    exit 2
  fi
  printf '%s' "${value}"
}

if [[ -z "${DEVICE}" ]]; then
  if [[ -n "${CUDA_VISIBLE_DEVICES:-}" ]]; then
    DEVICE="${CUDA_VISIBLE_DEVICES}"
  elif command -v nvidia-smi >/dev/null 2>&1; then
    DEVICE="$(nvidia-smi --query-gpu=index --format=csv,noheader | awk 'NR==1 {print $1}')"
  else
    DEVICE=0
  fi
fi
DEVICE="$(select_first_device "${DEVICE}")"

PORT="$(
  "${PYTHON_BIN}" - "${PORT}" <<'PY'
import secrets
import socket
import sys

requested = sys.argv[1]
reserved = {
    22, 80, 443, 3000, 3306, 5000, 5432, 6006, 6379, 8000, 8080, 8265,
    8888, 9090, 27017,
}


def is_free(port: int) -> bool:
    if port in reserved:
        return False
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        try:
            sock.bind(("127.0.0.1", port))
        except OSError:
            return False
    return True


if requested:
    try:
        port = int(requested)
    except ValueError as exc:
        raise SystemExit(f"Invalid port: {requested}") from exc
    if not 1024 <= port <= 65535:
        raise SystemExit(f"Port must be between 1024 and 65535: {port}")
    if port in reserved:
        raise SystemExit(f"Refusing reserved/core port: {port}")
    if not is_free(port):
        raise SystemExit(f"Port is already in use: {port}")
    print(port)
else:
    for _ in range(512):
        port = 18000 + secrets.randbelow(10000)
        if is_free(port):
            print(port)
            break
    else:
        raise SystemExit("Could not find a free port in 18000-27999")
PY
)"

echo "Starting Bidirectional-Parallel-Refresh vLLM server"
echo "  model:           ${MODEL_PATH}"
echo "  device:          ${DEVICE} (single GPU)"
echo "  endpoint:        http://127.0.0.1:${PORT}/v1"
echo "  model id:        ${SERVED_MODEL_NAME}"
echo "  refresh backend: ${REFRESH_BACKEND}"

export CUDA_VISIBLE_DEVICES="${DEVICE}"
export VLLM_USE_MODELSCOPE=False
export VLLM_USE_V2_MODEL_RUNNER=1
export DMTD_REFRESH_BACKEND="${REFRESH_BACKEND}"

exec "${VLLM_BIN}" serve "${MODEL_PATH}" \
  --host 0.0.0.0 \
  --port "${PORT}" \
  --served-model-name "${SERVED_MODEL_NAME}" \
  --tensor-parallel-size 1 \
  --dtype bfloat16 \
  --generation-config auto \
  --max-model-len "${MAX_MODEL_LEN}" \
  --gpu-memory-utilization "${GPU_MEMORY_UTILIZATION}" \
  --no-enable-prefix-caching \
  --enforce-eager \
  "${EXTRA_ARGS[@]}"