diff --git a/tools/bench/run_serving_bench.sh b/tools/bench/run_serving_bench.sh index 6fd972eb4..c8124607f 100755 --- a/tools/bench/run_serving_bench.sh +++ b/tools/bench/run_serving_bench.sh @@ -16,14 +16,22 @@ # GPU CUDA device ordinal [default: 0] # PORT server port [default: 8000] # RESULT_DIR output directory [default: ./bench-results] -# DATASET vllm-bench dataset: random | sharegpt | sonnet | speed-bench [default: random] +# DATASET vLLM dataset: random | sharegpt | hf | custom | ... [default: random] +# DATASET_PATH Local dataset path or Hugging Face dataset ID +# HF_SUBSET Hugging Face subset (for example, main for openai/gsm8k) +# HF_SPLIT Hugging Face split [default: test when DATASET=hf] +# HF_NO_STREAM load the Hugging Face dataset locally: 0 | 1 [default: 0] +# HF_OFFLINE force Hugging Face Hub and datasets offline: 0 | 1 [default: 0] # QPS_LIST space-separated QPS values [default: "1 2 4 8 10 12 16"] # CONCURRENCY_LIST space-separated concurrency values for spec sweep [default: "1 4 8"] # INPUT_LEN input length [default: 1024] # OUTPUT_LEN output length [default: 128] +# NUM_WARMUPS warm-up requests before each measured point [default: 0] +# SAVE_DETAILED save per-request timing/output arrays: 0 | 1 [default: 0] +# BINARY prebuilt pegainfer binary path # SEED base random seed; each point derives its own from SEED + axis/value [default: 42] # SECONDS_PER_RUN seconds per QPS run [default: 60] -# BENCH path to vllm-bench binary [default: vllm-bench on PATH] +# BENCH benchmark command; auto-detects `vllm bench serve` or vllm-bench # VLLM path to vllm binary for ENGINE=vllm [default: vllm on PATH] # VLLM_EXTRA_ARGS extra args passed to `vllm serve` [default: "--max-model-len 8192"] # LABEL engine label for result filenames [default: $ENGINE] @@ -48,13 +56,19 @@ GPU=${GPU:-0} PORT=${PORT:-8000} RESULT_DIR=${RESULT_DIR:-./bench-results} DATASET=${DATASET:-random} +DATASET_PATH=${DATASET_PATH:-} +HF_SUBSET=${HF_SUBSET:-} +HF_SPLIT=${HF_SPLIT:-} +HF_NO_STREAM=${HF_NO_STREAM:-0} +HF_OFFLINE=${HF_OFFLINE:-0} QPS_LIST=${QPS_LIST-"1 2 4 8 10 12 16"} CONCURRENCY_LIST=${CONCURRENCY_LIST:-"1 4 8"} INPUT_LEN=${INPUT_LEN:-1024} OUTPUT_LEN=${OUTPUT_LEN:-128} +NUM_WARMUPS=${NUM_WARMUPS:-0} +SAVE_DETAILED=${SAVE_DETAILED:-0} SEED=${SEED:-42} SECONDS_PER_RUN=${SECONDS_PER_RUN:-60} -BENCH=${BENCH:-vllm-bench} LABEL=${LABEL:-$ENGINE} SKIP_BUILD=${SKIP_BUILD:-0} @@ -79,10 +93,51 @@ RESULT_FILES=() mkdir -p "$RESULT_DIR" +if [[ -n "${BENCH:-}" ]]; then + read -r -a BENCH_CMD <<< "$BENCH" +elif command -v vllm >/dev/null 2>&1; then + BENCH_CMD=(vllm bench serve) +elif command -v vllm-bench >/dev/null 2>&1; then + BENCH_CMD=(vllm-bench) +else + echo "FATAL: neither 'vllm bench serve' nor vllm-bench is available" >&2 + exit 1 +fi + +if [[ "$HF_OFFLINE" == "1" ]]; then + export HF_HUB_OFFLINE=1 + export HF_DATASETS_OFFLINE=1 +fi + +DATASET_ARGS=(--dataset-name "$DATASET" --output-len "$OUTPUT_LEN") +if [[ -n "$DATASET_PATH" ]]; then + DATASET_ARGS+=(--dataset-path "$DATASET_PATH") +fi +if [[ "$DATASET" == "hf" ]]; then + if [[ -z "$DATASET_PATH" ]]; then + echo "FATAL: DATASET_PATH is required when DATASET=hf" >&2 + exit 1 + fi + DATASET_ARGS+=(--hf-split "${HF_SPLIT:-test}") + if [[ -n "$HF_SUBSET" ]]; then + DATASET_ARGS+=(--hf-subset "$HF_SUBSET") + fi + if [[ "$HF_NO_STREAM" == "1" ]]; then + DATASET_ARGS+=(--no-stream) + fi +fi +if [[ "$DATASET" == "random" ]]; then + DATASET_ARGS+=(--input-len "$INPUT_LEN") +fi +COMMON_BENCH_ARGS=(--num-warmups "$NUM_WARMUPS") +if [[ "$SAVE_DETAILED" == "1" ]]; then + COMMON_BENCH_ARGS+=(--save-detailed) +fi + # ---- launch server ---------------------------------------------------------- case "$ENGINE" in pegainfer) - BINARY="$REPO_ROOT/target/release/pegainfer" + BINARY=${BINARY:-"${CARGO_TARGET_DIR:-$REPO_ROOT/target}/release/pegainfer"} if [[ "$SKIP_BUILD" != "1" ]]; then echo "=== building pegainfer (SKIP_BUILD=1 to skip) ===" (cd "$REPO_ROOT" && CUDA_HOME=${CUDA_HOME:-/usr/local/cuda} cargo build --release -p pegainfer-server) @@ -105,6 +160,7 @@ case "$ENGINE" in vllm) VLLM=${VLLM:-vllm} VLLM_EXTRA_ARGS=${VLLM_EXTRA_ARGS:-"--max-model-len 8192"} + read -r -a VLLM_EXTRA_CMD <<< "$VLLM_EXTRA_ARGS" if [[ -n "$DRAFT_MODEL" ]]; then echo "WARN: DRAFT_MODEL is ignored for ENGINE=vllm" >&2 fi @@ -113,7 +169,7 @@ case "$ENGINE" in --port "$PORT" \ --served-model-name "$MODEL" \ --trust-remote-code \ - $VLLM_EXTRA_ARGS \ + "${VLLM_EXTRA_CMD[@]}" \ > "$RESULT_DIR/server-${ENGINE}-${MODEL_LABEL}.log" 2>&1 & SERVER_PID=$! # vLLM cold start (torch.compile) can take 70+ seconds @@ -164,18 +220,15 @@ fi # ---- QPS sweep -------------------------------------------------------------- if [[ -n "${QPS_LIST// /}" ]]; then echo "=== QPS sweep: qps=[$QPS_LIST] dataset=$DATASET ===" - DATASET_ARGS=(--dataset-name "$DATASET") - if [[ "$DATASET" == "random" ]]; then - DATASET_ARGS+=(--random-input-len "$INPUT_LEN" --random-output-len "$OUTPUT_LEN") - fi for QPS in $QPS_LIST; do NUM_PROMPTS=$(python3 -c "print(int($QPS * $SECONDS_PER_RUN))") point_seed qps "$QPS" echo "--- $LABEL $MODEL_LABEL qps=$QPS num_prompts=$NUM_PROMPTS dataset=$DATASET seed=$POINT_SEED ---" - "$BENCH" \ + "${BENCH_CMD[@]}" \ --backend openai --model "$MODEL" --port "$PORT" \ --base-url "http://localhost:$PORT" \ "${DATASET_ARGS[@]}" \ + "${COMMON_BENCH_ARGS[@]}" \ --num-prompts "$NUM_PROMPTS" \ --request-rate "$QPS" \ --seed "$POINT_SEED" \ @@ -193,18 +246,15 @@ fi # ---- Concurrency sweep (pegainfer only) ------------------------------------ if [[ "${ENGINE}" == "pegainfer" && -n "${CONCURRENCY_LIST// /}" ]]; then echo "=== spec concurrency sweep: c=[$CONCURRENCY_LIST] dataset=$DATASET ===" - DATASET_ARGS=(--dataset-name "$DATASET") - if [[ "$DATASET" == "random" ]]; then - DATASET_ARGS+=(--random-input-len "$INPUT_LEN" --random-output-len "$OUTPUT_LEN") - fi for C in $CONCURRENCY_LIST; do NUM_PROMPTS=$(python3 -c "print(int($C * $SECONDS_PER_RUN))") point_seed c "$C" echo "--- $LABEL $MODEL_LABEL c=$C num_prompts=$NUM_PROMPTS dataset=$DATASET seed=$POINT_SEED ---" - "$BENCH" \ + "${BENCH_CMD[@]}" \ --backend openai --model "$MODEL" --port "$PORT" \ --base-url "http://localhost:$PORT" \ "${DATASET_ARGS[@]}" \ + "${COMMON_BENCH_ARGS[@]}" \ --num-prompts "$NUM_PROMPTS" \ --max-concurrency "$C" \ --seed "$POINT_SEED" \ @@ -217,6 +267,14 @@ if [[ "${ENGINE}" == "pegainfer" && -n "${CONCURRENCY_LIST// /}" ]]; then done fi +METRICS_FILE="$RESULT_DIR/${LABEL}-${MODEL_LABEL}-${DATASET}-metrics.prom" +if curl -sf "http://localhost:$PORT/metrics" > "$METRICS_FILE"; then + echo "metrics saved to $METRICS_FILE" +else + echo "WARN: /metrics was unavailable; no metrics snapshot saved" >&2 + rm -f "$METRICS_FILE" +fi + # ---- summary --------------------------------------------------------------- echo "" echo "=== results summary ==="