Files
Gahow Wang 41a0c1c48f Migration correctness smoke tests: direct-read, partial-transfer, NIXL
Standalone smoke tests validating KV-migration correctness paths before
trace replay: full migrate-cache, partial-prefill transfer, and a
NIXL-connector variant, each with a runner.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-05-29 11:53:13 +08:00

76 lines
2.4 KiB
Bash

#!/usr/bin/env bash
# Fast iteration: start 2 vLLM kv_both, run smoke_test_migrate_cache, tear down.
#
# Usage: bash run_smoke_test.sh [WAIT_BETWEEN_S]
#
# Iteration overhead: ~3-4 min warmup + a few sec for the test. Cleanly
# tears everything down on exit so you can re-run repeatedly.
set -uo pipefail
PROJ_DIR="${PROJ_DIR:-/home/admin/cpfs/wjh/agentic-kv}"
MODEL="${MODEL:-/home/admin/cpfs/wjh/models/Qwen/Qwen3-Coder-30B-A3B-Instruct}"
VENV="$PROJ_DIR/.venv/bin"
LOGS_DIR="${LOGS_DIR:-$PROJ_DIR/outputs/smoke_test_$(date +%Y%m%d_%H%M%S)}"
mkdir -p "$LOGS_DIR"
# MOONCAKE_PROTOCOL controls Mooncake's C++ TransferEngine transport.
# Options exposed: rdma (default), tcp, nvlink_intra (NVLink intra-node).
PROTO="${MOONCAKE_PROTOCOL:-rdma}"
echo "[smoke] using Mooncake protocol: $PROTO"
cleanup() {
echo "[smoke] cleaning up vLLM..."
pkill -9 -f "vllm serve" 2>/dev/null || true
pkill -9 -f "EngineCore" 2>/dev/null || true
sleep 2
}
trap cleanup EXIT
cleanup
echo "[smoke] starting 2 vLLM kv_both on GPU 0,1"
for i in 0 1; do
port=$((8000 + i))
bp=$((8998 + i))
master=$((29500 + i))
PYTHONHASHSEED=42 \
VLLM_MOONCAKE_BOOTSTRAP_PORT=$bp \
CUDA_VISIBLE_DEVICES=$i \
MASTER_PORT=$master \
nohup "$VENV/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 \
--kv-transfer-config "{\"kv_connector\":\"MooncakeConnector\",\"kv_role\":\"kv_both\",\"kv_connector_extra_config\":{\"mooncake_protocol\":\"$PROTO\"}}" \
--enable-prompt-tokens-details \
> "$LOGS_DIR/vllm_inst_${i}_gpu${i}.log" 2>&1 &
disown
sleep 2
done
echo "[smoke] waiting for health on 8000 and 8001 ..."
for port in 8000 8001; do
tries=0
while ! curl -sf "http://127.0.0.1:$port/health" >/dev/null 2>&1; do
tries=$((tries+1))
if [ $tries -gt 180 ]; then
echo "[smoke] FATAL: $port not ready"; exit 1
fi
sleep 2
done
echo " port=$port ready"
done
echo "[smoke] running migration smoke test"
"$VENV/python" "$PROJ_DIR/microbench/connector_tax/cache_sweep/smoke_test_migrate_cache.py" \
--src-port 8000 --dst-port 8001 \
--src-bp 8998 --dst-bp 8999 \
${EXTRA_SMOKE_ARGS:-} \
2>&1 | tee "$LOGS_DIR/smoke_output.log"
ec=${PIPESTATUS[0]}
echo "[smoke] test exit=$ec, logs at $LOGS_DIR"
exit $ec