run_engines.sh 8.7 KB
Newer Older
Yan Ru Pei's avatar
Yan Ru Pei committed
1
2
3
4
5
6
#!/bin/bash

# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

# Parse command-line arguments
7
export DYNAMO_HOME=${DYNAMO_HOME:-"/workspace"}
Yan Ru Pei's avatar
Yan Ru Pei committed
8
9
NUM_WORKERS=8
MODEL_PATH="deepseek-ai/DeepSeek-R1-Distill-Llama-8B"
10
RECIPE_PATH="$DYNAMO_HOME/recipes/deepseek-r1-distill-llama-8b/trtllm"
Yan Ru Pei's avatar
Yan Ru Pei committed
11
TENSOR_PARALLEL_SIZE=1
Yan Ru Pei's avatar
Yan Ru Pei committed
12
DATA_PARALLEL_SIZE=1
Yan Ru Pei's avatar
Yan Ru Pei committed
13
USE_MOCKERS=false
14
15
USE_TRTLLM=false
MODE="agg"  # Options: agg (default), decode, prefill
Yan Ru Pei's avatar
Yan Ru Pei committed
16
BASE_GPU_OFFSET=0
Yan Ru Pei's avatar
Yan Ru Pei committed
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
EXTRA_ARGS=()

# Parse arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        --num-workers)
            NUM_WORKERS="$2"
            shift 2
            ;;
        --model-path)
            MODEL_PATH="$2"
            shift 2
            ;;
        --tensor-parallel-size)
            TENSOR_PARALLEL_SIZE="$2"
            shift 2
            ;;
Yan Ru Pei's avatar
Yan Ru Pei committed
34
35
36
37
        --data-parallel-size)
            DATA_PARALLEL_SIZE="$2"
            shift 2
            ;;
Yan Ru Pei's avatar
Yan Ru Pei committed
38
39
40
41
        --mockers)
            USE_MOCKERS=true
            shift
            ;;
42
43
44
45
46
47
48
49
50
51
        --trtllm)
            USE_TRTLLM=true
            shift
            ;;
        --prefill)
            MODE="prefill"
            shift
            ;;
        --decode)
            MODE="decode"
Yan Ru Pei's avatar
Yan Ru Pei committed
52
53
54
55
56
57
            shift
            ;;
        --base-gpu-offset)
            BASE_GPU_OFFSET="$2"
            shift 2
            ;;
Yan Ru Pei's avatar
Yan Ru Pei committed
58
59
60
61
62
63
        --)
            shift
            EXTRA_ARGS+=("$@")
            break
            ;;
        *)
64
            # Collect all other arguments as vLLM/mocker/trtllm arguments
Yan Ru Pei's avatar
Yan Ru Pei committed
65
66
67
68
69
70
            EXTRA_ARGS+=("$1")
            shift
            ;;
    esac
done

71
72
73
74
75
76
77
78
79
# Validate that only one engine type is selected
ENGINE_COUNT=0
[ "$USE_MOCKERS" = true ] && ((ENGINE_COUNT++))
[ "$USE_TRTLLM" = true ] && ((ENGINE_COUNT++))
if [ "$ENGINE_COUNT" -gt 1 ]; then
    echo "Error: Only one engine type (--mockers, --trtllm, or default vLLM) can be specified"
    exit 1
fi

Yan Ru Pei's avatar
Yan Ru Pei committed
80
81
82
83
84
85
86
# If no extra args provided, use defaults
if [ ${#EXTRA_ARGS[@]} -eq 0 ]; then
    if [ "$USE_MOCKERS" = true ]; then
        # Default args for mocker engine (only block-size needed as others are defaults)
        EXTRA_ARGS=(
            "--block-size" "64"
        )
87
88
    elif [ "$USE_TRTLLM" = true ]; then
        # Default args for TensorRT-LLM engine using predefined YAML configs
89
        # Config files located at: $RECIPE_PATH/{agg,decode,prefill}.yaml
90
        if [ "$MODE" = "prefill" ]; then
91
            ENGINE_CONFIG="$RECIPE_PATH/prefill.yaml"
92
        elif [ "$MODE" = "decode" ]; then
93
            ENGINE_CONFIG="$RECIPE_PATH/decode.yaml"
94
        else
95
            ENGINE_CONFIG="$RECIPE_PATH/agg.yaml"
96
97
98
99
100
101
        fi

        EXTRA_ARGS=(
            "--extra-engine-args" "$ENGINE_CONFIG"
            "--publish-events-and-metrics"
        )
Yan Ru Pei's avatar
Yan Ru Pei committed
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
    else
        # Default args for vLLM engine (explicitly include block-size)
        EXTRA_ARGS=(
            "--enforce-eager"
            "--max-num-batched-tokens" "16384"
            "--max-model-len" "32768"
            "--block-size" "64"
        )
    fi
fi

# Validate arguments
if ! [[ "$NUM_WORKERS" =~ ^[0-9]+$ ]] || [ "$NUM_WORKERS" -lt 1 ]; then
    echo "Error: NUM_WORKERS must be a positive integer"
    exit 1
fi

if ! [[ "$TENSOR_PARALLEL_SIZE" =~ ^[0-9]+$ ]] || [ "$TENSOR_PARALLEL_SIZE" -lt 1 ]; then
    echo "Error: TENSOR_PARALLEL_SIZE must be a positive integer"
    exit 1
fi

Yan Ru Pei's avatar
Yan Ru Pei committed
124
125
126
127
128
if ! [[ "$DATA_PARALLEL_SIZE" =~ ^[0-9]+$ ]] || [ "$DATA_PARALLEL_SIZE" -lt 1 ]; then
    echo "Error: DATA_PARALLEL_SIZE must be a positive integer"
    exit 1
fi

Yan Ru Pei's avatar
Yan Ru Pei committed
129
130
131
132
133
if ! [[ "$BASE_GPU_OFFSET" =~ ^[0-9]+$ ]]; then
    echo "Error: BASE_GPU_OFFSET must be a non-negative integer"
    exit 1
fi

Yan Ru Pei's avatar
Yan Ru Pei committed
134
135
136
# Calculate total GPUs needed (TP * DP per worker)
GPUS_PER_WORKER=$((TENSOR_PARALLEL_SIZE * DATA_PARALLEL_SIZE))
TOTAL_GPUS_NEEDED=$((NUM_WORKERS * GPUS_PER_WORKER))
Yan Ru Pei's avatar
Yan Ru Pei committed
137
LAST_GPU=$((BASE_GPU_OFFSET + TOTAL_GPUS_NEEDED - 1))
Yan Ru Pei's avatar
Yan Ru Pei committed
138
echo "Configuration:"
139
140
141
142
143
144
145
146
147
if [ "$USE_MOCKERS" = true ]; then
    ENGINE_TYPE="Mocker"
elif [ "$USE_TRTLLM" = true ]; then
    ENGINE_TYPE="TensorRT-LLM"
else
    ENGINE_TYPE="vLLM"
fi
echo "  Engine Type: $ENGINE_TYPE"
echo "  Mode: $MODE"
Yan Ru Pei's avatar
Yan Ru Pei committed
148
149
150
echo "  Workers: $NUM_WORKERS"
echo "  Model: $MODEL_PATH"
echo "  Tensor Parallel Size: $TENSOR_PARALLEL_SIZE"
Yan Ru Pei's avatar
Yan Ru Pei committed
151
152
echo "  Data Parallel Size: $DATA_PARALLEL_SIZE"
echo "  GPUs per worker: $GPUS_PER_WORKER"
Yan Ru Pei's avatar
Yan Ru Pei committed
153
echo "  Total GPUs needed: $TOTAL_GPUS_NEEDED"
Yan Ru Pei's avatar
Yan Ru Pei committed
154
echo "  GPU Range: $BASE_GPU_OFFSET-$LAST_GPU"
Yan Ru Pei's avatar
Yan Ru Pei committed
155
156
157
158
159
160
161
162
163
164
165
166
167
168
echo "  Engine args: ${EXTRA_ARGS[*]}"
echo ""

PIDS=()

cleanup() {
    echo -e "\nStopping all workers..."
    kill "${PIDS[@]}" 2>/dev/null
    wait
    exit 0
}

trap cleanup SIGINT SIGTERM

169
echo "Starting $NUM_WORKERS $MODE workers..."
Yan Ru Pei's avatar
Yan Ru Pei committed
170

171
172
173
174
175
if [ "$USE_MOCKERS" = true ]; then
    # For mockers, launch a single process with --num-workers
    # All workers share the same tokio runtime and thread pool
    MODE_CAPITALIZED=$(echo "$MODE" | sed 's/\(.\)/\U\1/')
    echo "[$MODE_CAPITALIZED Mocker] Starting $NUM_WORKERS workers in single process..."
Yan Ru Pei's avatar
Yan Ru Pei committed
176

177
178
179
    MOCKER_ARGS=()
    MOCKER_ARGS+=("--model-path" "$MODEL_PATH")
    MOCKER_ARGS+=("--num-workers" "$NUM_WORKERS")
Yan Ru Pei's avatar
Yan Ru Pei committed
180

181
182
183
184
185
186
187
188
189
190
    # Set endpoint based on worker mode
    if [ "$MODE" = "prefill" ]; then
        MOCKER_ARGS+=("--endpoint" "dyn://test.prefill.generate")
        MOCKER_ARGS+=("--is-prefill-worker")
    elif [ "$MODE" = "decode" ]; then
        MOCKER_ARGS+=("--endpoint" "dyn://test.mocker.generate")
        MOCKER_ARGS+=("--is-decode-worker")
    else
        MOCKER_ARGS+=("--endpoint" "dyn://test.mocker.generate")
    fi
Yan Ru Pei's avatar
Yan Ru Pei committed
191

192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
    if [ "$DATA_PARALLEL_SIZE" -gt 1 ]; then
        MOCKER_ARGS+=("--data-parallel-size" "$DATA_PARALLEL_SIZE")
    fi
    MOCKER_ARGS+=("${EXTRA_ARGS[@]}")

    python -m dynamo.mocker "${MOCKER_ARGS[@]}" &
    PIDS+=($!)
    echo "Started mocker with $NUM_WORKERS workers (PID: $!)"
else
    # For vLLM and TensorRT-LLM, use the original loop to launch separate processes
    for i in $(seq 1 $NUM_WORKERS); do
        {
            MODE_CAPITALIZED=$(echo "$MODE" | sed 's/\(.\)/\U\1/')
            echo "[$MODE_CAPITALIZED Worker-$i] Starting..."

            # Calculate GPU indices for this worker (with base offset)
            # Each worker needs TP * DP GPUs
            START_GPU=$(( BASE_GPU_OFFSET + (i - 1) * GPUS_PER_WORKER ))
            END_GPU=$(( START_GPU + GPUS_PER_WORKER - 1 ))
Yan Ru Pei's avatar
Yan Ru Pei committed
211

212
213
214
            # Build CUDA_VISIBLE_DEVICES string for all GPUs (TP * DP)
            if [ "$GPUS_PER_WORKER" -eq 1 ]; then
                GPU_DEVICES="$START_GPU"
Yan Ru Pei's avatar
Yan Ru Pei committed
215
            else
216
217
218
219
220
221
222
223
                GPU_DEVICES=""
                for gpu in $(seq $START_GPU $END_GPU); do
                    if [ -n "$GPU_DEVICES" ]; then
                        GPU_DEVICES="${GPU_DEVICES},$gpu"
                    else
                        GPU_DEVICES="$gpu"
                    fi
                done
Yan Ru Pei's avatar
Yan Ru Pei committed
224
225
            fi

226
227
            if [ "$USE_TRTLLM" = true ]; then
                echo "[$MODE_CAPITALIZED Worker-$i] Using GPUs: $GPU_DEVICES"
228
                # Run TensorRT-LLM engine
229
230
231
232
233
234
235
                TRTLLM_ARGS=()
                TRTLLM_ARGS+=("--model-path" "$MODEL_PATH")
                TRTLLM_ARGS+=("--tensor-parallel-size" "$TENSOR_PARALLEL_SIZE")
                if [ "$MODE" != "agg" ]; then
                    TRTLLM_ARGS+=("--disaggregation-mode" "$MODE")
                fi
                TRTLLM_ARGS+=("${EXTRA_ARGS[@]}")
Yan Ru Pei's avatar
Yan Ru Pei committed
236

237
                exec env CUDA_VISIBLE_DEVICES=$GPU_DEVICES trtllm-llmapi-launch python3 -m dynamo.trtllm \
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
                    "${TRTLLM_ARGS[@]}"
            else
                echo "[$MODE_CAPITALIZED Worker-$i] Using GPUs: $GPU_DEVICES"
                # Run vLLM engine with PYTHONHASHSEED=0 for deterministic event IDs in KV-aware routing
                VLLM_ARGS=()
                VLLM_ARGS+=("--model" "$MODEL_PATH")
                VLLM_ARGS+=("--tensor-parallel-size" "$TENSOR_PARALLEL_SIZE")
                if [ "$DATA_PARALLEL_SIZE" -gt 1 ]; then
                    VLLM_ARGS+=("--data-parallel-size" "$DATA_PARALLEL_SIZE")
                fi
                if [ "$MODE" = "prefill" ]; then
                    VLLM_ARGS+=("--is-prefill-worker")
                elif [ "$MODE" = "decode" ]; then
                    VLLM_ARGS+=("--is-decode-worker")
                fi
                VLLM_ARGS+=("${EXTRA_ARGS[@]}")
254

255
                exec env PYTHONHASHSEED=0 CUDA_VISIBLE_DEVICES=$GPU_DEVICES python3 -m dynamo.vllm \
256
                    "${VLLM_ARGS[@]}"
Yan Ru Pei's avatar
Yan Ru Pei committed
257
            fi
258
259
260
        } &
        PIDS+=($!)
        echo "Started $MODE worker $i (PID: $!)"
261
262
263
264
265
266

        # Add delay between TensorRT-LLM worker launches to avoid MPI initialization conflicts
        if [ "$USE_TRTLLM" = true ] && [ "$i" -lt "$NUM_WORKERS" ]; then
            echo "Waiting 2 seconds before launching next TensorRT-LLM worker..."
            sleep 2
        fi
267
268
    done
fi
Yan Ru Pei's avatar
Yan Ru Pei committed
269
270
271
272

echo "All workers started. Press Ctrl+C to stop."
wait
echo "All workers completed."