115 lines
4.2 KiB
Bash
115 lines
4.2 KiB
Bash
#!/usr/bin/env bash
|
|
# MB7 launcher (runs on dash0). Two 2-instance modes selected by MODE env:
|
|
# MODE=baseline : restore stock connector, no layerwise env
|
|
# MODE=layerwise : deploy mooncake_connector.LAYERWISE.py + MOONCAKE_LAYERWISE=1
|
|
#
|
|
# Chunked prefill is DISABLED (max-num-batched-tokens >= max prompt) so the
|
|
# producer prefill is a single forward and save_kv_layer fires once per layer
|
|
# in order — the layer-wise counter assumes this.
|
|
#
|
|
# The connector is always restored from .ORIG_BACKUP on exit.
|
|
#
|
|
# Usage (on dash0):
|
|
# MODE=baseline bash run_mb7.sh
|
|
# MODE=layerwise bash run_mb7.sh
|
|
|
|
set -uo pipefail
|
|
|
|
MODE="${MODE:-baseline}"
|
|
PROJ_DIR="${PROJ_DIR:-/home/admin/cpfs/wjh/agentic-kv}"
|
|
VENV="${VENV:-$PROJ_DIR/.venv}"
|
|
MODEL="${MODEL:-/home/admin/cpfs/wjh/models/Qwen/Qwen3-Coder-30B-A3B-Instruct}"
|
|
GPUS=(${GPUS:-0 1})
|
|
SIZES="${SIZES:-8192,16384,32768}"
|
|
REPEATS="${REPEATS:-3}"
|
|
BG_LOAD="${BG_LOAD:-0}"
|
|
CONCURRENT="${CONCURRENT:-1}"
|
|
MAX_BATCHED="${MAX_BATCHED:-40960}" # >= max prompt => no chunked prefill
|
|
DATE="$(date +%Y%m%d_%H%M)"
|
|
OUTDIR="${OUTDIR:-$PROJ_DIR/outputs/mb7_${MODE}_${DATE}}"
|
|
PYTHON="$VENV/bin/python"
|
|
MC_FILE="$VENV/lib/python3.12/site-packages/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py"
|
|
LW_SRC="${LW_SRC:-/tmp/mooncake_connector.LAYERWISE.py}"
|
|
DRIVER="$PROJ_DIR/microbench/connector_tax/layerwise/mb7_layerwise.py"
|
|
|
|
mkdir -p "$OUTDIR/logs"
|
|
PORTS=(8000 8001); BPS=(8998 8999)
|
|
|
|
echo "=== MB7 ($MODE) ==="
|
|
echo "Out: $OUTDIR ; connector: $MC_FILE"
|
|
|
|
restore_connector() {
|
|
if [ -f "$MC_FILE.ORIG_BACKUP" ]; then
|
|
cp -f "$MC_FILE.ORIG_BACKUP" "$MC_FILE"
|
|
echo "[restore] connector reset to ORIG"
|
|
fi
|
|
}
|
|
cleanup() {
|
|
pkill -9 -f "vllm serve" 2>/dev/null || true
|
|
pkill -9 -f "EngineCore" 2>/dev/null || true
|
|
sleep 4
|
|
restore_connector
|
|
}
|
|
trap cleanup EXIT
|
|
pkill -9 -f "vllm serve" 2>/dev/null || true; sleep 3
|
|
|
|
# Deploy the connector for the chosen mode.
|
|
if [ "$MODE" = "layerwise" ]; then
|
|
if [ ! -f "$LW_SRC" ]; then echo "FATAL: $LW_SRC not found (scp it first)"; exit 1; fi
|
|
cp -f "$LW_SRC" "$MC_FILE"
|
|
"$PYTHON" -c "import ast; ast.parse(open('$MC_FILE').read()); print('[deploy] LAYERWISE connector AST OK')" || exit 1
|
|
LW_ENV="MOONCAKE_LAYERWISE=1"
|
|
else
|
|
restore_connector
|
|
LW_ENV=""
|
|
fi
|
|
|
|
echo "[launch] 2 instances (max-num-batched-tokens=$MAX_BATCHED, chunked-prefill off)"
|
|
i=0
|
|
for gpu in "${GPUS[@]:0:2}"; do
|
|
port=${PORTS[$i]}; bp=${BPS[$i]}; master=$((29700 + i))
|
|
env $LW_ENV \
|
|
PYTHONHASHSEED=42 VLLM_MOONCAKE_BOOTSTRAP_PORT=$bp \
|
|
CUDA_VISIBLE_DEVICES=$gpu MASTER_PORT=$master \
|
|
nohup "$VENV/bin/vllm" serve "$MODEL" \
|
|
--host 0.0.0.0 --port "$port" --tensor-parallel-size 1 \
|
|
--trust-remote-code --enable-prefix-caching --dtype auto \
|
|
--gpu-memory-utilization 0.9 --max-model-len 200000 \
|
|
--max-num-batched-tokens "$MAX_BATCHED" \
|
|
--kv-transfer-config '{"kv_connector":"MooncakeConnector","kv_role":"kv_both"}' \
|
|
--enable-prompt-tokens-details \
|
|
> "$OUTDIR/logs/vllm_${i}_gpu${gpu}.log" 2>&1 &
|
|
disown; sleep 2; i=$((i + 1))
|
|
done
|
|
|
|
echo "[health] waiting ..."
|
|
for i in 0 1; do
|
|
port=${PORTS[$i]}; tries=0
|
|
while ! curl -sf "http://127.0.0.1:$port/health" >/dev/null 2>&1; do
|
|
tries=$((tries + 1)); [ $tries -gt 180 ] && { echo "FATAL inst_$i"; exit 1; }
|
|
sleep 2
|
|
done
|
|
echo " inst_$i ready"
|
|
done
|
|
for i in 0 1; do
|
|
bp=${BPS[$i]}; tries=0
|
|
while ! curl -sf "http://127.0.0.1:$bp/query" >/dev/null 2>&1; do
|
|
tries=$((tries+1)); [ $tries -gt 60 ] && { echo "WARN bp $bp"; break; }; sleep 2
|
|
done
|
|
done
|
|
|
|
echo "[run] mb7 --mode $MODE"
|
|
"$PYTHON" "$DRIVER" --mode "$MODE" \
|
|
--src-port "${PORTS[0]}" --dst-port "${PORTS[1]}" \
|
|
--src-bp "${BPS[0]}" --dst-bp "${BPS[1]}" \
|
|
--sizes "$SIZES" --repeats "$REPEATS" --bg-load "$BG_LOAD" \
|
|
--concurrent "$CONCURRENT" --out "$OUTDIR/mb7_result.json" \
|
|
2>&1 | tee "$OUTDIR/mb7_run.txt"
|
|
|
|
echo "[done] $OUTDIR"
|
|
# grep layerwise transfer logs from the producer (gpu0) for sanity
|
|
if [ "$MODE" = "layerwise" ]; then
|
|
echo "=== producer layerwise log lines ==="
|
|
grep -i "layerwise" "$OUTDIR/logs/vllm_0_gpu${GPUS[0]}.log" | tail -10 || true
|
|
fi
|