| 1 | #!/bin/bash |
| 2 | # |
| 3 | # Launch a Ray cluster inside Docker for vLLM inference. |
| 4 | # |
| 5 | # This script can start either a head node or a worker node, depending on the |
| 6 | # --head or --worker flag provided as the third positional argument. |
| 7 | # |
| 8 | # Usage: |
| 9 | # 1. Designate one machine as the head node and execute: |
| 10 | # bash run_cluster.sh \ |
| 11 | # vllm/vllm-openai \ |
| 12 | # <head_node_ip> \ |
| 13 | # --head \ |
| 14 | # /abs/path/to/huggingface/cache \ |
| 15 | # -e VLLM_HOST_IP=<head_node_ip> |
| 16 | # |
| 17 | # 2. On every worker machine, execute: |
| 18 | # bash run_cluster.sh \ |
| 19 | # vllm/vllm-openai \ |
| 20 | # <head_node_ip> \ |
| 21 | # --worker \ |
| 22 | # /abs/path/to/huggingface/cache \ |
| 23 | # -e VLLM_HOST_IP=<worker_node_ip> |
| 24 | # |
| 25 | # Each worker requires a unique VLLM_HOST_IP value. |
| 26 | # Keep each terminal session open. Closing a session stops the associated Ray |
| 27 | # node and thereby shuts down the entire cluster. |
| 28 | # Every machine must be reachable at the supplied IP address. |
| 29 | # |
| 30 | # The container is named "node-<random_suffix>". To open a shell inside |
| 31 | # a container after launch, use: |
| 32 | # docker exec -it node-<random_suffix> /bin/bash |
| 33 | # |
| 34 | # Then, you can execute vLLM commands on the Ray cluster as if it were a |
| 35 | # single machine, e.g. vllm serve ... |
| 36 | # |
| 37 | # To stop the container, use: |
| 38 | # docker stop node-<random_suffix> |
| 39 | |
| 40 | # Check for minimum number of required arguments. |
| 41 | if [ $# -lt 4 ]; then |
| 42 | echo "Usage: $0 docker_image head_node_ip --head|--worker path_to_hf_home [additional_args...]" |
| 43 | exit 1 |
| 44 | fi |
| 45 | |
| 46 | # Extract the mandatory positional arguments and remove them from $@. |
| 47 | DOCKER_IMAGE="$1" |
| 48 | HEAD_NODE_ADDRESS="$2" |
| 49 | NODE_TYPE="$3" # Should be --head or --worker. |
| 50 | PATH_TO_HF_HOME="$4" |
| 51 | shift 4 |
| 52 | |
| 53 | # Preserve any extra arguments so they can be forwarded to Docker. |
| 54 | ADDITIONAL_ARGS=("$@") |
| 55 | |
| 56 | # Validate the NODE_TYPE argument. |
| 57 | if [ "${NODE_TYPE}" != "--head" ] && [ "${NODE_TYPE}" != "--worker" ]; then |
| 58 | echo "Error: Node type must be --head or --worker" |
| 59 | exit 1 |
| 60 | fi |
| 61 | |
| 62 | # Extract VLLM_HOST_IP from ADDITIONAL_ARGS (e.g. "-e VLLM_HOST_IP=..."). |
| 63 | VLLM_HOST_IP="" |
| 64 | for ((i = 0; i < ${#ADDITIONAL_ARGS[@]}; i++)); do |
| 65 | arg="${ADDITIONAL_ARGS[$i]}" |
| 66 | case "${arg}" in |
| 67 | -e) |
| 68 | next="${ADDITIONAL_ARGS[$((i + 1))]:-}" |
| 69 | if [[ "${next}" == VLLM_HOST_IP=* ]]; then |
| 70 | VLLM_HOST_IP="${next#VLLM_HOST_IP=}" |
| 71 | break |
| 72 | fi |
| 73 | ;; |
| 74 | -eVLLM_HOST_IP=* | VLLM_HOST_IP=*) |
| 75 | VLLM_HOST_IP="${arg#*=}" |
| 76 | break |
| 77 | ;; |
| 78 | esac |
| 79 | done |
| 80 | |
| 81 | # For the head node, HEAD_NODE_ADDRESS and VLLM_HOST_IP should be consistent. |
| 82 | if [[ "${NODE_TYPE}" == "--head" && -n "${VLLM_HOST_IP}" ]]; then |
| 83 | if [[ "${VLLM_HOST_IP}" != "${HEAD_NODE_ADDRESS}" ]]; then |
| 84 | echo "Warning: VLLM_HOST_IP (${VLLM_HOST_IP}) differs from head_node_ip (${HEAD_NODE_ADDRESS})." |
| 85 | echo "Using VLLM_HOST_IP as the head node address." |
| 86 | HEAD_NODE_ADDRESS="${VLLM_HOST_IP}" |
| 87 | fi |
| 88 | fi |
| 89 | |
| 90 | # Generate a unique container name with random suffix. |
| 91 | # Docker container names must be unique on each host. |
| 92 | # The random suffix allows multiple Ray containers to run simultaneously on the same machine, |
| 93 | # for example, on a multi-GPU machine. |
| 94 | CONTAINER_NAME="node-${RANDOM}" |
| 95 | |
| 96 | # Define a cleanup routine that removes the container when the script exits. |
| 97 | # This prevents orphaned containers from accumulating if the script is interrupted. |
| 98 | cleanup() { |
| 99 | docker stop "${CONTAINER_NAME}" |
| 100 | docker rm "${CONTAINER_NAME}" |
| 101 | } |
| 102 | trap cleanup EXIT |
| 103 | |
| 104 | # Build the Ray start command based on the node role. |
| 105 | # The head node manages the cluster and accepts connections on port 6379, |
| 106 | # while workers connect to the head's address. |
| 107 | RAY_START_CMD="ray start --block" |
| 108 | if [ "${NODE_TYPE}" == "--head" ]; then |
| 109 | RAY_START_CMD+=" --head --node-ip-address=${HEAD_NODE_ADDRESS} --port=6379" |
| 110 | else |
| 111 | |
| 112 | RAY_START_CMD+=" --address=${HEAD_NODE_ADDRESS}:6379" |
| 113 | if [ -n "${VLLM_HOST_IP}" ]; then |
| 114 | RAY_START_CMD+=" --node-ip-address=${VLLM_HOST_IP}" |
| 115 | fi |
| 116 | fi |
| 117 | |
| 118 | # Launch the container with the assembled parameters. |
| 119 | # --network host: Allows Ray nodes to communicate directly via host networking |
| 120 | # --shm-size 10.24g: Increases shared memory |
| 121 | # --gpus all: Gives container access to all GPUs on the host |
| 122 | # -v HF_HOME: Mounts HuggingFace cache to avoid re-downloading models |
| 123 | docker run \ |
| 124 | --entrypoint /bin/bash \ |
| 125 | --network host \ |
| 126 | --name "${CONTAINER_NAME}" \ |
| 127 | --shm-size 10.24g \ |
| 128 | --gpus all \ |
| 129 | -v "${PATH_TO_HF_HOME}:/root/.cache/huggingface" \ |
| 130 | "${ADDITIONAL_ARGS[@]}" \ |
| 131 | "${DOCKER_IMAGE}" -c "${RAY_START_CMD}" |