@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
sh 118 lines 3.71 KB
Raw
1 #!/bin/bash
2 #
3 # Helper script to manually start or join a Ray cluster for online serving of vLLM models.
4 # This script is first executed on the head node, and then on each worker node with the IP address
5 # of the head node.
6 #
7 # Subcommands:
8 # leader: Launches a Ray head node and blocks until the cluster reaches the expected size (head + workers).
9 # worker: Starts a worker node that connects to an existing Ray head node.
10 #
11 # Example usage:
12 # On the head node machine, start the Ray head node process and run a vLLM server.
13 # ./multi-node-serving.sh leader --ray_port=6379 --ray_cluster_size=<SIZE> [<extra ray args>] && \
14 # vllm serve meta-llama/Meta-Llama-3.1-405B-Instruct --port 8080 --tensor-parallel-size 8 --pipeline-parallel-size 2 --distributed-executor-backend ray
15 #
16 # On each worker node, start the Ray worker node process.
17 # ./multi-node-serving.sh worker --ray_address=<HEAD_NODE_IP> --ray_port=6379 [<extra ray args>]
18 #
19 # About Ray:
20 # Ray is an open-source distributed execution framework that simplifies
21 # distributed computing. Learn more:
22 # https://ray.io/
23
24
25 subcommand=$1 # Either "leader" or "worker".
26 shift # Remove the subcommand from the argument list.
27
28 ray_port=6379 # Port used by the Ray head node.
29 ray_init_timeout=300 # Seconds to wait before timing out.
30 declare -a start_params # Parameters forwarded to the underlying 'ray start' command.
31
32 # Handle the worker subcommand.
33 case "$subcommand" in
34 worker)
35 ray_address=""
36 while [ $# -gt 0 ]; do
37 case "$1" in
38 --ray_address=*)
39 ray_address="${1#*=}"
40 ;;
41 --ray_port=*)
42 ray_port="${1#*=}"
43 ;;
44 --ray_init_timeout=*)
45 ray_init_timeout="${1#*=}"
46 ;;
47 *)
48 start_params+=("$1")
49 esac
50 shift
51 done
52
53 if [ -z "$ray_address" ]; then
54 echo "Error: Missing argument --ray_address"
55 exit 1
56 fi
57
58 # Retry until the worker node connects to the head node or the timeout expires.
59 for (( i=0; i < $ray_init_timeout; i+=5 )); do
60 if ray start --address="$ray_address":"$ray_port" --block "${start_params[@]}"; then
61 echo "Worker: Ray runtime started with head address $ray_address:$ray_port"
62 exit 0
63 fi
64 echo "Waiting until the ray worker is active..."
65 sleep 5s;
66 done
67 echo "Ray worker starts timeout, head address: $ray_address:$ray_port"
68 exit 1
69 ;;
70
71 # Handle the leader subcommand.
72 leader)
73 ray_cluster_size=""
74 while [ $# -gt 0 ]; do
75 case "$1" in
76 --ray_port=*)
77 ray_port="${1#*=}"
78 ;;
79 --ray_cluster_size=*)
80 ray_cluster_size="${1#*=}"
81 ;;
82 --ray_init_timeout=*)
83 ray_init_timeout="${1#*=}"
84 ;;
85 *)
86 start_params+=("$1")
87 esac
88 shift
89 done
90
91 if [ -z "$ray_cluster_size" ]; then
92 echo "Error: Missing argument --ray_cluster_size"
93 exit 1
94 fi
95
96 # Start the Ray head node.
97 ray start --head --port="$ray_port" "${start_params[@]}"
98
99 # Poll Ray until every worker node is active.
100 for (( i=0; i < $ray_init_timeout; i+=5 )); do
101 active_nodes=$(python3 -c 'import ray; ray.init(); print(sum(node["Alive"] for node in ray.nodes()))')
102 if [ "$active_nodes" -eq "$ray_cluster_size" ]; then
103 echo "All ray workers are active and the ray cluster is initialized successfully."
104 exit 0
105 fi
106 echo "Wait for all ray workers to be active. $active_nodes/$ray_cluster_size is active"
107 sleep 5s;
108 done
109
110 echo "Waiting for all ray workers to be active timed out."
111 exit 1
112 ;;
113
114 *)
115 echo "unknown subcommand: $subcommand"
116 exit 1
117 ;;
118 esac