Skip to content
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
88 changes: 73 additions & 15 deletions tools/bench/run_serving_bench.sh
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand All @@ -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}

Expand All @@ -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")

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Use the Rust flags for the vllm-bench fallback

When only the standalone Rust client is installed, or BENCH=vllm-bench is specified, this shared argument array invokes it with the Python vllm bench serve flags --output-len and --input-len. The maintained Rust client’s CLI contract uses --random-input-len and --random-output-len for random datasets (.agents/skills/vllm-bench/SKILL.md:100-106), so argument parsing fails before any requests are sent and the newly advertised fallback is unusable. Construct dataset-length arguments according to the selected client.

Useful? React with 👍 / 👎.

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)
Expand All @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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" \
Expand All @@ -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" \
Expand All @@ -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 ==="
Expand Down
Loading