minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[smg][ci]: migrate benchmarks to e2e_test/benchmarks/, use parent conftest (#16597)

Browse Source
This commit is contained in:
Simo Lin
2026-01-06 20:15:20 -08:00
committed by GitHub
parent 913b688f21
commit d8b8198192
14 changed files with 1043 additions and 1963 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/pr-test-pd-router.yml
-616
View File
@@ -1,616 +0,0 @@
name: PR Benchmark (SMG PD Router)
on:
push:
branches: [ main ]
paths:
- 'python/sglang/srt/disaggregation/**'
- 'scripts/ci/ci_start_disaggregation_servers.sh'
- 'sgl-model-gateway/**'
pull_request:
branches: [ main ]
types: [opened, synchronize, reopened, labeled]
paths:
- 'python/sglang/srt/disaggregation/**'
- 'scripts/ci/ci_start_disaggregation_servers.sh'
- 'sgl-model-gateway/**'
workflow_dispatch:
concurrency:
group: test-disaggregation-${{ github.ref }}
cancel-in-progress: true
env:
RUSTC_WRAPPER: sccache
SCCACHE_GHA_ENABLED: "true"
permissions:
contents: read
pull-requests: write
issues: write
jobs:
test-disaggregation:
# TEMPORARILY DISABLED: 8-gpu-h200-oracle node is broken
# TODO: Re-enable when node is fixed
if: false
# Original condition:
# if: |
# github.event_name != 'pull_request' ||
# (github.event.action != 'labeled' && contains(github.event.pull_request.labels.*.name, 'run-ci') && contains(github.event.pull_request.labels.*.name, 'router-benchmark')) ||
# (github.event.action == 'labeled' && github.event.label.name == 'run-ci' && contains(github.event.pull_request.labels.*.name, 'router-benchmark')) ||
# (github.event.action == 'labeled' && github.event.label.name == 'router-benchmark' && contains(github.event.pull_request.labels.*.name, 'run-ci'))
runs-on: [8-gpu-h200-oracle]
timeout-minutes: 45
steps:
- name: Checkout code
uses: actions/checkout@v4
with:
fetch-depth: 10
- name: Setup Python
uses: actions/setup-python@v4
with:
python-version: '3.12'
- name: Setup Rust
run: |
bash scripts/ci/ci_install_rust.sh
- name: Configure sccache
uses: mozilla-actions/sccache-action@v0.0.9
with:
version: "v0.12.0"
- name: Rust cache
uses: Swatinem/rust-cache@v2
with:
workspaces: sgl-model-gateway
shared-key: "rust-cache"
cache-all-crates: true
cache-on-failure: true
save-if: true
- name: Cache pip dependencies
uses: actions/cache@v4
with:
path: ~/.cache/pip
key: ${{ runner.os }}-pip-${{ hashFiles('python/pyproject.toml') }}
restore-keys: |
${{ runner.os }}-pip-
- name: Validate environment
run: |
echo "=== System Validation ==="
nvidia-smi
echo "GPU count: $(nvidia-smi -L | wc -l)"
if [ $(nvidia-smi -L | wc -l) -lt 8 ]; then
echo "Error: This test requires at least 8 GPUs"
exit 1
fi
echo "=== GPU Process Check ==="
# Fail fast if any GPU compute processes are active
if command -v nvidia-smi >/dev/null 2>&1; then
# Try to query compute apps first (preferred and concise)
gpu_procs=$(nvidia-smi --query-compute-apps=pid,process_name,gpu_uuid --format=csv,noheader 2>/dev/null | sed '/^$/d' || true)
# Fallback to detailed PIDS report if the query returns nothing but there might still be processes
if [ -z "$gpu_procs" ]; then
gpu_procs=$(nvidia-smi -q -d PIDS 2>/dev/null | awk '/Processes/{flag=1;next}/^$/{flag=0}flag' | sed '/^\s*Processes:/d' | sed '/^\s*$/d' || true)
fi
if [ -n "$gpu_procs" ]; then
echo "Error: Found active GPU processes using the device(s):"
echo "$gpu_procs"
exit 1
else
echo "No active GPU compute processes detected."
fi
else
echo "Error: nvidia-smi not found; skipping GPU process check."
exit 1
fi
echo "=== RDMA Validation ==="
if ! command -v ibv_devices >/dev/null 2>&1; then
echo "Error: InfiniBand tools not found"
exit 1
fi
# Check for active IB devices
found_active_device=false
for device in mlx5_{0..11}; do
if ibv_devinfo $device >/dev/null 2>&1; then
state=$(ibv_devinfo $device | grep "state:" | head -1 | awk '{print $2}')
if [[ "$state" == "PORT_ACTIVE" ]]; then
echo "✓ Found active device: $device"
found_active_device=true
break
fi
fi
done
if [ "$found_active_device" = false ]; then
echo "Error: No active IB devices found"
echo "Available devices:"
ibv_devices || true
exit 1
fi
echo "=== Model Validation ==="
if [ ! -d "/raid/models/meta-llama/Llama-3.1-8B-Instruct" ]; then
echo "Error: Model not found"
ls -la /raid/models/ || echo "No models directory"
exit 1
fi
echo "✓ Model found"
- name: Install SGLang dependencies
run: |
echo "Installing SGLang with all extras..."
python3 -m pip --no-cache-dir install --upgrade pip
python3 -m pip --no-cache-dir install torch==2.9.1 torchvision torchaudio --index-url https://download.pytorch.org/whl/test/cu128
python3 -m pip --no-cache-dir install -e "python[all]" --break-system-packages
python3 -m pip --no-cache-dir install mooncake-transfer-engine==0.3.8
python3 -m pip --no-cache-dir install --user --force-reinstall genai-bench==0.0.3
- name: Build and install sgl-model-gateway
run: |
source "$HOME/.cargo/env"
echo "Building sgl-model-gateway..."
cd sgl-model-gateway/bindings/python
pip install maturin
maturin build --release --out dist --features vendored-openssl
pip install --force-reinstall dist/*.whl
- name: Start disaggregation servers
id: start_servers
run: |
echo "Starting disaggregation servers..."
READY_FILE=".disagg_ready"
rm -f "$READY_FILE"
DISAGG_READY_FILE="$READY_FILE" bash scripts/ci/ci_start_disaggregation_servers.sh &
SERVER_PID=$!
echo "server_pid=$SERVER_PID" >> $GITHUB_OUTPUT
# Wait until script signals readiness (8/8 healthy) or timeout
TIMEOUT=300
ELAPSED=0
while [ $ELAPSED -lt $TIMEOUT ]; do
if [ -f "$READY_FILE" ]; then
echo "✓ All disaggregation servers are healthy (signal detected)"
break
fi
if ! ps -p $SERVER_PID > /dev/null; then
echo "Error: server bootstrap script exited prematurely"
exit 1
fi
sleep 5
ELAPSED=$((ELAPSED + 5))
done
if [ $ELAPSED -ge $TIMEOUT ]; then
echo "❌ Timeout waiting for disaggregation servers to be healthy"
exit 1
fi
echo "✓ Servers started (PID: $SERVER_PID)"
- name: Test all policies sequentially
timeout-minutes: 30
run: |
POLICIES=("random" "round_robin" "cache_aware" "power_of_two")
BASE_URL="http://127.0.0.9:8000"
# Free commonly used ports for router and metrics
echo "Freeing ports 29000 (metrics) and 8000 (API), if in use..."
fuser -k -n tcp 29000 2>/dev/null || true
fuser -k -n tcp 8000 2>/dev/null || true
sleep 1
for policy in "${POLICIES[@]}"; do
echo ""
echo "=================================================="
echo "Testing policy: $policy"
echo "=================================================="
# Free ports before starting router
fuser -k -n tcp 29000 2>/dev/null || true
fuser -k -n tcp 8000 2>/dev/null || true
# Start router with the current policy
echo "Starting router with policy: $policy..."
RUST_BACKTRACE=1 python3 -m sglang_router.launch_router \
--pd-disaggregation \
--policy "$policy" \
--prefill http://127.0.0.1:30001 9001 \
--prefill http://127.0.0.2:30002 9002 \
--prefill http://127.0.0.3:30003 9003 \
--prefill http://127.0.0.4:30004 9004 \
--decode http://127.0.0.5:30005 \
--decode http://127.0.0.6:30006 \
--decode http://127.0.0.7:30007 \
--decode http://127.0.0.8:30008 \
--host 127.0.0.9 \
--log-level warn \
--port 8000 &
ROUTER_PID=$!
# Wait for router to become healthy
echo "Waiting for router to become healthy..."
TIMEOUT=60
ELAPSED=0
while [ $ELAPSED -lt $TIMEOUT ]; do
if curl --connect-timeout 5 --silent http://127.0.0.9:8000 > /dev/null 2>&1; then
echo "✓ Router is reachable"
break
fi
if ! ps -p $ROUTER_PID > /dev/null; then
echo "Error: Router process died"
exit 1
fi
sleep 5
ELAPSED=$((ELAPSED + 5))
done
if [ $ELAPSED -ge $TIMEOUT ]; then
echo "Error: Router health check timeout"
kill $ROUTER_PID 2>/dev/null || true
exit 1
fi
# Test API functionality
echo "Testing API completions for $policy..."
response=$(curl -s -X POST "$BASE_URL/v1/chat/completions" \
-H "Content-Type: application/json" \
-H "Authorization: Bearer test-token" \
-d '{
"model": "/raid/models/meta-llama/Llama-3.1-8B-Instruct",
"messages": [
{"role": "user", "content": "Write a Python function to calculate fibonacci numbers recursively"}
],
"stream": false,
"max_completion_tokens": 100
}')
if echo "$response" | jq -e '.choices[0].message.content' > /dev/null 2>&1; then
echo "✓ API test passed for $policy"
else
echo "✗ API test failed for $policy: $response"
kill $ROUTER_PID 2>/dev/null || true
exit 1
fi
# Test streaming
echo "Testing streaming API for $policy..."
stream_response=$(timeout 30 curl -s -X POST "$BASE_URL/v1/chat/completions" \
-H "Content-Type: application/json" \
-H "Authorization: Bearer test-token" \
-d '{
"model": "/raid/models/meta-llama/Llama-3.1-8B-Instruct",
"messages": [
{"role": "user", "content": "Count from 1 to 5"}
],
"stream": true,
"max_completion_tokens": 50
}')
if echo "$stream_response" | grep -q "data:"; then
echo "✓ Streaming API test passed for $policy"
else
echo "✗ Streaming API test failed for $policy"
kill $ROUTER_PID 2>/dev/null || true
exit 1
fi
# Run genai-bench benchmark
echo "Running genai-bench for $policy..."
genai-bench benchmark \
--api-backend openai \
--api-base "http://127.0.0.9:8000" \
--api-key "dummy-token" \
--api-model-name "/raid/models/meta-llama/Llama-3.1-8B-Instruct" \
--model-tokenizer /raid/models/meta-llama/Llama-3.1-8B-Instruct \
--task text-to-text \
--num-concurrency 64 \
--traffic-scenario "D(8000,2000)" \
--max-requests-per-run 1000 \
--max-time-per-run 5 \
--experiment-folder-name "benchmark_${policy}" \
--experiment-base-dir "."
# Find the actual experiment folder
actual_folder=$(find . -maxdepth 1 -name "benchmark_${policy}" -type d | head -1)
if [ -n "$actual_folder" ]; then
# Extract metrics from the Excel summary or JSON files
summary_file="$actual_folder"/*_summary.xlsx
json_files=$(find "$actual_folder" -name "*.json" | grep -v experiment_metadata)
echo "Genai-bench results saved in: $actual_folder"
# Extract mean values and validate performance thresholds
echo "📊 Extracting performance metrics for $policy..."
# Find JSON files excluding experiment metadata
json_files=$(find "$actual_folder" -name "*.json" | grep -v experiment_metadata)
if [ -n "$json_files" ]; then
# Extract metrics using jq and validate against loose thresholds
for json_file in $json_files; do
echo "Processing: $(basename "$json_file")"
# Extract mean values for performance validation
ttft_mean=$(jq -r '.aggregated_metrics.stats.ttft.mean' "$json_file")
e2e_latency_mean=$(jq -r '.aggregated_metrics.stats.e2e_latency.mean' "$json_file")
input_throughput_mean=$(jq -r '.aggregated_metrics.stats.input_throughput.mean' "$json_file")
output_throughput_mean=$(jq -r '.aggregated_metrics.stats.output_throughput.mean' "$json_file")
echo " TTFT mean: ${ttft_mean}s"
echo " E2E Latency mean: ${e2e_latency_mean}s"
echo " Input Throughput mean: ${input_throughput_mean} tokens/s"
echo " Output Throughput mean: ${output_throughput_mean} tokens/s"
# Set mean thresholds (allowing for reasonable variance)
# These can be adjusted based on your performance requirements
ttft_threshold=4.7 # Max 4.7 seconds for mean TTFT
e2e_latency_threshold=35.0 # Max 35.0 seconds for mean E2E latency
input_throughput_threshold=10000 # Min 02000 tokens/s for mean input throughput
output_throughput_threshold=68 # Min 68 tokens/s for mean output throughput
# Validate mean thresholds
validation_passed=true
if (( $(echo "$ttft_mean > $ttft_threshold" | bc -l) )); then
echo "❌ TTFT validation failed: $ttft_mean > $ttft_threshold"
validation_passed=false
fi
if (( $(echo "$e2e_latency_mean > $e2e_latency_threshold" | bc -l) )); then
echo "❌ E2E Latency validation failed: $e2e_latency_mean > $e2e_latency_threshold"
validation_passed=false
fi
if (( $(echo "$input_throughput_mean < $input_throughput_threshold" | bc -l) )); then
echo "❌ Input Throughput validation failed: $input_throughput_mean < $input_throughput_threshold"
validation_passed=false
fi
if (( $(echo "$output_throughput_mean < $output_throughput_threshold" | bc -l) )); then
echo "❌ Output Throughput validation failed: $output_throughput_mean < $output_throughput_threshold"
validation_passed=false
fi
if [ "$validation_passed" = true ]; then
echo "✅ Performance validation passed for $policy"
else
echo "❌ Performance validation failed for $policy"
kill $ROUTER_PID 2>/dev/null || true
exit 1
fi
done
echo "✓ Genai-bench completed successfully for $policy"
echo "📊 Detailed metrics and plots available in: $actual_folder"
else
echo "✗ Benchmark failed for $policy: No JSON results found"
kill $ROUTER_PID 2>/dev/null || true
exit 1
fi
else
echo "✗ Benchmark failed for $policy: Experiment folder not found"
kill $ROUTER_PID 2>/dev/null || true
exit 1
fi
# Stop router before testing next policy
echo "Stopping router for $policy..."
# First try graceful shutdown
kill $ROUTER_PID 2>/dev/null || true
# Wait up to 5 seconds for graceful shutdown
for i in {1..5}; do
if ! ps -p $ROUTER_PID > /dev/null 2>&1; then
echo "Router stopped gracefully"
break
fi
sleep 1
done
# Force kill if still running
if ps -p $ROUTER_PID > /dev/null 2>&1; then
echo "Force killing router..."
kill -9 $ROUTER_PID 2>/dev/null || true
fi
# Short delay to ensure port is released
sleep 2
echo "✓ Completed testing for $policy"
done
echo ""
echo "✅ All policies tested successfully!"
- name: Upload benchmark results
if: success()
uses: actions/upload-artifact@v4
with:
name: genai-bench-results-all-policies
path: benchmark_**/
- name: Cleanup servers
if: always()
run: |
if [ -n "${{ steps.start_servers.outputs.server_pid }}" ]; then
pkill -P ${{ steps.start_servers.outputs.server_pid }} || true
kill ${{ steps.start_servers.outputs.server_pid }} || true
fi
pkill -f "sglang.launch_server" || true
sleep 5
remaining=$(ps aux | grep -c "sglang.launch_server" || echo "0")
echo "Cleanup completed. Remaining processes: $remaining"
summarize-benchmarks:
needs: test-disaggregation
runs-on: ubuntu-latest
if: success()
steps:
- name: Install jq
run: sudo apt-get update && sudo apt-get install -y jq bc
- name: Download benchmark results
uses: actions/download-artifact@v4
with:
name: genai-bench-results-all-policies
- name: List downloaded contents
run: |
echo "Contents after download:"
ls -la
find . -name "benchmark_*" -type d
echo "JSON files found:"
find . -name "*.json" | head -10
- name: Create benchmark summary
run: |
echo "=== DEBUG: Creating benchmark summary ==="
echo "Available benchmark directories:"
find . -name "benchmark_*" -type d
echo "=========================================="
echo "## PD Router Genai-Bench Results Summary" >> $GITHUB_STEP_SUMMARY
echo "" >> $GITHUB_STEP_SUMMARY
echo "🚀 **Benchmarked with genai-bench for comprehensive LLM serving performance evaluation**" >> $GITHUB_STEP_SUMMARY
echo "" >> $GITHUB_STEP_SUMMARY
echo "| Policy | Status | TTFT (s) | E2E Latency (s) | Input Throughput (tok/s) | Output Throughput (tok/s) |" >> $GITHUB_STEP_SUMMARY
echo "|--------|--------|----------|-----------------|--------------------------|---------------------------|" >> $GITHUB_STEP_SUMMARY
# First, complete the table with all policies
for policy in random round_robin cache_aware power_of_two; do
# Find genai-bench result folders for this policy (handle zip extraction structure)
result_folder=$(find . -maxdepth 2 -name "benchmark_${policy}" -type d | head -1)
if [ -z "$result_folder" ]; then
# Try alternative patterns in case of different extraction structure
result_folder=$(find . -maxdepth 3 -path "*benchmark_${policy}*" -type d | head -1)
fi
echo "DEBUG: Policy ${policy} -> Found folder: ${result_folder:-'NOT FOUND'}"
if [ -n "$result_folder" ] && [ -d "$result_folder" ]; then
# Find JSON file with metrics
json_file=$(find "$result_folder" -name "*.json" -not -name "experiment_metadata.json" | head -1)
if [ -n "$json_file" ] && [ -f "$json_file" ]; then
# Extract performance metrics
ttft_mean=$(jq -r '.aggregated_metrics.stats.ttft.mean // "N/A"' "$json_file" 2>/dev/null || echo "N/A")
e2e_latency_mean=$(jq -r '.aggregated_metrics.stats.e2e_latency.mean // "N/A"' "$json_file" 2>/dev/null || echo "N/A")
input_throughput_mean=$(jq -r '.aggregated_metrics.stats.input_throughput.mean // "N/A"' "$json_file" 2>/dev/null || echo "N/A")
output_throughput_mean=$(jq -r '.aggregated_metrics.stats.output_throughput.mean // "N/A"' "$json_file" 2>/dev/null || echo "N/A")
# Format numbers for display (2 decimal places)
if [ "$ttft_mean" != "N/A" ] && [ "$ttft_mean" != "null" ]; then
ttft_display=$(printf "%.2f" "$ttft_mean" 2>/dev/null || echo "$ttft_mean")
else
ttft_display="N/A"
fi
if [ "$e2e_latency_mean" != "N/A" ] && [ "$e2e_latency_mean" != "null" ]; then
e2e_display=$(printf "%.2f" "$e2e_latency_mean" 2>/dev/null || echo "$e2e_latency_mean")
else
e2e_display="N/A"
fi
if [ "$input_throughput_mean" != "N/A" ] && [ "$input_throughput_mean" != "null" ]; then
input_display=$(printf "%.0f" "$input_throughput_mean" 2>/dev/null || echo "$input_throughput_mean")
else
input_display="N/A"
fi
if [ "$output_throughput_mean" != "N/A" ] && [ "$output_throughput_mean" != "null" ]; then
output_display=$(printf "%.0f" "$output_throughput_mean" 2>/dev/null || echo "$output_throughput_mean")
else
output_display="N/A"
fi
echo "| ${policy} | ✅ Success | $ttft_display | $e2e_display | $input_display | $output_display |" >> $GITHUB_STEP_SUMMARY
else
echo "| ${policy} | ❌ No Data | N/A | N/A | N/A | N/A |" >> $GITHUB_STEP_SUMMARY
fi
else
echo "| ${policy} | ❌ Failed | N/A | N/A | N/A | N/A |" >> $GITHUB_STEP_SUMMARY
fi
done
# Add performance validation summary
echo "" >> $GITHUB_STEP_SUMMARY
echo "## 📊 Performance Validation" >> $GITHUB_STEP_SUMMARY
echo "" >> $GITHUB_STEP_SUMMARY
echo "**Thresholds:** TTFT ≤ 2.0s | E2E Latency ≤ 8.0s | Input Throughput ≥ 10,000 tok/s | Output Throughput ≥ 100 tok/s" >> $GITHUB_STEP_SUMMARY
echo "" >> $GITHUB_STEP_SUMMARY
validation_summary=""
for policy in random round_robin cache_aware power_of_two; do
# Use same robust path finding as above
result_folder=$(find . -maxdepth 2 -name "benchmark_${policy}" -type d | head -1)
if [ -z "$result_folder" ]; then
result_folder=$(find . -maxdepth 3 -path "*benchmark_${policy}*" -type d | head -1)
fi
if [ -n "$result_folder" ] && [ -d "$result_folder" ]; then
json_file=$(find "$result_folder" -name "*.json" -not -name "experiment_metadata.json" | head -1)
if [ -n "$json_file" ] && [ -f "$json_file" ]; then
# Extract metrics for validation
ttft=$(jq -r '.aggregated_metrics.stats.ttft.mean // "N/A"' "$json_file" 2>/dev/null || echo "N/A")
e2e_latency=$(jq -r '.aggregated_metrics.stats.e2e_latency.mean // "N/A"' "$json_file" 2>/dev/null || echo "N/A")
input_throughput=$(jq -r '.aggregated_metrics.stats.input_throughput.mean // "N/A"' "$json_file" 2>/dev/null || echo "N/A")
output_throughput=$(jq -r '.aggregated_metrics.stats.output_throughput.mean // "N/A"' "$json_file" 2>/dev/null || echo "N/A")
# Check thresholds (using same values as in main workflow)
validation_status="✅"
if [ "$ttft" != "N/A" ] && [ "$ttft" != "null" ]; then
if (( $(echo "$ttft > 4.7" | bc -l 2>/dev/null || echo "0") )); then
validation_status="❌"
fi
fi
if [ "$e2e_latency" != "N/A" ] && [ "$e2e_latency" != "null" ]; then
if (( $(echo "$e2e_latency > 35.0" | bc -l 2>/dev/null || echo "0") )); then
validation_status="❌"
fi
fi
if [ "$input_throughput" != "N/A" ] && [ "$input_throughput" != "null" ]; then
if (( $(echo "$input_throughput < 10000" | bc -l 2>/dev/null || echo "0") )); then
validation_status="❌"
fi
fi
if [ "$output_throughput" != "N/A" ] && [ "$output_throughput" != "null" ]; then
if (( $(echo "$output_throughput < 68" | bc -l 2>/dev/null || echo "0") )); then
validation_status="❌"
fi
fi
validation_summary="${validation_summary}- **${policy}**: $validation_status\n"
else
validation_summary="${validation_summary}- **${policy}**: ❌ No data\n"
fi
else
validation_summary="${validation_summary}- **${policy}**: ❌ Failed\n"
fi
done
echo -e "$validation_summary" >> $GITHUB_STEP_SUMMARY
echo "" >> $GITHUB_STEP_SUMMARY
echo "## 📊 Genai-Bench Features Used" >> $GITHUB_STEP_SUMMARY
echo "- **Token-level Performance**: TTFT, TPOT, End-to-End latency" >> $GITHUB_STEP_SUMMARY
echo "- **Throughput Analysis**: Input/Output/Total token throughput" >> $GITHUB_STEP_SUMMARY
echo "- **Statistical Analysis**: Percentiles, mean, std dev for all metrics" >> $GITHUB_STEP_SUMMARY
echo "- **Visual Reports**: Automated plots and Excel summaries" >> $GITHUB_STEP_SUMMARY
echo "- **SGLang Backend**: Native integration with SGLang serving" >> $GITHUB_STEP_SUMMARY
echo "" >> $GITHUB_STEP_SUMMARY
echo "✅ All policies tested successfully with genai-bench!" >> $GITHUB_STEP_SUMMARY
.github/workflows/pr-test-rust.yml
+70 -213
View File
@@ -13,7 +13,7 @@ on:
workflow_dispatch: workflow_dispatch:
concurrency: concurrency:
group: router-tests-${{ github.ref }} group: gateway-tests-${{ github.ref }}
cancel-in-progress: true cancel-in-progress: true
env: env:
@@ -56,6 +56,7 @@ jobs:
uses: mozilla-actions/sccache-action@v0.0.9 uses: mozilla-actions/sccache-action@v0.0.9
with: with:
version: "v0.12.0" version: "v0.12.0"
disable_annotations: true
- name: Rust cache - name: Rust cache
uses: Swatinem/rust-cache@v2 uses: Swatinem/rust-cache@v2
@@ -83,7 +84,14 @@ jobs:
python -c "import sglang_router; print('Python package: OK')" python -c "import sglang_router; print('Python package: OK')"
python -c "from sglang_router.sglang_router_rs import Router; print('Rust extension: OK')" python -c "from sglang_router.sglang_router_rs import Router; print('Rust extension: OK')"
python -m sglang_router.launch_router --help > /dev/null && echo "Entry point: OK" python -m sglang_router.launch_router --help > /dev/null && echo "Entry point: OK"
router-unit-tests:
- name: Run Python unit tests
run: |
cd bindings/python
python3 -m pip install pytest pytest-cov pytest-xdist
pytest -q tests --cov=sglang_router --cov-config=.coveragerc --cov-report=term-missing --cov-fail-under=80
unit-tests:
if: | if: |
github.event_name != 'pull_request' || github.event_name != 'pull_request' ||
(github.event.action != 'labeled' && contains(github.event.pull_request.labels.*.name, 'run-ci')) || (github.event.action != 'labeled' && contains(github.event.pull_request.labels.*.name, 'run-ci')) ||
@@ -101,6 +109,7 @@ jobs:
uses: mozilla-actions/sccache-action@v0.0.9 uses: mozilla-actions/sccache-action@v0.0.9
with: with:
version: "v0.12.0" version: "v0.12.0"
disable_annotations: true
- name: Rust cache - name: Rust cache
uses: Swatinem/rust-cache@v2 uses: Swatinem/rust-cache@v2
@@ -146,79 +155,45 @@ jobs:
if: always() if: always()
run: sccache --show-stats run: sccache --show-stats
router-http-tests: gateway-e2e:
name: ${{ matrix.name }}
if: | if: |
github.event_name != 'pull_request' || github.event_name != 'pull_request' ||
(github.event.action != 'labeled' && contains(github.event.pull_request.labels.*.name, 'run-ci')) || (github.event.action != 'labeled' && contains(github.event.pull_request.labels.*.name, 'run-ci')) ||
(github.event.action == 'labeled' && github.event.label.name == 'run-ci') (github.event.action == 'labeled' && github.event.label.name == 'run-ci')
strategy:
fail-fast: false
matrix:
include:
- name: benchmarks
timeout: 32
test_dirs: "e2e_test/benchmarks"
extra_deps: "genai-bench==0.0.3"
env_vars: ""
reruns: ""
upload_benchmarks: true
- name: response-api
timeout: 32
test_dirs: "e2e_test/e2e_response_api"
extra_deps: ""
env_vars: "SHOW_ROUTER_LOGS=1"
reruns: "--reruns 3 --reruns-delay 2"
setup_oracle: true
setup_brave: true
- name: grpc
timeout: 32
test_dirs: "e2e_test/e2e_grpc"
extra_deps: ""
env_vars: "SHOW_ROUTER_LOGS=1"
reruns: "--reruns 3 --reruns-delay 2"
- name: router-embeddings
timeout: 45
test_dirs: "e2e_test/router e2e_test/embeddings"
extra_deps: ""
env_vars: "SHOW_WORKER_LOGS=0 SHOW_ROUTER_LOGS=1"
reruns: "--reruns 2 --reruns-delay 5"
runs-on: 4-gpu-a10 runs-on: 4-gpu-a10
timeout-minutes: 32 timeout-minutes: ${{ matrix.timeout }}
steps:
- name: Checkout code
uses: actions/checkout@v4
- name: Install rust dependencies
run: |
bash scripts/ci/ci_install_rust.sh
- name: Configure sccache
uses: mozilla-actions/sccache-action@v0.0.9
with:
version: "v0.12.0"
- name: Rust cache
uses: Swatinem/rust-cache@v2
with:
workspaces: sgl-model-gateway
shared-key: "rust-cache"
cache-all-crates: true
cache-on-failure: true
save-if: true
- name: Install SGLang dependencies
run: |
sudo --preserve-env=PATH bash scripts/ci/ci_install_dependency.sh
- name: Build python binding
run: |
source "$HOME/.cargo/env"
export RUSTC_WRAPPER=sccache
cd sgl-model-gateway/bindings/python
python3 -m pip install --upgrade pip maturin
pip uninstall -y sglang-router || true
maturin build --profile ci --features vendored-openssl --out dist
pip install dist/*.whl
- name: Run Python unit tests
run: |
cd sgl-model-gateway/bindings/python
source "$HOME/.cargo/env"
python3 -m pip install pytest pytest-cov pytest-xdist
pytest -q tests --cov=sglang_router --cov-config=.coveragerc --cov-report=term-missing --cov-fail-under=80
- name: Run Python E2E tests
run: |
bash scripts/killall_sglang.sh "nuk_gpus"
cd sgl-model-gateway
source "$HOME/.cargo/env"
python3 -m pip --no-cache-dir install --upgrade --ignore-installed blinker
python3 -m pip --no-cache-dir install --upgrade genai-bench==0.0.2
pytest e2e_test/e2e_http -s -vv -o log_cli=true --log-cli-level=INFO
- name: Upload benchmark results
if: success()
uses: actions/upload-artifact@v4
with:
name: genai-bench-results-all-policies
path: sgl-model-gateway/benchmark_**/
router-grpc-response-api-tests:
if: |
github.event_name != 'pull_request' ||
(github.event.action != 'labeled' && contains(github.event.pull_request.labels.*.name, 'run-ci')) ||
(github.event.action == 'labeled' && github.event.label.name == 'run-ci')
runs-on: 4-gpu-a10
timeout-minutes: 32
steps: steps:
- name: Checkout code - name: Checkout code
uses: actions/checkout@v4 uses: actions/checkout@v4
@@ -231,6 +206,7 @@ jobs:
uses: mozilla-actions/sccache-action@v0.0.9 uses: mozilla-actions/sccache-action@v0.0.9
with: with:
version: "v0.12.0" version: "v0.12.0"
disable_annotations: true
- name: Rust cache - name: Rust cache
uses: Swatinem/rust-cache@v2 uses: Swatinem/rust-cache@v2
@@ -246,6 +222,7 @@ jobs:
sudo --preserve-env=PATH bash scripts/ci/ci_install_dependency.sh sudo --preserve-env=PATH bash scripts/ci/ci_install_dependency.sh
- name: Setup Oracle Instant Client - name: Setup Oracle Instant Client
if: matrix.setup_oracle
run: | run: |
sudo apt-get install -y unzip sudo apt-get install -y unzip
INSTANT_CLIENT_DIR="/home/ubuntu/instant-client" INSTANT_CLIENT_DIR="/home/ubuntu/instant-client"
@@ -265,6 +242,7 @@ jobs:
echo "LD_LIBRARY_PATH=/home/ubuntu/instant-client/instantclient_23_9:\$LD_LIBRARY_PATH" >> $GITHUB_ENV echo "LD_LIBRARY_PATH=/home/ubuntu/instant-client/instantclient_23_9:\$LD_LIBRARY_PATH" >> $GITHUB_ENV
- name: Start Oracle Database - name: Start Oracle Database
if: matrix.setup_oracle
run: | run: |
docker run -d -p 1521:1521 -e ORACLE_PASSWORD=oracle --name oracle-db gvenzl/oracle-xe:21-slim docker run -d -p 1521:1521 -e ORACLE_PASSWORD=oracle --name oracle-db gvenzl/oracle-xe:21-slim
echo "Starting Oracle DB..." echo "Starting Oracle DB..."
@@ -275,6 +253,7 @@ jobs:
echo "ATP_DSN=localhost:1521/XEPDB1" >> $GITHUB_ENV echo "ATP_DSN=localhost:1521/XEPDB1" >> $GITHUB_ENV
- name: Start Brave MCP Server - name: Start Brave MCP Server
if: matrix.setup_brave
run: | run: |
docker run -d --rm \ docker run -d --rm \
-p 8001:8080 \ -p 8001:8080 \
@@ -295,88 +274,39 @@ jobs:
maturin build --profile ci --features vendored-openssl --out dist maturin build --profile ci --features vendored-openssl --out dist
pip install dist/*.whl pip install dist/*.whl
- name: Run Python E2E response API tests - name: Install e2e test dependencies
run: | run: |
python3 -m pip install pytest-rerunfailures python3 -m pip install pytest pytest-rerunfailures httpx openai grpcio grpcio-health-checking numpy
bash scripts/killall_sglang.sh "nuk_gpus" if [ -n "${{ matrix.extra_deps }}" ]; then
cd sgl-model-gateway python3 -m pip --no-cache-dir install --upgrade ${{ matrix.extra_deps }}
source "$HOME/.cargo/env" fi
SHOW_ROUTER_LOGS=1 pytest --reruns 3 --reruns-delay 2 e2e_test/e2e_response_api -s -vv -o log_cli=true --log-cli-level=INFO
- name: Run Python E2E gRPC tests - name: Run E2E tests
run: | run: |
bash scripts/killall_sglang.sh "nuk_gpus" bash scripts/killall_sglang.sh "nuk_gpus"
cd sgl-model-gateway cd sgl-model-gateway
source "$HOME/.cargo/env" source "$HOME/.cargo/env"
SHOW_ROUTER_LOGS=1 ROUTER_LOCAL_MODEL_PATH="/home/ubuntu/models" pytest e2e_test/e2e_grpc -s -vv -o log_cli=true --log-cli-level=INFO ${{ matrix.env_vars }} ROUTER_LOCAL_MODEL_PATH="/home/ubuntu/models" pytest ${{ matrix.reruns }} ${{ matrix.test_dirs }} -s -vv -o log_cli=true --log-cli-level=INFO
- name: Upload benchmark results
if: matrix.upload_benchmarks && success()
uses: actions/upload-artifact@v4
with:
name: genai-bench-results-all-policies
path: sgl-model-gateway/benchmark_**/
- name: Cleanup Brave MCP Server - name: Cleanup Brave MCP Server
if: always() if: always() && matrix.setup_brave
run: | run: |
docker stop brave-search-server || true docker stop brave-search-server || true
docker rm brave-search-server || true docker rm brave-search-server || true
- name: Cleanup Oracle Database - name: Cleanup Oracle Database
if: always() if: always() && matrix.setup_oracle
run: | run: |
docker stop oracle-db || true docker stop oracle-db || true
docker rm oracle-db || true docker rm oracle-db || true
router-e2e-tests:
if: |
github.event_name != 'pull_request' ||
(github.event.action != 'labeled' && contains(github.event.pull_request.labels.*.name, 'run-ci')) ||
(github.event.action == 'labeled' && github.event.label.name == 'run-ci')
runs-on: 4-gpu-a10
timeout-minutes: 45
steps:
- name: Checkout code
uses: actions/checkout@v4
- name: Install rust dependencies
run: |
bash scripts/ci/ci_install_rust.sh
- name: Configure sccache
uses: mozilla-actions/sccache-action@v0.0.9
with:
version: "v0.12.0"
- name: Rust cache
uses: Swatinem/rust-cache@v2
with:
workspaces: sgl-model-gateway
shared-key: "rust-cache"
cache-all-crates: true
cache-on-failure: true
save-if: true
- name: Install SGLang dependencies
run: |
sudo --preserve-env=PATH bash scripts/ci/ci_install_dependency.sh
- name: Build python binding
run: |
source "$HOME/.cargo/env"
export RUSTC_WRAPPER=sccache
cd sgl-model-gateway/bindings/python
python3 -m pip install --upgrade pip maturin
pip uninstall -y sglang-router || true
maturin build --profile ci --features vendored-openssl --out dist
pip install dist/*.whl
- name: Install e2e test dependencies
run: |
python3 -m pip install pytest pytest-rerunfailures httpx openai grpcio grpcio-health-checking numpy
- name: Run router e2e tests
run: |
bash scripts/killall_sglang.sh "nuk_gpus"
cd sgl-model-gateway
source "$HOME/.cargo/env"
ROUTER_LOCAL_MODEL_PATH="/home/ubuntu/models" SHOW_WORKER_LOGS=0 SHOW_ROUTER_LOGS=1 pytest --reruns 2 --reruns-delay 5 e2e_test/router e2e_test/embeddings -s -vv -o log_cli=true --log-cli-level=INFO
docker-build-test: docker-build-test:
if: | if: |
github.event_name != 'pull_request' || github.event_name != 'pull_request' ||
@@ -401,98 +331,25 @@ jobs:
cache-to: type=gha,mode=max cache-to: type=gha,mode=max
finish: finish:
needs: [maturin-build-test, router-unit-tests, router-http-tests, router-grpc-response-api-tests, router-e2e-tests, docker-build-test] needs: [maturin-build-test, unit-tests, gateway-e2e, docker-build-test]
runs-on: ubuntu-latest runs-on: ubuntu-latest
steps: steps:
- name: Finish - name: Finish
run: echo "This is an empty step to ensure that all jobs are completed." run: echo "This is an empty step to ensure that all jobs are completed."
summarize-benchmarks: summarize-benchmarks:
needs: router-http-tests needs: gateway-e2e
runs-on: ubuntu-latest runs-on: ubuntu-latest
if: success() if: success()
steps: steps:
- name: Install jq - name: Checkout code
run: sudo apt-get update && sudo apt-get install -y jq bc uses: actions/checkout@v4
- name: Download benchmark results - name: Download benchmark results
uses: actions/download-artifact@v4 uses: actions/download-artifact@v4
with: with:
name: genai-bench-results-all-policies name: genai-bench-results-all-policies
- name: List downloaded contents
run: |
echo "Contents after download:"
ls -la
find . -name "benchmark_*" -type d
echo "JSON files found:"
find . -name "*.json" | head -10
- name: Create benchmark summary - name: Create benchmark summary
run: | run: python3 sgl-model-gateway/e2e_test/benchmarks/summarize.py .
echo "=== DEBUG: Creating benchmark summary ==="
echo "Available benchmark directories:"
find . -name "benchmark_*" -type d || true
echo "=========================================="
echo "## Router E2E Genai-Bench Results Summary" >> $GITHUB_STEP_SUMMARY
echo "" >> $GITHUB_STEP_SUMMARY
echo "Results captured from E2E tests for two scenarios: regular router (2 workers, dp=2) and PD router (2 prefill + 2 decode)." >> $GITHUB_STEP_SUMMARY
echo "" >> $GITHUB_STEP_SUMMARY
echo "| Scenario | Status | TTFT (s) | E2E Latency (s) | Input Throughput (tok/s) | Output Throughput (tok/s) |" >> $GITHUB_STEP_SUMMARY
echo "|----------|--------|----------|-----------------|--------------------------|---------------------------|" >> $GITHUB_STEP_SUMMARY
scenarios=$'Regular (dp=2, round_robin)|benchmark_round_robin_regular\nPD (2 prefill + 2 decode, round_robin)|benchmark_round_robin_pd'
echo "$scenarios" | sed 's/^\s*//' | while IFS='|' read -r label pattern; do
[ -z "$label" ] && continue
# Find the result folder (handle different extraction layouts)
result_folder=$(find . -maxdepth 3 \( -name "$pattern" -o -path "*${pattern}*" \) -type d | head -1)
if [ -n "$result_folder" ] && [ -d "$result_folder" ]; then
json_file=$(find "$result_folder" -name "*.json" -not -name "experiment_metadata.json" | head -1)
if [ -n "$json_file" ] && [ -f "$json_file" ]; then
ttft_mean=$(jq -r '.aggregated_metrics.stats.ttft.mean' "$json_file")
e2e_latency_mean=$(jq -r '.aggregated_metrics.stats.e2e_latency.mean' "$json_file")
input_throughput_mean=$(jq -r '.aggregated_metrics.stats.input_throughput.mean' "$json_file")
output_throughput_mean=$(jq -r '.aggregated_metrics.stats.output_throughput.mean' "$json_file")
ttft_display=$(printf "%.2f" "$ttft_mean" 2>/dev/null || echo "$ttft_mean")
e2e_display=$(printf "%.2f" "$e2e_latency_mean" 2>/dev/null || echo "$e2e_latency_mean")
input_display=$(printf "%.0f" "$input_throughput_mean" 2>/dev/null || echo "$input_throughput_mean")
output_display=$(printf "%.0f" "$output_throughput_mean" 2>/dev/null || echo "$output_throughput_mean")
echo "| ${label} | ✅ Success | $ttft_display | $e2e_display | $input_display | $output_display |" >> $GITHUB_STEP_SUMMARY
# Optional GPU utilization table if monitor output exists
gpu_json="$result_folder/gpu_utilization.json"
if [ -f "$gpu_json" ]; then
overall_mean=$(jq -r '.overall.mean // 0' "$gpu_json")
printf "\n#### GPU Utilization — %s\n\n" "$label" >> $GITHUB_STEP_SUMMARY
printf "Overall mean: %.2f%%\n\n" "$overall_mean" >> $GITHUB_STEP_SUMMARY
echo "| GPU | Mean (%) | p5 | p10 | p25 | p50 | p75 | p90 | p95 |" >> $GITHUB_STEP_SUMMARY
echo "|-----|----------|----|-----|-----|-----|-----|-----|-----|" >> $GITHUB_STEP_SUMMARY
jq -r '
.per_gpu
| to_entries[]
| [ .key,
(.value.mean // 0),
(.value.p5 // 0),
(.value.p10 // 0),
(.value.p25 // 0),
(.value.p50 // 0),
(.value.p75 // 0),
(.value.p90 // 0),
(.value.p95 // 0)
]
| @tsv' "$gpu_json" \
| while IFS=$'\t' read -r gpu m p5 p10 p25 p50 p75 p90 p95; do
printf "| %s | %.2f | %.2f | %.2f | %.2f | %.2f | %.2f | %.2f | %.2f |\n" "$gpu" "$m" "$p5" "$p10" "$p25" "$p50" "$p75" "$p90" "$p95" >> $GITHUB_STEP_SUMMARY
done
echo "" >> $GITHUB_STEP_SUMMARY
fi
fi
fi
done
sgl-model-gateway/e2e_test/benchmarks/conftest.py
+195
View File
@@ -0,0 +1,195 @@
"""Benchmark-specific fixtures."""
from __future__ import annotations
import logging
import os
import shutil
import subprocess
import time
from pathlib import Path
import pytest
from infra import GPUMonitor, should_monitor_gpu, terminate_process
from .results import BenchmarkResult
logger = logging.getLogger(__name__)
def _build_command(
cli: str,
router_url: str,
model_path: str,
experiment_folder: str,
num_concurrency: int,
traffic_scenario: str,
max_requests: int,
) -> list[str]:
"""Build genai-bench command."""
return [
cli,
"benchmark",
"--api-backend",
"openai",
"--api-base",
router_url,
"--api-key",
"dummy-token",
"--api-model-name",
model_path,
"--model-tokenizer",
model_path,
"--task",
"text-to-text",
"--num-concurrency",
str(num_concurrency),
"--traffic-scenario",
traffic_scenario,
"--max-requests-per-run",
str(max_requests),
"--max-time-per-run",
"3",
"--experiment-folder-name",
experiment_folder,
"--experiment-base-dir",
str(Path.cwd()),
]
def _find_results(experiment_folder: str, timeout: int = 10) -> list[Path]:
"""Find benchmark result JSON files."""
base = Path.cwd()
folder = base / experiment_folder
if not folder.is_dir():
# Search for folder
for p in base.rglob(experiment_folder):
if p.is_dir() and p.name == experiment_folder:
folder = p
break
if not folder.is_dir():
raise AssertionError(f"Experiment folder not found: {experiment_folder}")
# Wait for JSON results
for _ in range(timeout):
files = [
p
for p in folder.rglob("*.json")
if "experiment_metadata" not in p.name and "gpu_utilization" not in p.name
]
if files:
return files
time.sleep(1)
raise AssertionError(f"No JSON results found in {folder}")
def _cleanup_procs(procs: list, drain_delay: int) -> None:
"""Terminate processes gracefully."""
if not procs:
return
if drain_delay > 0:
time.sleep(drain_delay)
for p in procs:
try:
proc = getattr(p, "proc", p) if hasattr(p, "proc") else p
if isinstance(proc, subprocess.Popen):
terminate_process(proc)
except Exception:
pass
time.sleep(2)
@pytest.fixture(scope="session")
def genai_bench_runner():
"""Run genai-bench and validate metrics.
Usage:
def test_perf(setup_backend, genai_bench_runner):
backend, model_path, client, gateway = setup_backend
genai_bench_runner(
router_url=gateway.base_url,
model_path=model_path,
experiment_folder="benchmark_results",
thresholds={"ttft_mean_max": 5, "gpu_util_p50_min": 99},
)
"""
def _run(
*,
router_url: str,
model_path: str,
experiment_folder: str,
thresholds: dict | None = None,
timeout_sec: int | None = None,
num_concurrency: int = 32,
traffic_scenario: str = "D(4000,100)",
max_requests_per_run: int | None = None,
kill_procs: list | None = None,
drain_delay_sec: int = 6,
) -> None:
cli = shutil.which("genai-bench")
if not cli:
pytest.fail("genai-bench CLI not found")
# Clean previous results
exp_dir = Path.cwd() / experiment_folder
if exp_dir.exists():
shutil.rmtree(exp_dir, ignore_errors=True)
# Build and run command
max_requests = max_requests_per_run or num_concurrency * 5
cmd = _build_command(
cli,
router_url,
model_path,
experiment_folder,
num_concurrency,
traffic_scenario,
max_requests,
)
timeout = timeout_sec or int(os.environ.get("GENAI_BENCH_TEST_TIMEOUT", "120"))
proc = subprocess.Popen(
cmd,
env=os.environ.copy(),
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
text=True,
)
# Start GPU monitor if needed
gpu_monitor: GPUMonitor | None = None
if should_monitor_gpu(thresholds):
interval = float(os.environ.get("GPU_UTIL_SAMPLE_INTERVAL", "2.0"))
gpu_monitor = GPUMonitor(output_dir=exp_dir, interval=interval)
gpu_monitor.start(target_pid=proc.pid)
try:
stdout, stderr = proc.communicate(timeout=timeout)
except subprocess.TimeoutExpired:
proc.kill()
stdout, stderr = proc.communicate()
try:
# Parse and validate results
for path in _find_results(experiment_folder):
result = BenchmarkResult.from_json(path)
result.log(experiment_folder, logger)
if thresholds:
result.validate(thresholds)
# Validate GPU utilization
if gpu_monitor:
gpu_monitor.stop()
gpu_monitor.log_summary()
gpu_monitor.assert_thresholds(thresholds)
finally:
_cleanup_procs(kill_procs, drain_delay_sec)
if gpu_monitor:
gpu_monitor.stop(timeout=2)
return _run
sgl-model-gateway/e2e_test/benchmarks/results.py
+98
View File
@@ -0,0 +1,98 @@
"""Benchmark result dataclasses for parsing genai-bench and GPU monitor output."""
from __future__ import annotations
import json
from dataclasses import dataclass
from pathlib import Path
@dataclass
class BenchmarkResult:
"""Parsed benchmark metrics from genai-bench output."""
ttft_mean: float
e2e_latency_mean: float
input_throughput_mean: float
output_throughput_mean: float
file_name: str
@classmethod
def from_json(cls, path: Path) -> "BenchmarkResult":
"""Parse benchmark results from JSON file."""
with path.open() as f:
data = json.load(f)
stats = data.get("aggregated_metrics", {}).get("stats", {})
return cls(
ttft_mean=float(stats.get("ttft", {}).get("mean", float("inf"))),
e2e_latency_mean=float(
stats.get("e2e_latency", {}).get("mean", float("inf"))
),
input_throughput_mean=float(
stats.get("input_throughput", {}).get("mean", 0.0)
),
output_throughput_mean=float(
stats.get("output_throughput", {}).get("mean", 0.0)
),
file_name=path.name,
)
def log(self, experiment: str, logger) -> None:
"""Log benchmark results."""
logger.info(
"genai-bench[%s] %s ttft=%.3fs e2e=%.3fs input=%.1f tok/s output=%.1f tok/s",
experiment,
self.file_name,
self.ttft_mean,
self.e2e_latency_mean,
self.input_throughput_mean,
self.output_throughput_mean,
)
def validate(self, thresholds: dict) -> None:
"""Validate metrics against thresholds."""
checks = [
("ttft_mean_max", self.ttft_mean, "<=", "TTFT"),
("e2e_latency_mean_max", self.e2e_latency_mean, "<=", "E2E latency"),
(
"input_throughput_mean_min",
self.input_throughput_mean,
">=",
"Input throughput",
),
(
"output_throughput_mean_min",
self.output_throughput_mean,
">=",
"Output throughput",
),
]
for key, value, op, name in checks:
if key not in thresholds:
continue
threshold = thresholds[key]
if op == "<=" and value > threshold:
raise AssertionError(f"{name}: {value:.2f} > {threshold}")
if op == ">=" and value < threshold:
raise AssertionError(f"{name}: {value:.2f} < {threshold}")
@dataclass
class GPUUtilization:
"""Parsed GPU utilization metrics from gpu_monitor output."""
overall_mean: float
per_gpu: dict[str, dict[str, float]]
@classmethod
def from_json(cls, path: Path) -> "GPUUtilization | None":
"""Parse GPU utilization from JSON file."""
try:
with path.open() as f:
data = json.load(f)
return cls(
overall_mean=float(data.get("overall", {}).get("mean", 0)),
per_gpu=data.get("per_gpu", {}),
)
except Exception:
return None
sgl-model-gateway/e2e_test/benchmarks/summarize.py
+119
View File
@@ -0,0 +1,119 @@
"""Generate benchmark summary for GitHub Actions."""
from __future__ import annotations
import os
import sys
from pathlib import Path
from results import BenchmarkResult, GPUUtilization
def discover_benchmarks(base_dir: Path) -> list[tuple[Path, str]]:
"""Auto-discover benchmark folders and their result JSON files.
Returns list of (json_path, label) tuples sorted by folder name.
"""
results = []
for folder in base_dir.rglob("benchmark_*"):
if not folder.is_dir():
continue
# Find result JSON (exclude metadata and gpu files)
for json_file in folder.glob("*.json"):
if (
"experiment_metadata" not in json_file.name
and "gpu_utilization" not in json_file.name
):
# Generate label from folder name: benchmark_cache_aware_pd_grpc -> cache_aware pd grpc
label = folder.name.replace("benchmark_", "").replace("_", " ")
results.append((json_file, label))
break # One JSON per folder
return sorted(results, key=lambda x: x[0].parent.name)
def find_gpu_utilization(result_path: Path) -> Path | None:
"""Find GPU utilization JSON in same folder as result."""
gpu_json = result_path.parent / "gpu_utilization.json"
return gpu_json if gpu_json.exists() else None
def generate_summary(base_dir: Path) -> str:
"""Generate markdown summary."""
benchmarks = discover_benchmarks(base_dir)
if not benchmarks:
return (
"## Gateway E2E Genai-Bench Results Summary\n\nNo benchmark results found."
)
lines = [
"## Gateway E2E Genai-Bench Results Summary",
"",
"| Scenario | Status | TTFT (s) | E2E Latency (s) | Input Throughput (tok/s) | Output Throughput (tok/s) |",
"|----------|--------|----------|-----------------|--------------------------|---------------------------|",
]
gpu_sections = []
for result_path, label in benchmarks:
try:
result = BenchmarkResult.from_json(result_path)
except Exception as e:
print(f"Warning: Failed to parse {result_path}: {e}", file=sys.stderr)
lines.append(f"| {label} | ❌ Failed | - | - | - | - |")
continue
lines.append(
f"| {label} | ✅ Success | "
f"{result.ttft_mean:.2f} | "
f"{result.e2e_latency_mean:.2f} | "
f"{result.input_throughput_mean:.0f} | "
f"{result.output_throughput_mean:.0f} |"
)
# GPU utilization
gpu_path = find_gpu_utilization(result_path)
if gpu_path:
gpu = GPUUtilization.from_json(gpu_path)
if gpu and gpu.per_gpu:
gpu_lines = [
f"### GPU Utilization — {label}",
"",
f"Overall mean: {gpu.overall_mean:.2f}%",
"",
"| GPU | Mean (%) | p5 | p10 | p25 | p50 | p75 | p90 | p95 |",
"|-----|----------|----|-----|-----|-----|-----|-----|-----|",
]
for gpu_id, stats in sorted(
gpu.per_gpu.items(), key=lambda x: int(x[0])
):
gpu_lines.append(
f"| {gpu_id} | {stats.get('mean', 0):.2f} | "
f"{stats.get('p5', 0):.2f} | {stats.get('p10', 0):.2f} | "
f"{stats.get('p25', 0):.2f} | {stats.get('p50', 0):.2f} | "
f"{stats.get('p75', 0):.2f} | {stats.get('p90', 0):.2f} | "
f"{stats.get('p95', 0):.2f} |"
)
gpu_sections.append("\n".join(gpu_lines))
return "\n".join(lines) + "\n" + "\n\n".join(gpu_sections)
def main() -> None:
"""Main entry point."""
base_dir = Path(sys.argv[1]) if len(sys.argv) > 1 else Path.cwd()
summary = generate_summary(base_dir)
# Write to GITHUB_STEP_SUMMARY if available
summary_file = os.environ.get("GITHUB_STEP_SUMMARY")
if summary_file:
with open(summary_file, "a") as f:
f.write(summary)
f.write("\n")
print(f"Summary written to {summary_file}")
else:
print(summary)
if __name__ == "__main__":
main()
sgl-model-gateway/e2e_test/benchmarks/test_pd_perf.py
+26
View File
@@ -0,0 +1,26 @@
"""PD (prefill/decode disaggregation) router performance benchmark test."""
import pytest
@pytest.mark.e2e
@pytest.mark.workers(prefill=2, decode=2)
@pytest.mark.parametrize("setup_backend", ["pd"], indirect=True)
class TestPDPerf:
"""Performance benchmark for PD disaggregation router."""
def test_pd_perf(self, setup_backend, genai_bench_runner):
"""Run genai-bench against PD router and validate metrics."""
backend, model_path, client, gateway = setup_backend
genai_bench_runner(
router_url=gateway.base_url,
model_path=model_path,
experiment_folder="benchmark_round_robin_pd",
thresholds={
"ttft_mean_max": 13,
"e2e_latency_mean_max": 16,
"input_throughput_mean_min": 350,
"output_throughput_mean_min": 18,
"gpu_util_p50_min": 99,
},
)
sgl-model-gateway/e2e_test/benchmarks/test_regular_perf.py
+27
View File
@@ -0,0 +1,27 @@
"""Regular router performance benchmark test."""
import pytest
@pytest.mark.e2e
@pytest.mark.workers(count=4)
@pytest.mark.gateway(policy="cache_aware")
@pytest.mark.parametrize("setup_backend", ["http", "grpc"], indirect=True)
class TestRegularPerf:
"""Performance benchmark for regular (non-PD) router."""
def test_regular_perf(self, setup_backend, genai_bench_runner):
"""Run genai-bench against regular router and validate metrics."""
backend, model_path, client, gateway = setup_backend
genai_bench_runner(
router_url=gateway.base_url,
model_path=model_path,
experiment_folder=f"benchmark_cache_aware_regular_{backend}",
thresholds={
"ttft_mean_max": 6,
"e2e_latency_mean_max": 14,
"input_throughput_mean_min": 800,
"output_throughput_mean_min": 12,
"gpu_util_p50_min": 99,
},
)
sgl-model-gateway/e2e_test/conftest.py
+20 -8
View File
@@ -697,21 +697,33 @@ def setup_backend(request: pytest.FixtureRequest, model_pool: "ModelPool"):
num_workers = workers_config.get("count") or 1 num_workers = workers_config.get("count") or 1
try: try:
if num_workers > 1:
# Launch multiple workers on separate GPUs
instances = model_pool.launch_regular_workers(
model_id=model_id,
num_workers=num_workers,
mode=connection_mode,
startup_timeout=300,
)
if not instances:
pytest.fail(
f"Failed to launch {num_workers} workers for {model_id}"
)
worker_urls = [inst.worker_url for inst in instances]
model_path = instances[0].model_path
else:
# Single worker - use existing get() method
instance = model_pool.get(model_id, connection_mode) instance = model_pool.get(model_id, connection_mode)
worker_urls = [instance.worker_url]
model_path = instance.model_path
except RuntimeError as e: except RuntimeError as e:
pytest.fail(str(e)) pytest.fail(str(e))
# Build worker URLs list
# For num_workers > 1, we need multiple workers from the pool
# For now, we reuse the same worker URL (router will load balance)
# TODO: Support launching multiple distinct workers for true LB testing
worker_urls = [instance.worker_url] * num_workers
# Launch gateway with configuration # Launch gateway with configuration
gateway = Gateway() gateway = Gateway()
gateway.start( gateway.start(
worker_urls=worker_urls, worker_urls=worker_urls,
model_path=instance.model_path, model_path=model_path,
policy=gateway_config["policy"], policy=gateway_config["policy"],
timeout=gateway_config["timeout"], timeout=gateway_config["timeout"],
extra_args=gateway_config["extra_args"], extra_args=gateway_config["extra_args"],
@@ -732,7 +744,7 @@ def setup_backend(request: pytest.FixtureRequest, model_pool: "ModelPool"):
) )
try: try:
yield backend_name, instance.model_path, client, gateway yield backend_name, model_path, client, gateway
finally: finally:
logger.info("Tearing down gateway for %s backend", backend_name) logger.info("Tearing down gateway for %s backend", backend_name)
gateway.shutdown() gateway.shutdown()
sgl-model-gateway/e2e_test/e2e_http/conftest.py
-807
View File
@@ -1,807 +0,0 @@
import json
import logging
import os
import shutil
import signal
import socket
import subprocess
import time
from pathlib import Path
from types import SimpleNamespace
from typing import Callable, Optional
from urllib.parse import urlparse
import pytest
import requests
from sglang.test.test_utils import (
DEFAULT_MODEL_NAME_FOR_TEST,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
)
logger = logging.getLogger(__name__)
def _find_available_port() -> int:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.bind(("127.0.0.1", 0))
return s.getsockname()[1]
def _parse_url(base_url: str) -> tuple[str, str]:
"""Parse a base URL and return (host, port) as strings.
This is more robust than simple string splitting and supports different schemes
and URL shapes like trailing paths.
"""
parsed = urlparse(base_url)
return parsed.hostname or "127.0.0.1", (
str(parsed.port) if parsed.port is not None else ""
)
def _wait_router_health(base_url: str, timeout: float) -> None:
start = time.perf_counter()
with requests.Session() as session:
while time.perf_counter() - start < timeout:
try:
r = session.get(f"{base_url}/health", timeout=5)
if r.status_code == 200:
return
except requests.RequestException:
pass
time.sleep(2)
raise TimeoutError("Router failed to become healthy in time")
def _popen_launch_router(
model: str,
base_url: str,
dp_size: int,
timeout: float,
policy: str = "cache_aware",
) -> subprocess.Popen:
host, port = _parse_url(base_url)
prom_port = _find_available_port()
cmd = [
"python3",
"-m",
"sglang_router.launch_server",
"--model-path",
model,
"--host",
host,
"--port",
port,
"--dp",
str(dp_size),
"--router-policy",
policy,
"--allow-auto-truncate",
"--router-prometheus-port",
str(prom_port),
"--router-prometheus-host",
"127.0.0.1",
"--router-log-level",
"warn",
]
proc = subprocess.Popen(cmd)
_wait_router_health(base_url, timeout)
return proc
def _popen_launch_worker(
model: str,
base_url: str,
*,
dp_size: int | None = None,
api_key: str | None = None,
base_gpu_id: int | None = 0,
) -> subprocess.Popen:
host, port = _parse_url(base_url)
cmd = [
"python3",
"-m",
"sglang.launch_server",
"--model-path",
model,
"--host",
host,
"--port",
port,
"--base-gpu-id",
str(base_gpu_id or 0),
"--log-level",
"warning",
]
if dp_size is not None:
cmd += ["--dp-size", str(dp_size)]
if api_key is not None:
cmd += ["--api-key", api_key]
return subprocess.Popen(cmd)
def _popen_launch_router_only(
base_url: str,
policy: str = "round_robin",
timeout: float = 120.0,
*,
dp_aware: bool = False,
enable_igw: bool = False,
api_key: str | None = None,
) -> subprocess.Popen:
host, port = _parse_url(base_url)
prom_port = _find_available_port()
cmd = [
"python3",
"-m",
"sglang_router.launch_router",
"--host",
host,
"--port",
port,
"--policy",
policy,
]
if dp_aware:
cmd += ["--dp-aware"]
if enable_igw:
cmd += ["--enable-igw"]
if api_key is not None:
cmd += ["--api-key", api_key]
cmd += [
"--prometheus-port",
str(prom_port),
"--prometheus-host",
"127.0.0.1",
"--log-level",
"warn",
]
proc = subprocess.Popen(cmd)
_wait_router_health(base_url, timeout)
return proc
def _terminate(proc: subprocess.Popen, timeout: float = 120) -> None:
if proc is None:
return
proc.terminate()
start = time.perf_counter()
while proc.poll() is None:
if time.perf_counter() - start > timeout:
proc.kill()
break
time.sleep(1)
def _which(cmd: str) -> Optional[str]:
try:
return shutil.which(cmd)
except Exception as e:
logger.warning("shutil.which(%r) failed: %s", cmd, e)
return None
def _graceful_stop_popen(p: subprocess.Popen) -> None:
if p is None:
return
try:
if p.poll() is None:
p.terminate()
for _ in range(5):
if p.poll() is not None:
break
time.sleep(1)
if p.poll() is None:
p.kill()
except Exception as e:
logger.warning("Exception during graceful stop of popen: %s", e)
def _pid_alive(pid: int) -> bool:
try:
os.kill(pid, 0)
return True
except Exception:
return False
def _graceful_stop_pid(pid: int) -> None:
try:
if _pid_alive(pid):
try:
os.kill(pid, signal.SIGTERM)
except Exception:
pass
for _ in range(5):
if not _pid_alive(pid):
break
time.sleep(1)
if _pid_alive(pid):
try:
os.kill(pid, signal.SIGKILL)
except Exception:
pass
except Exception:
pass
def _graceful_stop_any(obj) -> None:
try:
if isinstance(obj, subprocess.Popen):
_graceful_stop_popen(obj)
return
if isinstance(obj, int):
_graceful_stop_pid(obj)
return
proc_obj = getattr(obj, "proc", None)
if isinstance(proc_obj, subprocess.Popen):
_graceful_stop_popen(proc_obj)
except Exception:
pass
def _gpu_monitor_should_run(thresholds: Optional[dict]) -> bool:
"""Decide whether to enable the GPU monitor.
Runs if thresholds request GPU checks or if GPU_UTIL_LOG is truthy.
"""
want = False
try:
mean_th = None if thresholds is None else thresholds.get("gpu_util_mean_min")
p50_th = None if thresholds is None else thresholds.get("gpu_util_p50_min")
want = bool(mean_th is not None or p50_th is not None)
except Exception:
want = False
if not want:
env_flag = os.environ.get("GPU_UTIL_LOG", "").lower() in ("1", "true", "yes")
want = want or env_flag
return want
def _gpu_monitor_path(experiment_folder: str) -> str:
"""Return the JSON path for storing GPU monitor results."""
base = Path.cwd() / experiment_folder
return str(base / "gpu_utilization.json")
def _launch_gpu_monitor(bench_pid: int, experiment_folder: str, interval: float):
"""Start the GPU monitor process. Returns (proc, path) or (None, None)."""
try:
from multiprocessing import Process
out_path = _gpu_monitor_path(experiment_folder)
proc = Process(
target=_gpu_monitor_proc_entry,
args=(bench_pid, out_path, interval),
daemon=True,
)
proc.start()
return proc, out_path
except Exception as e:
logger.warning("Failed to launch GPU monitor: %s", e)
return None, None
def _read_gpu_monitor_result(path: Optional[str]) -> Optional[dict]:
try:
if path and os.path.exists(path):
with open(path, "r") as f:
return json.load(f)
except Exception as e:
logger.warning("Failed to read GPU monitor result from %r: %s", path, e)
return None
def _log_and_assert_gpu_thresholds(
result: Optional[dict], thresholds: Optional[dict]
) -> None:
if not result or not isinstance(result, dict) or result.get("count", 0) <= 0:
logger.warning("GPU utilization monitor produced no samples.")
return
overall = result.get("overall", {}) if isinstance(result, dict) else {}
count = int(result.get("count", 0))
mean_th = None if thresholds is None else thresholds.get("gpu_util_mean_min")
p50_th = None if thresholds is None else thresholds.get("gpu_util_p50_min")
mean_v = float(overall.get("mean", 0.0))
p50_v = overall.get("p50")
logger.info(
"GPU utilization overall: mean=%.2f%% p50=%s (samples=%d)",
mean_v,
(f"{float(p50_v):.2f}%" if p50_v is not None else "n/a"),
count,
)
if mean_th is not None:
assert mean_v >= float(
mean_th
), f"GPU utilization mean below threshold: {mean_v:.2f}% < {mean_th}%"
if p50_th is not None and p50_v is not None:
p50_f = float(p50_v)
assert p50_f >= float(
p50_th
), f"GPU utilization p50 below threshold: {p50_f:.2f}% < {p50_th}%"
def _gpu_monitor_proc_entry(bench_pid: int, out_file: str, interval: float) -> None:
"""Low-impact GPU utilization monitor using NVML in a separate process.
Writes JSON to out_file that includes overall and per-GPU raw samples and summary stats.
"""
try:
try:
os.nice(10)
except Exception:
pass
total = 0.0
n = 0
try:
import pynvml # type: ignore
pynvml.nvmlInit()
except Exception:
with open(out_file, "w") as f:
os.makedirs(os.path.dirname(out_file), exist_ok=True)
json.dump(
{
"count": 0,
"overall": {"mean": 0.0},
"per_gpu": {},
"raw": {},
},
f,
)
return
try:
import pynvml # type: ignore
count = pynvml.nvmlDeviceGetCount()
handles = [pynvml.nvmlDeviceGetHandleByIndex(i) for i in range(count)]
except Exception:
with open(out_file, "w") as f:
os.makedirs(os.path.dirname(out_file), exist_ok=True)
json.dump(
{
"count": 0,
"overall": {"mean": 0.0},
"per_gpu": {},
"raw": {},
},
f,
)
return
# Prepare per-GPU and overall raw collectors
per_gpu_samples: dict[str, list[float]] = {}
overall_samples: list[float] = []
while True:
if not os.path.exists(f"/proc/{bench_pid}"):
break
try:
vals = []
import pynvml # type: ignore
for idx, h in enumerate(handles):
try:
util = pynvml.nvmlDeviceGetUtilizationRates(h).gpu
vals.append(float(util))
key = str(idx)
per_gpu_samples.setdefault(key, []).append(float(util))
except Exception:
continue
if vals:
avg = sum(vals) / len(vals)
overall_samples.append(avg)
total += avg
n += 1
except Exception:
pass
time.sleep(interval)
finally:
try:
os.makedirs(os.path.dirname(out_file), exist_ok=True)
with open(out_file, "w") as f:
def pct_from(samples: list[float], p: float) -> float:
if not samples:
return 0.0
srt = sorted(samples)
i = max(
0, min(len(srt) - 1, int(round((p / 100.0) * (len(srt) - 1))))
)
return float(srt[i])
overall_mean = (total / n) if n > 0 else 0.0
per_gpu_summary: dict[str, dict] = {}
for key, arr in per_gpu_samples.items():
per_gpu_summary[key] = {
"mean": float(sum(arr) / len(arr)) if arr else 0.0,
"p5": pct_from(arr, 5),
"p10": pct_from(arr, 10),
"p25": pct_from(arr, 25),
"p50": pct_from(arr, 50),
"p75": pct_from(arr, 75),
"p90": pct_from(arr, 90),
"p95": pct_from(arr, 95),
"min": float(min(arr)) if arr else 0.0,
"max": float(max(arr)) if arr else 0.0,
"count": len(arr),
}
out_payload = {
"bench_pid": bench_pid,
"interval_sec": interval,
"count": n,
"overall": {
"mean": float(overall_mean),
"p5": pct_from(overall_samples, 5),
"p10": pct_from(overall_samples, 10),
"p25": pct_from(overall_samples, 25),
"p50": pct_from(overall_samples, 50),
"p75": pct_from(overall_samples, 75),
"p90": pct_from(overall_samples, 90),
"p95": pct_from(overall_samples, 95),
"min": float(min(overall_samples)) if overall_samples else 0.0,
"max": float(max(overall_samples)) if overall_samples else 0.0,
},
"per_gpu": per_gpu_summary,
"raw": {
"overall": overall_samples,
"per_gpu": per_gpu_samples,
},
}
json.dump(out_payload, f)
except Exception:
pass
try:
import pynvml # type: ignore
pynvml.nvmlShutdown()
except Exception:
pass
@pytest.fixture(scope="session")
def genai_bench_runner() -> Callable[..., None]:
"""Provide a callable to run genai-bench and validate metrics.
Usage in tests:
def test(..., genai_bench_runner):
genai_bench_runner(router_url=..., model_path=..., experiment_folder=...)
"""
def _run(
*,
router_url: str,
model_path: str,
experiment_folder: str,
timeout_sec: int | None = None,
thresholds: dict | None = None,
extra_env: dict | None = None,
num_concurrency: int = 32,
traffic_scenario: str = "D(4000,100)",
max_requests_per_run: int | None = None,
clean_experiment: bool = True,
kill_procs: list | None = None,
drain_delay_sec: int = 6,
) -> None:
cli = _which("genai-bench")
if not cli:
pytest.fail(
"genai-bench CLI not found; please install it to run benchmarks"
)
# Clean previous experiment folder under current working directory
if clean_experiment:
exp_dir = Path.cwd() / experiment_folder
if exp_dir.exists():
shutil.rmtree(exp_dir, ignore_errors=True)
# Default requests per run if not provided
mrr = (
max_requests_per_run
if max_requests_per_run is not None
else num_concurrency * 5
)
cmd = [
cli,
"benchmark",
"--api-backend",
"openai",
"--api-base",
router_url,
"--api-key",
"dummy-token",
"--api-model-name",
model_path,
"--model-tokenizer",
model_path,
"--task",
"text-to-text",
"--num-concurrency",
str(num_concurrency),
"--traffic-scenario",
traffic_scenario,
"--max-requests-per-run",
str(mrr),
"--max-time-per-run",
"3",
"--experiment-folder-name",
experiment_folder,
"--experiment-base-dir",
str(Path.cwd()),
]
env = os.environ.copy()
if extra_env:
env.update(extra_env)
to = timeout_sec or int(os.environ.get("GENAI_BENCH_TEST_TIMEOUT", "120"))
proc = subprocess.Popen(
cmd, env=env, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True
)
# Optional GPU utilization monitor in a low-priority child process (pynvml only)
# Enabled only when gpu_util_mean_min is provided in thresholds.
monitor_path = None
monitor_proc = None
gpu_util_result: dict | None = None
want_gpu_monitor = _gpu_monitor_should_run(thresholds)
if want_gpu_monitor:
interval = float(os.environ.get("GPU_UTIL_SAMPLE_INTERVAL", "2.0"))
monitor_proc, monitor_path = _launch_gpu_monitor(
bench_pid=proc.pid,
experiment_folder=experiment_folder,
interval=interval,
)
stdout = stderr = ""
rc = None
try:
try:
stdout, stderr = proc.communicate(timeout=to)
except subprocess.TimeoutExpired:
# Simple: kill the CLI process if it doesn't exit in time
try:
proc.kill()
except Exception:
pass
stdout, stderr = proc.communicate()
rc = proc.returncode
# Prefer exact path under cwd; fallback to rglob search
base = Path.cwd()
direct = base / experiment_folder
candidates = [direct] if direct.is_dir() else []
if not candidates:
for p in base.rglob(experiment_folder):
if p.is_dir() and p.name == experiment_folder:
candidates = [p]
break
if not candidates:
raise AssertionError(
"Benchmark failed: experiment folder not found: "
f"{experiment_folder}\nExit code: {rc}\nSTDOUT (tail):\n{stdout[-1000:]}\nSTDERR (tail):\n{stderr[-1000:]}"
)
actual_folder = candidates[0]
json_files = []
for _ in range(10):
json_files = [
p
for p in actual_folder.rglob("*.json")
if "experiment_metadata" not in p.name
]
if json_files:
break
time.sleep(1)
if not json_files:
raise AssertionError(
"Benchmark failed: no JSON results found\n"
f"Exit code: {rc}\nSTDOUT (tail):\n{stdout[-1000:]}\nSTDERR (tail):\n{stderr[-1000:]}"
)
th = thresholds # None means "log only", no validation
for jf in json_files:
with jf.open("r") as f:
data = json.load(f)
stats = data.get("aggregated_metrics", {}).get("stats", {})
ttft_mean = float(stats.get("ttft", {}).get("mean", float("inf")))
e2e_latency_mean = float(
stats.get("e2e_latency", {}).get("mean", float("inf"))
)
input_tp_mean = float(
stats.get("input_throughput", {}).get("mean", 0.0)
)
output_tp_mean = float(
stats.get("output_throughput", {}).get("mean", 0.0)
)
logger.info(
"genai-bench[%s] %s ttft_mean=%.3fs e2e_latency_mean=%.3fs input_tp_mean=%.1f tok/s output_tp_mean=%.1f tok/s",
experiment_folder,
jf.name,
ttft_mean,
e2e_latency_mean,
input_tp_mean,
output_tp_mean,
)
if th is not None:
assert (
ttft_mean <= th["ttft_mean_max"]
), f"TTFT validation failed: {ttft_mean} > {th['ttft_mean_max']} (file={jf.name})"
assert (
e2e_latency_mean <= th["e2e_latency_mean_max"]
), f"E2E latency validation failed: {e2e_latency_mean} > {th['e2e_latency_mean_max']} (file={jf.name})"
assert (
input_tp_mean >= th["input_throughput_mean_min"]
), f"Input throughput validation failed: {input_tp_mean} < {th['input_throughput_mean_min']} (file={jf.name})"
assert (
output_tp_mean >= th["output_throughput_mean_min"]
), f"Output throughput validation failed: {output_tp_mean} < {th['output_throughput_mean_min']} (file={jf.name})"
# Validate optional GPU utilization threshold if provided
if want_gpu_monitor:
try:
if monitor_proc is not None:
monitor_proc.join(timeout=5)
except Exception:
pass
gpu_util_result = _read_gpu_monitor_result(monitor_path)
_log_and_assert_gpu_thresholds(gpu_util_result, thresholds)
finally:
# Always attempt to stop workers to avoid resource leakage
if kill_procs:
# Give router/workers a small grace period to finish any last drains
if drain_delay_sec > 0:
try:
time.sleep(drain_delay_sec)
except Exception:
pass
for p in kill_procs:
_graceful_stop_any(p)
try:
time.sleep(2)
except Exception:
pass
# Ensure GPU monitor process is cleaned up
if monitor_proc is not None and monitor_proc.is_alive():
try:
monitor_proc.terminate()
except Exception:
pass
return _run
def pytest_configure(config):
config.addinivalue_line("markers", "e2e: mark as end-to-end test")
@pytest.fixture(scope="session")
def e2e_model() -> str:
# Always use the default test model
return os.getenv("E2E_PRIMARY_MODEL", DEFAULT_MODEL_NAME_FOR_TEST)
@pytest.fixture
def e2e_router(e2e_model: str):
# Keep this available but tests below use router-only to avoid GPU contention
base_url = DEFAULT_URL_FOR_TEST
proc = _popen_launch_router(
e2e_model, base_url, dp_size=2, timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH
)
try:
yield SimpleNamespace(proc=proc, url=base_url)
finally:
_terminate(proc)
@pytest.fixture
def e2e_router_only_rr():
port = _find_available_port()
base_url = f"http://127.0.0.1:{port}"
proc = _popen_launch_router_only(base_url, policy="round_robin")
try:
yield SimpleNamespace(proc=proc, url=base_url)
finally:
_terminate(proc)
@pytest.fixture(scope="session")
def e2e_embedding_model() -> str:
"""Embedding model to use for E2E tests.
Defaults to an E5 Mistral model, can be overridden via E2E_EMBEDDING_MODEL env var.
"""
import os
return os.getenv("E2E_EMBEDDING_MODEL", "intfloat/e5-mistral-7b-instruct")
@pytest.fixture
def e2e_primary_embedding_worker(e2e_embedding_model: str):
"""Launch a single embedding worker using the specified model."""
port = _find_available_port()
base_url = f"http://127.0.0.1:{port}"
proc = _popen_launch_worker(e2e_embedding_model, base_url)
try:
yield SimpleNamespace(proc=proc, url=base_url)
finally:
_terminate(proc)
@pytest.fixture(scope="session")
def e2e_primary_worker(e2e_model: str):
port = _find_available_port()
base_url = f"http://127.0.0.1:{port}"
proc = _popen_launch_worker(e2e_model, base_url)
# Router health gate will handle worker readiness
try:
yield SimpleNamespace(proc=proc, url=base_url)
finally:
_terminate(proc)
@pytest.fixture
def e2e_router_only_rr_dp_aware_api():
"""Router-only with dp-aware enabled and an API key."""
port = _find_available_port()
base_url = f"http://127.0.0.1:{port}"
api_key = "secret"
proc = _popen_launch_router_only(
base_url, policy="round_robin", timeout=180.0, dp_aware=True, api_key=api_key
)
try:
yield SimpleNamespace(proc=proc, url=base_url, api_key=api_key)
finally:
_terminate(proc)
@pytest.fixture
def e2e_worker_dp2_api(e2e_model: str, e2e_router_only_rr_dp_aware_api):
"""Worker with dp-size=2 and the same API key as the dp-aware router."""
port = _find_available_port()
base_url = f"http://127.0.0.1:{port}"
api_key = e2e_router_only_rr_dp_aware_api.api_key
proc = _popen_launch_worker(e2e_model, base_url, dp_size=2, api_key=api_key)
try:
yield SimpleNamespace(proc=proc, url=base_url)
finally:
_terminate(proc)
@pytest.fixture(scope="session")
def e2e_two_workers_dp2(e2e_model: str):
"""Launch two workers, each with dp_size=2, mapped to GPUs [0,1] and [2,3]."""
workers = []
try:
# Worker A on GPUs 0-1
port_a = _find_available_port()
url_a = f"http://127.0.0.1:{port_a}"
proc_a = _popen_launch_worker(e2e_model, url_a, dp_size=2, base_gpu_id=0)
workers.append(SimpleNamespace(proc=proc_a, url=url_a))
# Worker B on GPUs 2-3
port_b = _find_available_port()
url_b = f"http://127.0.0.1:{port_b}"
proc_b = _popen_launch_worker(e2e_model, url_b, dp_size=2, base_gpu_id=2)
workers.append(SimpleNamespace(proc=proc_b, url=url_b))
yield workers
finally:
for w in workers:
_terminate(w.proc)
sgl-model-gateway/e2e_test/e2e_http/test_pd_router.py
-245
View File
@@ -1,245 +0,0 @@
import logging
import socket
import subprocess
import time
from types import SimpleNamespace
from typing import Optional
import pytest
import requests
logger = logging.getLogger(__name__)
def _find_available_port() -> int:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.bind(("127.0.0.1", 0))
return s.getsockname()[1]
def _wait_health(url: str, timeout: float = 180.0) -> None:
start = time.perf_counter()
with requests.Session() as session:
while time.perf_counter() - start < timeout:
try:
r = session.get(f"{url}/health", timeout=5)
if r.status_code == 200:
return
except requests.RequestException:
pass
time.sleep(1)
raise TimeoutError(f"Service at {url} failed to become healthy in time")
def _detect_ib_device() -> Optional[str]:
"""Return first active IB device name (e.g., mlx5_0) or None if unavailable."""
# Fast check that ibv_devinfo exists
try:
subprocess.run(
["ibv_devinfo", "-l"],
stdout=subprocess.DEVNULL,
stderr=subprocess.DEVNULL,
timeout=1,
)
except (FileNotFoundError, subprocess.TimeoutExpired):
return None
for i in range(12):
dev = f"mlx5_{i}"
try:
res = subprocess.run(
["ibv_devinfo", dev],
capture_output=True,
text=True,
timeout=2,
)
if res.returncode == 0 and ("state:" in res.stdout):
for line in res.stdout.splitlines():
if "state:" in line and "PORT_ACTIVE" in line:
return dev
except Exception:
pass
return None
def _popen_launch_prefill_worker(
model: str,
bootstrap_port: int,
ib_device: Optional[str] = None,
base_gpu_id: int = 0,
) -> SimpleNamespace:
port = _find_available_port()
url = f"http://127.0.0.1:{port}"
cmd = [
"python3",
"-m",
"sglang.launch_server",
"--model-path",
model,
"--disaggregation-mode",
"prefill",
"--host",
"127.0.0.1",
"--port",
str(port),
"--disaggregation-bootstrap-port",
str(bootstrap_port),
"--base-gpu-id",
str(base_gpu_id),
]
if ib_device:
cmd += ["--disaggregation-ib-device", ib_device]
proc = subprocess.Popen(cmd)
_wait_health(url, timeout=300.0)
return SimpleNamespace(proc=proc, url=url, bootstrap_port=bootstrap_port)
def _popen_launch_decode_worker(
model: str, ib_device: Optional[str] = None, base_gpu_id: int = 0
) -> SimpleNamespace:
port = _find_available_port()
url = f"http://127.0.0.1:{port}"
cmd = [
"python3",
"-m",
"sglang.launch_server",
"--model-path",
model,
"--disaggregation-mode",
"decode",
"--host",
"127.0.0.1",
"--port",
str(port),
"--base-gpu-id",
str(base_gpu_id),
]
if ib_device:
cmd += ["--disaggregation-ib-device", ib_device]
proc = subprocess.Popen(cmd)
_wait_health(url, timeout=300.0)
return SimpleNamespace(proc=proc, url=url)
def _terminate(proc: subprocess.Popen, timeout: float = 120) -> None:
if proc is None:
return
proc.terminate()
start = time.perf_counter()
while proc.poll() is None:
if time.perf_counter() - start > timeout:
proc.kill()
break
time.sleep(1)
@pytest.fixture(scope="module")
def pd_cluster(e2e_model: str):
"""Start 2 prefill + 2 decode workers and one PD router, once per module."""
# Environment capability checks: require sgl_kernel and GPU backend
try:
import sgl_kernel # noqa: F401
except Exception as e: # pragma: no cover - environment dependent
pytest.fail(f"PD e2e requires sgl_kernel but it is not available: {e}")
try:
import torch # noqa: F401
except Exception as e: # pragma: no cover - environment dependent
pytest.fail(
f"PD e2e requires torch but it is not available or misconfigured: {e}"
)
if not torch.cuda.is_available(): # pragma: no cover - environment dependent
pytest.fail("PD e2e requires CUDA backend, but CUDA is not available")
workers: list[SimpleNamespace] = []
router_proc = None
try:
ib_device = _detect_ib_device()
# Launch 4 workers across 4 GPUs: prefill on 0,1 and decode on 2,3
pf1 = _popen_launch_prefill_worker(
e2e_model,
bootstrap_port=_find_available_port(),
ib_device=ib_device,
base_gpu_id=0,
)
pf2 = _popen_launch_prefill_worker(
e2e_model,
bootstrap_port=_find_available_port(),
ib_device=ib_device,
base_gpu_id=1,
)
dc1 = _popen_launch_decode_worker(e2e_model, ib_device=ib_device, base_gpu_id=2)
dc2 = _popen_launch_decode_worker(e2e_model, ib_device=ib_device, base_gpu_id=3)
prefills = [pf1, pf2]
decodes = [dc1, dc2]
workers.extend(prefills + decodes)
# PD router with two prefill and two decode endpoints
rport = _find_available_port()
router_url = f"http://127.0.0.1:{rport}"
pport = _find_available_port()
prefill = [(pf.url, pf.bootstrap_port) for pf in prefills]
decode = [dc.url for dc in decodes]
cmd = [
"python3",
"-m",
"sglang_router.launch_router",
"--host",
"127.0.0.1",
"--port",
str(rport),
"--policy",
"round_robin",
"--pd-disaggregation",
"--log-level",
"warn",
]
for url, bport in prefill:
cmd += ["--prefill", url, str(bport)]
for url in decode:
cmd += ["--decode", url]
cmd += [
"--prometheus-port",
str(pport),
"--prometheus-host",
"127.0.0.1",
]
router_proc = subprocess.Popen(cmd)
_wait_health(router_url, timeout=180.0)
yield SimpleNamespace(
router_url=router_url, workers=workers, router_proc=router_proc
)
finally:
if router_proc is not None:
_terminate(router_proc)
for w in workers:
_terminate(w.proc)
@pytest.mark.e2e
def test_pd_genai_bench(e2e_model: str, pd_cluster, genai_bench_runner):
"""
Launch 4 workers, start a PD router (2 prefill + 2 decode), then run a
short genai-bench benchmark and validate aggregate metrics.
"""
# Run genai-bench against the shared router
policy_label = "benchmark_round_robin_pd"
genai_bench_runner(
router_url=pd_cluster.router_url,
model_path=e2e_model,
experiment_folder=policy_label,
thresholds={
"ttft_mean_max": 13,
"e2e_latency_mean_max": 16,
"input_throughput_mean_min": 350,
"output_throughput_mean_min": 18,
"gpu_util_p50_min": 99,
},
kill_procs=pd_cluster.workers,
)
sgl-model-gateway/e2e_test/e2e_http/test_regular_router.py
-56
View File
@@ -1,56 +0,0 @@
import time
import pytest
import requests
def _wait_for_workers(
base_url: str, expected_count: int, timeout: float = 60.0, headers: dict = None
) -> None:
"""Poll /workers endpoint until expected number of workers are registered."""
start = time.perf_counter()
with requests.Session() as session:
while time.perf_counter() - start < timeout:
try:
r = session.get(f"{base_url}/workers", headers=headers, timeout=5)
if r.status_code == 200:
workers = r.json().get("workers", [])
if len(workers) >= expected_count:
return
except requests.RequestException:
pass
time.sleep(0.5)
raise TimeoutError(
f"Expected {expected_count} workers at {base_url}, timed out after {timeout}s"
)
@pytest.mark.e2e
def test_genai_bench(
e2e_router_only_rr, e2e_two_workers_dp2, e2e_model, genai_bench_runner
):
"""Attach a worker to the regular router and run a short genai-bench."""
base = e2e_router_only_rr.url
for w in e2e_two_workers_dp2:
r = requests.post(f"{base}/workers", json={"url": w.url}, timeout=180)
assert (
r.status_code == 202
), f"Expected 202 ACCEPTED, got {r.status_code}: {r.text}"
# Wait for workers to be registered
_wait_for_workers(base, expected_count=2, timeout=60.0)
genai_bench_runner(
router_url=base,
model_path=e2e_model,
experiment_folder="benchmark_round_robin_regular",
thresholds={
"ttft_mean_max": 6,
"e2e_latency_mean_max": 14,
"input_throughput_mean_min": 800, # temp relax from 1000 to 800 for now
"output_throughput_mean_min": 12,
# Enforce GPU utilization p50 >= 99% during the run.
"gpu_util_p50_min": 99,
},
kill_procs=e2e_two_workers_dp2,
)
sgl-model-gateway/e2e_test/infra/__init__.py
+5
View File
@@ -35,6 +35,8 @@ from .gpu_allocator import (
nvml_context, nvml_context,
wait_for_gpu_memory_to_clear, wait_for_gpu_memory_to_clear,
) )
from .gpu_monitor import GPUMonitor
from .gpu_monitor import should_monitor as should_monitor_gpu
from .model_pool import ModelInstance, ModelPool from .model_pool import ModelInstance, ModelPool
from .model_specs import ( # Default model paths; Model groups from .model_specs import ( # Default model paths; Model groups
CHAT_MODELS, CHAT_MODELS,
@@ -104,6 +106,9 @@ __all__ = [
"wait_for_health", "wait_for_health",
"wait_for_workers_ready", "wait_for_workers_ready",
"detect_ib_device", "detect_ib_device",
# GPU monitoring
"GPUMonitor",
"should_monitor_gpu",
# Model management # Model management
"ModelInstance", "ModelInstance",
"ModelPool", "ModelPool",
sgl-model-gateway/e2e_test/infra/gpu_monitor.py
+329
View File
@@ -0,0 +1,329 @@
"""GPU utilization monitoring for benchmarks.
This module provides a low-impact GPU monitor that runs in a separate process
and collects utilization samples using NVML.
"""
from __future__ import annotations
import json
import logging
import os
import time
from multiprocessing import Process
from pathlib import Path
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from typing import Any
logger = logging.getLogger(__name__)
def _percentile(samples: list[float], p: float) -> float:
"""Calculate percentile from sorted samples."""
if not samples:
return 0.0
sorted_samples = sorted(samples)
idx = max(
0,
min(
len(sorted_samples) - 1, int(round((p / 100.0) * (len(sorted_samples) - 1)))
),
)
return float(sorted_samples[idx])
def _compute_stats(samples: list[float]) -> dict[str, float]:
"""Compute statistics for a list of samples."""
if not samples:
return {
"mean": 0.0,
"min": 0.0,
"max": 0.0,
"p5": 0.0,
"p10": 0.0,
"p25": 0.0,
"p50": 0.0,
"p75": 0.0,
"p90": 0.0,
"p95": 0.0,
"count": 0,
}
return {
"mean": sum(samples) / len(samples),
"min": min(samples),
"max": max(samples),
"p5": _percentile(samples, 5),
"p10": _percentile(samples, 10),
"p25": _percentile(samples, 25),
"p50": _percentile(samples, 50),
"p75": _percentile(samples, 75),
"p90": _percentile(samples, 90),
"p95": _percentile(samples, 95),
"count": len(samples),
}
def _monitor_loop(pid: int, output_path: str, interval: float) -> None:
"""Main monitoring loop - runs in separate process.
Monitors GPU utilization until the target process exits, then writes
results to output_path as JSON.
"""
# Lower process priority to minimize impact on benchmark
try:
os.nice(10)
except Exception:
pass
# Initialize NVML
try:
import pynvml
pynvml.nvmlInit()
except Exception as e:
logger.warning("Failed to initialize NVML: %s", e)
_write_empty_result(output_path)
return
# Get GPU handles
try:
device_count = pynvml.nvmlDeviceGetCount()
handles = [pynvml.nvmlDeviceGetHandleByIndex(i) for i in range(device_count)]
except Exception as e:
logger.warning("Failed to get GPU handles: %s", e)
_write_empty_result(output_path)
_shutdown_nvml()
return
# Collect samples
per_gpu_samples: dict[str, list[float]] = {str(i): [] for i in range(device_count)}
overall_samples: list[float] = []
try:
while _process_alive(pid):
try:
gpu_utils = []
for idx, handle in enumerate(handles):
try:
util = pynvml.nvmlDeviceGetUtilizationRates(handle).gpu
gpu_utils.append(float(util))
per_gpu_samples[str(idx)].append(float(util))
except Exception:
continue
if gpu_utils:
avg = sum(gpu_utils) / len(gpu_utils)
overall_samples.append(avg)
except Exception:
pass
time.sleep(interval)
finally:
# Write results
_write_result(output_path, pid, interval, overall_samples, per_gpu_samples)
_shutdown_nvml()
def _process_alive(pid: int) -> bool:
"""Check if process is still running."""
try:
os.kill(pid, 0)
return True
except (OSError, ProcessLookupError):
return False
def _write_empty_result(path: str) -> None:
"""Write empty result file."""
try:
os.makedirs(os.path.dirname(path), exist_ok=True)
with open(path, "w") as f:
json.dump(
{
"count": 0,
"overall": {"mean": 0.0},
"per_gpu": {},
"raw": {"overall": [], "per_gpu": {}},
},
f,
)
except Exception:
pass
def _write_result(
path: str,
pid: int,
interval: float,
overall_samples: list[float],
per_gpu_samples: dict[str, list[float]],
) -> None:
"""Write monitoring results to JSON file."""
try:
os.makedirs(os.path.dirname(path), exist_ok=True)
with open(path, "w") as f:
json.dump(
{
"bench_pid": pid,
"interval_sec": interval,
"count": len(overall_samples),
"overall": _compute_stats(overall_samples),
"per_gpu": {
k: _compute_stats(v) for k, v in per_gpu_samples.items()
},
"raw": {
"overall": overall_samples,
"per_gpu": per_gpu_samples,
},
},
f,
)
except Exception as e:
logger.warning("Failed to write GPU monitor results: %s", e)
def _shutdown_nvml() -> None:
"""Shutdown NVML."""
try:
import pynvml
pynvml.nvmlShutdown()
except Exception:
pass
class GPUMonitor:
"""GPU utilization monitor for benchmarks.
Usage:
monitor = GPUMonitor(output_dir="benchmark_results")
monitor.start(target_pid=12345)
# ... run benchmark ...
result = monitor.stop()
monitor.assert_thresholds({"gpu_util_p50_min": 99})
"""
def __init__(
self,
output_dir: str | Path = ".",
interval: float = 2.0,
):
self.output_dir = Path(output_dir)
self.interval = interval
self._process: Process | None = None
self._output_path: str | None = None
self._result: dict[str, Any] | None = None
@property
def output_path(self) -> str | None:
"""Path to the GPU utilization JSON file."""
return self._output_path
def start(self, target_pid: int) -> None:
"""Start monitoring GPU utilization for the target process."""
self._output_path = str(self.output_dir / "gpu_utilization.json")
self._result = None
self._process = Process(
target=_monitor_loop,
args=(target_pid, self._output_path, self.interval),
daemon=True,
)
self._process.start()
logger.debug("Started GPU monitor for PID %d", target_pid)
def stop(self, timeout: float = 5.0) -> dict[str, Any] | None:
"""Stop monitoring and return results."""
if self._process is None:
return None
try:
self._process.join(timeout=timeout)
except Exception:
pass
if self._process.is_alive():
try:
self._process.terminate()
except Exception:
pass
self._process = None
self._result = self._read_result()
return self._result
def _read_result(self) -> dict[str, Any] | None:
"""Read results from output file."""
if not self._output_path or not os.path.exists(self._output_path):
return None
try:
with open(self._output_path) as f:
return json.load(f)
except Exception as e:
logger.warning("Failed to read GPU monitor result: %s", e)
return None
def log_summary(self) -> None:
"""Log a summary of GPU utilization."""
result = self._result or self._read_result()
if not result or result.get("count", 0) <= 0:
logger.warning("GPU utilization monitor produced no samples")
return
overall = result.get("overall", {})
logger.info(
"GPU utilization: mean=%.2f%% p50=%.2f%% (samples=%d)",
overall.get("mean", 0.0),
overall.get("p50", 0.0),
result.get("count", 0),
)
def assert_thresholds(self, thresholds: dict[str, float] | None) -> None:
"""Assert GPU utilization meets thresholds.
Supported thresholds:
- gpu_util_mean_min: Minimum mean GPU utilization %
- gpu_util_p50_min: Minimum p50 GPU utilization %
"""
if not thresholds:
return
result = self._result or self._read_result()
if not result or result.get("count", 0) <= 0:
logger.warning("GPU utilization monitor produced no samples")
return
overall = result.get("overall", {})
mean_threshold = thresholds.get("gpu_util_mean_min")
if mean_threshold is not None:
mean_value = overall.get("mean", 0.0)
assert (
mean_value >= mean_threshold
), f"GPU utilization mean below threshold: {mean_value:.2f}% < {mean_threshold}%"
p50_threshold = thresholds.get("gpu_util_p50_min")
if p50_threshold is not None:
p50_value = overall.get("p50")
if p50_value is not None:
assert (
p50_value >= p50_threshold
), f"GPU utilization p50 below threshold: {p50_value:.2f}% < {p50_threshold}%"
def should_monitor(thresholds: dict[str, Any] | None) -> bool:
"""Check if GPU monitoring should be enabled.
Returns True if:
- thresholds contains gpu_util_mean_min or gpu_util_p50_min, OR
- GPU_UTIL_LOG environment variable is truthy
"""
if thresholds:
if thresholds.get("gpu_util_mean_min") is not None:
return True
if thresholds.get("gpu_util_p50_min") is not None:
return True
return os.environ.get("GPU_UTIL_LOG", "").lower() in ("1", "true", "yes")
sgl-model-gateway/e2e_test/infra/model_pool.py
+153 -17
View File
@@ -293,6 +293,7 @@ class ModelPool:
worker_type: WorkerType = WorkerType.REGULAR, worker_type: WorkerType = WorkerType.REGULAR,
bootstrap_port: int | None = None, bootstrap_port: int | None = None,
ib_device: str | None = None, ib_device: str | None = None,
instance_key: str | None = None,
) -> ModelInstance: ) -> ModelInstance:
"""Launch a model instance. """Launch a model instance.
@@ -303,6 +304,7 @@ class ModelPool:
worker_type: Worker type (REGULAR, PREFILL, or DECODE). worker_type: Worker type (REGULAR, PREFILL, or DECODE).
bootstrap_port: Bootstrap port for prefill workers in PD mode. bootstrap_port: Bootstrap port for prefill workers in PD mode.
ib_device: InfiniBand device for PD disaggregation. ib_device: InfiniBand device for PD disaggregation.
instance_key: Custom instance key, or None to auto-generate.
Returns: Returns:
The launched ModelInstance. The launched ModelInstance.
@@ -353,11 +355,15 @@ class ModelPool:
cmd.extend(["--disaggregation-ib-device", ib_device]) cmd.extend(["--disaggregation-ib-device", ib_device])
elif worker_type == WorkerType.DECODE: elif worker_type == WorkerType.DECODE:
cmd.extend(["--disaggregation-mode", "decode"]) cmd.extend(["--disaggregation-mode", "decode"])
# Base GPU ID 0 since CUDA_VISIBLE_DEVICES remaps the GPU
cmd.extend(["--base-gpu-id", "0"])
if ib_device: if ib_device:
cmd.extend(["--disaggregation-ib-device", ib_device]) cmd.extend(["--disaggregation-ib-device", ib_device])
# Build key based on worker type # Build key based on worker type (or use custom key)
if worker_type == WorkerType.REGULAR: if instance_key:
key = instance_key
elif worker_type == WorkerType.REGULAR:
key = f"{model_id}:{mode.value}" key = f"{model_id}:{mode.value}"
else: else:
key = f"{model_id}:{mode.value}:{worker_type.value}" key = f"{model_id}:{mode.value}:{worker_type.value}"
@@ -560,7 +566,11 @@ class ModelPool:
return instance return instance
def _evict_for_gpus( def _evict_for_gpus(
self, required_gpus: int, exclude_model_id: str | None = None self,
required_gpus: int,
exclude_model_id: str | None = None,
exclude_mode: ConnectionMode | None = None,
exclude_worker_types: set[WorkerType] | None = None,
) -> None: ) -> None:
"""Evict models until we have enough GPUs available. """Evict models until we have enough GPUs available.
@@ -570,29 +580,45 @@ class ModelPool:
Args: Args:
required_gpus: Number of GPUs needed. required_gpus: Number of GPUs needed.
exclude_model_id: Model ID to exclude from eviction (test may need exclude_model_id: Model ID to exclude from eviction.
multiple modes of the same model). exclude_mode: Connection mode to exclude from eviction (optional).
exclude_worker_types: Worker types to exclude from eviction.
If None, falls back to excluding by model_id only (backward compatible).
""" """
available = self.allocator.available_gpus() available = self.allocator.available_gpus()
if len(available) >= required_gpus: if len(available) >= required_gpus:
return # Already have enough return # Already have enough
# Sort by last_used descending (MRU eviction) - evict most recently used first # Sort by last_used descending (MRU eviction) - evict most recently used first
# Exclude instances of the same model_id (test may need multiple modes) # Store (dict_key, instance) tuples to preserve the actual key for eviction
evictable = [ evictable: list[tuple[str, ModelInstance]] = []
inst for dict_key, inst in self.instances.items():
for inst in self.instances.values() if exclude_worker_types is not None:
if exclude_model_id is None or inst.model_id != exclude_model_id # Precise matching with worker types
] # Must match model_id AND worker_type, mode is optional
evictable.sort(key=lambda x: x.last_used, reverse=True) if (
exclude_model_id is not None
and inst.model_id == exclude_model_id
and inst.worker_type in exclude_worker_types
):
# If mode is specified, also require mode match
if exclude_mode is None or inst.mode == exclude_mode:
continue
else:
# Backward compatible: exclude by model_id only
if exclude_model_id is not None and inst.model_id == exclude_model_id:
continue
evictable.append((dict_key, inst))
evictable.sort(key=lambda x: x[1].last_used, reverse=True)
freed_gpus = len(available) freed_gpus = len(available)
for inst in evictable: for dict_key, inst in evictable:
if freed_gpus >= required_gpus: if freed_gpus >= required_gpus:
break break
logger.info("Evicting model %s (MRU) to free GPUs", inst.key) logger.info("Evicting model %s (MRU) to free GPUs", dict_key)
self._evict_instance(inst.key) self._evict_instance(dict_key)
if inst.gpu_slot: if inst.gpu_slot:
freed_gpus += len(inst.gpu_slot.gpu_ids) freed_gpus += len(inst.gpu_slot.gpu_ids)
@@ -608,7 +634,13 @@ class ModelPool:
spec = get_model_spec(model_id) spec = get_model_spec(model_id)
required_gpus = spec.get("tp", 1) required_gpus = spec.get("tp", 1)
self._evict_for_gpus(required_gpus, exclude_model_id=model_id) # Exclude REGULAR workers of same model from eviction (keep them)
# but allow evicting PD workers (PREFILL/DECODE) to free GPUs
self._evict_for_gpus(
required_gpus,
exclude_model_id=model_id,
exclude_worker_types={WorkerType.REGULAR},
)
available = self.allocator.available_gpus() available = self.allocator.available_gpus()
if len(available) < required_gpus: if len(available) < required_gpus:
@@ -682,6 +714,102 @@ class ModelPool:
if inst.model_id == model_id and inst.worker_type == worker_type if inst.model_id == model_id and inst.worker_type == worker_type
] ]
def launch_regular_workers(
self,
model_id: str,
num_workers: int,
mode: ConnectionMode = ConnectionMode.HTTP,
startup_timeout: int = DEFAULT_STARTUP_TIMEOUT,
allow_eviction: bool = True,
) -> list[ModelInstance]:
"""Launch multiple regular workers for load balancing.
Args:
model_id: Model identifier from MODEL_SPECS.
num_workers: Number of workers to launch.
mode: Connection mode (HTTP or GRPC).
startup_timeout: Timeout for workers to become healthy.
allow_eviction: If True, evict MRU models to free GPUs.
Returns:
List of ModelInstance objects.
"""
self._startup_timeout = startup_timeout
if model_id not in MODEL_SPECS:
raise ValueError(f"Unknown model: {model_id}")
spec = get_model_spec(model_id)
tp = spec.get("tp", 1)
required_gpus = num_workers * tp
# Check if we have enough GPUs
available = self.allocator.available_gpus()
if len(available) < required_gpus:
if allow_eviction:
logger.info(
"Need %d GPUs for %d workers, only %d available. Evicting MRU models...",
required_gpus,
num_workers,
len(available),
)
# Exclude REGULAR workers of same model/mode from eviction
self._evict_for_gpus(
required_gpus,
exclude_model_id=model_id,
exclude_mode=mode,
exclude_worker_types={WorkerType.REGULAR},
)
else:
logger.info(
"Need %d GPUs for %d workers, only %d available. "
"Skipping (eviction not allowed).",
required_gpus,
num_workers,
len(available),
)
return []
# Build allocation specs for all workers
allocation_specs = {}
for i in range(num_workers):
key = f"{model_id}:{mode.value}:{i}"
allocation_specs[key] = {
"model": spec["model"],
"memory_gb": spec.get("memory_gb", 16),
"tp": tp,
}
# Allocate GPU slots
slots = self.allocator.allocate_slots(allocation_specs)
slot_map = {slot.assigned_model: slot for slot in slots}
if not slots:
raise RuntimeError(
f"Failed to allocate GPU slots for {num_workers} workers after eviction. "
f"Need {required_gpus} GPUs."
)
instances: list[ModelInstance] = []
# Launch workers
for i in range(num_workers):
key = f"{model_id}:{mode.value}:{i}"
gpu_slot = slot_map.get(key)
instance = self._launch_model(
model_id=model_id,
mode=mode,
gpu_slot=gpu_slot,
worker_type=WorkerType.REGULAR,
instance_key=key,
)
instances.append(instance)
# Wait for all to be healthy
self._wait_all_healthy()
return instances
def launch_pd_workers( def launch_pd_workers(
self, self,
model_id: str, model_id: str,
@@ -728,7 +856,13 @@ class ModelPool:
required_gpus, required_gpus,
len(available), len(available),
) )
self._evict_for_gpus(required_gpus, exclude_model_id=model_id) # Exclude PD workers of same model/mode, but evict REGULAR workers
self._evict_for_gpus(
required_gpus,
exclude_model_id=model_id,
exclude_mode=mode,
exclude_worker_types={WorkerType.PREFILL, WorkerType.DECODE},
)
else: else:
logger.info( logger.info(
"Need %d GPUs for PD workers, only %d available. " "Need %d GPUs for PD workers, only %d available. "
@@ -781,6 +915,7 @@ class ModelPool:
worker_type=WorkerType.PREFILL, worker_type=WorkerType.PREFILL,
bootstrap_port=bootstrap_port, bootstrap_port=bootstrap_port,
ib_device=ib_device, ib_device=ib_device,
instance_key=key,
) )
prefill_instances.append(instance) prefill_instances.append(instance)
@@ -794,6 +929,7 @@ class ModelPool:
gpu_slot=gpu_slot, gpu_slot=gpu_slot,
worker_type=WorkerType.DECODE, worker_type=WorkerType.DECODE,
ib_device=ib_device, ib_device=ib_device,
instance_key=key,
) )
decode_instances.append(instance) decode_instances.append(instance)