MB5 driver updates: PD-proxy + snapshot instrument + launcher tweaks

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
2026-05-29 11:53:27 +08:00
parent bad512d3c5
commit ee5db0b321
4 changed files with 153 additions and 17 deletions

View File

@@ -24,7 +24,9 @@
set -eo pipefail
FRESH_ROOT="/home/admin/cpfs/wjh/agentic-kv-fresh"
VENV="${FRESH_ROOT}/.venv"
# MB5_VENV lets a second host use an isolated venv clone (e.g. .venv_dash0) so
# two boxes can run in parallel without racing on the shared cpfs venv patch.
VENV="${MB5_VENV:-${FRESH_ROOT}/.venv}"
MODEL="${MODEL:-/home/admin/cpfs/wjh/models/Qwen/Qwen3-Coder-30B-A3B-Instruct}"
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
INSTRUMENT="${SCRIPT_DIR}/instrument_kv_snapshot.py"
@@ -86,14 +88,16 @@ case "${1:-start}" in
esac
# --- parse CONFIG into (prefill_gpus, decode_gpus) ----------------
USE_COLO_PROXY=0
case "${CONFIG}" in
8C) ROLES="combined"; P_GPUS=""; D_GPUS=""; COMBINED_GPUS="0,1,2,3,4,5,6,7" ;;
8C-proxy) ROLES="combined"; USE_COLO_PROXY=1; P_GPUS=""; D_GPUS=""; COMBINED_GPUS="0,1,2,3,4,5,6,7" ;;
6P+2D) ROLES="pd"; P_GPUS="0,1,2,3,4,5"; D_GPUS="6,7" ;;
5P+3D) ROLES="pd"; P_GPUS="0,1,2,3,4"; D_GPUS="5,6,7" ;;
4P+4D) ROLES="pd"; P_GPUS="0,1,2,3"; D_GPUS="4,5,6,7" ;;
3P+5D) ROLES="pd"; P_GPUS="0,1,2"; D_GPUS="3,4,5,6,7" ;;
2P+6D) ROLES="pd"; P_GPUS="0,1"; D_GPUS="2,3,4,5,6,7" ;;
*) echo "Unknown CONFIG=${CONFIG} (expected: 8C, 6P+2D, 5P+3D, 4P+4D, 3P+5D, 2P+6D)"; exit 1;;
*) echo "Unknown CONFIG=${CONFIG} (expected: 8C, 8C-proxy, 6P+2D, 5P+3D, 4P+4D, 3P+5D, 2P+6D)"; exit 1;;
esac
stop_all
@@ -137,6 +141,7 @@ launch_vllm() {
idx=0
proxy_args=()
colo_args=()
ENDPOINTS=""
case "${ROLES}" in
@@ -147,6 +152,7 @@ case "${ROLES}" in
bp=$((BASE_BP+idx))
launch_vllm "${idx}" "${gpu}" "${port}" "kv_both" "${bp}"
ENDPOINTS+="${ENDPOINTS:+,}http://127.0.0.1:${port}"
colo_args+=( --colo "http://127.0.0.1:${port}" )
idx=$((idx+1))
sleep 1
done
@@ -215,6 +221,25 @@ if [ "${ROLES}" = "pd" ]; then
ENDPOINTS="http://127.0.0.1:${PROXY_PORT}"
fi
if [ "${USE_COLO_PROXY}" = "1" ]; then
echo "[mb5] launching colo passthrough proxy on ${PROXY_PORT} (8 kv_both instances)"
nohup python "${PROXY_SRC}" "${colo_args[@]}" --port "${PROXY_PORT}" --host 0.0.0.0 \
> "${LOGS_DIR}/proxy.log" 2>&1 &
disown
tries=0
while ! curl -s -o /dev/null -w "%{http_code}" "http://127.0.0.1:${PROXY_PORT}/" 2>/dev/null | grep -qE "^[0-9]"; do
tries=$((tries+1))
if [ ${tries} -gt 60 ]; then
echo "[mb5] FATAL colo proxy did not come up in 2 min"
tail -40 "${LOGS_DIR}/proxy.log" || true
exit 1
fi
sleep 2
done
echo " colo proxy port=${PROXY_PORT} ready (HTTP responding)"
ENDPOINTS="http://127.0.0.1:${PROXY_PORT}"
fi
echo "[mb5] CONFIG=${CONFIG} RUN_LABEL=${RUN_LABEL} UP"
echo "ENDPOINTS=${ENDPOINTS}"
echo "RUN_ROOT=${RUN_ROOT}"