Repository navigation
test(e2e): add basic string stop and matched stop test coverage for MLX - #1538
zach-li-sudo wants to merge 3 commits into
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: ASSERTIVE Plan: Pro Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughThe MLX backend E2E test suite is extended with stop-string validation tests and streaming helpers. A new ChangesMLX stop-sequence validation tests
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![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 introduces end-to-end tests for stop sequence functionality in the MLX backend, covering both chat and completion endpoints in streaming and non-streaming modes. It also includes tests to verify that multi-token stop strings are correctly rejected with a 400 error. Feedback was provided to improve the robustness of the collect_streamed_completion helper by handling cases where no chunk contains a finish reason, preventing uninformative StopIteration errors and improving test failure diagnostics.
| final_choice = next( | ||
| c.choices[0] for c in reversed(chunks) if c.choices and c.choices[0].finish_reason | ||
| ) |
There was a problem hiding this comment.
The next() function is called on a generator that may be empty if no chunk in the stream contains a finish_reason. This would raise a StopIteration exception, which is less informative than an assertion failure. It is better to provide a default value to next() and then assert that a valid choice was found to improve test failure diagnostics.
| final_choice = next( | |
| c.choices[0] for c in reversed(chunks) if c.choices and c.choices[0].finish_reason | |
| ) | |
| final_choice = next( | |
| (c.choices[0] for c in reversed(chunks) if c.choices and c.choices[0].finish_reason), | |
| None, | |
| ) | |
| assert final_choice is not None, "No chunk with finish_reason found in stream" |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@e2e_test/mlx/test_mlx_backend.py`:
- Around line 39-42: The test helper assert_matched_stop currently compares
choice.matched_stop (a string) to an expected string, but the MLX router wrapper
sets matched_stop_token_id as a JSON number; update assert_matched_stop to
normalize types before comparing: fetch getattr(choice, "matched_stop", None)
and also check for getattr(choice, "matched_stop_token_id", None), coerce
numeric token ids to strings (or coerce expected to int) and then assert
equality; update all call-sites that pass expected values (lines around 203-204,
216-217, 252) to use the same normalization approach so the assertion treats "6"
and 6 as equivalent.
- Around line 191-217: Tests assume string stops are accepted but MLX currently
rejects non-empty stop strings via reject_stop_strings(...), so update the
failing tests to match backend behavior: in test_chat_stop_string_non_streaming
and test_completion_stop_string_non_streaming (and the similar cases at the
later block), remove or do not pass stop=[self.SINGLE_STRING_STOP] (use no stop
parameter or an accepted stop form), and update assertions accordingly (do not
assert finish_reason == "stop", remove
assert_stop_text_trimmed/assert_matched_stop) or mark the tests as
expected-to-fail/skip until MLX accepts string stops; target the two test
functions by name to locate and modify them.
- Around line 25-32: The helper collect_streamed_completion can raise
StopIteration when no chunk has a finish_reason; update it to explicitly check
for a final chunk before using next() — e.g., search for
collect_streamed_completion, compute final_choice_candidate by scanning
reversed(chunks) for a chunk with c.choices and c.choices[0].finish_reason,
assert that such a chunk exists (or raise a clear ValueError/AssertionError with
a diagnostic message) and then set final_choice from that candidate; this makes
failures explicit and debuggable instead of letting next() raise StopIteration.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 24ba503a-0b9a-4d8d-911a-ca0337f5fd4c
📒 Files selected for processing (1)
e2e_test/mlx/test_mlx_backend.py
| def assert_matched_stop(choice, expected): | ||
| actual = getattr(choice, "matched_stop", None) | ||
| assert actual == expected, f"Expected matched_stop={expected!r}, got {actual!r}" | ||
|
|
There was a problem hiding this comment.
matched_stop type assertion is inconsistent with MLX response mapping.
assert_matched_stop currently expects a stop string (e.g., "6"), but MLX is mapped via matched_stop_token_id to a JSON number in the router wrapper. These assertions are likely to fail even if stop handling works.
Suggested adjustment
-def assert_matched_stop(choice, expected):
+def assert_matched_stop(choice, expected):
actual = getattr(choice, "matched_stop", None)
assert actual == expected, f"Expected matched_stop={expected!r}, got {actual!r}"- assert_matched_stop(choice, self.SINGLE_STRING_STOP)
+ # MLX currently surfaces token-id matched_stop
+ assert isinstance(getattr(choice, "matched_stop", None), int)Also applies to: 203-204, 216-217, 252-252
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@e2e_test/mlx/test_mlx_backend.py` around lines 39 - 42, The test helper
assert_matched_stop currently compares choice.matched_stop (a string) to an
expected string, but the MLX router wrapper sets matched_stop_token_id as a JSON
number; update assert_matched_stop to normalize types before comparing: fetch
getattr(choice, "matched_stop", None) and also check for getattr(choice,
"matched_stop_token_id", None), coerce numeric token ids to strings (or coerce
expected to int) and then assert equality; update all call-sites that pass
expected values (lines around 203-204, 216-217, 252) to use the same
normalization approach so the assertion treats "6" and 6 as equivalent.
| def test_chat_stop_string_non_streaming(self, model, api_client): | ||
| response = api_client.chat.completions.create( | ||
| model=model, | ||
| messages=[{"role": "user", "content": self.STOP_SEQUENCE_TEST_PROMPT}], | ||
| max_tokens=160, | ||
| temperature=0, | ||
| stop=[self.SINGLE_STRING_STOP], | ||
| extra_body=self.NO_THINKING, | ||
| ) | ||
| choice = response.choices[0] | ||
| assert choice.finish_reason == "stop" | ||
| assert_stop_text_trimmed(choice.message.content or "", self.SINGLE_STRING_STOP) | ||
| assert_matched_stop(choice, self.SINGLE_STRING_STOP) | ||
|
|
||
| def test_completion_stop_string_non_streaming(self, model, api_client): | ||
| response = api_client.completions.create( | ||
| model=model, | ||
| prompt=self.STOP_SEQUENCE_TEST_PROMPT, | ||
| max_tokens=160, | ||
| temperature=0, | ||
| stop=[self.SINGLE_STRING_STOP], | ||
| ) | ||
| choice = response.choices[0] | ||
| assert choice.finish_reason == "stop" | ||
| assert_stop_text_trimmed(choice.text, self.SINGLE_STRING_STOP) | ||
| assert_matched_stop(choice, self.SINGLE_STRING_STOP) | ||
|
|
There was a problem hiding this comment.
Single-stop “success” expectations conflict with current MLX contract.
These tests assert HTTP 200 and finish_reason=="stop" for stop=["6"], but the current MLX gRPC request builder rejects any non-empty stop strings (reject_stop_strings(...)). This will make these cases fail until backend behavior is changed.
Suggested adjustment
- def test_chat_stop_string_non_streaming(self, model, api_client):
- response = api_client.chat.completions.create(
- model=model,
- messages=[{"role": "user", "content": self.STOP_SEQUENCE_TEST_PROMPT}],
- max_tokens=160,
- temperature=0,
- stop=[self.SINGLE_STRING_STOP],
- extra_body=self.NO_THINKING,
- )
- choice = response.choices[0]
- assert choice.finish_reason == "stop"
- assert_stop_text_trimmed(choice.message.content or "", self.SINGLE_STRING_STOP)
- assert_matched_stop(choice, self.SINGLE_STRING_STOP)
+ # Keep/enable success-path tests only once MLX stop-string support is implemented.
+ # For current behavior, assert 400 unsupported_stop_string.Also applies to: 239-252
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@e2e_test/mlx/test_mlx_backend.py` around lines 191 - 217, Tests assume string
stops are accepted but MLX currently rejects non-empty stop strings via
reject_stop_strings(...), so update the failing tests to match backend behavior:
in test_chat_stop_string_non_streaming and
test_completion_stop_string_non_streaming (and the similar cases at the later
block), remove or do not pass stop=[self.SINGLE_STRING_STOP] (use no stop
parameter or an accepted stop form), and update assertions accordingly (do not
assert finish_reason == "stop", remove
assert_stop_text_trimmed/assert_matched_stop) or mark the tests as
expected-to-fail/skip until MLX accepts string stops; target the two test
functions by name to locate and modify them.
Signed-off-by: Zhuo Li <zhuo.li.ca@outlook.com>
Signed-off-by: Zhuo Li <zhuo.li.ca@outlook.com>
Signed-off-by: Zhuo Li <zhuo.li.ca@outlook.com>
0077d46 to
0f69094
Compare
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 0f69094e22
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| text, final_choice = collect_streamed_completion(stream) | ||
| assert final_choice.finish_reason == "stop" | ||
| assert_stop_text_trimmed(text, self.SINGLE_STRING_STOP) | ||
| assert_matched_stop(final_choice, self.SINGLE_STRING_STOP) |
There was a problem hiding this comment.
When this new MLX streaming completion test runs, final_choice comes from the OpenAI completions streaming path, whose CompletionStreamChoice schema only serializes text, index, logprobs, and finish_reason; the final chunks built in model_gateway/src/routers/grpc/regular/streaming.rs likewise never include matched_stop. As a result getattr(final_choice, "matched_stop", None) is always None, so this test fails even when the backend correctly stops on the requested string.
Useful? React with 👍 / 👎.
|
This pull request has been automatically marked as stale because it has not had any activity within 14 days. It will be automatically closed if no further activity occurs within 16 days. Leave a comment if you feel this pull request should remain open. Thank you! |
|
This pull request has been automatically closed due to inactivity. Please feel free to reopen if you intend to continue working on it. Thank you! |
1 participant
Description
Problem
The MLX gRPC backend added support for string stop sequences and
matched_stopreporting, but had no end-to-end test coverage validating this through the full router → gRPC → MLX worker path.String stop support PR: #1524
Solution
Add 5 e2e tests to
e2e_test/mlx/test_mlx_backend.pycovering the stop sequence feature for the regular (non-Harmony) pipeline:finish_reason == "stop", stop text is excludedfrom the output, and
matched_stopechoes back the stop string.unsupported_stop_stringerror code.finish_reason == "stop"and the correctmatched_stopvalue.Helper functions
assert_stop_text_trimmed,assert_matched_stop,assert_api_error, andcollect_streamed_completionare extracted to keep thetest bodies concise.
Changes
e2e_test/mlx/test_mlx_backend.py: add 5 stop-sequence/matched-stop tests and supporting helper functions; addSTOP_SEQUENCE_TEST_PROMPTandSINGLE_STRING_STOPclass constants toTestMlxBackendTest Plan
Run on Apple Silicon (macOS arm64) with
mlx-community/Qwen3-0.6B-4bit:also run python format checks:
Summary by CodeRabbit