feat(mlx): add Python gRPC servicer for MLX backend

.github/workflows/pr-test-mlx.yml

@@ -0,0 +1,135 @@
+name: PR Test (MLX)
+
+on:
+  push:
+    branches: [main]
+    paths:
+      - "crates/grpc_client/proto/mlx_engine.proto"
+      - "crates/grpc_client/src/mlx_engine.rs"
+      - "crates/grpc_client/python/**"
+      - "grpc_servicer/smg_grpc_servicer/mlx/**"
+      - "grpc_servicer/pyproject.toml"
+      - "e2e_test/chat_completions/test_mlx_backend.py"
+      - "e2e_test/infra/model_specs.py"
+      - "e2e_test/infra/worker.py"
+      - "e2e_test/infra/constants.py"
+      - ".github/workflows/pr-test-mlx.yml"
+  pull_request:
+    branches: [main]
+    types: [opened, synchronize, reopened]
+    paths:
+      - "crates/grpc_client/proto/mlx_engine.proto"
+      - "crates/grpc_client/src/mlx_engine.rs"
+      - "crates/grpc_client/python/**"
+      - "grpc_servicer/smg_grpc_servicer/mlx/**"
+      - "grpc_servicer/pyproject.toml"
+      - "e2e_test/chat_completions/test_mlx_backend.py"
+      - "e2e_test/infra/model_specs.py"
+      - "e2e_test/infra/worker.py"
+      - "e2e_test/infra/constants.py"
+      - ".github/workflows/pr-test-mlx.yml"
+  workflow_dispatch:
+
+permissions:
+  contents: read
+
+concurrency:
+  group: mlx-tests-${{ github.ref }}
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
+
+jobs:
+  e2e-mlx:
+    name: E2E (MLX on Apple Silicon)
+    runs-on: macos-latest
+    timeout-minutes: 30
+    permissions:
+      contents: read
+    env:
+      E2E_RUNTIME: mlx
+      E2E_ENGINE: mlx
+      PYTHONUNBUFFERED: "1"
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v6
+
+      - name: Set up Python
+        uses: actions/setup-python@v6
+        with:
+          python-version: "3.12"
+
+      - name: Install Rust toolchain
+        uses: dtolnay/rust-toolchain@stable
+
+      - name: Install protoc
+        run: brew install protobuf
+
+      - name: Cache Cargo registry and target
+        uses: Swatinem/rust-cache@v2
+        with:
+          shared-key: mlx-pr-test
+          cache-on-failure: true
+
+      - name: Cache HuggingFace models
+        uses: actions/cache@v4
+        with:
+          path: ~/.cache/huggingface/hub
+          key: hf-mlx-${{ runner.os }}-Qwen3-0.6B-4bit
+          restore-keys: |
+            hf-mlx-${{ runner.os }}-
+
+      - name: Install uv (for openapi codegen)
+        run: pip install uv
+
+      - name: Build smg-grpc-proto Python package (proto codegen)
+        run: pip install -e ./crates/grpc_client/python
+
+      - name: Install grpc_servicer with MLX extra (mlx + mlx-lm)
+        run: pip install -e "./grpc_servicer[mlx]"
+
+      - name: Build and install SMG Python bindings (ci profile)
+        working-directory: bindings/python
+        run: |
+          pip install maturin
+          # `ci` profile (opt-level=2, thin LTO, 16 codegen-units) — faster
+          # to compile than release, runtime still plenty fast for a
+          # correctness E2E test.
+          # Use `maturin build` + `pip install` (not `maturin develop`)
+          # because the GitHub-hosted runner's Python is not in a virtualenv.
+          maturin build --profile ci --out dist
+          pip install dist/*.whl
+
+      - name: Generate Python client types (required by e2e_test/conftest.py)
+        run: make generate-python-types
+
+      - name: Install smg-client
+        run: pip install ./clients/python
+
+      - name: Install E2E test dependencies
+        run: pip install ./e2e_test
+
+      - name: Verify imports
+        run: |
+          python -c "from smg_grpc_proto import mlx_engine_pb2, mlx_engine_pb2_grpc; print('proto OK')"
+          python -c "from smg_grpc_servicer.mlx.servicer import MlxEngineServicer; print('servicer OK')"
+          python -c "import smg; print('smg OK')"
+          python -c "from smg_client import SmgClient; print('smg_client OK')"
+          python -c "import mlx_lm; print('mlx-lm OK')"
+
+      - name: Run MLX E2E tests
+        env:
+          SHOW_WORKER_LOGS: "1"
+          SHOW_ROUTER_LOGS: "1"
+          E2E_LOG_DIR: e2e-logs
+        run: |
+          pytest e2e_test/chat_completions/test_mlx_backend.py \
+            -s -vv \
+            --reruns 1 --reruns-delay 5
+
+      - name: Upload logs on failure
+        if: failure() || cancelled()
+        uses: actions/upload-artifact@v7
+        with:
+          name: e2e-mlx-logs
+          path: e2e-logs/
+          retention-days: 7
+          if-no-files-found: ignore

crates/grpc_client/python/smg_grpc_proto/__init__.py

@@ -1,4 +1,4 @@
-"""SMG gRPC Proto - Protocol definitions for SGLang, vLLM, and TRT-LLM."""
+"""SMG gRPC Proto - Protocol definitions for SGLang, vLLM, TRT-LLM, and MLX."""
 
 from importlib.metadata import version
 
@@ -8,6 +8,8 @@
 # These imports will work after the package is built (stubs generated at build time)
 try:
     from smg_grpc_proto.generated import (
+        mlx_engine_pb2,
+        mlx_engine_pb2_grpc,
         sglang_encoder_pb2,
         sglang_encoder_pb2_grpc,
         sglang_scheduler_pb2,
@@ -27,6 +29,8 @@
         "vllm_engine_pb2_grpc",
         "trtllm_service_pb2",
         "trtllm_service_pb2_grpc",
+        "mlx_engine_pb2",
+        "mlx_engine_pb2_grpc",
     ]
 except ImportError:
     # During development/build, generated modules may not exist yet

e2e_test/chat_completions/test_mlx_backend.py

@@ -0,0 +1,165 @@
+"""MLX backend E2E tests (Apple Silicon only).
+
+Tests the SMG router → gRPC → MLX worker pipeline using mlx-lm's
+BatchGenerator. The MLX backend only supports gRPC mode.
+
+Run locally with:
+    E2E_RUNTIME=mlx pytest e2e_test/chat_completions/test_mlx_backend.py -v
+
+CI runs on macos-14 (Apple Silicon) with the smallest Qwen3 model (~400 MB).
+See .github/workflows/pr-test-mlx.yml.
+"""
+
+from __future__ import annotations
+
+import json
+import platform
+import sys
+
+import pytest
+
+# MLX is Apple Silicon only — skip the entire module elsewhere.
+pytestmark = pytest.mark.skipif(
+    sys.platform != "darwin" or platform.machine() != "arm64",
+    reason="MLX backend requires Apple Silicon (macOS arm64)",
+)
+
+
+# ── Tools used for function calling tests ────────────────────────────────────
+
+WEATHER_TOOL = {
+    "type": "function",
+    "function": {
+        "name": "get_weather",
+        "description": "Get the current weather in a given location.",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "location": {
+                    "type": "string",
+                    "description": "The city name, e.g. 'Tokyo' or 'Paris'.",
+                },
+                "unit": {
+                    "type": "string",
+                    "enum": ["celsius", "fahrenheit"],
+                },
+            },
+            "required": ["location"],
+        },
+    },
+}
+
+
+@pytest.mark.engine("mlx")
+@pytest.mark.gpu(0)  # MLX uses unified memory, no discrete GPU
+@pytest.mark.model("mlx-community/Qwen3-0.6B-4bit")
+@pytest.mark.gateway(extra_args=["--tool-call-parser", "qwen"])
+@pytest.mark.parametrize("setup_backend", ["grpc"], indirect=True)
+class TestMlxBackend:
+    """End-to-end tests for the MLX gRPC backend.
+
+    Uses mlx-community/Qwen3-0.6B-4bit (~400 MB) — small enough to download
+    and run on a macos-latest GitHub Actions runner. Qwen3 supports both
+    native tool calling and thinking mode in a single model.
+    """
+
+    # Qwen3 enters thinking mode by default; disable for the basic chat tests
+    # so the small 0.6B model produces actual content within max_tokens budget.
+    NO_THINKING = {"chat_template_kwargs": {"enable_thinking": False}}
+
+    def test_basic_chat(self, model, api_client):
+        """Non-streaming chat completion goes through the full pipeline."""
+        response = api_client.chat.completions.create(
+            model=model,
+            messages=[{"role": "user", "content": "Reply with the single word 'OK'."}],
+            max_tokens=20,
+            temperature=0,
+            extra_body=self.NO_THINKING,
+        )
+        assert response.id
+        assert response.choices[0].message.role == "assistant"
+        assert isinstance(response.choices[0].message.content, str)
+        assert len(response.choices[0].message.content) > 0
+        assert response.usage.prompt_tokens > 0
+        assert response.usage.completion_tokens > 0
+
+    def test_streaming_chat(self, model, api_client):
+        """Streaming chat yields incremental delta chunks."""
+        stream = api_client.chat.completions.create(
+            model=model,
+            messages=[{"role": "user", "content": "Count: 1, 2, 3."}],
+            max_tokens=30,
+            temperature=0,
+            stream=True,
+            extra_body=self.NO_THINKING,
+        )
+        chunks = list(stream)
+        # First chunk has the role, subsequent chunks carry content deltas.
+        assert len(chunks) >= 2
+        text_pieces = [
+            c.choices[0].delta.content for c in chunks if c.choices and c.choices[0].delta.content
+        ]
+        assert len(text_pieces) > 0
+        assert "".join(text_pieces).strip()
+
+    def test_tool_calling(self, model, api_client):
+        """Qwen3 emits <tool_call> tags; SMG's qwen ToolParser extracts them."""
+        response = api_client.chat.completions.create(
+            model=model,
+            messages=[
+                {"role": "user", "content": "What's the weather in Tokyo right now?"},
+            ],
+            tools=[WEATHER_TOOL],
+            tool_choice="auto",
+            max_tokens=300,
+            temperature=0,
+        )
+        msg = response.choices[0].message
+        assert response.choices[0].finish_reason == "tool_calls", (
+            f"Expected finish_reason=tool_calls, got {response.choices[0].finish_reason!r} "
+            f"with content={msg.content!r}"
+        )
+        assert msg.tool_calls, "Expected tool_calls to be populated"
+        call = msg.tool_calls[0]
+        assert call.function.name == "get_weather"
+        # arguments is a JSON string per OpenAI spec
+        args = json.loads(call.function.arguments)
+        assert "tokyo" in args.get("location", "").lower()
+
+    def test_reasoning_thinking_mode(self, model, api_client):
+        """Qwen3 thinking mode populates reasoning_content separately from content."""
+        response = api_client.chat.completions.create(
+            model=model,
+            messages=[
+                {
+                    "role": "user",
+                    "content": (
+                        "Alice has 5 apples. She gives 2 to Bob. "
+                        "How many apples does Alice have? Think step by step."
+                    ),
+                }
+            ],
+            max_tokens=600,
+            temperature=0,
+        )
+        msg = response.choices[0].message
+        # Qwen3-0.6B may or may not always use thinking mode for trivial questions,
+        # but at minimum the final answer must mention 3.
+        content = msg.content or ""
+        reasoning = getattr(msg, "reasoning_content", None) or ""
+        full = content + reasoning
+        assert "3" in full, (
+            f"Expected '3' in response, got content={content!r} reasoning={reasoning!r}"
+        )
+
+    def test_max_tokens_finish_reason(self, model, api_client):
+        """When max_tokens is reached, finish_reason is 'length'."""
+        response = api_client.chat.completions.create(
+            model=model,
+            messages=[{"role": "user", "content": "Tell me a story."}],
+            max_tokens=10,
+            temperature=0,
+            extra_body=self.NO_THINKING,
+        )
+        assert response.choices[0].finish_reason == "length"
+        assert response.usage.completion_tokens == 10

e2e_test/infra/__init__.py

@@ -32,6 +32,7 @@
     Runtime,
     WorkerType,
     get_runtime,
+    is_mlx,
     is_sglang,
     is_trtllm,
     is_vllm,
@@ -108,6 +109,7 @@
     "is_vllm",
     "is_sglang",
     "is_trtllm",
+    "is_mlx",
     # Port utilities
     "get_open_port",
     "release_port",

e2e_test/infra/constants.py

@@ -25,6 +25,7 @@ class Runtime(StrEnum):
     SGLANG = "sglang"
     VLLM = "vllm"
     TRTLLM = "trtllm"
+    MLX = "mlx"
     OPENAI = "openai"
     XAI = "xai"
     GEMINI = "gemini"
@@ -33,7 +34,7 @@ class Runtime(StrEnum):
 
 # Convenience sets
 LOCAL_MODES = frozenset({ConnectionMode.HTTP, ConnectionMode.GRPC})
-LOCAL_RUNTIMES = frozenset({Runtime.SGLANG, Runtime.VLLM, Runtime.TRTLLM})
+LOCAL_RUNTIMES = frozenset({Runtime.SGLANG, Runtime.VLLM, Runtime.TRTLLM, Runtime.MLX})
 CLOUD_RUNTIMES = frozenset({Runtime.OPENAI, Runtime.XAI, Runtime.GEMINI, Runtime.ANTHROPIC})
 
 # Fixture parameter names (used in @pytest.mark.parametrize)
@@ -100,11 +101,21 @@ def is_trtllm() -> bool:
     return get_runtime() == "trtllm"
 
 
+def is_mlx() -> bool:
+    """Check if tests are running with MLX runtime (Apple Silicon only).
+
+    Returns:
+        True if E2E_RUNTIME is "mlx", False otherwise.
+    """
+    return get_runtime() == "mlx"
+
+
 # Runtime display labels
 RUNTIME_LABELS = {
     "sglang": "SGLang",
     "vllm": "vLLM",
     "trtllm": "TensorRT-LLM",
+    "mlx": "MLX",
 }
 
 ENV_SHOW_ROUTER_LOGS = "SHOW_ROUTER_LOGS"

e2e_test/infra/model_specs.py

@@ -184,6 +184,16 @@ def _resolve_model_path(hf_path: str) -> str:
             "--enable-chunked-prefill",
         ],
     },
+    # ── MLX models (Apple Silicon only) ──────────────────────────────────────
+    # Smallest Qwen3 with native tool calling + thinking mode (~400 MB).
+    # Used by CI on macos-latest runners. Qwen3 emits <tool_call> tags
+    # parsed by SMG's --tool-call-parser qwen, and uses <think> tags
+    # parsed by the reasoning parser.
+    "mlx-community/Qwen3-0.6B-4bit": {
+        "model": _resolve_model_path("mlx-community/Qwen3-0.6B-4bit"),
+        "tp": 1,
+        "features": ["chat", "streaming", "function_calling", "reasoning", "thinking"],
+    },
 }