From bcec6bcccd931179aa7e8cb1b808a1bb4d358826 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Wed, 24 Jun 2026 06:02:48 +0900 Subject: [PATCH 1/4] feat: add Phase 1 Python client package (mlxcel) over the server Add a pure-Python client package under python/ that drives the existing OpenAI-compatible mlxcel server. It either spawns and supervises a local `mlxcel serve` process (managed mode) or connects to a running one (connect mode), auto-discovers the served model id from /v1/models, and exposes the raw openai client as an escape hatch. No changes to the Rust inference core. Package contents: - src/mlxcel/_server.py: ManagedServer handles binary discovery (binary= / MLXCEL_BIN / PATH), transport selection (Unix socket default on POSIX with a short sun_path under /tmp, TCP elsewhere or on request), subprocess spawn, /health readiness polling with backoff and child-liveness checks, stderr forwarding to the mlxcel.server logger, and graceful SIGTERM-then-SIGKILL shutdown with socket cleanup, atexit, and a finalizer. - src/mlxcel/_client.py and src/mlxcel/_async_client.py: the synchronous LLM and asynchronous AsyncLLM, each wrapping the OpenAI SDK over a TCP or UDS httpx transport with explicit timeouts. Shared mode-selection, base-URL handling, and message-type narrowing live in src/mlxcel/_common.py. Methods: generate, stream, chat, chat_stream, models, tokenize, detokenize, plus model and openai_client properties and close(). tokenize/detokenize call the native /tokenize and /detokenize routes through the underlying httpx client. - src/mlxcel/_sampling.py: maps Python kwargs to OpenAI request fields and routes server-specific knobs (top_k, min_p, repetition_penalty, DRY) and unknown keys through extra_body, with response_format passthrough. - src/mlxcel/errors.py: MlxcelError, MlxcelServerError (carries stderr tail), MlxcelTimeoutError. HTTP and API errors propagate as native openai exceptions. Tests, CI, and docs: - tests/: stdlib-only fake_server.py (UDS or TCP, /health 503-then-200, canned /v1/* incl. SSE, /tokenize, /detokenize); test_client_mock.py uses httpx.MockTransport; test_lifecycle.py spawns the fake server via ManagedServer; test_e2e.py is marked e2e and skipped unless MLXCEL_BIN is set. - .github/workflows/python.yml runs ruff, ruff format --check, mypy, and pytest (unit + lifecycle) on ubuntu-latest, independent of the Rust CI and triggered only on python/** changes. - docs/python-client.md documents both modes, streaming, chat, structured output, the openai_client escape hatch, async usage, and the socket-path-length note; linked from README.md, docs/README.md, and the mkdocs nav. --- .github/workflows/python.yml | 50 ++++ README.md | 16 ++ docs/README.md | 1 + docs/python-client.md | 180 +++++++++++++ mkdocs.yml | 1 + python/.gitignore | 10 + python/README.md | 80 ++++++ python/examples/quickstart.py | 37 +++ python/examples/streaming.py | 32 +++ python/examples/structured_output.py | 64 +++++ python/pyproject.toml | 86 ++++++ python/src/mlxcel/__init__.py | 30 +++ python/src/mlxcel/_async_client.py | 279 ++++++++++++++++++++ python/src/mlxcel/_client.py | 278 ++++++++++++++++++++ python/src/mlxcel/_common.py | 107 ++++++++ python/src/mlxcel/_sampling.py | 88 +++++++ python/src/mlxcel/_server.py | 378 +++++++++++++++++++++++++++ python/src/mlxcel/errors.py | 34 +++ python/src/mlxcel/py.typed | 0 python/tests/fake_server.py | 279 ++++++++++++++++++++ python/tests/test_client_mock.py | 293 +++++++++++++++++++++ python/tests/test_e2e.py | 46 ++++ python/tests/test_lifecycle.py | 195 ++++++++++++++ 23 files changed, 2564 insertions(+) create mode 100644 .github/workflows/python.yml create mode 100644 docs/python-client.md create mode 100644 python/.gitignore create mode 100644 python/README.md create mode 100644 python/examples/quickstart.py create mode 100644 python/examples/streaming.py create mode 100644 python/examples/structured_output.py create mode 100644 python/pyproject.toml create mode 100644 python/src/mlxcel/__init__.py create mode 100644 python/src/mlxcel/_async_client.py create mode 100644 python/src/mlxcel/_client.py create mode 100644 python/src/mlxcel/_common.py create mode 100644 python/src/mlxcel/_sampling.py create mode 100644 python/src/mlxcel/_server.py create mode 100644 python/src/mlxcel/errors.py create mode 100644 python/src/mlxcel/py.typed create mode 100644 python/tests/fake_server.py create mode 100644 python/tests/test_client_mock.py create mode 100644 python/tests/test_e2e.py create mode 100644 python/tests/test_lifecycle.py diff --git a/.github/workflows/python.yml b/.github/workflows/python.yml new file mode 100644 index 00000000..8da82d78 --- /dev/null +++ b/.github/workflows/python.yml @@ -0,0 +1,50 @@ +# PR / push CI for the Python client package under python/. Runs ruff (lint + +# format), mypy (type check), and pytest (unit + lifecycle tests). It needs +# neither MLX nor the mlxcel binary, so it runs on plain ubuntu-latest and is +# independent of the Rust CI in ci.yml. The e2e test is skipped here because it +# requires a built binary (set MLXCEL_BIN to run it locally). + +name: Python + +on: + push: + branches: [main] + paths: + - 'python/**' + - '.github/workflows/python.yml' + pull_request: + branches: [main] + paths: + - 'python/**' + - '.github/workflows/python.yml' + +permissions: + contents: read + +jobs: + check: + name: lint, type-check, test + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v6 + with: + persist-credentials: false + + - uses: actions/setup-python@v5 + with: + python-version: '3.11' + + - name: Install package with dev extras + run: python -m pip install --upgrade pip && pip install -e "python[dev]" + + - name: Ruff lint + run: ruff check python + + - name: Ruff format check + run: ruff format --check python + + - name: Mypy + run: mypy python/src + + - name: Pytest (unit + lifecycle) + run: pytest python/tests -m "not e2e" -q diff --git a/README.md b/README.md index 7b6e86d6..33508696 100644 --- a/README.md +++ b/README.md @@ -265,6 +265,21 @@ mlxcel arch for the CLI summary, and see [Supported models](docs/supported-models.md) for the maintained architecture table, known limitations, and VLM coverage notes. +## Python + +`mlxcel` ships a pure-Python client that drives the OpenAI-compatible server from Python. It spawns and manages a local `mlxcel serve` process (managed mode) or connects to a running one (connect mode), auto-discovers the served model id, and exposes the raw `openai` client for the full API surface. + +```python +import mlxcel + +with mlxcel.LLM("mlx-community/Qwen3-4B-4bit") as llm: + print(llm.generate("def fib(n):", max_tokens=128)) + for delta in llm.stream("Write a haiku about autumn"): + print(delta, end="", flush=True) +``` + +Install with `pip install ./python`. See [Python client](docs/python-client.md) for managed and connect modes, streaming, structured output, async usage, and troubleshooting. The client lives in [`python/`](python) and builds entirely on the existing server (no native extension). + ## Optional GUI `mlxcel-server` can be used directly through HTTP clients. For a local graphical front-end, [Backend.AI Go](https://go.backend.ai) can be used as a companion UI for chat, model management, and multi-model routing. @@ -279,6 +294,7 @@ for the CLI summary, and see [Supported models](docs/supported-models.md) for th - [Tensor and pipeline parallelism](docs/distributed.md) - [TurboQuant KV cache](docs/turbo-kv-cache.md) - [OpenAI Responses API](docs/responses-api.md) +- [Python client](docs/python-client.md) - [Adding a new model](docs/adding-models.md) ## Contributing diff --git a/docs/README.md b/docs/README.md index 0c9d7919..e605f935 100644 --- a/docs/README.md +++ b/docs/README.md @@ -25,6 +25,7 @@ Current GitHub-facing docs: 10. `audio-api.md` — implemented `/v1/audio` endpoints: Whisper STT setup, request/response reference, WAV encoding details, and request validation order. 11. `adding-models.md` — contribution guide for new model architectures. 12. `block-diffusion.md` — DiffusionGemma block-diffusion generation: canvas denoising vs autoregressive, CLI flags, throughput, and phase 1 limitations. +13. `python-client.md`: the `mlxcel` Python client over the OpenAI-compatible server, covering managed and connect modes, streaming, chat, structured output, the `openai_client` escape hatch, async usage, and troubleshooting. ## Architecture Decision Records diff --git a/docs/python-client.md b/docs/python-client.md new file mode 100644 index 00000000..ca0c663e --- /dev/null +++ b/docs/python-client.md @@ -0,0 +1,180 @@ +# Python client + +`mlxcel` ships a pure-Python client that drives the OpenAI-compatible server (`mlxcel serve`) from Python. It either spawns and supervises a local server process (managed mode) or connects to an already-running one (connect mode), auto-discovers the served model id, and exposes the raw `openai` client for the full API surface (tools, `response_format`, logprobs, multimodal). + +This is Phase 1 of Python integration. It builds entirely on the existing HTTP server: no native extension, and no changes to the Rust inference core. The package lives in [`python/`](../python) and is published under the import name `mlxcel`. + +## Install + +```bash +pip install ./python # from a repo checkout +pip install ./python[dev] # adds pytest, ruff, mypy +``` + +Requires Python 3.9 or newer. Runtime dependencies are `openai>=1.40` and `httpx>=0.27`. Managed mode additionally needs the `mlxcel` binary; the client finds it via the `binary=` argument, the `MLXCEL_BIN` environment variable, or `mlxcel` on `PATH`, in that order. See [Installation](installation.md) for building the binary. + +## Managed mode + +Pass a model and the client spawns `mlxcel serve`, waits until the server reports ready, and stops it on exit. The model loads and warms up before the listener binds, so a first run that downloads weights can take a while; the client polls `/health` with a generous `startup_timeout` (default 600 seconds) and watches the child process for an early exit. + +```python +import mlxcel + +with mlxcel.LLM("mlx-community/Qwen3-4B-4bit") as llm: + text = llm.generate("def fib(n):", max_tokens=128, temperature=0.7) + print(text) + + reply = llm.chat([{"role": "user", "content": "Hello"}], max_tokens=64) + print(reply) + + print(llm.model) # resolved model id, auto-discovered from /v1/models + print(llm.models()) # [""] + ids = llm.tokenize("hello world") + print(llm.detokenize(ids)) +``` + +The context manager (`with`) is the recommended form; it guarantees shutdown. The client also registers an `atexit` cleanup and a best-effort finalizer so a leaked handle still stops the server. + +Useful managed-mode arguments: `binary=`, `host=`/`port=` (forces TCP and binds there), `socket=` (an explicit Unix socket path), `api_key=`, `ctx_size=`, `n_predict=`, `alias=`, `warmup=`, `extra_args=[...]` (forwarded verbatim to `mlxcel serve`), and `startup_timeout=`. + +## Connect mode + +Pass `base_url=` or `socket=` (but not a model) to talk to a server you started yourself. No subprocess is launched or managed. + +```python +# TCP +llm = mlxcel.LLM(base_url="http://localhost:8080/v1") + +# Unix domain socket +llm = mlxcel.LLM(socket="/tmp/mlxcel.sock") + +print(llm.generate("Hello")) +llm.close() +``` + +Start a matching server with the CLI: + +```bash +# TCP +mlxcel serve -m mlx-community/Qwen3-4B-4bit --host 127.0.0.1 --port 8080 + +# Unix domain socket: --port 0 reinterprets --host as the socket path +mlxcel serve -m mlx-community/Qwen3-4B-4bit --host /tmp/mlxcel.sock --port 0 +``` + +Passing both a model and a connect target raises `MlxcelError`, because the mode would be ambiguous. + +## Streaming + +`stream` and `chat_stream` yield text deltas as they arrive. + +```python +with mlxcel.LLM("mlx-community/Qwen3-4B-4bit") as llm: + for delta in llm.stream("Write a haiku about autumn"): + print(delta, end="", flush=True) + + for delta in llm.chat_stream([{"role": "user", "content": "List three uses for a Pi."}]): + print(delta, end="", flush=True) +``` + +## Chat + +`chat` returns the assistant message content as a string. + +```python +messages = [ + {"role": "system", "content": "You are concise."}, + {"role": "user", "content": "What is MLX?"}, +] +print(llm.chat(messages, max_tokens=128, temperature=0.3)) +``` + +## Sampling parameters + +Generation methods accept OpenAI sampling fields directly: `max_tokens`, `temperature`, `top_p`, `stop`, `seed`, `presence_penalty`, `frequency_penalty`, `logit_bias`, and `response_format`. Server-specific knobs that are not part of the OpenAI schema (`top_k`, `min_p`, `repetition_penalty`, and DRY settings) are forwarded in the request body. You can also pass an explicit `extra_body={...}` for arbitrary server fields; values you set there win on conflict. + +```python +llm.generate("Once upon a time", max_tokens=200, top_p=0.9, top_k=40, min_p=0.05) +llm.generate("...", extra_body={"repetition_penalty": 1.1}) +``` + +## Structured output + +The server's llguidance-backed constrained decoding honors `response_format`, so you can require schema-valid JSON. + +```python +import json + +schema = { + "type": "json_schema", + "json_schema": { + "name": "person", + "schema": { + "type": "object", + "properties": { + "name": {"type": "string"}, + "age": {"type": "integer"}, + }, + "required": ["name", "age"], + "additionalProperties": False, + }, + }, +} + +reply = llm.chat( + [{"role": "user", "content": "Invent a person as JSON."}], + response_format=schema, + max_tokens=128, +) +person = json.loads(reply) +``` + +## The `openai_client` escape hatch + +For anything the convenience methods do not cover (tools, logprobs, vision and audio inputs, the Responses API), reach for the configured OpenAI client directly. It is wired to the same transport, so TCP and Unix-socket setups both work. + +```python +oai = llm.openai_client # openai.OpenAI (or openai.AsyncOpenAI for AsyncLLM) +oai.chat.completions.create( + model=llm.model, + messages=[{"role": "user", "content": "Weather in SF?"}], + tools=[...], +) +``` + +## Async usage + +`mlxcel.AsyncLLM` mirrors the synchronous API with `async`/`await` and async iterators, backed by `openai.AsyncOpenAI` and `httpx.AsyncClient`. + +```python +import asyncio +import mlxcel + +async def main(): + async with mlxcel.AsyncLLM("mlx-community/Qwen3-4B-4bit") as llm: + print(await llm.generate("def fib(n):", max_tokens=128)) + async for delta in llm.stream("Write a haiku"): + print(delta, end="", flush=True) + +asyncio.run(main()) +``` + +The managed-server lifecycle (spawn and readiness polling) runs synchronously inside the constructor since it is a one-time blocking setup; every generation call is async. The model id resolves lazily on the first request, so read `llm.model` only after a call has run, or pass an explicit `model=` override. + +## Errors + +| Exception | Raised when | +|-----------|-------------| +| `MlxcelError` | base class for all client errors (also used for ambiguous arguments) | +| `MlxcelServerError` | a managed server fails to launch, become ready, or stay alive (carries the captured stderr tail) | +| `MlxcelTimeoutError` | the managed server does not become ready within `startup_timeout` | + +HTTP and API errors from the server propagate as native `openai` SDK exceptions (for example `openai.APIStatusError`), so status codes and response bodies stay visible. Only lifecycle concerns are wrapped in `Mlxcel*` types. + +## Troubleshooting + +- **Server never becomes ready.** The first run downloads weights before binding the listener. Increase `startup_timeout`, and enable logging to watch progress: `logging.getLogger("mlxcel.server").setLevel(logging.INFO)`. Server stderr (including download and load progress) is forwarded to that logger. +- **`MlxcelServerError` on startup with a stderr tail.** The child process exited before becoming ready. The attached stderr usually names the cause (model not found, out of memory, bad flag). +- **Binary not found.** Pass `binary="/path/to/mlxcel"`, set `MLXCEL_BIN`, or put `mlxcel` on `PATH`. +- **Unix socket path too long.** `sun_path` is about 104 bytes on macOS and 108 on Linux. On macOS, `$TMPDIR` resolves to a long `/var/folders/...` path, so the client defaults the socket to a short name under `/tmp` instead. If your explicit `socket=` path exceeds the limit, the client raises a clear error; pass a shorter path. +- **Port already in use (TCP).** In managed mode without an explicit `port=`, the client picks a free ephemeral port. Pass `port=` only when you need a fixed one. diff --git a/mkdocs.yml b/mkdocs.yml index 405398e8..306f42a4 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -143,6 +143,7 @@ nav: - user-guide/index.md - CLI Commands: user-guide/cli.md - Server Mode: user-guide/server.md + - Python Client: user-guide/python-client.md - Vision-Language Models: user-guide/vision-language-models.md - LoRA Adapters: user-guide/lora.md - Speculative Decoding: user-guide/speculative-decoding.md diff --git a/python/.gitignore b/python/.gitignore new file mode 100644 index 00000000..0c3ec942 --- /dev/null +++ b/python/.gitignore @@ -0,0 +1,10 @@ +# Python build / cache artifacts for the mlxcel client package. +.venv/ +__pycache__/ +*.py[cod] +.pytest_cache/ +.mypy_cache/ +.ruff_cache/ +*.egg-info/ +dist/ +build/ diff --git a/python/README.md b/python/README.md new file mode 100644 index 00000000..a9bbf3b6 --- /dev/null +++ b/python/README.md @@ -0,0 +1,80 @@ +# mlxcel (Python client) + +A thin, pure-Python client for the [mlxcel](https://github.com/lablup/mlxcel) OpenAI-compatible inference server. It spawns and manages a local `mlxcel serve` process (managed mode) or connects to an already-running one (connect mode), auto-discovers the served model id, and exposes the raw `openai` client for the full API surface. + +This is Phase 1 of Python integration: it builds entirely on the existing HTTP server, so it needs no native extension and no changes to the Rust inference core. + +## Install + +```bash +pip install ./python # from a repo checkout +pip install ./python[dev] # with pytest, ruff, mypy for development +``` + +Requires Python 3.9+. The client itself is pure Python (`openai>=1.40`, `httpx>=0.27`). Managed mode additionally needs the `mlxcel` binary on `PATH`, or pass `binary=` / set `MLXCEL_BIN`. + +## Usage + +```python +import mlxcel + +# Managed mode: spawn and supervise a local server. +with mlxcel.LLM("mlx-community/Qwen3-4B-4bit") as llm: + print(llm.generate("def fib(n):", max_tokens=128, temperature=0.7)) + + for delta in llm.stream("Write a haiku about autumn"): + print(delta, end="", flush=True) + + print(llm.chat([{"role": "user", "content": "Hello"}], max_tokens=64)) + + print(llm.model) # resolved model id (auto-discovered) + print(llm.models()) # [""] + ids = llm.tokenize("hello world") + print(llm.detokenize(ids)) + + # Escape hatch: the full OpenAI surface (tools, response_format, logprobs). + oai = llm.openai_client + oai.chat.completions.create(model=llm.model, messages=[...], tools=[...]) + +# Connect mode: talk to an already-running server. +llm = mlxcel.LLM(base_url="http://localhost:8080/v1") # TCP +llm = mlxcel.LLM(socket="/tmp/mlxcel.sock") # Unix socket +``` + +Async usage mirrors the sync API via `mlxcel.AsyncLLM` (`await llm.generate(...)`, `async for delta in llm.stream(...)`). + +## Modes + +- **Managed mode** (default when `model=` is given): the client spawns `mlxcel serve`, waits until `/health` returns ready, forwards server logs to the `mlxcel.server` Python logger, and stops the process on exit. +- **Connect mode** (when `base_url=` or `socket=` is given): no subprocess; the client just talks to the server. +- Passing both a model and a connect target raises `MlxcelError`. + +On POSIX, managed mode defaults to a Unix domain socket for low-overhead local IPC. Keep socket paths short (`sun_path` is about 104 bytes on macOS, 108 on Linux); the default lives under `/tmp`. Pass `socket=` to override. Windows uses TCP. + +## Sampling parameters + +`generate`, `stream`, `chat`, and `chat_stream` accept OpenAI sampling fields directly: `max_tokens`, `temperature`, `top_p`, `stop`, `seed`, `presence_penalty`, `frequency_penalty`, `logit_bias`, `response_format`. Server-specific knobs (`top_k`, `min_p`, `repetition_penalty`, DRY settings) are forwarded in the request body; you can also pass an explicit `extra_body={...}`. + +## Errors + +- `MlxcelError`: base class. +- `MlxcelServerError`: launch, readiness, or crash failures (carries the server stderr tail). +- `MlxcelTimeoutError`: readiness timeout. + +HTTP and API errors surface as native `openai` SDK exceptions (for example `openai.APIStatusError`), not as `Mlxcel*` types. + +## Tests + +```bash +pip install -e ./python[dev] +ruff check python +ruff format --check python +mypy python/src +pytest python/tests -m "not e2e" # unit + lifecycle, no binary needed +``` + +The end-to-end test (`-m e2e`) is skipped unless `MLXCEL_BIN` points at a built binary: + +```bash +MLXCEL_BIN=/path/to/mlxcel pytest python/tests/test_e2e.py -m e2e +``` diff --git a/python/examples/quickstart.py b/python/examples/quickstart.py new file mode 100644 index 00000000..89dd12cd --- /dev/null +++ b/python/examples/quickstart.py @@ -0,0 +1,37 @@ +"""Managed-mode quickstart: spawn a local server and generate text. + +Run with a model available locally or downloadable from Hugging Face:: + + python python/examples/quickstart.py +""" + +from __future__ import annotations + +import mlxcel + +MODEL = "mlx-community/Qwen3-4B-4bit" + + +def main() -> None: + # Managed mode: mlxcel spawns and supervises a local `mlxcel serve` process, + # waits until it is ready, and shuts it down on exit. + with mlxcel.LLM(MODEL) as llm: + print("resolved model id:", llm.model) + print("available models:", llm.models()) + + text = llm.generate("def fib(n):", max_tokens=128, temperature=0.7) + print("\ncompletion:\n", text) + + reply = llm.chat( + [{"role": "user", "content": "Give me one fact about Apple Silicon."}], + max_tokens=64, + ) + print("\nchat reply:\n", reply) + + ids = llm.tokenize("hello world") + print("\ntokens:", ids) + print("round-trip:", llm.detokenize(ids)) + + +if __name__ == "__main__": + main() diff --git a/python/examples/streaming.py b/python/examples/streaming.py new file mode 100644 index 00000000..6d947225 --- /dev/null +++ b/python/examples/streaming.py @@ -0,0 +1,32 @@ +"""Streaming example: print completion and chat deltas as they arrive. + +python python/examples/streaming.py +""" + +from __future__ import annotations + +import sys + +import mlxcel + +MODEL = "mlx-community/Qwen3-4B-4bit" + + +def main() -> None: + with mlxcel.LLM(MODEL) as llm: + print("completion stream:") + for delta in llm.stream("Write a haiku about autumn.", max_tokens=64): + sys.stdout.write(delta) + sys.stdout.flush() + print("\n") + + print("chat stream:") + messages = [{"role": "user", "content": "List three uses for a Raspberry Pi."}] + for delta in llm.chat_stream(messages, max_tokens=128): + sys.stdout.write(delta) + sys.stdout.flush() + print() + + +if __name__ == "__main__": + main() diff --git a/python/examples/structured_output.py b/python/examples/structured_output.py new file mode 100644 index 00000000..c379a196 --- /dev/null +++ b/python/examples/structured_output.py @@ -0,0 +1,64 @@ +"""Structured output: constrain decoding to a JSON schema via response_format. + +The server's llguidance-backed constrained decoding honors ``response_format``. +This example asks for a JSON object matching a schema and parses the result. + + python python/examples/structured_output.py +""" + +from __future__ import annotations + +import json + +import mlxcel + +MODEL = "mlx-community/Qwen3-4B-4bit" + +SCHEMA = { + "type": "json_schema", + "json_schema": { + "name": "person", + "schema": { + "type": "object", + "properties": { + "name": {"type": "string"}, + "age": {"type": "integer"}, + "city": {"type": "string"}, + }, + "required": ["name", "age", "city"], + "additionalProperties": False, + }, + }, +} + + +def main() -> None: + with mlxcel.LLM(MODEL) as llm: + reply = llm.chat( + [ + { + "role": "user", + "content": "Invent a fictional person. Reply as JSON with name, age, city.", + } + ], + max_tokens=128, + temperature=0.7, + response_format=SCHEMA, + ) + print("raw reply:", reply) + parsed = json.loads(reply) + print("parsed:", parsed) + + # The same call via the raw OpenAI client (escape hatch): + oai = llm.openai_client + completion = oai.chat.completions.create( + model=llm.model, + messages=[{"role": "user", "content": "Another fictional person as JSON."}], + response_format=SCHEMA, + max_tokens=128, + ) + print("via openai_client:", completion.choices[0].message.content) + + +if __name__ == "__main__": + main() diff --git a/python/pyproject.toml b/python/pyproject.toml new file mode 100644 index 00000000..c79854f7 --- /dev/null +++ b/python/pyproject.toml @@ -0,0 +1,86 @@ +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[project] +name = "mlxcel" +version = "0.1.0" +description = "Python client for the mlxcel OpenAI-compatible inference server" +readme = "README.md" +requires-python = ">=3.9" +license = "Apache-2.0" +authors = [ + { name = "Lablup Inc." }, + { name = "Jeongkyu Shin" }, +] +keywords = ["mlx", "llm", "inference", "openai", "apple-silicon"] +classifiers = [ + "Development Status :: 3 - Alpha", + "Intended Audience :: Developers", + "License :: OSI Approved :: Apache Software License", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "Topic :: Scientific/Engineering :: Artificial Intelligence", + "Typing :: Typed", +] +dependencies = [ + "openai>=1.40", + "httpx>=0.27", +] + +[project.optional-dependencies] +dev = [ + "pytest>=8.0", + "ruff>=0.6", + "mypy>=1.11", +] + +[project.urls] +Homepage = "https://github.com/lablup/mlxcel" +Repository = "https://github.com/lablup/mlxcel" +Issues = "https://github.com/lablup/mlxcel/issues" + +[tool.hatch.build.targets.wheel] +packages = ["src/mlxcel"] + +[tool.pytest.ini_options] +markers = [ + "e2e: end-to-end tests requiring a real mlxcel binary (set MLXCEL_BIN); skipped by default", +] +testpaths = ["tests"] + +[tool.ruff] +line-length = 100 +target-version = "py39" +src = ["src", "tests", "examples"] + +[tool.ruff.lint] +select = ["E", "F", "I", "UP", "B", "W"] +ignore = [ + "UP006", # keep typing.List etc. usable on 3.9 where needed + "UP007", # keep Optional[...] / Union[...] for 3.9 runtime compatibility + "UP035", # allow typing imports under requires-python 3.9 + "UP037", # keep quoted self-referential annotations + "UP045", # keep Optional[...] over X | None for 3.9 clarity +] + +[tool.ruff.lint.per-file-ignores] +"tests/fake_server.py" = ["B008"] + +[tool.mypy] +python_version = "3.9" +strict = true +warn_unused_configs = true +disallow_untyped_defs = true +namespace_packages = true +explicit_package_bases = true +mypy_path = "src" + +[[tool.mypy.overrides]] +module = ["tests.*", "examples.*"] +disallow_untyped_defs = false +ignore_missing_imports = true diff --git a/python/src/mlxcel/__init__.py b/python/src/mlxcel/__init__.py new file mode 100644 index 00000000..ead4375d --- /dev/null +++ b/python/src/mlxcel/__init__.py @@ -0,0 +1,30 @@ +"""mlxcel: a thin Python client over the mlxcel OpenAI-compatible server. + +Drive a local or remote ``mlxcel serve`` process with minimal boilerplate. The +primary entry point is :class:`LLM` (synchronous) and its async twin +:class:`AsyncLLM`. Both manage a local server subprocess (managed mode) or +connect to a running one (connect mode), auto-discover the served model id, and +expose the raw OpenAI client for the full API surface. + + import mlxcel + + with mlxcel.LLM("mlx-community/Qwen3-4B-4bit") as llm: + print(llm.generate("def fib(n):", max_tokens=128)) +""" + +from __future__ import annotations + +from ._async_client import AsyncLLM +from ._client import LLM +from .errors import MlxcelError, MlxcelServerError, MlxcelTimeoutError + +__version__ = "0.1.0" + +__all__ = [ + "LLM", + "AsyncLLM", + "MlxcelError", + "MlxcelServerError", + "MlxcelTimeoutError", + "__version__", +] diff --git a/python/src/mlxcel/_async_client.py b/python/src/mlxcel/_async_client.py new file mode 100644 index 00000000..7e0c8df2 --- /dev/null +++ b/python/src/mlxcel/_async_client.py @@ -0,0 +1,279 @@ +"""Asynchronous client (:class:`AsyncLLM`) for an mlxcel server. + +``AsyncLLM`` mirrors :class:`mlxcel.LLM` with ``async``/``await`` and async +iterators, backed by ``openai.AsyncOpenAI`` and ``httpx.AsyncClient``. The +managed-server lifecycle (spawn, readiness polling) runs synchronously inside +``__init__`` because it is a one-time blocking setup; every generation method is +async. The model id resolves lazily on the first request. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, AsyncIterator, Dict, Iterable, List, Optional, cast + +import httpx +from openai import AsyncOpenAI + +from ._common import ( + DEFAULT_TIMEOUT, + UDS_BASE, + ChatMessages, + as_messages, + connect_base_url, + is_managed, + native_base_url, + normalize_base_url, +) +from ._sampling import build_params +from ._server import ManagedServer +from .errors import MlxcelError + +if TYPE_CHECKING: + from openai import AsyncStream + from openai.types.chat import ChatCompletionChunk + from openai.types.completion import Completion + + +class AsyncLLM: + """Asynchronous mlxcel client mirroring :class:`mlxcel.LLM`. + + Use ``async with`` or call :meth:`close` (awaitable). + """ + + def __init__( + self, + model: Optional[str] = None, + *, + base_url: Optional[str] = None, + socket: Optional[str] = None, + api_key: Optional[str] = None, + binary: Optional[str] = None, + host: Optional[str] = None, + port: Optional[int] = None, + timeout: Optional[httpx.Timeout] = None, + startup_timeout: float = 600.0, + transport: Optional[httpx.AsyncBaseTransport] = None, + **server_kwargs: Any, + ) -> None: + managed = is_managed(model, base_url, socket, transport) + + self._server: Optional[ManagedServer] = None + self._http_client: Optional[httpx.AsyncClient] = None + self._model: Optional[str] = None + self._closed = False + timeout = timeout or DEFAULT_TIMEOUT + + model_override = server_kwargs.pop("model", None) + + if not managed: + resolved_base = self._connect(base_url, socket, transport, timeout) + else: + assert model is not None + resolved_base = self._spawn( + model, binary, host, port, socket, api_key, startup_timeout, timeout, server_kwargs + ) + + self._client = AsyncOpenAI( + base_url=resolved_base, + api_key=api_key or "-", + http_client=self._http_client, + ) + self._base_url = resolved_base + self._model_override = model_override + + # -- setup ------------------------------------------------------------- + + def _connect( + self, + base_url: Optional[str], + socket: Optional[str], + transport: Optional[httpx.AsyncBaseTransport], + timeout: httpx.Timeout, + ) -> str: + """Build the connect-mode async http client and return the base URL.""" + if transport is not None: + resolved_base = connect_base_url(base_url) + self._http_client = httpx.AsyncClient( + transport=transport, base_url=resolved_base, timeout=timeout + ) + elif socket is not None: + uds_transport = httpx.AsyncHTTPTransport(uds=socket) + self._http_client = httpx.AsyncClient( + transport=uds_transport, base_url=UDS_BASE, timeout=timeout + ) + resolved_base = f"{UDS_BASE}/v1" + else: + assert base_url is not None + resolved_base = normalize_base_url(base_url) + self._http_client = httpx.AsyncClient(base_url=resolved_base, timeout=timeout) + return resolved_base + + def _spawn( + self, + model: str, + binary: Optional[str], + host: Optional[str], + port: Optional[int], + socket: Optional[str], + api_key: Optional[str], + startup_timeout: float, + timeout: httpx.Timeout, + server_kwargs: Dict[str, Any], + ) -> str: + """Spawn and wait for a managed server, build its async http client, return the base URL.""" + self._server = ManagedServer( + model, + binary=binary, + host=host, + port=port, + socket_path=socket, + api_key=api_key, + startup_timeout=startup_timeout, + **server_kwargs, + ) + self._server.start() + resolved_base = self._server.base_url + if self._server.uds_path is not None: + uds_transport = httpx.AsyncHTTPTransport(uds=self._server.uds_path) + self._http_client = httpx.AsyncClient( + transport=uds_transport, base_url=UDS_BASE, timeout=timeout + ) + else: + self._http_client = httpx.AsyncClient(base_url=resolved_base, timeout=timeout) + return resolved_base + + async def _resolve_model(self) -> str: + if self._model is not None: + return self._model + if self._model_override is not None: + self._model = self._model_override + return self._model + data = await self._client.models.list() + items = list(data.data) + if not items: + raise MlxcelError("Server reported no models from /v1/models.") + self._model = items[0].id + return self._model + + # -- properties -------------------------------------------------------- + + @property + def model(self) -> str: + """The resolved model id. + + Raises: + MlxcelError: if accessed before any request has resolved it. Await a + generation call (or :meth:`models`) first, or pass an explicit + ``model=`` override. + """ + if self._model is None: + raise MlxcelError("Model id not resolved yet. Await a request first, or pass model=.") + return self._model + + @property + def openai_client(self) -> AsyncOpenAI: + """The configured ``openai.AsyncOpenAI`` instance.""" + return self._client + + # -- generation -------------------------------------------------------- + + async def generate(self, prompt: str, **sampling: Any) -> str: + """Generate a completion for ``prompt`` and return the text.""" + self._ensure_alive() + model = await self._resolve_model() + params = build_params(sampling) + resp = await self._client.completions.create(model=model, prompt=prompt, **params) + return resp.choices[0].text or "" + + async def stream(self, prompt: str, **sampling: Any) -> AsyncIterator[str]: + """Stream completion text deltas for ``prompt``.""" + self._ensure_alive() + model = await self._resolve_model() + params = build_params(sampling) + stream = cast( + "AsyncStream[Completion]", + await self._client.completions.create( + model=model, prompt=prompt, stream=True, **params + ), + ) + async for chunk in stream: + if chunk.choices and chunk.choices[0].text: + yield chunk.choices[0].text + + async def chat(self, messages: ChatMessages, **sampling: Any) -> str: + """Run a chat completion and return the assistant message content.""" + self._ensure_alive() + model = await self._resolve_model() + params = build_params(sampling) + resp = await self._client.chat.completions.create( + model=model, messages=as_messages(messages), **params + ) + return resp.choices[0].message.content or "" + + async def chat_stream(self, messages: ChatMessages, **sampling: Any) -> AsyncIterator[str]: + """Stream chat completion content deltas.""" + self._ensure_alive() + model = await self._resolve_model() + params = build_params(sampling) + stream = cast( + "AsyncStream[ChatCompletionChunk]", + await self._client.chat.completions.create( + model=model, messages=as_messages(messages), stream=True, **params + ), + ) + async for chunk in stream: + if chunk.choices and chunk.choices[0].delta.content: + yield chunk.choices[0].delta.content + + # -- model / tokenizer helpers ---------------------------------------- + + async def models(self) -> List[str]: + """List model ids advertised by the server.""" + self._ensure_alive() + data = await self._client.models.list() + return [m.id for m in data.data] + + async def tokenize(self, text: str, add_special: bool = False) -> List[int]: + """Tokenize ``text`` via the server's native ``/tokenize`` route.""" + self._ensure_alive() + resp = await self._raw_post("/tokenize", {"content": text, "add_special": add_special}) + return list(resp["tokens"]) + + async def detokenize(self, tokens: Iterable[int]) -> str: + """Decode token ids back to text via the native ``/detokenize`` route.""" + self._ensure_alive() + resp = await self._raw_post("/detokenize", {"tokens": list(tokens)}) + return str(resp["content"]) + + async def _raw_post(self, path: str, json: Dict[str, Any]) -> Dict[str, Any]: + assert self._http_client is not None + url = native_base_url(self._base_url) + path + response = await self._http_client.post(url, json=json) + response.raise_for_status() + result: Dict[str, Any] = response.json() + return result + + # -- lifecycle --------------------------------------------------------- + + def _ensure_alive(self) -> None: + if self._server is not None: + self._server.ensure_alive() + + async def close(self) -> None: + """Close the async HTTP client and stop the managed server, if any.""" + if self._closed: + return + self._closed = True + if self._http_client is not None: + await self._http_client.aclose() + if self._server is not None: + self._server.close() + + async def __aenter__(self) -> "AsyncLLM": + return self + + async def __aexit__(self, *exc: object) -> None: + await self.close() + + +__all__ = ["AsyncLLM"] diff --git a/python/src/mlxcel/_client.py b/python/src/mlxcel/_client.py new file mode 100644 index 00000000..885002da --- /dev/null +++ b/python/src/mlxcel/_client.py @@ -0,0 +1,278 @@ +"""Synchronous client (:class:`LLM`) for an mlxcel server. + +``LLM`` wraps the OpenAI SDK with an httpx transport that works over either TCP +or a Unix domain socket. In *managed mode* it spawns and supervises a local +``mlxcel serve`` process; in *connect mode* it talks to an already-running +server. After the server is ready the resolved model id is discovered once from +``/v1/models`` and cached so callers never pass a model string. The raw OpenAI +client is exposed via ``openai_client`` for the full API surface (tools, +``response_format``, logprobs, multimodal). The async twin lives in +:mod:`._async_client`. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Dict, Iterable, Iterator, List, Optional, cast + +import httpx +from openai import OpenAI + +from ._common import ( + DEFAULT_TIMEOUT, + UDS_BASE, + ChatMessages, + as_messages, + connect_base_url, + is_managed, + native_base_url, + normalize_base_url, +) +from ._sampling import build_params +from ._server import ManagedServer +from .errors import MlxcelError + +if TYPE_CHECKING: + from openai import Stream + from openai.types.chat import ChatCompletionChunk + from openai.types.completion import Completion + + +class LLM: + """Synchronous mlxcel client. + + Examples: + Managed mode (spawns a local server):: + + with mlxcel.LLM("mlx-community/Qwen3-4B-4bit") as llm: + print(llm.generate("def fib(n):", max_tokens=128)) + + Connect mode (existing server):: + + llm = mlxcel.LLM(base_url="http://localhost:8080/v1") + llm = mlxcel.LLM(socket="/tmp/mlxcel.sock") + """ + + def __init__( + self, + model: Optional[str] = None, + *, + base_url: Optional[str] = None, + socket: Optional[str] = None, + api_key: Optional[str] = None, + binary: Optional[str] = None, + host: Optional[str] = None, + port: Optional[int] = None, + timeout: Optional[httpx.Timeout] = None, + startup_timeout: float = 600.0, + transport: Optional[httpx.BaseTransport] = None, + **server_kwargs: Any, + ) -> None: + managed = is_managed(model, base_url, socket, transport) + + self._server: Optional[ManagedServer] = None + self._http_client: Optional[httpx.Client] = None + self._model: Optional[str] = None + self._closed = False + timeout = timeout or DEFAULT_TIMEOUT + + model_override = server_kwargs.pop("model", None) + + if not managed: + resolved_base = self._connect(base_url, socket, transport, timeout) + else: + assert model is not None + resolved_base = self._spawn( + model, binary, host, port, socket, api_key, startup_timeout, timeout, server_kwargs + ) + + self._client = OpenAI( + base_url=resolved_base, + api_key=api_key or "-", + http_client=self._http_client, + ) + self._base_url = resolved_base + self._model = model_override or self._discover_model() + + # -- setup ------------------------------------------------------------- + + def _connect( + self, + base_url: Optional[str], + socket: Optional[str], + transport: Optional[httpx.BaseTransport], + timeout: httpx.Timeout, + ) -> str: + """Build the connect-mode http client (no subprocess) and return the base URL.""" + if transport is not None: + # Injected transport (e.g. httpx.MockTransport in tests). + resolved_base = connect_base_url(base_url) + self._http_client = httpx.Client( + transport=transport, base_url=resolved_base, timeout=timeout + ) + elif socket is not None: + uds_transport = httpx.HTTPTransport(uds=socket) + self._http_client = httpx.Client( + transport=uds_transport, base_url=UDS_BASE, timeout=timeout + ) + resolved_base = f"{UDS_BASE}/v1" + else: + assert base_url is not None + resolved_base = normalize_base_url(base_url) + self._http_client = httpx.Client(base_url=resolved_base, timeout=timeout) + return resolved_base + + def _spawn( + self, + model: str, + binary: Optional[str], + host: Optional[str], + port: Optional[int], + socket: Optional[str], + api_key: Optional[str], + startup_timeout: float, + timeout: httpx.Timeout, + server_kwargs: Dict[str, Any], + ) -> str: + """Spawn and wait for a managed server, build its http client, return the base URL.""" + self._server = ManagedServer( + model, + binary=binary, + host=host, + port=port, + socket_path=socket, + api_key=api_key, + startup_timeout=startup_timeout, + **server_kwargs, + ) + self._server.start() + resolved_base = self._server.base_url + if self._server.uds_path is not None: + uds_transport = httpx.HTTPTransport(uds=self._server.uds_path) + self._http_client = httpx.Client( + transport=uds_transport, base_url=UDS_BASE, timeout=timeout + ) + else: + self._http_client = httpx.Client(base_url=resolved_base, timeout=timeout) + return resolved_base + + def _discover_model(self) -> str: + data = self._client.models.list() + items = list(data.data) + if not items: + raise MlxcelError("Server reported no models from /v1/models.") + return items[0].id + + # -- properties -------------------------------------------------------- + + @property + def model(self) -> str: + """The resolved model id, auto-discovered from ``/v1/models``.""" + assert self._model is not None + return self._model + + @property + def openai_client(self) -> OpenAI: + """The configured ``openai.OpenAI`` instance for the full API surface.""" + return self._client + + # -- generation -------------------------------------------------------- + + def generate(self, prompt: str, **sampling: Any) -> str: + """Generate a completion for ``prompt`` and return the text.""" + self._ensure_alive() + params = build_params(sampling) + resp = self._client.completions.create(model=self.model, prompt=prompt, **params) + return resp.choices[0].text or "" + + def stream(self, prompt: str, **sampling: Any) -> Iterator[str]: + """Stream completion text deltas for ``prompt``.""" + self._ensure_alive() + params = build_params(sampling) + stream = cast( + "Stream[Completion]", + self._client.completions.create(model=self.model, prompt=prompt, stream=True, **params), + ) + for chunk in stream: + if chunk.choices and chunk.choices[0].text: + yield chunk.choices[0].text + + def chat(self, messages: ChatMessages, **sampling: Any) -> str: + """Run a chat completion and return the assistant message content.""" + self._ensure_alive() + params = build_params(sampling) + resp = self._client.chat.completions.create( + model=self.model, messages=as_messages(messages), **params + ) + return resp.choices[0].message.content or "" + + def chat_stream(self, messages: ChatMessages, **sampling: Any) -> Iterator[str]: + """Stream chat completion content deltas.""" + self._ensure_alive() + params = build_params(sampling) + stream = cast( + "Stream[ChatCompletionChunk]", + self._client.chat.completions.create( + model=self.model, messages=as_messages(messages), stream=True, **params + ), + ) + for chunk in stream: + if chunk.choices and chunk.choices[0].delta.content: + yield chunk.choices[0].delta.content + + # -- model / tokenizer helpers ---------------------------------------- + + def models(self) -> List[str]: + """List model ids advertised by the server.""" + self._ensure_alive() + return [m.id for m in self._client.models.list().data] + + def tokenize(self, text: str, add_special: bool = False) -> List[int]: + """Tokenize ``text`` via the server's native ``/tokenize`` route.""" + self._ensure_alive() + resp = self._raw_post("/tokenize", {"content": text, "add_special": add_special}) + return list(resp["tokens"]) + + def detokenize(self, tokens: Iterable[int]) -> str: + """Decode token ids back to text via the native ``/detokenize`` route.""" + self._ensure_alive() + resp = self._raw_post("/detokenize", {"tokens": list(tokens)}) + return str(resp["content"]) + + def _raw_post(self, path: str, json: Dict[str, Any]) -> Dict[str, Any]: + assert self._http_client is not None + url = native_base_url(self._base_url) + path + response = self._http_client.post(url, json=json) + response.raise_for_status() + result: Dict[str, Any] = response.json() + return result + + # -- lifecycle --------------------------------------------------------- + + def _ensure_alive(self) -> None: + if self._server is not None: + self._server.ensure_alive() + + def close(self) -> None: + """Close the HTTP client and stop the managed server, if any.""" + if self._closed: + return + self._closed = True + if self._http_client is not None: + self._http_client.close() + if self._server is not None: + self._server.close() + + def __enter__(self) -> "LLM": + return self + + def __exit__(self, *exc: object) -> None: + self.close() + + def __del__(self) -> None: + try: + self.close() + except Exception: + pass + + +__all__ = ["LLM"] diff --git a/python/src/mlxcel/_common.py b/python/src/mlxcel/_common.py new file mode 100644 index 00000000..3c523524 --- /dev/null +++ b/python/src/mlxcel/_common.py @@ -0,0 +1,107 @@ +"""Shared helpers for the sync and async clients. + +Mode selection, base-URL normalization, message-type narrowing, and the +transport constants live here so :mod:`._client` (sync ``LLM``) and +:mod:`._async_client` (``AsyncLLM``) stay small and share one source of truth. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Iterable, List, Mapping, Optional, cast + +import httpx + +from .errors import MlxcelError + +if TYPE_CHECKING: + from openai.types.chat import ChatCompletionMessageParam + +# Placeholder host used for Unix-socket transports; routing is via the uds +# transport, so the host name is never resolved over the network. +UDS_BASE = "http://mlxcel" + +# Default timeouts applied when a custom http_client is injected, because the +# SDK's built-in timeouts are dropped in that case. Generous read timeout so a +# slow first token does not abort a long generation. +DEFAULT_TIMEOUT = httpx.Timeout(connect=10.0, read=600.0, write=30.0, pool=10.0) + +ChatMessages = Iterable[Mapping[str, Any]] + + +def is_managed( + model: Optional[str], + base_url: Optional[str], + socket: Optional[str], + transport: Optional[object] = None, +) -> bool: + """Decide managed vs connect mode and validate the argument combination. + + A ``model`` means managed mode (spawn a server); ``socket=`` may accompany it + as the bind path. A connect target (``base_url=`` or ``socket=`` without a + model, or an injected ``transport=``) means connect mode (no subprocess). + + Returns: + True for managed mode, False for connect mode. + + Raises: + MlxcelError: if the arguments are ambiguous or insufficient. + """ + if model is not None: + # Managed mode. A base_url or injected transport would be contradictory. + if base_url is not None or transport is not None: + raise MlxcelError( + "Pass either a model (managed mode) or a connect target " + "(base_url=/transport=), not both. In managed mode, socket= is " + "the bind path for the spawned server." + ) + return True + + # No model: connect mode, which needs a target. + if base_url is not None and socket is not None: + raise MlxcelError("Pass either base_url= or socket=, not both.") + if base_url is None and socket is None and transport is None: + raise MlxcelError( + "A model is required in managed mode. Pass model=, or use connect " + "mode with base_url=, socket=, or transport=." + ) + return False + + +def normalize_base_url(base_url: str) -> str: + """Ensure the OpenAI base URL ends with ``/v1``.""" + trimmed = base_url.rstrip("/") + if trimmed.endswith("/v1"): + return trimmed + return f"{trimmed}/v1" + + +def native_base_url(openai_base_url: str) -> str: + """Derive the server root (no ``/v1``) for native routes like ``/tokenize``.""" + return openai_base_url[: -len("/v1")] if openai_base_url.endswith("/v1") else openai_base_url + + +def connect_base_url(base_url: Optional[str]) -> str: + """Resolved OpenAI base URL for a connect-mode client (UDS placeholder otherwise).""" + return normalize_base_url(base_url) if base_url else f"{UDS_BASE}/v1" + + +def as_messages(messages: ChatMessages) -> "List[ChatCompletionMessageParam]": + """Materialize caller messages as the OpenAI message-param list type. + + The public API accepts loosely-typed mappings (plain dicts) for ergonomics; + the OpenAI SDK wants its precise TypedDict union. The shapes are compatible + at runtime, so this only narrows the static type. + """ + return cast("List[ChatCompletionMessageParam]", list(messages)) + + +__all__ = [ + "UDS_BASE", + "DEFAULT_TIMEOUT", + "ChatMessages", + "is_managed", + "normalize_base_url", + "native_base_url", + "connect_base_url", + "as_messages", +] diff --git a/python/src/mlxcel/_sampling.py b/python/src/mlxcel/_sampling.py new file mode 100644 index 00000000..0b5b054d --- /dev/null +++ b/python/src/mlxcel/_sampling.py @@ -0,0 +1,88 @@ +"""Map Python keyword arguments to OpenAI request fields. + +The OpenAI Python SDK accepts a fixed set of sampling parameters on +``completions.create`` and ``chat.completions.create``. mlxcel and the +underlying llama-server expose a few extra knobs (``top_k``, ``min_p``, +``repetition_penalty``, DRY sampling) that are not part of the OpenAI schema; +those are forwarded through the SDK's ``extra_body`` escape hatch so the server +receives them verbatim in the JSON request body. +""" + +from __future__ import annotations + +from typing import Any, Dict, Tuple + +# Sampling kwargs that map directly onto native OpenAI request fields. +_OPENAI_FIELDS: Tuple[str, ...] = ( + "max_tokens", + "temperature", + "top_p", + "stop", + "seed", + "presence_penalty", + "frequency_penalty", + "logit_bias", + "n", + "logprobs", + "top_logprobs", + "response_format", +) + +# Server-specific sampling knobs forwarded via ``extra_body``. These are not in +# the OpenAI schema; the server reads them from the raw request body. +_EXTRA_BODY_FIELDS: Tuple[str, ...] = ( + "top_k", + "min_p", + "repetition_penalty", + "repeat_penalty", + "dry_multiplier", + "dry_base", + "dry_allowed_length", + "dry_penalty_last_n", +) + + +def build_params(kwargs: Dict[str, Any]) -> Dict[str, Any]: + """Split caller kwargs into OpenAI request fields plus an ``extra_body`` payload. + + Known OpenAI fields are passed through as top-level request parameters. + Recognized server-specific knobs are collected into ``extra_body``. An + explicit ``extra_body=`` mapping supplied by the caller is merged in and + takes precedence, so advanced users can send arbitrary server fields. Any + remaining unknown kwargs are also routed into ``extra_body`` rather than + silently dropped, since the OpenAI SDK rejects unexpected top-level keys. + + Args: + kwargs: Raw keyword arguments from a generate/chat call. + + Returns: + A new dict suitable for splatting into the OpenAI ``create`` call. + """ + params: Dict[str, Any] = {} + extra_body: Dict[str, Any] = {} + + # A caller-provided extra_body is applied last so it wins on conflicts. + caller_extra_body = kwargs.pop("extra_body", None) + + for key, value in kwargs.items(): + if value is None: + continue + if key in _OPENAI_FIELDS: + params[key] = value + else: + # Recognized server knob or an unknown key: send it in the body so + # the server can interpret it instead of the SDK rejecting it. + extra_body[key] = value + + if caller_extra_body: + if not isinstance(caller_extra_body, dict): + raise TypeError("extra_body must be a mapping") + extra_body.update(caller_extra_body) + + if extra_body: + params["extra_body"] = extra_body + + return params + + +__all__ = ["build_params"] diff --git a/python/src/mlxcel/_server.py b/python/src/mlxcel/_server.py new file mode 100644 index 00000000..6196073c --- /dev/null +++ b/python/src/mlxcel/_server.py @@ -0,0 +1,378 @@ +"""Lifecycle management for a local ``mlxcel serve`` subprocess. + +``ManagedServer`` discovers the binary, picks a transport (Unix domain socket on +POSIX, TCP elsewhere or on request), spawns the server, waits until ``/health`` +reports ready, forwards the server's stderr to the ``mlxcel.server`` logger, and +shuts the process down cleanly. The server loads and warms up the model *before* +it binds its listener, so a first-run weight download can keep the socket +unavailable for minutes; readiness is therefore driven by polling ``/health`` +with a generous timeout while watching the child process for an early exit. +""" + +from __future__ import annotations + +import atexit +import logging +import os +import secrets +import shutil +import socket +import subprocess +import sys +import threading +import time +from collections import deque +from typing import Deque, List, Optional + +import httpx + +from .errors import MlxcelServerError, MlxcelTimeoutError + +logger = logging.getLogger("mlxcel.server") + +# Conservative Unix domain socket path limit. The kernel struct sun_path is +# 104 bytes on macOS and 108 on Linux; 100 leaves headroom for the trailing NUL +# and keeps a single rule across platforms. +_MAX_SOCKET_PATH = 100 + +# Number of trailing stderr lines retained for crash/error reporting. +_STDERR_TAIL_LINES = 50 + +_IS_WINDOWS = sys.platform.startswith("win") + + +def _find_binary(binary: Optional[str]) -> str: + """Resolve the ``mlxcel`` executable. + + Order: explicit ``binary`` argument, then the ``MLXCEL_BIN`` environment + variable, then ``mlxcel`` on ``PATH``. + + Raises: + MlxcelServerError: if no usable binary is found. + """ + candidate = binary or os.environ.get("MLXCEL_BIN") + if candidate: + resolved = shutil.which(candidate) or (candidate if os.path.isfile(candidate) else None) + if resolved: + return resolved + raise MlxcelServerError( + f"mlxcel binary not found at {candidate!r}. " + "Pass binary=... or set MLXCEL_BIN to the built executable." + ) + + found = shutil.which("mlxcel") + if found: + return found + + raise MlxcelServerError( + "Could not find the 'mlxcel' executable. Install it (e.g. via Homebrew: " + "'brew install lablup/tap/mlxcel'), build it from source with " + "'make release' and add target/release to PATH, pass binary=..., or set " + "the MLXCEL_BIN environment variable." + ) + + +def _free_tcp_port() -> int: + """Bind an ephemeral TCP port, then release it so the child can claim it. + + There is an unavoidable race between releasing the port and the child + rebinding it, but it is small for a local single-launch flow. + """ + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + return int(sock.getsockname()[1]) + + +def _default_socket_path() -> str: + """A short, unique Unix socket path under ``/tmp``. + + ``/tmp`` is used deliberately rather than ``tempfile.gettempdir()``: on + macOS the latter resolves to a long ``/var/folders/...`` path that can blow + past the sun_path limit. + """ + return f"/tmp/mlxcel-{os.getpid()}-{secrets.token_hex(3)}.sock" + + +class ManagedServer: + """Spawn and supervise a local ``mlxcel serve`` process. + + Either a Unix domain socket (default on POSIX) or TCP is used. After + construction, call :meth:`start` to launch and block until ready, then read + :attr:`base_url`, :attr:`uds_path`, and :attr:`api_key` to build a client. + """ + + def __init__( + self, + model: str, + *, + binary: Optional[str] = None, + host: Optional[str] = None, + port: Optional[int] = None, + socket_path: Optional[str] = None, + api_key: Optional[str] = None, + ctx_size: Optional[int] = None, + n_predict: Optional[int] = None, + alias: Optional[str] = None, + warmup: Optional[bool] = None, + extra_args: Optional[List[str]] = None, + startup_timeout: float = 600.0, + shutdown_grace: float = 5.0, + ) -> None: + self.model = model + self._binary_arg = binary + self.api_key = api_key + self._ctx_size = ctx_size + self._n_predict = n_predict + self._alias = alias + self._warmup = warmup + self._extra_args = list(extra_args or []) + self._startup_timeout = startup_timeout + self._shutdown_grace = shutdown_grace + + self._proc: Optional[subprocess.Popen[str]] = None + self._log_thread: Optional[threading.Thread] = None + self._stderr_tail: Deque[str] = deque(maxlen=_STDERR_TAIL_LINES) + self._closed = False + self._atexit_registered = False + + # Transport resolution. TCP is forced on Windows or when host/port given. + self.uds_path: Optional[str] = None + self.host: Optional[str] = None + self.port: Optional[int] = None + + use_tcp = _IS_WINDOWS or host is not None or port is not None + if use_tcp: + self.host = host or "127.0.0.1" + self.port = port if port is not None else _free_tcp_port() + else: + path = socket_path or _default_socket_path() + if len(path) > _MAX_SOCKET_PATH: + raise MlxcelServerError( + f"Unix socket path is too long ({len(path)} > {_MAX_SOCKET_PATH} bytes): " + f"{path!r}. Pass a shorter socket=... path (e.g. under /tmp)." + ) + self.uds_path = path + + @property + def base_url(self) -> str: + """OpenAI base URL (with the ``/v1`` suffix) for this transport. + + For a Unix socket the host portion is a placeholder; the actual routing + happens through the httpx ``uds`` transport. + """ + if self.uds_path is not None: + return "http://mlxcel/v1" + return f"http://{self.host}:{self.port}/v1" + + @property + def _health_url(self) -> str: + if self.uds_path is not None: + return "http://mlxcel/health" + return f"http://{self.host}:{self.port}/health" + + def _build_command(self) -> List[str]: + binary = _find_binary(self._binary_arg) + cmd: List[str] = [binary, "serve", "-m", self.model] + + if self.uds_path is not None: + # UDS mode: --port 0 and --host reinterpreted as the socket path. + cmd += ["--host", self.uds_path, "--port", "0"] + else: + assert self.host is not None and self.port is not None + cmd += ["--host", self.host, "--port", str(self.port)] + + if self.api_key: + cmd += ["--api-key", self.api_key] + if self._ctx_size is not None: + cmd += ["--ctx-size", str(self._ctx_size)] + if self._n_predict is not None: + cmd += ["--n-predict", str(self._n_predict)] + if self._alias is not None: + cmd += ["-a", self._alias] + if self._warmup is True: + cmd += ["--warmup"] + elif self._warmup is False: + cmd += ["--no-warmup"] + cmd += self._extra_args + return cmd + + def start(self) -> None: + """Spawn the server and block until ``/health`` returns HTTP 200. + + Raises: + MlxcelServerError: if the binary cannot be found, the process exits + before becoming ready, or another launch failure occurs. + MlxcelTimeoutError: if readiness is not reached within + ``startup_timeout``. + """ + if self._proc is not None: + raise MlxcelServerError("Server already started.") + + cmd = self._build_command() + logger.debug("Launching mlxcel server: %s", " ".join(cmd)) + + # Remove a stale socket file so bind does not fail on EADDRINUSE. + if self.uds_path is not None and os.path.exists(self.uds_path): + try: + os.unlink(self.uds_path) + except OSError: + pass + + try: + self._proc = subprocess.Popen( # noqa: S603 - command built from trusted args + cmd, + stdout=subprocess.DEVNULL, + stderr=subprocess.PIPE, + text=True, + bufsize=1, + ) + except OSError as exc: + raise MlxcelServerError(f"Failed to launch mlxcel server: {exc}") from exc + + self._start_log_forwarding() + self._register_atexit() + + try: + self._wait_until_ready() + except BaseException: + # Any failure (including KeyboardInterrupt) during startup must not + # leak the child process. + self.close() + raise + + def _start_log_forwarding(self) -> None: + proc = self._proc + assert proc is not None and proc.stderr is not None + stderr = proc.stderr + + def _pump() -> None: + try: + for raw in stderr: + line = raw.rstrip("\n") + self._stderr_tail.append(line) + logger.info("%s", line) + except (ValueError, OSError): + # Pipe closed during shutdown; expected, never raise. + return + + thread = threading.Thread(target=_pump, name="mlxcel-server-logs", daemon=True) + thread.start() + self._log_thread = thread + + def _stderr_snapshot(self) -> str: + return "\n".join(self._stderr_tail) + + def _make_probe_client(self) -> httpx.Client: + if self.uds_path is not None: + transport = httpx.HTTPTransport(uds=self.uds_path) + return httpx.Client(transport=transport, timeout=5.0) + return httpx.Client(timeout=5.0) + + def _wait_until_ready(self) -> None: + proc = self._proc + assert proc is not None + deadline = time.monotonic() + self._startup_timeout + backoff = 0.05 + url = self._health_url + + with self._make_probe_client() as client: + while True: + exit_code = proc.poll() + if exit_code is not None: + raise MlxcelServerError( + f"mlxcel server exited with code {exit_code} before becoming ready.", + stderr=self._stderr_snapshot(), + ) + + ready = self._probe_once(client, url) + if ready: + return + + if time.monotonic() >= deadline: + raise MlxcelTimeoutError( + f"mlxcel server did not become ready within {self._startup_timeout:.0f}s.", + stderr=self._stderr_snapshot(), + ) + + time.sleep(backoff) + backoff = min(backoff * 1.5, 1.0) + + @staticmethod + def _probe_once(client: httpx.Client, url: str) -> bool: + """Return True once ``/health`` answers HTTP 200. + + Connection errors (socket file not yet present, connection refused) and + HTTP 503 (loading / no slot) are treated as "still starting". + """ + try: + resp = client.get(url) + except (httpx.ConnectError, httpx.ConnectTimeout, httpx.ReadError): + return False + if resp.status_code == 200: + return True + if resp.status_code == 503: + return False + # Any other status is unexpected for /health; keep polling rather than + # failing hard, the next iteration's process-liveness check will catch a + # dead server. + return False + + def ensure_alive(self) -> None: + """Raise if the supervised process has died. + + Call before delegating work so a crashed server surfaces a clear + ``MlxcelServerError`` (with stderr context) instead of an opaque + connection error. + """ + proc = self._proc + if proc is None: + raise MlxcelServerError("Server is not running.") + exit_code = proc.poll() + if exit_code is not None: + raise MlxcelServerError( + f"mlxcel server has exited with code {exit_code}.", + stderr=self._stderr_snapshot(), + ) + + def close(self) -> None: + """Terminate the process (SIGTERM, then SIGKILL) and clean up the socket.""" + if self._closed: + return + self._closed = True + + proc = self._proc + if proc is not None and proc.poll() is None: + proc.terminate() + try: + proc.wait(timeout=self._shutdown_grace) + except subprocess.TimeoutExpired: + proc.kill() + try: + proc.wait(timeout=self._shutdown_grace) + except subprocess.TimeoutExpired: + logger.warning("mlxcel server did not exit after SIGKILL.") + + if self._log_thread is not None and self._log_thread.is_alive(): + self._log_thread.join(timeout=1.0) + + if self.uds_path is not None and os.path.exists(self.uds_path): + try: + os.unlink(self.uds_path) + except OSError: + pass + + def _register_atexit(self) -> None: + if not self._atexit_registered: + atexit.register(self.close) + self._atexit_registered = True + + def __del__(self) -> None: + # Best-effort cleanup for a leaked handle. Interpreter shutdown may have + # already torn down modules, so swallow everything. + try: + self.close() + except Exception: + pass + + +__all__ = ["ManagedServer"] diff --git a/python/src/mlxcel/errors.py b/python/src/mlxcel/errors.py new file mode 100644 index 00000000..7bd26c2b --- /dev/null +++ b/python/src/mlxcel/errors.py @@ -0,0 +1,34 @@ +"""Exception types raised by the mlxcel client. + +Lifecycle concerns (binary discovery, server startup, readiness, crash) raise +the ``Mlxcel*`` exceptions defined here. HTTP and API errors from the server +propagate as native ``openai`` SDK exceptions so callers keep full visibility +into status codes and response bodies. +""" + +from __future__ import annotations + + +class MlxcelError(Exception): + """Base class for all mlxcel client errors.""" + + +class MlxcelServerError(MlxcelError): + """A managed ``mlxcel serve`` process failed to launch, become ready, or stay alive. + + When the failure was observed against a running child process, ``stderr`` + holds the tail of the captured server log to aid debugging. + """ + + def __init__(self, message: str, stderr: str | None = None) -> None: + self.stderr = stderr + if stderr: + message = f"{message}\n\n--- mlxcel server stderr (tail) ---\n{stderr}" + super().__init__(message) + + +class MlxcelTimeoutError(MlxcelServerError): + """The managed server did not become ready within the configured timeout.""" + + +__all__ = ["MlxcelError", "MlxcelServerError", "MlxcelTimeoutError"] diff --git a/python/src/mlxcel/py.typed b/python/src/mlxcel/py.typed new file mode 100644 index 00000000..e69de29b diff --git a/python/tests/fake_server.py b/python/tests/fake_server.py new file mode 100644 index 00000000..863ac88a --- /dev/null +++ b/python/tests/fake_server.py @@ -0,0 +1,279 @@ +"""A tiny stdlib-only fake of the mlxcel HTTP server for lifecycle tests. + +Binds either a Unix domain socket or a TCP port and serves the subset of routes +the client touches: ``/health`` (HTTP 503 for a short warmup window, then 200), +``/v1/models``, ``/v1/completions`` (plain and SSE), ``/v1/chat/completions`` +(plain and SSE), ``/tokenize``, and ``/detokenize``. Responses are canned. + +Run as a script for the lifecycle tests:: + + python fake_server.py --uds /tmp/x.sock [--ready-after 0.3] [--model alias] + python fake_server.py --host 127.0.0.1 --port 8080 + python fake_server.py --uds /tmp/x.sock --crash # exit before binding + +Uses only the standard library so it can stand in for the real binary in CI. +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +import time +from http.server import BaseHTTPRequestHandler +from socketserver import TCPServer, ThreadingMixIn, UnixStreamServer +from typing import Any, Dict + +MODEL_ID = "fake-model" +START_TIME = time.monotonic() +READY_AFTER = 0.0 + + +def _now() -> int: + return int(time.time()) + + +class Handler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def log_message(self, *args: Any) -> None: # noqa: D401 - silence default logging + return + + def _send_json(self, status: int, payload: Dict[str, Any]) -> None: + body = json.dumps(payload).encode("utf-8") + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def _read_body(self) -> Dict[str, Any]: + length = int(self.headers.get("Content-Length", 0)) + if length == 0: + return {} + raw = self.rfile.read(length) + try: + parsed: Dict[str, Any] = json.loads(raw.decode("utf-8")) + return parsed + except (json.JSONDecodeError, UnicodeDecodeError): + return {} + + def _is_ready(self) -> bool: + return (time.monotonic() - START_TIME) >= READY_AFTER + + # -- routing ----------------------------------------------------------- + + def do_GET(self) -> None: # noqa: N802 - http.server API + if self.path == "/health": + if not self._is_ready(): + self._send_json(503, {"status": "loading model"}) + return + self._send_json( + 200, + {"status": "ok", "model": MODEL_ID, "context_size": 0, "tool_call_parser": None}, + ) + return + if self.path == "/v1/models": + self._send_json( + 200, + { + "object": "list", + "data": [ + { + "id": MODEL_ID, + "object": "model", + "created": _now(), + "owned_by": "user", + } + ], + }, + ) + return + self._send_json(404, {"error": "not found"}) + + def do_POST(self) -> None: # noqa: N802 - http.server API + body = self._read_body() + stream = bool(body.get("stream")) + + if self.path == "/v1/completions": + if stream: + self._stream_completions() + else: + self._send_json(200, self._completion_payload()) + return + if self.path == "/v1/chat/completions": + if stream: + self._stream_chat() + else: + self._send_json(200, self._chat_payload()) + return + if self.path == "/tokenize": + content = str(body.get("content", "")) + tokens = [ord(c) % 256 for c in content] + self._send_json(200, {"tokens": tokens}) + return + if self.path == "/detokenize": + tokens = body.get("tokens", []) + content = "".join(chr(int(t)) for t in tokens) + self._send_json(200, {"content": content}) + return + self._send_json(404, {"error": "not found"}) + + # -- canned payloads --------------------------------------------------- + + def _completion_payload(self) -> Dict[str, Any]: + return { + "id": "cmpl-fake", + "object": "text_completion", + "created": _now(), + "model": MODEL_ID, + "choices": [ + {"index": 0, "text": "hello world", "finish_reason": "stop", "logprobs": None} + ], + "usage": {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3}, + } + + def _chat_payload(self) -> Dict[str, Any]: + return { + "id": "chatcmpl-fake", + "object": "chat.completion", + "created": _now(), + "model": MODEL_ID, + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "hi there"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3}, + } + + def _write_sse(self, chunks: list[Dict[str, Any]]) -> None: + self.send_response(200) + self.send_header("Content-Type", "text/event-stream") + self.send_header("Cache-Control", "no-cache") + self.send_header("Connection", "close") + self.end_headers() + for chunk in chunks: + self.wfile.write(f"data: {json.dumps(chunk)}\n\n".encode()) + self.wfile.write(b"data: [DONE]\n\n") + self.wfile.flush() + + def _stream_completions(self) -> None: + base = { + "id": "cmpl-fake", + "object": "text_completion", + "created": _now(), + "model": MODEL_ID, + } + chunks = [ + {**base, "choices": [{"index": 0, "text": "hello ", "finish_reason": None}]}, + {**base, "choices": [{"index": 0, "text": "world", "finish_reason": "stop"}]}, + ] + self._write_sse(chunks) + + def _stream_chat(self) -> None: + base = { + "id": "chatcmpl-fake", + "object": "chat.completion.chunk", + "created": _now(), + "model": MODEL_ID, + } + chunks = [ + {**base, "choices": [{"index": 0, "delta": {"role": "assistant"}}]}, + {**base, "choices": [{"index": 0, "delta": {"content": "hi "}}]}, + { + **base, + "choices": [{"index": 0, "delta": {"content": "there"}, "finish_reason": "stop"}], + }, + ] + self._write_sse(chunks) + + +class ThreadingUnixServer(ThreadingMixIn, UnixStreamServer): + daemon_threads = True + allow_reuse_address = True + + +class ThreadingTCPServer(ThreadingMixIn, TCPServer): + daemon_threads = True + allow_reuse_address = True + + +def main() -> None: + global READY_AFTER, MODEL_ID + + parser = argparse.ArgumentParser() + parser.add_argument("--uds") + parser.add_argument("--host") + parser.add_argument("--port", type=int) + parser.add_argument("--ready-after", type=float, default=0.0) + parser.add_argument("--model", default="fake-model") + parser.add_argument("--crash", action="store_true", help="exit before binding") + # Accept and ignore the flags the client always passes so this can stand in + # for `mlxcel serve` invoked by ManagedServer. + parser.add_argument("serve", nargs="?") + parser.add_argument("-m", "--model-path") + parser.add_argument("--api-key") + parser.add_argument("--ctx-size") + parser.add_argument("--n-predict") + parser.add_argument("-a", "--alias") + parser.add_argument("--warmup", action="store_true") + parser.add_argument("--no-warmup", action="store_true") + args, _unknown = parser.parse_known_args() + + READY_AFTER = args.ready_after + MODEL_ID = args.alias or args.model + + if args.crash: + # Emit some stderr so the lifecycle test can assert it is captured. + sys.stderr.write("fatal: simulated startup failure\n") + sys.stderr.flush() + sys.exit(3) + + # Honor UDS mode signalled either by --uds or by `--port 0 --host ` + # (the way ManagedServer invokes the real binary). + uds = args.uds + if uds is None and args.port == 0 and args.host: + uds = args.host + + # Simulate the real server's startup ordering: warm up (sleep) before bind. + if READY_AFTER > 0: + time.sleep(0) # binding happens immediately; /health gates readiness. + + if uds: + if os.path.exists(uds): + os.unlink(uds) + server: Any = ThreadingUnixServer(uds, Handler) + sys.stderr.write(f"Starting mlxcel server on unix:{uds}\n") + else: + host = args.host or "127.0.0.1" + port = args.port or 8080 + server = ThreadingTCPServer((host, port), Handler) + sys.stderr.write(f"Starting mlxcel server on {host}:{port}\n") + sys.stderr.flush() + + try: + server.serve_forever(poll_interval=0.05) + except KeyboardInterrupt: + pass + finally: + server.server_close() + if uds and os.path.exists(uds): + try: + os.unlink(uds) + except OSError: + pass + + +if __name__ == "__main__": + # Reduce the chance of a noisy traceback on broken-pipe during shutdown. + try: + main() + except (BrokenPipeError, ConnectionResetError): + pass + except OSError as exc: + sys.stderr.write(f"fake_server bind error: {exc}\n") + sys.exit(4) diff --git a/python/tests/test_client_mock.py b/python/tests/test_client_mock.py new file mode 100644 index 00000000..8558dba7 --- /dev/null +++ b/python/tests/test_client_mock.py @@ -0,0 +1,293 @@ +"""Unit tests for LLM / AsyncLLM using httpx.MockTransport (no real server).""" + +from __future__ import annotations + +import json +from typing import Any, Dict, List + +import httpx +import pytest + +import mlxcel +from mlxcel._sampling import build_params + +MODEL_ID = "mock-model" + +# Captured request bodies, keyed by path, so tests can assert sampling mapping. +_LAST_BODY: Dict[str, Dict[str, Any]] = {} + + +def _models_response() -> httpx.Response: + return httpx.Response( + 200, + json={ + "object": "list", + "data": [{"id": MODEL_ID, "object": "model", "created": 1, "owned_by": "user"}], + }, + ) + + +def _completion_response() -> httpx.Response: + return httpx.Response( + 200, + json={ + "id": "cmpl-1", + "object": "text_completion", + "created": 1, + "model": MODEL_ID, + "choices": [{"index": 0, "text": "generated text", "finish_reason": "stop"}], + }, + ) + + +def _chat_response() -> httpx.Response: + return httpx.Response( + 200, + json={ + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 1, + "model": MODEL_ID, + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "chat reply"}, + "finish_reason": "stop", + } + ], + }, + ) + + +def _sse(chunks: List[Dict[str, Any]]) -> httpx.Response: + body = "".join(f"data: {json.dumps(c)}\n\n" for c in chunks) + "data: [DONE]\n\n" + return httpx.Response( + 200, content=body.encode("utf-8"), headers={"content-type": "text/event-stream"} + ) + + +def _handler(request: httpx.Request) -> httpx.Response: + path = request.url.path + body: Dict[str, Any] = {} + if request.content: + try: + body = json.loads(request.content) + except json.JSONDecodeError: + body = {} + _LAST_BODY[path] = body + stream = bool(body.get("stream")) + + if path == "/v1/models": + return _models_response() + if path == "/v1/completions": + if stream: + return _sse( + [ + {"choices": [{"index": 0, "text": "gen ", "finish_reason": None}]}, + {"choices": [{"index": 0, "text": "text", "finish_reason": "stop"}]}, + ] + ) + return _completion_response() + if path == "/v1/chat/completions": + if stream: + return _sse( + [ + {"choices": [{"index": 0, "delta": {"content": "chat "}}]}, + { + "choices": [ + {"index": 0, "delta": {"content": "reply"}, "finish_reason": "stop"} + ] + }, + ] + ) + return _chat_response() + if path == "/tokenize": + return httpx.Response(200, json={"tokens": [1, 2, 3]}) + if path == "/detokenize": + return httpx.Response(200, json={"content": "decoded"}) + return httpx.Response(404, json={"error": "not found"}) + + +def _error_handler(request: httpx.Request) -> httpx.Response: + """Like ``_handler`` but returns HTTP 500 on chat completions.""" + if request.url.path == "/v1/models": + return _models_response() + if request.url.path == "/v1/chat/completions": + return httpx.Response(500, json={"error": {"message": "boom"}}) + return httpx.Response(404, json={"error": "not found"}) + + +@pytest.fixture() +def llm() -> Any: + _LAST_BODY.clear() + transport = httpx.MockTransport(_handler) + client = mlxcel.LLM(transport=transport) + yield client + client.close() + + +def test_model_auto_discovery(llm: Any) -> None: + assert llm.model == MODEL_ID + + +def test_generate(llm: Any) -> None: + assert llm.generate("hello") == "generated text" + + +def test_generate_passes_sampling(llm: Any) -> None: + llm.generate("hello", max_tokens=42, temperature=0.5) + body = _LAST_BODY["/v1/completions"] + assert body["max_tokens"] == 42 + assert body["temperature"] == 0.5 + assert body["model"] == MODEL_ID + + +def test_generate_extra_body_passthrough(llm: Any) -> None: + llm.generate("hello", top_k=20, min_p=0.05, repetition_penalty=1.1) + body = _LAST_BODY["/v1/completions"] + # Server-specific knobs land in the request body (extra_body merged in). + assert body["top_k"] == 20 + assert body["min_p"] == 0.05 + assert body["repetition_penalty"] == 1.1 + + +def test_stream(llm: Any) -> None: + assert "".join(llm.stream("hello")) == "gen text" + + +def test_chat(llm: Any) -> None: + assert llm.chat([{"role": "user", "content": "hi"}]) == "chat reply" + + +def test_chat_stream(llm: Any) -> None: + out = "".join(llm.chat_stream([{"role": "user", "content": "hi"}])) + assert out == "chat reply" + + +def test_models(llm: Any) -> None: + assert llm.models() == [MODEL_ID] + + +def test_tokenize(llm: Any) -> None: + assert llm.tokenize("hi", add_special=True) == [1, 2, 3] + assert _LAST_BODY["/tokenize"] == {"content": "hi", "add_special": True} + + +def test_detokenize(llm: Any) -> None: + assert llm.detokenize([1, 2, 3]) == "decoded" + assert _LAST_BODY["/detokenize"] == {"tokens": [1, 2, 3]} + + +def test_response_format_passthrough(llm: Any) -> None: + schema = {"type": "json_object"} + llm.chat([{"role": "user", "content": "hi"}], response_format=schema) + body = _LAST_BODY["/v1/chat/completions"] + assert body["response_format"] == schema + + +def test_openai_client_escape_hatch(llm: Any) -> None: + from openai import OpenAI + + assert isinstance(llm.openai_client, OpenAI) + resp = llm.openai_client.completions.create(model=llm.model, prompt="x") + assert resp.choices[0].text == "generated text" + + +def test_http_error_propagates_as_openai_exception() -> None: + import openai + + transport = httpx.MockTransport(_error_handler) + client = mlxcel.LLM(transport=transport) + try: + with pytest.raises(openai.APIStatusError): + client.chat([{"role": "user", "content": "hi"}]) + finally: + client.close() + + +# -- mode-selection / validation -------------------------------------------- + + +def test_both_model_and_connect_target_is_error() -> None: + with pytest.raises(mlxcel.MlxcelError): + mlxcel.LLM("some-model", base_url="http://localhost:8080/v1") + + +def test_no_args_is_error() -> None: + with pytest.raises(mlxcel.MlxcelError): + mlxcel.LLM() + + +def test_base_url_and_socket_is_error() -> None: + with pytest.raises(mlxcel.MlxcelError): + mlxcel.LLM(base_url="http://x/v1", socket="/tmp/x.sock") + + +# -- sampling unit ---------------------------------------------------------- + + +def test_build_params_splits_extra_body() -> None: + params = build_params({"max_tokens": 10, "top_k": 5, "unknown_knob": "v"}) + assert params["max_tokens"] == 10 + assert params["extra_body"]["top_k"] == 5 + assert params["extra_body"]["unknown_knob"] == "v" + + +def test_build_params_caller_extra_body_wins() -> None: + params = build_params({"top_k": 5, "extra_body": {"top_k": 99}}) + assert params["extra_body"]["top_k"] == 99 + + +def test_build_params_drops_none() -> None: + params = build_params({"max_tokens": None, "temperature": 0.0}) + assert "max_tokens" not in params + assert params["temperature"] == 0.0 + + +# -- async ------------------------------------------------------------------ + + +def _run(coro: Any) -> Any: + import asyncio + + return asyncio.run(coro) + + +def test_async_generate_and_discovery() -> None: + _LAST_BODY.clear() + transport = httpx.MockTransport(_handler) + + async def go() -> tuple[str, str, str]: + client = mlxcel.AsyncLLM(transport=transport) + try: + text = await client.generate("hello", max_tokens=7) + chat = await client.chat([{"role": "user", "content": "hi"}]) + return text, chat, client.model + finally: + await client.close() + + text, chat, model = _run(go()) + assert text == "generated text" + assert chat == "chat reply" + assert model == MODEL_ID + assert _LAST_BODY["/v1/completions"]["max_tokens"] == 7 + + +def test_async_stream_and_tokenize() -> None: + _LAST_BODY.clear() + transport = httpx.MockTransport(_handler) + + async def go() -> tuple[str, list[int], str]: + client = mlxcel.AsyncLLM(transport=transport) + try: + deltas = [d async for d in client.stream("hi")] + tokens = await client.tokenize("hi") + decoded = await client.detokenize([1, 2, 3]) + return "".join(deltas), tokens, decoded + finally: + await client.close() + + streamed, tokens, decoded = _run(go()) + assert streamed == "gen text" + assert tokens == [1, 2, 3] + assert decoded == "decoded" diff --git a/python/tests/test_e2e.py b/python/tests/test_e2e.py new file mode 100644 index 00000000..4f15832a --- /dev/null +++ b/python/tests/test_e2e.py @@ -0,0 +1,46 @@ +"""End-to-end test against a real ``mlxcel serve`` process. + +Skipped unless ``MLXCEL_BIN`` points at a built mlxcel binary. Set +``MLXCEL_E2E_MODEL`` to choose the model (defaults to a small 4-bit checkpoint). +The first run may download weights from Hugging Face, so the startup timeout is +generous. + +Run explicitly with:: + + MLXCEL_BIN=/path/to/mlxcel pytest python/tests/test_e2e.py -m e2e +""" + +from __future__ import annotations + +import os + +import pytest + +import mlxcel + +pytestmark = pytest.mark.e2e + +MLXCEL_BIN = os.environ.get("MLXCEL_BIN") +MODEL = os.environ.get("MLXCEL_E2E_MODEL", "mlx-community/Qwen3-0.6B-4bit") + +skip_no_bin = pytest.mark.skipif( + not MLXCEL_BIN, reason="MLXCEL_BIN not set; skipping real-binary e2e test" +) + + +@skip_no_bin +def test_managed_generate_returns_text() -> None: + with mlxcel.LLM(MODEL, binary=MLXCEL_BIN, startup_timeout=900.0) as llm: + assert llm.model + text = llm.generate("The capital of France is", max_tokens=16, temperature=0.0) + assert isinstance(text, str) + assert text.strip() != "" + + +@skip_no_bin +def test_managed_chat_and_stream() -> None: + with mlxcel.LLM(MODEL, binary=MLXCEL_BIN, startup_timeout=900.0) as llm: + reply = llm.chat([{"role": "user", "content": "Say hello."}], max_tokens=16) + assert reply.strip() != "" + streamed = "".join(llm.stream("Count: 1 2 3", max_tokens=16)) + assert isinstance(streamed, str) diff --git a/python/tests/test_lifecycle.py b/python/tests/test_lifecycle.py new file mode 100644 index 00000000..b20d0c57 --- /dev/null +++ b/python/tests/test_lifecycle.py @@ -0,0 +1,195 @@ +"""Lifecycle tests: spawn the stdlib fake_server via ManagedServer. + +These exercise binary discovery, spawn, /health readiness polling, ready state, +stderr log capture, graceful shutdown, and the early-exit failure path, all +without the real Rust binary or a model. +""" + +from __future__ import annotations + +import logging +import os +import stat +import sys +import time +from pathlib import Path + +import httpx +import pytest + +import mlxcel +from mlxcel._server import ManagedServer +from mlxcel.errors import MlxcelServerError + +HERE = Path(__file__).resolve().parent +FAKE_SERVER = HERE / "fake_server.py" + + +@pytest.fixture() +def fake_binary(tmp_path: Path) -> str: + """An executable shim that forwards args to the stdlib fake server. + + ManagedServer invokes `` serve -m --host --port 0``. + The shim re-execs the fake server with the same argv tail so the fake server + sees the real flag layout the manager produces. + """ + shim = tmp_path / "mlxcel-shim" + shim.write_text( + "#!/usr/bin/env python3\n" + "import os, sys\n" + f"os.execv({sys.executable!r}, [{sys.executable!r}, {str(FAKE_SERVER)!r}] + sys.argv[1:])\n" + ) + shim.chmod(shim.stat().st_mode | stat.S_IEXEC | stat.S_IXGRP | stat.S_IXOTH) + return str(shim) + + +def _short_socket(tmp_path: Path) -> str: + # Keep well under the sun_path limit; tmp_path can be long, so use /tmp. + return f"/tmp/mlxcel-test-{os.getpid()}-{tmp_path.name[:6]}.sock" + + +@pytest.mark.skipif(sys.platform.startswith("win"), reason="UDS not supported on Windows") +def test_managed_uds_lifecycle( + fake_binary: str, tmp_path: Path, caplog: pytest.LogCaptureFixture +) -> None: + sock = _short_socket(tmp_path) + caplog.set_level(logging.INFO, logger="mlxcel.server") + + server = ManagedServer( + "fake/model", + binary=fake_binary, + socket_path=sock, + extra_args=["--ready-after", "0.3"], + startup_timeout=15.0, + ) + server.start() + try: + assert server.uds_path == sock + assert os.path.exists(sock) + # /health must answer 200 after readiness. + transport = httpx.HTTPTransport(uds=sock) + with httpx.Client(transport=transport, base_url="http://mlxcel") as client: + resp = client.get("http://mlxcel/health") + assert resp.status_code == 200 + assert resp.json()["status"] == "ok" + # stderr log forwarding captured the startup line. + assert any("Starting mlxcel server" in rec.message for rec in caplog.records) + server.ensure_alive() + finally: + server.close() + + # Socket file removed and process gone after shutdown. + assert not os.path.exists(sock) + assert server._proc is not None and server._proc.poll() is not None + + +def test_managed_tcp_lifecycle(fake_binary: str) -> None: + server = ManagedServer( + "fake/model", + binary=fake_binary, + host="127.0.0.1", + extra_args=["--ready-after", "0.2"], + startup_timeout=15.0, + ) + server.start() + try: + assert server.port is not None and server.port > 0 + with httpx.Client(timeout=5.0) as client: + resp = client.get(f"http://127.0.0.1:{server.port}/health") + assert resp.status_code == 200 + finally: + server.close() + assert server._proc is not None and server._proc.poll() is not None + + +def test_early_exit_raises_with_stderr(fake_binary: str, tmp_path: Path) -> None: + sock = _short_socket(tmp_path) + server = ManagedServer( + "fake/model", + binary=fake_binary, + socket_path=sock, + extra_args=["--crash"], + startup_timeout=15.0, + ) + with pytest.raises(MlxcelServerError) as excinfo: + server.start() + # The captured stderr tail must be attached. + assert "simulated startup failure" in str(excinfo.value) + server.close() + + +def test_binary_not_found() -> None: + server = ManagedServer("fake/model", binary="/nonexistent/mlxcel-binary-xyz") + with pytest.raises(MlxcelServerError): + server.start() + + +def test_socket_path_too_long(tmp_path: Path) -> None: + long_path = "/tmp/" + ("x" * 200) + ".sock" + with pytest.raises(MlxcelServerError): + ManagedServer("fake/model", socket_path=long_path) + + +@pytest.mark.skipif(sys.platform.startswith("win"), reason="UDS not supported on Windows") +def test_llm_managed_mode_end_to_end(fake_binary: str, tmp_path: Path) -> None: + sock = _short_socket(tmp_path) + with mlxcel.LLM( + "fake/model", + binary=fake_binary, + socket=sock, + extra_args=["--ready-after", "0.2"], + startup_timeout=15.0, + ) as llm: + assert llm.model == "fake-model" + assert llm.generate("hello") == "hello world" + assert llm.chat([{"role": "user", "content": "hi"}]) == "hi there" + assert "".join(llm.stream("x")) == "hello world" + assert llm.tokenize("ab") == [ord("a"), ord("b")] + assert llm.detokenize([104, 105]) == "hi" + # After the context manager exits the socket is cleaned up. + assert not os.path.exists(sock) + + +@pytest.mark.skipif(sys.platform.startswith("win"), reason="UDS not supported on Windows") +def test_connect_mode_to_running_uds_server(fake_binary: str, tmp_path: Path) -> None: + # Start a server with the manager, then connect to it without managing it. + sock = _short_socket(tmp_path) + server = ManagedServer( + "fake/model", + binary=fake_binary, + socket_path=sock, + extra_args=["--ready-after", "0.2"], + startup_timeout=15.0, + ) + server.start() + try: + client = mlxcel.LLM(socket=sock) + try: + assert client.model == "fake-model" + assert client.generate("hi") == "hello world" + finally: + client.close() + # The connect-mode client must NOT have stopped the server it joined. + server.ensure_alive() + finally: + server.close() + + +def test_ensure_alive_after_crash(fake_binary: str, tmp_path: Path) -> None: + server = ManagedServer( + "fake/model", + binary=fake_binary, + host="127.0.0.1", + extra_args=["--ready-after", "0.2"], + startup_timeout=15.0, + ) + server.start() + proc = server._proc + assert proc is not None + proc.kill() + proc.wait(timeout=5.0) + # Give the manager's view a moment; ensure_alive should now raise. + time.sleep(0.05) + with pytest.raises(MlxcelServerError): + server.ensure_alive() + server.close() From 029188764c560991f49490bcd16c4caf560c47da Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Wed, 24 Jun 2026 06:33:32 +0900 Subject: [PATCH 2/4] fix(python): harden auth handling and response_format routing Pass the server API key through the LLAMA_API_KEY environment variable instead of argv so it is not exposed via ps or /proc//cmdline, and attach a Bearer header on the native /tokenize and /detokenize posts that bypass the OpenAI SDK auth injection. Route response_format through extra_body for the plain completions endpoint (the SDK rejects it as a top-level field there) while keeping it top-level for chat. Add regression tests for auth-header presence and absence, response_format routing, and the API key staying out of argv. --- python/src/mlxcel/_async_client.py | 19 +++++-- python/src/mlxcel/_client.py | 19 +++++-- python/src/mlxcel/_sampling.py | 27 ++++++++- python/src/mlxcel/_server.py | 17 +++++- python/tests/test_client_mock.py | 89 ++++++++++++++++++++++++++++++ python/tests/test_lifecycle.py | 44 +++++++++++++++ 6 files changed, 200 insertions(+), 15 deletions(-) diff --git a/python/src/mlxcel/_async_client.py b/python/src/mlxcel/_async_client.py index 7e0c8df2..11770132 100644 --- a/python/src/mlxcel/_async_client.py +++ b/python/src/mlxcel/_async_client.py @@ -61,6 +61,10 @@ def __init__( self._http_client: Optional[httpx.AsyncClient] = None self._model: Optional[str] = None self._closed = False + # Resolved API key, used to authorize native (/tokenize, /detokenize) + # routes that bypass the OpenAI SDK's own Authorization injection. The + # empty-or-None case stays unauthenticated to preserve the no-auth path. + self._api_key = api_key timeout = timeout or DEFAULT_TIMEOUT model_override = server_kwargs.pop("model", None) @@ -181,7 +185,7 @@ async def generate(self, prompt: str, **sampling: Any) -> str: """Generate a completion for ``prompt`` and return the text.""" self._ensure_alive() model = await self._resolve_model() - params = build_params(sampling) + params = build_params(sampling, chat=False) resp = await self._client.completions.create(model=model, prompt=prompt, **params) return resp.choices[0].text or "" @@ -189,7 +193,7 @@ async def stream(self, prompt: str, **sampling: Any) -> AsyncIterator[str]: """Stream completion text deltas for ``prompt``.""" self._ensure_alive() model = await self._resolve_model() - params = build_params(sampling) + params = build_params(sampling, chat=False) stream = cast( "AsyncStream[Completion]", await self._client.completions.create( @@ -204,7 +208,7 @@ async def chat(self, messages: ChatMessages, **sampling: Any) -> str: """Run a chat completion and return the assistant message content.""" self._ensure_alive() model = await self._resolve_model() - params = build_params(sampling) + params = build_params(sampling, chat=True) resp = await self._client.chat.completions.create( model=model, messages=as_messages(messages), **params ) @@ -214,7 +218,7 @@ async def chat_stream(self, messages: ChatMessages, **sampling: Any) -> AsyncIte """Stream chat completion content deltas.""" self._ensure_alive() model = await self._resolve_model() - params = build_params(sampling) + params = build_params(sampling, chat=True) stream = cast( "AsyncStream[ChatCompletionChunk]", await self._client.chat.completions.create( @@ -248,7 +252,12 @@ async def detokenize(self, tokens: Iterable[int]) -> str: async def _raw_post(self, path: str, json: Dict[str, Any]) -> Dict[str, Any]: assert self._http_client is not None url = native_base_url(self._base_url) + path - response = await self._http_client.post(url, json=json) + # The OpenAI SDK injects Authorization on /v1/* routes, but these native + # routes go through the bare httpx client, so add the bearer token here + # when a key was supplied. With no key, omit the header to keep the + # no-auth path working against servers started without --api-key. + headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None + response = await self._http_client.post(url, json=json, headers=headers) response.raise_for_status() result: Dict[str, Any] = response.json() return result diff --git a/python/src/mlxcel/_client.py b/python/src/mlxcel/_client.py index 885002da..b4f6702f 100644 --- a/python/src/mlxcel/_client.py +++ b/python/src/mlxcel/_client.py @@ -73,6 +73,10 @@ def __init__( self._http_client: Optional[httpx.Client] = None self._model: Optional[str] = None self._closed = False + # Resolved API key, used to authorize native (/tokenize, /detokenize) + # routes that bypass the OpenAI SDK's own Authorization injection. The + # empty-or-None case stays unauthenticated to preserve the no-auth path. + self._api_key = api_key timeout = timeout or DEFAULT_TIMEOUT model_override = server_kwargs.pop("model", None) @@ -180,14 +184,14 @@ def openai_client(self) -> OpenAI: def generate(self, prompt: str, **sampling: Any) -> str: """Generate a completion for ``prompt`` and return the text.""" self._ensure_alive() - params = build_params(sampling) + params = build_params(sampling, chat=False) resp = self._client.completions.create(model=self.model, prompt=prompt, **params) return resp.choices[0].text or "" def stream(self, prompt: str, **sampling: Any) -> Iterator[str]: """Stream completion text deltas for ``prompt``.""" self._ensure_alive() - params = build_params(sampling) + params = build_params(sampling, chat=False) stream = cast( "Stream[Completion]", self._client.completions.create(model=self.model, prompt=prompt, stream=True, **params), @@ -199,7 +203,7 @@ def stream(self, prompt: str, **sampling: Any) -> Iterator[str]: def chat(self, messages: ChatMessages, **sampling: Any) -> str: """Run a chat completion and return the assistant message content.""" self._ensure_alive() - params = build_params(sampling) + params = build_params(sampling, chat=True) resp = self._client.chat.completions.create( model=self.model, messages=as_messages(messages), **params ) @@ -208,7 +212,7 @@ def chat(self, messages: ChatMessages, **sampling: Any) -> str: def chat_stream(self, messages: ChatMessages, **sampling: Any) -> Iterator[str]: """Stream chat completion content deltas.""" self._ensure_alive() - params = build_params(sampling) + params = build_params(sampling, chat=True) stream = cast( "Stream[ChatCompletionChunk]", self._client.chat.completions.create( @@ -241,7 +245,12 @@ def detokenize(self, tokens: Iterable[int]) -> str: def _raw_post(self, path: str, json: Dict[str, Any]) -> Dict[str, Any]: assert self._http_client is not None url = native_base_url(self._base_url) + path - response = self._http_client.post(url, json=json) + # The OpenAI SDK injects Authorization on /v1/* routes, but these native + # routes go through the bare httpx client, so add the bearer token here + # when a key was supplied. With no key, omit the header to keep the + # no-auth path working against servers started without --api-key. + headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None + response = self._http_client.post(url, json=json, headers=headers) response.raise_for_status() result: Dict[str, Any] = response.json() return result diff --git a/python/src/mlxcel/_sampling.py b/python/src/mlxcel/_sampling.py index 0b5b054d..0210ba28 100644 --- a/python/src/mlxcel/_sampling.py +++ b/python/src/mlxcel/_sampling.py @@ -12,7 +12,8 @@ from typing import Any, Dict, Tuple -# Sampling kwargs that map directly onto native OpenAI request fields. +# Sampling kwargs accepted as top-level fields by BOTH the completions and the +# chat-completions endpoints. _OPENAI_FIELDS: Tuple[str, ...] = ( "max_tokens", "temperature", @@ -25,9 +26,14 @@ "n", "logprobs", "top_logprobs", - "response_format", ) +# Fields the OpenAI SDK accepts only on ``chat.completions.create``. For the +# plain ``completions.create`` endpoint the SDK raises ``TypeError`` on these, +# so they are routed through ``extra_body`` instead (the mlxcel/llguidance +# server reads constrained-decoding settings from the raw request body). +_CHAT_ONLY_FIELDS: Tuple[str, ...] = ("response_format",) + # Server-specific sampling knobs forwarded via ``extra_body``. These are not in # the OpenAI schema; the server reads them from the raw request body. _EXTRA_BODY_FIELDS: Tuple[str, ...] = ( @@ -42,7 +48,7 @@ ) -def build_params(kwargs: Dict[str, Any]) -> Dict[str, Any]: +def build_params(kwargs: Dict[str, Any], *, chat: bool = True) -> Dict[str, Any]: """Split caller kwargs into OpenAI request fields plus an ``extra_body`` payload. Known OpenAI fields are passed through as top-level request parameters. @@ -52,8 +58,16 @@ def build_params(kwargs: Dict[str, Any]) -> Dict[str, Any]: remaining unknown kwargs are also routed into ``extra_body`` rather than silently dropped, since the OpenAI SDK rejects unexpected top-level keys. + Chat-only fields (see :data:`_CHAT_ONLY_FIELDS`, e.g. ``response_format``) + are valid top-level parameters for ``chat.completions.create`` but not for + ``completions.create``. When ``chat`` is False they are routed through + ``extra_body`` so the server still receives them while the SDK does not + raise ``TypeError`` on an unexpected keyword. + Args: kwargs: Raw keyword arguments from a generate/chat call. + chat: True when building params for ``chat.completions.create``; False + when building for the plain ``completions.create`` endpoint. Returns: A new dict suitable for splatting into the OpenAI ``create`` call. @@ -69,6 +83,13 @@ def build_params(kwargs: Dict[str, Any]) -> Dict[str, Any]: continue if key in _OPENAI_FIELDS: params[key] = value + elif key in _CHAT_ONLY_FIELDS: + # Top-level for chat; the completions endpoint only accepts it in + # the raw body, so route it through extra_body there. + if chat: + params[key] = value + else: + extra_body[key] = value else: # Recognized server knob or an unknown key: send it in the body so # the server can interpret it instead of the SDK rejecting it. diff --git a/python/src/mlxcel/_server.py b/python/src/mlxcel/_server.py index 6196073c..700db365 100644 --- a/python/src/mlxcel/_server.py +++ b/python/src/mlxcel/_server.py @@ -171,6 +171,13 @@ def _health_url(self) -> str: return f"http://{self.host}:{self.port}/health" def _build_command(self) -> List[str]: + """Build the ``mlxcel serve`` argv. + + The API key is intentionally NOT placed on the command line: argv is + world-readable via ``ps`` / ``/proc//cmdline`` on multi-user hosts. + It is passed to the child through the ``LLAMA_API_KEY`` environment + variable in :meth:`start` instead. + """ binary = _find_binary(self._binary_arg) cmd: List[str] = [binary, "serve", "-m", self.model] @@ -181,8 +188,6 @@ def _build_command(self) -> List[str]: assert self.host is not None and self.port is not None cmd += ["--host", self.host, "--port", str(self.port)] - if self.api_key: - cmd += ["--api-key", self.api_key] if self._ctx_size is not None: cmd += ["--ctx-size", str(self._ctx_size)] if self._n_predict is not None: @@ -209,8 +214,15 @@ def start(self) -> None: raise MlxcelServerError("Server already started.") cmd = self._build_command() + # Safe to log: the API key is passed via the environment, not argv. logger.debug("Launching mlxcel server: %s", " ".join(cmd)) + # Pass the API key through the environment so it never appears in argv + # (visible via ps / /proc//cmdline) or in the launch log above. + env = os.environ.copy() + if self.api_key: + env["LLAMA_API_KEY"] = self.api_key + # Remove a stale socket file so bind does not fail on EADDRINUSE. if self.uds_path is not None and os.path.exists(self.uds_path): try: @@ -225,6 +237,7 @@ def start(self) -> None: stderr=subprocess.PIPE, text=True, bufsize=1, + env=env, ) except OSError as exc: raise MlxcelServerError(f"Failed to launch mlxcel server: {exc}") from exc diff --git a/python/tests/test_client_mock.py b/python/tests/test_client_mock.py index 8558dba7..6148b548 100644 --- a/python/tests/test_client_mock.py +++ b/python/tests/test_client_mock.py @@ -16,6 +16,9 @@ # Captured request bodies, keyed by path, so tests can assert sampling mapping. _LAST_BODY: Dict[str, Dict[str, Any]] = {} +# Captured request headers, keyed by path, so tests can assert auth behavior. +_LAST_HEADERS: Dict[str, Dict[str, str]] = {} + def _models_response() -> httpx.Response: return httpx.Response( @@ -75,6 +78,7 @@ def _handler(request: httpx.Request) -> httpx.Response: except json.JSONDecodeError: body = {} _LAST_BODY[path] = body + _LAST_HEADERS[path] = dict(request.headers) stream = bool(body.get("stream")) if path == "/v1/models": @@ -120,6 +124,7 @@ def _error_handler(request: httpx.Request) -> httpx.Response: @pytest.fixture() def llm() -> Any: _LAST_BODY.clear() + _LAST_HEADERS.clear() transport = httpx.MockTransport(_handler) client = mlxcel.LLM(transport=transport) yield client @@ -185,6 +190,42 @@ def test_response_format_passthrough(llm: Any) -> None: assert body["response_format"] == schema +def test_response_format_on_generate_uses_extra_body(llm: Any) -> None: + # response_format is chat-only for the OpenAI SDK; on the completions + # endpoint it must be routed through extra_body (no TypeError) and still + # reach the server inside the request body. + schema = {"type": "json_object"} + assert llm.generate("hi", response_format=schema) == "generated text" + body = _LAST_BODY["/v1/completions"] + assert body["response_format"] == schema + + +# -- native-route authorization (regression: missing Bearer on /tokenize) ---- + + +def test_native_routes_carry_bearer_when_api_key_set() -> None: + _LAST_BODY.clear() + _LAST_HEADERS.clear() + transport = httpx.MockTransport(_handler) + client = mlxcel.LLM(transport=transport, api_key="secret") + try: + client.tokenize("hi") + client.detokenize([1, 2, 3]) + finally: + client.close() + assert _LAST_HEADERS["/tokenize"].get("authorization") == "Bearer secret" + assert _LAST_HEADERS["/detokenize"].get("authorization") == "Bearer secret" + + +def test_native_routes_omit_auth_without_api_key(llm: Any) -> None: + # The default fixture client has no api_key; the no-auth path must stay + # intact so servers started without --api-key keep working. + llm.tokenize("hi") + llm.detokenize([1, 2, 3]) + assert "authorization" not in _LAST_HEADERS["/tokenize"] + assert "authorization" not in _LAST_HEADERS["/detokenize"] + + def test_openai_client_escape_hatch(llm: Any) -> None: from openai import OpenAI @@ -244,6 +285,18 @@ def test_build_params_drops_none() -> None: assert params["temperature"] == 0.0 +def test_build_params_response_format_routing() -> None: + schema = {"type": "json_object"} + # Chat: top-level field. + chat_params = build_params({"response_format": schema}, chat=True) + assert chat_params["response_format"] == schema + assert "extra_body" not in chat_params + # Completions: routed through extra_body so completions.create won't raise. + cmpl_params = build_params({"response_format": schema}, chat=False) + assert "response_format" not in cmpl_params + assert cmpl_params["extra_body"]["response_format"] == schema + + # -- async ------------------------------------------------------------------ @@ -291,3 +344,39 @@ async def go() -> tuple[str, list[int], str]: assert streamed == "gen text" assert tokens == [1, 2, 3] assert decoded == "decoded" + + +def test_async_native_routes_carry_bearer_when_api_key_set() -> None: + _LAST_BODY.clear() + _LAST_HEADERS.clear() + transport = httpx.MockTransport(_handler) + + async def go() -> None: + client = mlxcel.AsyncLLM(transport=transport, api_key="secret") + try: + await client.tokenize("hi") + await client.detokenize([1, 2, 3]) + finally: + await client.close() + + _run(go()) + assert _LAST_HEADERS["/tokenize"].get("authorization") == "Bearer secret" + assert _LAST_HEADERS["/detokenize"].get("authorization") == "Bearer secret" + + +def test_async_native_routes_omit_auth_without_api_key() -> None: + _LAST_BODY.clear() + _LAST_HEADERS.clear() + transport = httpx.MockTransport(_handler) + + async def go() -> None: + client = mlxcel.AsyncLLM(transport=transport) + try: + await client.tokenize("hi") + await client.detokenize([1, 2, 3]) + finally: + await client.close() + + _run(go()) + assert "authorization" not in _LAST_HEADERS["/tokenize"] + assert "authorization" not in _LAST_HEADERS["/detokenize"] diff --git a/python/tests/test_lifecycle.py b/python/tests/test_lifecycle.py index b20d0c57..069a7358 100644 --- a/python/tests/test_lifecycle.py +++ b/python/tests/test_lifecycle.py @@ -175,6 +175,50 @@ def test_connect_mode_to_running_uds_server(fake_binary: str, tmp_path: Path) -> server.close() +def test_api_key_not_in_argv(fake_binary: str) -> None: + # Regression: the API key must never appear on the command line (visible via + # ps / /proc//cmdline). It is passed through LLAMA_API_KEY instead. + server = ManagedServer("fake/model", binary=fake_binary, host="127.0.0.1", api_key="secret") + cmd = server._build_command() + assert "--api-key" not in cmd + assert "secret" not in " ".join(cmd) + + +def test_api_key_passed_via_env( + fake_binary: str, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + # The key reaches the child through LLAMA_API_KEY, not argv. Capture what + # Popen receives without launching a real process. + captured_env = {} + captured_cmd = [] + + class _StubPopen: + def __init__(self, cmd, *args, **kwargs): + captured_env.update(kwargs.get("env") or {}) + captured_cmd.extend(cmd) + # Empty stderr stream so log forwarding starts and finishes cleanly. + self.stderr = iter(()) + + def poll(self): + # Report an immediate non-None exit so _wait_until_ready raises a + # clean MlxcelServerError instead of polling /health forever. + return 1 + + monkeypatch.setattr("mlxcel._server.subprocess.Popen", _StubPopen) + + sock = _short_socket(tmp_path) + server = ManagedServer("fake/model", binary=fake_binary, socket_path=sock, api_key="secret") + # start() raises because the stubbed process "exits" immediately; we only + # care that Popen was handed the key via env and not via argv. + with pytest.raises(MlxcelServerError): + server.start() + + assert captured_env.get("LLAMA_API_KEY") == "secret" + joined = " ".join(captured_cmd) + assert "secret" not in joined + assert "--api-key" not in joined + + def test_ensure_alive_after_crash(fake_binary: str, tmp_path: Path) -> None: server = ManagedServer( "fake/model", From 2218fe76bce72f9d3d25a83727cfcf267344f525 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Wed, 24 Jun 2026 06:39:42 +0900 Subject: [PATCH 3/4] fix(python): release client and managed server on constructor failure Wrap the resource-creating section of LLM.__init__ and AsyncLLM.__init__ so a failure after the http client is built (or the managed subprocess is spawned), for example an empty /v1/models discovery response, deterministically tears those resources down instead of leaking them until garbage collection. The sync client reuses its idempotent close(); the async client does synchronous cleanup because it cannot await in __init__. Add a best-effort __del__ to AsyncLLM mirroring the sync client so a never-awaited close() still stops the managed server and drops the async http client pool. await close() and async-with remain the correct API. --- python/src/mlxcel/_async_client.py | 65 +++++++++++++++++++++++------- python/src/mlxcel/_client.py | 40 +++++++++++------- 2 files changed, 77 insertions(+), 28 deletions(-) diff --git a/python/src/mlxcel/_async_client.py b/python/src/mlxcel/_async_client.py index 11770132..279c8d69 100644 --- a/python/src/mlxcel/_async_client.py +++ b/python/src/mlxcel/_async_client.py @@ -69,21 +69,42 @@ def __init__( model_override = server_kwargs.pop("model", None) - if not managed: - resolved_base = self._connect(base_url, socket, transport, timeout) - else: - assert model is not None - resolved_base = self._spawn( - model, binary, host, port, socket, api_key, startup_timeout, timeout, server_kwargs + try: + if not managed: + resolved_base = self._connect(base_url, socket, transport, timeout) + else: + assert model is not None + resolved_base = self._spawn( + model, + binary, + host, + port, + socket, + api_key, + startup_timeout, + timeout, + server_kwargs, + ) + + self._client = AsyncOpenAI( + base_url=resolved_base, + api_key=api_key or "-", + http_client=self._http_client, ) - - self._client = AsyncOpenAI( - base_url=resolved_base, - api_key=api_key or "-", - http_client=self._http_client, - ) - self._base_url = resolved_base - self._model_override = model_override + self._base_url = resolved_base + self._model_override = model_override + except BaseException: + # close() is a coroutine and cannot be awaited from __init__, so + # perform synchronous cleanup directly. The managed subprocess is + # the critical resource; stopping it here ensures deterministic + # reaping even when the caller never awaits close(). + if self._server is not None: + self._server.close() + # Drop the async http client reference. It has never been used so + # no active connections exist; the pool will be released when the + # object is collected. + self._http_client = None + raise # -- setup ------------------------------------------------------------- @@ -284,5 +305,21 @@ async def __aenter__(self) -> "AsyncLLM": async def __aexit__(self, *exc: object) -> None: await self.close() + def __del__(self) -> None: + # Best-effort synchronous cleanup at GC time. The async connection pool + # is released by dropping the reference; the managed subprocess is reaped + # via ManagedServer.close() which is synchronous. Do NOT attempt to run + # the event loop or await close() here: the loop may already be gone at + # interpreter shutdown. The correct API remains `await close()` or + # `async with`. + if self._closed: + return + try: + if self._server is not None: + self._server.close() + self._http_client = None + except Exception: + pass + __all__ = ["AsyncLLM"] diff --git a/python/src/mlxcel/_client.py b/python/src/mlxcel/_client.py index b4f6702f..e7ca831a 100644 --- a/python/src/mlxcel/_client.py +++ b/python/src/mlxcel/_client.py @@ -81,21 +81,33 @@ def __init__( model_override = server_kwargs.pop("model", None) - if not managed: - resolved_base = self._connect(base_url, socket, transport, timeout) - else: - assert model is not None - resolved_base = self._spawn( - model, binary, host, port, socket, api_key, startup_timeout, timeout, server_kwargs + try: + if not managed: + resolved_base = self._connect(base_url, socket, transport, timeout) + else: + assert model is not None + resolved_base = self._spawn( + model, + binary, + host, + port, + socket, + api_key, + startup_timeout, + timeout, + server_kwargs, + ) + + self._client = OpenAI( + base_url=resolved_base, + api_key=api_key or "-", + http_client=self._http_client, ) - - self._client = OpenAI( - base_url=resolved_base, - api_key=api_key or "-", - http_client=self._http_client, - ) - self._base_url = resolved_base - self._model = model_override or self._discover_model() + self._base_url = resolved_base + self._model = model_override or self._discover_model() + except BaseException: + self.close() + raise # -- setup ------------------------------------------------------------- From 2932d5fa6a039dacc2dbe2907d5992d3fc9dabf2 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Wed, 24 Jun 2026 06:47:31 +0900 Subject: [PATCH 4/4] fix(python): set _closed before is_managed() to guard __del__ on early failure AsyncLLM.__init__ called is_managed() before self._closed = False, so __del__ raised AttributeError when the constructor failed at argument validation (ambiguous-args path). Move _closed initialization first. Add tests for async chat_stream, models(), openai_client escape hatch, model property before resolution, and mode-validation errors on AsyncLLM. --- docs/python-client.md | 18 ++++++++ mkdocs.ko.yml | 1 + python/src/mlxcel/_async_client.py | 7 +-- python/tests/test_client_mock.py | 73 ++++++++++++++++++++++++++++++ 4 files changed, 96 insertions(+), 3 deletions(-) diff --git a/docs/python-client.md b/docs/python-client.md index ca0c663e..1569338b 100644 --- a/docs/python-client.md +++ b/docs/python-client.md @@ -64,6 +64,24 @@ mlxcel serve -m mlx-community/Qwen3-4B-4bit --host /tmp/mlxcel.sock --port 0 Passing both a model and a connect target raises `MlxcelError`, because the mode would be ambiguous. +## Security: multi-user hosts + +On a shared machine, the default socket path under `/tmp` is world-readable: any local user can connect to the server you spawned and send requests. If that matters for your deployment, pass an explicit `socket=` path under a directory only you can read, for example one under `$XDG_RUNTIME_DIR` (mode `0700`, owned by your uid): + +```python +import os, pathlib, mlxcel + +runtime_dir = pathlib.Path(os.environ.get("XDG_RUNTIME_DIR", f"/run/user/{os.getuid()}")) +runtime_dir.mkdir(mode=0o700, parents=True, exist_ok=True) + +with mlxcel.LLM("mlx-community/Qwen3-4B-4bit", socket=str(runtime_dir / "mlxcel.sock")) as llm: + print(llm.generate("hello")) +``` + +The CLI equivalent is `mlxcel serve --host "$XDG_RUNTIME_DIR/mlxcel.sock" --port 0`. + +On macOS, `$TMPDIR` already expands to a per-user path under `/var/folders`, so the default socket is private there. On Linux without an active login session, `$XDG_RUNTIME_DIR` may be absent; fall back to a `0700` subdirectory under your home directory if needed. + ## Streaming `stream` and `chat_stream` yield text deltas as they arrive. diff --git a/mkdocs.ko.yml b/mkdocs.ko.yml index 11f9e25e..95e32260 100644 --- a/mkdocs.ko.yml +++ b/mkdocs.ko.yml @@ -146,6 +146,7 @@ nav: - user-guide/index.md - CLI 명령어: user-guide/cli.md - 서버 모드: user-guide/server.md + - Python 클라이언트: user-guide/python-client.md - 비전-언어 모델: user-guide/vision-language-models.md - LoRA 어댑터: user-guide/lora.md - 추측적 디코딩: user-guide/speculative-decoding.md diff --git a/python/src/mlxcel/_async_client.py b/python/src/mlxcel/_async_client.py index 279c8d69..9eb7b491 100644 --- a/python/src/mlxcel/_async_client.py +++ b/python/src/mlxcel/_async_client.py @@ -55,18 +55,19 @@ def __init__( transport: Optional[httpx.AsyncBaseTransport] = None, **server_kwargs: Any, ) -> None: - managed = is_managed(model, base_url, socket, transport) - + # Set _closed before any call that can raise so __del__ never sees a + # missing attribute even when __init__ fails early (e.g. bad arg combo). + self._closed = False self._server: Optional[ManagedServer] = None self._http_client: Optional[httpx.AsyncClient] = None self._model: Optional[str] = None - self._closed = False # Resolved API key, used to authorize native (/tokenize, /detokenize) # routes that bypass the OpenAI SDK's own Authorization injection. The # empty-or-None case stays unauthenticated to preserve the no-auth path. self._api_key = api_key timeout = timeout or DEFAULT_TIMEOUT + managed = is_managed(model, base_url, socket, transport) model_override = server_kwargs.pop("model", None) try: diff --git a/python/tests/test_client_mock.py b/python/tests/test_client_mock.py index 6148b548..19abbb14 100644 --- a/python/tests/test_client_mock.py +++ b/python/tests/test_client_mock.py @@ -380,3 +380,76 @@ async def go() -> None: _run(go()) assert "authorization" not in _LAST_HEADERS["/tokenize"] assert "authorization" not in _LAST_HEADERS["/detokenize"] + + +def test_async_chat_stream() -> None: + transport = httpx.MockTransport(_handler) + + async def go() -> str: + client = mlxcel.AsyncLLM(transport=transport) + try: + chunks = [d async for d in client.chat_stream([{"role": "user", "content": "hi"}])] + return "".join(chunks) + finally: + await client.close() + + assert _run(go()) == "chat reply" + + +def test_async_models() -> None: + transport = httpx.MockTransport(_handler) + + async def go() -> list[str]: + client = mlxcel.AsyncLLM(transport=transport) + try: + return await client.models() + finally: + await client.close() + + assert _run(go()) == [MODEL_ID] + + +def test_async_openai_client_escape_hatch() -> None: + from openai import AsyncOpenAI + + transport = httpx.MockTransport(_handler) + + async def go() -> bool: + client = mlxcel.AsyncLLM(transport=transport) + try: + return isinstance(client.openai_client, AsyncOpenAI) + finally: + await client.close() + + assert _run(go()) + + +def test_async_model_property_raises_before_resolution() -> None: + transport = httpx.MockTransport(_handler) + + async def go() -> None: + client = mlxcel.AsyncLLM(transport=transport) + try: + # model property raises before any request has resolved the id + with pytest.raises(mlxcel.MlxcelError): + _ = client.model + finally: + await client.close() + + _run(go()) + + +def test_async_ambiguous_args_is_error() -> None: + async def go() -> None: + with pytest.raises(mlxcel.MlxcelError): + mlxcel.AsyncLLM("some-model", base_url="http://localhost:8080/v1") + + _run(go()) + + +def test_async_no_args_is_error() -> None: + async def go() -> None: + with pytest.raises(mlxcel.MlxcelError): + mlxcel.AsyncLLM() + + _run(go())