gahow/agentic-kvc
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add unified_nixl_both policy: NIXL connector isolation control

Browse Source 
Adds a NIXL-backed counterpart to unified_kv_both so we can attribute
the kv_both substrate overhead measured in the elastic_migration_v2
section to either Mooncake-specific code or a generic v1-connector
cost shared by all connectors.

- scripts/cache_aware_proxy.py: register --policy unified_nixl_both.
  Picker is identical to unified (and unified_kv_both); routing
  decisions never go through the PD-sep branch. Differs only at the
  vLLM launch layer.
- scripts/b3_isolated_policy.sh: new KV_CONNECTOR env var
  (Mooncake|Nixl), auto-set based on POLICY. NIXL launch path uses
  --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_both"}'
  with no VLLM_MOONCAKE_BOOTSTRAP_PORT (NIXL uses UCX side-channels).
- Health-check timeout: 90 iterations * 2s -> 180 iterations * 2s
  (180s -> 360s). Empirically NIXL needs ~100-150s per instance to
  initialize the UCX agent and register KV cache memory; 8
  concurrent NIXL launches frequently overshoot the previous 180s
  budget. Mooncake is unaffected (still finishes well inside the new
  budget). The 8-vLLM unified_nixl_both first launch tripped the
  old timeout despite 7/8 instances reaching startup-complete.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
Gahow Wang
2026-05-26 14:57:54 +08:00
parent 06dd175441
commit 151bf33541
2 changed files with 51 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

40
scripts/b3_isolated_policy.sh
View File

@@ -31,8 +31,15 @@ TRACE="${2:?usage: $0 <policy> <trace> <rundir>}"
RUNDIR="${3:?usage: $0 <policy> <trace> <rundir>}"
# Auto-enable kv_both when the policy requires it.
# KV_CONNECTOR (Mooncake|Nixl) selects the underlying connector when KV_BOTH=1.
KV_CONNECTOR="${KV_CONNECTOR:-Mooncake}"
if [ "$POLICY" = "unified_v2" ] || [ "$POLICY" = "unified_kv_both" ]; then
ENABLE_KV_BOTH=1
KV_CONNECTOR="Mooncake"
fi
if [ "$POLICY" = "unified_nixl_both" ]; then
ENABLE_KV_BOTH=1
KV_CONNECTOR="Nixl"
fi
mkdir -p "$RUNDIR/engine_state" "$RUNDIR/logs"
@@ -49,17 +56,13 @@ trap cleanup EXIT
# Hard reset first
cleanup
echo "[isolated] launching $N_INSTANCES vLLM on GPUs $GPU_INDICES ENABLE_KV_BOTH=$ENABLE_KV_BOTH ..."
echo "[isolated] launching $N_INSTANCES vLLM on GPUs $GPU_INDICES ENABLE_KV_BOTH=$ENABLE_KV_BOTH KV_CONNECTOR=$KV_CONNECTOR ..."
i=0
kv_both_extra=""
if [ "$ENABLE_KV_BOTH" = "1" ]; then
kv_both_extra="--kv-transfer-config {\"kv_connector\":\"MooncakeConnector\",\"kv_role\":\"kv_both\"}"
fi
for gpu in $GPU_INDICES; do
port=$((BASE_PORT + i))
master=$((29500 + i))
bp=$((BOOTSTRAP_BASE_PORT + i))
if [ "$ENABLE_KV_BOTH" = "1" ]; then
if [ "$ENABLE_KV_BOTH" = "1" ] && [ "$KV_CONNECTOR" = "Mooncake" ]; then
PYTHONHASHSEED=42 \
VLLM_MOONCAKE_BOOTSTRAP_PORT=$bp \
AGENTIC_STEP_LOG_PATH="$RUNDIR/engine_state/engine_${i}.jsonl" \
@@ -75,6 +78,23 @@ for gpu in $GPU_INDICES; do
--kv-transfer-config '{"kv_connector":"MooncakeConnector","kv_role":"kv_both"}' \
$EXTRA_VLLM_ARGS \
> "$RUNDIR/logs/vllm_inst_${i}_gpu${gpu}.log" 2>&1 &
elif [ "$ENABLE_KV_BOTH" = "1" ] && [ "$KV_CONNECTOR" = "Nixl" ]; then
# NixlConnector uses UCX side-channels for handshake (no bootstrap
# port needed). Side-channel host defaults to NIC IP discovered
# at register_kv_caches time.
AGENTIC_STEP_LOG_PATH="$RUNDIR/engine_state/engine_${i}.jsonl" \
AGENTIC_WORKER_ID="engine_${i}" \
CUDA_VISIBLE_DEVICES=$gpu \
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":"NixlConnector","kv_role":"kv_both"}' \
$EXTRA_VLLM_ARGS \
> "$RUNDIR/logs/vllm_inst_${i}_gpu${gpu}.log" 2>&1 &
else
AGENTIC_STEP_LOG_PATH="$RUNDIR/engine_state/engine_${i}.jsonl" \
AGENTIC_WORKER_ID="engine_${i}" \
@@ -95,13 +115,17 @@ for gpu in $GPU_INDICES; do
done
echo "[isolated] waiting for vLLM health ..."
# NIXL init takes ~100-150s per instance even with concurrent launches;
# Mooncake is closer to ~30-60s. Use a generous 360s timeout to cover
# both (90s -> 360s vs the previous 180s).
HEALTH_MAX_TRIES=180
for i in $(seq 0 $((N_INSTANCES - 1))); do
port=$((BASE_PORT + i))
tries=0
while ! curl -sf "http://127.0.0.1:$port/health" >/dev/null 2>&1; do
tries=$((tries + 1))
if [ $tries -gt 90 ]; then
echo "[isolated] FATAL: inst_$i not healthy after 180s"
if [ $tries -gt $HEALTH_MAX_TRIES ]; then
echo "[isolated] FATAL: inst_$i not healthy after $((HEALTH_MAX_TRIES * 2))s"
exit 1
fi
sleep 2