Add TP2 prefill serving-path smoke experiment
This commit is contained in:
163
runs/frontier-tp2-prefill-serving-v0/run_prefill_profile.sh
Normal file
163
runs/frontier-tp2-prefill-serving-v0/run_prefill_profile.sh
Normal file
@@ -0,0 +1,163 @@
|
||||
#!/usr/bin/env bash
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
TP="${TP:?TP is required}"
|
||||
OUTPUT_ROOT="${OUTPUT_ROOT:?OUTPUT_ROOT is required}"
|
||||
SERVER_PORT="${SERVER_PORT:?SERVER_PORT is required}"
|
||||
VENV_ROOT="${VENV_ROOT:-/home/admin/cpfs/wjh/venvs/vllm-0.20.0-cu129-workload-regime-v2}"
|
||||
MODEL_ROOT="${MODEL_ROOT:-/home/admin/cpfs/wjh/models/Qwen/Qwen3-30B-A3B}"
|
||||
GPU_MEMORY_UTILIZATION="${GPU_MEMORY_UTILIZATION:-0.92}"
|
||||
SERVER_READY_ATTEMPTS="${SERVER_READY_ATTEMPTS:-900}"
|
||||
WARMUP_REQUESTS="${WARMUP_REQUESTS:-2}"
|
||||
SCRIPT_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
MATERIALIZER="${SCRIPT_ROOT}/materialize_prefill_request.py"
|
||||
CLIENT="${SCRIPT_ROOT}/run_request.py"
|
||||
SERVER_PID=""
|
||||
|
||||
IFS=',' read -r -a GPU_IDS <<< "${CUDA_VISIBLE_DEVICES:?GPU allocation is required}"
|
||||
if [[ "${#GPU_IDS[@]}" -ne "${TP}" ]]; then
|
||||
echo "ERROR: TP=${TP}, but CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES}" >&2
|
||||
exit 1
|
||||
fi
|
||||
if [[ ! "${WARMUP_REQUESTS}" =~ ^[1-9][0-9]*$ ]]; then
|
||||
echo "ERROR: WARMUP_REQUESTS must be a positive integer" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
mkdir -p "${OUTPUT_ROOT}"
|
||||
OUTPUT_ROOT="$(cd "${OUTPUT_ROOT}" && pwd)"
|
||||
mkdir -p "${OUTPUT_ROOT}/logs" "${OUTPUT_ROOT}/provenance" \
|
||||
"${OUTPUT_ROOT}/requests" "${OUTPUT_ROOT}/results" \
|
||||
"${OUTPUT_ROOT}/trace-staging" "${OUTPUT_ROOT}/traces/profile"
|
||||
exec > >(tee -a "${OUTPUT_ROOT}/logs/controller.log") 2>&1
|
||||
|
||||
cleanup() {
|
||||
if [[ -n "${SERVER_PID}" ]] && kill -0 "${SERVER_PID}" 2>/dev/null; then
|
||||
kill -TERM -- "-${SERVER_PID}" 2>/dev/null || true
|
||||
for _ in $(seq 1 30); do
|
||||
kill -0 "${SERVER_PID}" 2>/dev/null || break
|
||||
sleep 1
|
||||
done
|
||||
kill -KILL -- "-${SERVER_PID}" 2>/dev/null || true
|
||||
fi
|
||||
SERVER_PID=""
|
||||
}
|
||||
trap cleanup EXIT INT TERM
|
||||
|
||||
sha256sum "${BASH_SOURCE[0]}" "${MATERIALIZER}" "${CLIENT}" \
|
||||
"${MODEL_ROOT}/config.json" > "${OUTPUT_ROOT}/provenance/inputs.sha256"
|
||||
"${VENV_ROOT}/bin/python" -c \
|
||||
'import torch, transformers, vllm; print(f"torch={torch.__version__}"); print(f"transformers={transformers.__version__}"); print(f"vllm={vllm.__version__}")' \
|
||||
> "${OUTPUT_ROOT}/provenance/runtime.versions"
|
||||
nvidia-smi --query-gpu=index,name,uuid,driver_version,memory.total \
|
||||
--format=csv,noheader > "${OUTPUT_ROOT}/provenance/gpus.before.csv"
|
||||
ps -eo user,pid,ppid,etimes,pcpu,pmem,args --sort=pid \
|
||||
> "${OUTPUT_ROOT}/provenance/processes.before.txt"
|
||||
uptime > "${OUTPUT_ROOT}/provenance/uptime.before.txt"
|
||||
env | sort > "${OUTPUT_ROOT}/provenance/environment.txt"
|
||||
|
||||
PROFILE_CONFIG="$("${VENV_ROOT}/bin/python" - "${OUTPUT_ROOT}/trace-staging" <<'PY'
|
||||
import json
|
||||
import sys
|
||||
print(json.dumps({
|
||||
"profiler": "torch",
|
||||
"torch_profiler_dir": sys.argv[1],
|
||||
"torch_profiler_with_stack": False,
|
||||
"torch_profiler_record_shapes": True,
|
||||
"torch_profiler_use_gzip": True,
|
||||
"ignore_frontend": True,
|
||||
"wait_iterations": 0,
|
||||
"warmup_iterations": 0,
|
||||
"active_iterations": 2,
|
||||
}, separators=(",", ":")))
|
||||
PY
|
||||
)"
|
||||
printf '%s\n' "${PROFILE_CONFIG}" > "${OUTPUT_ROOT}/provenance/profiler-config.json"
|
||||
|
||||
export TOKENIZERS_PARALLELISM=false
|
||||
export VLLM_USE_V1=1
|
||||
export HF_HUB_OFFLINE=1
|
||||
export TRANSFORMERS_OFFLINE=1
|
||||
export FLASHINFER_WORKSPACE_BASE="${FLASHINFER_WORKSPACE_BASE:-${OUTPUT_ROOT}/flashinfer-workspace}"
|
||||
mkdir -p "${FLASHINFER_WORKSPACE_BASE}"
|
||||
ulimit -n 65536
|
||||
|
||||
"${VENV_ROOT}/bin/python" "${MATERIALIZER}" --model "${MODEL_ROOT}" \
|
||||
--input-tokens 8192 --output-tokens 2 \
|
||||
--output "${OUTPUT_ROOT}/requests/q8192-o2.json"
|
||||
|
||||
printf 'LAUNCH host=%s tp=%s gpus=%s output=%s warmups=%s\n' \
|
||||
"$(hostname)" "${TP}" "${CUDA_VISIBLE_DEVICES}" "${OUTPUT_ROOT}" \
|
||||
"${WARMUP_REQUESTS}"
|
||||
setsid "${VENV_ROOT}/bin/vllm" serve "${MODEL_ROOT}" \
|
||||
--host 127.0.0.1 --port "${SERVER_PORT}" \
|
||||
--served-model-name qwen30-prefill-profile \
|
||||
--tensor-parallel-size "${TP}" \
|
||||
--gpu-memory-utilization "${GPU_MEMORY_UTILIZATION}" \
|
||||
--max-model-len 40960 --max-num-batched-tokens 8192 --max-num-seqs 16 \
|
||||
--no-enable-prefix-caching --enable-chunked-prefill --no-enable-log-requests \
|
||||
--enable-logging-iteration-details --profiler-config "${PROFILE_CONFIG}" \
|
||||
> "${OUTPUT_ROOT}/logs/server.log" 2>&1 &
|
||||
SERVER_PID=$!
|
||||
|
||||
READY=0
|
||||
for _ in $(seq 1 "${SERVER_READY_ATTEMPTS}"); do
|
||||
if curl -fsS --max-time 2 \
|
||||
"http://127.0.0.1:${SERVER_PORT}/v1/models" \
|
||||
> "${OUTPUT_ROOT}/results/models.json" 2>/dev/null; then
|
||||
READY=1
|
||||
break
|
||||
fi
|
||||
if ! kill -0 "${SERVER_PID}" 2>/dev/null; then
|
||||
tail -200 "${OUTPUT_ROOT}/logs/server.log"
|
||||
exit 1
|
||||
fi
|
||||
sleep 3
|
||||
done
|
||||
if [[ "${READY}" -ne 1 ]]; then
|
||||
echo "ERROR: server readiness timeout" >&2
|
||||
tail -200 "${OUTPUT_ROOT}/logs/server.log"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
for index in $(seq 1 "${WARMUP_REQUESTS}"); do
|
||||
"${VENV_ROOT}/bin/python" "${CLIENT}" --port "${SERVER_PORT}" \
|
||||
--request "${OUTPUT_ROOT}/requests/q8192-o2.json" \
|
||||
--output "${OUTPUT_ROOT}/results/warmup-${index}.json"
|
||||
done
|
||||
|
||||
curl -fsS -X POST "http://127.0.0.1:${SERVER_PORT}/start_profile" \
|
||||
> "${OUTPUT_ROOT}/logs/start-profile.txt"
|
||||
"${VENV_ROOT}/bin/python" "${CLIENT}" --port "${SERVER_PORT}" \
|
||||
--request "${OUTPUT_ROOT}/requests/q8192-o2.json" \
|
||||
--output "${OUTPUT_ROOT}/results/profile.json"
|
||||
curl -fsS -X POST "http://127.0.0.1:${SERVER_PORT}/stop_profile" \
|
||||
> "${OUTPUT_ROOT}/logs/stop-profile.txt"
|
||||
|
||||
deadline=$((SECONDS + 120))
|
||||
while (( SECONDS < deadline )); do
|
||||
trace_count="$(find "${OUTPUT_ROOT}/trace-staging" -maxdepth 1 -type f \
|
||||
-name '*.pt.trace.json*' | wc -l)"
|
||||
if (( trace_count >= TP )); then
|
||||
break
|
||||
fi
|
||||
sleep 2
|
||||
done
|
||||
trace_count="$(find "${OUTPUT_ROOT}/trace-staging" -maxdepth 1 -type f \
|
||||
-name '*.pt.trace.json*' | wc -l)"
|
||||
if (( trace_count < TP )); then
|
||||
echo "ERROR: expected ${TP} rank traces, found ${trace_count}" >&2
|
||||
exit 1
|
||||
fi
|
||||
find "${OUTPUT_ROOT}/trace-staging" -maxdepth 1 -type f \
|
||||
-name '*.pt.trace.json*' -exec mv -t "${OUTPUT_ROOT}/traces/profile" {} +
|
||||
|
||||
cleanup
|
||||
nvidia-smi --query-gpu=index,name,uuid,driver_version,memory.total \
|
||||
--format=csv,noheader > "${OUTPUT_ROOT}/provenance/gpus.after.csv"
|
||||
find "${OUTPUT_ROOT}" -type f \
|
||||
! -path '*/provenance/artifacts.sha256' -print0 \
|
||||
| sort -z | xargs -0 sha256sum \
|
||||
> "${OUTPUT_ROOT}/provenance/artifacts.sha256"
|
||||
echo PREFILL_PROFILE_COMPLETE
|
||||
Reference in New Issue
Block a user