@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
sh 151 lines 5.06 KB
Raw
1 #!/bin/bash
2
3 set -euox pipefail
4
5 # To detect ROCm
6 # Check multiple indicators:
7 if [ -e /dev/kfd ] || \
8 [ -d /opt/rocm ] || \
9 command -v rocm-smi &> /dev/null || \
10 [ -n "${ROCM_HOME:-}" ]; then
11 IS_ROCM=1
12 else
13 IS_ROCM=0
14 fi
15
16 if [[ $# -lt 4 ]]; then
17 echo "Usage: .buildkite/scripts/run-multi-node-test.sh WORKING_DIR NUM_NODES NUM_GPUS DOCKER_IMAGE COMMAND1 COMMAND2 ... COMMANDN"
18 exit 1
19 fi
20
21 WORKING_DIR=$1
22 NUM_NODES=$2
23 NUM_GPUS=$3
24 DOCKER_IMAGE=$4
25
26 shift 4
27 COMMANDS=("$@")
28 if [ ${#COMMANDS[@]} -ne "$NUM_NODES" ]; then
29 echo "The number of commands must be equal to the number of nodes."
30 echo "Number of nodes: $NUM_NODES"
31 echo "Number of commands: ${#COMMANDS[@]}"
32 exit 1
33 fi
34
35 echo "List of commands"
36 for command in "${COMMANDS[@]}"; do
37 echo "$command"
38 done
39
40
41 STARTED_CONTAINERS=()
42 NETWORK_CREATED=0
43 RUN_DIR=$(mktemp -d)
44 NETWORK_NAME="docker-net-${RUN_DIR##*/}"
45
46 start_network() {
47 if docker network inspect docker-net > /dev/null 2>&1; then
48 if [ "$(docker network inspect --format '{{len .Containers}}' docker-net)" != 0 ]; then
49 echo "docker-net still has attached containers; refusing to reuse it." >&2
50 return 1
51 fi
52 docker network rm docker-net
53 fi
54 docker network create --subnet=192.168.10.0/24 "$NETWORK_NAME"
55 NETWORK_CREATED=1
56 }
57
58 start_nodes() {
59 for node in $(seq 0 $(($NUM_NODES-1))); do
60 DEVICE_LIST=""
61 for node_gpu in $(seq 0 $(($NUM_GPUS - 1))); do
62 DEVICE_NUM=$(($node * $NUM_GPUS + $node_gpu))
63 DEVICE_LIST+=$(($DEVICE_NUM))
64 if [ "$node_gpu" -lt $(($NUM_GPUS - 1)) ]; then
65 DEVICE_LIST+=','
66 fi
67 done
68 if [ "$IS_ROCM" -eq 1 ]; then
69 GPU_DEVICES=(--device /dev/kfd --device /dev/dri -e "HIP_VISIBLE_DEVICES=${DEVICE_LIST}")
70 else
71 # The literal quotes around device=... are required by docker's
72 # --gpus value parser when the device list itself contains commas.
73 GPU_DEVICES=(--gpus "\"device=${DEVICE_LIST}\"")
74 fi
75
76 # start the container in detached mode
77 # things to note:
78 # 1. --shm-size=10.24gb is required. don't use --ipc=host
79 # 2. pass HF_TOKEN to the container
80 # 3. map the huggingface cache directory to the container
81 # 3. assign ip addresses to the containers (head node: 192.168.10.10, worker nodes:
82 # starting from 192.168.10.11)
83 CONTAINER_ID=$(docker run -d "${GPU_DEVICES[@]}" --shm-size=10.24gb -e HF_TOKEN \
84 -v ~/.cache/huggingface:/root/.cache/huggingface --name "$NETWORK_NAME-node$node" \
85 --network "$NETWORK_NAME" --ip 192.168.10.$((10 + $node)) --rm "$DOCKER_IMAGE" \
86 /bin/bash -c "tail -f /dev/null")
87 STARTED_CONTAINERS+=("$CONTAINER_ID")
88
89 # organize containers into a ray cluster
90 if [ "$node" -eq 0 ]; then
91 # start the ray head node
92 docker exec -d "$CONTAINER_ID" /bin/bash -c "ray start --head --port=6379 --block"
93 # wait for the head node to be ready
94 sleep 10
95 else
96 # start the ray worker nodes, and connect them to the head node
97 docker exec -d "$CONTAINER_ID" /bin/bash -c "ray start --address=192.168.10.10:6379 --block"
98 fi
99 done
100
101 # wait for the cluster to be ready
102 sleep 10
103
104 # print the cluster status
105 docker exec "${STARTED_CONTAINERS[0]}" /bin/bash -c "ray status"
106 }
107
108 run_nodes() {
109 # important: iterate in reverse order to start the head node last
110 # we start the worker nodes first, in detached mode, and then start the head node
111 # in the foreground, so that the output of the head node is visible in the buildkite logs
112 for node in $(seq $(($NUM_NODES - 1)) -1 0); do
113 DEVICE_LIST=""
114 for node_gpu in $(seq 0 $(($NUM_GPUS - 1))); do
115 DEVICE_NUM=$(($node * $NUM_GPUS + $node_gpu))
116 DEVICE_LIST+=$(($DEVICE_NUM))
117 if [ "$node_gpu" -lt $(($NUM_GPUS - 1)) ]; then
118 DEVICE_LIST+=','
119 fi
120 done
121 echo "Running node$node with GPU devices: $DEVICE_LIST"
122 if [ "$node" -ne 0 ]; then
123 docker exec -d "${STARTED_CONTAINERS[$node]}" /bin/bash -c "cd $WORKING_DIR ; ${COMMANDS[$node]}"
124 else
125 # Allocate a TTY (-t -i) for the foreground head node so its output
126 # keeps ANSI color in the Buildkite log (see run-amd-test.sh).
127 docker exec -t -i "${STARTED_CONTAINERS[$node]}" /bin/bash -c "cd $WORKING_DIR ; ${COMMANDS[$node]}"
128 fi
129 done
130 }
131 cleanup() {
132 local status=$?
133 local cleanup_status=0
134 if [ "${#STARTED_CONTAINERS[@]}" -gt 0 ]; then
135 for container_id in "${STARTED_CONTAINERS[@]}"; do
136 docker stop "$container_id" || cleanup_status=$?
137 done
138 fi
139 if [ "$NETWORK_CREATED" -eq 1 ]; then
140 docker network rm "$NETWORK_NAME" || cleanup_status=$?
141 fi
142 rmdir "$RUN_DIR" || cleanup_status=$?
143 if [ "$status" -ne 0 ]; then
144 return "$status"
145 fi
146 return "$cleanup_status"
147 }
148 trap cleanup EXIT
149 start_network
150 start_nodes
151 run_nodes