# 下載鏡像，鏡像比較大
docker pull vllm/vllm-openai:v0.7.2

#!/bin/bash

# Check for minimum number of required arguments
if [ $# -lt 4 ]; then
    echo "Usage: $0 docker_image head_node_address --head|--worker path_to_hf_home [additional_args...]"
    exit 1
fi

# Assign the first three arguments and shift them away
DOCKER_IMAGE="$1"
HEAD_NODE_ADDRESS="$2"
NODE_TYPE="$3"  # Should be --head or --worker
PATH_TO_HF_HOME="$4"
shift 4

# Additional arguments are passed directly to the Docker command
ADDITIONAL_ARGS=("$@")

# Validate node type
if [ "${NODE_TYPE}" != "--head" ] && [ "${NODE_TYPE}" != "--worker" ]; then
    echo "Error: Node type must be --head or --worker"
    exit 1
fi

# Define a function to cleanup on EXIT signal
cleanup() {
    docker stop node
    docker rm node
}
trap cleanup EXIT

# Command setup for head or worker node
RAY_START_CMD="ray start --block"
if [ "${NODE_TYPE}" == "--head" ]; then
    RAY_START_CMD+=" --head --port=6379"
else
    RAY_START_CMD+=" --address=${HEAD_NODE_ADDRESS}:6379"
fi

# Run the docker command with the user specified parameters and additional arguments
docker run \
    --entrypoint /bin/bash \
    --network host \
    --name node \
    --shm-size 10.24g \
    --gpus all \
    -v "${PATH_TO_HF_HOME}:/root/.cache/huggingface" \
    "${ADDITIONAL_ARGS[@]}" \
    "${DOCKER_IMAGE}" -c "${RAY_START_CMD}"

# ip_of_head_node：主節(jié)點容器所在宿主機的IP地址
# /path/to/the/huggingface/home/in/this/node: 映射到到容器中的路徑
# ip_of_this_node：當前節(jié)點所在宿主機的IP地址
# --head：表示主節(jié)點
bash run_cluster.sh \
                vllm/vllm-openai \
                ip_of_head_node \
                --head \
                /path/to/the/huggingface/home/in/this/node \
                -e VLLM_HOST_IP=ip_of_this_node

bash run_cluster.sh \
    vllm/vllm-openai:v0.7.2 \
    192.168.108.100 \
    --head \
    /home/vllm \
    -e VLLM_HOST_IP=192.168.108.100 \
    > nohup.log 2>&1 &

# ip_of_head_node：主節(jié)點容器所在宿主機的IP地址
# /path/to/the/huggingface/home/in/this/node: 映射到到容器中的路徑
# ip_of_this_node：當前節(jié)點所在宿主機的IP地址
# --worker：表示工作節(jié)點
bash run_cluster.sh \
                vllm/vllm-openai \
                ip_of_head_node \
                --worker \
                /path/to/the/huggingface/home/in/this/node \
                -e VLLM_HOST_IP=ip_of_this_node

bash run_cluster.sh \
    vllm/vllm-openai:v0.7.2 \
    192.168.108.100 \
    --worker \
    /home/vllm \
    -e VLLM_HOST_IP192.168.108.101 \
    > nohup.log 2>&1 &

# 進入容器
docker exec -it node /bin/bash

# 查看集群信息
ray status
# 返回值中有GPU數(shù)量、CPU配置和內(nèi)存大小等
======== Autoscaler status: 2025-02-13 20:18:13.886242 ========
Node status
---------------------------------------------------------------
Active:
 1 node_89c804d654976b3c606850c461e8dc5c6366de5e0ccdb360fcaa1b1c
 1 node_4b794efd101bc393da41f0a45bd72eeb3fb78e8e507d72b5fdfb4c1b
Pending:
 (no pending nodes)
Recent failures:
 (no failures)

Resources
---------------------------------------------------------------
Usage:
 0.0/128.0 CPU
 0.0/4.0 GPU
 0B/20 GiB memory
 0B/19.46GiB object_store_memory

Demands:
 (no resource demands)

# 啟動模型服務，可根據(jù)情況設置模型參數(shù)
# /path/to/the/model/in/the/container：模型路徑
# tensor-parallel-size：張量并行數(shù)量，模型層內(nèi)拆分后并行計算；
# pipeline-parallel-size：管道并行數(shù)量，模型不同層拆分后并行計算，在單個顯存不夠時可以設置此參數(shù)
vllm serve /path/to/the/model/in/the/container \
     --tensor-parallel-size 8 \
     --pipeline-parallel-size 2

# 進入節(jié)點，主節(jié)點和工作節(jié)點都可以
docker exec -it node /bin/bash

# 執(zhí)行命令參數(shù)
nohup vllm serve /root/.cache/huggingface/Qwen2.5-7B-Instruct \
    --served-model-name qwen2.5-7b \
    --tensor-parallel-size 2 \
    --pipeline-parallel-size 2 \
    > nohup.log 2>&1 &

curl http://localhost:8000/v1/chat/completions \
-X POST \
-H "Content-Type: application/json" \
-d '{
    "model": "qwen2.5-7b",
    "messages": [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "介紹一下中國，不少于10000字"}
    ],
    "stream": true
}'