Skip to content

Commit 3ee3642

Browse files
committed
kvbm: sandbox scripts
Signed-off-by: Ryan Olson <[email protected]>
1 parent 3b8e61e commit 3ee3642

File tree

4 files changed

+310
-0
lines changed

4 files changed

+310
-0
lines changed
Lines changed: 126 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,126 @@
1+
#!/bin/bash
2+
# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
3+
# SPDX-License-Identifier: Apache-2.0
4+
5+
# Script to launch vLLM with the DynamoConnector for testing
6+
7+
set -e
8+
9+
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
10+
REPO_ROOT="$(dirname "$SCRIPT_DIR")"
11+
12+
# =============================================================================
13+
# Cleanup: Kill any previous vLLM processes
14+
# =============================================================================
15+
echo "🧹 Cleaning up previous processes..."
16+
pkill -9 -f "vllm.entrypoints" 2>/dev/null || true
17+
pkill -9 -f "EngineCore" 2>/dev/null || true
18+
sleep 2
19+
20+
# Check if port 8000 is in use
21+
if lsof -i :8000 >/dev/null 2>&1; then
22+
echo "⚠️ Port 8000 is still in use, killing processes..."
23+
fuser -k 8000/tcp 2>/dev/null || true
24+
sleep 1
25+
fi
26+
27+
# =============================================================================
28+
# Environment Setup
29+
# =============================================================================
30+
31+
# Find virtual environment
32+
if [ -d "$SCRIPT_DIR/venv" ]; then
33+
VENV_PATH="$SCRIPT_DIR/venv"
34+
elif [ -d "$SCRIPT_DIR/.venv" ] && [ ! -L "$SCRIPT_DIR/.venv" ]; then
35+
VENV_PATH="$SCRIPT_DIR/.venv"
36+
else
37+
echo "❌ No virtual environment found!"
38+
exit 1
39+
fi
40+
41+
# Activate the virtual environment
42+
source "$VENV_PATH/bin/activate"
43+
44+
# Ensure our kvbm package is in the path
45+
export PYTHONPATH="$REPO_ROOT/lib/bindings/kvbm/python:$PYTHONPATH"
46+
47+
# =============================================================================
48+
# Rust Logging Configuration
49+
# =============================================================================
50+
# Set RUST_LOG if not already set - enables tracing from Rust components
51+
# export RUST_LOG="${RUST_LOG:-dynamo_kvbm=debug,dynamo_nova=info,warn}"
52+
export RUST_LOG="${RUST_LOG:-debug}"
53+
export DYN_LOG="${DYN_LOG:-debug}"
54+
55+
# =============================================================================
56+
# vLLM Configuration
57+
# =============================================================================
58+
59+
# Default model - use gpt2 for quick testing
60+
MODEL="${MODEL:-gpt2}"
61+
62+
# Configure KV transfer for DynamoConnector
63+
# The connector is at: kvbm.v2.vllm.schedulers.connector.DynamoConnector
64+
kv_transfer_config='{
65+
"kv_connector": "DynamoConnector",
66+
"kv_role": "kv_both",
67+
"kv_connector_module_path": "kvbm.v2.vllm.schedulers.connector",
68+
"kv_connector_extra_config": {
69+
"leader": {
70+
"cache": { "host": { "cache_size_gb": 1.0 } },
71+
"tokio": { "worker_threads": 2 }
72+
},
73+
"worker": {
74+
"nixl": { "backends": { "UCX": {}, "POSIX": {} } },
75+
"tokio": { "worker_threads": 1 }
76+
}
77+
}
78+
}'
79+
80+
# =============================================================================
81+
# Display Configuration
82+
# =============================================================================
83+
echo ""
84+
echo "🚀 Launching vLLM with DynamoConnector"
85+
echo "=============================================================================="
86+
echo ""
87+
echo "Environment:"
88+
echo " SCRIPT_DIR: $SCRIPT_DIR"
89+
echo " REPO_ROOT: $REPO_ROOT"
90+
echo " VENV_PATH: $VENV_PATH"
91+
echo " Python: $(which python)"
92+
echo " vLLM version: $(python -c 'import vllm; print(vllm.__version__)')"
93+
echo " kvbm version: $(python -c 'import kvbm; print(kvbm.__version__)' 2>/dev/null || echo 'N/A')"
94+
echo ""
95+
echo "PYTHONPATH:"
96+
echo " $PYTHONPATH" | tr ':' '\n' | head -3
97+
echo ""
98+
echo "Rust Logging:"
99+
echo " RUST_LOG: $RUST_LOG"
100+
echo ""
101+
echo "vLLM Configuration:"
102+
echo " Model: $MODEL"
103+
echo " Host: 127.0.0.1"
104+
echo " Port: 8000"
105+
echo " Enforce Eager: true"
106+
echo ""
107+
echo "KV Transfer Config:"
108+
echo "$kv_transfer_config" | python -m json.tool 2>/dev/null || echo "$kv_transfer_config"
109+
echo ""
110+
echo "GPU Status:"
111+
nvidia-smi --query-gpu=index,name,memory.free,memory.total --format=csv,noheader 2>/dev/null || echo " nvidia-smi not available"
112+
echo ""
113+
echo "=============================================================================="
114+
echo ""
115+
116+
# =============================================================================
117+
# Launch vLLM
118+
# =============================================================================
119+
# Using --enforce-eager to avoid CUDA graph compilation overhead during testing
120+
exec python -m vllm.entrypoints.openai.api_server \
121+
--model "$MODEL" \
122+
--kv-transfer-config "$kv_transfer_config" \
123+
--enforce-eager \
124+
--host 127.0.0.1 \
125+
--port 8000 \
126+
"$@"

.sandbox/rebuild.sh

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
#!/bin/bash
2+
# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
3+
# SPDX-License-Identifier: Apache-2.0
4+
5+
# Quick rebuild script for the kvbm Python bindings
6+
# Uses local venv if exists, otherwise tries default location
7+
8+
set -e
9+
10+
echo "🔧 Rebuilding KVBM Python bindings..."
11+
echo "=================================="
12+
13+
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
14+
REPO_ROOT="$(dirname "$SCRIPT_DIR")"
15+
KVBM_BINDINGS_DIR="$REPO_ROOT/lib/bindings/kvbm"
16+
17+
# Find virtual environment
18+
# Priority: local venv > local .venv > symlinked .venv
19+
if [ -d "$SCRIPT_DIR/venv" ]; then
20+
VENV_PATH="$SCRIPT_DIR/venv"
21+
elif [ -d "$SCRIPT_DIR/.venv" ] && [ ! -L "$SCRIPT_DIR/.venv" ]; then
22+
VENV_PATH="$SCRIPT_DIR/.venv"
23+
elif [ -L "$SCRIPT_DIR/.venv" ] && [ -d "$(readlink -f "$SCRIPT_DIR/.venv")" ]; then
24+
VENV_PATH="$(readlink -f "$SCRIPT_DIR/.venv")"
25+
else
26+
echo "❌ No virtual environment found!"
27+
echo " Expected locations:"
28+
echo " - $SCRIPT_DIR/venv"
29+
echo " - $SCRIPT_DIR/.venv"
30+
echo ""
31+
echo " Create one with: python -m venv $SCRIPT_DIR/venv"
32+
exit 1
33+
fi
34+
35+
echo "📦 Using venv: $VENV_PATH"
36+
37+
# Activate the virtual environment
38+
source "$VENV_PATH/bin/activate"
39+
40+
# Navigate to the kvbm Python bindings directory
41+
cd "$KVBM_BINDINGS_DIR"
42+
43+
echo "📦 Building with maturin..."
44+
echo "Build mode: dev (default)"
45+
echo ""
46+
47+
# Build with maturin develop
48+
maturin develop
49+
50+
echo ""
51+
echo "✅ Build complete!"
52+
echo "Package: kvbm"
53+
54+
# Verify the build
55+
python -c "from kvbm._core import v2; print(f'kvbm._core.v2 loaded: {[a for a in dir(v2) if not a.startswith(\"_\")]}')"

.sandbox/test_cmpl_1.sh

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
#!/bin/bash
2+
3+
# Script to send a completion request to vLLM
4+
# Endpoint: http://127.0.0.1:8000/v1/completions
5+
6+
curl -X POST http://127.0.0.1:8000/v1/completions \
7+
-H "Content-Type: application/json" \
8+
-d '{
9+
"model": "gpt2",
10+
"prompt": "what is a dynamo? and how did it signify the start of the industrial revolution?",
11+
"max_tokens": 16,
12+
"temperature": 0
13+
}'

.sandbox/test_e2e.sh

Lines changed: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,116 @@
1+
#!/bin/bash
2+
# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
3+
# SPDX-License-Identifier: Apache-2.0
4+
5+
# All-in-one test script that:
6+
# 1. Rebuilds Python bindings
7+
# 2. Launches vLLM with connector
8+
# 3. Waits for server to be ready
9+
# 4. Sends test completion request
10+
# 5. Captures all output
11+
12+
set -e
13+
14+
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
15+
LOG_FILE="$SCRIPT_DIR/test_e2e_$(date +%Y%m%d_%H%M%S).log"
16+
17+
echo "===================================================================="
18+
echo "End-to-End KVBM Connector Test"
19+
echo "===================================================================="
20+
echo ""
21+
echo "Log file: $LOG_FILE"
22+
echo ""
23+
24+
# Function to cleanup on exit
25+
cleanup() {
26+
echo ""
27+
echo "Cleaning up..."
28+
pkill -9 -f "vllm.entrypoints" 2>/dev/null || true
29+
pkill -9 -f "EngineCore" 2>/dev/null || true
30+
}
31+
32+
trap cleanup EXIT
33+
34+
# Step 1: Rebuild bindings
35+
echo "Step 1: Rebuilding Python bindings..."
36+
echo "---------------------------------------------------------------------"
37+
if ! "$SCRIPT_DIR/rebuild.sh" 2>&1 | tee -a "$LOG_FILE"; then
38+
echo "❌ Rebuild failed!"
39+
exit 1
40+
fi
41+
42+
echo "✅ Rebuild complete"
43+
echo ""
44+
45+
# Step 2: Launch vLLM in background
46+
echo "Step 2: Launching vLLM with DynamoConnector..."
47+
echo "---------------------------------------------------------------------"
48+
49+
# Launch vLLM in background using the existing script
50+
"$SCRIPT_DIR/launch_vllm_with_connector.sh" 2>&1 >> "$LOG_FILE" &
51+
VLLM_PID=$!
52+
53+
echo "vLLM launched with PID: $VLLM_PID"
54+
echo ""
55+
56+
# Step 3: Wait for server to be ready
57+
echo "Step 3: Waiting for vLLM server to be ready on port 8000..."
58+
echo "---------------------------------------------------------------------"
59+
60+
MAX_WAIT=60
61+
WAIT_COUNT=0
62+
63+
while ! curl -s http://127.0.0.1:8000/health &>/dev/null; do
64+
if [ $WAIT_COUNT -ge $MAX_WAIT ]; then
65+
echo "❌ Timeout waiting for vLLM server (waited ${MAX_WAIT}s)"
66+
exit 1
67+
fi
68+
69+
# Check if process is still running
70+
if ! kill -0 $VLLM_PID 2>/dev/null; then
71+
echo "❌ vLLM process died during startup"
72+
exit 1
73+
fi
74+
75+
echo -n "."
76+
sleep 1
77+
WAIT_COUNT=$((WAIT_COUNT + 1))
78+
done
79+
80+
echo ""
81+
echo "✅ vLLM server is ready!"
82+
echo ""
83+
84+
# Step 4: Send test request
85+
echo "Step 4: Sending test completion request..."
86+
echo "---------------------------------------------------------------------"
87+
88+
RESPONSE=$("$SCRIPT_DIR/test_cmpl_1.sh" 2>&1 | tee -a "$LOG_FILE")
89+
90+
echo "$RESPONSE"
91+
echo ""
92+
93+
# Check if response contains error
94+
if echo "$RESPONSE" | grep -q '"error"'; then
95+
echo "❌ Request failed with error"
96+
97+
echo ""
98+
echo "Recent logs from vLLM:"
99+
echo "---------------------------------------------------------------------"
100+
tail -50 "$LOG_FILE" | grep -A 10 -i "error\|exception\|traceback" || tail -50 "$LOG_FILE"
101+
102+
exit 1
103+
else
104+
echo "✅ Request completed successfully"
105+
fi
106+
107+
echo ""
108+
echo "===================================================================="
109+
echo "Test Summary"
110+
echo "===================================================================="
111+
echo "✅ Bindings rebuilt"
112+
echo "✅ vLLM launched and initialized"
113+
echo "✅ Completion request processed"
114+
echo ""
115+
echo "Full logs saved to: $LOG_FILE"
116+
echo ""

0 commit comments

Comments
 (0)