feat(mlx): add Python gRPC servicer for MLX backend - #1099
Conversation
Implements MlxEngineServicer (8 RPCs) backed by mlx-lm's BatchGenerator for Apple Silicon inference. Wraps the BatchGenerator behind the MlxEngine proto from #1034. Components: - MlxEngineServicer: Generate (streaming + non-streaming), HealthCheck, Abort, GetModelInfo, GetServerInfo, GetTokenizer; Embed and SubscribeKvEvents return UNIMPLEMENTED - MlxHealthServicer: standard grpc.health.v1.Health protocol - server.py: standalone gRPC server entrypoint with model loading - 42 unit + integration tests covering all RPCs Generation loop: - Background daemon thread runs BatchGenerator.next() in tight loop - Dispatches tokens to per-request asyncio.Queues via loop.call_soon_threadsafe - Auto-removes completed/aborted requests from BatchGenerator Sampling param mapping: - temperature, top_p, top_k, min_p -> make_sampler() - repetition/frequency/presence_penalty, logit_bias -> make_logits_processors() - max_tokens, stop_token_ids (via SequenceStateMachine), ignore_eos - logprobs (output only, top-k extraction from generate_step logprobs) - seed (best-effort, mx.random.seed global state) Also exports mlx_engine_pb2/mlx_engine_pb2_grpc from smg-grpc-proto package and adds [mlx] optional dependency to grpc_servicer. Signed-off-by: key4ng <rukeyang@gmail.com>
vllm and sglang servicers have no unit tests in this repo — testing is done via E2E pytest in e2e_test/. Remove MLX unit tests for consistency. Signed-off-by: key4ng <rukeyang@gmail.com>
Adds Apple-Silicon-only E2E test using mlx-community/Qwen3-0.6B-4bit (~400 MB) which fits comfortably on a macos-latest GitHub Actions runner. Qwen3 supports both native tool calling and thinking mode in a single small model. Test coverage: - Basic non-streaming chat completion - Streaming chat with delta chunks - Tool calling via SMG's qwen ToolParser parsing <tool_call> tags - Reasoning content via thinking mode - max_tokens finish_reason Infrastructure additions: - Runtime.MLX enum + is_mlx() helper - _build_mlx_cmd() in worker.py (gRPC-only, no HTTP variant) - mlx-community/Qwen3-0.6B-4bit entry in MODEL_SPECS Tests skip on non-Darwin/non-arm64 platforms via module-level pytestmark. Verified locally: 15/15 passes across 3 runs. Signed-off-by: key4ng <rukeyang@gmail.com>
Adds .github/workflows/pr-test-mlx.yml that runs the MLX E2E test on macos-14 (Apple Silicon) GitHub-hosted runners. Triggered on changes to MLX-related paths (proto, Rust client, Python servicer, MLX test file, or the workflow itself). Steps: 1. Install protoc, Rust toolchain, Python 3.12 2. Cache Cargo + HuggingFace model 3. Build smg-grpc-proto Python package (proto codegen) 4. Install grpc_servicer[mlx] (pulls mlx + mlx-lm wheels) 5. Build smg Python bindings via maturin develop --release 6. Generate Python client types (cargo + datamodel-code-generator) 7. Install smg-client + e2e_test (deps) 8. Verify imports 9. Run pytest with E2E_RUNTIME=mlx, --reruns 1 for flake tolerance 10. Upload logs on failure The test downloads mlx-community/Qwen3-0.6B-4bit (~400 MB) on first run and caches it for subsequent runs. Signed-off-by: key4ng <rukeyang@gmail.com>
--release build takes 6-8 min on macos-14 3-core runners. Since this is a correctness test (not a perf test), debug is sufficient and drops first-run time to ~7 min, cached runs to ~3 min. Signed-off-by: key4ng <rukeyang@gmail.com>
github-actions
Bot
added
dependencies
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds MLX backend support across CI, proto re-exports, e2e infra (runtime/model/worker), a new MLX gRPC servicer (health, server, engine), packaging extras, and macOS arm64 E2E tests. Changes
Sequence DiagramsequenceDiagram
autonumber
participant Client as Client / Test
participant Router as SMG Router
participant Servicer as MlxEngineServicer
participant BatchGen as BatchGenerator
participant MLX as MLX Model
participant Health as MlxHealthServicer
Client->>Router: Send request
Router->>Servicer: gRPC Generate(request)
Servicer->>Servicer: Validate request & build sampler/logits
Servicer->>BatchGen: Insert request (uid)
Servicer->>Health: Ensure serving state
loop generation batches
BatchGen->>MLX: next() -> batch
MLX->>BatchGen: Token results / finish state
BatchGen->>Servicer: Route per-uid results
Servicer->>Client: Stream GenerateStreamChunk(s) / final GenerateComplete
end
Servicer->>BatchGen: Cleanup finished uid
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
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 a new MLX inference backend for Apple Silicon, featuring a gRPC servicer, a standalone server entry point, and comprehensive E2E tests using the Qwen3 model. The implementation leverages the mlx-lm library for batch generation. The review feedback primarily addresses critical thread-safety concerns, as the BatchGenerator is accessed concurrently by the background generation thread and gRPC tasks without synchronization. Other recommendations include fixing hardcoded token usage counts in streaming responses, using more robust event loop retrieval methods, and following standard Python practices for import placement.
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Good work — clean implementation of the MLX gRPC servicer. The architecture (event loop for insert, background daemon thread for BatchGenerator.next(), call_soon_threadsafe dispatch) is sound and matches the design from PR #1034.
Review summary:
- 🔴 Important: 0
- 🟡 Nit: 3
- 🟣 Pre-existing: 0
The nits flagged are: (1) batch_generator.remove() called from both the event-loop thread (Abort) and the generation thread — potential thread-safety concern, (2) _build_output_logprobs with num_logprobs=0 returns entire vocabulary due to [-0:] semantics, (3) asyncio.get_event_loop() deprecated in favor of get_running_loop().
![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: 6ee5cce360
ℹ️ About Codex in GitHub
Your team has set up Codex to 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
…/format - Use 'maturin build --out dist && pip install dist/*.whl' instead of 'maturin develop' — GitHub-hosted runner Python is not in a venv, so 'maturin develop' errors out with 'Couldn't find a virtualenv'. - Fix ruff E731: replace lambda with def _stream_ctx in generation loop - Apply ruff format to all mlx python files Signed-off-by: key4ng <rukeyang@gmail.com>
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@grpc_servicer/smg_grpc_servicer/mlx/server.py`:
- Around line 116-118: The code currently ignores the return value of
server.add_insecure_port(listen_addr) so the process can continue even if the
socket failed to bind; update the startup in server.py to capture the return
(e.g., bound_port = server.add_insecure_port(listen_addr)) and if it is 0, log
an error and fail fast (raise an exception or call sys.exit(1)) before calling
await server.start(), ensuring you reference listen_addr,
server.add_insecure_port, and await server.start() in your change.
In `@grpc_servicer/smg_grpc_servicer/mlx/servicer.py`:
- Around line 89-104: In _build_output_logprobs: guard the case num_logprobs ==
0 (or top_k <= 0) before computing top_k and slicing with [-top_k:] to avoid
selecting the whole vocabulary; if num_logprobs is 0 return None (or an empty
OutputLogProbs) early, otherwise proceed to compute top_k, top_indices =
mx.argpartition(...)[-top_k:], top_values, sort_order, etc., so logprobs_array
and the top-k logic are only run when top_k > 0.
- Around line 374-383: Abort currently removes uid mappings but never signals
the waiting Generate coroutine, causing it to hang on queue.get(); modify Abort
(mlx_grpc_servicer.ml x.servicer Abort) to look up the uid via
self._request_uid_map and, when removing self._uid_queues.pop(uid, None), put a
terminal sentinel (e.g., None) into the corresponding queue (or call
put_nowait(None)) so any blocked Generate waiter wakes; also update Generate to
check the dequeued value (r = await queue.get()) for the sentinel (if r is None:
break/return) before accessing r.token so the stream can finish and
_active_requests can decrement properly.
- Around line 194-210: The GetModelInfoResponse currently recomputes
eos_token_ids from config and returns [] when config lacks eos_token_id; instead
reuse the resolved EOS IDs that the server already computes (the
tokenizer-derived fallback implemented in server.py lines ~56-64). Update the
code that builds GetModelInfoResponse to use the pre-resolved EOS ID container
(e.g., the attribute or variable the server module exposes such as the
resolved_eos_ids / self.eos_token_ids or the same helper used in server.py)
rather than recomputing from config, so the response matches the actual
generation stop behavior.
🪄 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: 13fbb568-26da-40b8-b8a6-fe0814769626
📒 Files selected for processing (13)
.github/workflows/pr-test-mlx.ymlcrates/grpc_client/python/smg_grpc_proto/__init__.pye2e_test/chat_completions/test_mlx_backend.pye2e_test/infra/__init__.pye2e_test/infra/constants.pye2e_test/infra/model_specs.pye2e_test/infra/worker.pygrpc_servicer/pyproject.tomlgrpc_servicer/smg_grpc_servicer/mlx/__init__.pygrpc_servicer/smg_grpc_servicer/mlx/__main__.pygrpc_servicer/smg_grpc_servicer/mlx/health_servicer.pygrpc_servicer/smg_grpc_servicer/mlx/server.pygrpc_servicer/smg_grpc_servicer/mlx/servicer.py
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 17389ce93a
ℹ️ About Codex in GitHub
Your team has set up Codex to 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
There was a problem hiding this comment.
Actionable comments posted: 2
♻️ Duplicate comments (4)
grpc_servicer/smg_grpc_servicer/mlx/servicer.py (3)
201-228:⚠️ Potential issue | 🟡 MinorReuse the resolved EOS IDs for consistency.
GetModelInforecomputeseos_token_idsfrom config only (lines 208-214), returning[]when config lackseos_token_id. However,_build_state_machineusesself._eos_token_idswhich includes the tokenizer fallback computed inserver.py. This inconsistency meansGetModelInfocan report different EOS IDs than what generation actually uses.🔧 Proposed fix
async def GetModelInfo( self, request: mlx_engine_pb2.GetModelInfoRequest, context: grpc.aio.ServicerContext, ) -> mlx_engine_pb2.GetModelInfoResponse: config = self.model_config - - eos = config.get("eos_token_id") - if isinstance(eos, int): - eos_token_ids = [eos] - elif isinstance(eos, list): - eos_token_ids = eos - else: - eos_token_ids = [] return mlx_engine_pb2.GetModelInfoResponse( model_path=self.model_path, is_generation=True, max_context_length=config.get("max_position_embeddings", 0), vocab_size=config.get("vocab_size", 0), served_model_name=self.model_path, model_type=config.get("model_type", ""), architectures=config.get("architectures", []), - eos_token_ids=eos_token_ids, + eos_token_ids=self._eos_token_ids, pad_token_id=config.get("pad_token_id") or 0, bos_token_id=config.get("bos_token_id") or 0, max_req_input_len=config.get("max_position_embeddings", 0), )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@grpc_servicer/smg_grpc_servicer/mlx/servicer.py` around lines 201 - 228, GetModelInfo currently recomputes eos_token_ids from model_config, causing mismatch with the tokenizer-resolved EOS IDs used by _build_state_machine; update GetModelInfo to use the already-resolved self._eos_token_ids (or fall back to converting config values only if self._eos_token_ids is empty) and ensure the value passed (eos_token_ids) is a list of ints to preserve consistency between reported EOS IDs and those used by generation.
401-410:⚠️ Potential issue | 🔴 CriticalAbort must wake the blocked
Generatewaiter.
Generate()is blocked onqueue.get(). WhenAbortpops the queue from_uid_queueswithout putting a terminal sentinel, theGeneratecoroutine hangs indefinitely,_active_requestsnever decrements, and resources leak.🔧 Proposed fix
async def Abort(self, request, context): for request_id in request.request_ids: uid = self._request_uid_map.pop(request_id, None) if uid is not None: - self._uid_queues.pop(uid, None) + queue = self._uid_queues.pop(uid, None) + if queue is not None and self._loop is not None: + self._loop.call_soon_threadsafe(queue.put_nowait, None) try: self.batch_generator.remove([uid]) except Exception: logger.warning("Failed to remove uid %d for request %s", uid, request_id) return mlx_engine_pb2.AbortResponse()Then handle the sentinel in
Generate():if request.stream: while True: r = await queue.get() + if r is None: + break output_logprobs = self._build_output_logprobs((And similarly in the non-streaming branch.)
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@grpc_servicer/smg_grpc_servicer/mlx/servicer.py` around lines 401 - 410, Abort currently pops entries from _uid_queues (via Abort) which leaves any awaiting Generate coroutines blocked on queue.get(); modify Abort (function Abort) to put a terminal sentinel into the corresponding per-uid queue after removing the uid (or instead of removing) so any blocked Generate waiter is unblocked, and ensure you still call batch_generator.remove([uid]) and pop from _request_uid_map; then update Generate (function Generate) to detect the sentinel from _uid_queues[uid] (and the non-streaming branch counterpart), treat it as a termination signal by decrementing _active_requests and returning/cleaning up instead of waiting indefinitely.
90-116:⚠️ Potential issue | 🟠 MajorGuard
num_logprobs <= 0before computing top-k slice.When
num_logprobsis0,top_kbecomes0and[-0:]selects the entire vocabulary, causing unnecessary computation and payload bloat.🔧 Proposed fix
`@staticmethod` def _build_output_logprobs(token_id, logprobs_array, num_logprobs): """Build OutputLogProbs proto from an mlx logprobs array.""" - if num_logprobs is None: + if num_logprobs is None or num_logprobs <= 0: return None🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@grpc_servicer/smg_grpc_servicer/mlx/servicer.py` around lines 90 - 116, The _build_output_logprobs function should guard against non-positive num_logprobs to avoid slicing with [-0:] which returns the entire vocab; modify _build_output_logprobs to check if num_logprobs is None or num_logprobs <= 0 and return None immediately, otherwise continue computing top_k and the top_indices/top_values; update references inside _build_output_logprobs (token_id, logprobs_array, num_logprobs, top_k, top_indices, top_values, top_logprobs) accordingly so no further processing happens when num_logprobs <= 0.grpc_servicer/smg_grpc_servicer/mlx/server.py (1)
131-134:⚠️ Potential issue | 🟠 MajorFail fast if the listen socket cannot be bound.
server.add_insecure_port(...)returns0when the address cannot be bound. Without checking this return value, the server proceeds into startup without actually listening.🔧 Proposed fix
listen_addr = f"{args.host}:{args.port}" - server.add_insecure_port(listen_addr) + bound_port = server.add_insecure_port(listen_addr) + if bound_port == 0: + raise RuntimeError(f"Failed to bind gRPC server to {listen_addr}") await server.start()🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@grpc_servicer/smg_grpc_servicer/mlx/server.py` around lines 131 - 134, The server currently ignores the return value of server.add_insecure_port(listen_addr) which yields 0 on bind failure; update the startup path in the block using listen_addr, server.add_insecure_port, server.start and logger.info to check the returned port value, and if it's 0 log an error via logger.error (including listen_addr), and fail fast (raise an exception or exit) instead of calling await server.start() and logger.info.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@e2e_test/chat_completions/test_mlx_backend.py`:
- Around line 104-128: The json import is inside test_tool_calling; move the
import statement (import json) to the module-level imports at the top of the
file and update the test to call json.loads(...) instead of _json.loads(...);
update references in test_tool_calling where call.function.arguments is parsed
to use json.loads(call.function.arguments) so the behavior and names match the
module-level import.
In `@grpc_servicer/smg_grpc_servicer/mlx/servicer.py`:
- Around line 241-253: start_generation_loop currently calls
asyncio.get_event_loop() which can be ambiguous when invoked from an async
context (serve_grpc); change this to use asyncio.get_running_loop() to retrieve
the active loop and assign it to self._loop, ensuring start_generation_loop and
the thread target _generation_loop use that running loop; update any related
tests or callers if they relied on get_event_loop behavior and consider a
fallback to asyncio.get_event_loop() only if get_running_loop() raises
RuntimeError in non-async call sites.
---
Duplicate comments:
In `@grpc_servicer/smg_grpc_servicer/mlx/server.py`:
- Around line 131-134: The server currently ignores the return value of
server.add_insecure_port(listen_addr) which yields 0 on bind failure; update the
startup path in the block using listen_addr, server.add_insecure_port,
server.start and logger.info to check the returned port value, and if it's 0 log
an error via logger.error (including listen_addr), and fail fast (raise an
exception or exit) instead of calling await server.start() and logger.info.
In `@grpc_servicer/smg_grpc_servicer/mlx/servicer.py`:
- Around line 201-228: GetModelInfo currently recomputes eos_token_ids from
model_config, causing mismatch with the tokenizer-resolved EOS IDs used by
_build_state_machine; update GetModelInfo to use the already-resolved
self._eos_token_ids (or fall back to converting config values only if
self._eos_token_ids is empty) and ensure the value passed (eos_token_ids) is a
list of ints to preserve consistency between reported EOS IDs and those used by
generation.
- Around line 401-410: Abort currently pops entries from _uid_queues (via Abort)
which leaves any awaiting Generate coroutines blocked on queue.get(); modify
Abort (function Abort) to put a terminal sentinel into the corresponding per-uid
queue after removing the uid (or instead of removing) so any blocked Generate
waiter is unblocked, and ensure you still call batch_generator.remove([uid]) and
pop from _request_uid_map; then update Generate (function Generate) to detect
the sentinel from _uid_queues[uid] (and the non-streaming branch counterpart),
treat it as a termination signal by decrementing _active_requests and
returning/cleaning up instead of waiting indefinitely.
- Around line 90-116: The _build_output_logprobs function should guard against
non-positive num_logprobs to avoid slicing with [-0:] which returns the entire
vocab; modify _build_output_logprobs to check if num_logprobs is None or
num_logprobs <= 0 and return None immediately, otherwise continue computing
top_k and the top_indices/top_values; update references inside
_build_output_logprobs (token_id, logprobs_array, num_logprobs, top_k,
top_indices, top_values, top_logprobs) accordingly so no further processing
happens when num_logprobs <= 0.
🪄 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: 0689f1d8-d194-48e5-8540-3b4d09414aa8
📒 Files selected for processing (5)
.github/workflows/pr-test-mlx.ymle2e_test/chat_completions/test_mlx_backend.pygrpc_servicer/smg_grpc_servicer/mlx/health_servicer.pygrpc_servicer/smg_grpc_servicer/mlx/server.pygrpc_servicer/smg_grpc_servicer/mlx/servicer.py
Previous build used maturin's default (release) despite comment saying debug. Switch to the ci profile already defined in bindings/python/ Cargo.toml (opt-level=2, thin LTO, 16 codegen-units) for faster compilation on Apple Silicon runners. Signed-off-by: key4ng <rukeyang@gmail.com>
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In @.github/workflows/pr-test-mlx.yml:
- Line 43: The workflow job currently uses a floating runner image "runs-on:
macos-latest" which can drift; change that to an explicit Apple Silicon image
(e.g., "runs-on: macos-14") so the MLX/Apple-Silicon-specific job always runs on
an arm64 macOS host. Update the runs-on value in the job definition to the
pinned macOS 14 image and verify any matrix or conditionals that expect arm64
continue to work.
🪄 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: cd3db0f9-bee6-4296-8e35-de324e949e08
📒 Files selected for processing (1)
.github/workflows/pr-test-mlx.yml
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: f202dc2849
ℹ️ About Codex in GitHub
Your team has set up Codex to 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
- Default temperature to 1.0 when unset, matching vLLM/SGLang/TRT-LLM behavior (was forcing greedy decoding via 0.0). - Guard logprobs=0 in _build_output_logprobs; top_k=0 previously caused [-0:] to slice the full vocabulary. - Switch asyncio.get_event_loop() to get_running_loop() (deprecated in 3.10+). - Abort now puts a sentinel None into the request queue so a blocked Generate waiter unblocks instead of hanging until transport cancel. - Streaming Generate emits a running completion_tokens counter on both chunks and the final GenerateComplete (was hardcoded to 1 / 0). - Non-streaming Generate populates output_logprobs on GenerateComplete when sampling_params.logprobs is set. - GetModelInfo now returns the resolved self._eos_token_ids (which falls back to tokenizer-derived IDs in server.py) instead of recomputing from config.json only. - Add TODO(mlx-threadsafety) on the servicer class documenting the known BatchGenerator insert/next/remove multi-thread concern and deferring the fix (lock-around-next would kill throughput; needs a single-threaded request-queue refactor). Signed-off-by: key4ng <rukeyang@gmail.com>
|
Thanks for the comprehensive update, @key4ng. The consolidation of fixes in |
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@grpc_servicer/smg_grpc_servicer/mlx/servicer.py`:
- Around line 389-414: The code currently overwrites last_logprobs each
iteration, so when yielding the non-streaming response you must accumulate
per-token logprobs into a sequence-level payload instead of passing only the
final token's logprobs; change the loop that uses last_logprobs and
self._build_output_logprobs to collect each token's OutputLogProbs (e.g., create
output_logprobs_list = [] before the while, on each r call
self._build_output_logprobs(r.token, r.logprobs, num_logprobs) and append the
result to output_logprobs_list while handling possible None values), and then
pass the combined sequence-level logprobs (output_logprobs_list or an aggregated
structure expected by _complete_response) as output_logprobs to
self._complete_response along with all_output_ids, finish_reason, etc.; keep
variable names last_logprobs, _build_output_logprobs, _complete_response,
all_output_ids, and queue to locate the change.
- Around line 327-339: The background generator can observe a new uid from
batch_generator.insert before the per-uid asyncio.Queue is registered, causing
race drops; fix by serializing insert/next/remove with a shared lock (e.g.,
self._batch_lock) so uid visibility is atomic with queue publication: acquire
self._batch_lock before calling self.batch_generator.insert(...) in the handler,
create and assign self._uid_queues[uid] and self._request_uid_map[request_id],
then release the lock; likewise, acquire the same self._batch_lock around calls
to self.batch_generator.next(...) and any remove() logic inside _generation_loop
to ensure _generation_loop cannot observe a uid until its queue is registered.
Include the new lock initialization (e.g., in __init__) and use the same lock in
insert/next/remove paths.
🪄 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: 5434ae22-7db1-4fd0-85f7-39606d6c8375
📒 Files selected for processing (1)
grpc_servicer/smg_grpc_servicer/mlx/servicer.py
- server.py: check server.add_insecure_port return value (0 means bind failed) and raise instead of starting a server that never listened. - servicer.py: in non-streaming Generate, aggregate per-token logprobs across the whole sequence into the final GenerateComplete instead of keeping only the last step's logprobs. Signed-off-by: key4ng <rukeyang@gmail.com>
- GetModelInfo now reports self._ctx_limit (resolved via the 4-key fallback in __init__) instead of reading only max_position_embeddings. Previously reported 0 for model families that advertise context via max_seq_len / n_positions / seq_length. - snapshot_download allow_patterns now includes tiktoken.model and *.tiktoken so tiktoken-only models can actually ship a usable tokenizer bundle via GetTokenizer (the earlier servicer whitelist fix only covered files that happened to be on disk). Signed-off-by: key4ng <rukeyang@gmail.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 95866f6c45
ℹ️ About Codex in GitHub
Your team has set up Codex to 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
HF transformers>=4.43 ships chat_template as a standalone chat_template.json / chat_template.jinja file for newer models (Gemma 4, Llama 3.1+) instead of embedding it in tokenizer_config.json. Our GetTokenizer whitelist dropped these, so the router-side discover_chat_template_in_dir path found nothing and ChatTemplateState::apply failed at chat-request time for affected models. Fix in two places (same invariant as the tiktoken patch): - _TOKENIZER_FILES in servicer.py includes both filenames. - snapshot_download allow_patterns in server.py includes both so HF-repo model paths actually download them. Qwen3 (our CI model) embeds chat_template in tokenizer_config.json and would not have caught this; Gemma 4 is the canonical affected model. Signed-off-by: key4ng <rukeyang@gmail.com>
Move the following out of PR #1099 to keep the review diff focused on the servicer itself: - e2e_test/chat_completions/test_mlx_backend.py (MLX E2E test suite) - .github/workflows/pr-test-mlx.yml (CI workflow on macos-latest) - MLX-specific additions to e2e_test/infra/{__init__,constants,model_specs,worker}.py These will land in a follow-up PR from keyang/mlx-e2e-tests after this servicer PR merges. The servicer code itself remains fully locally testable via 'python -m smg_grpc_servicer.mlx.server --model ...'. Signed-off-by: key4ng <rukeyang@gmail.com>
CI and e2e_test/infra/worker.py both invoke the servicer as 'python -m smg_grpc_servicer.mlx.server'. Nothing uses the shorter 'python -m smg_grpc_servicer.mlx' form that __main__.py enables. Matches vllm/sglang conventions — neither ships a package-level __main__.py either. Signed-off-by: key4ng <rukeyang@gmail.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 000f74cc32
ℹ️ About Codex in GitHub
Your team has set up Codex to 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| bound_port = server.add_insecure_port(listen_addr) | ||
| if bound_port == 0: | ||
| raise RuntimeError(f"Failed to bind gRPC server to {listen_addr}") | ||
| await server.start() |
There was a problem hiding this comment.
serve_grpc begins serving at await server.start() before _warmup(...) and servicer.start_generation_loop(), so Generate can accept requests in that startup window and then block forever on await queue.get() because no background generation thread is dispatching responses yet. This is practically reachable because the MLX client’s readiness path uses the service HealthCheck RPC (which always reports healthy), so traffic can be routed before the loop is initialized.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 3b47bf5. Reordered serve_grpc so _warmup() + servicer.start_generation_loop() run before await server.start(). The window where the servicer accepts RPCs but has no gen thread dispatching is eliminated. Good catch — HealthCheck always returning OK meant the router couldn't detect this state on its own.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: a05ac9cca5
ℹ️ About Codex in GitHub
Your team has set up Codex to 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| self._request_uid_map.pop(request_id, None) | ||
| self._uid_queues.pop(uid, None) |
There was a problem hiding this comment.
When Generate exits through cancellation (for example client disconnects or deadline expiry), the finally block immediately removes request_id -> uid and queue entries but never removes the UID from BatchGenerator. If the follow-up Abort RPC arrives after this cleanup, Abort cannot find the UID and skips batch_generator.remove, so the cancelled request keeps generating tokens until its own stop/max-tokens condition. This wastes decode capacity and can throttle other requests under load.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 3b47bf5. Generate's finally block now calls batch_generator.remove([uid]) under _gen_lock, with a swallow-on-exception guard since the gen thread's finish-path or Abort may have already removed it. Client disconnects / deadline cancellations now reclaim the batch slot immediately instead of letting the decode run to natural completion.
Four lazy imports that were inside function bodies are now module-level: - server.py: mlx_lm.load, huggingface_hub.snapshot_download, mlx_lm.generate.BatchGenerator - servicer.py: smg_grpc_proto.generated.common_pb2 All of these are hard dependencies of the module path (mlx_lm is installed via the [mlx] extra; common_pb2 is in smg_grpc_proto which is already imported at top-level), so the lazy form was redundant protection. Matches the convention already used for mlx.core and mlx_lm.generate.generation_stream. Signed-off-by: key4ng <rukeyang@gmail.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 2d3638ed32
ℹ️ About Codex in GitHub
Your team has set up Codex to 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| [project.optional-dependencies] | ||
| vllm = ["vllm>=0.19.0"] | ||
| sglang = ["sglang>=0.5.10"] | ||
| mlx = ["mlx>=0.22.0", "mlx-lm>=0.22.0"] |
There was a problem hiding this comment.
Adding the mlx extra here exposes the MLX servicer path, but smg-grpc-servicer still allows any smg-grpc-proto>=0.4.6; that range does not guarantee mlx_engine_pb2 is present, while smg_grpc_servicer.mlx.server now imports it unconditionally. In environments that resolve to a pre-MLX proto build that still satisfies this spec, the MLX server fails at import/startup time. Please raise the minimum smg-grpc-proto version to the first release that includes MLX stubs.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 3b47bf5. Bumped crates/grpc_client/python/pyproject.toml to 0.4.7 (first version that ships mlx_engine_pb2 — the main-branch 0.4.6 does not) and pinned [mlx] extra in grpc_servicer/pyproject.toml to require smg-grpc-proto>=0.4.7. Installing smg-grpc-servicer[mlx] against an older proto build now fails at dependency resolution instead of at runtime import.
1. server.py: reorder startup. Warmup + start_generation_loop now run BEFORE await server.start(), so an incoming Generate RPC can never land in the window where the servicer is listening but no gen thread is dispatching (which would hang the RPC forever since HealthCheck always returns OK and can't signal this state). 2. servicer.py Generate finally: call batch_generator.remove(uid) on any Generate exit (client disconnect, deadline, CancelledError). Previously only Abort and natural finish removed from the batch, so a cancelled Generate would keep decoding until its own stop/max condition, wasting batch slots. 3. pyproject.toml: bump smg-grpc-proto to 0.4.7 (first version that ships mlx_engine_pb2) and pin [mlx] extra to require >=0.4.7 so installing against an older proto build can't crash at import. 4. servicer.py Abort: drain already-buffered tokens from the per-request queue before putting the None sentinel. Without this, a slow consumer with a backlog would keep receiving stale tokens after abort. Signed-off-by: key4ng <rukeyang@gmail.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 3b47bf5558
ℹ️ About Codex in GitHub
Your team has set up Codex to 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
HealthCheck previously hardcoded healthy=True, so the router kept routing traffic to an MLX worker even if: - the generation thread crashed (uncaught exception in _gen_loop) - the servicer is mid-shutdown - start_generation_loop() hasn't run yet (startup window) Now derive healthy from three cheap state checks: - _shutdown_event.is_set() → shutting down - _gen_thread is None → not initialized - not _gen_thread.is_alive() → thread died Message field carries the specific reason for operator debugging. Signed-off-by: key4ng <rukeyang@gmail.com>
|
e2e test will be in a seperate pr |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 5c24ce2b92
ℹ️ About Codex in GitHub
Your team has set up Codex to 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| # smg-grpc-proto>=0.4.7 is the first release that ships mlx_engine_pb2; | ||
| # without this floor, installing [mlx] against an older proto build would | ||
| # crash at import time when smg_grpc_servicer.mlx.server runs. | ||
| mlx = ["smg-grpc-proto>=0.4.7", "mlx>=0.22.0", "mlx-lm>=0.22.0"] |
There was a problem hiding this comment.
mlx extra currently allows mlx-lm>=0.22.0, but this servicer imports BatchGenerator, SequenceStateMachine, and generation_stream from mlx_lm.generate (mlx/server.py and mlx/servicer.py). In mlx-lm 0.22.x those symbols are not available, so a resolver that pins 0.22.* (still valid under this constraint) will fail at import/startup time before serving any RPCs. Please bump the mlx-lm floor to the first version that exports this batching/state-machine API.
Useful? React with 👍 / 👎.
Adds the full E2E test surface for the MLX gRPC servicer that landed in #1099. Split off from #1099 to keep that review focused on the servicer itself. What's new ---------- - e2e_test/chat_completions/test_mlx_backend.py — 5 tests covering basic chat, streaming, tool calling (Qwen3 native <tool_call>), reasoning_content (Qwen3 thinking mode), and max_tokens finish reason. Module skips on non-Apple-Silicon hosts. - .github/workflows/pr-test-mlx.yml — runs the test suite on macos-latest (Apple Silicon) with HF model + Cargo target caching. Triggers only on MLX-related path changes. Plumbed through existing infra ------------------------------ - e2e_test/infra/constants.py: Runtime.MLX, is_mlx() helper, RUNTIME_LABELS entry, MLX added to LOCAL_RUNTIMES. - e2e_test/infra/__init__.py: re-exports is_mlx. - e2e_test/infra/model_specs.py: mlx-community/Qwen3-0.6B-4bit spec (~400 MB; smallest model that has both tool calling and thinking mode for combined coverage). - e2e_test/infra/worker.py: _build_mlx_cmd() — gRPC-only (rejects HTTP since the MLX servicer doesn't expose one). Test plan --------- Verified locally on Apple Silicon: 5/5 pass against the MLX servicer already on main. Signed-off-by: key4ng <rukeyang@gmail.com>
Adds the full E2E test surface for the MLX gRPC servicer that landed in #1099. Split off from #1099 to keep that review focused on the servicer itself. What's new ---------- - e2e_test/chat_completions/test_mlx_backend.py — 5 tests covering basic chat, streaming, tool calling (Qwen3 native <tool_call>), reasoning_content (Qwen3 thinking mode), and max_tokens finish reason. Module skips on non-Apple-Silicon hosts. - .github/workflows/pr-test-mlx.yml — runs the test suite on macos-latest (Apple Silicon) with HF model + Cargo target caching. Triggers only on MLX-related path changes. Plumbed through existing infra ------------------------------ - e2e_test/infra/constants.py: Runtime.MLX, is_mlx() helper, RUNTIME_LABELS entry, MLX added to LOCAL_RUNTIMES. - e2e_test/infra/__init__.py: re-exports is_mlx. - e2e_test/infra/model_specs.py: mlx-community/Qwen3-0.6B-4bit spec (~400 MB; smallest model that has both tool calling and thinking mode for combined coverage). - e2e_test/infra/worker.py: _build_mlx_cmd() — gRPC-only (rejects HTTP since the MLX servicer doesn't expose one). Test plan --------- Verified locally on Apple Silicon: 5/5 pass against the MLX servicer already on main. Signed-off-by: key4ng <rukeyang@gmail.com>
…smg-project#1099) Signed-off-by: Zhuo Li <zhuo.li.ca@outlook.com>
Adds the full E2E test surface for the MLX gRPC servicer that landed in #1099. Split off from #1099 to keep that review focused on the servicer itself. What's new ---------- - e2e_test/chat_completions/test_mlx_backend.py — 5 tests covering basic chat, streaming, tool calling (Qwen3 native <tool_call>), reasoning_content (Qwen3 thinking mode), and max_tokens finish reason. Module skips on non-Apple-Silicon hosts. - .github/workflows/pr-test-mlx.yml — runs the test suite on macos-latest (Apple Silicon) with HF model + Cargo target caching. Triggers only on MLX-related path changes. Plumbed through existing infra ------------------------------ - e2e_test/infra/constants.py: Runtime.MLX, is_mlx() helper, RUNTIME_LABELS entry, MLX added to LOCAL_RUNTIMES. - e2e_test/infra/__init__.py: re-exports is_mlx. - e2e_test/infra/model_specs.py: mlx-community/Qwen3-0.6B-4bit spec (~400 MB; smallest model that has both tool calling and thinking mode for combined coverage). - e2e_test/infra/worker.py: _build_mlx_cmd() — gRPC-only (rejects HTTP since the MLX servicer doesn't expose one). Test plan --------- Verified locally on Apple Silicon: 5/5 pass against the MLX servicer already on main. Signed-off-by: key4ng <rukeyang@gmail.com>
…smg-project#1099) Signed-off-by: Zhuo Li <zhuo.li.ca@outlook.com>
…smg-project#1099) Signed-off-by: Zhuo Li <zhuo.li.ca@outlook.com>
…smg-project#1099) Signed-off-by: Zhuo Li <zhuo.li.ca@outlook.com>
…smg-project#1099) Signed-off-by: Zhuo Li <zhuo.li.ca@outlook.com>
2 participants
Description
Problem
#1034 added the MlxEngine proto and Rust client. This PR adds the corresponding Python servicer that actually runs mlx-lm on Apple Silicon, completing the end-to-end pipeline.
Solution
A new
grpc_servicer/smg_grpc_servicer/mlx/module implementing theMlxEngineproto service. Wraps mlx-lm'sBatchGeneratorbehind the gRPC contract from PR #1034, enabling SMG to route requests to Apple Silicon workers the same way it routes to vLLM/SGLang/TRT-LLM.Changes
Python servicer (
grpc_servicer/smg_grpc_servicer/mlx/)MlxEngineServicer— implements all 8 RPCs. Generate supports both streaming and non-streaming. A background daemon thread runsBatchGenerator.next()in a tight loop and dispatches tokens to per-requestasyncio.Queues vialoop.call_soon_threadsafe.MlxHealthServicer— standardgrpc.health.v1.Healthprotocol for liveness/readiness probes.server.py— CLI entrypoint with model loading, warmup, graceful shutdown. Run viapython -m smg_grpc_servicer.mlx.server --model <path>.Sampling parameter mapping
temperature,top_p,top_k,min_pmake_sampler(...)repetition_penalty,frequency_penalty,presence_penalty,logit_biasmake_logits_processors(...)max_tokens,stop_token_ids,ignore_eosBatchGenerator.insert+SequenceStateMachinelogprobsgenerate_steplogprobs arrayseedmx.random.seed()(best-effort, global state)Packaging
mlxoptional dependency ingrpc_servicer/pyproject.toml: install withpip install smg-grpc-servicer[mlx]smg-grpc-protoPython package exportsmlx_engine_pb2/mlx_engine_pb2_grpcE2E test (
e2e_test/chat_completions/test_mlx_backend.py)5 tests exercising the full SMG router → gRPC → MLX pipeline:
ToolParser(parses<tool_call>tags)max_tokensfinish reasonUses
mlx-community/Qwen3-0.6B-4bit(~400 MB). Skipped on non-Darwin/non-arm64 via module-levelpytestmark.Infrastructure additions:
Runtime.MLXenum +is_mlx()helper ine2e_test/infra/constants.py_build_mlx_cmd()ine2e_test/infra/worker.py(gRPC-only, rejects HTTP mode)mlx-community/Qwen3-0.6B-4bitCI (
.github/workflows/pr-test-mlx.yml)Runs the E2E test on
macos-14(Apple Silicon) GitHub-hosted runners.Swatinem/rust-cache) and HF model (~/.cache/huggingface/hub)Test Plan
Verified locally on Apple Silicon with
mlx-community/Qwen3-0.6B-4bit: 15/15 test passes across 3 consecutive runs. Also verified earlier withmlx-community/Qwen3-8B-4bitfor production-quality output.# Local run E2E_RUNTIME=mlx pytest e2e_test/chat_completions/test_mlx_backend.py -vChecklist
cargo +nightly fmtpassescargo clippy --all-targets --all-features -- -D warningspassesSummary by CodeRabbit
New Features
Tests
Chores