Skip to content

Commit 551e90e

Browse files
authored
examples/llm_server: add OpenAI-compatible LLM server (pytorch#20308)
Add the reusable OpenAI-compatible control plane and worker protocol for local ExecuTorch LLM serving under examples/llm_server. The Python side owns HTTP, chat templating, tool parsing, validation, session routing, warm-resume transcript splicing, and streaming. The C++ side provides the shared JSONL worker loop and warm-resume prefill planner that model-specific worker binaries can reuse while keeping model execution out of Python. This intentionally lands as one server PR rather than a staged extension/llm/server stack. Model workers live under examples/models and include the shared worker_loop.h; this PR does not reintroduce the postponed generic TextLLM worker. pytorch#20001
1 parent 3b77c08 commit 551e90e

40 files changed

Lines changed: 8678 additions & 0 deletions

.github/workflows/_llm_server.yml

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
name: LLM Server Tests
2+
3+
on:
4+
workflow_call:
5+
inputs:
6+
docker-image:
7+
description: Docker image to use for Linux tests.
8+
required: false
9+
type: string
10+
default: ci-image:executorch-ubuntu-22.04-clang12
11+
12+
jobs:
13+
linux:
14+
uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
15+
permissions:
16+
id-token: write
17+
contents: read
18+
with:
19+
runner: linux.2xlarge
20+
docker-image: ${{ inputs.docker-image }}
21+
submodules: recursive
22+
ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
23+
timeout: 60
24+
script: |
25+
set -eux
26+
27+
eval "$(conda shell.bash hook)"
28+
CONDA_ENV=$(conda env list --json | jq -r ".envs | .[-1]")
29+
conda activate "${CONDA_ENV}"
30+
31+
python -m pip install --progress-bar off \
32+
-r examples/llm_server/python/requirements.txt \
33+
httpx \
34+
pytest
35+
36+
export PYTHONPATH="$(dirname "${PWD}"):${PYTHONPATH:-}"
37+
export PYTHONDONTWRITEBYTECODE=1
38+
39+
python -m pytest -q examples/llm_server/python/tests
40+
41+
cmake -S . -B cmake-out \
42+
-DCMAKE_BUILD_TYPE=Release \
43+
-DCMAKE_INSTALL_PREFIX=cmake-out \
44+
-DEXECUTORCH_BUILD_EXTENSION_LLM=ON \
45+
-DEXECUTORCH_BUILD_EXTENSION_LLM_RUNNER=ON
46+
cmake --build cmake-out --target install -j "$(nproc)"
47+
48+
cmake -S examples/llm_server/cpp -B cmake-out/examples/llm_server/cpp \
49+
-DCMAKE_PREFIX_PATH="${PWD}/cmake-out"
50+
cmake --build cmake-out/examples/llm_server/cpp \
51+
--target test_worker_prefill_plan test_worker_loop \
52+
-j "$(nproc)"
53+
ctest --test-dir cmake-out/examples/llm_server/cpp --output-on-failure

.github/workflows/trunk.yml

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1045,6 +1045,20 @@ jobs:
10451045
build-tool: cmake
10461046
docker-image: ci-image:executorch-ubuntu-22.04-clang12
10471047

1048+
llm-server:
1049+
needs: [changed-files, run-decision]
1050+
if: |
1051+
contains(needs.changed-files.outputs.changed-files, 'examples/llm_server') ||
1052+
contains(needs.changed-files.outputs.changed-files, 'extension/llm/runner') ||
1053+
contains(needs.changed-files.outputs.changed-files, 'extension/llm/tokenizers') ||
1054+
contains(needs.changed-files.outputs.changed-files, '.github/workflows/_llm_server.yml') ||
1055+
contains(needs.changed-files.outputs.changed-files, '.github/workflows/trunk.yml') ||
1056+
needs.run-decision.outputs.is-full-run == 'true'
1057+
uses: ./.github/workflows/_llm_server.yml
1058+
permissions:
1059+
id-token: write
1060+
contents: read
1061+
10481062
test-models-windows:
10491063
needs: run-decision
10501064
if: |

examples/llm_server/README.md

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,99 @@
1+
# ExecuTorch LLM Server
2+
3+
OpenAI-compatible serving for ExecuTorch LLMs, so any OpenAI-compatible agent
4+
harness (pi, opencode, ...) can use ExecuTorch as a local backend.
5+
6+
```
7+
examples/llm_server/
8+
spec/ # language-neutral OpenAI contract ExecuTorch targets
9+
conformance/ # one test suite every language server must pass
10+
python/ # Python server implementation (current)
11+
# cpp/ # future: no-Python single-binary server
12+
```
13+
14+
**Which entry point:** `examples.llm_server.python.server` is the reusable
15+
OpenAI control plane for workers that accept the generic `--model_path` /
16+
`--tokenizer_path` flags. Model-specific examples can wrap the same control
17+
plane when they need extra flags or a different tool parser.
18+
19+
Why this layout: the OpenAI contract is identical across languages, so the
20+
**spec** and **conformance** suite are shared, and each language gets its own
21+
implementation directory. The real cross-language reuse comes from the C++
22+
`LLMEngine`/`LLMSession` primitives underneath, packaged as process-isolated
23+
**worker binaries** that any control plane drives over a small JSONL protocol.
24+
See `python/README.md` to run it.
25+
26+
Status: experimental, reliability-first and deliberately narrow. Implemented:
27+
`/health`, `/v1/models`, `/v1/chat/completions` (streaming + non-streaming),
28+
Hugging Face chat templates (`--hf-tokenizer`), `temperature` / `max_tokens` /
29+
`max_completion_tokens` / `stop`, Hermes tool calling by default
30+
(`<tool_call>...</tool_call>` JSON, complete calls only; model-specific launchers
31+
may select the Qwen XML format) with `tool_choice="none"`,
32+
structured API errors, and best-effort cancellation. One worker process with
33+
serialized execution; a worker can host isolated sessions on one weight load when its engine reports
34+
capacity > 1 (with warm append-only resume across turns). KV/prefix state lives inside the
35+
worker/session, not the control plane. Unsupported params (including `top_p`,
36+
`seed`, `n>1`, `reasoning_effort`, penalties, `logit_bias`, `response_format`,
37+
`logprobs`, and `tool_choice="required"`) are rejected with a structured 400
38+
rather than silently ignored. See `python/README.md` to run it and
39+
`spec/README.md` for the exact contract.
40+
41+
## Use from pi (or any OpenAI-compatible harness)
42+
43+
Point pi at the server to use ExecuTorch as a local backend for tool-use
44+
workflows. Launch the server:
45+
46+
```bash
47+
python -m executorch.examples.llm_server.python.server \
48+
--worker-bin <model-worker> \
49+
--model-path <model.pte> \
50+
--tokenizer-path <tokenizer.model-or-json> \
51+
--hf-tokenizer <hf-model-or-local-dir> \
52+
--model-id <model-id> \
53+
--host 127.0.0.1 \
54+
--port 8000
55+
```
56+
57+
Useful optional flags (full reference in `python/README.md`):
58+
59+
- `--no-think` — default `enable_thinking=false` for templates that support it
60+
(e.g. Qwen3-style).
61+
- `--max-context N` — reject over-long prompts cleanly; use the export-time
62+
context length.
63+
- `--allow-chatml-fallback` — approximate ChatML when the model has no HF
64+
`chat_template`; experimentation only, not recommended for reliable tool use.
65+
66+
Point pi at the server via `~/.pi/agent/models.json`:
67+
68+
```json
69+
{ "providers": { "executorch": {
70+
"baseUrl": "http://127.0.0.1:8000/v1", "api": "openai-completions",
71+
"apiKey": "x", "models": [ { "id": "<model-id>",
72+
"compat": { "sendSessionAffinityHeaders": true } } ] } } }
73+
```
74+
75+
Other OpenAI-compatible clients use their own schema — generically: base URL
76+
`http://127.0.0.1:8000/v1`, the model id you passed to `--model-id`, and a dummy
77+
API key if one is required.
78+
79+
Supported contract for pi:
80+
81+
- Endpoint `POST /v1/chat/completions`; streaming supported.
82+
- Tool calls: the model's Hermes-style `<tool_call>...</tool_call>` output is
83+
parsed and returned as OpenAI `tool_calls`. This generic server uses Hermes by
84+
default; a model-specific server may select the Qwen XML format.
85+
- `tool_choice`: only `"auto"`, `"none"`, or unset.
86+
- Rejected with a structured 400 (`unsupported_parameter`), not silently
87+
ignored: `tool_choice="required"` or specific-function forcing,
88+
`response_format` JSON/constrained output, `logprobs`, `top_p` other than
89+
`1.0`, and `seed`.
90+
91+
Reliability guidance:
92+
93+
- Use the model's real HF `chat_template` (`--hf-tokenizer`) for tool use, kept
94+
aligned with the exported tokenizer/model.
95+
- If tool calls come back as plain text, confirm the model is emitting the
96+
configured tool-call format's markers (Hermes for the generic server) and that
97+
`tools` were included in the request.
98+
- If a request fails with `unsupported_parameter`, remove or disable that
99+
OpenAI knob in your pi/client config.

examples/llm_server/__init__.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
# Copyright (c) Meta Platforms, Inc. and affiliates.
2+
# All rights reserved.
3+
#
4+
# This source code is licensed under the BSD-style license found in the
5+
# LICENSE file in the root directory of this source tree.
Lines changed: 207 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,207 @@
1+
# Copyright (c) Meta Platforms, Inc. and affiliates.
2+
# All rights reserved.
3+
#
4+
# This source code is licensed under the BSD-style license found in the
5+
# LICENSE file in the root directory of this source tree.
6+
7+
"""Language-neutral OpenAI-contract conformance tests.
8+
9+
Runs against any base URL (ExecuTorch, llama.cpp, mlx-lm, ...) so every server
10+
implementation is validated against one shared spec. Point it at a running
11+
server:
12+
13+
OPENAI_BASE_URL=http://127.0.0.1:8000/v1 pytest test_openai_contract.py
14+
15+
Skips automatically if no server is reachable.
16+
"""
17+
18+
import json
19+
import os
20+
import urllib.error
21+
import urllib.request
22+
23+
import pytest
24+
25+
BASE_URL = os.environ.get("OPENAI_BASE_URL", "http://127.0.0.1:8000/v1").rstrip("/")
26+
MODEL = os.environ.get("OPENAI_MODEL", "executorch")
27+
28+
29+
def _post(path: str, body: dict, stream: bool = False):
30+
req = urllib.request.Request(
31+
f"{BASE_URL}{path}",
32+
data=json.dumps(body).encode(),
33+
headers={"Content-Type": "application/json"},
34+
method="POST",
35+
)
36+
return urllib.request.urlopen(req, timeout=120)
37+
38+
39+
def _server_up() -> bool:
40+
try:
41+
urllib.request.urlopen(f"{BASE_URL}/models", timeout=5)
42+
return True
43+
except Exception:
44+
return False
45+
46+
47+
pytestmark = pytest.mark.skipif(
48+
not _server_up(), reason="no OpenAI server at OPENAI_BASE_URL"
49+
)
50+
51+
52+
def test_models_listing():
53+
with urllib.request.urlopen(f"{BASE_URL}/models", timeout=10) as r:
54+
data = json.loads(r.read())
55+
assert data["object"] == "list"
56+
assert any("id" in m for m in data["data"])
57+
58+
59+
def test_chat_completion_nonstreaming():
60+
body = {
61+
"model": MODEL,
62+
"messages": [{"role": "user", "content": "Say hello in one word."}],
63+
"max_tokens": 16,
64+
"temperature": 0.0,
65+
}
66+
with _post("/chat/completions", body) as r:
67+
data = json.loads(r.read())
68+
assert data["object"] == "chat.completion"
69+
assert data["choices"][0]["message"]["role"] == "assistant"
70+
assert isinstance(data["choices"][0]["message"]["content"], str)
71+
assert data["choices"][0]["finish_reason"] is not None
72+
73+
74+
def test_chat_completion_streaming():
75+
body = {
76+
"model": MODEL,
77+
"messages": [{"role": "user", "content": "Count to three."}],
78+
"max_tokens": 32,
79+
"stream": True,
80+
}
81+
saw_role = saw_content = saw_done = False
82+
with _post("/chat/completions", body, stream=True) as r:
83+
for raw in r:
84+
line = raw.decode().strip()
85+
if not line.startswith("data:"):
86+
continue
87+
payload = line[len("data:") :].strip()
88+
if payload == "[DONE]":
89+
saw_done = True
90+
break
91+
chunk = json.loads(payload)
92+
assert chunk["object"] == "chat.completion.chunk"
93+
delta = chunk["choices"][0]["delta"]
94+
saw_role = saw_role or delta.get("role") == "assistant"
95+
saw_content = saw_content or bool(delta.get("content"))
96+
assert saw_role and saw_content and saw_done
97+
98+
99+
def test_multibyte_streaming_integrity():
100+
# Byte-level BPE can split a multi-byte character across tokens; the stream
101+
# must reassemble it, not abort with a UTF-8 decode error.
102+
body = {
103+
"model": MODEL,
104+
"messages": [
105+
{"role": "user", "content": "Reply with exactly: 你好世界 🌍 café"}
106+
],
107+
"max_tokens": 32,
108+
"temperature": 0.0,
109+
"stream": True,
110+
}
111+
content, saw_done, saw_error = "", False, False
112+
with _post("/chat/completions", body, stream=True) as r:
113+
for raw in r:
114+
line = raw.decode().strip()
115+
if not line.startswith("data:"):
116+
continue
117+
payload = line[len("data:") :].strip()
118+
if payload == "[DONE]":
119+
saw_done = True
120+
break
121+
chunk = json.loads(payload)
122+
if "error" in chunk:
123+
saw_error = True
124+
content += (
125+
chunk["choices"][0]["delta"].get("content", "")
126+
if chunk.get("choices")
127+
else ""
128+
)
129+
assert saw_done and not saw_error
130+
assert isinstance(content, str) and content # reassembled, valid UTF-8
131+
132+
133+
def test_usage_chunk_in_stream():
134+
body = {
135+
"model": MODEL,
136+
"messages": [{"role": "user", "content": "Say hi."}],
137+
"max_tokens": 16,
138+
"stream": True,
139+
"stream_options": {"include_usage": True},
140+
}
141+
usage = None
142+
with _post("/chat/completions", body, stream=True) as r:
143+
for raw in r:
144+
line = raw.decode().strip()
145+
if not line.startswith("data:"):
146+
continue
147+
payload = line[len("data:") :].strip()
148+
if payload == "[DONE]":
149+
break
150+
chunk = json.loads(payload)
151+
if chunk.get("usage"):
152+
usage = chunk["usage"]
153+
assert usage is not None, "no usage chunk emitted with include_usage"
154+
assert usage["prompt_tokens"] > 0 and usage["completion_tokens"] > 0
155+
assert usage["total_tokens"] == usage["prompt_tokens"] + usage["completion_tokens"]
156+
157+
158+
WEATHER_TOOL = {
159+
"type": "function",
160+
"function": {
161+
"name": "get_weather",
162+
"description": "Get the current weather for a city.",
163+
"parameters": {
164+
"type": "object",
165+
"properties": {"city": {"type": "string"}},
166+
"required": ["city"],
167+
},
168+
},
169+
}
170+
171+
172+
def test_tool_call_response_shape():
173+
body = {
174+
"model": MODEL,
175+
"messages": [
176+
{"role": "user", "content": "What is the weather in Paris? Use the tool."}
177+
],
178+
"tools": [WEATHER_TOOL],
179+
"max_tokens": 128,
180+
"temperature": 0.0,
181+
}
182+
with _post("/chat/completions", body) as r:
183+
data = json.loads(r.read())
184+
calls = data["choices"][0]["message"].get("tool_calls")
185+
assert calls, "expected tool_calls in response"
186+
tc = calls[0]
187+
assert tc["type"] == "function"
188+
assert tc["id"]
189+
assert tc["function"]["name"] == "get_weather"
190+
json.loads(tc["function"]["arguments"]) # arguments is a JSON string
191+
assert data["choices"][0]["finish_reason"] == "tool_calls"
192+
193+
194+
def test_error_body_shape():
195+
# Over-long prompt -> structured 400 (OpenAI error envelope), not a 500/drop.
196+
body = {
197+
"model": MODEL,
198+
"messages": [{"role": "user", "content": "word " * 40000}],
199+
"max_tokens": 8,
200+
}
201+
try:
202+
_post("/chat/completions", body)
203+
raise AssertionError("expected an HTTP error for over-long prompt")
204+
except urllib.error.HTTPError as e:
205+
assert 400 <= e.code < 500
206+
err = json.loads(e.read())["error"]
207+
assert err["message"] and err["type"]

0 commit comments

Comments
 (0)