diff --git a/.github/workflows/pr-test-pd-router.yml b/.github/workflows/pr-test-pd-router.yml deleted file mode 100644 index edc46de82..000000000 --- a/.github/workflows/pr-test-pd-router.yml +++ /dev/null @@ -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 diff --git a/.github/workflows/pr-test-rust.yml b/.github/workflows/pr-test-rust.yml index 6e4f5fc38..7f2e94bf1 100644 --- a/.github/workflows/pr-test-rust.yml +++ b/.github/workflows/pr-test-rust.yml @@ -13,7 +13,7 @@ on: workflow_dispatch: concurrency: - group: router-tests-${{ github.ref }} + group: gateway-tests-${{ github.ref }} cancel-in-progress: true env: @@ -56,6 +56,7 @@ jobs: uses: mozilla-actions/sccache-action@v0.0.9 with: version: "v0.12.0" + disable_annotations: true - name: Rust cache uses: Swatinem/rust-cache@v2 @@ -83,7 +84,14 @@ jobs: python -c "import sglang_router; print('Python package: 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" - 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: | github.event_name != 'pull_request' || (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 with: version: "v0.12.0" + disable_annotations: true - name: Rust cache uses: Swatinem/rust-cache@v2 @@ -146,79 +155,45 @@ jobs: if: always() run: sccache --show-stats - router-http-tests: + gateway-e2e: + name: ${{ matrix.name }} 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') + 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 - timeout-minutes: 32 - 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 + timeout-minutes: ${{ matrix.timeout }} steps: - name: Checkout code uses: actions/checkout@v4 @@ -231,6 +206,7 @@ jobs: uses: mozilla-actions/sccache-action@v0.0.9 with: version: "v0.12.0" + disable_annotations: true - name: Rust cache uses: Swatinem/rust-cache@v2 @@ -246,6 +222,7 @@ jobs: sudo --preserve-env=PATH bash scripts/ci/ci_install_dependency.sh - name: Setup Oracle Instant Client + if: matrix.setup_oracle run: | sudo apt-get install -y unzip 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 - name: Start Oracle Database + if: matrix.setup_oracle run: | docker run -d -p 1521:1521 -e ORACLE_PASSWORD=oracle --name oracle-db gvenzl/oracle-xe:21-slim echo "Starting Oracle DB..." @@ -275,6 +253,7 @@ jobs: echo "ATP_DSN=localhost:1521/XEPDB1" >> $GITHUB_ENV - name: Start Brave MCP Server + if: matrix.setup_brave run: | docker run -d --rm \ -p 8001:8080 \ @@ -295,88 +274,39 @@ jobs: maturin build --profile ci --features vendored-openssl --out dist pip install dist/*.whl - - name: Run Python E2E response API tests + - name: Install e2e test dependencies run: | - python3 -m pip install pytest-rerunfailures - bash scripts/killall_sglang.sh "nuk_gpus" - cd sgl-model-gateway - source "$HOME/.cargo/env" - 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 + python3 -m pip install pytest pytest-rerunfailures httpx openai grpcio grpcio-health-checking numpy + if [ -n "${{ matrix.extra_deps }}" ]; then + python3 -m pip --no-cache-dir install --upgrade ${{ matrix.extra_deps }} + fi - - name: Run Python E2E gRPC tests + - name: Run E2E tests run: | bash scripts/killall_sglang.sh "nuk_gpus" cd sgl-model-gateway 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 - if: always() + if: always() && matrix.setup_brave run: | docker stop brave-search-server || true docker rm brave-search-server || true - name: Cleanup Oracle Database - if: always() + if: always() && matrix.setup_oracle run: | docker stop 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: if: | github.event_name != 'pull_request' || @@ -401,98 +331,25 @@ jobs: cache-to: type=gha,mode=max 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 steps: - name: Finish run: echo "This is an empty step to ensure that all jobs are completed." summarize-benchmarks: - needs: router-http-tests + needs: gateway-e2e runs-on: ubuntu-latest if: success() steps: - - name: Install jq - run: sudo apt-get update && sudo apt-get install -y jq bc + - name: Checkout code + uses: actions/checkout@v4 - 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 || 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 + run: python3 sgl-model-gateway/e2e_test/benchmarks/summarize.py . diff --git a/sgl-model-gateway/e2e_test/benchmarks/conftest.py b/sgl-model-gateway/e2e_test/benchmarks/conftest.py new file mode 100644 index 000000000..6d99e4a3f --- /dev/null +++ b/sgl-model-gateway/e2e_test/benchmarks/conftest.py @@ -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 diff --git a/sgl-model-gateway/e2e_test/benchmarks/results.py b/sgl-model-gateway/e2e_test/benchmarks/results.py new file mode 100644 index 000000000..4f60c8758 --- /dev/null +++ b/sgl-model-gateway/e2e_test/benchmarks/results.py @@ -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 diff --git a/sgl-model-gateway/e2e_test/benchmarks/summarize.py b/sgl-model-gateway/e2e_test/benchmarks/summarize.py new file mode 100644 index 000000000..dbd0b838c --- /dev/null +++ b/sgl-model-gateway/e2e_test/benchmarks/summarize.py @@ -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() diff --git a/sgl-model-gateway/e2e_test/benchmarks/test_pd_perf.py b/sgl-model-gateway/e2e_test/benchmarks/test_pd_perf.py new file mode 100644 index 000000000..8075e12c5 --- /dev/null +++ b/sgl-model-gateway/e2e_test/benchmarks/test_pd_perf.py @@ -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, + }, + ) diff --git a/sgl-model-gateway/e2e_test/benchmarks/test_regular_perf.py b/sgl-model-gateway/e2e_test/benchmarks/test_regular_perf.py new file mode 100644 index 000000000..3b71b1619 --- /dev/null +++ b/sgl-model-gateway/e2e_test/benchmarks/test_regular_perf.py @@ -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, + }, + ) diff --git a/sgl-model-gateway/e2e_test/conftest.py b/sgl-model-gateway/e2e_test/conftest.py index 7d1db2ea1..96eac4db5 100644 --- a/sgl-model-gateway/e2e_test/conftest.py +++ b/sgl-model-gateway/e2e_test/conftest.py @@ -697,21 +697,33 @@ def setup_backend(request: pytest.FixtureRequest, model_pool: "ModelPool"): num_workers = workers_config.get("count") or 1 try: - instance = model_pool.get(model_id, connection_mode) + 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) + worker_urls = [instance.worker_url] + model_path = instance.model_path except RuntimeError as 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 gateway = Gateway() gateway.start( worker_urls=worker_urls, - model_path=instance.model_path, + model_path=model_path, policy=gateway_config["policy"], timeout=gateway_config["timeout"], extra_args=gateway_config["extra_args"], @@ -732,7 +744,7 @@ def setup_backend(request: pytest.FixtureRequest, model_pool: "ModelPool"): ) try: - yield backend_name, instance.model_path, client, gateway + yield backend_name, model_path, client, gateway finally: logger.info("Tearing down gateway for %s backend", backend_name) gateway.shutdown() diff --git a/sgl-model-gateway/e2e_test/e2e_http/conftest.py b/sgl-model-gateway/e2e_test/e2e_http/conftest.py deleted file mode 100644 index 4e0b241c2..000000000 --- a/sgl-model-gateway/e2e_test/e2e_http/conftest.py +++ /dev/null @@ -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) diff --git a/sgl-model-gateway/e2e_test/e2e_http/test_pd_router.py b/sgl-model-gateway/e2e_test/e2e_http/test_pd_router.py deleted file mode 100644 index ef7c593ab..000000000 --- a/sgl-model-gateway/e2e_test/e2e_http/test_pd_router.py +++ /dev/null @@ -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, - ) diff --git a/sgl-model-gateway/e2e_test/e2e_http/test_regular_router.py b/sgl-model-gateway/e2e_test/e2e_http/test_regular_router.py deleted file mode 100644 index 34a4eced0..000000000 --- a/sgl-model-gateway/e2e_test/e2e_http/test_regular_router.py +++ /dev/null @@ -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, - ) diff --git a/sgl-model-gateway/e2e_test/infra/__init__.py b/sgl-model-gateway/e2e_test/infra/__init__.py index 6b92a9885..6bb4d8448 100644 --- a/sgl-model-gateway/e2e_test/infra/__init__.py +++ b/sgl-model-gateway/e2e_test/infra/__init__.py @@ -35,6 +35,8 @@ from .gpu_allocator import ( nvml_context, 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_specs import ( # Default model paths; Model groups CHAT_MODELS, @@ -104,6 +106,9 @@ __all__ = [ "wait_for_health", "wait_for_workers_ready", "detect_ib_device", + # GPU monitoring + "GPUMonitor", + "should_monitor_gpu", # Model management "ModelInstance", "ModelPool", diff --git a/sgl-model-gateway/e2e_test/infra/gpu_monitor.py b/sgl-model-gateway/e2e_test/infra/gpu_monitor.py new file mode 100644 index 000000000..870eec2b0 --- /dev/null +++ b/sgl-model-gateway/e2e_test/infra/gpu_monitor.py @@ -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") diff --git a/sgl-model-gateway/e2e_test/infra/model_pool.py b/sgl-model-gateway/e2e_test/infra/model_pool.py index 32697d827..5b01abf12 100644 --- a/sgl-model-gateway/e2e_test/infra/model_pool.py +++ b/sgl-model-gateway/e2e_test/infra/model_pool.py @@ -293,6 +293,7 @@ class ModelPool: worker_type: WorkerType = WorkerType.REGULAR, bootstrap_port: int | None = None, ib_device: str | None = None, + instance_key: str | None = None, ) -> ModelInstance: """Launch a model instance. @@ -303,6 +304,7 @@ class ModelPool: worker_type: Worker type (REGULAR, PREFILL, or DECODE). bootstrap_port: Bootstrap port for prefill workers in PD mode. ib_device: InfiniBand device for PD disaggregation. + instance_key: Custom instance key, or None to auto-generate. Returns: The launched ModelInstance. @@ -353,11 +355,15 @@ class ModelPool: cmd.extend(["--disaggregation-ib-device", ib_device]) elif worker_type == WorkerType.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: cmd.extend(["--disaggregation-ib-device", ib_device]) - # Build key based on worker type - if worker_type == WorkerType.REGULAR: + # Build key based on worker type (or use custom key) + if instance_key: + key = instance_key + elif worker_type == WorkerType.REGULAR: key = f"{model_id}:{mode.value}" else: key = f"{model_id}:{mode.value}:{worker_type.value}" @@ -560,7 +566,11 @@ class ModelPool: return instance 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: """Evict models until we have enough GPUs available. @@ -570,29 +580,45 @@ class ModelPool: Args: required_gpus: Number of GPUs needed. - exclude_model_id: Model ID to exclude from eviction (test may need - multiple modes of the same model). + exclude_model_id: Model ID to exclude from eviction. + 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() if len(available) >= required_gpus: return # Already have enough # Sort by last_used descending (MRU eviction) - evict most recently used first - # Exclude instances of the same model_id (test may need multiple modes) - evictable = [ - inst - for inst in self.instances.values() - if exclude_model_id is None or inst.model_id != exclude_model_id - ] - evictable.sort(key=lambda x: x.last_used, reverse=True) + # Store (dict_key, instance) tuples to preserve the actual key for eviction + evictable: list[tuple[str, ModelInstance]] = [] + for dict_key, inst in self.instances.items(): + if exclude_worker_types is not None: + # Precise matching with worker types + # Must match model_id AND worker_type, mode is optional + 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) - for inst in evictable: + for dict_key, inst in evictable: if freed_gpus >= required_gpus: break - logger.info("Evicting model %s (MRU) to free GPUs", inst.key) - self._evict_instance(inst.key) + logger.info("Evicting model %s (MRU) to free GPUs", dict_key) + self._evict_instance(dict_key) if inst.gpu_slot: freed_gpus += len(inst.gpu_slot.gpu_ids) @@ -608,7 +634,13 @@ class ModelPool: spec = get_model_spec(model_id) 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() if len(available) < required_gpus: @@ -682,6 +714,102 @@ class ModelPool: 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( self, model_id: str, @@ -728,7 +856,13 @@ class ModelPool: required_gpus, 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: logger.info( "Need %d GPUs for PD workers, only %d available. " @@ -781,6 +915,7 @@ class ModelPool: worker_type=WorkerType.PREFILL, bootstrap_port=bootstrap_port, ib_device=ib_device, + instance_key=key, ) prefill_instances.append(instance) @@ -794,6 +929,7 @@ class ModelPool: gpu_slot=gpu_slot, worker_type=WorkerType.DECODE, ib_device=ib_device, + instance_key=key, ) decode_instances.append(instance)