feat(batch): add multi-endpoint body validation and testing#1982
feat(batch): add multi-endpoint body validation and testing#1982NJX-njx wants to merge 2 commits intovllm-project:mainfrom
Conversation
Expand batch API testing and validation to cover all supported endpoints: /v1/chat/completions, /v1/completions, /v1/embeddings, and /v1/rerank. Changes: - Add endpoint-specific body validation in _validate_batch_input_file: - chat/completions: requires model + messages (array) - completions: requires model + prompt (string or array) - embeddings: requires model + input (string or array) - rerank: requires model + query (string) + documents (array) - Add multi-endpoint input data generators with sample bodies for each type - Add parametrized e2e test (test_openai_batch_api_multi_endpoint) covering all 4 endpoints - Add TestEndpointBodyValidation test class with 20+ unit tests for body validation - Update e2e test generators to accept endpoint parameter Closes vllm-project#1653
|
You have reached your Codex usage limits for code reviews. You can see your limits in the Codex usage dashboard. |
Summary of ChangesHello, 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! This pull request significantly improves the robustness and reliability of the batch API by introducing comprehensive validation for request bodies across all supported AI endpoints. It ensures that incoming batch job requests adhere to the specific requirements of each endpoint, preventing malformed data from proceeding to job execution. Furthermore, the changes include extensive testing to verify the correct functioning of this new validation logic and the overall batch workflow for multi-endpoint scenarios. Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Changelog
Activity
Using Gemini Code AssistThe 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 by creating a comment using either
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 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. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. 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
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request significantly improves the batch processing system by introducing robust body validation for API endpoints using the new _validate_request_body_for_endpoint function, along with comprehensive unit and end-to-end tests. However, a critical security flaw allows attackers to bypass this validation by providing a full URL in the input file instead of just the path, potentially leading to malformed requests reaching the inference engine. To address this, it is crucial to use the validated endpoint path from the batch request for body validation.
| request["body"], request_url, line_num | ||
| ) |
There was a problem hiding this comment.
The _validate_batch_input_file function passes the untrusted request_url from the input file to _validate_request_body_for_endpoint instead of using the validated endpoint from the batch request. Since _validate_batch_input_file only checks if request_url ends with the expected endpoint (line 187), an attacker can provide a full URL (e.g., http://attacker.com/v1/chat/completions) that passes the check but fails to match any key in ENDPOINT_REQUIRED_BODY_FIELDS (which contains only paths like /v1/chat/completions). This causes _validate_request_body_for_endpoint to skip validation (line 74), allowing malformed request bodies to bypass the early validation check.
body_error = _validate_request_body_for_endpoint(
request["body"], endpoint, line_num
)There was a problem hiding this comment.
Pull request overview
Expands the batch API to validate request bodies per endpoint and broadens automated coverage to exercise all supported batch endpoints end-to-end.
Changes:
- Added endpoint-specific request body validation during batch input file validation.
- Added multi-endpoint batch input generators for e2e tests.
- Added parametrized e2e test and new unit tests covering endpoint-specific body validation.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 4 comments.
| File | Description |
|---|---|
| python/aibrix/aibrix/metadata/api/v1/batch.py | Adds endpoint-aware request body validation during batch input parsing. |
| python/aibrix/tests/batch/test_e2e_openai_batch_api.py | Adds multi-endpoint generators and a parametrized e2e workflow test across endpoints. |
| python/aibrix/tests/e2e/test_batch_api.py | Generalizes batch input generator to support multiple endpoints. |
| python/aibrix/tests/batch/test_batch_endpoints.py | Adds unit tests for endpoint-specific body validation logic. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| # Step 4: Download and verify output | ||
| output_response = client.get(f"/v1/files/{output_file_id}/content") | ||
| assert output_response.status_code == 200 | ||
|
|
||
| output_content = output_response.content.decode("utf-8") | ||
| assert output_content, f"[{endpoint}] Output file is empty" | ||
| assert verify_batch_output_content(output_content, num_requests), ( | ||
| f"[{endpoint}] Output verification failed" | ||
| ) |
There was a problem hiding this comment.
verify_batch_output_content() (per provided context) asserts response.body contains ["model", "choices"], which will not match typical responses for /v1/embeddings (usually data) or /v1/rerank (usually results). This makes the new parametrized test fail for non-chat endpoints. Update verification to be endpoint-aware (e.g., branch assertions by endpoint), or generalize the verifier to validate only common fields plus endpoint-specific payload shape.
| required_fields = ENDPOINT_REQUIRED_BODY_FIELDS.get(endpoint) | ||
| if required_fields is None: | ||
| # Unknown endpoint, skip body validation | ||
| return None |
There was a problem hiding this comment.
Body validation is skipped entirely when endpoint doesn’t exactly match a key in ENDPOINT_REQUIRED_BODY_FIELDS. Since the value passed in comes from the batch input line (request_url), any benign variation (e.g., trailing slash, query string) would bypass validation. Consider normalizing the endpoint before lookup (e.g., strip query/fragment, strip trailing /) and/or validating against the canonical batch job endpoint (the endpoint parameter passed to _validate_batch_input_file) rather than the per-line URL string.
| from aibrix.metadata.api.v1.batch import _validate_request_body_for_endpoint | ||
|
|
||
|
|
There was a problem hiding this comment.
Tests are importing and depending on a private function (_validate_request_body_for_endpoint). This increases coupling and makes refactors harder (renames/moves break tests even if behavior is unchanged). Consider promoting this validator to a non-underscored helper (or moving it into a small shared validation module) and importing that public surface from tests.
| from aibrix.metadata.api.v1.batch import _validate_request_body_for_endpoint | |
| import aibrix.metadata.api.v1.batch as batch_module | |
| _validate_request_body_for_endpoint = batch_module._validate_request_body_for_endpoint |
| return f"Line {line_num}: 'messages' must be an array for {endpoint}" | ||
| elif endpoint == BatchJobEndpoint.COMPLETIONS.value: | ||
| prompt = body.get("prompt") | ||
| if not isinstance(prompt, (str, list)): | ||
| return f"Line {line_num}: 'prompt' must be a string or array for {endpoint}" | ||
| elif endpoint == BatchJobEndpoint.EMBEDDINGS.value: | ||
| input_val = body.get("input") | ||
| if not isinstance(input_val, (str, list)): | ||
| return f"Line {line_num}: 'input' must be a string or array for {endpoint}" | ||
| elif endpoint == BatchJobEndpoint.RERANK.value: | ||
| if not isinstance(body.get("query"), str): | ||
| return f"Line {line_num}: 'query' must be a string for {endpoint}" | ||
| if not isinstance(body.get("documents"), list): | ||
| return f"Line {line_num}: 'documents' must be an array for {endpoint}" |
There was a problem hiding this comment.
The error messages use “array”, which is more JavaScript/JSON terminology; in Python-facing validation errors, “list” is usually clearer (especially since the check is isinstance(..., list)). Consider rewording to “list” (and similarly for other messages) to make failures more actionable/debuggable.
| return f"Line {line_num}: 'messages' must be an array for {endpoint}" | |
| elif endpoint == BatchJobEndpoint.COMPLETIONS.value: | |
| prompt = body.get("prompt") | |
| if not isinstance(prompt, (str, list)): | |
| return f"Line {line_num}: 'prompt' must be a string or array for {endpoint}" | |
| elif endpoint == BatchJobEndpoint.EMBEDDINGS.value: | |
| input_val = body.get("input") | |
| if not isinstance(input_val, (str, list)): | |
| return f"Line {line_num}: 'input' must be a string or array for {endpoint}" | |
| elif endpoint == BatchJobEndpoint.RERANK.value: | |
| if not isinstance(body.get("query"), str): | |
| return f"Line {line_num}: 'query' must be a string for {endpoint}" | |
| if not isinstance(body.get("documents"), list): | |
| return f"Line {line_num}: 'documents' must be an array for {endpoint}" | |
| return f"Line {line_num}: 'messages' must be a list for {endpoint}" | |
| elif endpoint == BatchJobEndpoint.COMPLETIONS.value: | |
| prompt = body.get("prompt") | |
| if not isinstance(prompt, (str, list)): | |
| return f"Line {line_num}: 'prompt' must be a string or list for {endpoint}" | |
| elif endpoint == BatchJobEndpoint.EMBEDDINGS.value: | |
| input_val = body.get("input") | |
| if not isinstance(input_val, (str, list)): | |
| return f"Line {line_num}: 'input' must be a string or list for {endpoint}" | |
| elif endpoint == BatchJobEndpoint.RERANK.value: | |
| if not isinstance(body.get("query"), str): | |
| return f"Line {line_num}: 'query' must be a string for {endpoint}" | |
| if not isinstance(body.get("documents"), list): | |
| return f"Line {line_num}: 'documents' must be a list for {endpoint}" |
|
Can you sign the commit and help address copilot comments. |
Summary
Expands batch API to fully validate and test all supported endpoints: /v1/chat/completions, /v1/completions, /v1/embeddings, and /v1/rerank.
Changes
1. Endpoint-specific body validation (\�atch.py)
Added _validate_request_body_for_endpoint()\ that checks request bodies have the required fields for each endpoint type:
This validation runs during _validate_batch_input_file()\ to catch malformed requests early before job execution.
2. Multi-endpoint input data generators
Both e2e test files (\ est_e2e_openai_batch_api.py\ and \ ests/e2e/test_batch_api.py) now have:
3. Parametrized e2e test
Added \ est_openai_batch_api_multi_endpoint\ that exercises the full batch workflow (upload create poll verify) for each of the 4 endpoint types.
4. Body validation unit tests
Added \TestEndpointBodyValidation\ class in \ est_batch_endpoints.py\ with 20+ test cases covering:
Closes #1653