[CI] Add stress testing for router #633
Conversation
- Implements multi-threaded stress testing with configurable parameters - Supports different workload patterns (sustained, spike, gradual) - Includes comprehensive metrics collection and reporting - Provides cleanup mechanisms and proper error handling - Enables scalability testing for production environments Signed-off-by: Kobe Chen <[email protected]>
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Summary of Changes
Hello @kobe0938, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!
I'm adding a new stress test script for the router to validate its round-robin routing logic under high concurrent loads. This test is designed for internal use, not regular CI, due to its execution time.
Highlights
- New Stress Test Script: I've introduced a new shell script, tests/e2e/stress-test.sh, dedicated to stress testing the router.
- Round-Robin Routing Validation: The script specifically targets and validates the router's round-robin routing logic, ensuring requests are evenly distributed across multiple backend servers.
- High Concurrency Simulation: It leverages Apache Bench (ab) to simulate high concurrent request loads, mimicking real-world heavy traffic scenarios.
- Configurable Test Parameters: The test is highly configurable, allowing users to adjust parameters such as concurrent requests, total requests, router port, backend ports, and the model used for testing.
- Internal Use Only: This stress test is intended for internal usage between releases and is not integrated into the regular CI pipeline due to its extended execution time.
Using Gemini Code Assist
The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.
Invoking Gemini
You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| Feature | Command | Description |
|---|---|---|
| Code Review | /gemini review |
Performs a code review for the current pull request in its current state. |
| Pull Request Summary | /gemini summary |
Provides a summary of the current pull request in its current state. |
| Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in issue comments and review comments. |
| Help | /gemini help |
Displays a list of available commands. |
Customization
To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.
Limitations & Feedback
Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments or fill out our survey to provide feedback.
You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.
Footnotes
-
Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩
There was a problem hiding this comment.
Code Review
This pull request introduces a new stress testing script for the router's round-robin logic. The script is well-structured and provides comprehensive testing capabilities. My review includes several suggestions to improve its robustness, security, and maintainability. Key points include using mktemp for secure temporary file creation, safer process cleanup using PIDs instead of pkill, more robust command-line argument parsing, and fixing a minor inconsistency in the help documentation.
| # Function to cleanup processes | ||
| cleanup() { | ||
| print_status "Cleaning up router processes..." | ||
| pkill -f "python3 -m src.vllm_router.app" || true |
There was a problem hiding this comment.
Using pkill -f "python3 -m src.vllm_router.app" is risky as it can terminate other unrelated vllm_router processes running on the same machine. The script already stores the router's process ID in the ROUTER_PID variable. It would be safer to kill the specific process using its PID.
| pkill -f "python3 -m src.vllm_router.app" || true | |
| if [[ -n "${ROUTER_PID-}" ]] && ps -p "$ROUTER_PID" > /dev/null; then | |
| print_status "Stopping router with PID $ROUTER_PID..." | |
| kill "$ROUTER_PID" | |
| else | |
| # Fallback in case PID is not available or process is already gone | |
| pkill -f "python3 -m src.vllm_router.app" || true | |
| fi |
| local payload_file="/tmp/stress_payload.json" | ||
| cat > "$payload_file" << EOF | ||
| { | ||
| "model": "$MODEL", | ||
| "messages": [ | ||
| {"role": "user", "content": "Test message for stress testing"} | ||
| ], | ||
| "max_tokens": 10, | ||
| "temperature": 0.7 | ||
| } | ||
| EOF | ||
|
|
||
| # Run Apache Bench | ||
| ab -c "$CONCURRENT" \ | ||
| -n "$REQUESTS" \ | ||
| -p "$payload_file" \ | ||
| -T "application/json" \ | ||
| -H "Authorization: Bearer test" \ | ||
| -H "x-user-id: stress-test-user" \ | ||
| "http://localhost:$ROUTER_PORT/v1/chat/completions" | ||
|
|
||
| # Clean up payload file | ||
| rm -f "$payload_file" |
There was a problem hiding this comment.
The payload file is created at a hardcoded path /tmp/stress_payload.json. This can lead to race conditions and conflicts if multiple instances of this stress test are run concurrently. It's also a security best practice to use mktemp to create a temporary file with a unique name.
To implement this correctly and ensure cleanup, I recommend the following changes:
- Declare
PAYLOAD_FILE=""at the top of the script with other global variables. - In
run_stress_test, usePAYLOAD_FILE=$(mktemp /tmp/stress_payload.XXXXXX.json)and assign it to a global variable, and remove therm -fat the end of the function. - In the
cleanupfunction, addrm -f "$PAYLOAD_FILE"to ensure the temporary file is deleted on script exit.
| local payload_file="/tmp/stress_payload.json" | |
| cat > "$payload_file" << EOF | |
| { | |
| "model": "$MODEL", | |
| "messages": [ | |
| {"role": "user", "content": "Test message for stress testing"} | |
| ], | |
| "max_tokens": 10, | |
| "temperature": 0.7 | |
| } | |
| EOF | |
| # Run Apache Bench | |
| ab -c "$CONCURRENT" \ | |
| -n "$REQUESTS" \ | |
| -p "$payload_file" \ | |
| -T "application/json" \ | |
| -H "Authorization: Bearer test" \ | |
| -H "x-user-id: stress-test-user" \ | |
| "http://localhost:$ROUTER_PORT/v1/chat/completions" | |
| # Clean up payload file | |
| rm -f "$payload_file" | |
| PAYLOAD_FILE=$(mktemp /tmp/stress_payload.XXXXXX.json) | |
| cat > "$PAYLOAD_FILE" << EOF | |
| { | |
| "model": "$MODEL", | |
| "messages": [ | |
| {"role": "user", "content": "Test message for stress testing"} | |
| ], | |
| "max_tokens": 10, | |
| "temperature": 0.7 | |
| } | |
| EOF | |
| # Run Apache Bench | |
| ab -c "$CONCURRENT" \ | |
| -n "$REQUESTS" \ | |
| -p "$PAYLOAD_FILE" \ | |
| -T "application/json" \ | |
| -H "Authorization: Bearer test" \ | |
| -H "x-user-id: stress-test-user" \ | |
| "http://localhost:$ROUTER_PORT/v1/chat/completions" |
tests/e2e/stress-test.sh
Outdated
| local backend1_count=$(grep -c "to http://localhost:$BACKEND1_PORT" "$log_file" || echo "0") | ||
| local backend2_count=$(grep -c "to http://localhost:$BACKEND2_PORT" "$log_file" || echo "0") |
There was a problem hiding this comment.
The || echo "0" is not necessary here. grep -c will print 0 to standard output if no lines match, and the command substitution $(...) will capture this 0. While grep -c exits with status 1 in this case, set -e is not triggered for commands inside $(...). The || echo "0" construct can also mask real errors from grep (which would exit with 2).
| local backend1_count=$(grep -c "to http://localhost:$BACKEND1_PORT" "$log_file" || echo "0") | |
| local backend2_count=$(grep -c "to http://localhost:$BACKEND2_PORT" "$log_file" || echo "0") | |
| local backend1_count=$(grep -c "to http://localhost:$BACKEND1_PORT" "$log_file") | |
| local backend2_count=$(grep -c "to http://localhost:$BACKEND2_PORT" "$log_file") |
| } | ||
|
|
||
| # Parse command line arguments | ||
| while [[ $# -gt 0 ]]; do |
There was a problem hiding this comment.
The argument parsing loop is not robust against missing option arguments. For example, if the script is called with -c as the last argument, it will fail with an "unbound variable" error because of set -u. You should add checks to ensure that options requiring a value are provided with one.
Here's an example of how you could improve the handling for the --concurrent option. Similar checks should be added for all other options that take arguments.
-c|--concurrent)
if [[ -z "${2-}" ]]; then
print_error "Option '$1' requires an argument."
show_usage
exit 1
fi
CONCURRENT="$2"
shift 2
;;
Separate variable declarations from assignments to avoid masking return values in backend count extraction logic. Signed-off-by: YuhanLiu11 <[email protected]>
Stress testing of the router's round-robin routing logic under high concurrent load. Related PRs: #448 and #589
Not included in the regular CI as the stress testing takes too long to complete. For internal usage only between releases.
Example usage:
Example output:
BEFORE SUBMITTING, PLEASE READ THE CHECKLIST BELOW AND FILL IN THE DESCRIPTION ABOVE
-swhen doinggit commit[Bugfix],[Feat], and[CI].Detailed Checklist (Click to Expand)
Thank you for your contribution to production-stack! Before submitting the pull request, please ensure the PR meets the following criteria. This helps us maintain the code quality and improve the efficiency of the review process.
PR Title and Classification
Please try to classify PRs for easy understanding of the type of changes. The PR title is prefixed appropriately to indicate the type of change. Please use one of the following:
[Bugfix]for bug fixes.[CI/Build]for build or continuous integration improvements.[Doc]for documentation fixes and improvements.[Feat]for new features in the cluster (e.g., autoscaling, disaggregated prefill, etc.).[Router]for changes to thevllm_router(e.g., routing algorithm, router observability, etc.).[Misc]for PRs that do not fit the above categories. Please use this sparingly.Note: If the PR spans more than one category, please include all relevant prefixes.
Code Quality
The PR need to meet the following code quality standards:
pre-committo format your code. SeeREADME.mdfor installation.DCO and Signed-off-by
When contributing changes to this project, you must agree to the DCO. Commits must include a
Signed-off-by:header which certifies agreement with the terms of the DCO.Using
-swithgit commitwill automatically add this header.What to Expect for the Reviews
We aim to address all PRs in a timely manner. If no one reviews your PR within 5 days, please @-mention one of YuhanLiu11
, Shaoting-Feng or ApostaC.