373 lines
11 KiB
Bash
Executable File
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

#!/usr/bin/env bash
# TP×DP matrix benchmark for DeepSeek-V4-Flash on vLLM.
set -Eeuo pipefail
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
EXPERIMENT_NAME="$(basename "$SCRIPT_DIR")"
# shellcheck source=/dev/null
source "${SCRIPT_DIR}/../../scripts/common/lib.sh"
# shellcheck source=/dev/null
source "${SCRIPT_DIR}/../../scripts/common/platform.sh"
# shellcheck source=/dev/null
source "${SCRIPT_DIR}/config.env"
RUN_ID="${RUN_ID:-$(date '+%Y%m%d-%H%M%S')}"
RESULT_BASE="${SCRIPT_DIR}/results"
MATRIX_FILE="${MATRIX_FILE:-${SCRIPT_DIR}/matrix.json}"
MATRIX_MODE="${MATRIX_MODE:-Y}"
SCENARIO_TIMEOUT_S="${SCENARIO_TIMEOUT_S:-1800}"
GPU_MEM_SAMPLE_INTERVAL_S="${GPU_MEM_SAMPLE_INTERVAL_S:-1}"
PYTHON="${VENV_CLIENT}/bin/python"
log_dir_global="${RESULT_BASE}/${RUN_ID}/logs"
mkdir -p "$log_dir_global"
log_init "${log_dir_global}/orchestrator.log"
log "experiment=${EXPERIMENT_NAME} run_id=${RUN_ID} platform=${PLATFORM} hardware=${HARDWARE}"
log "matrix_mode=${MATRIX_MODE} matrix_file=${MATRIX_FILE}"
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
health_port_for() {
local dp="$1"
if [[ "$dp" -gt 1 ]]; then
echo "$VLLM_DP_SUPERVISOR_PORT"
else
echo "$VLLM_PORT"
fi
}
is_server_healthy() {
local port="$1"
curl --fail --silent --show-error --max-time 5 "http://127.0.0.1:${port}/health" >/dev/null 2>&1
}
stop_server() {
local tp="$1"
local dp="$2"
local pid_file="/data/user1/yy/${EXPERIMENT}_vllm_tp${tp}_dp${dp}.pid"
if [[ -f "$pid_file" ]]; then
local pid
pid="$(cat "$pid_file")"
if kill -0 "$pid" 2>/dev/null; then
log "stopping vllm server pid=${pid} (tp=${tp}, dp=${dp})"
kill "$pid" 2>/dev/null || true
sleep 5
kill -9 "$pid" 2>/dev/null || true
fi
rm -f "$pid_file"
fi
# Fallback: kill any vllm serve processes for this model.
pkill -9 -f "vllm serve.*${MODEL_NAME}" 2>/dev/null || true
pkill -9 -f "vllm serve.*${MODEL_PATH}" 2>/dev/null || true
sleep 2
}
build_server_args() {
local tp="$1"
local dp="$2"
local args=(
"vllm serve" "$MODEL_PATH"
--trust-remote-code
--kv-cache-dtype fp8
--block-size 256
--tensor-parallel-size "$tp"
--no-enable-flashinfer-autotune
--port "$VLLM_PORT"
)
if [[ "$dp" -gt 1 ]]; then
args+=(
--data-parallel-size "$dp"
--data-parallel-size-local "$dp"
--data-parallel-multi-port-external-lb
)
fi
printf '%s ' "${args[@]}"
}
start_server() {
local tp="$1"
local dp="$2"
log "starting vllm server tp=${tp} dp=${dp}"
bash "${SCRIPT_DIR}/start_vllm_dp.sh" "$tp" "$dp" \
>> "${log_dir_global}/vllm_tp${tp}_dp${dp}.server.outer.log" 2>&1
local hport
hport="$(health_port_for "$dp")"
if ! is_server_healthy "$hport"; then
log "error: vllm server tp=${tp} dp=${dp} failed health check on port ${hport}"
return 1
fi
log "vllm server tp=${tp} dp=${dp} is healthy on port ${hport}"
}
run_warmup() {
local tp="$1"
local dp="$2"
local input_len="$3"
local output_len="$4"
local hport
hport="$(health_port_for "$dp")"
log "warming up tp=${tp} dp=${dp} (input=${input_len}, output=${output_len}, num=1)"
"$PYTHON" "${SCRIPT_DIR}/../../scripts/common/warmup.py" \
--backend vllm \
--host 127.0.0.1 \
--port "$hport" \
--input-len "$input_len" \
--output-len "$output_len" \
--num 1 \
--env-python "$PYTHON" \
>> "${log_dir_global}/vllm_tp${tp}_dp${dp}.warmup.log" 2>&1
log "warmup completed"
}
scenario_already_completed() {
local output_file="$1"
local expected="$2"
[[ -s "$output_file" ]] || return 1
local completed
completed="$("$PYTHON" -c "
import json, sys
path = sys.argv[1]
try:
with open(path, 'r', encoding='utf-8') as f:
for line in f:
line = line.strip()
if line:
data = json.loads(line)
print(data.get('completed', 0))
break
except Exception:
print(0)
" "$output_file")"
[[ "${completed:-0}" -ge "$expected" ]]
}
start_gpu_monitor() {
local csv_path="$1"
mkdir -p "$(dirname "$csv_path")"
nvidia-smi \
--query-gpu=timestamp,index,memory.used,memory.total,utilization.gpu \
--format=csv \
-l "$GPU_MEM_SAMPLE_INTERVAL_S" \
> "$csv_path" 2>/dev/null &
echo $!
}
stop_gpu_monitor() {
local pid="$1"
if kill -0 "$pid" 2>/dev/null; then
kill "$pid" 2>/dev/null || true
sleep 1
kill -9 "$pid" 2>/dev/null || true
fi
}
append_scenario_record() {
local result_root="$1"
local json_path="$result_root/results.json"
shift
# Remaining args are key=value pairs.
local scenario_json
scenario_json="$("$PYTHON" -c "
import json, sys
pairs = [a.split('=', 1) for a in sys.argv[1:]]
d = {k: json.loads(v) for k, v in pairs}
print(json.dumps(d, ensure_ascii=False))
" "$@")"
PYTHON="$PYTHON" append_scenario_to_json "$json_path" "$scenario_json"
}
# ---------------------------------------------------------------------------
# Per-configuration runner
# ---------------------------------------------------------------------------
run_parallel_config() {
local tp="$1"
local dp="$2"
local config_label="tp${tp}_dp${dp}"
local result_root="${RESULT_BASE}/${RUN_ID}/${config_label}"
local raw_dir="${result_root}/raw_outputs"
local gpu_log_dir="${result_root}/gpu_logs"
local phase_log_dir="${result_root}/logs"
mkdir -p "$raw_dir" "$gpu_log_dir" "$phase_log_dir"
log "===== ${config_label} START ====="
# Generate scenario list for this config.
local scenario_tsv="${result_root}/scenarios.tsv"
"$PYTHON" "${SCRIPT_DIR}/generate_scenarios.py" \
--matrix "$MATRIX_FILE" \
--mode "$MATRIX_MODE" \
> "$scenario_tsv"
local total_scenarios
total_scenarios="$(tail -n +2 "$scenario_tsv" | wc -l)"
log "generated ${total_scenarios} scenarios for ${config_label}"
if [[ "$total_scenarios" -eq 0 ]]; then
log "no scenarios for ${config_label}; skipping"
return 0
fi
# Write metadata.
ensure_result_root "$result_root"
write_metadata_json \
"${result_root}/results.json" \
"${EXPERIMENT_NAME}_${config_label}" \
"$RUN_ID" \
"$MODEL_PATH" \
"vllm" \
"vllm" \
"$HARDWARE" \
"$ACCELERATOR" \
"$CHIP" \
"experiments/${EXPERIMENT_NAME}/run_bench.sh" \
"$VENV_VLLM" \
"H200 vLLM TP×DP matrix for DeepSeek-V4-Flash"
# Embed static config and the exact server command line used.
local server_args_str
server_args_str="$(build_server_args "$tp" "$dp")"
jq --arg tp "$tp" --arg dp "$dp" --arg cuda "$CUDA_VISIBLE_DEVICES" --arg args "$server_args_str" \
'.config = {
"tp": ($tp | tonumber),
"dp": ($dp | tonumber),
"cuda_visible_devices": $cuda,
"backend": "vllm",
"server_start_script": "experiments/'${EXPERIMENT_NAME}'/start_vllm_dp.sh",
"server_args": $args
}' "${result_root}/results.json" > "${result_root}/results.json.tmp" && \
mv "${result_root}/results.json.tmp" "${result_root}/results.json"
# Group scenarios by input_len. We restart the server for each ISL to keep
# memory state isolated and to warm up per context length.
local current_isl=""
local group_started=false
tail -n +2 "$scenario_tsv" | while IFS=$'\t' read -r mark isl dsl conc num; do
# When the input length changes, restart the server for the new group.
if [[ "$isl" != "$current_isl" ]]; then
if [[ "$group_started" == true ]]; then
stop_server "$tp" "$dp"
fi
current_isl="$isl"
stop_server "$tp" "$dp"
if ! start_server "$tp" "$dp"; then
log "ERROR: ${config_label} failed to start for ISL=${isl}; skipping this group"
group_started=false
continue
fi
group_started=true
# Warmup with the shortest output for this ISL.
run_warmup "$tp" "$dp" "$isl" 128
fi
# If the server for this group did not start, skip all scenarios in it.
if [[ "$group_started" != true ]]; then
log "skipping ${config_label} ISL=${isl} scenario (server not started)"
continue
fi
local output_file="${raw_dir}/vllm_main_c${conc}_i${isl}_o${dsl}.jsonl"
local detail_log="${phase_log_dir}/vllm_${config_label}_c${conc}_i${isl}_o${dsl}.log"
local gpu_csv="${gpu_log_dir}/gpu_mem_c${conc}_i${isl}_o${dsl}.csv"
if scenario_already_completed "$output_file" "$num"; then
log "skipping already-completed ${config_label} scenario: c=${conc} i=${isl} o=${dsl}"
continue
fi
log "running ${config_label} scenario: mark=${mark} c=${conc} i=${isl} o=${dsl} n=${num}"
local gpu_pid
gpu_pid="$(start_gpu_monitor "$gpu_csv")"
local bench_rc=0
timeout "$SCENARIO_TIMEOUT_S" \
"$PYTHON" -m sglang.bench_serving \
--backend vllm \
--host 127.0.0.1 \
--port "$(health_port_for "$dp")" \
--dataset-name random \
--random-input-len "$isl" \
--random-output-len "$dsl" \
--num-prompts "$num" \
--max-concurrency "$conc" \
--request-rate 10000 \
--output-file "$output_file" \
--output-details \
> "$detail_log" 2>&1 || bench_rc=$?
stop_gpu_monitor "$gpu_pid"
if [[ "$bench_rc" -ne 0 ]]; then
log "ERROR: ${config_label} scenario c=${conc} i=${isl} o=${dsl} failed (rc=${bench_rc}); see ${detail_log}"
if [[ "$mark" == "P" ]]; then
log "optional (P) scenario failed; recording as skipped_oom and continuing"
append_scenario_record "$result_root" \
"name=c${conc}_i${isl}_o${dsl}" \
"config=$(jq -n --arg phase main --argjson c "$conc" --argjson i "$isl" --argjson o "$dsl" --arg dataset random --argjson n "$num" '{phase: $phase, concurrency: $c, input_len: $i, output_len: $o, dataset: $dataset, num_prompts: $n}')" \
"status=\"skipped_oom\"" \
"note=\"bench command failed or timed out (rc=${bench_rc})\""
continue
else
log "mandatory (Y) scenario failed; aborting current ISL group"
stop_server "$tp" "$dp"
group_started=false
continue
fi
fi
log "finished ${config_label} scenario: output=${output_file}"
done
stop_server "$tp" "$dp"
# Parse results.
log "parsing ${config_label} results"
"$PYTHON" "${SCRIPT_DIR}/../../scripts/common/parse_backend.py" "$result_root" --backend vllm \
>> "${phase_log_dir}/parse.log" 2>&1 || {
log "WARNING: parser failed for ${config_label}; see ${phase_log_dir}/parse.log"
}
log "===== ${config_label} DONE ====="
}
# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
# Cleanup any leftovers.
for cfg in "${PARALLEL_CONFIGS[@]}"; do
read -r tp dp <<< "$cfg"
stop_server "$tp" "$dp"
done
# Run each parallel configuration.
for cfg in "${PARALLEL_CONFIGS[@]}"; do
read -r tp dp <<< "$cfg"
run_parallel_config "$tp" "$dp"
done
# Generate cross-configuration comparison.
log "generating comparison report"
"$PYTHON" "${SCRIPT_DIR}/compare.py" \
--run-root "${RESULT_BASE}/${RUN_ID}" \
--output "${RESULT_BASE}/${RUN_ID}/comparison.md" \
>> "${log_dir_global}/compare.log" 2>&1 || {
log "WARNING: comparison script failed; see ${log_dir_global}/compare.log"
}
log "all results saved to ${RESULT_BASE}/${RUN_ID}"